Add 1-node GB200 functional test for checkpoint load with load_optim=False

[asolergi-nv]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds a 1-node / 4-GPU GB200 functional test that loads a checkpoint with the optimizer state dropped (--no-load-optim) — the path post-training (SFT) and the MLPerf DeepSeek submission take.

Why

Every GB200 mr-github test today either trains from scratch or resumes with the optimizer state intact. A numerical regression that only appears when a checkpoint is loaded without it can pass the entire unit and functional suite and still land in main. This closes that gap.

Details

Derived from the 2-node Nemotron 3.5 Lightning nightly test (#6312), scaled to a single node:

• TEST_TYPE: frozen-start — loads the frozen checkpoint from ${CHECKPOINT_LOAD_PATH} in a single training run.
• --no-load-optim is what drops the optimizer state. --finetune resets the iteration counter so the run logs steps 1..20 and the golden values line up; without it the extractor pads steps 1..N with "nan".
• TP1 / PP1 / CP1 / EP4 on 4 GPUs. The 52-layer architecture is scaled to 14 layers (MEMEM*EMEMEM*E/*E/*E), preserving its 6:6:2 Mamba/MoE/attention ratio and both MTP depths.
• alltoall rather than the flex/HybridEP dispatcher. No GB200 1-node test uses HybridEP, and its allgather does not survive the evaluation loop that this test deliberately runs (HYBRID-EP ALLGATHER TIMEOUT, reproduced on 2xGB200).
• eval-interval 10 < train-iters 20, so validation runs during training rather than only at the end.
• the_pile shard00 + GPT-2 BPE, matching the other 1-node GB200 tests.
• The run loads a pre-staged checkpoint artifact from the CI artifacts store; the test does not generate it.

mtp_2 loss is added to the metric extractor and to CHECK_THRESHOLDS. The same two lines also appear in #6312, so whichever lands second will have a trivial conflict.

Validation

Run on OCI GB200 through the real run_ci_test.sh harness with N_REPEAT=5:

• 5/5 repeats passed against the golden values, exit 0, 793 s wall clock.
• Every compared metric green: lm loss, mtp_1 loss, mtp_2 loss, num-zeros, mem-allocated-bytes, mem-max-allocated-bytes.
• Two independent runs differ by 0.002% on lm loss, 0.005% on mtp_2 loss and 1.65% on num-zeros — well inside tolerance.
• Checkpoint load verified numerically: iteration 1 starts at loss 6.73 (the checkpoint's level) rather than ~10.9 (fresh init).

iteration-time is intentionally excluded from METRICS: its median varies ~3.6% run-to-run against a 5% tolerance, and 0 of the 58 existing GB200 1-node tests compare it.

Issue tracking

Linked issue: N/A

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.