[ckpt] fix: Reject Energon resume after DP changes

[yaoyu-33]

Summary

Reject Energon dataloader-state restoration when the saved and current data-parallel sizes differ.

Supported trigger and impact

Native distributed checkpoints can be resharded across a changed DP size. Energon rank-local state cannot be restored that way without redistribution, so Bridge documents that direct Energon restore requires the same DP size.

When a DP=2 checkpoint was resumed at DP=1, the existing guard saw rank 0's old file and restored it. The extra rank-1 state was silently dropped. Energon then rebuilt dataset sharding for the new world size, so continuing from only rank 0's old position could skip or repeat samples and change the training stream while checkpoint load reported success.

Expansion already failed when a current rank file was missing; shrinkage bypassed that asymmetric check.

Root cause and minimal fix

maybe_load_dataloader_state() validated only the current rank's filename. A completed Energon generation contains exactly one deterministic file per saved pure-DP rank, so the rank-file cardinality identifies the saved DP size.

Before deserializing state, the loader now counts those files through the existing local Path.glob or MSC glob path and compares the count with the current DP process-group size. Any mismatch raises with the existing fail-loud data-order rationale. The per-rank file check remains as a malformed-generation guard.

Regression evidence

An exact-function deterministic reproducer executes the maintained helper with a completed two-file generation and current DP size 1.

Fail before on d0debd0bd308f3b4aef6aa0904952d655cb31eb5:

uv run --no-sync python reproduce_energon_dp_shrink.py src/megatron/bridge/training/checkpointing.py
exit 1: DP-shrink restore did not fail; restore_calls=[{'position': 25}]

Pass after, including unchanged DP=1 local and MSC-shaped restore controls:

uv run --no-sync python reproduce_energon_dp_shrink.py src/megatron/bridge/training/checkpointing.py
exit 0

The focused repository unit test was also added. Local collection was attempted but the CPU workstation stops before collection on unavailable pinned GPU import dependencies; that environment failure is not claimed as passing evidence.

Additional validation:

uv run --no-sync ruff check src/megatron/bridge/training/checkpointing.py tests/unit_tests/training/test_checkpointing.py
passed

uv run --project "$BRIDGE_PROJECT" --no-sync pre-commit run --all-files
passed

git diff --check
passed

Scope

This patch changes only Energon dataloader-state restore validation and adds one focused CPU regression test. It does not redistribute dataloader state across DP sizes, change native model/optimizer resharding, add configuration or dependencies, or claim GPU, convergence, performance, model-quality, or full-suite validation.

[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 cc7016a

[claude]

Light review - LGTM

The DP-size guard in maybe_load_dataloader_state is a correct, fail-loud addition: it counts saved train_dataloader_dprank*.pt files and refuses to resume when that count differs from the current DP size, closing the case where a smaller current DP world would silently discard saved rank state (which the pre-existing per-rank missing-file check did not catch). Both the MSC (msc.glob) and local (Path.glob) branches are consistent with existing usage elsewhere in the file, and the new check sits ahead of the per-rank file check without breaking it.

Observations (non-blocking):

• The existing test_existing_dir_missing_file_raises (empty dir, 0 files) now trips the new mismatch branch first rather than the per-rank check. It still asserts on "data-parallel size" so it passes, but the message reads "saved for data-parallel size 0" - slightly odd wording for the empty-dir case. Not worth changing.
• Coverage is good for the DP-shrink direction (2 saved / 1 current). A DP-grow case with a genuine partial file set (1 saved file, current dp_size=2, this rank=1) would exercise the increase path more explicitly.

Suggested test cases

• No perf tests impacted.

[yaoyu-33]

/ok to test cc7016a

[yaoyu-33]

/ok to test 18a16ac

[yaoyu-33]

/ok to test 18a16ac