fix(diloco test): set enable_checkpointing explicitly when loading state

[gulsumgudukbay]

Description

test_diloco_checkpoint_saving_and_normal_resume builds a third config to verify normal (non-DiLoCo) full-state restoration. It sets load_full_state_path but relies on the base config file for enable_checkpointing, and config validation
rejects that combination: "You must set enable_checkpointing=True to load a checkpoint."

base.yml defaults the flag to true, so the test passes today. But get_test_config_path() resolves to decoupled_base_test.yml when DECOUPLE_GCLOUD=TRUE, and that file sets enable_checkpointing: false, so the
test fails validation there. The first config in the same test already passes the flag explicitly; this makes the third one do the same instead of depending on which base file happens to be selected.

Tests

Reproduce the failure on main and the fix on this branch with:

DECOUPLE_GCLOUD=TRUE python -m pytest \
  tests/integration/diloco_test.py -k test_diloco_checkpoint_saving_and_normal_resume

Verified on 4x AMD Instinct MI355X (ROCm 10, jax 0.11.1) with DECOUPLE_GCLOUD=TRUE: the test fails on main
with the validation error above and passes on this branch.

Checklist

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gemini-code-assist]

Code Review

This pull request updates the integration test test_diloco_checkpoint_saving_and_normal_resume in tests/integration/diloco_test.py by adding the enable_checkpointing=true configuration parameter. There are no review comments, and I have no feedback to provide.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!