[ckpt] fix: Restore checkpoints without Energon state

[yaoyu-33]

Problem

Selecting an older model checkpoint that predates Energon dataloader-state saving can fail if the same checkpoint root contains state for a newer generation.

For example, a supported checkpoint.ckpt_step=10 resume may select iter_0000010 while energon/iter_0000020 exists. The restore path checked only the shared energon root, then treated the missing energon/iter_0000010/train_dataloader_dprank000.pt as a data-parallel-size mismatch and aborted. This prevents otherwise valid older model-only checkpoints from resuming.

Root cause and fix

The loader did not distinguish two different conditions:

• the selected checkpoint has no Energon generation because it predates state saving;
• the selected generation exists but its current DP-rank file is missing.

Check the selected iter_N directory before resolving the DP-rank file. A missing selected generation now follows the existing pre-feature compatibility behavior and starts the dataloader fresh. A missing rank file inside an existing selected generation remains a hard error to avoid silently changing data order. The checkpointing documentation now describes this boundary.

Validation

Fail before, on the unmodified fbb7570cf7eec94fd2e6064454d84f7fad07fbfa base:

uv run python -m pytest -p no:cacheprovider tests/unit_tests/training/test_checkpointing.py::TestMaybeLoadDataloaderState::test_missing_selected_iteration_warns_and_skips -q --tb=short
1 failed: RuntimeError claimed a data-parallel-size change for an absent selected generation

Pass after:

uv run python -m pytest -p no:cacheprovider tests/unit_tests/training/test_checkpointing.py::TestMaybeLoadDataloaderState::test_missing_selected_iteration_warns_and_skips -q --tb=short
1 passed

uv run python -m pytest -p no:cacheprovider tests/unit_tests/training/test_checkpointing.py::TestRecordDataloaderStateDir tests/unit_tests/training/test_checkpointing.py::TestMaybeLoadDataloaderState tests/unit_tests/training/test_checkpointing.py::TestMaybeSaveDataloaderState -q --tb=short
24 passed

uv run pre-commit run --all-files
passed

git diff --check
passed

The adjacent set includes a negative control proving that an existing selected generation with a missing DP-rank file still raises. This is a focused CPU path; no GPU, distributed, full-suite, convergence, or performance validation was run.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 685fd32

[claude]

LGTM

The change correctly reorders the restore logic so the missing iter_N directory case returns early before computing dp_rank, while a present iter_N dir with a missing per-DP-rank file still raises loudly. Docstring, docs, and unit tests are updated consistently.

Suggested test cases

• test_missing_selected_iteration_warns_and_skips
• test_existing_dir_missing_file_raises
• test_missing_dir_warns_and_skips

No perf tests impacted.