[Example] Benchmark DCP on Nemotron Math v3

[ilml]

Summary

• Add a runnable fixed-CP versus Dynamic Context Parallelism benchmark using the normal pretrain_gpt.py training stack.
• Support both the built-in variable-length mock workload and nvidia/Nemotron-SFT-Math-v3 through the new variable-length dataset loader.
• Run both cases with the same model, global batch, seed, dataset, checkpoint, and --cuda-graph-impl none.
• Require complete iteration/timing records, align every iteration and consumed-sample count, compare loss trajectories, and report average, median, trimmed-mean, and samples-per-second throughput.
• Write a credential-safe reproducibility manifest with the commit, topology, software/GPU versions, data/checkpoint identity, and benchmark controls.

This draft provides the reproducible harness. It does not claim current Nemotron Math v3 loss-parity or throughput results yet; those require a complete GPU run from this rebuilt stack and an immutable dataset snapshot.

Compared cases

Case	Relevant flags
Fixed-CP baseline	--sequence-packing-scheduler dp_balanced --context-parallel-size N
Dynamic CP	--dynamic-context-parallel --sequence-packing-scheduler default_dynamic_cp --min-dynamic-context-parallel-size N --context-parallel-size N

Both cases use --use-varlen-dataset, --calculate-per-token-loss, --max-seqlen-per-dp-cp-rank, the single dataloader, and CUDA graphs disabled.

User-facing controls

Environment variable	Purpose
DATASET_PATH / TOKENIZER_MODEL	Select Nemotron Math v3, another supported dataset, or a local Parquet/JSON/JSONL snapshot and its chat-template tokenizer. If unset, the benchmark uses mock data.
LOAD_PATH	Loads the same Megatron distributed checkpoint into both cases; otherwise both start from the same seeded initialization.
GPUS_PER_NODE, TP_SIZE, CP_SIZE, PP_SIZE	Configure the topology. The current harness deliberately supports one node only.
MAX_SEQLEN_PER_DP_CP_RANK	Set the packed-token capacity per scheduler rank.
MIN_DYNAMIC_CONTEXT_PARALLEL_SIZE	Set the minimum DCP group size, subject to the DP x CP topology constraints.
TRAIN_ITERS / WARMUP_ITERS	Control run length and the timing warmup excluded from throughput statistics.
CASE_ORDER	Run baseline_first or dcp_first to expose order-dependent effects.
LOSS_ATOL / LOSS_RTOL / CHECK_LOSS_PARITY	Configure loss comparison tolerances and whether out-of-tolerance finite losses fail. Missing, duplicate, skipped, NaN, or sample-misaligned iterations always fail.
ANALYZE_ONLY=1	Recompute the report from existing logs without launching training or overwriting the manifest.

Stack and attribution

• Base: [Training] Add variable-length THD datasets #5681, which adds the variable-length dataset and Nemotron Math v3 loader.
• [Training] Add variable-length THD datasets #5681 is itself stacked on [Core] Add dynamic context-parallel THD scheduling #5679, which provides the DCP runtime and schedulers.
• The recipe incorporates the setup demonstrated in xiaoyao0115's training branch.

Current limitations

• One node only; multi-node log and TensorBoard collection is intentionally rejected.
• Dense models and the single dataloader only for DCP.
• Megatron FSDP and CUDA graphs are unsupported.
• Transformer Engine 2.9 or newer is required.
• The default prompt format uses full-sequence causal loss rather than assistant-only masking.

Validation

• bash -n examples/dynamic_context_parallel/benchmark_dcp.sh
• git diff --check
• Synthetic complete-log analysis covering timings, samples/second, alignment, and loss parity
• Failure-path checks for missing iterations, invalid DCP topology/capacity, and invalid Transformer Engine determinism settings
• Reproducibility-manifest smoke test
• Signed commit verified locally: 55095561c62188742de2a1051f281ac689843cc8

Real GPU training was not run in the available host environment because PyTorch is not installed there. The PR remains a draft pending a recorded Nemotron Math v3 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.

[ilml]

/ok to test 2ca3f72

[ilml]

/ok to test 9571b3a

[ilml]

/ok to test 514e7e2

[ilml]

/ok to test 1ca1090

[ilml]

/ok to test 5509556