fix(megatron): preserve diffusion resume state

[gphuang]

Summary

Supports Issue #220, Task 9 Option 6.

Preserves rank-local Energon state during Megatron diffusion resume and fails closed when a required state file is absent or malformed. Resets warmup RNG accounting and adds opt-in batch-fingerprint plus instantiated-linear-class audit markers so checkpoint continuity and MXFP4-to-FP8 routing can be proven before a full run.

Test plan

• Changed Primus modules and both new unit-test files compile
• tests/unit_tests/backends/megatron/test_dataloader_checkpoint.py passes in pinned v26.5
• tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py passes in pinned v26.5
• Combined pinned-image run: 25 tests passed
• Build the pinned Flux model and exercise the four-leg 5+5 resume smoke after review

[gphuang]

Note: audit markers were dropped on stdout

The batch-fingerprint and linear-class-census markers introduced in this PR were initially emitted with bare print(). On the training path, the run log preserves fd 2 and Python logging but not fd 1, so markers were discarded and the continuity audit reported zero markers on an otherwise healthy run (10 steps, finite loss, rc=0).

A probe that emitted the same token to stdout, stderr, and logging from both callsites confirmed this path behavior: logging and stderr were preserved, stdout was not. Routing both markers through the module logger moved a Flux 12B MXFP4 smoke from 0 markers to 8 census and 80 fingerprints, and the continuity gate passed. Pushed as 7a7e4f07 and 7e4810a5.

Two follow-ups for the next runner:

• The runner hardcodes --local-ranks-filter 0, so only rank 0 reaches the log while the audit needs fingerprints from all ranks. Because it appends $LOCAL_RANKS after its own flag and torchrun honors the last occurrence, LOCAL_RANKS='--local-ranks-filter 0,1,2,3,4,5,6,7' re-admits all ranks without patching the runner.
• Once the marker gate passed, the smoke reached first checkpoint save and hit an unrelated pre-existing failure: torch_dist save asserts on adaLN_modulation because Flux's 19 double-stream and 38 single-stream blocks register different global shapes under one stacked key. That blocks resume testing but is out of scope for this PR.

[gphuang]

Note: review follow-ups landed

Pushed 1a272bc4 to the PR branch with two review-gap follow-ups:

• Updated tests/unit_tests/backends/megatron/test_diffusion_audit_markers.py to assert PRIMUS_BATCH_FINGERPRINT via caplog (logger capture) instead of capsys stdout, matching the current implementation in forward_step.py.
• Added explicit CUDA runtime coverage in tests/unit_tests/backends/megatron/diffusion/test_flux_dist_checkpoint.py via test_torch_dist_save_load_round_trip_cuda (@pytest.mark.skipif(not torch.cuda.is_available(), reason=\"Requires CUDA\")) while keeping the existing CPU/CUDA fallback round-trip test.

Validation in this shell:

• python3 -m compileall on both updated test files passes.
• Full pytest remains blocked in this shell (No module named pytest, and this host also lacks torch).

Remaining non-blocking integration gaps are unchanged: GPU execution of the checkpoint round-trip path and the end-to-end stage1->control continuity smoke.

[gphuang]

Pushed commit 67e52597 to address the review note: _publish_and_capture now catches Exception instead of BaseException, so test failures are still captured without swallowing system-level exits.