From e7540cb753f91a4257c95fa726d7fe347e3f8721 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 28 Jul 2026 06:25:25 +0000 Subject: [PATCH 01/85] =?UTF-8?q?[=F0=9F=A4=96]:=20Update=20docs-versions?= =?UTF-8?q?=20after=20code-freeze=20for=20r0.6.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/versions1.json | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/versions1.json b/docs/versions1.json index 77f566734db..471dbd33202 100644 --- a/docs/versions1.json +++ b/docs/versions1.json @@ -3,6 +3,10 @@ "version": "nightly", "url": "https://docs.nvidia.com/nemo/megatron-bridge/nightly/" }, + { + "version": "0.6.0", + "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.6.0/" + }, { "name": "0.4.2 (latest) · 26.04.01", "version": "0.4.2", From 2b56a25bacf1ccc19d85417c476a00dde3c5771a Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 28 Jul 2026 06:25:26 +0000 Subject: [PATCH 02/85] Bump docs version to r0.6.0 --- docs/conf.py | 2 +- docs/project.json | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index b67b384da09..d40d47241fa 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -37,7 +37,7 @@ project = "Megatron Bridge" copyright = "2026, NVIDIA Corporation" author = "NVIDIA Corporation" -release = "nightly" +release = "0.6.0" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration diff --git a/docs/project.json b/docs/project.json index 240d43ad959..e07a4ab80c0 100644 --- a/docs/project.json +++ b/docs/project.json @@ -1,4 +1,4 @@ { - "name": "megatron-bridge", - "version": "nightly" + "version": "0.6.0", + "name": "megatron-bridge" } From 6c0ad5301390f598473b356bbb5497053e45b231 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 29 Jul 2026 22:59:04 +0200 Subject: [PATCH 03/85] cp: `Add SequentialMLP mappings for MTP heads in nemotron_h_bridge.py (5140)` into `r0.6.0` (#5152) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: NeMo Bot Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py b/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py index 0d8f5300212..0b67d8b9568 100644 --- a/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py +++ b/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py @@ -462,8 +462,12 @@ def mapping_registry(self) -> MegatronMappingRegistry: "mtp.layers.*.mtp_model_layer.layers.*.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.*.norm.weight", "mtp.layers.*.mtp_model_layer.layers.*.mlp.router.weight": "mtp.layers.*.mixer.gate.weight", "mtp.layers.*.mtp_model_layer.layers.*.mlp.router.expert_bias": "mtp.layers.*.mixer.gate.e_score_correction_bias", + # GroupedMLP (moe_grouped_gemm=True): expert weights are stored as weight0, weight1, ... "mtp.layers.*.mtp_model_layer.layers.*.mlp.experts.linear_fc1.weight*": "mtp.layers.*.mixer.experts.*.up_proj.weight", "mtp.layers.*.mtp_model_layer.layers.*.mlp.experts.linear_fc2.weight*": "mtp.layers.*.mixer.experts.*.down_proj.weight", + # SequentialMLP (moe_grouped_gemm=False): expert weights are stored per local_expert + "mtp.layers.*.mtp_model_layer.layers.*.mlp.experts.local_experts.*.linear_fc1.weight": "mtp.layers.*.mixer.experts.*.up_proj.weight", + "mtp.layers.*.mtp_model_layer.layers.*.mlp.experts.local_experts.*.linear_fc2.weight": "mtp.layers.*.mixer.experts.*.down_proj.weight", "mtp.layers.*.mtp_model_layer.layers.*.mlp.fc1_latent_proj.weight": "mtp.layers.*.mixer.fc1_latent_proj.weight", "mtp.layers.*.mtp_model_layer.layers.*.mlp.fc2_latent_proj.weight": "mtp.layers.*.mixer.fc2_latent_proj.weight", "mtp.layers.*.mtp_model_layer.layers.*.mlp.shared_experts.linear_fc1.weight": "mtp.layers.*.mixer.shared_experts.up_proj.weight", From 0720b4984cb3a858f956df9ce9bebcdb620209d8 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Thu, 30 Jul 2026 21:04:32 +0200 Subject: [PATCH 04/85] cp: `[test] fix: repair Nemotron-H CI regressions (5143)` into `r0.6.0` (#5146) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../recipes/test_nemotronh_recipes_pretrain.py | 1 + .../unit_tests/models/nemotronh/test_nemotron_h_bridge.py | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py b/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py index be96bef9e9b..c07e5549285 100644 --- a/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py +++ b/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py @@ -36,6 +36,7 @@ "num_moe_experts": 16, "moe_token_dispatcher_type": "alltoall", "moe_shared_expert_overlap": True, + "sequence_parallel": True, }, ), ] diff --git a/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py b/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py index 269ebbdae4e..44a3093b329 100644 --- a/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py +++ b/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py @@ -368,14 +368,14 @@ def test_provider_bridge_mtp_config(self, active_nemotronh_config_dict): @pytest.mark.parametrize("num_nextn_predict_layers", [None, 0]) def test_provider_bridge_disables_mtp_naturally( self, - nemotron_nano_v2_config_dict, + active_nemotronh_config_dict, num_nextn_predict_layers, ): """HF configs without enabled MTP produce a normal non-MTP provider.""" from types import SimpleNamespace config_dict = { - **nemotron_nano_v2_config_dict, + **active_nemotronh_config_dict, "n_routed_experts": 0, "mtp_hybrid_override_pattern": "*E", "keep_mtp_spec_in_bf16": True, @@ -393,14 +393,14 @@ def test_provider_bridge_disables_mtp_naturally( assert result.mtp_use_repeated_layer is False assert result.keep_mtp_spec_in_bf16 is False - def test_provider_bridge_rejects_incomplete_mtp_config(self, nemotron_nano_v2_config_dict): + def test_provider_bridge_rejects_incomplete_mtp_config(self, active_nemotronh_config_dict): """An enabled HF MTP config must describe its hybrid block.""" from types import SimpleNamespace mock_pretrained = Mock(spec=PreTrainedCausalLM) mock_pretrained.config = SimpleNamespace( **{ - **nemotron_nano_v2_config_dict, + **active_nemotronh_config_dict, "n_routed_experts": 0, "num_nextn_predict_layers": 1, } From 08acf5609e71b1390cba33912a30201195dcda96 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 31 Jul 2026 00:59:06 +0200 Subject: [PATCH 05/85] cp: Nemotron 3.5 Nano verification card and final-checkpoint refresh (#5124, #5199) into r0.6.0 (#5178) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- .../nemotron-3.5-nano/card.yaml | 719 ++++++++++++++++++ scripts/training/README.md | 10 +- scripts/training/setup_experiment.py | 23 +- .../create-model-verification-card/SKILL.md | 141 +++- .../scripts/validate_card.py | 245 ++++-- .../perf_recipes/nemotronh/h100/nemotronh.py | 4 + .../bridge/recipes/nemotronh/__init__.py | 6 +- .../recipes/nemotronh/gb200/__init__.py | 6 +- .../nemotronh/gb200/nemotron_3_nano.py | 65 +- .../recipes/nemotronh/h100/nemotron_3_nano.py | 11 + .../recipes/nemotronh/test_nemotron_3_nano.py | 6 + .../test_nemotron_3_5_nano_perf_recipes.py | 20 +- .../recipes/test_nemotronh_recipes.py | 134 +++- .../test_validate_model_verification_card.py | 278 +++++++ .../scripts/training/test_setup_experiment.py | 56 ++ 15 files changed, 1602 insertions(+), 122 deletions(-) create mode 100644 examples/model_verification_cards/nemotron-3.5-nano/card.yaml create mode 100644 tests/unit_tests/scripts/test_validate_model_verification_card.py diff --git a/examples/model_verification_cards/nemotron-3.5-nano/card.yaml b/examples/model_verification_cards/nemotron-3.5-nano/card.yaml new file mode 100644 index 00000000000..89a68a3ea3e --- /dev/null +++ b/examples/model_verification_cards/nemotron-3.5-nano/card.yaml @@ -0,0 +1,719 @@ +# Agent-readable model verification card for Nemotron 3.5 Nano. +# status: unverified | verified | unsupported | not_applicable + +title: nemotron_3_5_nano +summary: > + Performance scope: pretrain_performance.H100 and pretrain_performance.GB200 + use tuned canonical performance recipes; timing and throughput metrics from + functional training items remain sanity checks rather than optimized + performance results. The convergence pretraining recipes retain real + tokenization, real data, natural routing, safety checks, and checkpointing + while adopting the corresponding hardware performance recipe's execution + stack. Nemotron 3.5 Nano support verification covers conversion, inference, + and training on H100 and GB200. The + model.hf_id/--hf-model and model.hf_revision/--hf-revision values are + non-functional placeholders pending publication. +verification_index: + model_level: + verified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - megatron_to_hf_cpu + - megatron_to_hf_gpu + - manual_forward_pass + - inference + training: + H100: + verified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] + GB200: + verified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] + performance: + H100: verified + GB200: verified + fsdp: + GB200: verified +model: + hf_id: nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + hf_revision: PUBLIC_HF_REVISION_PENDING + architecture: NemotronHForCausalLM + min_transformers_version: "5.8.1" +verification_environment: + base_container: nvcr.io/nvidian/nemo:26.08.rc2 + bridge_commit: 619a2df1da9517b4408ea475b428bf2aba618c8f # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + bridge_commit: 7edd7c7f68bcdd2b30cbaff426ece2ee0da9deb5 # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron-path work/model-verification/nemotron-3.5-nano/cpu-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-23 + expected_result: > + The pinned checkpoint imports on CPU, creates a reloadable iter_0000000, + and preserves every MTP tensor for the paired CPU export. After export, + all 6,513 tensors and 32,913,266,240 parameters match the pinned source + exactly in name, shape, dtype, and value. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 --tp 1 --pp 1 --ep 8 --etp 1 + --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron-path work/model-verification/nemotron-3.5-nano/imported-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-23 + expected_result: > + The command exits successfully and creates a reloadable iter_0000000. + After the paired GPU export, all 6,513 tensors and 32,913,266,240 + parameters match the pinned source exactly in name, shape, dtype, and + value, including the MTP layers. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + bridge_commit: 7edd7c7f68bcdd2b30cbaff426ece2ee0da9deb5 # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron-path work/model-verification/nemotron-3.5-nano/cpu-megatron/iter_0000000 + --hf-path work/model-verification/nemotron-3.5-nano/cpu-hf-export + --torch-dtype bfloat16 + last_verified: 2026-07-23 + expected_result: > + Strict CPU export writes a native NemotronHForCausalLM checkpoint that + reloads with Transformers. All 6,513 tensors and 32,913,266,240 + parameters match the pinned source bitwise with maximum difference zero. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 --tp 1 --pp 1 --ep 8 --etp 1 + --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron-path work/model-verification/nemotron-3.5-nano/imported-megatron/iter_0000000 + --hf-path work/model-verification/nemotron-3.5-nano/hf-export + --torch-dtype bfloat16 --distributed-save + last_verified: 2026-07-23 + expected_result: > + Strict export exits successfully; all 6,513 tensors match the pinned + source bitwise with original dtypes preserved, maximum difference zero, + and Transformers reloads the output natively as NemotronHForCausalLM. + + manual_forward_pass: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron_model_path work/model-verification/nemotron-3.5-nano/imported-megatron/iter_0000000 + --tp 1 --pp 1 --ep 8 --etp 1 + --prompt "The capital of France is the city of" + last_verified: 2026-07-23 + expected_result: > + The pinned-revision one-step comparison exits successfully. Hugging Face + and Megatron next-token predictions match at token ID 6993 (" Paris"); + cosine similarity is 0.998581, above the 0.99 gate. Maximum and mean + absolute logit differences are 0.796875 and 0.116128, respectively. + + inference: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/hf_to_megatron_generate_text.py + --hf_model_path nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron_model_path work/model-verification/nemotron-3.5-nano/imported-megatron/iter_0000000 + --tp 1 --pp 1 --ep 8 --etp 1 + --prompt "The capital of France is" --max_new_tokens 2 + last_verified: 2026-07-23 + expected_result: > + Two independent runs produce byte-identical output with this exact + 2-token completion, including its leading space: " Paris.". + + pretrain: + H100: + status: verified + precision: bf16 + bridge_commit: 33ede31740aa44183a253a62a4bcf968f176e161 # pragma: allowlist secret + enabled_features: + moe_dispatcher: hybridep + context_parallel_size: 2 + cuda_graph: + implementation: transformer_engine + scopes: [mamba] + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_5_nano_pretrain_config + --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 20 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 + 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3.5-nano-h100/rp2 + rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=20 + validation.eval_iters=0 validation.eval_interval=0 + model.moe_router_force_load_balancing=false + checkpoint.load=null + --save_dir work/model-verification/nemotron-3.5-nano-h100/pretrain-reference + --save_interval 10 checkpoint.async_save=false + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 + last_verified: 2026-07-28 + metrics: + initial_loss: 12.18927 + final_loss: 7.239111 + last_10_steps_step_time_ms_avg: 36869.270 + # Real RP2 convergence-verification run, not a mock-data benchmark. + last_10_steps_model_tflops_per_gpu_avg: 200.180 + peak_allocated_memory_gib: 68.547 + peak_reserved_memory_gib: 73.893 + expected_result: > + On two complete eight-H100 nodes, the real-data RP2 CP2/p2p run + completes exactly 20 sequence-length-8,192, GBS/MBS 512/1 optimizer + steps with natural routing and finite LM/MTP-1/MTP-2 losses. LM loss is + 12.18927 -> 7.239111, steps 11-20 average 36,869.270 ms / 200.180 + TFLOP/s/GPU, no iteration is skipped or NaN, and complete 20-file, + 430-GiB checkpoints are written at steps 10 and 20. + + GB200: + status: verified + precision: bf16 + bridge_commit: edca4a7d868cd1214bf5292f28a515139470dd99 # pragma: allowlist secret + enabled_features: + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 + --additional-slurm-params segment=2 + --recipe nemotron_3_5_nano_pretrain_8k_config + --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 + 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3.5-nano-gb200/rp2 + rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + model.moe_router_force_load_balancing=false + checkpoint.load=null checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-reference + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 + last_verified: 2026-07-28 + metrics: + initial_loss: 12.171970 + final_loss: 5.798433 + last_10_steps_step_time_ms_avg: 21483.630 + # Real RP2 convergence-verification run, not a mock-data benchmark. + last_10_steps_model_tflops_per_gpu_avg: 686.920 + peak_allocated_memory_gib: 158.220 + peak_reserved_memory_gib: 162.290 + expected_result: > + On two complete four-GPU GB200 nodes in one segment-2 NVLink domain, + the sequence-length-8,192, GBS/MBS 512/2 real-data RP2 + convergence-verification run completes exactly 100 finite-loss + optimizer steps with natural routing, Transformer Engine fused cross + entropy, CUDA graphs disabled, finite LM/MTP-1/MTP-2 losses, and no + skipped or NaN iterations. LM loss is 12.171970 -> 5.798433; + real-data steps 91-100 average 21,483.630 ms / 686.920 TFLOP/s/GPU + with 158.220 GiB peak allocated and 162.290 GiB peak reserved memory. + Complete 15-file, 430-GiB checkpoints with distributed metadata are + written at steps 50 and 100. + + pretrain_fsdp: + GB200: + status: verified + variants: + bf16: + status: verified + precision: bf16 + bridge_commit: edca4a7d868cd1214bf5292f28a515139470dd99 # pragma: allowlist secret + enabled_features: + moe_dispatcher: hybridep + megatron_fsdp: optim_grads_params + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 + --additional-slurm-params segment=2 + --recipe nemotron_3_5_nano_pretrain_8k_fsdp_config + --mode pretrain --dataset megatron-indexed --seq_length 8192 + --max_steps 100 --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 + 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3.5-nano-gb200/rp2 + rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + model.moe_router_force_load_balancing=false + checkpoint.load=null checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-fsdp + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 + last_verified: 2026-07-28 + metrics: + initial_loss: 12.171970 + final_loss: 5.797082 + last_10_steps_step_time_ms_avg: 26665.830 + # Real RP2 convergence-verification run, not a mock-data benchmark. + last_10_steps_model_tflops_per_gpu_avg: 553.430 + peak_allocated_memory_gib: 160.760 + peak_reserved_memory_gib: 166.230 + expected_result: > + On two complete four-GPU GB200 nodes in one segment-2 NVLink + domain, the sequence-length-8,192, GBS/MBS 512/2 BF16 real-data + RP2 FSDP convergence-verification run completes exactly 100 + finite-loss optimizer steps with natural routing, Transformer + Engine fused cross entropy, active MTP-1/MTP-2 losses, and no + skipped or NaN iterations. LM loss is 12.171970 -> 5.797082; + real-data steps 91-100 average 26,665.830 ms / 553.430 + TFLOP/s/GPU with 160.760 GiB peak allocated and 166.230 GiB peak + reserved memory. Complete 14-file, 369-GiB FSDP checkpoints with + distributed metadata are written at steps 50 and 100. Do not use + this run to claim convergence neutrality against ordinary BF16 + pretraining: 188 of 300 shared LM/MTP loss values exceed the 1% + loss gate, with a maximum relative difference of 12.95%, despite + the similar final loss. + + fp8_mx: + status: verified + precision: fp8_mx + bridge_commit: 968ab25dd0e8b76801e757ee9976d9806af9c11a # pragma: allowlist secret + enabled_features: + moe_dispatcher: hybridep + megatron_fsdp: optim_grads_params + command: > + ./scripts/training/train.sh --wait --nodes 2 --gpus-per-node 4 + --additional-slurm-params segment=2 + --recipe nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config + --mode pretrain --max_steps 20 + last_verified: 2026-07-24 + metrics: + initial_loss: 12.19034 + final_loss: 3.913218 + last_10_steps_step_time_ms_avg: 13917.000 + # Mock-data performance run; not a convergence measurement. + last_10_steps_model_tflops_per_gpu_avg: 795.390 + peak_allocated_memory_gib: 169.54 + peak_reserved_memory_gib: 173.86 + expected_result: > + On two complete four-GPU GB200 nodes in one segment-2 NVLink + domain, the sequence-length-8,192 MXFP8 FSDP recipe completes + exactly 20 finite-loss optimizer steps with GBS/MBS 384/3, active + MTP-1 and MTP-2 losses, and no skipped or NaN iterations. LM loss + is 12.19034 -> 3.913218, and steps 11-20 average 13,917.000 ms / + 795.390 TFLOP/s/GPU with 169.54 GiB peak allocated and 173.86 GiB + peak reserved memory. This run must not be compared with the BF16 + real-data pretrain or BF16 FSDP runs: it uses MXFP8, mock data, + forced routing, GBS/MBS 384/3, and no CUDA graphs. + + sft: + H100: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_5_nano_sft_openmathinstruct2_packed_config + --mode sft --max_steps 100 + --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-h100/imported-megatron/iter_0000000 + --save_dir work/model-verification/nemotron-3.5-nano-h100/sft-checkpoints + last_verified: 2026-07-30 + metrics: + initial_loss: 0.4407026 + final_loss: 0.2756689 + last_10_steps_step_time_ms_avg: 9856.900 + last_10_steps_model_tflops_per_gpu_avg: 79.470 + expected_result: > + Packed OpenMathInstruct-2 full SFT completes 100 finite-loss steps with + active MTP gradients, LM loss 0.4407026 -> 0.2756689, finite MTP-1 + loss 0.7634735 -> 0.4076096, and no skipped or NaN iterations. Steps + 91-100 average 9,856.900 ms / 79.470 TFLOP/s/GPU, and the complete + reloadable full-model checkpoint is written at step 100. + + GB200: + status: verified + precision: bf16 + bridge_commit: edca4a7d868cd1214bf5292f28a515139470dd99 # pragma: allowlist secret + enabled_features: + sequence_packing: offline + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 + --additional-slurm-params segment=2 + --recipe nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config + --mode sft --max_steps 100 + --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-gb200/imported-megatron/iter_0000000 + --save_dir work/model-verification/nemotron-3.5-nano-gb200/sft-checkpoints + last_verified: 2026-07-30 + metrics: + initial_loss: 0.4407797 + final_loss: 0.2760494 + last_10_steps_step_time_ms_avg: 6697.760 + last_10_steps_model_tflops_per_gpu_avg: 260.770 + peak_allocated_memory_gib: 101.640 + peak_reserved_memory_gib: 104.100 + expected_result: > + On two complete four-GPU GB200 nodes in one segment-2 NVLink domain, + packed OpenMathInstruct-2 full SFT completes exactly 100 finite-loss + optimizer steps with active MTP gradients. LM loss is 0.4407797 -> + 0.2760494, MTP-1 loss is 0.7641921 -> 0.4087451, and MTP-2 loss is + 0.8651200 -> 0.4649736. Steps 91-100 average 6,697.760 ms / 260.770 + TFLOP/s/GPU with 101.640 GiB peak allocated and 104.100 GiB peak + reserved memory. No iteration is skipped or NaN, and the complete + 12-file reloadable full-model checkpoint is written at step 100. + + sft_export_inference: + H100: + status: verified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 8 --tp 2 --pp 1 --ep 8 --etp 1 + --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron-path work/model-verification/nemotron-3.5-nano-h100/sft-checkpoints/iter_0000100 + --hf-path work/model-verification/nemotron-3.5-nano-h100/sft-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/nemotron-3.5-nano-h100/sft-hf-export + --prompt "Solve briefly: If 3x + 5 = 20, what is x?" + --max-new-tokens 45 --chat-template --disable-thinking + last_verified: 2026-07-30 + expected_result: "The step-100 SFT checkpoint exports 6,513 tensors, including all 270 MTP tensors, to native Hugging Face format and reloads as NemotronHForCausalLM. The deterministic verifier performs one greedy generation and produces exactly 45 new tokens. The literal completion is: \"To solve for \u0024x\u0024, follow these steps:\n\n1. Subtract 5 from both sides of the equation:\n\\[ 3x + 5 - 5 = 20 - 5 \\Rightarrow 3x\"." + + GB200: + status: verified + precision: bf16 + bridge_commit: edca4a7d868cd1214bf5292f28a515139470dd99 # pragma: allowlist secret + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 4 --tp 1 --pp 1 --ep 8 --etp 1 + --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 + --hf-revision PUBLIC_HF_REVISION_PENDING + --megatron-path work/model-verification/nemotron-3.5-nano-gb200/sft-checkpoints/iter_0000100 + --hf-path work/model-verification/nemotron-3.5-nano-gb200/sft-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/nemotron-3.5-nano-gb200/sft-hf-export + --prompt "Solve briefly: If 3x + 5 = 20, what is x?" + --max-new-tokens 45 --chat-template --disable-thinking + last_verified: 2026-07-30 + expected_result: "The step-100 SFT checkpoint exports 6,513 tensors, including all 270 MTP tensors, to native Hugging Face format and reloads as NemotronHForCausalLM. The deterministic verifier performs one greedy generation and produces exactly 45 new tokens. The literal completion is: \"To solve for \u0024x\u0024, follow these steps:\n\n1. Subtract 5 from both sides of the equation:\n\\[ 3x + 5 - 5 = 20 - 5 \\Rightarrow 3x\"." + + sft_long_context: + H100: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_5_nano_sft_config + --mode sft --dataset openmathinstruct2 --max_steps 100 --seq_length 32768 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + --tensor_model_parallel_size 2 --context_parallel_size 2 + --expert_tensor_parallel_size 1 --expert_model_parallel_size 8 + --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-h100/imported-megatron/iter_0000000 + 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:32768,pad_seq_to_mult:4}' + model.sequence_parallel=true model.cp_comm_type=p2p + model.cross_entropy_loss_fusion=true model.calculate_per_token_loss=true + model.recompute_granularity=full model.recompute_modules=null + model.recompute_method=uniform model.recompute_num_layers=1 + model.moe_router_force_load_balancing=false ddp.average_in_collective=false + ddp.overlap_param_gather=false optimizer.overlap_param_gather=false + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null checkpoint.save_optim=false checkpoint.save_rng=false + checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-h100/long-sft-checkpoints + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-30 + metrics: + initial_loss: 0.4268321 + final_loss: 0.2649768 + last_10_steps_step_time_ms_avg: 40995.130 + last_10_steps_model_tflops_per_gpu_avg: 188.280 + expected_result: > + Packed OpenMathInstruct-2 CP2 SFT completes exactly 100 finite-loss 32K + steps with LM loss 0.4268321 -> 0.2649768 and MTP-1 loss 0.7375677 -> + 0.4385137. All 100 steps have finite LM/MTP losses and no skipped or + NaN iterations. Steps 91-100 average 40,995.130 ms / 188.280 + TFLOP/s/GPU. Complete step-50 and step-100 checkpoints each contain 20 + files, distributed metadata, saved run configuration and training + state, with latest marker 100. + + GB200: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 + --recipe nemotron_3_5_nano_sft_config + --mode sft --dataset openmathinstruct2 --max_steps 100 --seq_length 32768 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + --context_parallel_size 2 --expert_tensor_parallel_size 1 + --expert_model_parallel_size 8 + --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-gb200/imported-megatron/iter_0000000 + model.moe_flex_dispatcher_backend=hybridep + 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:32768,pad_seq_to_mult:4}' + model.cp_comm_type=a2a model.cross_entropy_loss_fusion=false + model.calculate_per_token_loss=true model.recompute_granularity=selective + 'model.recompute_modules=[moe,layernorm,core_attn,mlp]' + model.recompute_method=null model.recompute_num_layers=null + model.moe_router_force_load_balancing=false ddp.average_in_collective=false + ddp.overlap_param_gather=false optimizer.overlap_param_gather=false + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null checkpoint.save_optim=false checkpoint.save_rng=false + checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-gb200/long-sft-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-30 + metrics: + initial_loss: 0.4268995 + final_loss: 0.2649511 + last_10_steps_step_time_ms_avg: 100358.390 + last_10_steps_model_tflops_per_gpu_avg: 153.810 + expected_result: > + Packed OpenMathInstruct-2 CP2 SFT completes exactly 100 finite-loss 32K + steps with LM loss 0.4268995 -> 0.2649511 and MTP-1 loss 0.7376469 -> + 0.4384451. All 100 steps have finite LM/MTP losses and no skipped or + NaN iterations. The complete step-100 checkpoint has distributed + metadata, 12 files, a latest marker of 100, and all four metrics + recorded. + + peft: + H100: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 + --recipe nemotron_3_5_nano_peft_config + --mode lora --dataset openmathinstruct2 --max_steps 100 --seq_length 4096 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + --tensor_model_parallel_size 2 --expert_tensor_parallel_size 1 + --expert_model_parallel_size 8 + --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-h100/imported-megatron/iter_0000000 + 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:4096,pad_seq_to_mult:2}' + model.sequence_parallel=true + ddp.overlap_param_gather=false optimizer.overlap_param_gather=false + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-h100/peft-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-30 + metrics: + initial_loss: 0.4406879 + final_loss: 0.2856177 + last_10_steps_step_time_ms_avg: 16927.070 + last_10_steps_model_tflops_per_gpu_avg: 92.650 + expected_result: > + Packed rank-32, alpha-32, zero-dropout LoRA completes 100 finite-loss + steps with adapters on the MTP layers, LM loss 0.4406879 -> 0.2856177, + finite MTP-1 loss 0.7634883 -> 0.4294895, and no skipped or NaN + iterations. The complete adapter checkpoint is written at step 100 and + reloads over the pinned base checkpoint at step 100. + + GB200: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 + --recipe nemotron_3_5_nano_peft_config + --mode lora --dataset openmathinstruct2 --max_steps 100 --seq_length 4096 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + --tensor_model_parallel_size 2 --expert_tensor_parallel_size 1 + --expert_model_parallel_size 8 + --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-gb200/imported-megatron/iter_0000000 + model.moe_flex_dispatcher_backend=hybridep + 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:4096,pad_seq_to_mult:2}' + model.sequence_parallel=true + ddp.overlap_param_gather=false optimizer.overlap_param_gather=false + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-gb200/peft-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-30 + metrics: + initial_loss: 0.4405894 + final_loss: 0.2855471 + last_10_steps_step_time_ms_avg: 19163.540 + last_10_steps_model_tflops_per_gpu_avg: 82.080 + expected_result: > + Packed rank-32, alpha-32, zero-dropout LoRA completes 100 finite-loss + steps with 48 rank-local adapter attachment records covering MTP + attention, expert, and shared-expert layers, finite MTP losses, no + skipped or NaN iterations, and all four metrics recorded. The complete + adapter checkpoint is written at step 100 and reloads over the pinned + base checkpoint at step 100. + + checkpoint_resume: + H100: + status: verified + precision: bf16 + bridge_commit: 33ede31740aa44183a253a62a4bcf968f176e161 # pragma: allowlist secret + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_5_nano_pretrain_config + --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 20 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 + 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3.5-nano-h100/rp2 + rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=20 + validation.eval_iters=0 validation.eval_interval=0 + model.moe_router_force_load_balancing=false + --load_dir work/model-verification/nemotron-3.5-nano-h100/pretrain-reference + checkpoint.ckpt_step=10 checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-h100/pretrain-resumed + --save_interval 10 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 + last_verified: 2026-07-28 + metrics: + initial_loss: 7.981691 + final_loss: 7.238846 + last_10_steps_step_time_ms_avg: 56671.640 + last_10_steps_model_tflops_per_gpu_avg: 176.630 + resume_comparison: + reference_item: pretrain + sentinel_steps: [11, 20] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: true + expected_result: > + The continuation restores optimizer, scheduler, data-order, and RNG + state from step 10, begins at step 11, finishes at step 20 in a distinct + output root, and writes a complete step-20 checkpoint. Resumed losses + 7.981691 and 7.238846 match uninterrupted losses 7.981684 and 7.239111 + at steps 11 and 20 within the declared tolerance. All 30 shared + LM/MTP-1/MTP-2 values pass with a maximum relative difference of + 0.0851%, with no skipped or NaN iterations, and the 20-file, 430-GiB + step-20 checkpoint is complete. + + GB200: + status: verified + precision: bf16 + bridge_commit: edca4a7d868cd1214bf5292f28a515139470dd99 # pragma: allowlist secret + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 + --additional-slurm-params segment=2 + --recipe nemotron_3_5_nano_pretrain_8k_config + --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 + 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3.5-nano-gb200/rp2 + rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + model.moe_router_force_load_balancing=false + --load_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-reference + checkpoint.ckpt_step=50 checkpoint.async_save=false + --save_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-resumed + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 + last_verified: 2026-07-28 + metrics: + initial_loss: 6.166278 + final_loss: 5.796845 + last_10_steps_step_time_ms_avg: 22233.450 + last_10_steps_model_tflops_per_gpu_avg: 663.860 + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: true + expected_result: > + The continuation restores optimizer, scheduler, data-order, and RNG + state from the matched 8K step-50 reference, begins at step 51, and + finishes at step 100 in a distinct output root. Resumed losses + 6.166278 and 5.796845 match uninterrupted losses 6.166278 and 5.798433 + at steps 51 and 100. All 150 shared LM/MTP-1/MTP-2 values pass with a + maximum relative difference of 0.127978%, no iteration is skipped or + NaN, and the 15-file, 430-GiB step-100 checkpoint is complete. + + pretrain_performance: + H100: + status: verified + precision: bf16 + bridge_commit: 9b0c2dff42e6175e91d34cd2a504529f803ebde9 # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config + --max_steps 50 + last_verified: 2026-07-28 + metrics: + initial_loss: 12.16155 + final_loss: 0.01361997 + last_10_steps_step_time_ms_avg: 20826.370 + last_10_steps_model_tflops_per_gpu_avg: 354.300 + peak_allocated_memory_gib: 71.949 + peak_reserved_memory_gib: 74.631 + expected_result: > + On 16x H100, the exact mock-data BF16 performance recipe completes + exactly 50 steps with finite LM/MTP-1/MTP-2 losses and no skipped or + NaN iterations. Transformer Engine CUDA graphs cover attention and + Mamba scopes. Steps 41-50 average 20,826.370 ms and 354.300 model + TFLOP/s/GPU; peak rank-0 allocated/reserved memory is 71.949/74.631 + GiB. + + GB200: + status: verified + precision: bf16 + bridge_commit: 9b0c2dff42e6175e91d34cd2a504529f803ebde9 # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 + --recipe nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config + --max_steps 50 + last_verified: 2026-07-28 + metrics: + initial_loss: 12.18250 + final_loss: 0.03325584 + last_10_steps_step_time_ms_avg: 18236.060 + last_10_steps_model_tflops_per_gpu_avg: 809.220 + peak_allocated_memory_gib: 156.850 + peak_reserved_memory_gib: 161.260 + expected_result: > + On 8x GB200, the exact mock-data BF16 performance recipe completes + exactly 50 steps with finite LM/MTP-1/MTP-2 losses and no skipped or NaN + iterations. Transformer Engine CUDA graphs cover attention, Mamba, + MoE-router, and MoE-preprocess scopes. Steps 41-50 average 18,236.060 ms + and 809.220 model TFLOP/s/GPU; peak rank-0 allocated/reserved memory is + 156.850/161.260 GiB. diff --git a/scripts/training/README.md b/scripts/training/README.md index 00b76a86053..158c1c0c2e2 100644 --- a/scripts/training/README.md +++ b/scripts/training/README.md @@ -254,13 +254,17 @@ target cluster. For example, a Pyxis/Enroot cluster may use: The `=` form is required when `ARG` begins with `-`. +Pass target-specific sbatch settings through +`--additional-slurm-params 'KEY=VALUE;OTHER_KEY=OTHER_VALUE'`. These values +populate the executor's `additional_parameters` and are not forwarded to the +training command or its srun step. + Use `-lmc FREQ`, `--peak-mem-clk FREQ`, or `--peak_mem_clk FREQ` to lock the GPU memory clock to a fixed frequency in MHz once per node before training. VR200 benchmark recipes default to `4752`; other recipes leave the memory clock unlocked. Pass `-lmc -1` to disable the VR200 default explicitly. -The compact launcher does not add rank-command prefixes or extra sbatch parameters such as `segment`. Configure those -through the target cluster integration, or use `scripts/performance/setup_experiment.py` when its compatibility NUMA -and segment policies are required. +The compact launcher does not add rank-command prefixes or infer target-specific +allocation and NUMA policies. Configure those explicitly for the target. The launcher submits the experiment in detached mode and returns after Slurm accepts the job. Inspect its state and logs with the cluster's normal `squeue`, `sacct`, and log-file workflow. diff --git a/scripts/training/setup_experiment.py b/scripts/training/setup_experiment.py index 0cb7e1b9995..c5ef6d000a4 100755 --- a/scripts/training/setup_experiment.py +++ b/scripts/training/setup_experiment.py @@ -75,6 +75,13 @@ def _build_parser() -> argparse.ArgumentParser: execution.add_argument("--partition", default=os.environ.get("SLURM_PARTITION"), help="Slurm partition.") execution.add_argument("--time", default="04:00:00", help="Slurm time limit.") execution.add_argument("--gres", help="Optional Slurm GRES value.") + execution.add_argument( + "--additional-slurm-params", + "--additional_slurm_params", + type=_parse_additional_slurm_params, + default={}, + help="Additional sbatch parameters as semicolon-separated KEY=VALUE pairs.", + ) execution.add_argument( "--no-gpu-resource-request", action="store_true", @@ -163,6 +170,17 @@ def _parse_mounts(values: list[str]) -> list[str]: return mounts +def _parse_additional_slurm_params(value: str) -> dict[str, str]: + """Parse semicolon-separated Slurm executor parameters.""" + parameters: dict[str, str] = {} + for item in value.split(";"): + key, separator, parameter_value = item.partition("=") + if not separator or not key or not parameter_value: + raise argparse.ArgumentTypeError("--additional-slurm-params expects semicolon-separated KEY=VALUE pairs.") + parameters[key] = parameter_value + return parameters + + def _validate_args( args: argparse.Namespace, benchmark_metadata: BenchmarkRecipeMetadata | None = None, @@ -273,7 +291,10 @@ def _build_executor( # Keep Slurm control commands available to the batch script without # forwarding the host PATH into the training container. slurm_env_names = list(dict.fromkeys(["PATH", *env_names])) - executor.additional_parameters = {"export": ",".join(slurm_env_names)} + executor.additional_parameters = { + **args.additional_slurm_params, + "export": ",".join(slurm_env_names), + } executor.srun_args = srun_args return executor diff --git a/skills/create-model-verification-card/SKILL.md b/skills/create-model-verification-card/SKILL.md index 39a1a3160f0..af01ff5f02a 100644 --- a/skills/create-model-verification-card/SKILL.md +++ b/skills/create-model-verification-card/SKILL.md @@ -69,6 +69,40 @@ exported, omit the item instead of adding an unverified placeholder. Once a canonical recipe exists, keep the item in the card even if its run is still unverified. +Add `pretrain_fsdp` as an optional hardware-scoped item when a first-class +Megatron FSDP performance recipe exists for the exact model variant. Record +the FSDP result under `pretrain_fsdp.`. Keep it separate from +checkpoint-resume and tuned non-FSDP `pretrain_performance` results. Do not +embed a baseline, control, computed delta, or relative speedup under another +item; agents can derive valid comparisons from the standalone commands, +resolved recipes, and raw metrics. The FSDP item does not claim checkpoint +save/load unless a separate functional item records that evidence. + +When the same hardware has multiple first-class FSDP runs with different +precisions or convergence contracts, use one aggregate hardware container and +key its standalone leaves by precision: + +```yaml +pretrain_fsdp: + GB200: + status: verified + variants: + bf16: + status: verified + precision: bf16 + # complete standalone leaf + fp8_mx: + status: verified + precision: fp8_mx + # complete standalone leaf +``` + +The container status is `verified` only when every precision variant is +verified; otherwise it is `unverified`. Keep the ordinary direct hardware-leaf +shape when there is only one FSDP run. A precision variant repeats its +precision as a scalar so agents do not have to infer workload facts from a +mapping key. + A concrete `pretrain_performance.` leaf means a tuned canonical performance recipe exists for that hardware. Its item status states whether the card's benchmark run has been verified; an `unverified` leaf still records @@ -106,19 +140,19 @@ items: The hardware-scoped names are `pretrain`, `sft`, `sft_export_inference`, `sft_long_context`, `peft`, `checkpoint_resume`, and optional -`pretrain_performance`. Use canonical public accelerator identifiers such as -`H100`, `B200`, or `GB200`, never a private cluster name. The validator's -public-hardware allowlist is authoritative and must be updated when a new -accelerator target is introduced. The hardware key replaces the old `gpu_type` -field. Each hardware leaf is independent and must carry its own status plus the -command or commands, date, metrics, features, and optional commit override that -apply to that item. Dependencies resolve within the same hardware key: -`checkpoint_resume.H100` consumes `pretrain.H100`, and -`sft_export_inference.H100` consumes `sft.H100`. Never fall back across -hardware targets. Use the reserved key `all` only as the sole leaf for a -model-wide `unsupported` or `not_applicable` limitation. A terminal dependency -leaf still names its logical dependency but does not require a matching `all` -or concrete-hardware dependency leaf. +`pretrain_performance` and `pretrain_fsdp`. Use canonical public accelerator +identifiers such as `H100`, `B200`, or `GB200`, never a private cluster name. +The validator's public-hardware allowlist is authoritative and must be updated +when a new accelerator target is introduced. The hardware key replaces the old +`gpu_type` field. Each hardware leaf is independent and must carry its own +status plus the command or commands, date, metrics, features, and optional +commit override that apply to that item. Dependencies resolve within the same +hardware key: `checkpoint_resume.H100` consumes `pretrain.H100`, and +`sft_export_inference.H100` consumes `sft.H100`. Never fall back across hardware +targets. Use the reserved key `all` only as the sole leaf for a model-wide +`unsupported` or `not_applicable` limitation. A terminal dependency leaf still +names its logical dependency but does not require a matching `all` or +concrete-hardware dependency leaf. Use only `unverified`, `verified`, `unsupported`, or `not_applicable`. Do not add `smoke` or an evidence field. @@ -135,7 +169,9 @@ Use `model_level` for the six direct items: the four conversion directions, hardware-scoped items: `pretrain`, `sft`, `sft_export_inference`, `sft_long_context`, `peft`, and `checkpoint_resume`. Keep the optional `pretrain_performance` item separate under `performance`; omit `performance` -when the card has no canonical performance recipe. +when the card has no canonical performance recipe. Keep optional +`pretrain_fsdp` leaves separate under `fsdp`; omit `fsdp` when the card has no +FSDP recipe. Group item names under the same four status names used by the detailed items. For an explicitly indexed hardware target with no corresponding item leaf, @@ -174,6 +210,22 @@ When a canonical performance recipe exists, mirror only its concrete leaves: H100: verified ``` +When an FSDP recipe exists, mirror only its concrete leaves: + +```yaml + fsdp: + GB200: verified +``` + +For a multi-variant FSDP hardware container, this scalar mirrors the aggregate +container status rather than duplicating the per-precision inventory. + +Do not add prose comparisons, control payloads, computed deltas, or relative +speedups to performance leaves. Agents can compare standalone runs +automatically after resolving their recipes. The only exception is a concise +warning that names two superficially similar runs that must **not** be compared +and the exact convergence-contract differences that make them incompatible. + The index may declare an allowlisted public hardware target such as `GB200` before detailed evidence exists. It must also include every concrete hardware target present in the detailed items. Do not use a private cluster name or @@ -196,6 +248,10 @@ Use `bf16` for BF16. Training items may instead use `fp8_mx` for MXFP8 or `nvfp4` for NVFP4. Keep MXFP8 and NVFP4 training-only, and do not list either until that exact item has completed in that mode. +The outer hardware container of a multi-variant `pretrain_fsdp` item is the +only exception: each variant owns the scalar precision and the container owns +only aggregate `status` plus `variants`. + ### 3. Use the public Slurm launchers Assume the caller supplies the account, partition, concrete runtime image, @@ -281,11 +337,29 @@ loss sentinels for every layout and do not claim step-by-step numerical parity. Performance settings are **intended** to preserve training semantics, not guaranteed to be bitwise neutral. Parallel reductions, fusions, recompute, and -dispatcher implementations can change floating-point order. After changing -them, require finite loss, no skipped iterations, and compatible loss sentinels -before calling the mapping verified. Anything that changes arithmetic -precision, forced router balancing, token dropping, packing, or effective batch -construction is a convergence change, even when introduced to improve speed. +dispatcher implementations can change floating-point order. Every paired run +used to claim that a performance-related feature preserves convergence must +start from the same weights and use the same clean Bridge commit, dataset and +revision, sample order, tokenizer, sequence length, global batch size, token +budget, optimizer steps, seeds, objective, routing, precision, optimizer, and +learning-rate schedule. Change only the feature under test and unavoidable +execution settings. Keep micro batch size and gradient accumulation unchanged +when possible; if either changes, require the same loss check but do not claim +step-by-step numerical identity. If a verified run deliberately changes any +convergence field, make the exact deviation explicit in its command or +`expected_result`. + +Compare LM and auxiliary-loss values at every shared optimizer step. Each +candidate value must satisfy +`abs(candidate - reference) <= 1e-6 + 0.01 * abs(reference)`. Require finite +losses, zero skipped or NaN iterations, and the same qualitative loss trend. +If any convergence field differs or any loss falls outside this bound, do not +claim that the feature is convergence-neutral. Treat the result as standalone +performance evidence, investigate the discrepancy, and explicitly identify +the pair as not comparable only when the card would otherwise invite a false +comparison. Anything that changes arithmetic precision, forced router +balancing, token dropping, packing, or effective batch construction is a +convergence change, even when introduced to improve speed. The **benchmark-only configuration** may deliberately change semantics to find an upper throughput bound. It includes mock data, forced MoE load balancing, @@ -575,6 +649,17 @@ For every verified training item, record: optimizer steps; - `last_10_steps_model_tflops_per_gpu_avg`: arithmetic mean over the same rows. +Optionally record `peak_allocated_memory_gib` and +`peak_reserved_memory_gib`. They are required on verified `pretrain_fsdp` +leaves, including every verified precision variant. + +Record raw metrics for each run only. Do not store comparison payloads, +throughput or memory deltas, speedups, or prose ranking one performance run +against another. Agents can calculate those from standalone leaves after +checking that their resolved convergence contracts match. A concise +not-comparable warning is allowed only when it prevents a misleading +comparison and names the specific mismatched fields. + Parse all fields from each complete keyed optimizer-step line. Do not collect loss, time, and throughput independently and zip them. Reject missing, duplicate, skipped, NaN, or non-finite rows rather than excluding them. @@ -592,8 +677,9 @@ reproducible functional run; keep the physical dataset root private. ### 7. Record only important enabled features -Use `enabled_features` only on pretrain, SFT, long-context SFT, and PEFT. Keep -it empty when none of these are central to the verification. +Use `enabled_features` only on pretrain, SFT, long-context SFT, PEFT, and +`pretrain_fsdp`. Keep it empty when none of these are central to the +verification. | Key | Allowed value | | --- | --- | @@ -602,6 +688,7 @@ it empty when none of these are central to the verification. | `cuda_graph.scopes` | `full_iteration`, `attn`, `mlp`, `moe`, `moe_router`, `moe_preprocess`, or `mamba` | | `context_parallel_size` | integer greater than one | | `moe_dispatcher` | `deepep` or `hybridep` | +| `megatron_fsdp` | `optim_grads_params` (only on `pretrain_fsdp`) | Do not list routine TP/PP/DP sizes, Transformer Engine, fused loss, distributed optimizer, ordinary communication overlap, LoRA, or DoRA. @@ -630,12 +717,16 @@ an item verified merely to make validation pass. - Keep all twelve core inventory items and use only the four statuses. Include `pretrain_performance` only when the exact variant has a canonical public - performance recipe. + performance recipe, and `pretrain_fsdp` only when an exact FSDP performance + recipe has a completed standalone verification run. Use precision-keyed + variants under one hardware container when multiple FSDP runs exist for the + same hardware, and make the container status summarize every variant. - Start the summary with the exact untuned performance disclaimer unless at least one concrete `pretrain_performance` hardware leaf exists; never use an `all` placeholder, and scope any tuned claim to the exact concrete leaf. - Put the verified workload precision on every direct item or hardware leaf; use `fp8_mx` and `nvfp4` only for training leaves that ran in those modes. + For a multi-variant FSDP container, put precision on each variant. - Pin a public immutable HF revision, minimum Transformers version, public base container, and exact Bridge verification commit; use an item override only for a verified workload run from a different clean commit. @@ -659,7 +750,11 @@ an item verified merely to make validation pass. with `qwen3_30b_a3b_convergence_v2` or record the exception and classify the result as support verification rather than cross-model convergence evidence. - Change only the execution/performance contract while tuning throughput, and - recheck loss sentinels after numerically non-bitwise changes. + require every shared-step loss to match within the declared 1% relative plus + `1e-6` absolute bound after numerically non-bitwise changes. +- Keep performance leaves free of control payloads, computed deltas, speedups, + and comparison prose. Mention two runs together only to warn that differing + convergence contracts make them unsuitable for comparison. - Leave recipe global and micro batch sizes unchanged in card commands. - Save full SFT, export it to HF, and record an exact deterministic N-token HF completion in a two-command ordered list. diff --git a/skills/create-model-verification-card/scripts/validate_card.py b/skills/create-model-verification-card/scripts/validate_card.py index 4d9c6dfbaa5..cf0ba992144 100644 --- a/skills/create-model-verification-card/scripts/validate_card.py +++ b/skills/create-model-verification-card/scripts/validate_card.py @@ -52,7 +52,7 @@ "peft", "checkpoint_resume", ) -OPTIONAL_ITEM_NAMES = ("pretrain_performance",) +OPTIONAL_ITEM_NAMES = ("pretrain_performance", "pretrain_fsdp") ITEM_NAMES = REQUIRED_ITEM_NAMES + OPTIONAL_ITEM_NAMES MODEL_LEVEL_INDEX_SCOPE = ( "hf_to_megatron_cpu", @@ -71,7 +71,15 @@ "checkpoint_resume", ) TRAINING_ITEMS = frozenset( - {"pretrain", "sft", "sft_long_context", "peft", "checkpoint_resume", "pretrain_performance"} + { + "pretrain", + "sft", + "sft_long_context", + "peft", + "checkpoint_resume", + "pretrain_performance", + "pretrain_fsdp", + } ) HARDWARE_SCOPED_ITEMS = TRAINING_ITEMS | {"sft_export_inference"} PUBLIC_HARDWARE_KEYS = frozenset( @@ -91,8 +99,8 @@ } ) CONVERSION_ITEMS = frozenset({"hf_to_megatron_cpu", "hf_to_megatron_gpu", "megatron_to_hf_cpu", "megatron_to_hf_gpu"}) -FEATURE_ITEMS = frozenset({"pretrain", "sft", "sft_long_context", "peft"}) -METRIC_NAMES = frozenset( +FEATURE_ITEMS = frozenset({"pretrain", "sft", "sft_long_context", "peft", "pretrain_fsdp"}) +REQUIRED_METRIC_NAMES = frozenset( { "initial_loss", "final_loss", @@ -100,11 +108,16 @@ "last_10_steps_model_tflops_per_gpu_avg", } ) -FEATURE_KEYS = frozenset({"sequence_packing", "cuda_graph", "context_parallel_size", "moe_dispatcher"}) +OPTIONAL_METRIC_NAMES = frozenset({"peak_allocated_memory_gib", "peak_reserved_memory_gib"}) +METRIC_NAMES = REQUIRED_METRIC_NAMES | OPTIONAL_METRIC_NAMES +FEATURE_KEYS = frozenset( + {"sequence_packing", "cuda_graph", "context_parallel_size", "moe_dispatcher", "megatron_fsdp"} +) PACKING_VALUES = frozenset({"offline", "in_batch"}) CUDA_GRAPH_IMPLEMENTATIONS = frozenset({"local", "transformer_engine"}) CUDA_GRAPH_SCOPES = frozenset({"full_iteration", "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba"}) MOE_DISPATCHERS = frozenset({"deepep", "hybridep"}) +MEGATRON_FSDP_STRATEGIES = frozenset({"optim_grads_params"}) MANUAL_FORWARD_COSINE_THRESHOLD = 0.99 MANUAL_FORWARD_REVISION_PINNING_DATE = dt.date(2026, 7, 20) UNTUNED_PERFORMANCE_DISCLAIMER = ( @@ -113,7 +126,7 @@ ) TOP_LEVEL_KEYS = frozenset({"title", "model", "verification_environment", "summary", "verification_index", "items"}) -VERIFICATION_INDEX_KEYS = frozenset({"model_level", "training", "performance"}) +VERIFICATION_INDEX_KEYS = frozenset({"model_level", "training", "performance", "fsdp"}) MODEL_KEYS = frozenset({"hf_id", "hf_revision", "architecture", "min_transformers_version"}) ENVIRONMENT_KEYS = frozenset({"base_container", "bridge_commit"}) ITEM_KEYS = frozenset( @@ -129,6 +142,7 @@ "enabled_features", "metrics", "resume_comparison", + "variants", } ) RESUME_KEYS = frozenset( @@ -140,7 +154,6 @@ "sentinels_match", } ) - FORBIDDEN_KEY_FRAGMENTS = ( "account", "cluster", @@ -314,7 +327,24 @@ def _iter_item_leaves( if item_name in HARDWARE_SCOPED_ITEMS and "status" not in value: for hardware, leaf in value.items(): if isinstance(leaf, Mapping): - yield item_name, hardware, leaf, ("items", item_name, hardware) + variants = leaf.get("variants") if item_name == "pretrain_fsdp" else None + if isinstance(variants, Mapping): + for variant_name, variant in variants.items(): + if isinstance(variant, Mapping): + yield ( + item_name, + hardware, + variant, + ( + "items", + item_name, + hardware, + "variants", + str(variant_name), + ), + ) + else: + yield item_name, hardware, leaf, ("items", item_name, hardware) else: yield item_name, None, value, ("items", item_name) @@ -479,48 +509,45 @@ def _validate_verification_index( errors=errors, ) - performance_path = (*path, "performance") - performance_variants = { - hardware: item - for hardware, item in hardware_groups.get("pretrain_performance", {}).items() - if hardware != "all" - } - if not performance_variants: - if "performance" in verification_index: - errors.append( - f"{_pointer(*performance_path)}: omit performance when pretrain_performance has no concrete leaves" - ) - return - if "performance" not in verification_index: - errors.append(f"{_pointer(*performance_path)}: required to mirror pretrain_performance concrete leaves") - return + for index_name, item_name in (("performance", "pretrain_performance"), ("fsdp", "pretrain_fsdp")): + index_path = (*path, index_name) + variants = { + hardware: item for hardware, item in hardware_groups.get(item_name, {}).items() if hardware != "all" + } + if not variants: + if index_name in verification_index: + errors.append(f"{_pointer(*index_path)}: omit {index_name} when {item_name} has no concrete leaves") + continue + if index_name not in verification_index: + errors.append(f"{_pointer(*index_path)}: required to mirror {item_name} concrete leaves") + continue - performance = _as_mapping(verification_index.get("performance"), path=performance_path, errors=errors) - if performance is None: - return - expected_hardware = set(performance_variants) - actual_hardware = set(performance) - for hardware in sorted(actual_hardware): - if hardware not in PUBLIC_HARDWARE_KEYS: - errors.append( - f"{_pointer(*performance_path, str(hardware))}: expected a supported public hardware key; " - f"choose from {sorted(PUBLIC_HARDWARE_KEYS)}" - ) - for hardware in sorted(expected_hardware - actual_hardware): - errors.append(f"{_pointer(*performance_path, hardware)}: required to mirror pretrain_performance.{hardware}") - for hardware in sorted(actual_hardware - expected_hardware): - errors.append(f"{_pointer(*performance_path, hardware)}: no matching pretrain_performance.{hardware} leaf") - for hardware in sorted(expected_hardware & actual_hardware): - indexed_status = performance.get(hardware) - if not isinstance(indexed_status, str) or indexed_status not in STATUSES: - errors.append(f"{_pointer(*performance_path, hardware)}: expected one of {sorted(STATUSES)}") + index = _as_mapping(verification_index.get(index_name), path=index_path, errors=errors) + if index is None: continue - expected_status = _item_status(performance_variants[hardware]) - if expected_status is not None and indexed_status != expected_status: - errors.append( - f"{_pointer(*performance_path, hardware)}: indexed as {indexed_status} but " - f"pretrain_performance.{hardware} is {expected_status}" - ) + expected_hardware = set(variants) + actual_hardware = set(index) + for hardware in sorted(actual_hardware): + if hardware not in PUBLIC_HARDWARE_KEYS: + errors.append( + f"{_pointer(*index_path, str(hardware))}: expected a supported public hardware key; " + f"choose from {sorted(PUBLIC_HARDWARE_KEYS)}" + ) + for hardware in sorted(expected_hardware - actual_hardware): + errors.append(f"{_pointer(*index_path, hardware)}: required to mirror {item_name}.{hardware}") + for hardware in sorted(actual_hardware - expected_hardware): + errors.append(f"{_pointer(*index_path, hardware)}: no matching {item_name}.{hardware} leaf") + for hardware in sorted(expected_hardware & actual_hardware): + indexed_status = index.get(hardware) + if not isinstance(indexed_status, str) or indexed_status not in STATUSES: + errors.append(f"{_pointer(*index_path, hardware)}: expected one of {sorted(STATUSES)}") + continue + expected_status = _item_status(variants[hardware]) + if expected_status is not None and indexed_status != expected_status: + errors.append( + f"{_pointer(*index_path, hardware)}: indexed as {indexed_status} but " + f"{item_name}.{hardware} is {expected_status}" + ) def _is_iso_date(value: Any) -> bool: @@ -569,7 +596,13 @@ def _command_values(item: Mapping[str, Any]) -> list[str]: return [] -def _validate_enabled_features(value: Any, *, item_path: tuple[str, ...], errors: list[str]) -> None: +def _validate_enabled_features( + value: Any, + *, + item_name: str, + item_path: tuple[str, ...], + errors: list[str], +) -> None: path = (*item_path, "enabled_features") features = _as_mapping(value, path=path, errors=errors) if features is None: @@ -615,10 +648,18 @@ def _validate_enabled_features(value: Any, *, item_path: tuple[str, ...], errors if dispatcher is not None and (not isinstance(dispatcher, str) or dispatcher not in MOE_DISPATCHERS): errors.append(f"{_pointer(*path, 'moe_dispatcher')}: expected deepep or hybridep") + megatron_fsdp = features.get("megatron_fsdp") + if item_name == "pretrain_fsdp": + if not isinstance(megatron_fsdp, str) or megatron_fsdp not in MEGATRON_FSDP_STRATEGIES: + errors.append(f"{_pointer(*path, 'megatron_fsdp')}: expected one of {sorted(MEGATRON_FSDP_STRATEGIES)}") + elif megatron_fsdp is not None: + errors.append(f"{_pointer(*path, 'megatron_fsdp')}: allowed only on pretrain_fsdp") + def _validate_metrics( item: Mapping[str, Any], *, + item_name: str, item_path: tuple[str, ...], status: str, errors: list[str], @@ -627,10 +668,11 @@ def _validate_metrics( metrics = _as_mapping(item.get("metrics"), path=path, errors=errors) if metrics is None: return - _check_keys(metrics, allowed=METRIC_NAMES, required=METRIC_NAMES, path=path, errors=errors) + required_names = METRIC_NAMES if item_name == "pretrain_fsdp" and status == "verified" else REQUIRED_METRIC_NAMES + _check_keys(metrics, allowed=METRIC_NAMES, required=required_names, path=path, errors=errors) if status == "verified": - for name in sorted(METRIC_NAMES): + for name in sorted(REQUIRED_METRIC_NAMES | (OPTIONAL_METRIC_NAMES & metrics.keys())): value = metrics.get(name) if not _is_finite_number(value): errors.append(f"{_pointer(*path, name)}: verified metrics must be finite numbers") @@ -639,6 +681,8 @@ def _validate_metrics( in { "last_10_steps_step_time_ms_avg", "last_10_steps_model_tflops_per_gpu_avg", + "peak_allocated_memory_gib", + "peak_reserved_memory_gib", } and float(value) <= 0 ): @@ -1331,6 +1375,8 @@ def _validate_item( item = _as_mapping(value, path=path, errors=errors) if item is None: return + if "variants" in item: + errors.append(f"{_pointer(*path, 'variants')}: allowed only on a pretrain_fsdp hardware container") required = frozenset({"status", "precision", "last_verified", "expected_result"}) if item_name == "sft_export_inference": required |= frozenset({"commands"}) @@ -1438,7 +1484,13 @@ def _validate_item( ) if item_name in TRAINING_ITEMS: - _validate_metrics(item, item_path=path, status=status, errors=errors) + _validate_metrics( + item, + item_name=item_name, + item_path=path, + status=status, + errors=errors, + ) _validate_training_window( item, item_name=item_name, @@ -1454,7 +1506,12 @@ def _validate_item( errors.append(f"{_pointer(*path)}: metrics are allowed only on training items") if item_name in FEATURE_ITEMS: - _validate_enabled_features(item.get("enabled_features"), item_path=path, errors=errors) + _validate_enabled_features( + item.get("enabled_features"), + item_name=item_name, + item_path=path, + errors=errors, + ) if item_name == "sft_long_context" and status == "verified": features = item.get("enabled_features") if isinstance(features, Mapping): @@ -1482,6 +1539,65 @@ def _validate_item( _validate_inference(item, item_name=item_name, item_path=path, status=status, errors=errors) +def _validate_fsdp_variant_group( + value: Any, + *, + path: tuple[str, ...], + model_revision: str | None, + errors: list[str], +) -> None: + """Validate multiple precision variants under one FSDP hardware target.""" + group = _as_mapping(value, path=path, errors=errors) + if group is None: + return + _check_keys( + group, + allowed=frozenset({"status", "variants"}), + required=frozenset({"status", "variants"}), + path=path, + errors=errors, + ) + + status = group.get("status") + if not isinstance(status, str) or status not in {"verified", "unverified"}: + errors.append(f"{_pointer(*path, 'status')}: FSDP variant groups must be verified or unverified") + + variants = _as_mapping(group.get("variants"), path=(*path, "variants"), errors=errors) + if variants is None: + return + if not variants: + errors.append(f"{_pointer(*path, 'variants')}: expected at least one precision variant") + return + + variant_statuses: list[str] = [] + for precision_name, variant in variants.items(): + variant_path = (*path, "variants", precision_name) + if precision_name not in PRECISIONS: + errors.append(f"{_pointer(*variant_path)}: expected one of {sorted(PRECISIONS)}") + continue + _validate_item( + "pretrain_fsdp", + variant, + errors, + path=variant_path, + model_revision=model_revision, + ) + if isinstance(variant, Mapping): + if variant.get("precision") != precision_name: + errors.append(f"{_pointer(*variant_path, 'precision')}: must match the variant key") + variant_status = variant.get("status") + if isinstance(variant_status, str) and variant_status in {"verified", "unverified"}: + variant_statuses.append(variant_status) + + expected_status = ( + "verified" + if len(variant_statuses) == len(variants) and all(value == "verified" for value in variant_statuses) + else "unverified" + ) + if status in {"verified", "unverified"} and status != expected_status: + errors.append(f"{_pointer(*path, 'status')}: must be {expected_status} to summarize the precision variants") + + def _walk_keys(value: Any, path: tuple[str, ...] = ()) -> Iterable[tuple[tuple[str, ...], str]]: if isinstance(value, Mapping): for key, child in value.items(): @@ -1631,13 +1747,22 @@ def _validate_card(card: Mapping[str, Any], raw: str, deny_terms: tuple[str, ... variants = _hardware_variants(items[name], item_name=name, errors=errors) hardware_groups[name] = variants for hardware, item in variants.items(): - _validate_item( - name, - item, - errors, - path=("items", name, hardware), - model_revision=model_revision, - ) + item_path = ("items", name, hardware) + if name == "pretrain_fsdp" and "variants" in item: + _validate_fsdp_variant_group( + item, + path=item_path, + model_revision=model_revision, + errors=errors, + ) + else: + _validate_item( + name, + item, + errors, + path=item_path, + model_revision=model_revision, + ) else: _validate_item( name, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py index 39e9cb9df3b..22fac9f9a3b 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py @@ -175,6 +175,10 @@ def nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer: def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: """Nemotron 3.5 Nano pretrain: 16× H100, BF16.""" cfg = nemotron_3_nano_pretrain_16gpu_h100_bf16_config() + # Keep the benchmark workload aligned with the GB200 BF16 recipe. The + # hardware recipes may tune execution-only knobs such as microbatch size, + # recompute, and CUDA graph coverage independently. + cfg.train.global_batch_size = 512 cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 85aea82d065..1540e438a80 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -15,7 +15,8 @@ # Nemotron Nano v2 models # Nemotron 3 and 3.5 Nano models from megatron.bridge.recipes.nemotronh.gb200 import ( - nemotron_3_5_nano_pretrain_4k_config, + nemotron_3_5_nano_pretrain_8k_config, + nemotron_3_5_nano_pretrain_8k_fsdp_config, nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, @@ -96,7 +97,8 @@ "nemotron_nano_12b_v2_peft_config", # Nemotron 3 and 3.5 Nano models "nemotron_3_5_nano_peft_config", - "nemotron_3_5_nano_pretrain_4k_config", + "nemotron_3_5_nano_pretrain_8k_config", + "nemotron_3_5_nano_pretrain_8k_fsdp_config", "nemotron_3_5_nano_pretrain_config", "nemotron_3_5_nano_sft_config", "nemotron_3_5_nano_sft_openmathinstruct2_packed_config", diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py index 907da5ecfa1..2fc637f14fb 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py @@ -13,7 +13,8 @@ # limitations under the License. from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano import ( - nemotron_3_5_nano_pretrain_4k_config, + nemotron_3_5_nano_pretrain_8k_config, + nemotron_3_5_nano_pretrain_8k_fsdp_config, nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, @@ -21,7 +22,8 @@ __all__ = [ - "nemotron_3_5_nano_pretrain_4k_config", + "nemotron_3_5_nano_pretrain_8k_config", + "nemotron_3_5_nano_pretrain_8k_fsdp_config", "nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py index e56fd5c2fce..136a7ded2f1 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py @@ -19,6 +19,7 @@ from megatron.bridge import AutoBridge from megatron.bridge.recipes.common import _pretrain_common from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import ( + nemotron_3_5_nano_pretrain_config, nemotron_3_5_nano_sft_openmathinstruct2_packed_config, ) from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS @@ -29,8 +30,6 @@ _NEMOTRON_3_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" -# Placeholder until the public Nemotron 3.5 Nano repository is released. -_NEMOTRON_3_5_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: @@ -158,16 +157,51 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_4k_config() -> ConfigContainer: - """Return the Nemotron 3.5 Nano BF16 4K pretraining config.""" - cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() - cfg.model.mtp_num_layers = 2 - cfg.model.mtp_hybrid_override_pattern = "*E" - cfg.model.mtp_use_repeated_layer = True - cfg.model.keep_mtp_spec_in_bf16 = True - cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID +def nemotron_3_5_nano_pretrain_8k_config() -> ConfigContainer: + """Return the Nemotron 3.5 Nano BF16 8K pretraining config.""" + cfg = nemotron_3_5_nano_pretrain_config() + + # Keep the H100 convergence contract intact and change only GB200 + # execution/performance settings. + cfg.train.micro_batch_size = 2 + cfg.model.context_parallel_size = 1 + cfg.model.cp_comm_type = None + cfg.model.cross_entropy_fusion_impl = "te" + cfg.model.cuda_graph_impl = "none" + set_cuda_graph_modules(cfg.model, []) + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.ddp.average_in_collective = False + cfg.checkpoint.save_interval = 50 + cfg.env_vars["NVLINK_DOMAIN_SIZE"] = 72 + cfg.env_vars["USE_MNNVL"] = 1 + return cfg + + +def nemotron_3_5_nano_pretrain_8k_fsdp_config() -> ConfigContainer: + """Return the Nemotron 3.5 Nano BF16 8K Megatron FSDP config.""" + cfg = nemotron_3_5_nano_pretrain_8k_config() + + # Megatron FSDP module hooks are incompatible with Transformer Engine CUDA + # graph capture. Keep all convergence settings inherited from the BF16 + # reference recipe and change only FSDP and its required execution knobs. + cfg.model.cuda_graph_impl = "none" + set_cuda_graph_modules(cfg.model, []) + + cfg.dist.use_megatron_fsdp = True + cfg.ddp.use_megatron_fsdp = True + cfg.ddp.num_distributed_optimizer_instances = 1 + cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params" + cfg.ddp.outer_dp_sharding_strategy = "no_shard" + cfg.ddp.average_in_collective = False + cfg.ddp.megatron_fsdp_main_params_dtype = torch.float32 + cfg.ddp.megatron_fsdp_main_grads_dtype = torch.float32 + cfg.ddp.megatron_fsdp_grad_comm_dtype = torch.bfloat16 + + cfg.checkpoint.load = None + cfg.checkpoint.ckpt_format = "fsdp_dtensor" return cfg @@ -186,10 +220,6 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContain cfg.model.recompute_num_layers = None cfg.train.empty_unused_memory_level = 0 - # Match the validated GB200 BF16 pretraining reduction policy. - cfg.mixed_precision = get_mixed_precision_config(cfg.mixed_precision) - cfg.mixed_precision.grad_reduce_in_fp32 = False - cfg.ddp.grad_reduce_in_fp32 = False cfg.ddp.overlap_param_gather = True cfg.optimizer.overlap_param_gather = True @@ -215,7 +245,8 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContain __all__ = [ - "nemotron_3_5_nano_pretrain_4k_config", + "nemotron_3_5_nano_pretrain_8k_config", + "nemotron_3_5_nano_pretrain_8k_fsdp_config", "nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", diff --git a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py index 782989d51bc..15651cef5ee 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py @@ -255,6 +255,14 @@ def nemotron_3_nano_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: def nemotron_3_5_nano_pretrain_config() -> ConfigContainer: """Return the Nemotron 3.5 Nano BF16 pretraining config.""" cfg = nemotron_3_nano_pretrain_8gpu_h100_bf16_config() + cfg.train.global_batch_size = 512 + # Split the 8K sequence across two ranks so each MTP head materializes only + # half of its vocabulary-loss workspace on an 80-GiB H100. P2P retains the + # fused-attention path for this model's grouped-query layout. + cfg.model.context_parallel_size = 2 + cfg.model.cp_comm_type = "p2p" + cfg.model.recompute_modules = ["moe", "layernorm", "core_attn"] + set_cuda_graph_modules(cfg.model, ["mamba"]) cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True @@ -440,6 +448,9 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_config() -> ConfigContainer: cfg.train.manual_gc_interval = 100 cfg.train.empty_unused_memory_level = 2 + cfg.mixed_precision = get_mixed_precision_config(cfg.mixed_precision) + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = False cfg.optimizer.lr = 5e-6 cfg.optimizer.min_lr = 0.0 cfg.optimizer.overlap_param_gather = False diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py index 6aca30cec7c..f63d744e944 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py @@ -42,6 +42,7 @@ nemotron_3_nano_sft_config, ) from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.utils.cuda_graph import cuda_graph_module_names @pytest.mark.unit @@ -108,6 +109,11 @@ def test_nemotron_3_5_pretrain_config(self): assert recipe_module._NEMOTRON_3_5_NANO_MODEL_ID == ("nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16") assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID assert config.tokenizer.tokenizer_model == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID + assert config.train.global_batch_size == 512 + assert config.model.context_parallel_size == 2 + assert config.model.cp_comm_type == "p2p" + assert config.model.recompute_modules == ["moe", "layernorm", "core_attn"] + assert cuda_graph_module_names(config.model) == ["mamba"] def test_pretrain_recipes_do_not_expose_mtp_flag(self): """Nemotron 3 and 3.5 pretraining use distinct parameterless factories.""" diff --git a/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py b/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py index 01fe65bc28c..3675747c53a 100644 --- a/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py +++ b/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py @@ -189,8 +189,10 @@ def test_h100_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer] cfg = recipe_factory() assert cfg.model.expert_model_parallel_size == 8 - assert cfg.train.global_batch_size == 1024 + expected_global_batch_size = 512 if recipe_factory is nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config else 1024 + assert cfg.train.global_batch_size == expected_global_batch_size assert cfg.train.micro_batch_size == 1 + assert cfg.model.context_parallel_size == 1 assert cfg.model.recompute_granularity == "selective" assert cfg.model.seq_length == 8192 assert cfg.dataset.seq_length == 8192 @@ -200,6 +202,22 @@ def test_h100_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer] assert cfg.env_vars["USE_MNNVL"] == 0 +def test_bf16_perf_recipes_share_training_workload() -> None: + """H100 and GB200 BF16 recipes differ only in execution tuning.""" + h100_cfg = nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() + gb200_cfg = nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config() + + gb200_cfg.train.micro_batch_size = h100_cfg.train.micro_batch_size + gb200_cfg.model.recompute_granularity = h100_cfg.model.recompute_granularity + gb200_cfg.model.recompute_modules = h100_cfg.model.recompute_modules + gb200_cfg.model.cuda_graph_impl = h100_cfg.model.cuda_graph_impl + gb200_cfg.model.cuda_graph_scope = h100_cfg.model.cuda_graph_scope + gb200_cfg.model.cuda_graph_modules = h100_cfg.model.cuda_graph_modules + gb200_cfg.env_vars = h100_cfg.env_vars + + assert gb200_cfg == h100_cfg + + @pytest.mark.parametrize("recipe_factory", _GB200_RECIPES, ids=lambda recipe: recipe.__name__) def test_gb200_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer]) -> None: """GB200 Nemotron 3.5 variants retain the existing Nano performance topology.""" diff --git a/tests/unit_tests/recipes/test_nemotronh_recipes.py b/tests/unit_tests/recipes/test_nemotronh_recipes.py index 6ae9ff58049..2622512e199 100644 --- a/tests/unit_tests/recipes/test_nemotronh_recipes.py +++ b/tests/unit_tests/recipes/test_nemotronh_recipes.py @@ -173,17 +173,19 @@ def test_nemotron_3_5_nano_h100_convergence_recipe_uses_perf_execution_policy(): assert cfg.model.sequence_parallel is False assert cfg.model.expert_tensor_parallel_size == 1 assert cfg.model.expert_model_parallel_size == 8 - assert cfg.train.global_batch_size == 1024 + assert cfg.train.global_batch_size == 512 assert cfg.train.micro_batch_size == 1 assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.moe_router_force_load_balancing is False assert cfg.model.cross_entropy_fusion_impl == "native" assert cfg.model.cuda_graph_impl == "transformer_engine" - assert cuda_graph_module_names(cfg.model) == ["attn", "mamba"] + assert cuda_graph_module_names(cfg.model) == ["mamba"] assert cfg.model.recompute_granularity == "selective" - assert cfg.model.recompute_modules == ["moe", "layernorm"] + assert cfg.model.recompute_modules == ["moe", "layernorm", "core_attn"] + assert cfg.mixed_precision.bf16 is True + assert cfg.mixed_precision.fp16 is False assert cfg.mixed_precision.grad_reduce_in_fp32 is False assert cfg.optimizer.use_precision_aware_optimizer is False assert cfg.optimizer.main_grads_dtype == torch.float32 @@ -204,15 +206,15 @@ def test_nemotron_3_5_nano_h100_convergence_recipe_uses_perf_execution_policy(): assert cfg.env_vars["USE_MNNVL"] == 0 -def test_nemotron_3_5_nano_4k_convergence_recipe_uses_perf_execution_policy(): - """The 4K convergence recipe keeps safety checks while using the measured fast path.""" - from megatron.bridge.recipes.nemotronh import nemotron_3_5_nano_pretrain_4k_config +def test_nemotron_3_5_nano_8k_convergence_recipe_uses_perf_execution_policy(): + """The 8K convergence recipe keeps safety checks while using the measured fast path.""" + from megatron.bridge.recipes.nemotronh import nemotron_3_5_nano_pretrain_8k_config from megatron.bridge.utils.cuda_graph import cuda_graph_module_names - cfg = nemotron_3_5_nano_pretrain_4k_config() + cfg = nemotron_3_5_nano_pretrain_8k_config() - assert cfg.model.seq_length == 4096 - assert cfg.dataset.seq_length == 4096 + assert cfg.model.seq_length == 8192 + assert cfg.dataset.seq_length == 8192 assert cfg.model.tensor_model_parallel_size == 1 assert cfg.model.pipeline_model_parallel_size == 1 assert cfg.model.sequence_parallel is False @@ -222,18 +224,21 @@ def test_nemotron_3_5_nano_4k_convergence_recipe_uses_perf_execution_policy(): assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.moe_router_force_load_balancing is False - assert cfg.model.cross_entropy_fusion_impl == "native" - assert cfg.model.cuda_graph_impl == "transformer_engine" - assert cuda_graph_module_names(cfg.model) == ["attn", "mamba", "moe_router", "moe_preprocess"] + assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.cuda_graph_impl == "none" + assert cuda_graph_module_names(cfg.model) == [] assert cfg.model.recompute_granularity is None assert cfg.model.recompute_modules is None + assert cfg.mixed_precision.bf16 is True + assert cfg.mixed_precision.fp16 is False assert cfg.mixed_precision.grad_reduce_in_fp32 is False assert cfg.ddp.grad_reduce_in_fp32 is False assert cfg.ddp.check_for_nan_in_grad is True assert cfg.ddp.check_for_large_grads is True + assert cfg.ddp.average_in_collective is False assert cfg.rerun_state_machine.check_for_nan_in_loss is True - assert cfg.checkpoint.save_interval == 200 + assert cfg.checkpoint.save_interval == 50 assert cfg.checkpoint.async_save is False assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" @@ -243,6 +248,73 @@ def test_nemotron_3_5_nano_4k_convergence_recipe_uses_perf_execution_policy(): assert cfg.env_vars["USE_MNNVL"] == 1 +def test_nemotron_3_5_nano_h100_and_gb200_differ_only_in_execution_policy(): + """H100 and GB200 must share one training contract.""" + from megatron.bridge.recipes.nemotronh import ( + nemotron_3_5_nano_pretrain_8k_config, + nemotron_3_5_nano_pretrain_config, + ) + + h100 = nemotron_3_5_nano_pretrain_config() + gb200 = nemotron_3_5_nano_pretrain_8k_config() + + # Normalize the explicit GB200 execution/performance overrides. Any other + # difference, including a convergence-setting drift, must fail this test. + gb200.train.micro_batch_size = h100.train.micro_batch_size + gb200.model.context_parallel_size = h100.model.context_parallel_size + gb200.model.cp_comm_type = h100.model.cp_comm_type + gb200.model.cross_entropy_fusion_impl = h100.model.cross_entropy_fusion_impl + gb200.model.cuda_graph_impl = h100.model.cuda_graph_impl + gb200.model.cuda_graph_modules = h100.model.cuda_graph_modules + gb200.model.recompute_granularity = h100.model.recompute_granularity + gb200.model.recompute_modules = h100.model.recompute_modules + gb200.model.recompute_method = h100.model.recompute_method + gb200.model.recompute_num_layers = h100.model.recompute_num_layers + gb200.ddp.average_in_collective = h100.ddp.average_in_collective + gb200.checkpoint.save_interval = h100.checkpoint.save_interval + gb200.env_vars = h100.env_vars + + assert gb200 == h100 + + +def test_nemotron_3_5_nano_8k_fsdp_recipe_preserves_convergence_contract(): + """The BF16 FSDP recipe changes only FSDP-required execution settings.""" + from megatron.bridge.recipes.nemotronh import ( + nemotron_3_5_nano_pretrain_8k_config, + nemotron_3_5_nano_pretrain_8k_fsdp_config, + ) + from megatron.bridge.utils.cuda_graph import cuda_graph_module_names + + reference = nemotron_3_5_nano_pretrain_8k_config() + cfg = nemotron_3_5_nano_pretrain_8k_fsdp_config() + + assert cfg.model.seq_length == reference.model.seq_length == 8192 + assert cfg.dataset.seq_length == reference.dataset.seq_length == 8192 + assert cfg.train.global_batch_size == reference.train.global_batch_size == 512 + assert cfg.train.micro_batch_size == reference.train.micro_batch_size == 2 + assert cfg.model.moe_router_force_load_balancing is reference.model.moe_router_force_load_balancing is False + assert cfg.optimizer == reference.optimizer + assert cfg.scheduler == reference.scheduler + assert cfg.mixed_precision == reference.mixed_precision + assert cfg.rng == reference.rng + + assert cfg.model.cuda_graph_impl == "none" + assert cuda_graph_module_names(cfg.model) == [] + assert getattr(cfg.model, "init_model_with_meta_device", False) is False + assert getattr(reference.model, "init_model_with_meta_device", False) is False + assert cfg.dist.use_megatron_fsdp is True + assert cfg.ddp.use_megatron_fsdp is True + assert cfg.ddp.num_distributed_optimizer_instances == 1 + assert cfg.ddp.data_parallel_sharding_strategy == "optim_grads_params" + assert cfg.ddp.outer_dp_sharding_strategy == "no_shard" + assert cfg.ddp.average_in_collective is False + assert cfg.ddp.megatron_fsdp_main_params_dtype == torch.float32 + assert cfg.ddp.megatron_fsdp_main_grads_dtype == torch.float32 + assert cfg.ddp.megatron_fsdp_grad_comm_dtype == torch.bfloat16 + assert cfg.checkpoint.load is None + assert cfg.checkpoint.ckpt_format == "fsdp_dtensor" + + def test_nemotron_3_5_nano_openmath_sft_tp1_recipe_uses_tuned_defaults(): """The TP1 packed SFT recipe owns the measured topology and run contract.""" cfg = _nemotronh_module.nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() @@ -282,6 +354,42 @@ def test_nemotron_3_5_nano_openmath_sft_tp1_recipe_uses_tuned_defaults(): assert cfg.env_vars["USE_MNNVL"] == 1 +def test_nemotron_3_5_nano_h100_and_gb200_sft_differ_only_in_execution_policy(): + """H100 and GB200 packed SFT must share one training contract.""" + from megatron.bridge.recipes.nemotronh import ( + nemotron_3_5_nano_sft_openmathinstruct2_packed_config, + nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, + ) + + h100 = nemotron_3_5_nano_sft_openmathinstruct2_packed_config() + gb200 = nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() + + # Remove the explicit GB200 execution/performance model overrides. Any + # other difference, including a convergence-setting drift, must fail. + model_execution_fields = { + "tensor_model_parallel_size", + "sequence_parallel", + "moe_flex_dispatcher_backend", + "moe_flex_dispatcher_num_sms", + "moe_hybridep_num_sms", + "recompute_granularity", + "recompute_modules", + "recompute_method", + "recompute_num_layers", + } + h100_model = {key: value for key, value in vars(h100.model).items() if key not in model_execution_fields} + gb200_model = {key: value for key, value in vars(gb200.model).items() if key not in model_execution_fields} + assert gb200_model == h100_model + + gb200.model = h100.model + gb200.train.empty_unused_memory_level = h100.train.empty_unused_memory_level + gb200.ddp.overlap_param_gather = h100.ddp.overlap_param_gather + gb200.optimizer.overlap_param_gather = h100.optimizer.overlap_param_gather + gb200.env_vars = h100.env_vars + + assert gb200 == h100 + + def test_nemotron_nano_9b_v2_lora_defaults(): """Test that Nemotron Nano 9B v2 LoRA has correct default parallelism.""" from megatron.bridge.recipes.nemotronh import nemotron_nano_9b_v2_peft_config diff --git a/tests/unit_tests/scripts/test_validate_model_verification_card.py b/tests/unit_tests/scripts/test_validate_model_verification_card.py new file mode 100644 index 00000000000..1d7df997954 --- /dev/null +++ b/tests/unit_tests/scripts/test_validate_model_verification_card.py @@ -0,0 +1,278 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +from pathlib import Path + +import pytest + + +pytestmark = pytest.mark.unit + + +def _load_validator(): + script = ( + Path(__file__).resolve().parents[3] + / "skills" + / "create-model-verification-card" + / "scripts" + / "validate_card.py" + ) + spec = importlib.util.spec_from_file_location("test_model_verification_card_validator", script) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _complete_index_inputs(module): + items = {item_name: {"status": "verified"} for item_name in module.MODEL_LEVEL_INDEX_SCOPE} + hardware_groups = {item_name: {"H100": {"status": "verified"}} for item_name in module.TRAINING_INDEX_SCOPE} + verification_index = { + "model_level": {"verified": list(module.MODEL_LEVEL_INDEX_SCOPE)}, + "training": {"H100": {"verified": list(module.TRAINING_INDEX_SCOPE)}}, + } + return verification_index, items, hardware_groups + + +def _fsdp_metrics(): + return { + "initial_loss": 12.19034, + "final_loss": 3.913218, + "last_10_steps_step_time_ms_avg": 13917.0, + "last_10_steps_model_tflops_per_gpu_avg": 795.39, + "peak_allocated_memory_gib": 169.54, + "peak_reserved_memory_gib": 173.86, + } + + +def test_fsdp_index_mirrors_concrete_hardware_leaf(): + module = _load_validator() + verification_index, items, hardware_groups = _complete_index_inputs(module) + hardware_groups["pretrain_fsdp"] = {"GB200": {"status": "verified"}} + verification_index["fsdp"] = {"GB200": "verified"} + errors = [] + + module._validate_verification_index( + verification_index, + items=items, + hardware_groups=hardware_groups, + errors=errors, + ) + + assert errors == [] + + +def test_fsdp_index_is_required_for_concrete_leaf(): + module = _load_validator() + verification_index, items, hardware_groups = _complete_index_inputs(module) + hardware_groups["pretrain_fsdp"] = {"GB200": {"status": "verified"}} + errors = [] + + module._validate_verification_index( + verification_index, + items=items, + hardware_groups=hardware_groups, + errors=errors, + ) + + assert "/verification_index/fsdp: required to mirror pretrain_fsdp concrete leaves" in errors + + +def test_fsdp_index_rejects_mismatched_hardware_and_status(): + module = _load_validator() + verification_index, items, hardware_groups = _complete_index_inputs(module) + hardware_groups["pretrain_fsdp"] = {"GB200": {"status": "verified"}} + verification_index["fsdp"] = {"H100": "unverified"} + errors = [] + + module._validate_verification_index( + verification_index, + items=items, + hardware_groups=hardware_groups, + errors=errors, + ) + + assert "/verification_index/fsdp/GB200: required to mirror pretrain_fsdp.GB200" in errors + assert "/verification_index/fsdp/H100: no matching pretrain_fsdp.H100 leaf" in errors + + +def test_fsdp_index_is_omitted_for_terminal_all_leaf(): + module = _load_validator() + verification_index, items, hardware_groups = _complete_index_inputs(module) + hardware_groups["pretrain_fsdp"] = {"all": {"status": "not_applicable"}} + errors = [] + + module._validate_verification_index( + verification_index, + items=items, + hardware_groups=hardware_groups, + errors=errors, + ) + + assert errors == [] + + +def test_verified_fsdp_metrics_are_valid(): + module = _load_validator() + fsdp_item = {"status": "verified", "precision": "fp8_mx", "metrics": _fsdp_metrics()} + errors = [] + + module._validate_metrics( + fsdp_item, + item_name="pretrain_fsdp", + item_path=("items", "pretrain_fsdp", "GB200"), + status="verified", + errors=errors, + ) + + assert errors == [] + + +def test_fsdp_hardware_leaf_accepts_multiple_precision_variants(): + module = _load_validator() + errors = [] + variant = { + "status": "verified", + "precision": "bf16", + "command": "./scripts/training/train.sh --nodes 2 --gpus-per-node 4 --max_steps 20", + "last_verified": "2026-07-27", + "expected_result": "The run completes with finite losses.", + "enabled_features": {"megatron_fsdp": "optim_grads_params"}, + "metrics": _fsdp_metrics(), + } + + module._validate_fsdp_variant_group( + { + "status": "verified", + "variants": { + "bf16": variant, + "fp8_mx": {**variant, "precision": "fp8_mx"}, + }, + }, + path=("items", "pretrain_fsdp", "GB200"), + model_revision=None, + errors=errors, + ) + + assert errors == [] + + +def test_fsdp_variant_group_status_summarizes_all_variants(): + module = _load_validator() + errors = [] + variant = { + "status": "unverified", + "precision": "bf16", + "command": None, + "last_verified": None, + "expected_result": None, + "enabled_features": {"megatron_fsdp": "optim_grads_params"}, + "metrics": {}, + } + + module._validate_fsdp_variant_group( + { + "status": "verified", + "variants": {"bf16": variant}, + }, + path=("items", "pretrain_fsdp", "GB200"), + model_revision=None, + errors=errors, + ) + + assert "/items/pretrain_fsdp/GB200/status: must be unverified to summarize the precision variants" in errors + + +def test_fsdp_variant_precision_matches_mapping_key(): + module = _load_validator() + errors = [] + variant = { + "status": "verified", + "precision": "fp8_mx", + "command": "./scripts/training/train.sh --nodes 2 --gpus-per-node 4 --max_steps 20", + "last_verified": "2026-07-27", + "expected_result": "The run completes with finite losses.", + "enabled_features": {"megatron_fsdp": "optim_grads_params"}, + "metrics": _fsdp_metrics(), + } + + module._validate_fsdp_variant_group( + { + "status": "verified", + "variants": {"bf16": variant}, + }, + path=("items", "pretrain_fsdp", "GB200"), + model_revision=None, + errors=errors, + ) + + assert "/items/pretrain_fsdp/GB200/variants/bf16/precision: must match the variant key" in errors + + +def test_performance_comparison_payload_is_not_an_item_key(): + module = _load_validator() + errors = [] + + module._check_keys( + {"matched_non_fsdp_comparison": {}}, + allowed=module.ITEM_KEYS, + required=frozenset(), + path=("items", "pretrain", "GB200"), + errors=errors, + ) + + assert errors == ["/items/pretrain/GB200/matched_non_fsdp_comparison: unknown key"] + + +def test_verified_fsdp_metrics_require_peak_memory(): + module = _load_validator() + item = {"metrics": _fsdp_metrics()} + del item["metrics"]["peak_allocated_memory_gib"] + del item["metrics"]["peak_reserved_memory_gib"] + errors = [] + + module._validate_metrics( + item, + item_name="pretrain_fsdp", + item_path=("items", "pretrain_fsdp", "GB200"), + status="verified", + errors=errors, + ) + + assert "/items/pretrain_fsdp/GB200/metrics/peak_allocated_memory_gib: required key is missing" in errors + assert "/items/pretrain_fsdp/GB200/metrics/peak_reserved_memory_gib: required key is missing" in errors + + +def test_megatron_fsdp_feature_is_scoped_to_fsdp_item(): + module = _load_validator() + features = {"moe_dispatcher": "hybridep", "megatron_fsdp": "optim_grads_params"} + fsdp_errors = [] + pretrain_errors = [] + + module._validate_enabled_features( + features, + item_name="pretrain_fsdp", + item_path=("items", "pretrain_fsdp", "GB200"), + errors=fsdp_errors, + ) + module._validate_enabled_features( + features, + item_name="pretrain", + item_path=("items", "pretrain", "GB200"), + errors=pretrain_errors, + ) + + assert fsdp_errors == [] + assert pretrain_errors == ["/items/pretrain/GB200/enabled_features/megatron_fsdp: allowed only on pretrain_fsdp"] diff --git a/tests/unit_tests/scripts/training/test_setup_experiment.py b/tests/unit_tests/scripts/training/test_setup_experiment.py index 8556f2af271..e7b5b2d2e0b 100644 --- a/tests/unit_tests/scripts/training/test_setup_experiment.py +++ b/tests/unit_tests/scripts/training/test_setup_experiment.py @@ -212,6 +212,29 @@ def test_parser_consumes_repeatable_srun_args(): assert training_args == ["--recipe", "gpt_oss_20b_pretrain_config"] +def test_parser_consumes_additional_slurm_parameters(): + module = _load_setup_experiment_module() + + args, training_args = module.parse_args( + [ + "--additional-slurm-params", + "segment=8;reservation=testing", + "--recipe", + "gpt_oss_20b_pretrain_config", + ] + ) + + assert args.additional_slurm_params == {"segment": "8", "reservation": "testing"} + assert training_args == ["--recipe", "gpt_oss_20b_pretrain_config"] + + +def test_additional_slurm_parameters_require_key_value_pairs(): + module = _load_setup_experiment_module() + + with pytest.raises(SystemExit): + module.parse_args(["--additional-slurm-params", "segment"]) + + @pytest.mark.parametrize("option", ["-lmc", "--peak-mem-clk", "--peak_mem_clk"]) def test_parser_consumes_peak_mem_clk(option): module = _load_setup_experiment_module() @@ -486,6 +509,39 @@ def __init__(self, **kwargs): assert set(executor.container_env) == {"PYTHONPATH"} +def test_slurm_executor_forwards_additional_slurm_parameters(tmp_path, monkeypatch): + module = _load_setup_experiment_module() + + class _SlurmExecutor: + def __init__(self, **kwargs): + self.kwargs = kwargs + + module.run.LocalTunnel = lambda **kwargs: types.SimpleNamespace(**kwargs) + module.run.Packager = object + module.run.SlurmExecutor = _SlurmExecutor + monkeypatch.setattr(module, "get_nemorun_home", lambda: str(tmp_path)) + args, _ = module.parse_args( + [ + "--nodes", + "8", + "--gpus-per-node", + "4", + "--account", + "account", + "--partition", + "partition", + "--container-image", + "image.sqsh", + "--additional_slurm_params", + "segment=8", + ] + ) + + executor = module._build_executor(args, [], []) + + assert executor.additional_parameters == {"segment": "8", "export": "PATH"} + + def test_training_task_environment_does_not_inject_benchmark_offline_defaults(): module = _load_setup_experiment_module() environment = module._task_environment() From ff3795155a582cc1b7acc5c48c10098cafc123fe Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 31 Jul 2026 20:12:18 +0200 Subject: [PATCH 06/85] cp: `fix(docker): update dependencies in Dockerfile and fw_pyproject.toml for compatibility and security (5202)` into `r0.6.0` (#5208) Signed-off-by: Ajay Balasa Signed-off-by: NeMo Bot Co-authored-by: Ajay --- docker/Dockerfile.ci | 3 +-- docker/Dockerfile.fw_final | 28 ++++++++++++++++------------ docker/common/fw_pyproject.toml | 11 +++++++---- 3 files changed, 24 insertions(+), 18 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 019e3317113..46c7c8b2f39 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -65,8 +65,7 @@ ENV PATH="/opt/venv/bin:/opt/venv/bin:/root/.local/bin:$PATH" \ RUN --mount=type=bind,source=docker/patches/deepep.patch,target=/opt/deepep.patch \ curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \ uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages && \ - # Address CVE-2025-68973 - apt-get update && apt install -y --only-upgrade gnupg && \ + apt-get update && apt install -y --only-upgrade gnupg openssl libssl3t64 && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* && \ if [ "$INSTALL_DEEPEP" = "True" ]; then \ diff --git a/docker/Dockerfile.fw_final b/docker/Dockerfile.fw_final index dae1913424e..74a3132f8d6 100644 --- a/docker/Dockerfile.fw_final +++ b/docker/Dockerfile.fw_final @@ -71,24 +71,26 @@ RUN \ ## ############################################################################## -RUN pip install "starlette==0.49.1" \ +RUN pip install "fastapi==0.139.1" \ + "starlette==1.3.1" \ + "prometheus-fastapi-instrumentator>=8.0.0" \ "urllib3>=2.7.0" \ "aiohttp>=3.13.3" \ "jaraco-context>=6.1.0" \ "wheel>=0.46.2" \ "python-multipart>=0.0.22" \ "setuptools>=80.10.2" \ - "pillow>=12.1.1" \ + "pillow>=12.3.0" \ "protobuf==6.33.5" \ "xgrammar>=0.1.32" \ - "tornado==6.5.5" \ + "tornado==6.5.6" \ "black==26.3.1" \ "onnx==1.21.0" \ "notebook>=7.5.6" \ - "jupyterlab>=4.5.7" \ - "jupyter-server>=2.18.0" \ - "mistune>=3.2.1" \ - "GitPython>=3.1.50" && \ + "jupyterlab>=4.5.10" \ + "jupyter-server>=2.20.0" \ + "mistune>=3.3.0" \ + "GitPython>=3.1.55" && \ # Address CVE from transitive dependency pip uninstall -y opencv-python-headless nvidia-dali-cuda130 @@ -96,10 +98,10 @@ RUN pip install "starlette==0.49.1" \ # The actual onnx Python package is managed above via pip/uv; this stale copy triggers CVE scanners. RUN rm -rf /opt/pytorch/pytorch/third_party/onnx -# Patch wandb-core: bump google.golang.org/grpc from 1.79.2 to 1.79.3 (CVE fix) +# Patch wandb-core with fixed gRPC-Go and Go standard-library versions. ARG TARGETARCH -RUN GRPC_VERSION=1.79.3 && \ - GO_VERSION=1.26.1 && \ +RUN GRPC_VERSION=1.82.1 && \ + GO_VERSION=1.26.5 && \ WANDB_VERSION=$(python3 -c "import wandb; print(wandb.__version__)") && \ WANDB_CORE_BIN=/opt/venv/lib/python3.12/site-packages/wandb/bin/wandb-core && \ curl -fsSL "https://dl.google.com/go/go${GO_VERSION}.linux-${TARGETARCH}.tar.gz" | tar -C /tmp -xz && \ @@ -114,12 +116,13 @@ RUN GRPC_VERSION=1.79.3 && \ -o "${WANDB_CORE_BIN}" ./cmd/wandb-core/ && \ rm -rf /tmp/wandb-src /tmp/go /tmp/gopath -# Patch wandb-core (dist-packages copy): bump github.com/go-git/go-git/v5 to 5.19.0 (CVE fix). +# Patch wandb-core (dist-packages copy) with fixed go-git, gRPC-Go, and Go versions. # This stock upstream binary ships with go-git v5.17.2 and is a separate wandb install from the # /opt/venv copy rebuilt above, so it must be rebuilt against its own wandb version (resolved via # the system interpreter that owns dist-packages) to avoid a python/wandb-core version mismatch. RUN GO_GIT_VERSION=5.19.0 && \ - GO_VERSION=1.26.1 && \ + GRPC_VERSION=1.82.1 && \ + GO_VERSION=1.26.5 && \ WANDB_VERSION=$(/usr/bin/python3 -c "import wandb; print(wandb.__version__)") && \ WANDB_CORE_BIN=/usr/local/lib/python3.12/dist-packages/wandb/bin/wandb-core && \ curl -fsSL "https://dl.google.com/go/go${GO_VERSION}.linux-${TARGETARCH}.tar.gz" | tar -C /tmp -xz && \ @@ -128,6 +131,7 @@ RUN GO_GIT_VERSION=5.19.0 && \ git clone --depth 1 --branch "v${WANDB_VERSION}" https://github.com/wandb/wandb.git /tmp/wandb-src && \ cd /tmp/wandb-src/core && \ go get github.com/go-git/go-git/v5@v${GO_GIT_VERSION} && \ + go get google.golang.org/grpc@v${GRPC_VERSION} && \ go mod tidy && \ go mod vendor && \ CGO_ENABLED=0 GOOS=linux GOARCH=${TARGETARCH} go build -trimpath -ldflags="-s -w" \ diff --git a/docker/common/fw_pyproject.toml b/docker/common/fw_pyproject.toml index 5d889a5943c..ca9f9e83d9e 100644 --- a/docker/common/fw_pyproject.toml +++ b/docker/common/fw_pyproject.toml @@ -67,9 +67,9 @@ override-dependencies = [ "packaging>=25.0", "nvidia-pytriton>=0.7.0", "urllib3>=2.0.7", - "sh==2.2.2", + "sh==2.2.4", "jupyter-core==5.8.1", - "pillow==12.2.0", + "pillow==12.3.0", "protobuf==6.33.5", "cuda-python>=13.0.0", "multi-storage-client!=0.36.0", @@ -115,12 +115,15 @@ override-dependencies = [ "setuptools>=82.0.0", # Required by vllm "open-clip-torch>=3.2.0", # Required by megatron-bridge, override nemo-toolkit's 2.24.0 constraint "lilcom!=1.8.2", - "cryptography>=43.0.0,<47", # Required by megatron-bridge, override nemo-run's constraint + "cryptography>=48.0.1,<49", # Required by megatron-bridge, override nemo-run's constraint "opencv-python; sys_platform == 'never'", # Use the dependency from the base pytorch container "langchain~=1.0.0", "langchain-core~=1.3.3", # GHSA-pjwx-r37v-7724 "decord; sys_platform == 'never'", # Optional dependency. Resolves CVE "av; sys_platform == 'never'", # Optional dependency. Resolves CVE "opencv-python-headless; sys_platform == 'never'", # Optional dependency. Resolves CVE - "fastapi<0.139.2" # 0.139.2 fails with ray 2.55.1 at least + "ray>=2.56.0", + "starlette==1.3.1", + "prometheus-fastapi-instrumentator>=8.0.0", # 7.x requires starlette<1.0.0 + "fastapi<0.139.2" # 0.139.2 fails with ray serve ] From 9544c9385e121f0cce13cd9a460df530f1a7c85b Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Mon, 3 Aug 2026 19:25:23 +0200 Subject: [PATCH 07/85] cp: `fix(model): map GLM-5 routed expert count (5220)` into `r0.6.0` (#5248) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- .../bridge/models/glm_moe_dsa/glm5_bridge.py | 13 +++++++ .../models/glm_moe_dsa/test_glm5_bridge.py | 37 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py index 3b776e3e979..56ff48fb08c 100644 --- a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py +++ b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +from typing import Any from megatron.core.models.gpt.gpt_model import GPTModel from transformers import GlmMoeDsaForCausalLM @@ -53,6 +54,18 @@ class GLM5Bridge(MegatronModelBridge): >>> provider = bridge.to_megatron_provider() """ + def _should_map_hf_config_field(self, hf_config: Any, hf_name: str, megatron_name: str, value: Any) -> bool: + """Return whether an HF config field should be mapped to the Megatron provider.""" + # Transformers 5.11-5.12 has an upstream GLM config bug: ``GlmMoeDsaConfig`` declares an independent + # ``num_experts=256`` default alongside the model's authoritative ``n_routed_experts`` field. The generic + # config mapping is first-wins, so that injected default can shadow the checkpoint's routed-expert count. + # Transformers 5.13 removed the duplicate field and redirects legacy ``num_experts`` input to + # ``n_routed_experts``. Detect the affected config shape instead of version-gating so this compatibility + # workaround applies only while both fields exist and does not become the default GLM mapping behavior. + if hf_name == "num_experts" and hasattr(hf_config, "num_experts") and hasattr(hf_config, "n_routed_experts"): + return False + return super()._should_map_hf_config_field(hf_config, hf_name, megatron_name, value) + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider: provider = super().provider_bridge(hf_pretrained) hf_config = hf_pretrained.config diff --git a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py index 3ea8ceba72b..2f58aabfb68 100644 --- a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py +++ b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py @@ -17,6 +17,7 @@ from types import SimpleNamespace import pytest +from transformers import GlmMoeDsaConfig from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping, QKVMapping @@ -60,6 +61,42 @@ def _provider_from_hf_config(monkeypatch: pytest.MonkeyPatch, **config_overrides return GLM5Bridge().provider_bridge(SimpleNamespace(config=SimpleNamespace(**config))) +def test_hf_config_ignores_upstream_num_experts_default() -> None: + """GLM-5 ignores the num_experts default injected by affected Transformers versions.""" + hf_config = GlmMoeDsaConfig(n_routed_experts=8) + + provider_kwargs = GLM5Bridge().hf_config_to_provider_kwargs(hf_config) + + assert hf_config.num_experts == 256 + assert provider_kwargs["num_moe_experts"] == hf_config.n_routed_experts == 8 + + +@pytest.mark.parametrize( + ("hf_config", "should_map"), + [ + (SimpleNamespace(num_experts=256, n_routed_experts=8), False), + (SimpleNamespace(num_experts=8), True), + (SimpleNamespace(n_routed_experts=8), True), + ], +) +def test_num_experts_workaround_is_config_shape_specific(hf_config: SimpleNamespace, should_map: bool) -> None: + """The workaround applies only to the conflicting upstream config shape.""" + bridge = GLM5Bridge() + + result = bridge._should_map_hf_config_field(hf_config, "num_experts", "num_moe_experts", 256) + + assert result is should_map + + +def test_megatron_config_export_keeps_generic_moe_aliases() -> None: + """GLM-5 keeps the generic reverse mappings rather than establishing a special default.""" + mapped_config = GLM5Bridge.megatron_to_hf_config(SimpleNamespace(num_moe_experts=8)) + + assert mapped_config["num_experts"] == 8 + assert mapped_config["num_local_experts"] == 8 + assert mapped_config["n_routed_experts"] == 8 + + @pytest.mark.parametrize( ("config_overrides", "expected_topk_freq", "expected_skip_topk_offset"), [ From c2b85a362dfb47a4aa54cbe9a1875316538966d4 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Mon, 3 Aug 2026 19:25:52 +0200 Subject: [PATCH 08/85] cp: `feat(sft): DSv4 Flash offline-packed SFT with THD + CP (5126)` into `r0.6.0` (#5247) Signed-off-by: weijia chen Signed-off-by: NeMo Bot Co-authored-by: weijiac0619 Co-authored-by: Claude Sonnet 4.6 (1M context) --- scripts/training/recipe_metadata.py | 2 +- scripts/training/recipe_runner.py | 2 + .../models/deepseek/deepseek_v4_step.py | 227 ++++++++++++++++++ .../bridge/recipes/deepseek/__init__.py | 2 + .../bridge/recipes/deepseek/deepseek_v4.py | 24 ++ src/megatron/bridge/training/gpt_step.py | 9 +- .../bridge/training/utils/packed_seq_utils.py | 24 ++ .../recipes/test_deepseek_recipes_finetune.py | 2 + tests/unit_tests/models/deepseek/__init__.py | 0 .../models/deepseek/test_deepseek_v4_step.py | 115 +++++++++ 10 files changed, 404 insertions(+), 3 deletions(-) create mode 100644 src/megatron/bridge/models/deepseek/deepseek_v4_step.py create mode 100644 tests/unit_tests/models/deepseek/__init__.py create mode 100644 tests/unit_tests/models/deepseek/test_deepseek_v4_step.py diff --git a/scripts/training/recipe_metadata.py b/scripts/training/recipe_metadata.py index 29c3cd9e2a6..6e296e51731 100644 --- a/scripts/training/recipe_metadata.py +++ b/scripts/training/recipe_metadata.py @@ -63,7 +63,7 @@ LIBRARY_RECIPE_PRECEDENCE_COLLISIONS: frozenset[str] = frozenset() PUBLIC_MODES = frozenset({"pretrain", "sft", "lora", "dora"}) -TEXT_FORWARD_STEPS = frozenset({"gpt_step", "llm_step"}) +TEXT_FORWARD_STEPS = frozenset({"dsv4_step", "gpt_step", "llm_step"}) # Put specific multimodal families before the text default. This registry is # source-agnostic: library and benchmark recipes with the same identity use diff --git a/scripts/training/recipe_runner.py b/scripts/training/recipe_runner.py index dde0f902f99..1980e53419b 100644 --- a/scripts/training/recipe_runner.py +++ b/scripts/training/recipe_runner.py @@ -52,6 +52,7 @@ STEP_FUNCTIONS: dict[str, StepFunctionEntry] = { "audio_lm_step": ("megatron.bridge.training.audio_lm_step", "forward_step"), + "dsv4_step": ("megatron.bridge.models.deepseek.deepseek_v4_step", "forward_step"), "gpt_step": ("megatron.bridge.training.gpt_step", "forward_step"), "llm_step": ("megatron.bridge.training.gpt_step", "forward_step"), "vlm_step": ("megatron.bridge.training.vlm_step", "forward_step"), @@ -66,6 +67,7 @@ STEP_MODALITIES = { "audio_lm_step": "audio", + "dsv4_step": "text", "gpt_step": "text", "llm_step": "text", "vlm_step": "vlm", diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_step.py b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py new file mode 100644 index 00000000000..5e2cefd5549 --- /dev/null +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_step.py @@ -0,0 +1,227 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DSv4-specific training step with contiguous CP partition support. + +DSv4 hybrid attention uses a CSA (Compressed Sparse Attention) compressor that +exchanges boundary hidden states between adjacent CP ranks. This requires +contiguous token assignment (each rank gets a consecutive slice), unlike the +default zigzag interleaved assignment used by standard causal models. + +MCore enforces cp_partition_mode='contiguous' is only valid with dsv4_hybrid +attention (see TransformerConfig validation). Use --step_func dsv4_step for +DSv4 SFT/pretrain with CP > 1. +""" + +import logging +from typing import Iterable + +import torch +from megatron.core import parallel_state +from megatron.core.models.gpt import GPTModel +from megatron.core.pipeline_parallel.utils import ( + is_pp_first_stage, + is_pp_last_stage, + is_vp_first_stage, + is_vp_last_stage, +) +from megatron.core.utils import get_batch_on_this_cp_rank + +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.gpt_step import ( + _create_loss_function, + _current_stage_needs_mtp_inputs_from_layout, + _forward_step_common, + _has_packed_sequence_metadata, + _middle_pp_stage_needs_batch, + _partition_packed_batch_for_cp, + get_batch_from_iterator, +) +from megatron.bridge.training.state import GlobalState + + +logger = logging.getLogger(__name__) + +# DSv4 offline-packed SFT passes cp_partition_mode +# through the batch dict so get_packed_seq_params can forward it to PackedSeqParams. +# These fields are MCore-dev-only, so they live here rather than in generic gpt_step.py. +_DSV4_CURRENT_PACKED_SEQ_PARAM_KEYS = ( + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + "max_seqlen_q", + "max_seqlen_kv", + "total_tokens", + "cp_partition_mode", +) +_DSV4_LEGACY_PACKED_SEQ_PARAM_KEYS = ( + "cu_seqlens", + "cu_seqlens_unpadded", + "cu_seqlens_argmin", + "max_seqlen", + "cu_seqlens_unpadded_argmin", + "total_tokens", + "cp_partition_mode", +) + + +def _packed_metadata_for_forward(batch: dict) -> dict | None: + """Extract packed-sequence metadata for DSv4, including CP partition fields.""" + if batch.get("cu_seqlens_q") is not None: + return {k: batch[k] for k in _DSV4_CURRENT_PACKED_SEQ_PARAM_KEYS if batch.get(k) is not None} + if batch.get("cu_seqlens") is not None: + return {k: batch[k] for k in _DSV4_LEGACY_PACKED_SEQ_PARAM_KEYS if batch.get(k) is not None} + return None + + +# Sequence-length metadata keys — excluded from token-dimension slicing. +_SEQLEN_KEYS = frozenset( + { + "cu_seqlens", + "cu_seqlens_unpadded", + "cu_seqlens_argmin", + "cu_seqlens_unpadded_argmin", + "max_seqlen", + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + "max_seqlen_q", + "max_seqlen_kv", + "token_count", + "attention_mask", + } +) + + +def _partition_packed_batch_contiguous( + batch: dict[str, torch.Tensor], + cp_size: int, +) -> dict[str, torch.Tensor]: + """Slice a consecutive [start, end) token window for this CP rank. + + Only data tensors (tokens, labels, loss_mask, position_ids, etc.) are sliced. + Sequence-length metadata (cu_seqlens, max_seqlen, ...) is intentionally kept + at global values — the DSv4 CSA compressor needs global sequence boundaries to + correctly exchange boundary hidden states between adjacent CP ranks. + This mirrors how zigzag mode leaves cu_seqlens untouched. + + The packed sequence length must be divisible by cp_size — ensure + packed_sequence_size = N * cp_size when running pack_sft_data. + """ + cp_rank = parallel_state.get_context_parallel_rank() + + _data_val = next((v for k, v in batch.items() if v is not None and k not in _SEQLEN_KEYS), None) + if _data_val is None: + return batch # middle PP stage with no data tensors — nothing to slice + + total_tokens = _data_val.size(1) + if total_tokens % cp_size != 0: + raise RuntimeError( + f"Contiguous CP partitioning requires packed sequence length ({total_tokens}) " + f"to be divisible by cp_size ({cp_size}). " + "Set packed_sequence_size to a multiple of cp_size when running pack_sft_data." + ) + local_len = total_tokens // cp_size + start = cp_rank * local_len + end = start + local_len + + for key, val in batch.items(): + if val is None or key in _SEQLEN_KEYS: + continue + batch[key] = val[:, start:end].contiguous() + + return batch + + +def get_batch( # pragma: no cover + data_iterator: Iterable, + cfg: ConfigContainer, + use_mtp: bool = False, + *, + pg_collection, + vp_stage: int | None = None, +): + """get_batch with DSv4 contiguous CP partition support. + + Identical to gpt_step.get_batch but dispatches to contiguous partitioning + when cfg.model.cp_partition_mode == 'contiguous', and injects cp_partition_mode + into the batch so get_packed_seq_params can forward it to PackedSeqParams. + """ + model_cfg = getattr(cfg, "model", None) + vp_size = getattr(model_cfg, "virtual_pipeline_model_parallel_size", None) + is_first = is_pp_first_stage(pg_collection.pp) and ( + vp_stage is None or is_vp_first_stage(vp_stage=vp_stage, vp_size=vp_size) + ) + is_last = is_pp_last_stage(pg_collection.pp) and ( + vp_stage is None or is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size) + ) + is_middle = (not is_first) and (not is_last) + include_full_batch_fields = is_middle and _middle_pp_stage_needs_batch(cfg) + include_mtp_inputs = use_mtp and _current_stage_needs_mtp_inputs_from_layout( + cfg, pg_collection=pg_collection, is_last=is_last, vp_stage=vp_stage + ) + if is_middle and not include_full_batch_fields and not include_mtp_inputs: + return None, None, None, None, None, None + + batch = get_batch_from_iterator( + data_iterator, + include_mtp_inputs=include_mtp_inputs, + skip_getting_attention_mask_from_dataset=getattr( + cfg.dataset, "skip_getting_attention_mask_from_dataset", True + ), + is_first_pp_stage=is_first, + is_last_pp_stage=is_last, + include_full_batch_fields=include_full_batch_fields, + ) + + cp_size = pg_collection.cp.size() + has_packed = _has_packed_sequence_metadata(batch) + if has_packed and cp_size > 1: + cp_mode = getattr(cfg.model, "cp_partition_mode", "zigzag") + if cp_mode == "contiguous": + batch = _partition_packed_batch_contiguous(batch, cp_size) + else: + batch = _partition_packed_batch_for_cp(batch, pg_collection.cp) + # Inject cp_partition_mode so get_packed_seq_params forwards it to PackedSeqParams. + batch["cp_partition_mode"] = cp_mode + else: + batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=pg_collection.cp) + + return ( + batch["tokens"], + batch["labels"], + batch["loss_mask"], + batch.get("attention_mask"), + batch["position_ids"], + _packed_metadata_for_forward(batch), + ) + + +def forward_step( # pragma: no cover + state: GlobalState, + data_iterator: Iterable, + model: GPTModel, + return_schedule_plan: bool = False, +): + """Forward training step for DSv4 with contiguous CP partition support.""" + output, loss_mask = _forward_step_common( + state, data_iterator, model, return_schedule_plan, _get_batch_fn=get_batch + ) + return output, _create_loss_function( + loss_mask, + check_for_nan_in_loss=state.cfg.rerun_state_machine.check_for_nan_in_loss, + check_for_spiky_loss=state.cfg.rerun_state_machine.check_for_spiky_loss, + ) diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 65a744be192..1cf2ac77dd9 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -41,6 +41,7 @@ deepseek_v4_flash_pretrain_mxfp8_config, deepseek_v4_flash_pretrain_mxfp8_gb200_config, deepseek_v4_flash_sft_config, + deepseek_v4_flash_sft_openmath_thinking_packed_config, deepseek_v4_pro_pretrain_config, deepseek_v4_pro_pretrain_mxfp8_config, set_deepseek_v4_pipeline_model_parallel_layout, @@ -72,6 +73,7 @@ "deepseek_v4_flash_pretrain_mxfp8_config", "deepseek_v4_flash_pretrain_mxfp8_gb200_config", "deepseek_v4_flash_sft_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_config", "deepseek_v4_flash_no_mtp_sft_config", "deepseek_v4_pro_pretrain_config", "deepseek_v4_pro_pretrain_mxfp8_config", diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index 34e3d5af138..5ad648049bf 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py @@ -60,6 +60,8 @@ from megatron.bridge.recipes.deepseek.h100.deepseek_v4 import ( deepseek_v4_flash_sft_32gpu_h100_bf16_config as deepseek_v4_flash_sft_config, ) +from megatron.bridge.recipes.utils.dataset_utils import default_openmathinstruct2_thinking_config +from megatron.bridge.training.config import ConfigContainer __all__ = [ @@ -74,9 +76,31 @@ "deepseek_v4_flash_pretrain_mxfp8_gb200_config", "deepseek_v4_flash_pretrain_muon_gb200_config", "deepseek_v4_flash_sft_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_config", "deepseek_v4_pro_pretrain_config", "deepseek_v4_pro_pretrain_mxfp8_config", "DEEPSEEK_V4_PRO_HF_PATH", "DEEPSEEK_V4_FLASH_HF_PATH", "set_deepseek_v4_pipeline_model_parallel_layout", ] + + +def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: + """DSv4 Flash SFT on OpenMathInstruct-2 with thinking channel and offline-packed sequences. + + CoT reasoning goes into the assistant thinking field and the final answer into the + content field. Uses packed sequences for efficient training. + Pre-pack data with ``prepare_gpt_sft_packed_data.py`` before running SFT. + When using CP>1, pass ``model.cp_partition_mode=contiguous`` (required for DSv4 CSA + attention) and ``pad_seq_to_mult=4`` to ensure divisibility by cp_size. + """ + cfg = deepseek_v4_flash_sft_config() + # DSv4 hybrid attention requires contiguous CP partition when CP > 1; + # setting it unconditionally is safe (no-op when context_parallel_size=1). + cfg.model.cp_partition_mode = "contiguous" + cfg.dataset = default_openmathinstruct2_thinking_config( + seq_length=cfg.model.seq_length, + enable_offline_packing=True, + pad_seq_to_mult=2 * cfg.model.context_parallel_size, + ) + return cfg diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 94ff6f1fee7..5d1d42da534 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -348,7 +348,12 @@ def get_batch( def _forward_step_common( - state: GlobalState, data_iterator: Iterable, model: GPTModel, return_schedule_plan: bool = False + state: GlobalState, + data_iterator: Iterable, + model: GPTModel, + return_schedule_plan: bool = False, + *, + _get_batch_fn=None, ) -> tuple[torch.Tensor, torch.Tensor]: """Forward training step. @@ -378,7 +383,7 @@ def _forward_step_common( attention_mask, position_ids, packed_seq_metadata, - ) = get_batch( + ) = (_get_batch_fn or get_batch)( data_iterator, state.cfg, use_mtp, diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py index 0341b65ea38..25323af3773 100644 --- a/src/megatron/bridge/training/utils/packed_seq_utils.py +++ b/src/megatron/bridge/training/utils/packed_seq_utils.py @@ -265,6 +265,14 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen_kv if max_seqlen_kv is not None else max_seqlen_q, total_tokens=batch.get("total_tokens"), qkv_format="thd", + # cp_partition_mode available in dev MCore only. + **( + { + "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), + } + if hasattr(PackedSeqParams, "cp_partition_mode") + else {} + ), ) cu_seqlens_padded = batch["cu_seqlens"].squeeze() @@ -303,6 +311,14 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen, total_tokens=total_tokens, qkv_format="thd", + # cp_partition_mode available in dev MCore only. + **( + { + "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), + } + if hasattr(PackedSeqParams, "cp_partition_mode") + else {} + ), ) else: return PackedSeqParams( @@ -312,4 +328,12 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar max_seqlen_kv=max_seqlen, total_tokens=total_tokens, qkv_format="thd", + # cp_partition_mode available in dev MCore only. + **( + { + "cp_partition_mode": batch.get("cp_partition_mode", "zigzag"), + } + if hasattr(PackedSeqParams, "cp_partition_mode") + else {} + ), ) diff --git a/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py b/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py index 0b4a1f7a000..5510b5afb99 100644 --- a/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py +++ b/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_finetune.py @@ -31,6 +31,7 @@ from megatron.bridge.recipes.deepseek import ( deepseek_v4_flash_no_mtp_sft_config, deepseek_v4_flash_sft_config, + deepseek_v4_flash_sft_openmath_thinking_packed_config, ) from megatron.bridge.recipes.deepseek.h100 import deepseek_v4 as deepseek_v4_h100_module @@ -89,6 +90,7 @@ def _deepseek_v4_toy_model_path() -> str: DEEPSEEK_V4_SFT_RECIPES = [ (deepseek_v4_flash_sft_config, "deepseek_v4_flash_sft", False), (deepseek_v4_flash_no_mtp_sft_config, "deepseek_v4_flash_no_mtp_sft", False), + (deepseek_v4_flash_sft_openmath_thinking_packed_config, "deepseek_v4_flash_sft_openmath_thinking_packed", False), ] diff --git a/tests/unit_tests/models/deepseek/__init__.py b/tests/unit_tests/models/deepseek/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py new file mode 100644 index 00000000000..6ef9b8ebca0 --- /dev/null +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_step.py @@ -0,0 +1,115 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for deepseek_v4_step.py contiguous CP partition logic.""" + +import pytest +import torch + +from megatron.bridge.models.deepseek.deepseek_v4_step import _partition_packed_batch_contiguous + + +def _make_batch(tokens=None, cu_seqlens=None, **extra): + batch = {} + if tokens is not None: + batch["tokens"] = tokens + if cu_seqlens is not None: + batch["cu_seqlens"] = cu_seqlens + batch.update(extra) + return batch + + +class TestPartitionPackedBatchContiguous: + """Tests for _partition_packed_batch_contiguous.""" + + def _run(self, monkeypatch, batch, cp_rank=0, cp_size=2): + monkeypatch.setattr( + "megatron.bridge.models.deepseek.deepseek_v4_step.parallel_state.get_context_parallel_rank", + lambda: cp_rank, + ) + return _partition_packed_batch_contiguous(batch, cp_size) + + def test_rank0_receives_first_half(self, monkeypatch): + """Rank 0 of 2 receives the first half of each data tensor.""" + tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) + labels = torch.arange(8, dtype=torch.long).unsqueeze(0) + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = _make_batch(tokens=tokens, labels=labels, cu_seqlens=cu_seqlens) + result = self._run(monkeypatch, batch, cp_rank=0, cp_size=2) + assert result["tokens"].shape == (1, 4) + assert torch.equal(result["tokens"].squeeze(), torch.arange(4, dtype=torch.long)) + # cu_seqlens kept global — not partitioned (CSA needs global sequence boundaries) + assert result["cu_seqlens"].squeeze().tolist() == [0, 4, 8] + + def test_rank1_receives_second_half(self, monkeypatch): + """Rank 1 of 2 receives the second half.""" + tokens = torch.arange(8, dtype=torch.long).unsqueeze(0) + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = _make_batch(tokens=tokens, cu_seqlens=cu_seqlens) + result = self._run(monkeypatch, batch, cp_rank=1, cp_size=2) + assert result["tokens"].shape == (1, 4) + assert torch.equal(result["tokens"].squeeze(), torch.arange(4, 8, dtype=torch.long)) + # cu_seqlens kept global — not partitioned (CSA needs global sequence boundaries) + assert result["cu_seqlens"].squeeze().tolist() == [0, 4, 8] + + def test_rejects_non_divisible_length(self, monkeypatch): + """Raises RuntimeError when total_tokens is not divisible by cp_size.""" + tokens = torch.arange(5, dtype=torch.long).unsqueeze(0) + batch = _make_batch(tokens=tokens, cu_seqlens=torch.tensor([[0, 5]], dtype=torch.int32)) + with pytest.raises(RuntimeError, match="divisible by cp_size"): + self._run(monkeypatch, batch, cp_size=2) + + def test_middle_pp_stage_returns_unchanged(self, monkeypatch): + """Middle PP stage (all data tensors None) returns batch unchanged.""" + cu_seqlens = torch.tensor([[0, 4, 8]], dtype=torch.int32) + batch = {"tokens": None, "labels": None, "loss_mask": None, "cu_seqlens": cu_seqlens} + result = self._run(monkeypatch, batch, cp_size=2) + assert result is batch + + +class TestPackedMetadataForForward: + """Tests for _packed_metadata_for_forward.""" + + def test_returns_none_for_empty_batch(self): + batch = {"tokens": None, "labels": None} + from megatron.bridge.models.deepseek.deepseek_v4_step import _packed_metadata_for_forward + + assert _packed_metadata_for_forward(batch) is None + + def test_legacy_path_extracts_cu_seqlens_and_cp_partition_mode(self): + from megatron.bridge.models.deepseek.deepseek_v4_step import _packed_metadata_for_forward + + batch = { + "cu_seqlens": torch.tensor([[0, 4, 8]], dtype=torch.int32), + "max_seqlen": torch.tensor([[8]]), + "cp_partition_mode": "contiguous", + "total_tokens": 8, + } + meta = _packed_metadata_for_forward(batch) + assert meta is not None + assert meta["cp_partition_mode"] == "contiguous" + assert "cu_seqlens" in meta + + def test_current_path_with_cu_seqlens_q(self): + from megatron.bridge.models.deepseek.deepseek_v4_step import _packed_metadata_for_forward + + batch = { + "cu_seqlens_q": torch.tensor([[0, 4]], dtype=torch.int32), + "max_seqlen_q": torch.tensor([[4]]), + "cp_partition_mode": "contiguous", + } + meta = _packed_metadata_for_forward(batch) + assert meta is not None + assert meta.get("cp_partition_mode") == "contiguous" + assert "cu_seqlens_q" in meta From 359609f602f131d22901dcb396cce02f1a6cde0c Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 01:28:27 +0200 Subject: [PATCH 09/85] cp: `[Doc] Document latest Mcore toggle `moe_hybridep_pad_uneven_dispatch_inputs` for HybridEP THD training (4213)` into `r0.6.0` (#5159) Signed-off-by: Zhongbo Zhu Signed-off-by: NeMo Bot Co-authored-by: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> --- docs/parallelisms.md | 20 ++++++++++++++++++++ docs/training/packed-sequences.md | 9 +++++++++ 2 files changed, 29 insertions(+) diff --git a/docs/parallelisms.md b/docs/parallelisms.md index 191ff049b0e..660f9ebd17e 100644 --- a/docs/parallelisms.md +++ b/docs/parallelisms.md @@ -223,6 +223,26 @@ model_config = GPTModelProvider( # Apply HybridEP optimization apply_flex_dispatcher_backend(model_config, moe_flex_dispatcher_backend="hybridep") + +# Opt in when local dispatch token counts can differ across HybridEP ranks. +model_config.moe_hybridep_pad_uneven_dispatch_inputs = True +``` + +`moe_hybridep_pad_uneven_dispatch_inputs` is an explicit opt-in. Enable it +when ranks in the HybridEP communication group can enter dispatch with different +token counts, such as dynamically packed THD batches. HybridEP then takes the +group-wide maximum token count, aligns it for the kernel, pads each rank before +dispatch, and removes the padding after combine. + +Leave the option disabled when every rank already supplies the same number of +tokens. Sequence packing alone does not imply that dispatcher inputs are uneven, +and enabling this path unnecessarily adds a MAX all-reduce and padding overhead. +The option does not enable HybridEP or sequence packing; configure those +separately. With `scripts/training/run_recipe.py`, the equivalent trailing CLI +override is: + +```bash +model.moe_hybridep_pad_uneven_dispatch_inputs=true ``` **GPU Architecture Requirements:** diff --git a/docs/training/packed-sequences.md b/docs/training/packed-sequences.md index cb74c923309..c0f09d433dd 100644 --- a/docs/training/packed-sequences.md +++ b/docs/training/packed-sequences.md @@ -113,11 +113,20 @@ The durable constraints for packed sequences in Bridge are: evaluation CP constraints and `CP * TP` when sequence parallelism is enabled - for fine-tuning with CP enabled, per-token loss behavior and reduction settings matter +- HybridEP users whose local dispatch token counts differ across ranks must + explicitly set `model.moe_hybridep_pad_uneven_dispatch_inputs=True`; this + pads only to the group-wide aligned maximum before dispatch and trims the + padding after combine - CUDA-graph-friendly packed metadata requires additional padding constraints Model-family support is not universal. Some families and recipe paths explicitly opt out of packed sequences or related packing modes. +The HybridEP option is not enabled automatically by in-batch packing. Some +packed datasets already produce equal dispatcher sizes, while variable-sized +inputs can also occur without sequence packing. Keeping the option explicit +avoids an unnecessary synchronization and padding cost for equal-sized inputs. + ## Relationship to Long-Sequence Training Packed sequences and long-sequence training are often mentioned together because From aad71778210f7f4917ce1cb972992219ed18b3f1 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 01:42:22 +0200 Subject: [PATCH 10/85] cp: `fix(perf): port r0.5.0 VR200 performance configs (5196)` into `r0.6.0` (#5197) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- scripts/performance/utils/executors.py | 2 +- scripts/performance/utils/overrides.py | 2 +- .../deepseek/vr200/deepseek_v3.py | 6 ++++ .../perf_recipes/gpt_oss/vr200/gpt_oss.py | 13 +++++-- .../bridge/perf_recipes/kimi/vr200/kimi_k2.py | 3 ++ .../bridge/perf_recipes/llama/vr200/llama3.py | 3 ++ .../perf_recipes/llama/vr200/llama31.py | 6 ++-- .../recipes/test_perf_recipe_environment.py | 18 ++++++++++ .../scripts/performance/test_executors.py | 14 ++++++++ .../performance/test_recipe_environment.py | 35 +++++++++++++++++++ 10 files changed, 95 insertions(+), 7 deletions(-) diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py index f520140817f..072bf3b56d5 100644 --- a/scripts/performance/utils/executors.py +++ b/scripts/performance/utils/executors.py @@ -179,7 +179,7 @@ def slurm_executor( log_repo_status_cmd = "bash /opt/Megatron-Bridge/docker/common/print_sha.sh /nemo_run/configs/repo_status.json" custom_bash_cmds.append(log_repo_status_cmd) - numa_divisor = 2 if gpu.lower() in ["gb200", "gb300"] else 4 + numa_divisor = 2 if gpu.lower() in ["gb200", "gb300", "vr200"] else 4 numa_cmd = f"numactl --cpunodebind=$((SLURM_LOCALID/{numa_divisor})) --membind=$((SLURM_LOCALID/{numa_divisor}))" if gpu.lower() in ["b300"] and enable_pct_binding: numa_cmd += " -C $((SLURM_LOCALID * 16)),$((SLURM_LOCALID * 16 + 1))" diff --git a/scripts/performance/utils/overrides.py b/scripts/performance/utils/overrides.py index 11b0981ef68..2ddd07bf70d 100644 --- a/scripts/performance/utils/overrides.py +++ b/scripts/performance/utils/overrides.py @@ -250,7 +250,7 @@ def _apply_flat_cli_environment_compatibility( if connection_override: moe_a2a_overlap = args.moe_a2a_overlap if args.moe_a2a_overlap is not None else base_moe_a2a_overlap max_connections = 32 if base_dispatcher_backend in {"deepep", "hybridep"} else 8 - if gpu in {"b200", "b300", "gb200", "gb300"}: + if gpu in {"b200", "b300", "gb200", "gb300", "vr200"}: max_connections = 32 elif (tp_size > 1 or cp_size > 1) and not moe_a2a_overlap: max_connections = 1 diff --git a/src/megatron/bridge/perf_recipes/deepseek/vr200/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/vr200/deepseek_v3.py index 4d270ca2c77..89844ca409b 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vr200/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vr200/deepseek_v3.py @@ -181,6 +181,9 @@ def deepseek_v3_pretrain_128gpu_vr200_fp8mx_config() -> ConfigContainer: "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, # Keep DeepSeek kernel selection aligned with the measured baseline. "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } @@ -320,6 +323,9 @@ def deepseek_v3_pretrain_256gpu_vr200_fp8mx_config() -> ConfigContainer: "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, # Keep DeepSeek kernel selection aligned with the measured baseline. "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } diff --git a/src/megatron/bridge/perf_recipes/gpt_oss/vr200/gpt_oss.py b/src/megatron/bridge/perf_recipes/gpt_oss/vr200/gpt_oss.py index 846f854e955..f17602c2875 100644 --- a/src/megatron/bridge/perf_recipes/gpt_oss/vr200/gpt_oss.py +++ b/src/megatron/bridge/perf_recipes/gpt_oss/vr200/gpt_oss.py @@ -37,7 +37,7 @@ def gpt_oss_20b_pretrain_8gpu_vr200_nvfp4_config() -> ConfigContainer: cfg.env_vars = { **COMMON_PERF_ENV_VARS, # CUDA stream scheduling for this model and parallel layout. - "CUDA_DEVICE_MAX_CONNECTIONS": 1, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, # CUDA graph and allocator behavior for this recipe. "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. @@ -51,6 +51,9 @@ def gpt_oss_20b_pretrain_8gpu_vr200_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, } @@ -98,6 +101,9 @@ def gpt_oss_20b_pretrain_8gpu_vr200_fp8mx_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, } return cfg @@ -129,7 +135,7 @@ def gpt_oss_20b_pretrain_64gpu_vr200_nvfp4_config() -> ConfigContainer: cfg.env_vars = { **COMMON_PERF_ENV_VARS, # CUDA stream scheduling for this model and parallel layout. - "CUDA_DEVICE_MAX_CONNECTIONS": 1, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, # CUDA graph and allocator behavior for this recipe. "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. @@ -143,6 +149,9 @@ def gpt_oss_20b_pretrain_64gpu_vr200_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, } diff --git a/src/megatron/bridge/perf_recipes/kimi/vr200/kimi_k2.py b/src/megatron/bridge/perf_recipes/kimi/vr200/kimi_k2.py index b156453222a..a0eb466f4a7 100644 --- a/src/megatron/bridge/perf_recipes/kimi/vr200/kimi_k2.py +++ b/src/megatron/bridge/perf_recipes/kimi/vr200/kimi_k2.py @@ -71,5 +71,8 @@ def kimi_k2_pretrain_256gpu_vr200_fp8mx_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, } return cfg diff --git a/src/megatron/bridge/perf_recipes/llama/vr200/llama3.py b/src/megatron/bridge/perf_recipes/llama/vr200/llama3.py index e2058a65b8e..4cea157b66e 100644 --- a/src/megatron/bridge/perf_recipes/llama/vr200/llama3.py +++ b/src/megatron/bridge/perf_recipes/llama/vr200/llama3.py @@ -130,6 +130,9 @@ def llama3_70b_pretrain_64gpu_vr200_bf16_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, } return cfg diff --git a/src/megatron/bridge/perf_recipes/llama/vr200/llama31.py b/src/megatron/bridge/perf_recipes/llama/vr200/llama31.py index 38767268c13..c82a1369814 100644 --- a/src/megatron/bridge/perf_recipes/llama/vr200/llama31.py +++ b/src/megatron/bridge/perf_recipes/llama/vr200/llama31.py @@ -31,7 +31,7 @@ def llama31_405b_pretrain_256gpu_vr200_bf16_config() -> ConfigContainer: cfg.env_vars = { **COMMON_PERF_ENV_VARS, # CUDA stream scheduling for this model and parallel layout. - "CUDA_DEVICE_MAX_CONNECTIONS": 1, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, # CUDA graph and allocator behavior for this recipe. "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. @@ -50,7 +50,7 @@ def llama31_405b_pretrain_256gpu_vr200_fp8mx_config() -> ConfigContainer: cfg.env_vars = { **COMMON_PERF_ENV_VARS, # CUDA stream scheduling for this model and parallel layout. - "CUDA_DEVICE_MAX_CONNECTIONS": 1, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, # CUDA graph and allocator behavior for this recipe. "NCCL_GRAPH_REGISTER": 0, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", @@ -73,7 +73,7 @@ def llama31_405b_pretrain_256gpu_vr200_nvfp4_config() -> ConfigContainer: cfg.env_vars = { **COMMON_PERF_ENV_VARS, # CUDA stream scheduling for this model and parallel layout. - "CUDA_DEVICE_MAX_CONNECTIONS": 1, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, # CUDA graph and allocator behavior for this recipe. "NCCL_GRAPH_REGISTER": 0, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index d6447610005..ebf5c178cdb 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -50,6 +50,18 @@ ("b200", "deepseek_v3_pretrain_256gpu_b200_fp8mx_config"), ("b200", "deepseek_v3_pretrain_256gpu_b200_nvfp4_config"), } +_VR200_CUDNN_LAYERNORM_RECIPES = { + "deepseek_v3_pretrain_128gpu_vr200_fp8mx_config", + "deepseek_v3_pretrain_256gpu_vr200_fp8mx_config", + "gpt_oss_20b_pretrain_8gpu_vr200_fp8mx_config", + "gpt_oss_20b_pretrain_8gpu_vr200_nvfp4_config", + "gpt_oss_20b_pretrain_64gpu_vr200_nvfp4_config", + "kimi_k2_pretrain_256gpu_vr200_fp8mx_config", + "llama3_70b_pretrain_64gpu_vr200_bf16_config", + "nemotron_3_nano_pretrain_8gpu_vr200_bf16_config", + "nemotron_3_nano_pretrain_8gpu_vr200_fp8mx_config", + "nemotron_3_nano_pretrain_8gpu_vr200_nvfp4_config", +} def _function(path: Path, function_name: str) -> ast.FunctionDef: @@ -153,6 +165,12 @@ def test_explicit_environment_invariants_across_all_flat_recipes(): cudnn_names = {"NVTE_NORM_BWD_USE_CUDNN", "NVTE_NORM_FWD_USE_CUDNN"} assert environment.keys().isdisjoint(cudnn_names) or environment.keys() >= cudnn_names + if path.parent.name == "vr200": + assert environment["CUDA_DEVICE_MAX_CONNECTIONS"] == 32 + uses_cudnn_layernorm = function_name in _VR200_CUDNN_LAYERNORM_RECIPES + assert (environment.keys() >= cudnn_names) == uses_cudnn_layernorm + if uses_cudnn_layernorm: + assert all(environment[name] == 1 for name in cudnn_names) hybrid_ep_names = environment.keys() & _HYBRID_EP_ENV_NAMES assert not hybrid_ep_names or hybrid_ep_names == _HYBRID_EP_ENV_NAMES diff --git a/tests/unit_tests/scripts/performance/test_executors.py b/tests/unit_tests/scripts/performance/test_executors.py index 8aee642b5dd..8c9f35e2b7a 100644 --- a/tests/unit_tests/scripts/performance/test_executors.py +++ b/tests/unit_tests/scripts/performance/test_executors.py @@ -190,6 +190,20 @@ def test_recipe_env_vars_are_exported_and_forced_into_container(tmp_path): assert set(recipe_env_vars) <= set(executor.container_env or []) +@pytest.mark.skipif(not HAS_NEMO_RUN, reason="nemo_run not installed") +def test_vr200_slurm_executor_uses_two_gpus_per_numa_node(tmp_path): + executor = slurm_executor( + gpu="vr200", + account="test", + partition="test", + log_dir=str(tmp_path), + nodes=1, + num_gpus_per_node=4, + ) + + assert "SLURM_LOCALID/2" in executor.launcher.template_vars["pre_cmds"] + + @pytest.mark.skipif(not HAS_NEMO_RUN, reason="nemo_run not installed") def test_recipe_env_vars_are_added_to_kubeflow_trainer_environment(monkeypatch): """Kubeflow workers should inherit launcher-resolved recipe variables.""" diff --git a/tests/unit_tests/scripts/performance/test_recipe_environment.py b/tests/unit_tests/scripts/performance/test_recipe_environment.py index b0564b54de2..4890a3b1ec7 100644 --- a/tests/unit_tests/scripts/performance/test_recipe_environment.py +++ b/tests/unit_tests/scripts/performance/test_recipe_environment.py @@ -494,6 +494,41 @@ def test_flat_default_args_leave_inline_recipe_environment_unchanged(): assert recipe.env_vars == original_env +def test_vr200_argparse_overrides_keep_sm100_cuda_connection_count(): + from utils.overrides import _apply_flat_cli_environment_compatibility + + recipe = SimpleNamespace( + env_vars={"CUDA_DEVICE_MAX_CONNECTIONS": 1}, + model=SimpleNamespace( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + ), + ) + args = SimpleNamespace( + gpu="vr200", + moe_a2a_overlap=None, + tensor_model_parallel_size=2, + pipeline_model_parallel_size=None, + context_parallel_size=None, + expert_model_parallel_size=None, + nccl_ub=None, + model_family_name="llama", + model_recipe_name="llama31_405b", + task="pretrain", + ) + + _apply_flat_cli_environment_compatibility( + recipe, + args, + base_dispatcher_backend=None, + base_moe_a2a_overlap=False, + ) + + assert recipe.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] == 32 + + def test_flat_explicit_argparse_compatibility_changes_only_legacy_coupled_values(): from utils.overrides import _apply_flat_cli_environment_compatibility From 9472ca28a447c9202b0828de3e9ecb0ee66247d7 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 01:42:43 +0200 Subject: [PATCH 11/85] cp: `feat(misc): add memory snapshot analysis skill (5150)` into `r0.6.0` (#5198) Signed-off-by: Jaemin Choi Signed-off-by: NeMo Bot Co-authored-by: Jaemin Choi --- docs/fern/versions/nightly.yml | 2 + .../nightly/pages/training/profiling.mdx | 8 +- docs/skills-index.md | 1 + docs/training/profiling.md | 8 +- .../SKILL.md | 395 +++++++++++++ .../evals/evals.json | 31 ++ .../scripts/common.py | 463 ++++++++++++++++ .../scripts/compare_snapshots.py | 499 +++++++++++++++++ .../scripts/parse_snapshot.py | 224 ++++++++ .../scripts/plot_timeline.py | 524 ++++++++++++++++++ .../scripts/replay_step.py | 201 +++++++ .../scripts/replay_to_time.py | 290 ++++++++++ .../skill-card.md | 67 +++ 13 files changed, 2711 insertions(+), 2 deletions(-) create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/SKILL.md create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/evals/evals.json create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/scripts/common.py create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/scripts/parse_snapshot.py create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py create mode 100644 skills/nemo-mbridge-memory-snapshot-analysis/skill-card.md diff --git a/docs/fern/versions/nightly.yml b/docs/fern/versions/nightly.yml index e411c820661..961d6542cbf 100644 --- a/docs/fern/versions/nightly.yml +++ b/docs/fern/versions/nightly.yml @@ -298,6 +298,8 @@ navigation: path: ../../../skills/nemo-mbridge-multi-node-slurm/SKILL.md - page: Templates path: ../../../skills/nemo-mbridge-multi-node-slurm/references/templates.md + - page: Memory Snapshot Analysis + path: ../../../skills/nemo-mbridge-memory-snapshot-analysis/SKILL.md - page: Resiliency path: ../../../skills/nemo-mbridge-resiliency/SKILL.md section: Agent Skills Reference diff --git a/docs/fern/versions/nightly/pages/training/profiling.mdx b/docs/fern/versions/nightly/pages/training/profiling.mdx index e3d080fb9d9..12596139a06 100644 --- a/docs/fern/versions/nightly/pages/training/profiling.mdx +++ b/docs/fern/versions/nightly/pages/training/profiling.mdx @@ -164,7 +164,13 @@ cfg.profiling = ProfilingConfig( ### Analyze Memory Usage -After the run completes, memory snapshots for each specified rank are saved to the designated path. Load these traces using the PyTorch Memory Viz tool to plot memory usage over time and detect bottlenecks or leaks in your training pipeline. +After the run completes, memory snapshots for each specified rank are saved to the designated path, with the rank inserted before the extension (for example, `memory_trace_0.pickle`). Load these traces using the PyTorch Memory Viz tool to plot memory usage over time and detect bottlenecks or leaks in your training pipeline. + +For comparing two runs, attributing peak memory to specific allocation sites, or drilling into a captured out-of-memory snapshot, see the Memory Snapshot Analysis skill, which bundles scripts for replaying and diffing these files. + + +`record_memory_history` and `use_pytorch_profiler` are independent settings. `ProfilerStep#N` annotations inside the snapshot come from the PyTorch profiler, so a snapshot captured with `record_memory_history=True` alone contains no step boundaries. Enable both if you want per-step memory analysis. + ## Optimize Profiling Accuracy diff --git a/docs/skills-index.md b/docs/skills-index.md index d68b13482de..5efff72a0a4 100644 --- a/docs/skills-index.md +++ b/docs/skills-index.md @@ -60,6 +60,7 @@ skills/nemo-mbridge-perf-activation-recompute/SKILL skills/nemo-mbridge-multi-node-slurm/SKILL skills/nemo-mbridge-multi-node-slurm/references/templates +skills/nemo-mbridge-memory-snapshot-analysis/SKILL ``` ```{toctree} diff --git a/docs/training/profiling.md b/docs/training/profiling.md index 1520263cbfd..ef34adb1cff 100644 --- a/docs/training/profiling.md +++ b/docs/training/profiling.md @@ -167,7 +167,13 @@ cfg.profiling = ProfilingConfig( ### Analyze Memory Usage -After the run completes, memory snapshots for each specified rank are saved to the designated path. Load these traces using the PyTorch Memory Viz tool to plot memory usage over time and detect bottlenecks or leaks in your training pipeline. +After the run completes, memory snapshots for each specified rank are saved to the designated path, with the rank inserted before the extension (for example, `memory_trace_0.pickle`). Load these traces using the PyTorch Memory Viz tool to plot memory usage over time and detect bottlenecks or leaks in your training pipeline. + +For comparing two runs, attributing peak memory to specific allocation sites, or drilling into a captured out-of-memory snapshot, see the [Memory Snapshot Analysis skill](../skills/nemo-mbridge-memory-snapshot-analysis/SKILL.md), which bundles scripts for replaying and diffing these files. + +```{note} +`record_memory_history` and `use_pytorch_profiler` are independent settings. `ProfilerStep#N` annotations inside the snapshot come from the PyTorch profiler, so a snapshot captured with `record_memory_history=True` alone contains no step boundaries. Enable both if you want per-step memory analysis. +``` ## Optimize Profiling Accuracy diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/SKILL.md b/skills/nemo-mbridge-memory-snapshot-analysis/SKILL.md new file mode 100644 index 00000000000..807664d001e --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/SKILL.md @@ -0,0 +1,395 @@ +--- +name: nemo-mbridge-memory-snapshot-analysis +description: Analyze and compare PyTorch CUDA memory snapshots produced by Megatron Bridge's ProfilingConfig(record_memory_history=True). Replays the recorded allocation timeline to plot memory over time, compare two runs, and attribute peak memory to source code locations. +license: Apache-2.0 +when_to_use: Debugging an OOM or peak-memory regression from a captured snapshot pickle, plotting GPU memory over time, comparing memory between two configurations, or attributing memory to specific allocations; 'snapshot.pickle', 'snapshot_0.pickle', 'memory snapshot', 'record_memory_history', 'memory_snapshot_path', 'device_traces', 'plot memory usage', 'graph memory over time', 'memory timeline', 'visualize GPU memory', 'memory usage chart', 'why does this config use more memory', 'peak memory analysis', 'OOM pickle'. +--- + +# Memory Snapshot Analysis + +Stable documentation: @docs/training/profiling.md +Related skills: @skills/nemo-mbridge-perf-memory-tuning/SKILL.md (how to *fix* +the memory problem this skill helps you *find*) + +## What It Is + +Bridge writes PyTorch CUDA memory snapshots when `record_memory_history=True` +is set on `ProfilingConfig`. Each snapshot pickle contains: + +| Key | Contents | +|---|---| +| `device_traces` | Chronological alloc/free events with Python stack context — the bulk of the file | +| `segments` | Point-in-time allocator state at the moment of the dump | +| `external_annotations` | `ProfilerStep#N` START/END markers — **only present when the PyTorch profiler also ran** | +| `allocator_settings` | `PYTORCH_CUDA_ALLOC_CONF` and related allocator config | + +The bundled scripts replay `device_traces` to reconstruct memory over time, +group allocations by source location, and diff two snapshots to explain *which +tensors* account for a memory difference. + +PyTorch's own `memory_viz` also renders a timeline from these files and is worth +using for a single snapshot. What it cannot do is put **two runs on one axis** — +which is the whole question when you are asking why config B peaks higher than +config A. `plot_timeline.py` overlays them, and `compare_snapshots.py` then +attributes the gap to specific allocation sites. Add scriptable JSON output on +top and that is the reason this skill exists alongside `memory_viz`. + +**No dependencies** — Python stdlib only, and they run on any `python3` back to +3.9, so a login node or laptop system interpreter is fine (macOS ships 3.9; +`from __future__ import annotations` keeps the newer type syntax from being +evaluated at import). No virtualenv, no `uv sync`, no GPU. The HTML plot loads +Plotly.js from a CDN in the browser, so it adds nothing to the Python +environment. + +## Pick a Script + +| If the question is… | Run | Output | +|---|---|---| +| "what's even in this file?" | `parse_snapshot.py` | text / JSON | +| "**plot / graph / visualize** memory over time" | `plot_timeline.py` | **standalone HTML** | +| "why does B peak higher than A?" | `compare_snapshots.py` | text / JSON | +| "what's live at *this* moment?" | `replay_to_time.py` | text / JSON | +| "what happened during step N?" | `replay_step.py` (needs step markers) | text / JSON | + +`plot_timeline.py` is the only script that writes a file; the rest print to +stdout. It also takes **two** snapshots and overlays them, which is usually the +fastest way to see *where* two runs diverge before asking `compare_snapshots.py` +*why*. + +A good default loop for a memory regression: overview both files, overlay them +on a timeline, read off the time where they split, then drill in at that moment. + +```bash +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py A_0.pickle B_0.pickle --labels A B -o compare.html +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py A_0.pickle B_0.pickle --time 83.0 +``` + +## How Bridge Produces These Files + +Read this before interpreting any numbers. Bridge's capture path differs from a +hand-rolled `_record_memory_history()` call in ways that change what the +numbers mean. + +```python +from megatron.bridge.training.config import ProfilingConfig + +cfg.profiling = ProfilingConfig( + record_memory_history=True, + memory_snapshot_path="/nemo_run/snapshot.pickle", + profile_step_end=15, + profile_ranks=[0], +) +``` + +Five consequences worth internalizing: + +**1. `external_annotations` is populated, but never with `ProfilerStep`.** +This is the single most misleading thing about these files. The list is *not* +empty — a real Bridge pretrain snapshot carries thousands of entries, but all of +them are allocator-level `record_function` markers: + +```text + 1448 nccl:all_gather_into_tensor_coalesced + 588 nccl:all_reduce + 504 nccl:reduce_scatter_tensor_coalesced + 462 nccl:_all_gather_base + 96 enumerate(DataLoader)#_SingleProcessDataLoaderIter.__next__ + 44 Optimizer.step#FusedAdam.step + 24 CustomFSDP.forward +``` + +`ProfilerStep#N` markers come from `torch.profiler`'s step boundaries, and +`record_memory_history` and `use_pytorch_profiler` are independent flags. So a +snapshot can be rich in annotations and still have zero step boundaries. +Consequences: + +- `replay_step.py` cannot run — it needs step boundaries, and exits 1. +- `compare_snapshots.py` automatically falls back to full-trace replay. +- `plot_timeline.py` and `replay_to_time.py` are unaffected — they never needed steps. +- The NCCL/optimizer/dataloader markers *are* still useful for orienting + yourself in the timeline, and `replay_step.py` surfaces them per step when + step boundaries do exist. + +To get per-step drill-down, set `use_pytorch_profiler=True` alongside +`record_memory_history=True`. Note `finalize()` forbids combining it with +`use_nsys_profiler`. + +**2. Where the trace starts depends on the Bridge version.** Since +`fix: record CUDA memory history before snapshot` (`9e9e58b9f`, 2026-04-23), +`start_memory_history_recording()` runs at `setup.py:330`, *before* +`_build_distributed_model()`. Weights and optimizer state are then allocated +**inside** the trace, so they appear as ordinary `alloc` events attributable to +their construction site and `baseline_at_start` is small. + +Snapshots captured before that commit — or by an external callback that enables +recording at train start — begin after model construction instead, and their +`baseline_at_start` is the whole model. A verified pre-fix capture shows a +26.29 GiB baseline for exactly this reason. Check the earliest frames in the +trace to tell which regime you are in. + +**3. The trace is capped at 100k entries.** `trace_alloc_max_entries=100_000` +is hardcoded in `start_memory_history_recording()`. PyTorch keeps the most +recent entries once the cap is hit, so a long enough run silently loses its +beginning. Symptom: a large `unmatched frees` count and a `baseline_at_start` +that no longer reconciles with `nvidia-smi`. If you need the start of the run, +dump earlier by lowering `profile_step_end`. Snapshots captured outside this +code path are not subject to the cap and can be much larger — a real 113-second +capture contained 184,679 events. + +**4. One dump per rank, at `profile_step_end`.** The snapshot is written when +`iteration == profile_step_end`, with the rank appended before the extension. +`memory_snapshot_path="/nemo_run/snapshot.pickle"` with `profile_ranks=[0]` +yields `/nemo_run/snapshot_0.pickle`. Compare like for like — rank 0 against +rank 0. + +**5. OOM produces its own snapshot automatically.** An out-of-memory observer is +attached at capture time and dumps `snapshot_oom_rank-{N}.pickle` at the moment +of failure. This file is usually the most valuable one you have after a crash: +it captures exactly what was live when the allocator gave up. Analyze it with +`parse_snapshot.py` and `replay_to_time.py`. + +## Workflows + +Scripts live in `skills/nemo-mbridge-memory-snapshot-analysis/scripts/`. + +> **Load trusted snapshots only.** Every script below starts by unpickling the +> file, and `pickle.load` executes arbitrary code embedded in it. Analyze +> snapshots produced by your own training runs or by someone you trust — never a +> pickle from an untrusted issue attachment, bucket, or download. There is no +> safe-mode parse: the format requires full deserialization. +> +> Snapshots also carry absolute paths and stack frames from the machine that +> produced them, so check before sharing one outside your organization. + +### 1. Single snapshot overview + +Start here to see what a file actually contains — especially whether it has step +annotations and how many events survived the ring buffer. + +```bash +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/parse_snapshot.py snapshot_0.pickle +``` + +Shows allocator settings, segments totals and fragmentation, baseline memory, +detected steps, event counts by action, and the top 15 allocation sources by +throughput. Add `--json` for machine-readable output. + +### 2. Compare two runs + +The primary use case: explaining why config B uses more memory than config A. + +```bash +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py A_0.pickle B_0.pickle +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py A_0.pickle B_0.pickle --top 20 --frame-depth 2 +``` + +Produces an allocator-settings diff, a side-by-side segments baseline, a +per-step table (or a full-trace summary when annotations are absent), a +drill-down at the largest divergence grouped by source and sorted by delta, and +automatic observations. + +### 3. Visualize the timeline + +```bash +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py snapshot_0.pickle -o timeline.html +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py A_0.pickle B_0.pickle --labels "TP2" "TP4" -o compare.html +``` + +Standalone HTML with a WebGL Plotly timeline. Two toggles matter: + +- **X-axis**: "Wall Clock" vs "Event Index". Event-index mode matches + `memory_viz`, stretching busy iterations and compressing idle gaps so + sawtooth allocation patterns become visible. +- **Y-axis**: "Absolute" (baseline + delta) vs "Delta from trace start". + +Wall-clock mode buckets by time (default 1 ms, `--precision`) preserving min/max +per bucket, so busy traces produce more points. + +### 4. Drill into a specific moment + +The natural follow-up to the plot: spot a divergence visually, then find out +what is responsible. + +```bash +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py snapshot_0.pickle --time 83.0 +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py A_0.pickle B_0.pickle --time 83.0 --frame-depth 2 +``` + +Replays to the given elapsed second and reports baseline/delta/absolute memory, +live allocation counts, and live allocations grouped by source. In comparison +mode it sorts by absolute delta and separately totals only-in-A and only-in-B +sources. + +### 5. Per-step drill-down (requires annotations) + +Only works when the snapshot was captured with `use_pytorch_profiler=True`. + +```bash +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py snapshot_0.pickle --step 7 +python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py snapshot_0.pickle --all-steps +``` + +Per step: alloc/free counts, throughput, peak delta and absolute peak, live +allocations at peak grouped by source, and active annotation phases. + +## Example Output + +```text +--- Segments Baseline --- + Snapshot A Snapshot B Delta + ──────────────────── ────────────── ────────────── ────────────── + Segments 24 24 0 + Reserved (total) 53.21 GiB 73.49 GiB + 20.27 GiB + Allocated 45.10 GiB 62.28 GiB + 17.18 GiB + Active 45.10 GiB 62.28 GiB + 17.18 GiB + Awaiting free 0 B 0 B 0 B + Inactive (reusable) 8.12 GiB 11.21 GiB + 3.09 GiB + Baseline at start 39.61 GiB 52.77 GiB + 13.15 GiB + +--- Full-Trace Summary (no ProfilerStep annotations) --- + Snapshot A Snapshot B Delta(B-A) + ─────────────────────────────────── ────────────── ────────────── ────────────── + Peak delta (from trace start) 9.65 GiB 19.15 GiB + 9.50 GiB + Absolute peak (baseline+delta) 49.26 GiB 71.92 GiB + 22.65 GiB + Alloc count 760 1,840 + 1,080 + +--- Drill-Down: Full Trace (peak diff: +9.50 GiB) --- + Source Peak A Peak B Delta(B-A) + ──────────────────────────────────────────── ────────── ────────── ──────────── + router_forward@router.py:145 1.81 GiB 4.51 GiB + 2.70 GiB + mlp_forward@mlp.py:301 1.59 GiB 4.28 GiB + 2.69 GiB + attention_forward@attention.py:88 1.77 GiB 3.99 GiB + 2.21 GiB +``` + +## Interpreting Results + +All sizes are **binary** — KiB/MiB/GiB, 1024-based — so they line up directly +with `nvidia-smi` and `torch.cuda.memory_allocated`. A decimal "GB" reads about +7.4% larger for the same bytes, which is enough to make a correct cross-check +look like a discrepancy. + +### Segments vs device traces + +`segments` is the allocator's state at dump time — use it for total memory +accounting. `device_traces` is the event log — use it to understand what changes +over time. The two are cross-referenced to derive `baseline_at_start`, which is +what turns relative deltas into real GPU memory figures. A timeline alone cannot +tell you how much memory the model was actually using. + +### Reserved, allocated, active, inactive + +The four segment numbers are distinct and easy to conflate: + +| Term | Meaning | +|---|---| +| Reserved (total) | Memory the allocator holds from the driver | +| Allocated | Blocks currently handed out to tensors | +| Active | Allocated **plus** `active_awaiting_free` — freed by the caller but still held pending stream sync | +| Inactive (reusable) | `reserved - active`; free blocks the allocator can hand out again | + +Fragmentation is reported as `inactive / reserved`. It is deliberately computed +from *active*, not *allocated*: awaiting-free blocks are not reusable yet and +are not fragmentation, so counting them as such over-reports the problem. When +`Awaiting free` is non-zero, expect the two definitions to disagree. + +### Absolute peak + +`absolute_peak = baseline_at_start + cumulative_delta_at_step_start + step_peak_delta`. +Compare this against `nvidia-smi` or the `memory/` metrics Bridge logs to +TensorBoard to sanity-check that the trace covers what you think it covers. + +Per-step replays seed their live set with whatever was already live when the +step opened, so the source table accounts for memory carried into the step +(weights, optimizer state, graph pools, activations held across the boundary) +and not just what the step itself allocated. On a real annotated capture this +recovered about 3.4 GiB per step that the table previously omitted. + +**The table will not sum to `absolute_peak`, and the gap is not a single clean +quantity.** The table shows allocations the trace can see and attribute to a +stack frame; `absolute_peak` also includes memory that predates the trace and +has no frames, and the two are related through `step_start_delta`, which nets +frees of pre-trace memory against later allocations. Treat the table as "what +is attributable and how it ranks", not as a decomposition of the peak. Use the +segments numbers for total accounting. + +### Unmatched frees + +`free_completed` events with no matching `alloc` are allocations that predate +the visible trace. A small count is normal and harmless. A *large* count means +the trace does not reach back to the allocations it is freeing — either the +100k cap dropped the beginning of the run, or recording started after model +construction (see point 2 above). Either way, `baseline_at_start` is doing more +guessing than measuring, so cross-check it against `nvidia-smi` before trusting +absolute numbers. + +### Source grouping + +Allocations are grouped by their first non-internal stack frame, rendered as +`function@file.py:line`. PyTorch, CUDA, Hydra, and stdlib frames are skipped. +When results look over-aggregated (everything attributed to one wrapper), raise +`--frame-depth 2` to split by two levels of user code. + +### CUDA graphs + +A snapshot with dramatically fewer allocs indicates CUDA graph replay. Graphs +pre-allocate at capture time and reuse on replay, so those allocations stop +appearing in `device_traces`. The memory did not disappear — it moved into the +graph's pinned pool, which is why enabling graphs can *raise* peak memory even +as the trace gets quieter. + +A matched pair of Bridge captures — same model and config, graphs off vs on — +shows both halves of this: + +| | Graphs off | Graphs on | +|---|---|---| +| Alloc count | 281,949 | 62,271 | +| Total alloc throughput | 23,683 GiB | 7,682 GiB | +| Absolute peak | 85.89 GiB | 119.80 GiB | + +The trace got 4.5x quieter and peak memory went *up* 34 GiB. Running +`compare_snapshots.py` on the pair attributes most of that increase to the graph +memory pool rather than to any model tensor — which is the signal you want, and +the reason to diff rather than eyeball a single file. See +@skills/nemo-mbridge-perf-cuda-graphs/SKILL.md. + +## Options + +| Flag | Scripts | Description | +|---|---|---| +| `--frame-depth N` | compare, replay_step, replay_to_time | Stack frames per source key (default 1) | +| `--top N` | compare, replay_step, replay_to_time | Top sources to show | +| `--step N` | compare, replay_step | Specific step; compare defaults to largest divergence | +| `--all-steps` | replay_step | Replay every complete step | +| `--time T` | replay_to_time | Elapsed seconds from trace start | +| `--labels L1 L2` | plot_timeline | Trace display labels (default: filenames) | +| `--output PATH` | plot_timeline | Output HTML path | +| `--precision MS` | plot_timeline | Wall-clock bucket width in ms (default 1.0) | +| `--device D` | replay_step | Force device index (default: auto-detect) | +| `--json` | all except plot_timeline | Machine-readable output | + +Diagnostics go to stderr via `logging`; report bodies go to stdout, so `--json` +output pipes cleanly. + +## Troubleshooting + +**"no ProfilerStep annotations found"** — Expected for a Bridge run with +`record_memory_history=True` and no PyTorch profiler, and not a sign of a broken +snapshot: the file still has NCCL and optimizer annotations, just no step +boundaries. Use `compare_snapshots.py` (auto-falls back), `plot_timeline.py`, or +`replay_to_time.py`. Add `use_pytorch_profiler=True` on the next run if you need +per-step analysis. + +**"no device traces found"** — The pickle has only `segments`. Either +`record_memory_history` was False, or the snapshot came from somewhere other +than Bridge's capture path. `parse_snapshot.py` still summarizes segments. + +**Numbers don't match `nvidia-smi`** — Check the unmatched-free count first; a +wrapped ring buffer makes `baseline_at_start` meaningless. Also confirm you are +looking at the rank that actually peaked — with pipeline parallelism the first +PP stage usually holds the most memory, and `profile_ranks` defaults to `[0]`. + +**Steps don't line up between two snapshots** — Matching is by step number, so +different warmup counts pair the wrong steps. Force alignment with `--step N`. + +**Large files** — The whole pickle is deserialized into memory; budget roughly +10-20x the file size in RAM. Loading progress is logged to stderr. diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/evals/evals.json b/skills/nemo-mbridge-memory-snapshot-analysis/evals/evals.json new file mode 100644 index 00000000000..05283d1e2c1 --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/evals/evals.json @@ -0,0 +1,31 @@ +[ + { + "id": "memory-snapshot-analysis-compare-two-runs", + "question": "Use the nemo-mbridge-memory-snapshot-analysis skill. I ran the same Megatron Bridge recipe twice with different parallelism and set ProfilingConfig(record_memory_history=True). I have snapshot_0.pickle from each run and the second one peaks about 20 GB higher. Which script do I run to find out which allocations account for the difference, and how do I make the attribution finer-grained if everything lands on one wrapper function?", + "expected_skill": "nemo-mbridge-memory-snapshot-analysis", + "expected_script": "skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py", + "ground_truth": "Run compare_snapshots.py with the two per-rank pickles: python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py A_0.pickle B_0.pickle. It emits an allocator-settings diff, a side-by-side segments baseline, a per-step or full-trace summary, and a drill-down at the largest divergence with live allocations grouped by source and sorted by delta. When sources are over-aggregated onto a single wrapper, pass --frame-depth 2 to split by two levels of user code, and --top N to widen the list. Compare the same rank in both runs, since profile_ranks defaults to [0] and the files are written as snapshot_{rank}.pickle.", + "expected_behavior": [ + "Read the nemo-mbridge-memory-snapshot-analysis skill before answering.", + "Recommend compare_snapshots.py as the script for the two-snapshot comparison.", + "Explain that --frame-depth 2 gives finer source attribution when grouping is too coarse.", + "Note that the comparison must be rank-matched because snapshots are written per rank.", + "Do not suggest that the user write a new snapshot parsing script from scratch." + ] + }, + { + "id": "memory-snapshot-analysis-missing-step-annotations", + "question": "Use the nemo-mbridge-memory-snapshot-analysis skill. I set ProfilingConfig(record_memory_history=True, memory_snapshot_path='/nemo_run/snapshot.pickle', profile_ranks=[0]) on a Megatron Bridge pretrain job, but when I try the per-step replay on the resulting snapshot_0.pickle it says there are no ProfilerStep annotations. Is my snapshot broken, and what can I still do with it?", + "expected_skill": "nemo-mbridge-memory-snapshot-analysis", + "expected_script": null, + "ground_truth": "The snapshot is not broken. In Megatron Bridge, record_memory_history and use_pytorch_profiler are independent flags, and ProfilerStep#N markers specifically come from the PyTorch profiler's step boundaries. The external_annotations list is not empty in this case — it still carries allocator-level record_function markers such as nccl:all_reduce, nccl:all_gather_into_tensor_coalesced, Optimizer.step#FusedAdam.step, and DataLoader iteration — but it contains zero ProfilerStep entries, so replay_step.py cannot run and exits 1. The other workflows still work: compare_snapshots.py automatically falls back to a full-trace replay, and plot_timeline.py and replay_to_time.py never needed step boundaries. parse_snapshot.py gives the overview. To get per-step drill-down on a future run, additionally set use_pytorch_profiler=True, which cannot be combined with use_nsys_profiler.", + "expected_behavior": [ + "Read the nemo-mbridge-memory-snapshot-analysis skill before answering.", + "State that the snapshot is valid and the missing step markers are expected for this config.", + "Explain that ProfilerStep markers require use_pytorch_profiler=True, not just record_memory_history=True.", + "Avoid claiming that external_annotations is empty; note that NCCL, optimizer, and dataloader markers are still present.", + "Name the workflows that still work without step markers: compare_snapshots fallback, plot_timeline, replay_to_time, parse_snapshot.", + "Mention that use_pytorch_profiler and use_nsys_profiler are mutually exclusive." + ] + } +] diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/scripts/common.py b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/common.py new file mode 100644 index 00000000000..b7e58f67c5b --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/common.py @@ -0,0 +1,463 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared utilities for CUDA memory snapshot analysis.""" + +# Annotations are lazy so the PEP 604 `X | None` syntax below does not need to +# evaluate at import time; these scripts stay runnable on the system python3 +# (3.9 on macOS) even though the repo itself targets 3.10+. +from __future__ import annotations + +import logging +import os +import pickle +import sys +from collections import defaultdict +from dataclasses import dataclass, field + + +# --------------------------------------------------------------------------- +# Frame filtering – paths considered "internal" (not user code) +# --------------------------------------------------------------------------- +_INTERNAL_PATH_FRAGMENTS = ( + "/torch/", + "/torch_", + "torch/nn/modules/module.py", + "torch/autograd/", + "torch/amp/", + "torch/cuda/", + "torch/distributed/", + "/cuda/", + "/hydra/", + "/python3.", + "/importlib/", + "/runpy.py", +) + +_INTERNAL_FUNC_NAMES = {"", "", "_call_impl", "_wrapped_call_impl", "inner"} + + +logger = logging.getLogger(__name__) + +TRUST_EPILOG = ( + "Security: snapshots are unpickled, and pickle.load executes arbitrary code. " + "Only analyze snapshot files from training runs you trust." +) + + +_KIB = 1024 +_MIB = 1024**2 +_GIB = 1024**3 + + +def format_size(num_bytes: int) -> str: + """Human-readable byte size in binary units, e.g. '37.91 GiB', '58.72 MiB'. + + Binary rather than decimal so the numbers line up with ``nvidia-smi`` and + ``torch.cuda.memory_allocated`` reporting. A decimal "GB" reads ~7.4% larger + than the same quantity in GiB, which is enough to make a cross-check against + ``nvidia-smi`` look like a real discrepancy when nothing is wrong. + """ + abs_b = abs(num_bytes) + sign = "-" if num_bytes < 0 else "" + if abs_b >= _GIB: + return f"{sign}{abs_b / _GIB:.2f} GiB" + if abs_b >= _MIB: + return f"{sign}{abs_b / _MIB:.2f} MiB" + if abs_b >= _KIB: + return f"{sign}{abs_b / _KIB:.2f} KiB" + return f"{sign}{abs_b} B" + + +def load_snapshot(path: str) -> dict: + """Load a CUDA memory snapshot pickle file. + + Snapshots are unpickled, and ``pickle.load`` executes arbitrary code, so + only load files produced by a training run you trust. + """ + if not os.path.isfile(path): + logger.error(f"file not found: {path}") + sys.exit(1) + size_mib = os.path.getsize(path) / _MIB + logger.info(f"Loading {path} ({size_mib:.1f} MiB) — unpickling executes code; load trusted files only.") + with open(path, "rb") as f: + data = pickle.load(f) + if not isinstance(data, dict): + logger.error(f"expected dict, got {type(data).__name__}") + sys.exit(1) + return data + + +def find_active_device(snapshot: dict) -> int | None: + """Return the index of the first device_traces entry with events.""" + traces = snapshot.get("device_traces", []) + for i, trace in enumerate(traces): + if trace and len(trace) > 0: + return i + return None + + +def get_profiler_steps(annotations: list) -> dict[int, dict[str, int | None]]: + """Parse ProfilerStep#N annotations into {step_num: {'start': us, 'end': us}}. + + Incomplete steps (missing END) are included with end=None. + """ + steps: dict[int, dict[str, int | None]] = {} + for ann in annotations: + name = ann.get("name", "") + if "ProfilerStep#" not in name: + continue + try: + step_num = int(name.split("#")[1]) + except (IndexError, ValueError): + continue + if step_num not in steps: + steps[step_num] = {"start": None, "end": None} + stage = ann.get("stage", "") + time_us = ann.get("time_us", 0) + if stage == "START": + steps[step_num]["start"] = time_us + elif stage == "END" and time_us > 0: + steps[step_num]["end"] = time_us + return steps + + +def get_step_events(traces: list, start_us: int, end_us: int | None) -> list: + """Filter trace events to a time range [start_us, end_us].""" + if end_us is None: + return [e for e in traces if e["time_us"] >= start_us] + return [e for e in traces if start_us <= e["time_us"] <= end_us] + + +def get_source_key(frames: list, depth: int = 1) -> str: + """Extract a meaningful source key from stack frames. + + Walks the frames top-to-bottom, skipping PyTorch/CUDA internals, + and returns `function@file.py:line` for the first `depth` user-level frames. + """ + relevant = [] + for f in frames: + fn = f.get("filename", "") + name = f.get("name", "") + # Skip internal frames + if any(frag in fn for frag in _INTERNAL_PATH_FRAGMENTS): + continue + if name in _INTERNAL_FUNC_NAMES: + continue + short_file = fn.rsplit("/", 1)[-1] if "/" in fn else fn + relevant.append(f"{name}@{short_file}:{f.get('line', '?')}") + if len(relevant) >= depth: + break + if not relevant: + # Fallback: use top frame even if internal + for f in frames[:depth]: + fn = f.get("filename", "") + short_file = fn.rsplit("/", 1)[-1] if "/" in fn else fn + relevant.append(f"{f.get('name', '?')}@{short_file}:{f.get('line', '?')}") + return " -> ".join(relevant) if relevant else "" + + +# --------------------------------------------------------------------------- +# Baseline computation +# --------------------------------------------------------------------------- + + +@dataclass +class BaselineInfo: + """Baseline memory at the start of profiling. + + Derived by cross-referencing segments (point-in-time allocator state) + with device_traces (alloc/free event log during profiling). + """ + + segments_active: int = 0 # Active bytes in segments at snapshot time + pre_existing_freed: int = 0 # Pre-existing allocs freed during profiling + profiling_still_live: int = 0 # Profiling-period allocs still live at snapshot + baseline_at_start: int = 0 # Estimated memory at profiling start + + +def compute_baseline(snapshot: dict, device_idx: int) -> BaselineInfo: + """Estimate memory allocated before profiling started. + + At snapshot time, segments_active = pre_existing_still_live + profiling_still_live. + At profiling start, memory = pre_existing_still_live + pre_existing_freed + (the freed ones hadn't been freed yet). + """ + segments = snapshot.get("segments", []) + # active_size, not allocated_size: the trace replay below only releases a + # block on free_completed, so a block in active_awaiting_free is still live + # from the replay's point of view. Using allocated_size here would subtract + # those blocks on one side of the equation but not the other. + segments_active = sum(s.get("active_size", s.get("allocated_size", 0)) for s in segments) + + traces = snapshot.get("device_traces", []) + device_trace = traces[device_idx] if device_idx < len(traces) else [] + + allocated: dict[int, int] = {} # addr -> size for profiling-period allocs + pre_existing_freed = 0 + + for ev in device_trace: + action = ev.get("action") + if action == "alloc": + allocated[ev["addr"]] = ev["size"] + elif action == "free_completed": + if ev["addr"] in allocated: + del allocated[ev["addr"]] + else: + pre_existing_freed += ev["size"] + + profiling_still_live = sum(allocated.values()) + baseline_at_start = segments_active - profiling_still_live + pre_existing_freed + + return BaselineInfo( + segments_active=segments_active, + pre_existing_freed=pre_existing_freed, + profiling_still_live=profiling_still_live, + baseline_at_start=baseline_at_start, + ) + + +def compute_step_start_deltas(traces: list, steps: dict) -> dict[int, int]: + """Compute the cumulative memory delta at the start of each profiler step. + + Replays the full trace chronologically, recording the running net delta + (allocs minus frees) at the moment each step begins. Combined with + baseline_at_start, this gives the absolute memory at each step boundary. + """ + step_starts = sorted( + [(info["start"], num) for num, info in steps.items() if info["start"]], + key=lambda x: x[0], + ) + if not step_starts: + return {} + + current_delta = 0 + result: dict[int, int] = {} + step_idx = 0 + + for ev in traces: + t = ev.get("time_us", 0) + while step_idx < len(step_starts) and t >= step_starts[step_idx][0]: + result[step_starts[step_idx][1]] = current_delta + step_idx += 1 + action = ev.get("action") + if action == "alloc": + current_delta += ev["size"] + elif action == "free_completed": + current_delta -= ev["size"] + + # Handle steps that start after the last event + while step_idx < len(step_starts): + result[step_starts[step_idx][1]] = current_delta + step_idx += 1 + + return result + + +# --------------------------------------------------------------------------- +# Replay engine +# --------------------------------------------------------------------------- + + +@dataclass +class ReplayResult: + """Result of replaying device_traces events.""" + + peak_delta: int = 0 + peak_live_set: dict[int, tuple[int, list]] = field(default_factory=dict) + end_delta: int = 0 + total_alloc_bytes: int = 0 + alloc_count: int = 0 + free_count: int = 0 + unmatched_free_count: int = 0 + unmatched_free_bytes: int = 0 + + +def live_set_before(traces: list, cutoff_us: int) -> dict[int, tuple[int, list]]: + """Return allocations still live immediately before *cutoff_us*. + + Used to seed a windowed replay so that memory allocated earlier in the run + and still held (weights, optimizer state, graph pools, activations carried + across a step boundary) is attributable, not just the allocations that + happen to occur inside the window. + """ + live: dict[int, tuple[int, list]] = {} + for e in traces: + if e.get("time_us", 0) >= cutoff_us: + break + action = e.get("action") + if action == "alloc": + live[e["addr"]] = (e["size"], e.get("frames", [])) + elif action == "free_completed": + live.pop(e["addr"], None) + return live + + +def replay_events(events: list, initial_live: dict[int, tuple[int, list]] | None = None) -> ReplayResult: + """Replay alloc/free_completed events to track memory over time. + + Returns peak delta from start, the set of live allocations at peak, + and summary statistics. + + Args: + events: Trace events to replay, in chronological order. + initial_live: Allocations already live when the window opens, as + returned by :func:`live_set_before`. These are carried into + ``peak_live_set`` so source attribution accounts for them, but they + do not contribute to ``peak_delta``, which stays a delta measured + from the start of the window. + """ + live: dict[int, tuple[int, list]] = dict(initial_live) if initial_live else {} # addr -> (size, frames) + current_delta = 0 + peak_delta = 0 + # Seeded state is the peak until an alloc beats it, so a window that only + # frees still reports the allocations that were live when it opened. + peak_snapshot: dict[int, tuple[int, list]] = dict(live) + total_alloc = 0 + alloc_count = 0 + free_count = 0 + unmatched_frees = 0 + unmatched_free_bytes = 0 + + for e in events: + action = e["action"] + if action == "alloc": + addr = e["addr"] + size = e["size"] + frames = e.get("frames", []) + live[addr] = (size, frames) + current_delta += size + total_alloc += size + alloc_count += 1 + if current_delta > peak_delta: + peak_delta = current_delta + # Snapshot the live set at peak — copy the dict + peak_snapshot = dict(live) + elif action == "free_completed": + addr = e["addr"] + size = e["size"] + free_count += 1 + if addr in live: + current_delta -= live[addr][0] + del live[addr] + else: + # Pre-existing allocation freed during trace + current_delta -= size + unmatched_frees += 1 + unmatched_free_bytes += size + + return ReplayResult( + peak_delta=peak_delta, + peak_live_set=peak_snapshot, + end_delta=current_delta, + total_alloc_bytes=total_alloc, + alloc_count=alloc_count, + free_count=free_count, + unmatched_free_count=unmatched_frees, + unmatched_free_bytes=unmatched_free_bytes, + ) + + +def group_by_source(live_set: dict[int, tuple[int, list]], depth: int = 1) -> list[tuple[str, int, int]]: + """Group a live allocation set by source key. + + Returns sorted list of (source_key, total_bytes, count), descending by bytes. + """ + groups: dict[str, list[int]] = defaultdict(lambda: [0, 0]) + for addr, (size, frames) in live_set.items(): + key = get_source_key(frames, depth) + groups[key][0] += size + groups[key][1] += 1 + result = [(key, total, count) for key, (total, count) in groups.items()] + result.sort(key=lambda x: -x[1]) + return result + + +@dataclass +class TimeReplayResult: + """Result of replaying device_traces up to a specific elapsed time.""" + + live_set: dict[int, tuple[int, list]] = field(default_factory=dict) + current_delta: int = 0 + alloc_count: int = 0 + free_count: int = 0 + unmatched_free_count: int = 0 + + +def replay_to_time(events: list, target_elapsed_us: int) -> TimeReplayResult: + """Replay alloc/free events up to *target_elapsed_us* from profiling start. + + *Profiling start* is defined as the timestamp of the very first event in + the trace (which is typically a ``segment_map``), so ``--time T`` means + T seconds from when profiling began — matching external logs and the + timeline plot's x-axis. + + Returns the live allocation set (addr -> (size, frames)) and cumulative + delta at that point, suitable for source grouping with ``group_by_source``. + """ + if not events: + return TimeReplayResult() + t0 = events[0].get("time_us", 0) + + cutoff = t0 + target_elapsed_us + live: dict[int, tuple[int, list]] = {} + current_delta = 0 + alloc_count = 0 + free_count = 0 + unmatched = 0 + + for e in events: + if e.get("time_us", 0) > cutoff: + break + action = e["action"] + if action == "alloc": + addr = e["addr"] + size = e["size"] + live[addr] = (size, e.get("frames", [])) + current_delta += size + alloc_count += 1 + elif action == "free_completed": + addr = e["addr"] + free_count += 1 + if addr in live: + current_delta -= live[addr][0] + del live[addr] + else: + current_delta -= e["size"] + unmatched += 1 + + return TimeReplayResult( + live_set=live, + current_delta=current_delta, + alloc_count=alloc_count, + free_count=free_count, + unmatched_free_count=unmatched, + ) + + +def get_step_annotations(annotations: list, start_us: int, end_us: int | None) -> dict[str, int]: + """Count annotation names active during a step's time range.""" + counts: dict[str, int] = defaultdict(int) + for ann in annotations: + if ann.get("name", "").startswith("ProfilerStep"): + continue + t = ann.get("time_us", 0) + if t < start_us: + continue + if end_us is not None and t > end_us: + continue + if ann.get("stage") == "START": + counts[ann["name"]] += 1 + return dict(counts) diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py new file mode 100644 index 00000000000..c9456d40e3c --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py @@ -0,0 +1,499 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare two PyTorch CUDA memory snapshots side-by-side.""" + +# Annotations are lazy so the PEP 604 `X | None` syntax below does not need to +# evaluate at import time; these scripts stay runnable on the system python3 +# (3.9 on macOS) even though the repo itself targets 3.10+. +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys + + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from common import ( + TRUST_EPILOG, + compute_baseline, + compute_step_start_deltas, + find_active_device, + format_size, + get_profiler_steps, + get_step_events, + group_by_source, + live_set_before, + load_snapshot, + replay_events, +) + + +logger = logging.getLogger(__name__) + + +def _normalize_setting(value: object) -> object: + """Render a setting in a stable, comparable form. + + ``roundup_power2_divisions`` arrives as a dict whose key order is not + meaningful, which makes a naive ``!=`` report spurious differences and reads + badly in the report. Sorting it gives a stable representation so the setting + can be compared and shown rather than suppressed — it changes allocator + rounding and is a legitimate explanation for two runs differing in reserved + memory or fragmentation. + """ + if isinstance(value, dict): + return dict(sorted(value.items(), key=lambda kv: str(kv[0]))) + return value + + +def diff_settings(settings_a: dict, settings_b: dict) -> list: + """Return list of (key, val_a, val_b) for differing settings.""" + all_keys = sorted(set(list(settings_a.keys()) + list(settings_b.keys()))) + diffs = [] + for k in all_keys: + va = _normalize_setting(settings_a.get(k)) + vb = _normalize_setting(settings_b.get(k)) + if va != vb: + diffs.append((k, va, vb)) + return diffs + + +def compare( + path_a: str, + path_b: str, + step: int | None = None, + top_n: int = 15, + frame_depth: int = 1, + as_json: bool = False, +) -> None: + """Compare two snapshots and drill into the step with the largest divergence. + + Falls back to a full-trace replay when neither snapshot carries matching + ``ProfilerStep`` annotations, which is the common case for Megatron Bridge + runs that enable ``record_memory_history`` without the PyTorch profiler. + + Args: + path_a: Path to the first snapshot pickle file. + path_b: Path to the second snapshot pickle file. + step: Step to drill into. Defaults to the largest peak difference. + top_n: Number of top allocation sources to show. + frame_depth: Stack frames to include in each source key. + as_json: Emit machine-readable JSON instead of a text report. + """ + snap_a = load_snapshot(path_a) + snap_b = load_snapshot(path_b) + + # --- Config diff --- + settings_a = snap_a.get("allocator_settings", {}) + settings_b = snap_b.get("allocator_settings", {}) + config_diffs = diff_settings(settings_a, settings_b) + + # --- Segments --- + def seg_summary(snap): + # allocated_size counts blocks currently handed out; active_size also + # counts active_awaiting_free (freed by the user, still held pending + # stream sync). Reusable memory is total - active, so computing inactive + # from allocated would report awaiting-free blocks as fragmentation. + segs = snap.get("segments", []) + total = sum(s.get("total_size", 0) for s in segs) + allocated = sum(s.get("allocated_size", 0) for s in segs) + active = sum(s.get("active_size", s.get("allocated_size", 0)) for s in segs) + return { + "count": len(segs), + "total": total, + "allocated": allocated, + "active": active, + "awaiting_free": active - allocated, + "inactive": total - active, + } + + seg_a = seg_summary(snap_a) + seg_b = seg_summary(snap_b) + + # --- Steps --- + ann_a = snap_a.get("external_annotations", []) + ann_b = snap_b.get("external_annotations", []) + steps_a = get_profiler_steps(ann_a) + steps_b = get_profiler_steps(ann_b) + + # Find matching complete steps + complete_a = {n for n, s in steps_a.items() if s["start"] and s["end"]} + complete_b = {n for n, s in steps_b.items() if s["start"] and s["end"]} + matching = sorted(complete_a & complete_b) + + # --- Device traces --- + dev_a = find_active_device(snap_a) + dev_b = find_active_device(snap_b) + traces_a = snap_a["device_traces"][dev_a] if dev_a is not None else [] + traces_b = snap_b["device_traces"][dev_b] if dev_b is not None else [] + + # --- Baseline memory --- + baseline_a = compute_baseline(snap_a, dev_a) if dev_a is not None else None + baseline_b = compute_baseline(snap_b, dev_b) if dev_b is not None else None + base_a = baseline_a.baseline_at_start if baseline_a else 0 + base_b = baseline_b.baseline_at_start if baseline_b else 0 + + # --- Per-step starting deltas --- + deltas_a = compute_step_start_deltas(traces_a, steps_a) + deltas_b = compute_step_start_deltas(traces_b, steps_b) + + # --- Replay each matching step --- + step_results = [] + for num in matching: + events_a = get_step_events(traces_a, steps_a[num]["start"], steps_a[num]["end"]) + events_b = get_step_events(traces_b, steps_b[num]["start"], steps_b[num]["end"]) + # Seed each replay with what was already live when the step opened, so + # the source table accounts for memory carried into the step rather than + # only what the step itself allocated. + replay_a = replay_events(events_a, initial_live=live_set_before(traces_a, steps_a[num]["start"])) + replay_b = replay_events(events_b, initial_live=live_set_before(traces_b, steps_b[num]["start"])) + abs_peak_a = base_a + deltas_a.get(num, 0) + replay_a.peak_delta + abs_peak_b = base_b + deltas_b.get(num, 0) + replay_b.peak_delta + step_results.append( + { + "step": num, + "peak_a": replay_a.peak_delta, + "peak_b": replay_b.peak_delta, + "peak_diff": replay_b.peak_delta - replay_a.peak_delta, + "abs_peak_a": abs_peak_a, + "abs_peak_b": abs_peak_b, + "abs_peak_diff": abs_peak_b - abs_peak_a, + "allocs_a": replay_a.alloc_count, + "allocs_b": replay_b.alloc_count, + "throughput_a": replay_a.total_alloc_bytes, + "throughput_b": replay_b.total_alloc_bytes, + "replay_a": replay_a, + "replay_b": replay_b, + } + ) + + # --- Full-trace fallback when no step annotations --- + full_trace_mode = not matching and traces_a and traces_b + full_trace_result = None + if full_trace_mode: + logger.info(" No matching ProfilerStep annotations; falling back to full-trace replay.") + replay_a_full = replay_events(traces_a) + replay_b_full = replay_events(traces_b) + abs_peak_a_full = base_a + replay_a_full.peak_delta + abs_peak_b_full = base_b + replay_b_full.peak_delta + full_trace_result = { + "peak_a": replay_a_full.peak_delta, + "peak_b": replay_b_full.peak_delta, + "peak_diff": replay_b_full.peak_delta - replay_a_full.peak_delta, + "abs_peak_a": abs_peak_a_full, + "abs_peak_b": abs_peak_b_full, + "allocs_a": replay_a_full.alloc_count, + "allocs_b": replay_b_full.alloc_count, + "throughput_a": replay_a_full.total_alloc_bytes, + "throughput_b": replay_b_full.total_alloc_bytes, + "unmatched_a": replay_a_full.unmatched_free_count, + "unmatched_b": replay_b_full.unmatched_free_count, + "replay_a": replay_a_full, + "replay_b": replay_b_full, + } + + # --- Pick drill-down step --- + if step is not None: + drill_step = step + if step not in [r["step"] for r in step_results]: + logger.warning(f"step {step} not in matching complete steps {matching}") + drill_step = None + elif step_results: + # Rank by absolute peak, not the within-step delta: a step can allocate + # modestly yet still sit at the run's memory high-water mark because of + # what it inherited, and that is the step worth drilling into. + drill_step = max(step_results, key=lambda r: abs(r["abs_peak_diff"]))["step"] + else: + drill_step = None + + # --- Drill-down source comparison --- + drill_data = None + + # Helper to build drill_data from two replay results + def _build_drill_data(label, replay_a_res, replay_b_res, peak_diff): + sources_a = group_by_source(replay_a_res.peak_live_set, depth=frame_depth) + sources_b = group_by_source(replay_b_res.peak_live_set, depth=frame_depth) + src_map_a = {k: (total, count) for k, total, count in sources_a} + src_map_b = {k: (total, count) for k, total, count in sources_b} + all_keys = sorted( + set(list(src_map_a.keys()) + list(src_map_b.keys())), + key=lambda k: -abs(src_map_b.get(k, (0, 0))[0] - src_map_a.get(k, (0, 0))[0]), + ) + return { + "step": label, + "peak_diff": peak_diff, + "sources": [ + { + "source": k, + "peak_a": src_map_a.get(k, (0, 0))[0], + "peak_b": src_map_b.get(k, (0, 0))[0], + "delta": src_map_b.get(k, (0, 0))[0] - src_map_a.get(k, (0, 0))[0], + } + for k in all_keys + ], + } + + if drill_step is not None: + dr = next((r for r in step_results if r["step"] == drill_step), None) + if dr: + drill_data = _build_drill_data( + drill_step, + dr["replay_a"], + dr["replay_b"], + dr["peak_diff"], + ) + elif full_trace_result: + drill_data = _build_drill_data( + "full-trace", + full_trace_result["replay_a"], + full_trace_result["replay_b"], + full_trace_result["peak_diff"], + ) + + # --- Observations --- + observations = [] + + # Alloc ratio observations (per-step or full-trace) + obs_items = step_results or ([full_trace_result] if full_trace_result else []) + for r in obs_items: + if r["allocs_a"] > 0 and r["allocs_b"] > 0: + ratio = max(r["allocs_a"], r["allocs_b"]) / min(r["allocs_a"], r["allocs_b"]) + if ratio > 10: + fewer = "A" if r["allocs_a"] < r["allocs_b"] else "B" + scope = f"step {r['step']}" if "step" in r and isinstance(r.get("step"), int) else "the full trace" + observations.append( + f"Snapshot {fewer} has ~{ratio:.0f}x fewer allocs in {scope} " + f"({min(r['allocs_a'], r['allocs_b']):,} vs {max(r['allocs_a'], r['allocs_b']):,}), " + f"suggesting CUDA graph replay is active." + ) + if drill_data and drill_data["sources"]: + top5_delta = sum(abs(s["delta"]) for s in drill_data["sources"][:5]) + scope = f"step {drill_data['step']}" if isinstance(drill_data["step"], int) else "the full trace" + observations.append(f"Top 5 sources account for {format_size(top5_delta)} of the peak difference in {scope}.") + + # --- Output --- + if as_json: + result = { + "file_a": path_a, + "file_b": path_b, + "config_diffs": config_diffs, + "segments_a": seg_a, + "segments_b": seg_b, + "baseline_a": baseline_a.baseline_at_start if baseline_a else None, + "baseline_b": baseline_b.baseline_at_start if baseline_b else None, + "matching_steps": matching, + "step_results": [{k: v for k, v in r.items() if k not in ("replay_a", "replay_b")} for r in step_results], + "full_trace": ( + {k: v for k, v in full_trace_result.items() if k not in ("replay_a", "replay_b")} + if full_trace_result + else None + ), + "drill_down": drill_data, + "observations": observations, + } + print(json.dumps(result, indent=2, default=str)) + return + + # --- Text output --- + name_a = os.path.basename(path_a) + name_b = os.path.basename(path_b) + + print("=" * 70) + print(" CUDA Memory Snapshot Comparison") + print("=" * 70) + print(f"\n A: {name_a}") + print(f" B: {name_b}") + + # Config diff + print("\n--- Allocator Settings Diff ---") + if not config_diffs: + print(" (identical)") + else: + for k, va, vb in config_diffs: + print(f" {k}:") + print(f" A: {va}") + print(f" B: {vb}") + + # Segments + print("\n--- Segments Baseline ---") + print(f" {'':20s} {'Snapshot A':>14} {'Snapshot B':>14} {'Delta':>14}") + print(f" {'─' * 20} {'─' * 14} {'─' * 14} {'─' * 14}") + for label, ka, kb in [ + ("Segments", seg_a["count"], seg_b["count"]), + ]: + delta = kb - ka + sign = "+" if delta > 0 else "" + print(f" {label:20s} {ka:>14} {kb:>14} {sign}{delta:>13}") + for label, ka, kb in [ + ("Reserved (total)", seg_a["total"], seg_b["total"]), + ("Allocated", seg_a["allocated"], seg_b["allocated"]), + ("Active", seg_a["active"], seg_b["active"]), + ("Awaiting free", seg_a["awaiting_free"], seg_b["awaiting_free"]), + ("Inactive (reusable)", seg_a["inactive"], seg_b["inactive"]), + ("Baseline at start", base_a, base_b), + ]: + delta = kb - ka + sign = "+" if delta > 0 else "" + print(f" {label:20s} {format_size(ka):>14} {format_size(kb):>14} {sign}{format_size(delta):>13}") + + # Full-trace summary (when no step annotations) + if full_trace_mode and full_trace_result: + ftr = full_trace_result + print("\n--- Full-Trace Summary (no ProfilerStep annotations) ---") + print(f" {'':35s} {'Snapshot A':>14} {'Snapshot B':>14} {'Delta(B-A)':>14}") + print(f" {'─' * 35} {'─' * 14} {'─' * 14} {'─' * 14}") + for label, va, vb, is_size in [ + ("Peak delta (from trace start)", ftr["peak_a"], ftr["peak_b"], True), + ("Absolute peak (baseline+delta)", ftr["abs_peak_a"], ftr["abs_peak_b"], True), + ("Total alloc throughput", ftr["throughput_a"], ftr["throughput_b"], True), + ]: + diff = vb - va + sign = "+" if diff > 0 else "" + if is_size: + print(f" {label:35s} {format_size(va):>14} {format_size(vb):>14} {sign}{format_size(diff):>13}") + else: + print(f" {label:35s} {va:>14,} {vb:>14,} {sign}{diff:>13,}") + for label, va, vb in [ + ("Alloc count", ftr["allocs_a"], ftr["allocs_b"]), + ("Unmatched frees (pre-existing)", ftr["unmatched_a"], ftr["unmatched_b"]), + ]: + diff = vb - va + sign = "+" if diff > 0 else "" + print(f" {label:35s} {va:>14,} {vb:>14,} {sign}{diff:>13,}") + + # Per-step summary + print("\n--- Per-Step Summary ---") + if not step_results and not full_trace_mode: + print(" (no matching complete steps)") + elif not step_results and full_trace_mode: + print(" (skipped — see full-trace summary above)") + else: + has_baseline = base_a > 0 or base_b > 0 + if has_baseline: + print( + f" {'Step':>5} {'Peak Δ A':>12} {'Peak Δ B':>12} " + f"{'Diff(B-A)':>12} {'Abs Peak A':>12} {'Abs Peak B':>12} " + f"{'Allocs A':>10} {'Allocs B':>10}" + ) + print(f" {'─' * 5} {'─' * 12} {'─' * 12} {'─' * 12} {'─' * 12} {'─' * 12} {'─' * 10} {'─' * 10}") + else: + print( + f" {'Step':>5} {'Peak A':>12} {'Peak B':>12} " + f"{'Diff(B-A)':>12} {'Allocs A':>10} {'Allocs B':>10} " + f"{'Thruput A':>12} {'Thruput B':>12}" + ) + print(f" {'─' * 5} {'─' * 12} {'─' * 12} {'─' * 12} {'─' * 10} {'─' * 10} {'─' * 12} {'─' * 12}") + for r in step_results: + diff = r["peak_diff"] + sign = "+" if diff > 0 else "" + if has_baseline: + print( + f" {r['step']:>5} " + f"{format_size(r['peak_a']):>12} " + f"{format_size(r['peak_b']):>12} " + f"{sign}{format_size(diff):>11} " + f"{format_size(r['abs_peak_a']):>12} " + f"{format_size(r['abs_peak_b']):>12} " + f"{r['allocs_a']:>10,} " + f"{r['allocs_b']:>10,}" + ) + else: + print( + f" {r['step']:>5} " + f"{format_size(r['peak_a']):>12} " + f"{format_size(r['peak_b']):>12} " + f"{sign}{format_size(diff):>11} " + f"{r['allocs_a']:>10,} " + f"{r['allocs_b']:>10,} " + f"{format_size(r['throughput_a']):>12} " + f"{format_size(r['throughput_b']):>12}" + ) + + # Excluded steps + excluded_a = sorted(set(steps_a.keys()) - complete_a) + excluded_b = sorted(set(steps_b.keys()) - complete_b) + only_a = sorted(complete_a - complete_b) + only_b = sorted(complete_b - complete_a) + notes = [] + if excluded_a or excluded_b: + all_exc = sorted(set(excluded_a) | set(excluded_b)) + notes.append(f"Step(s) {all_exc} excluded (incomplete, no END annotation)") + if only_a: + notes.append(f"Step(s) {only_a} only in A") + if only_b: + notes.append(f"Step(s) {only_b} only in B") + for n in notes: + print(f" * {n}") + + # Drill-down + if drill_data: + step_label = f"Step {drill_data['step']}" if isinstance(drill_data["step"], int) else "Full Trace" + print( + f"\n--- Drill-Down: {step_label} " + f"(peak diff: {'+' if drill_data['peak_diff'] > 0 else ''}" + f"{format_size(drill_data['peak_diff'])}) ---" + ) + sources = drill_data["sources"][:top_n] + if sources: + print(f" {'Source':50s} {'Peak A':>12} {'Peak B':>12} {'Delta(B-A)':>12}") + print(f" {'─' * 50} {'─' * 12} {'─' * 12} {'─' * 12}") + for s in sources: + delta = s["delta"] + sign = "+" if delta > 0 else "" + src_display = s["source"][:50] + print( + f" {src_display:50s} " + f"{format_size(s['peak_a']):>12} " + f"{format_size(s['peak_b']):>12} " + f"{sign}{format_size(delta):>11}" + ) + print(f" ({len(drill_data['sources'])} total sources, showing top {top_n})") + + # Observations + if observations: + print("\n--- Observations ---") + for i, obs in enumerate(observations, 1): + print(f" {i}. {obs}") + + print() + + +def main() -> None: + """Parse arguments and run the two-snapshot comparison.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + parser = argparse.ArgumentParser( + description="Compare two CUDA memory snapshots.", + epilog=TRUST_EPILOG, + ) + parser.add_argument("pickle_a", help="Path to snapshot A") + parser.add_argument("pickle_b", help="Path to snapshot B") + parser.add_argument("--step", type=int, help="Specific step to drill into (default: largest divergence)") + parser.add_argument("--top", type=int, default=15, help="Number of top sources in drill-down (default: 15)") + parser.add_argument( + "--frame-depth", type=int, default=1, help="Stack frame depth for source grouping (default: 1)" + ) + parser.add_argument("--json", action="store_true", help="Output as JSON") + args = parser.parse_args() + + compare( + args.pickle_a, args.pickle_b, step=args.step, top_n=args.top, frame_depth=args.frame_depth, as_json=args.json + ) + + +if __name__ == "__main__": + main() diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/scripts/parse_snapshot.py b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/parse_snapshot.py new file mode 100644 index 00000000000..3ccda065edf --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/parse_snapshot.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-file overview of a PyTorch CUDA memory snapshot.""" + +# Annotations are lazy so the PEP 604 `X | None` syntax below does not need to +# evaluate at import time; these scripts stay runnable on the system python3 +# (3.9 on macOS) even though the repo itself targets 3.10+. +from __future__ import annotations + +import argparse +import json +import os +import sys +from collections import defaultdict + + +# Allow running as script from any directory +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from common import ( + TRUST_EPILOG, + compute_baseline, + find_active_device, + format_size, + get_profiler_steps, + get_source_key, + load_snapshot, +) + + +def print_overview(snapshot: dict, path: str, as_json: bool = False) -> None: + """Print a one-file overview: settings, segments, baseline, steps, top sources. + + Args: + snapshot: Deserialized snapshot dictionary. + path: Path the snapshot was loaded from, used for the file-size line. + as_json: Emit machine-readable JSON instead of a text report. + """ + file_size = os.path.getsize(path) + + # --- Segments --- + # allocated_size counts blocks currently handed out; active_size also counts + # active_awaiting_free (freed by the user, still held pending stream sync). + # Reusable memory is total - active, so computing inactive from allocated + # would report awaiting-free blocks as fragmentation. + segments = snapshot.get("segments", []) + total_size = sum(s.get("total_size", 0) for s in segments) + allocated_size = sum(s.get("allocated_size", 0) for s in segments) + active_size = sum(s.get("active_size", s.get("allocated_size", 0)) for s in segments) + awaiting_free_size = active_size - allocated_size + inactive_size = total_size - active_size + + # --- Profiler steps --- + annotations = snapshot.get("external_annotations", []) + steps = get_profiler_steps(annotations) + + # --- Device traces --- + device_idx = find_active_device(snapshot) + device_traces = snapshot.get("device_traces", []) + traces = device_traces[device_idx] if device_idx is not None and device_idx < len(device_traces) else [] + action_counts = defaultdict(int) + for e in traces: + action_counts[e["action"]] += 1 + + # --- Top allocation sources by throughput --- + source_throughput = defaultdict(lambda: [0, 0]) # key -> [bytes, count] + for e in traces: + if e["action"] == "alloc": + key = get_source_key(e.get("frames", [])) + source_throughput[key][0] += e["size"] + source_throughput[key][1] += 1 + top_sources = sorted(source_throughput.items(), key=lambda x: -x[1][0])[:15] + + # --- Baseline memory --- + baseline = compute_baseline(snapshot, device_idx) if device_idx is not None else None + + if as_json: + result = { + "file": path, + "file_size_bytes": file_size, + "segments": { + "count": len(segments), + "total_size": total_size, + "allocated_size": allocated_size, + "active_size": active_size, + "awaiting_free_size": awaiting_free_size, + "inactive_size": inactive_size, + }, + "baseline": { + "baseline_at_start": baseline.baseline_at_start, + "segments_active": baseline.segments_active, + "pre_existing_freed": baseline.pre_existing_freed, + "profiling_still_live": baseline.profiling_still_live, + } + if baseline + else None, + "profiler_steps": { + num: { + "start_us": info["start"], + "end_us": info["end"], + "duration_ms": round((info["end"] - info["start"]) / 1000, 1) + if info["start"] and info["end"] + else None, + "complete": info["end"] is not None, + } + for num, info in sorted(steps.items()) + }, + "device_traces": { + "device": device_idx, + "event_count": len(traces), + "actions": dict(action_counts), + }, + "top_sources": [{"source": k, "throughput_bytes": v[0], "alloc_count": v[1]} for k, v in top_sources], + } + print(json.dumps(result, indent=2)) + return + + # --- Print report --- + print("=" * 70) + print(" CUDA Memory Snapshot Overview") + print("=" * 70) + + print(f"\n File: {path}") + print(f" Size: {format_size(file_size)}") + + # Allocator settings + settings = snapshot.get("allocator_settings", {}) + if settings: + print("\n--- Allocator Settings ---") + for k, v in sorted(settings.items()): + # roundup_power2_divisions is usually an empty dict, but when set it + # changes allocation rounding and can explain a reserved-memory or + # fragmentation difference, so report it rather than hiding it. + if isinstance(v, dict): + v = dict(sorted(v.items(), key=lambda kv: str(kv[0]))) or "(unset)" + print(f" {k}: {v}") + + # Segments + print("\n--- Segments ---") + print(f" Count: {len(segments)}") + print(f" Reserved (total): {format_size(total_size)}") + print(f" Allocated: {format_size(allocated_size)}") + print(f" Active: {format_size(active_size)}") + print(f" Awaiting free: {format_size(awaiting_free_size)}") + print(f" Inactive (reusable): {format_size(inactive_size)}") + if total_size > 0: + frag_pct = inactive_size / total_size * 100 + print(f" Fragmentation: {frag_pct:.1f}% (inactive / reserved)") + + # Baseline memory + if baseline: + print("\n--- Baseline Memory (at profiling start) ---") + print(f" Baseline: {format_size(baseline.baseline_at_start)}") + print(f" Pre-existing freed: {format_size(baseline.pre_existing_freed)}") + print(f" Profiling live: {format_size(baseline.profiling_still_live)}") + + # Profiler steps + print("\n--- Profiler Steps ---") + if not steps: + print(" (none detected)") + else: + print(f" {'Step':>6} {'Duration':>10} {'Status'}") + print(f" {'─' * 6} {'─' * 10} {'─' * 12}") + for num, info in sorted(steps.items()): + if info["start"] and info["end"]: + dur_ms = (info["end"] - info["start"]) / 1000 + status = "complete" + dur_str = f"{dur_ms:.1f} ms" + else: + status = "incomplete" + dur_str = "—" + print(f" {num:>6} {dur_str:>10} {status}") + + # Device traces + print("\n--- Device Traces ---") + if device_idx is None: + print(" (no traces found)") + else: + print(f" Device: {device_idx}") + print(f" Total events: {len(traces):,}") + for action, count in sorted(action_counts.items()): + print(f" {action:20s}: {count:>8,}") + + # Top sources + print("\n--- Top Allocation Sources (by throughput) ---") + if not top_sources: + print(" (no allocations)") + else: + print(f" {'#':>3} {'Throughput':>12} {'Allocs':>8} Source") + print(f" {'─' * 3} {'─' * 12} {'─' * 8} {'─' * 40}") + for i, (key, (total, count)) in enumerate(top_sources, 1): + print(f" {i:>3} {format_size(total):>12} {count:>8,} {key}") + + print() + + +def main() -> None: + """Parse arguments and print the snapshot overview.""" + parser = argparse.ArgumentParser( + description="Overview of a CUDA memory snapshot pickle file.", + epilog=TRUST_EPILOG, + ) + parser.add_argument("pickle_path", help="Path to the snapshot .pickle file") + parser.add_argument("--json", action="store_true", help="Output as JSON") + args = parser.parse_args() + + snapshot = load_snapshot(args.pickle_path) + print_overview(snapshot, args.pickle_path, as_json=args.json) + + +if __name__ == "__main__": + main() diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py new file mode 100644 index 00000000000..817d1c5355d --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py @@ -0,0 +1,524 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generate an interactive HTML timeline of CUDA memory usage from snapshots.""" + +# Annotations are lazy so the PEP 604 `X | None` syntax below does not need to +# evaluate at import time; these scripts stay runnable on the system python3 +# (3.9 on macOS) even though the repo itself targets 3.10+. +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys + + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from common import TRUST_EPILOG, compute_baseline, find_active_device, load_snapshot + + +COLORS = [ + "#00d2ff", + "#ff6b6b", + "#2ecc71", + "#f39c12", + "#9b59b6", + "#e74c3c", + "#1abc9c", + "#e67e22", +] + + +# --------------------------------------------------------------------------- +# Replay helpers (local – only needed for timeline generation) +# --------------------------------------------------------------------------- + + +logger = logging.getLogger(__name__) + + +def _replay_timeseries(traces): + """Replay trace events into parallel lists for both x-axis modes. + + Uses ``free_completed`` to match PyTorch memory_viz behaviour. + + Returns (times_us, event_indices, delta_bytes) where: + - *times_us*: wall-clock timestamp of each state change + - *event_indices*: synthetic counter (1 per alloc/free) for event-index mode + - *delta_bytes*: cumulative allocated bytes at each point + """ + times = [] + indices = [] + deltas = [] + live = {} + current = 0 + idx = 0 + + # Seed the state at trace start. Without this the series begins at the first + # alloc/free rather than at delta 0, so a trace whose earliest visible event + # is an unmatched free (older captures, or a ring buffer that dropped the + # matching allocs) yields only negative deltas and reports a peak below the + # reconstructed baseline — when the true peak was the baseline itself. + if traces: + times.append(traces[0].get("time_us", 0)) + indices.append(idx) + deltas.append(0) + idx += 1 + + for e in traces: + action = e["action"] + if action == "alloc": + addr = e["addr"] + live[addr] = e["size"] + current += e["size"] + times.append(e["time_us"]) + indices.append(idx) + deltas.append(current) + idx += 1 + elif action == "free_completed": + addr = e["addr"] + if addr in live: + current -= live.pop(addr) + else: + current -= e["size"] + times.append(e["time_us"]) + indices.append(idx) + deltas.append(current) + idx += 1 + + return times, indices, deltas + + +def _downsample_by_time(times_us, values, precision_us): + """Downsample by fixed-width time buckets, preserving min/max per bucket.""" + n = len(times_us) + if n == 0: + return [], [] + + out_t, out_v = [], [] + bucket_end = times_us[0] + precision_us + min_idx = max_idx = 0 + + for i in range(1, n): + if times_us[i] >= bucket_end: + # Emit current bucket + first, second = (min_idx, max_idx) if min_idx <= max_idx else (max_idx, min_idx) + out_t.append(times_us[first]) + out_v.append(values[first]) + if first != second: + out_t.append(times_us[second]) + out_v.append(values[second]) + # Start new bucket at this event + bucket_end = times_us[i] + precision_us + min_idx = max_idx = i + else: + if values[i] < values[min_idx]: + min_idx = i + if values[i] > values[max_idx]: + max_idx = i + + # Emit last bucket + first, second = (min_idx, max_idx) if min_idx <= max_idx else (max_idx, min_idx) + out_t.append(times_us[first]) + out_v.append(values[first]) + if first != second: + out_t.append(times_us[second]) + out_v.append(values[second]) + + return out_t, out_v + + +def _downsample_by_count(xs, values, target): + """Downsample to a target point count via min/max bucketing.""" + target = max(target, 4) + n = len(xs) + if n <= target: + return xs, values + + bucket_size = max(1, n // (target // 2)) + out_x, out_v = [], [] + + for i in range(0, n, bucket_size): + end = min(i + bucket_size, n) + min_idx = max_idx = i + for j in range(i + 1, end): + if values[j] < values[min_idx]: + min_idx = j + if values[j] > values[max_idx]: + max_idx = j + + first, second = (min_idx, max_idx) if min_idx <= max_idx else (max_idx, min_idx) + out_x.append(xs[first]) + out_v.append(values[first]) + if first != second: + out_x.append(xs[second]) + out_v.append(values[second]) + + return out_x, out_v + + +# --------------------------------------------------------------------------- +# HTML generation +# --------------------------------------------------------------------------- + +_HTML_TEMPLATE = r""" + + + +CUDA Memory Timeline + + + + +
+

CUDA Memory Timeline

+

+
+
+
+ + +
+
+ + +
+
+ + + + +
+
+
+
+ + +""" + + +def generate(pickle_paths: list[str], labels: list[str], output: str, precision_ms: float) -> None: + """Write a standalone interactive HTML timeline for one or more snapshots. + + Args: + pickle_paths: Snapshot pickle files to plot, one trace per file. + labels: Display label for each trace, parallel to ``pickle_paths``. + output: Path the HTML file is written to. + precision_ms: Time bucket width in milliseconds for the wall-clock view. + """ + GIB = 1024**3 + precision_us = int(precision_ms * 1000) + traces_data = [] + + for i, path in enumerate(pickle_paths): + snap = load_snapshot(path) + dev = find_active_device(snap) + if dev is None: + logger.warning(f"no device traces in {path}, skipping.") + continue + trace = snap["device_traces"][dev] + baseline = compute_baseline(snap, dev) + base_gib = baseline.baseline_at_start / GIB + + # Anchor to profiling start (first event of any kind) + t0 = trace[0].get("time_us", 0) if trace else 0 + + logger.info(f" Replaying {len(trace):,} events ...") + raw_times, raw_indices, raw_deltas = _replay_timeseries(trace) + if not raw_times: + logger.warning(f"empty trace in {path}, skipping.") + continue + + # Wall-clock view: downsample by time precision + ds_times, ds_deltas_t = _downsample_by_time(raw_times, raw_deltas, precision_us) + + # Event-index view: match the wall-clock point count for similar file size + time_points = len(ds_times) + ds_indices, ds_deltas_e = _downsample_by_count( + raw_indices, + raw_deltas, + time_points, + ) + + times_s = [(t - t0) / 1e6 for t in ds_times] + delta_gib_time = [d / GIB for d in ds_deltas_t] + event_idx = list(ds_indices) + delta_gib_event = [d / GIB for d in ds_deltas_e] + + peak_delta_gib = max(d / GIB for d in raw_deltas) if raw_deltas else 0 + peak_abs_gib = peak_delta_gib + base_gib + + label = labels[i] if i < len(labels) else os.path.basename(path) + color = COLORS[i % len(COLORS)] + alloc_count = sum(1 for e in trace if e["action"] == "alloc") + + logger.info(f" {len(times_s):,} time points ({precision_ms}ms), {len(event_idx):,} event-index points") + + traces_data.append( + { + "label": label, + "color": color, + "baseline_gib": round(base_gib, 4), + "alloc_count": alloc_count, + "peak_delta_gib": round(peak_delta_gib, 4), + "peak_abs_gib": round(peak_abs_gib, 4), + "points": len(times_s), + # Time-based x-axis data + "times_s": times_s, + "delta_gib": delta_gib_time, + # Event-index x-axis data + "event_idx": event_idx, + "delta_gib_event": delta_gib_event, + } + ) + + if not traces_data: + logger.error("no valid traces to plot.") + sys.exit(1) + + html = _HTML_TEMPLATE.replace("__DATA_PLACEHOLDER__", json.dumps(traces_data)).replace( + "__PRECISION_MS__", str(precision_ms) + ) + + os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True) + with open(output, "w") as f: + f.write(html) + + size_kib = os.path.getsize(output) / 1024 + logger.info(f"Written to {output} ({size_kib:.0f} KiB)") + + +def main() -> None: + """Parse arguments and write the interactive timeline HTML.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + parser = argparse.ArgumentParser( + description="Generate an interactive HTML timeline of CUDA memory usage.", + epilog=TRUST_EPILOG, + ) + parser.add_argument("pickles", nargs="+", help="Snapshot pickle file(s)") + parser.add_argument( + "--labels", + nargs="+", + default=[], + help="Display labels for each snapshot (default: filenames)", + ) + parser.add_argument( + "--output", + "-o", + default=None, + help="Output HTML path (default: cuda_memory_timeline.html in working dir)", + ) + + def _positive_float(value): + fval = float(value) + if fval <= 0: + raise argparse.ArgumentTypeError(f"--precision must be strictly positive, got {value}") + return fval + + parser.add_argument( + "--precision", + type=_positive_float, + default=1.0, + help="Time bucket width in milliseconds for wall-clock mode (default: 1.0). " + "Smaller values produce more plot points and finer detail.", + ) + args = parser.parse_args() + + if args.labels and len(args.labels) != len(args.pickles): + parser.error("Number of --labels must match number of pickle files") + + output = args.output or "cuda_memory_timeline.html" + generate(args.pickles, args.labels, output, args.precision) + + +if __name__ == "__main__": + main() diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py new file mode 100644 index 00000000000..7c9a6929a24 --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Replay device_traces for a specific training step and show source-grouped allocations.""" + +# Annotations are lazy so the PEP 604 `X | None` syntax below does not need to +# evaluate at import time; these scripts stay runnable on the system python3 +# (3.9 on macOS) even though the repo itself targets 3.10+. +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys + + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from common import ( + TRUST_EPILOG, + compute_baseline, + compute_step_start_deltas, + find_active_device, + format_size, + get_profiler_steps, + get_step_annotations, + get_step_events, + group_by_source, + live_set_before, + load_snapshot, + replay_events, +) + + +logger = logging.getLogger(__name__) + + +def replay_one_step( + traces: list, + annotations: list, + step_num: int, + step_info: dict, + top_n: int = 15, + frame_depth: int = 1, + as_json: bool = False, + baseline_at_start: int = 0, + step_start_delta: int = 0, +) -> dict: + """Replay a single step and return results.""" + start = step_info["start"] + end = step_info["end"] + events = get_step_events(traces, start, end) + # Seed with what was already live when the step opened, so the source table + # accounts for memory carried into the step, not only what it allocated. + result = replay_events(events, initial_live=live_set_before(traces, start)) + sources = group_by_source(result.peak_live_set, depth=frame_depth) + ann_counts = get_step_annotations(annotations, start, end) + + complete = end is not None + duration_ms = (end - start) / 1000 if complete else None + + absolute_peak = baseline_at_start + step_start_delta + result.peak_delta + + step_result = { + "step": step_num, + "complete": complete, + "duration_ms": duration_ms, + "alloc_count": result.alloc_count, + "free_count": result.free_count, + "total_throughput": result.total_alloc_bytes, + "peak_delta": result.peak_delta, + "absolute_peak": absolute_peak, + "end_delta": result.end_delta, + "unmatched_frees": result.unmatched_free_count, + "unmatched_free_bytes": result.unmatched_free_bytes, + "top_sources_at_peak": sources[:top_n], + "annotations": ann_counts, + } + + if not as_json: + print(f"\n{'=' * 70}") + print(f" Step {step_num}" + (" (incomplete)" if not complete else "")) + print(f"{'=' * 70}") + if duration_ms: + print(f" Duration: {duration_ms:.1f} ms") + print(f" Allocs: {result.alloc_count:,}") + print(f" Frees: {result.free_count:,}") + print(f" Throughput: {format_size(result.total_alloc_bytes)}") + print(f" Peak delta: {format_size(result.peak_delta)}") + if baseline_at_start > 0: + print(f" Absolute peak: {format_size(absolute_peak)}") + print(f" End delta: {format_size(result.end_delta)}") + if result.unmatched_free_count > 0: + print(f" Pre-existing frees: {result.unmatched_free_count} ({format_size(result.unmatched_free_bytes)})") + + if ann_counts: + print("\n --- Active Annotations ---") + for name, count in sorted(ann_counts.items(), key=lambda x: -x[1]): + print(f" {count:>5}x {name}") + + print(f"\n --- Top {min(top_n, len(sources))} Sources at Peak (by size) ---") + if sources: + print(f" {'#':>3} {'Size':>12} {'Count':>6} Source") + print(f" {'─' * 3} {'─' * 12} {'─' * 6} {'─' * 40}") + for i, (key, total, count) in enumerate(sources[:top_n], 1): + print(f" {i:>3} {format_size(total):>12} {count:>6} {key}") + else: + print(" (no live allocations at peak)") + + return step_result + + +def main() -> None: + """Parse arguments and replay the requested training step(s).""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + parser = argparse.ArgumentParser( + description="Replay device_traces for a training step.", + epilog=TRUST_EPILOG, + ) + parser.add_argument("pickle_path", help="Path to the snapshot .pickle file") + parser.add_argument("--step", type=int, help="Step number to replay") + parser.add_argument("--all-steps", action="store_true", help="Replay all complete steps") + parser.add_argument("--top", type=int, default=15, help="Number of top sources to show (default: 15)") + parser.add_argument( + "--frame-depth", type=int, default=1, help="Stack frame depth for source grouping (default: 1)" + ) + parser.add_argument("--device", type=int, help="Device index (default: auto-detect)") + parser.add_argument("--json", action="store_true", help="Output as JSON") + args = parser.parse_args() + + if args.step is None and not args.all_steps: + parser.error("Specify --step N or --all-steps") + + snapshot = load_snapshot(args.pickle_path) + annotations = snapshot.get("external_annotations", []) + steps = get_profiler_steps(annotations) + + if not steps: + logger.error("no ProfilerStep annotations found.") + sys.exit(1) + + device_idx = args.device if args.device is not None else find_active_device(snapshot) + if device_idx is None: + logger.error("no device traces found.") + sys.exit(1) + traces = snapshot["device_traces"][device_idx] + + # Compute baseline and per-step starting deltas + baseline = compute_baseline(snapshot, device_idx) + step_deltas = compute_step_start_deltas(traces, steps) + + if not args.json: + print(f"\n Baseline at profiling start: {format_size(baseline.baseline_at_start)}") + + if args.all_steps: + step_nums = sorted(steps.keys()) + else: + if args.step not in steps: + available = sorted(steps.keys()) + logger.error(f"step {args.step} not found. Available: {available}") + sys.exit(1) + step_nums = [args.step] + + all_results = [] + for num in step_nums: + info = steps[num] + if not info["start"]: + continue + r = replay_one_step( + traces, + annotations, + num, + info, + top_n=args.top, + frame_depth=args.frame_depth, + as_json=args.json, + baseline_at_start=baseline.baseline_at_start, + step_start_delta=step_deltas.get(num, 0), + ) + all_results.append(r) + + if args.json: + print(json.dumps(all_results, indent=2, default=str)) + else: + print() + + +if __name__ == "__main__": + main() diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py new file mode 100644 index 00000000000..48b52a3a9f3 --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Drill down into CUDA memory state at a specific point in time.""" + +# Annotations are lazy so the PEP 604 `X | None` syntax below does not need to +# evaluate at import time; these scripts stay runnable on the system python3 +# (3.9 on macOS) even though the repo itself targets 3.10+. +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys + + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from common import ( + TRUST_EPILOG, + compute_baseline, + find_active_device, + format_size, + group_by_source, + load_snapshot, + replay_to_time, +) + + +logger = logging.getLogger(__name__) + + +def _trace_duration_s(traces): + """Return elapsed seconds between first and last event.""" + first = last = None + for e in traces: + t = e.get("time_us", 0) + if first is None: + first = t + last = t + if first is None or last is None: + return 0.0 + return (last - first) / 1e6 + + +def drill_single(path: str, time_s: float, top_n: int, frame_depth: int, as_json: bool) -> None: + """Report what is live in a single snapshot at a given elapsed time. + + Args: + path: Path to the snapshot pickle file. + time_s: Elapsed seconds from the first trace event to replay up to. + top_n: Number of top allocation sources to show. + frame_depth: Stack frames to include in each source key. + as_json: Emit machine-readable JSON instead of a text report. + """ + snap = load_snapshot(path) + dev = find_active_device(snap) + if dev is None: + logger.error("no device traces found.") + sys.exit(1) + + traces = snap["device_traces"][dev] + duration = _trace_duration_s(traces) + baseline = compute_baseline(snap, dev) + base = baseline.baseline_at_start + + target_us = int(time_s * 1e6) + result = replay_to_time(traces, target_us) + groups = group_by_source(result.live_set, depth=frame_depth) + + abs_mem = base + result.current_delta + + if as_json: + out = { + "file": path, + "time_s": time_s, + "trace_duration_s": round(duration, 2), + "baseline": base, + "delta": result.current_delta, + "absolute": abs_mem, + "alloc_count": result.alloc_count, + "free_count": result.free_count, + "live_count": len(result.live_set), + "sources": [{"source": src, "bytes": total, "count": cnt} for src, total, cnt in groups[:top_n]], + } + print(json.dumps(out, indent=2, default=str)) + return + + name = os.path.basename(path) + print("=" * 80) + print(f" Memory State at t = {time_s:.2f}s") + print("=" * 80) + print(f"\n File: {name}") + print(f" Duration: {duration:.2f}s total trace") + print(f"\n Baseline: {format_size(base):>12}") + print(f" Delta at t={time_s:.2f}s: {format_size(result.current_delta):>12}") + print(f" Absolute: {format_size(abs_mem):>12}") + print(f" Live allocations: {len(result.live_set):>12,}") + print(f" Allocs so far: {result.alloc_count:>12,}") + print(f" Frees so far: {result.free_count:>12,}") + + print(f"\n--- Live Allocations by Source (top {top_n}) ---\n") + if groups: + print(f" {'Source':<60s} {'Size':>12s} {'Count':>6s}") + print(f" {'─' * 60} {'─' * 12} {'─' * 6}") + for src, total, cnt in groups[:top_n]: + display = src if len(src) <= 60 else src[:57] + "..." + print(f" {display:<60s} {format_size(total):>12s} {cnt:>6,}") + print(f"\n ({len(groups)} total sources, showing top {top_n})") + else: + print(" (no live allocations)") + print() + + +def drill_compare(path_a: str, path_b: str, time_s: float, top_n: int, frame_depth: int, as_json: bool) -> None: + """Compare what is live in two snapshots at the same elapsed time. + + Sources are sorted by absolute delta so the allocations responsible for a + divergence surface first, with only-in-A and only-in-B totals reported + separately. + + Args: + path_a: Path to the first snapshot pickle file. + path_b: Path to the second snapshot pickle file. + time_s: Elapsed seconds from each trace's first event to replay up to. + top_n: Number of top allocation sources to show. + frame_depth: Stack frames to include in each source key. + as_json: Emit machine-readable JSON instead of a text report. + """ + snap_a = load_snapshot(path_a) + snap_b = load_snapshot(path_b) + + dev_a = find_active_device(snap_a) + dev_b = find_active_device(snap_b) + if dev_a is None or dev_b is None: + logger.error("one or both snapshots have no device traces.") + sys.exit(1) + + traces_a = snap_a["device_traces"][dev_a] + traces_b = snap_b["device_traces"][dev_b] + baseline_a = compute_baseline(snap_a, dev_a) + baseline_b = compute_baseline(snap_b, dev_b) + base_a = baseline_a.baseline_at_start + base_b = baseline_b.baseline_at_start + + target_us = int(time_s * 1e6) + res_a = replay_to_time(traces_a, target_us) + res_b = replay_to_time(traces_b, target_us) + + groups_a = group_by_source(res_a.live_set, depth=frame_depth) + groups_b = group_by_source(res_b.live_set, depth=frame_depth) + + abs_a = base_a + res_a.current_delta + abs_b = base_b + res_b.current_delta + + # Merge sources + map_a = {src: (total, cnt) for src, total, cnt in groups_a} + map_b = {src: (total, cnt) for src, total, cnt in groups_b} + all_sources = set(map_a.keys()) | set(map_b.keys()) + + rows = [] + for src in all_sources: + ta, ca = map_a.get(src, (0, 0)) + tb, cb = map_b.get(src, (0, 0)) + rows.append((abs(ta - tb), ta - tb, ta, ca, tb, cb, src)) + rows.sort(reverse=True) + + if as_json: + out = { + "file_a": path_a, + "file_b": path_b, + "time_s": time_s, + "baseline_a": base_a, + "baseline_b": base_b, + "delta_a": res_a.current_delta, + "delta_b": res_b.current_delta, + "absolute_a": abs_a, + "absolute_b": abs_b, + "live_count_a": len(res_a.live_set), + "live_count_b": len(res_b.live_set), + "sources": [ + { + "source": src, + "bytes_a": ta, + "count_a": ca, + "bytes_b": tb, + "count_b": cb, + "delta": ta - tb, + } + for _, _, ta, ca, tb, cb, src in rows[:top_n] + ], + } + print(json.dumps(out, indent=2, default=str)) + return + + name_a = os.path.basename(path_a) + name_b = os.path.basename(path_b) + + print("=" * 100) + print(f" Memory Comparison at t = {time_s:.2f}s") + print("=" * 100) + print(f"\n A: {name_a}") + print(f" B: {name_b}") + + print(f"\n {'':45s} {'A':>14s} {'B':>14s} {'Delta(A-B)':>14s}") + print(f" {'─' * 45} {'─' * 14} {'─' * 14} {'─' * 14}") + print( + f" {'Baseline':45s} {format_size(base_a):>14s} {format_size(base_b):>14s} {format_size(base_a - base_b):>14s}" + ) + print( + f" {'Delta at t=' + f'{time_s:.2f}s':45s} {format_size(res_a.current_delta):>14s} {format_size(res_b.current_delta):>14s} {format_size(res_a.current_delta - res_b.current_delta):>14s}" + ) + print( + f" {'Absolute':45s} {format_size(abs_a):>14s} {format_size(abs_b):>14s} {format_size(abs_a - abs_b):>14s}" + ) + print( + f" {'Live allocations':45s} {len(res_a.live_set):>14,} {len(res_b.live_set):>14,} {len(res_a.live_set) - len(res_b.live_set):>+14,}" + ) + + print(f"\n--- Live Allocations at t={time_s:.2f}s (top {top_n} by |delta|) ---\n") + print(f" {'Source':<60s} {'A':>12s} {'#A':>5s} {'B':>12s} {'#B':>5s} {'Delta(A-B)':>12s}") + print(f" {'─' * 60} {'─' * 12} {'─' * 5} {'─' * 12} {'─' * 5} {'─' * 12}") + + for _, diff, ta, ca, tb, cb, src in rows[:top_n]: + sign = "+" if diff > 0 else "" + display = src if len(src) <= 60 else src[:57] + "..." + print( + f" {display:<60s} " + f"{format_size(ta):>12s} {ca:>5,} " + f"{format_size(tb):>12s} {cb:>5,} " + f"{sign}{format_size(diff):>11s}" + ) + + print(f"\n ({len(rows)} total sources, showing top {top_n})") + + # Only-in-A / only-in-B summaries + only_a = [(src, ta, ca) for _, _, ta, ca, tb, cb, src in rows if tb == 0 and ta > 1024 * 1024] + only_b = [(src, tb, cb) for _, _, ta, ca, tb, cb, src in rows if ta == 0 and tb > 1024 * 1024] + if only_a: + total = sum(t for _, t, _ in only_a) + print(f"\n Only in A (> 1 MiB): {len(only_a)} sources, {format_size(total)} total") + if only_b: + total = sum(t for _, t, _ in only_b) + print(f" Only in B (> 1 MiB): {len(only_b)} sources, {format_size(total)} total") + print() + + +def main() -> None: + """Parse arguments and run the point-in-time drill-down.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + parser = argparse.ArgumentParser( + description="Drill down into CUDA memory state at a specific elapsed time.", + epilog=TRUST_EPILOG, + ) + parser.add_argument("pickles", nargs="+", help="One or two snapshot pickle files") + parser.add_argument( + "--time", + "-t", + type=float, + required=True, + help="Elapsed seconds from profiling start to inspect", + ) + parser.add_argument("--top", type=int, default=25, help="Number of top sources (default: 25)") + parser.add_argument("--frame-depth", type=int, default=1, help="Stack frame depth (default: 1)") + parser.add_argument("--json", action="store_true", help="Output as JSON") + args = parser.parse_args() + + if len(args.pickles) > 2: + parser.error("At most two pickle files supported") + + if len(args.pickles) == 1: + drill_single(args.pickles[0], args.time, args.top, args.frame_depth, args.json) + else: + drill_compare(args.pickles[0], args.pickles[1], args.time, args.top, args.frame_depth, args.json) + + +if __name__ == "__main__": + main() diff --git a/skills/nemo-mbridge-memory-snapshot-analysis/skill-card.md b/skills/nemo-mbridge-memory-snapshot-analysis/skill-card.md new file mode 100644 index 00000000000..7cdf3a17f07 --- /dev/null +++ b/skills/nemo-mbridge-memory-snapshot-analysis/skill-card.md @@ -0,0 +1,67 @@ +## Description:
+Analyze and compare PyTorch CUDA memory snapshots produced by Megatron Bridge's `ProfilingConfig(record_memory_history=True)`. Replays the recorded allocation timeline to plot memory over time, compare two runs, and attribute peak memory to source code locations.
+ +This skill is ready for commercial/non-commercial use.
+ +## Owner +NVIDIA
+ +### License/Terms of Use:
+Apache 2.0
+ +## Use Case:
+Developers and engineers debugging out-of-memory failures and peak-memory regressions in Megatron Bridge GPU training workloads, comparing memory behavior between two configurations, or attributing GPU memory to specific allocation sites.
+ +### Deployment Geography for Use:
+Global
+ +## Requirements / Dependencies:
+**Requires API Key or External Credential:** [Not Specified]
+**Credential Type(s):** [None identified]
+ +The bundled scripts use only the Python standard library. The generated HTML timeline loads Plotly.js from a public CDN when opened in a browser; it is not a Python dependency.
+ +Do not include secrets in prompts/logs/output; use least-privilege credentials; rotate keys as appropriate.
+ +## Known Risks and Mitigations:
+Risk: Review before execution as proposals could introduce incorrect or misleading guidance into skills.
+Mitigation: Review and scan skill before deployment.
+ +Risk: Memory snapshot pickle files are deserialized with `pickle.load`, which executes arbitrary code for untrusted input.
+Mitigation: Only analyze snapshot files produced by your own training runs or from a trusted source.
+ +Risk: Snapshot files embed absolute source paths and Python stack frames from the machine that produced them, which may reveal internal directory layouts.
+Mitigation: Review before sharing snapshot files or script output outside your organization.
+ +## Reference(s):
+- [Profiling Documentation](../../docs/training/profiling.md)
+- [Memory Tuning Skill](../nemo-mbridge-perf-memory-tuning/SKILL.md)
+- [CUDA Graphs Skill](../nemo-mbridge-perf-cuda-graphs/SKILL.md)
+- [Understanding GPU Memory (PyTorch blog)](https://pytorch.org/blog/understanding-gpu-memory-1/)
+ +## Skill Output:
+**Output Type(s):** [Analysis, Shell commands, Generated HTML report, JSON]
+**Output Format:** [Markdown with inline bash and Python code blocks; scripts emit formatted text tables, JSON, or a standalone HTML timeline]
+**Output Parameters:** [1D]
+**Other Properties Related to Output:** [None]
+ +## Evaluation Agents Used:
+- `claude-code`
+- `codex`
+ +## Evaluation Tasks:
+Evaluated against 2 tasks in the NVSkills-Eval external profile (astra-sandbox environment).
+ +## Evaluation Metrics Used:
+Reported benchmark dimensions:
+- Security: Checks whether skill-assisted execution avoids unsafe behavior such as secret leakage, destructive commands, or unauthorized access.
+- Correctness: Checks whether the agent follows the expected workflow and produces the correct final output.
+- Discoverability: Checks whether the agent loads the skill when relevant and avoids using it when irrelevant.
+- Effectiveness: Checks whether the agent performs measurably better with the skill than without it.
+- Efficiency: Checks whether the agent uses fewer tokens and avoids redundant work.
+ +Underlying evaluation signals used in this run:
+- `security`: Checks for unsafe operations, secret leakage, and unauthorized access.
+- `skill_execution`: Verifies that the agent loaded the expected skill and workflow.
+- `skill_efficiency`: Checks routing quality, decoy avoidance, and redundant tool usage.
+- `accuracy`: Grades final-answer correctness against the reference answer.
From a1e812ed0f5dbc9b44c7f80f825f548ba7b3ab6b Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 01:43:18 +0200 Subject: [PATCH 12/85] cp: `fix(model): add Gemma 3 VL verification coverage (5084)` into `r0.6.0` (#5204) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../compare_hf_and_megatron/compare.py | 81 +++--- .../gemma-3-4b-it/card.yaml | 275 ++++++++++++++++++ .../bridge/models/gemma/gemma3_provider.py | 9 +- .../models/gemma_vl/gemma3_vl_provider.py | 13 + .../models/gemma_vl/modeling_gemma3_vl.py | 4 +- .../recipes/gemma3_vl/h100/gemma3_vl.py | 25 +- .../models/gemma_vl/test_gemma3_vl_bridge.py | 67 ++++- .../gemma_vl/test_gemma3_vl_modeling.py | 40 +++ .../gemma_vl/test_gemma3_vl_provider.py | 84 ++++++ .../recipes/test_gemma3_vl_recipes.py | 3 + .../unit_tests/test_compare_mask_handling.py | 84 ++++++ 11 files changed, 634 insertions(+), 51 deletions(-) create mode 100644 examples/model_verification_cards/gemma-3-4b-it/card.yaml create mode 100644 tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 610e14b5b4f..0ce406fa2d1 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -93,6 +93,7 @@ import argparse import gc import importlib +import io import os import sys from typing import Optional @@ -103,23 +104,8 @@ from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.utils import InferenceMode from megatron.core.pipeline_parallel.schedules import get_forward_backward_func -from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer - - -try: - from qwen_vl_utils import process_vision_info - - QWEN_VL_UTILS_AVAILABLE = True -except ImportError: - QWEN_VL_UTILS_AVAILABLE = False - process_vision_info = None -import io -import os - -# Import debugger module from same directory -import sys - from PIL import Image +from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer from megatron.bridge import AutoBridge from megatron.bridge.models.hf_pretrained.utils import is_safe_repo @@ -245,6 +231,15 @@ def is_vision_language_model( arch = getattr(config, "architectures", []) arch_str = " ".join(arch).lower() if arch else "" + # Some VLMs, including Gemma 3, use a family name without an explicit + # vision marker. Prefer the structured multimodal config over name + # heuristics so their image inputs are not silently discarded. + has_vision_config = getattr(config, "vision_config", None) is not None + has_text_config = getattr(config, "text_config", None) is not None + is_conditional_generation = "forconditionalgeneration" in arch_str + if has_vision_config and (has_text_config or is_conditional_generation): + return True + # Common patterns for VL models vl_indicators = [ "vl", @@ -446,40 +441,36 @@ def process_inputs(tokenizer, processor, image_path: Optional[str], prompt: str, tp_size: Tensor parallel size for padding sequence length Returns: - Tuple of (input_ids, pixel_values, image_grid_thw, messages) + Tuple of (input_ids, pixel_values, image_grid_thw, token_type_ids) """ if is_vl_model and image_path: - if not QWEN_VL_UTILS_AVAILABLE: - raise ImportError("qwen_vl_utils is required for vision-language models but not installed") - - # Create messages with image and text messages = [ { "role": "user", "content": [ - {"type": "image", "image": image_path}, + {"type": "image", "image": load_image(image_path)}, {"type": "text", "text": prompt}, ], } ] - - # Process vision info - image_inputs, video_inputs = process_vision_info(messages) - - # Apply chat template - text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - - # Process inputs - inputs = processor( - text=[text], - images=image_inputs, - videos=video_inputs, - padding=True, + inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=True, return_tensors="pt", ) - input_ids = pad_input_ids_to_tp_multiple(inputs.input_ids, tp_size, tokenizer.pad_token_id or 0) - return input_ids, inputs.pixel_values, inputs.image_grid_thw, messages + input_ids = pad_input_ids_to_tp_multiple(inputs["input_ids"], tp_size, tokenizer.pad_token_id or 0) + token_type_ids = inputs.get("token_type_ids") + if token_type_ids is not None: + token_type_ids = pad_input_ids_to_tp_multiple(token_type_ids, tp_size, 0) + return ( + input_ids, + inputs.get("pixel_values"), + inputs.get("image_grid_thw"), + token_type_ids, + ) else: # Text-only processing for both VL models without images and regular LLMs if is_vl_model and processor: @@ -573,7 +564,7 @@ def _export_and_load_roundtrip_hf_model(args, is_vl_model: bool, megatron_model, return None -def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokenizer): +def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokenizer, *, token_type_ids=None): """Run HuggingFace model inference and return results. Args: @@ -582,6 +573,7 @@ def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokeniz pixel_values: Pixel values for vision models (optional). image_grid_thw: Image grid dimensions (optional). tokenizer: Tokenizer for decoding. + token_type_ids: Multimodal token type IDs (optional). Returns: Tuple of (hf_logits, hf_next_token, hf_logits_stats, hf_top5_info, logits_shape). @@ -600,6 +592,8 @@ def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokeniz hf_inputs["pixel_values"] = pixel_values if image_grid_thw is not None: hf_inputs["image_grid_thw"] = image_grid_thw + if token_type_ids is not None: + hf_inputs["token_type_ids"] = token_type_ids hf_output = hf_model(**hf_inputs) @@ -804,7 +798,7 @@ def compare_models_one_step(args) -> None: # Process inputs print_rank_0(f"Processing inputs - Prompt: '{args.prompt}', Image: {args.image_path}") - input_ids, pixel_values, image_grid_thw, messages = process_inputs( + input_ids, pixel_values, image_grid_thw, token_type_ids = process_inputs( tokenizer, processor, args.image_path, args.prompt, is_vl_model, args.tp ) @@ -814,13 +808,20 @@ def compare_models_one_step(args) -> None: pixel_values = pixel_values.cuda() if image_grid_thw is not None: image_grid_thw = image_grid_thw.cuda() + if token_type_ids is not None: + token_type_ids = token_type_ids.cuda() print_rank_0(f"Input shape: {input_ids.shape}") print_rank_0(f"Pixel values shape: {pixel_values.shape if pixel_values is not None else 'None'}") # Run HF model forward pass hf_logits, hf_next_token, hf_logits_stats, hf_top5_info, logits_shape = _run_hf_inference( - hf_model, input_ids, pixel_values, image_grid_thw, tokenizer + hf_model, + input_ids, + pixel_values, + image_grid_thw, + tokenizer, + token_type_ids=token_type_ids, ) del hf_model diff --git a/examples/model_verification_cards/gemma-3-4b-it/card.yaml b/examples/model_verification_cards/gemma-3-4b-it/card.yaml new file mode 100644 index 00000000000..d2f19aa0d8f --- /dev/null +++ b/examples/model_verification_cards/gemma-3-4b-it/card.yaml @@ -0,0 +1,275 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: gemma_3_4b_it +summary: > + Performance disclaimer: this model has not been performance-tuned; reported + timing and throughput metrics are sanity checks, not optimized performance + results. Gemma 3 4B Instruct exact-checkpoint CPU and TP2/PP2 GPU import, + TP2/PP2 GPU export, manual multimodal forward parity, and deterministic + TP2/PP2 inference are verified. CPU export remains unverified because its + vision namespace gate does not pass; SFT, post-SFT export/inference, and + PEFT remain pending. Pretraining and checkpoint resume are unsupported + because there is no public exact-model pretraining recipe. Long-context SFT + is unsupported because the required arbitrary multimodal mask currently + excludes context parallelism and in-batch THD packing. Existing toy-model + tests are not used as exact-model evidence. +verification_index: + model_level: + verified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - megatron_to_hf_gpu + - manual_forward_pass + - inference + unverified: + - megatron_to_hf_cpu + training: + H100: + unverified: [sft, sft_export_inference, peft] + unsupported: [pretrain, checkpoint_resume, sft_long_context] +model: + hf_id: google/gemma-3-4b-it + hf_revision: 093f9f388b31de276ce2de164bdc2081324b9767 # pragma: allowlist secret + architecture: Gemma3ForConditionalGeneration + min_transformers_version: "5.8.0" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.04-py3 + bridge_commit: 2d0175b116af6cb785bb8efab7ee053435496266 # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model google/gemma-3-4b-it + --hf-revision 093f9f388b31de276ce2de164bdc2081324b9767 + --megatron-path work/model-verification/gemma-3-4b-it/cpu-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-24 + expected_result: > + Exact-revision BF16 CPU import loaded all 781 Bridge tensors and saved + one reloadable DCP shard. An independent Gloo reload mapped all 883 + source tensors with exact keys, shapes, dtypes, and torch.equal values, + including all 437 legacy vision-tower keys. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 4 --hf-model google/gemma-3-4b-it + --hf-revision 093f9f388b31de276ce2de164bdc2081324b9767 + --megatron-path work/model-verification/gemma-3-4b-it/imported-megatron + --torch-dtype bfloat16 --tp 2 --pp 2 --ep 1 --etp 1 + last_verified: 2026-07-24 + expected_result: > + Exact-revision distributed BF16 import completed with TP2/PP2/EP1/ETP1. + All 781 Bridge tensors loaded and all four ranks saved a reloadable + checkpoint with the vision tower and multimodal projector on the owning + pipeline stage. + + megatron_to_hf_cpu: + status: unverified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model google/gemma-3-4b-it + --hf-revision 093f9f388b31de276ce2de164bdc2081324b9767 + --megatron-path work/model-verification/gemma-3-4b-it/cpu-megatron/iter_0000000 + --hf-path work/model-verification/gemma-3-4b-it/cpu-hf-export + --torch-dtype bfloat16 + last_verified: null + expected_result: > + A fresh exact-revision CPU import and export from Bridge commit + 2d0175b116af6cb785bb8efab7ee053435496266 completed on 2026-07-24, but + the source/export key gate reported 437 missing legacy + vision_tower.vision_model.* keys and 437 extra flat vision_tower.* keys. + This leaf remains unverified pending a source-namespace-preserving design; + changing the general config-only export path would regress pruned and + custom checkpoint exports. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 4 --hf-model google/gemma-3-4b-it + --hf-revision 093f9f388b31de276ce2de164bdc2081324b9767 + --megatron-path work/model-verification/gemma-3-4b-it/imported-megatron/iter_0000000 + --hf-path work/model-verification/gemma-3-4b-it/hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 2 --pp 2 --ep 1 --etp 1 + last_verified: 2026-07-24 + expected_result: > + TP2/PP2 distributed export completed without a missing PP owner. Source + and export contained the same 883 keys with exact shapes, dtypes, and + torch.equal values, including all 437 legacy + vision_tower.vision_model.* tensors. Native + Gemma3ForConditionalGeneration strict reload reported no missing, + unexpected, mismatched, or errored tensors. + + manual_forward_pass: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=1 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path google/gemma-3-4b-it + --hf-revision 093f9f388b31de276ce2de164bdc2081324b9767 + --megatron_model_path work/model-verification/gemma-3-4b-it/imported-megatron/iter_0000000 + --model_class Gemma3ForConditionalGeneration + --image_path work/model-verification/gemma-3-4b-it/candy.JPG + --prompt "What animal is on the candy?" + --tp 1 --pp 1 --ep 1 --etp 1 + last_verified: 2026-07-24 + expected_result: > + The exact-revision image-and-text comparison used TP1 without padding. + Token match: true; Hugging Face and Megatron both produced next token ID + 22515 ("Based"). Cosine similarity: 0.990209, passing the required 0.99 + threshold. Maximum and mean absolute logit differences were 6.046631 + and 1.016037 respectively and are report-only diagnostics. + + inference: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=4 + examples/conversion/hf_to_megatron_generate_vlm.py + --hf_model_path google/gemma-3-4b-it + --hf-revision 093f9f388b31de276ce2de164bdc2081324b9767 + --megatron_model_path work/model-verification/gemma-3-4b-it/imported-megatron/iter_0000000 + --image_path work/model-verification/gemma-3-4b-it/candy.JPG + --prompt "What animal is on the candy?" + --max_new_tokens 32 + --tp 2 --pp 2 --ep 1 --etp 1 + last_verified: 2026-07-24 + expected_result: |- + One exact-revision TP2/PP2 BF16 greedy multimodal run generated exactly + 32 tokens without reaching EOS. The public input was p-blog/candy.JPG at + documentation-images revision + 44c3249461c5a5336585d1e31e0a4ce7fce06a5a with SHA-256 + fc417c899e94f8df465b7541c5a70f0eebb85c414d06345f0b290c061eccc84c. + The JSON-escaped literal completion was: "Based on the image, the animal + on the candy is a **turtle**. \n\nYou can see a clear depiction of a + turtle shell on each of the candies" + + pretrain: + all: + status: unsupported + precision: null + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + Megatron Bridge does not provide a public Gemma 3 VL 4B pretraining + recipe. The available exact-model recipes cover SFT and PEFT, so a + supported pretraining command and checkpoint contract cannot be stated. + + sft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 1 + --recipe gemma3_vl_4b_sft_config --mode sft --dataset cord-v2 + --pretrained_checkpoint work/model-verification/gemma-3-4b-it/imported-megatron/iter_0000000 + --max_steps 10 --seq_length 4096 + 'dataset.source.load_kwargs={revision:"7f0115a4b758a71d6473b8d085751692da2fef98"}' + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/gemma-3-4b-it/sft-checkpoints + --save_interval 10 logger.log_interval=1 logger.log_throughput=true + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + Pending exact-revision CORD-v2 full-SFT verification with the public + Gemma 3 VL 4B H100 recipe. The run must complete at least ten optimizer + steps with finite loss, no skipped or NaN iterations, all four metrics, + and a reloadable final full-model checkpoint. + + sft_export_inference: + H100: + status: unverified + precision: bf16 + depends_on: sft + commands: null + last_verified: null + expected_result: > + Pending the full-SFT checkpoint. Export must reload natively with exact + keys, shapes, and dtypes, then one deterministic greedy multimodal + Hugging Face generation must record its exact generated-token count and + literal completion. + + sft_long_context: + H100: + status: unsupported + precision: null + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + Gemma 3 VL image-bidirectional attention currently requires an + arbitrary unfused mask. The provider rejects context parallelism + greater than one, and the public recipes disable in-batch THD packing, + so there is no supported exact-model long-context SFT contract. + + peft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 1 + --recipe gemma3_vl_4b_peft_config --mode lora --dataset cord-v2 + --pretrained_checkpoint work/model-verification/gemma-3-4b-it/imported-megatron/iter_0000000 + --max_steps 10 --seq_length 4096 + 'dataset.source.load_kwargs={revision:"7f0115a4b758a71d6473b8d085751692da2fef98"}' + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/gemma-3-4b-it/peft-checkpoints + --save_interval 10 logger.log_interval=1 logger.log_throughput=true + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + Pending exact-revision CORD-v2 LoRA verification with the public Gemma 3 + VL 4B H100 recipe. The run must complete at least ten optimizer steps + with finite loss, no skipped or NaN iterations, all four metrics, and a + reloadable adapter checkpoint. + + checkpoint_resume: + all: + status: unsupported + precision: null + depends_on: pretrain + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + resume_comparison: null + expected_result: > + Checkpoint-resume verification depends on a supported pretraining + baseline. Because Megatron Bridge has no public Gemma 3 VL 4B + pretraining recipe, there is no supported reference run or resume + checkpoint contract for this model. diff --git a/src/megatron/bridge/models/gemma/gemma3_provider.py b/src/megatron/bridge/models/gemma/gemma3_provider.py index e52e9a4545c..d8f3eb24327 100644 --- a/src/megatron/bridge/models/gemma/gemma3_provider.py +++ b/src/megatron/bridge/models/gemma/gemma3_provider.py @@ -142,12 +142,13 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> "MCoreG def gemma3_layer_spec(config) -> ModuleSpec: """Gemma3 custom layer spec.""" + attn_mask_type = AttnMaskType.arbitrary if config.is_vision_language else AttnMaskType.causal return ModuleSpec( module=TransformerLayer, submodules=TransformerLayerSubmodules( self_attention=ModuleSpec( module=Gemma3SelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, + params={"attn_mask_type": attn_mask_type}, submodules=SelfAttentionSubmodules( linear_qkv=TELayerNormColumnParallelLinear, core_attention=Gemma3TEDotProductAttention, # mixed gloabl/local attn @@ -235,7 +236,11 @@ def __init__( config = copy.deepcopy(config) if _is_local_attn_layer(layer_number, config.interleaved_attn_pattern): # local attention, (q, k) - config.window_size = (config.window_size - 1, 0) + window_size = config.window_size - 1 + # The VL mask already blocks future text tokens. Keep the right + # window open so image tokens in the same block stay bidirectional. + right_window_size = window_size if config.is_vision_language else 0 + config.window_size = (window_size, right_window_size) else: # global attention config.window_size = None diff --git a/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py b/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py index 9478df0af3f..0e63c088697 100644 --- a/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py +++ b/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py @@ -15,6 +15,7 @@ from dataclasses import dataclass, field from megatron.core.models.gpt import GPTModel as MCoreGPTModel +from megatron.core.transformer.enums import AttnBackend from transformers import SiglipVisionConfig from megatron.bridge.models.gemma.gemma3_provider import Gemma3ModelProvider @@ -27,6 +28,11 @@ class Gemma3VLModelProvider(Gemma3ModelProvider): Base model provider for Gemma VL Models. """ + # These are model invariants, not user-selectable settings. Marking them as + # non-init also drops stale causal/flash values from older checkpoint configs. + is_vision_language: bool = field(default=True, init=False) + attention_backend: AttnBackend = field(default=AttnBackend.unfused, init=False) + # VL models shouldn't scatter embeddings across sequence parallel regions because # the vision embeddings are going to be inserted into the language embeddings. scatter_embedding_sequence_parallel: bool = False @@ -53,6 +59,13 @@ class Gemma3VLModelProvider(Gemma3ModelProvider): freeze_vision_projection: bool = False def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Gemma3VLModel: + if not self.is_vision_language: + raise ValueError("Gemma 3 VL requires is_vision_language=True.") + if self.attention_backend is not AttnBackend.unfused: + raise ValueError("Gemma 3 VL requires the unfused attention backend for arbitrary attention masks.") + if self.context_parallel_size != 1: + raise ValueError("Gemma 3 VL does not support context parallelism with arbitrary attention masks.") + model = Gemma3VLModel(self, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) # Apply freeze options if any are enabled diff --git a/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py b/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py index 7f130261af3..3acdc912829 100644 --- a/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py +++ b/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py @@ -231,9 +231,9 @@ def freeze(self, freeze_language_model: bool, freeze_vision_model: bool, freeze_ def _compute_attention_mask( self, - input_ids: torch.Tensor, + input_ids: Optional[torch.Tensor], ) -> Optional[torch.Tensor]: - if not self.pre_process: + if input_ids is None: return None batch_size, seq_len = input_ids.shape causal_mask = torch.tril(torch.ones((batch_size, 1, seq_len, seq_len))).to(input_ids.device) diff --git a/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py b/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py index 7421f4b1cb7..81319727619 100644 --- a/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py +++ b/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py @@ -18,6 +18,7 @@ """ import torch +from megatron.core.transformer.enums import AttnBackend from megatron.bridge import AutoBridge from megatron.bridge.peft.base import PEFT @@ -69,7 +70,7 @@ def gemma3_vl_4b_sft_1gpu_h100_bf16_config() -> ConfigContainer: cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = "flash" + cfg.model.attention_backend = AttnBackend.unfused cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -111,6 +112,8 @@ def gemma3_vl_4b_sft_1gpu_h100_bf16_config() -> ConfigContainer: # Dataset configuration cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = hf_path + # Gemma 3 VL's image-bidirectional mask is not yet compatible with THD packing. + cfg.dataset.enable_in_batch_packing = False # DDP settings - VLMs require no overlap cfg.ddp.overlap_grad_reduce = False @@ -179,7 +182,7 @@ def gemma3_vl_12b_sft_4gpu_h100_bf16_config() -> ConfigContainer: cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = "flash" + cfg.model.attention_backend = AttnBackend.unfused cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -221,6 +224,8 @@ def gemma3_vl_12b_sft_4gpu_h100_bf16_config() -> ConfigContainer: # Dataset configuration cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = hf_path + # Gemma 3 VL's image-bidirectional mask is not yet compatible with THD packing. + cfg.dataset.enable_in_batch_packing = False # DDP settings - VLMs require no overlap cfg.ddp.overlap_grad_reduce = False @@ -289,7 +294,7 @@ def gemma3_vl_27b_sft_16gpu_h100_bf16_config() -> ConfigContainer: cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = "flash" + cfg.model.attention_backend = AttnBackend.unfused cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -331,6 +336,8 @@ def gemma3_vl_27b_sft_16gpu_h100_bf16_config() -> ConfigContainer: # Dataset configuration cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = hf_path + # Gemma 3 VL's image-bidirectional mask is not yet compatible with THD packing. + cfg.dataset.enable_in_batch_packing = False # DDP settings - VLMs require no overlap cfg.ddp.overlap_grad_reduce = False @@ -408,7 +415,7 @@ def gemma3_vl_4b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = "flash" + cfg.model.attention_backend = AttnBackend.unfused cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -450,6 +457,8 @@ def gemma3_vl_4b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> # Dataset configuration cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = hf_path + # Gemma 3 VL's image-bidirectional mask is not yet compatible with THD packing. + cfg.dataset.enable_in_batch_packing = False # DDP settings - VLMs require no overlap cfg.ddp.overlap_grad_reduce = False @@ -527,7 +536,7 @@ def gemma3_vl_12b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = "flash" + cfg.model.attention_backend = AttnBackend.unfused cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -569,6 +578,8 @@ def gemma3_vl_12b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - # Dataset configuration cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = hf_path + # Gemma 3 VL's image-bidirectional mask is not yet compatible with THD packing. + cfg.dataset.enable_in_batch_packing = False # DDP settings - VLMs require no overlap cfg.ddp.overlap_grad_reduce = False @@ -646,7 +657,7 @@ def gemma3_vl_27b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = "flash" + cfg.model.attention_backend = AttnBackend.unfused cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -688,6 +699,8 @@ def gemma3_vl_27b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - # Dataset configuration cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = hf_path + # Gemma 3 VL's image-bidirectional mask is not yet compatible with THD packing. + cfg.dataset.enable_in_batch_packing = False # DDP settings - VLMs require no overlap cfg.ddp.overlap_grad_reduce = False diff --git a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_bridge.py b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_bridge.py index edaf09c0ad6..74f041fe3d1 100644 --- a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_bridge.py +++ b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_bridge.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest import torch @@ -288,6 +288,71 @@ def test_mapping_registry_accepts_legacy_checkpoint_vision_tower_keys(self, gemm assert str(mapping.megatron_param) == "vision_tower.embeddings.patch_embedding.bias" assert str(mapping.hf_param) == "vision_tower.vision_model.embeddings.patch_embedding.bias" + export_mapping = registry.megatron_to_hf_lookup("vision_tower.embeddings.patch_embedding.bias") + + assert export_mapping is not None + assert str(export_mapping.megatron_param) == "vision_tower.embeddings.patch_embedding.bias" + assert str(export_mapping.hf_param) == "vision_tower.vision_model.embeddings.patch_embedding.bias" + + def test_build_conversion_tasks_preserves_legacy_vision_tower_owner(self, gemma3_vl_bridge): + """Test legacy HF keys do not turn the owning PP task into a receiver-only task.""" + megatron_name = "vision_tower.embeddings.patch_embedding.bias" + legacy_hf_name = "vision_tower.vision_model.embeddings.patch_embedding.bias" + weight = torch.nn.Parameter(torch.ones(2)) + module = Mock() + model = Mock() + model.config = Mock() + model.named_parameters.return_value = [(megatron_name, weight)] + + hf_pretrained = Mock() + hf_pretrained.config = Mock() + hf_pretrained.state.source.get_all_keys.return_value = {legacy_hf_name} + + with ( + patch( + "megatron.bridge.models.conversion.model_bridge.unwrap_model", + return_value=[model], + ), + patch( + "megatron.bridge.models.conversion.model_bridge.persistent_buffers", + return_value=[], + ), + patch( + "megatron.bridge.models.conversion.model_bridge._get_pg_collection_from_model", + return_value=None, + ), + patch( + "megatron.bridge.models.conversion.model_bridge._get_pp_rank", + return_value=0, + ), + patch( + "megatron.bridge.models.conversion.model_bridge._megatron_local_name_to_global", + return_value=megatron_name, + ), + patch( + "megatron.bridge.models.conversion.model_bridge.get_module_and_param_from_name", + return_value=(module, weight), + ), + patch.object( + gemma3_vl_bridge, + "_megatron_global_param_names_all_pp_ranks", + return_value=[megatron_name], + ), + patch.object( + gemma3_vl_bridge, + "_share_embeddings_and_output_weights", + return_value=False, + ), + ): + tasks = gemma3_vl_bridge.build_conversion_tasks(hf_pretrained, [model]) + + assert len(tasks) == 1 + task = tasks[0] + assert task is not None + assert task.param_weight is weight + assert task.megatron_module is module + assert str(task.mapping.hf_param) == legacy_hf_name + def test_mapping_registry_multimodal_projector_params(self, gemma3_vl_bridge): """Test mapping_registry handles multimodal projector parameters correctly.""" registry = gemma3_vl_bridge.mapping_registry() diff --git a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py new file mode 100644 index 00000000000..e87f08bdcb7 --- /dev/null +++ b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py @@ -0,0 +1,40 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import Mock + +import torch + +from megatron.bridge.models.gemma_vl.modeling_gemma3_vl import Gemma3VLModel + + +def test_attention_mask_is_available_on_later_pipeline_stages(): + """Later PP stages must receive the same image-bidirectional mask.""" + model = object.__new__(Gemma3VLModel) + model.config = Mock(image_token_id=99) + model.pre_process = False + input_ids = torch.tensor([[1, 99, 99, 2]]) + + mask = model._compute_attention_mask(input_ids) + + assert mask is not None + assert mask[0, 0, 1, 2].item() is False + assert mask[0, 0, 0, 3].item() is True + + +def test_attention_mask_requires_input_ids(): + """Mask construction is skipped only when token IDs are unavailable.""" + model = object.__new__(Gemma3VLModel) + + assert model._compute_attention_mask(None) is None diff --git a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py index ed9257af104..726341b7c7b 100644 --- a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py +++ b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py @@ -13,12 +13,18 @@ # limitations under the License. import math +from unittest.mock import patch +import pytest import torch +from megatron.core.transformer.enums import AttnBackend, AttnMaskType from transformers import SiglipVisionConfig +from megatron.bridge.models.gemma.gemma3_provider import Gemma3TEDotProductAttention, gemma3_layer_spec from megatron.bridge.models.gemma_vl.gemma3_vl_provider import Gemma3VLModelProvider from megatron.bridge.models.gemma_vl.modeling_gemma3_vl import Gemma3VLMultimodalProjectorConfig +from megatron.bridge.training.utils.config_utils import apply_run_config_backward_compat +from megatron.bridge.utils.instantiate_utils import instantiate class TestGemma3VLModelProvider: @@ -63,6 +69,8 @@ def test_gemma3_vl_vl_specific_defaults(self): ) # Check VL-specific defaults + assert provider.is_vision_language is True + assert provider.attention_backend is AttnBackend.unfused assert provider.scatter_embedding_sequence_parallel is False assert isinstance(provider.vision_config, SiglipVisionConfig) assert isinstance(provider.vision_projector_config, Gemma3VLMultimodalProjectorConfig) @@ -80,6 +88,46 @@ def test_gemma3_vl_vl_specific_defaults(self): assert provider.freeze_vision_model is False assert provider.freeze_vision_projection is False + def test_gemma3_vl_attention_consumes_custom_mask(self): + """Test that Gemma3 VL attention selects the arbitrary-mask path.""" + provider = Gemma3VLModelProvider( + num_layers=28, + hidden_size=2560, + num_attention_heads=10, + window_size=1024, + ) + + layer_spec = gemma3_layer_spec(provider) + assert layer_spec.submodules.self_attention.params["attn_mask_type"] is AttnMaskType.arbitrary + + with patch.object(Gemma3TEDotProductAttention.__bases__[0], "__init__", return_value=None) as init: + Gemma3TEDotProductAttention( + config=provider, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + ) + + assert init.call_args.kwargs["attn_mask_type"] is AttnMaskType.arbitrary + assert init.call_args.kwargs["config"].window_size == (1023, 1023) + + def test_legacy_checkpoint_attention_settings_are_dropped(self): + """Test that old checkpoint values cannot disable the VL mask.""" + legacy_config = { + "_target_": "megatron.bridge.models.gemma_vl.gemma3_vl_provider.Gemma3VLModelProvider", + "num_layers": 28, + "hidden_size": 2560, + "num_attention_heads": 10, + "is_vision_language": False, + "attention_backend": AttnBackend.flash, + } + + sanitized_config = apply_run_config_backward_compat(legacy_config) + provider = instantiate(sanitized_config) + + assert provider.is_vision_language is True + assert provider.attention_backend is AttnBackend.unfused + def test_gemma3_vl_custom_vision_config(self): """Test Gemma3VLModelProvider with custom vision configuration.""" custom_vision_config = SiglipVisionConfig( @@ -105,6 +153,42 @@ def test_gemma3_vl_custom_vision_config(self): assert provider.vision_config.image_size == 448 assert provider.vision_config.patch_size == 14 + def test_gemma3_vl_rejects_context_parallelism(self): + """Test that unsupported context parallelism fails before model construction.""" + provider = Gemma3VLModelProvider( + num_layers=28, + hidden_size=2560, + num_attention_heads=10, + context_parallel_size=2, + ) + + with pytest.raises(ValueError, match="does not support context parallelism"): + provider.provide() + + def test_gemma3_vl_rejects_causal_override(self): + """Test that runtime overrides cannot disable the VL attention mask.""" + provider = Gemma3VLModelProvider( + num_layers=28, + hidden_size=2560, + num_attention_heads=10, + ) + provider.is_vision_language = False + + with pytest.raises(ValueError, match="requires is_vision_language=True"): + provider.provide() + + def test_gemma3_vl_rejects_flash_attention_override(self): + """Test that runtime overrides cannot select a backend that drops the VL mask.""" + provider = Gemma3VLModelProvider( + num_layers=28, + hidden_size=2560, + num_attention_heads=10, + ) + provider.attention_backend = AttnBackend.flash + + with pytest.raises(ValueError, match="requires the unfused attention backend"): + provider.provide() + def test_gemma3_vl_custom_vision_projector_config(self): """Test Gemma3VLModelProvider with custom vision projector configuration.""" custom_projector_config = Gemma3VLMultimodalProjectorConfig( diff --git a/tests/unit_tests/recipes/test_gemma3_vl_recipes.py b/tests/unit_tests/recipes/test_gemma3_vl_recipes.py index a22f9ab682e..7a6dfbdf945 100644 --- a/tests/unit_tests/recipes/test_gemma3_vl_recipes.py +++ b/tests/unit_tests/recipes/test_gemma3_vl_recipes.py @@ -25,6 +25,7 @@ import pytest import torch +from megatron.core.transformer.enums import AttnBackend from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -96,6 +97,8 @@ def _assert_basic_config(cfg): assert cfg.tokenizer is not None assert cfg.checkpoint is not None assert cfg.rng is not None + assert cfg.model.attention_backend is AttnBackend.unfused + assert cfg.dataset.enable_in_batch_packing is False assert cfg.train.global_batch_size >= 1 assert cfg.train.micro_batch_size >= 1 diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index 8d1f80dd93c..9d33dd71879 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -20,6 +20,7 @@ import os import sys +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -90,6 +91,63 @@ class TestCompareMaskHandling: """Tests for attention_mask handling in compare.py Megatron and HF paths.""" + def test_gemma3_config_is_detected_as_vision_language_model(self): + """Structured vision and text configs identify Gemma 3 as a VLM.""" + config = SimpleNamespace( + model_type="gemma3", + architectures=["Gemma3ForConditionalGeneration"], + text_config=object(), + vision_config=object(), + ) + + with patch.object(compare.AutoConfig, "from_pretrained", return_value=config): + assert compare.is_vision_language_model("google/gemma-3-4b-it") is True + + def test_vlm_inputs_preserve_image_token_types_and_tp_padding(self): + """VLM preprocessing keeps Gemma image inputs without requiring a grid.""" + input_ids = torch.tensor([[1, 2, 3]]) + token_type_ids = torch.tensor([[0, 1, 0]]) + pixel_values = torch.randn(1, 3, 4, 4) + processed = { + "input_ids": input_ids, + "pixel_values": pixel_values, + "token_type_ids": token_type_ids, + } + processor = MagicMock() + processor.apply_chat_template.return_value = processed + tokenizer = SimpleNamespace(pad_token_id=7) + image = object() + + with patch.object(compare, "load_image", return_value=image): + actual_input_ids, actual_pixels, image_grid_thw, actual_token_type_ids = compare.process_inputs( + tokenizer, + processor, + "/tmp/example.png", + "Describe this image.", + is_vl_model=True, + tp_size=2, + ) + + torch.testing.assert_close(actual_input_ids, torch.tensor([[1, 2, 3, 7]])) + assert actual_pixels is pixel_values + assert image_grid_thw is None + torch.testing.assert_close(actual_token_type_ids, torch.tensor([[0, 1, 0, 0]])) + processor.apply_chat_template.assert_called_once_with( + [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": "Describe this image."}, + ], + } + ], + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt", + ) + def test_single_batch_iterator_stores_none_attention_mask(self): """Test that SingleBatchIterator preserves None attention_mask in batch dict.""" input_ids = torch.tensor([[1, 2, 3]]) @@ -210,6 +268,32 @@ def test_hf_path_receives_ones_like_attention_mask(self): assert call_kwargs["attention_mask"].shape == input_ids.shape assert torch.equal(call_kwargs["attention_mask"], expected_mask) + def test_hf_path_receives_multimodal_token_type_ids(self): + """Gemma 3 token types reach HF so its image attention mask matches Megatron.""" + mock_hf_model = MagicMock() + mock_output = MagicMock() + mock_output.logits = torch.randn(1, 3, 100) + mock_hf_model.return_value = mock_output + input_ids = torch.tensor([[1, 2, 3]]) + token_type_ids = torch.tensor([[0, 1, 0]]) + mock_tokenizer = MagicMock() + mock_tokenizer.decode.return_value = "test" + + with ( + patch.object(compare, "_is_rank_0", return_value=True), + patch.object(compare, "print_rank_0"), + ): + _run_hf_inference( + mock_hf_model, + input_ids, + pixel_values=torch.randn(1, 3, 4, 4), + image_grid_thw=None, + tokenizer=mock_tokenizer, + token_type_ids=token_type_ids, + ) + + assert mock_hf_model.call_args.kwargs["token_type_ids"] is token_type_ids + def test_hf_broadcast_uses_model_output_vocab_size(self): """Test that non-rank-0 buffers use the HF logits size instead of tokenizer vocab size.""" broadcast_shapes = [] From 77e25d2dab00453d82439b9d593b13de3b29cf40 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 01:45:04 +0200 Subject: [PATCH 13/85] cp: `fix(training): preserve alltoall flex fallback (5207)` into `r0.6.0` (#5214) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../training/flex_dispatcher_backend.py | 12 ++++- tests/unit_tests/training/test_deepep.py | 44 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/training/flex_dispatcher_backend.py b/src/megatron/bridge/training/flex_dispatcher_backend.py index 4445dc74a37..7431c24214c 100644 --- a/src/megatron/bridge/training/flex_dispatcher_backend.py +++ b/src/megatron/bridge/training/flex_dispatcher_backend.py @@ -24,6 +24,12 @@ logger: logging.Logger = logging.getLogger(__name__) +def _fallback_to_alltoall(model_config: TransformerConfig) -> None: + """Clear flex dispatcher state when the requested backend cannot be used.""" + model_config.moe_token_dispatcher_type = "alltoall" + model_config.moe_flex_dispatcher_backend = None + + def apply_flex_dispatcher_backend( model_config: TransformerConfig, moe_flex_dispatcher_backend: str | None = None, @@ -51,16 +57,18 @@ def apply_flex_dispatcher_backend( if get_rank_safe() == 0: logger.warning( f"DeepEP is only applicable to Ampere, Hopper, and Blackwell (B200/B300) GPUs. " - f"Current GPU: {device_properties.name}. Skipping DeepEP configuration." + f"Current GPU: {device_properties.name}. Falling back to alltoall." ) + _fallback_to_alltoall(model_config) return elif moe_flex_dispatcher_backend == "hybridep": if not device_properties.major in [8, 9, 10]: if get_rank_safe() == 0: logger.warning( f"HybridEP is only applicable for GB200, GB300 with NVL72 and for Ampere, Hopper, B200 and B300 GPUs. " - f"Current GPU: {device_properties.name}. Skipping HybridEP configuration." + f"Current GPU: {device_properties.name}. Falling back to alltoall." ) + _fallback_to_alltoall(model_config) return else: if get_rank_safe() == 0: diff --git a/tests/unit_tests/training/test_deepep.py b/tests/unit_tests/training/test_deepep.py index 50209a071bf..2a2d9c0c21b 100644 --- a/tests/unit_tests/training/test_deepep.py +++ b/tests/unit_tests/training/test_deepep.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from megatron.core.transformer import TransformerConfig +from scripts.performance.utils.utils import finalize_config_overrides from megatron.bridge.training.flex_dispatcher_backend import ( apply_flex_dispatcher_backend, @@ -185,6 +187,48 @@ def test_apply_flex_dispatcher_backend_warns_for_unsupported_gpu_pascal( assert config.moe_shared_expert_overlap != False +class TestFlexDispatcherFallback: + """Test unsupported flex backends stay disabled after config finalization.""" + + @pytest.mark.parametrize( + ("backend", "major", "device_name"), + [ + pytest.param("deepep", 10, "NVIDIA GB200", id="deepep-gb200"), + pytest.param("hybridep", 11, "NVIDIA X200", id="hybridep-unsupported"), + ], + ) + @patch("torch.cuda.get_device_properties") + def test_unsupported_backend_falls_back_to_alltoall( + self, + mock_get_device_properties, + backend, + major, + device_name, + ): + mock_properties = MagicMock() + mock_properties.major = major + mock_properties.name = device_name + mock_get_device_properties.return_value = mock_properties + config = SimpleNamespace( + model=SimpleNamespace( + num_moe_experts=8, + moe_token_dispatcher_type="alltoall", + moe_flex_dispatcher_backend=backend, + moe_shared_expert_overlap=True, + ), + ddp=SimpleNamespace(nccl_ub=False, fsdp_manual_registration=False), + ) + + apply_flex_dispatcher_backend(config.model, moe_flex_dispatcher_backend=backend) + finalize_config_overrides(config) + validate_flex_dispatcher_backend(config.model) + + assert config.model.moe_token_dispatcher_type == "alltoall" + assert config.model.moe_flex_dispatcher_backend is None + assert config.model.moe_shared_expert_overlap is True + mock_get_device_properties.assert_called_once_with(0) + + class TestValidateDeepEP: """Test the validate_flex_dispatcher_backend function.""" From 4c10b8e8f2da158fdf243365d34c836e607fd292 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 01:45:35 +0200 Subject: [PATCH 14/85] cp: `fix(model): align MiniMax evidence and inference task routing (5194)` into `r0.6.0` (#5235) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../conversion/hf_to_megatron_generate_vlm.py | 2 + examples/conversion/vlm_generate_utils.py | 36 ++ .../minimax-m3/card.yaml | 30 +- scripts/inference/README.md | 50 +- scripts/inference/setup_inference.py | 39 +- scripts/inference/vlm_generation.py | 427 ++++++++++++++++++ scripts/inference/vlm_generation_utils.py | 379 ++++++++++++++++ .../create-model-verification-card/SKILL.md | 57 ++- .../scripts/validate_card.py | 43 +- .../scripts/inference/test_setup_inference.py | 25 +- .../test_validate_model_verification_card.py | 68 +++ .../utils/test_vlm_generate_utils.py | 67 ++- 12 files changed, 1162 insertions(+), 61 deletions(-) create mode 100644 scripts/inference/vlm_generation.py create mode 100644 scripts/inference/vlm_generation_utils.py diff --git a/examples/conversion/hf_to_megatron_generate_vlm.py b/examples/conversion/hf_to_megatron_generate_vlm.py index 1358cfaaf88..44ad8e2279d 100644 --- a/examples/conversion/hf_to_megatron_generate_vlm.py +++ b/examples/conversion/hf_to_megatron_generate_vlm.py @@ -166,6 +166,7 @@ def main(args) -> None: if is_kimi and image_token_id is None: image_token_id = 163605 is_gemma4 = "gemma4" in model_type + is_minimax = model_type == "minimax_m3_vl" is_mistral3 = model_type == "mistral3" # ------------------------------------------------------------------ @@ -278,6 +279,7 @@ def _disable_mtp(m): args.prompt, is_gemma4=is_gemma4, is_kimi=is_kimi, + is_minimax=is_minimax, is_mistral3=is_mistral3, image_token_id=image_token_id, ) diff --git a/examples/conversion/vlm_generate_utils.py b/examples/conversion/vlm_generate_utils.py index c429e4936a0..bfe1574b8b7 100644 --- a/examples/conversion/vlm_generate_utils.py +++ b/examples/conversion/vlm_generate_utils.py @@ -195,6 +195,7 @@ def process_image_inputs( *, is_gemma4: bool = False, is_kimi: bool = False, + is_minimax: bool = False, is_mistral3: bool = False, image_token_id: Optional[int] = None, ): @@ -216,6 +217,9 @@ def process_image_inputs( if is_kimi: return _process_kimi_inputs(processor, image_path, prompt, image_token_id) + if is_minimax: + return _process_minimax_inputs(processor, image_path, prompt) + return _process_default_inputs(processor, image_path, prompt, allow_raw_prompt=is_mistral3) @@ -236,6 +240,38 @@ def _process_kimi_inputs(processor, image_path, prompt, image_token_id): return input_ids, inputs.pixel_values, grid_thws, None, None, None +def _process_minimax_inputs(processor, image_path, prompt): + """Format MiniMax multimodal input with its tokenizer chat template.""" + tokenizer = getattr(processor, "tokenizer", None) + if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: + raise ValueError("The MiniMax tokenizer has no chat template") + + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": prompt}, + ], + } + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = processor( + text=[text], + images=[load_image(image_path).convert("RGB")], + padding=True, + return_tensors="pt", + ) + return ( + inputs.input_ids, + inputs.get("pixel_values"), + inputs.get("image_grid_thw"), + inputs.get("image_sizes"), + inputs.get("mm_token_type_ids"), + inputs.get("image_position_ids"), + ) + + def _process_default_inputs(processor, image_path, prompt, *, allow_raw_prompt=False): if getattr(processor, "chat_template", None) is None: if not allow_raw_prompt: diff --git a/examples/model_verification_cards/minimax-m3/card.yaml b/examples/model_verification_cards/minimax-m3/card.yaml index 87a70c64768..1e869a1d5cf 100644 --- a/examples/model_verification_cards/minimax-m3/card.yaml +++ b/examples/model_verification_cards/minimax-m3/card.yaml @@ -7,10 +7,13 @@ summary: > timing and throughput metrics are sanity checks, not optimized performance results. Model-level verification covers the complete MiniMax-M3 vision-language checkpoint, including the language model, vision tower, - projector paths, and lightning-indexer state. The pinned 32-H100 GPU import, - GPU export, and one deterministic multimodal Megatron generation completed - their verification gates. All 59 exported safetensors shards are - byte-for-byte identical to the pinned source, and Transformers 5.12.1 + projector paths, and lightning-indexer state. The pinned 32-H100 GPU import + and GPU export completed their verification gates. A current clean 32-H100 + run also reloaded the persisted full-VLM checkpoint through the public + VLM generation entry point, processed the pinned repository image, generated a + coherent completion containing the exact image-only OCR target, and reached + EOS without token-0 or NUL degeneration. All 59 exported safetensors shards + are byte-for-byte identical to the pinned source, and Transformers 5.12.1 natively reloads the export without loading errors. CPU import and export remain unverified because the CPU-only attempt exited during CUDA-dependent model initialization before producing a checkpoint. Full-VLM manual Hugging @@ -135,6 +138,7 @@ items: inference: status: verified precision: bf16 + bridge_commit: 9bd762e26525ab40703f9dac0b8a935b2fc78c19 # pragma: allowlist secret command: > ./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 --task vlm-generation @@ -149,17 +153,19 @@ items: explain or reason." --max_new_tokens 64 --trust_remote_code - last_verified: 2026-07-23 - expected_result: >- + last_verified: 2026-07-31 + expected_result: > One deterministic greedy generation on 32 H100s reloaded the persisted full-VLM checkpoint at TP1/PP1/EP32/ETP1, processed docs/images/tp1.png, - generated 33 tokens under a 64-token maximum, reached EOS naturally, and - exited successfully. The exact completion, represented with JSON Unicode - escapes for the model tags and quotes, was + generated 50 tokens under a 64-token maximum, reached EOS naturally, and + exited successfully. The first five selected token IDs were 200059, 758, + 3100, 9210, and 640; the completion contained zero NUL bytes and did not + exhibit token-0/NUL degeneration. The exact literal completion, + represented with JSON Unicode escapes for the model tags and quotes, was "\u003cmm:think\u003eThe user wants me to read the green caption near the - top of the image and respond with only the exact quoted - term.\u003c\u002fmm:think\u003e\u0022Pre-LayerNorm\u0022". The completion - contained the image-only OCR target "pre-layernorm". + top of the image. The green text says \u0022Pre-LayerNorm\u0022 Transformer + Block. The exact quoted term would be + \u0022Pre-LayerNorm\u0022.\u003c\u002fmm:think\u003e\u0022Pre-LayerNorm\u0022". pretrain: H100: diff --git a/scripts/inference/README.md b/scripts/inference/README.md index 02453500893..e83f022e587 100644 --- a/scripts/inference/README.md +++ b/scripts/inference/README.md @@ -1,9 +1,9 @@ # Inference launcher -`infer.sh` submits Bridge-backed offline text generation from a Slurm login -node. It uses NeMo Run's Slurm executor and launches one srun-native process per -GPU; users should not enter an allocation or wrap it in `srun`, `torchrun`, or -`sbatch`. +`infer.sh` submits Bridge-backed offline text or vision-language generation and +HF/Megatron model comparison from a Slurm login node. It uses NeMo Run's Slurm +executor and launches one srun-native process per GPU; users should not enter an +allocation or wrap it in `srun`, `torchrun`, or `sbatch`. ```bash ./scripts/inference/infer.sh \ @@ -22,8 +22,41 @@ GPU; users should not enter an allocation or wrap it in `srun`, `torchrun`, or The launcher owns only Slurm resources, the container, mounts, explicitly forwarded environment variables, and submission behavior. All other arguments -are forwarded unchanged to `text_generation.py`, which owns model loading, -parallelism, prompts, sampling, and inference-engine configuration. +are forwarded unchanged to the selected inference entry point. Text generation +uses `text_generation.py`; `--task vlm-generation` selects +`vlm_generation.py` for multimodal generation, and `--task model-comparison` +selects `examples/conversion/compare_hf_and_megatron/compare.py` for forward-logit +parity. + +```bash +./scripts/inference/infer.sh \ + --task vlm-generation \ + --nodes 1 --gpus-per-node 4 \ + --account ACCOUNT --partition PARTITION \ + --container-image /path/to/megatron-bridge.sqsh \ + --mount /path/to/Megatron-Bridge:/opt/Megatron-Bridge \ + --env HF_TOKEN \ + --hf_model_path Qwen/Qwen2.5-VL-3B-Instruct \ + --image_path /shared/image.jpg \ + --prompt "Describe this image." \ + --tp 4 +``` + +For a one-step Hugging Face/Megatron forward-logit comparison, select the +comparison task and pass the comparison entry point's arguments unchanged: + +```bash +./scripts/inference/infer.sh \ + --task model-comparison \ + --nodes 1 --gpus-per-node 2 \ + --account ACCOUNT --partition PARTITION \ + --container-image /path/to/megatron-bridge.sqsh \ + --mount /path/to/Megatron-Bridge:/opt/Megatron-Bridge \ + --env HF_TOKEN \ + --hf_model_path Qwen/Qwen3-1.7B \ + --prompt "Hello world" \ + --tp 2 +``` ## Model and checkpoint inputs @@ -51,9 +84,8 @@ visible at the same path from every allocated compute node. Repeat `--prompt` for inline prompts, or pass a mounted line-oriented or JSONL file with `--prompt-file`. JSONL records may use a `text`, `prompt`, or `input` field. -Run `uv run python scripts/inference/text_generation.py --help` in a configured -Megatron Bridge environment for the full model, checkpoint, prompt, sampling, -parallelism, and engine CLI. +Run the selected entry point with `--help` in a configured Megatron Bridge +environment for its full CLI. ## Environment and cluster options diff --git a/scripts/inference/setup_inference.py b/scripts/inference/setup_inference.py index fd4bd17786c..ad57f32ad63 100755 --- a/scripts/inference/setup_inference.py +++ b/scripts/inference/setup_inference.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Submit Bridge-backed offline text generation through Slurm.""" +"""Submit Bridge-backed offline inference and comparison through Slurm.""" from __future__ import annotations @@ -30,12 +30,17 @@ logger = logging.getLogger(__name__) CONTAINER_REPO_ROOT = Path("/opt/Megatron-Bridge") +INFERENCE_TASKS = { + "text-generation": Path("scripts/inference/text_generation.py"), + "vlm-generation": Path("scripts/inference/vlm_generation.py"), + "model-comparison": Path("examples/conversion/compare_hf_and_megatron/compare.py"), +} def _build_parser() -> argparse.ArgumentParser: """Build the lightweight head-node parser.""" parser = argparse.ArgumentParser( - description="Launch Megatron Bridge offline text generation through Slurm.", + description="Launch Megatron Bridge offline inference or model comparison through Slurm.", formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False, epilog=""" @@ -46,12 +51,18 @@ def _build_parser() -> argparse.ArgumentParser: --hf-model-path meta-llama/Llama-3.2-1B \\ --prompt "Megatron Bridge inference is" --max_new_tokens 32 -Arguments not owned by this launcher are forwarded unchanged to -scripts/inference/text_generation.py. Run that entry point with --help to see -model, checkpoint, prompt, sampling, and inference-engine options. +Use --task vlm-generation for multimodal generation or --task model-comparison +for a one-step Hugging Face/Megatron comparison. Arguments not owned by this +launcher are forwarded unchanged to the selected repository entry point. """, ) execution = parser.add_argument_group("Execution") + execution.add_argument( + "--task", + choices=tuple(INFERENCE_TASKS), + default="text-generation", + help="Repository inference task to launch (default: text-generation).", + ) execution.add_argument("--nodes", type=int, default=1, help="Number of Slurm nodes (default: 1).") execution.add_argument( "--gpus-per-node", @@ -196,10 +207,11 @@ def _build_executor(args: argparse.Namespace, env_names: list[str], mounts: list return executor -def _build_task(inference_args: list[str]) -> run.Script: - """Build the existing Bridge text-generation task for the submitted container.""" +def _build_task(task_name: str, inference_args: list[str]) -> run.Script: + """Build the selected Bridge inference task for the submitted container.""" + task_path = CONTAINER_REPO_ROOT / INFERENCE_TASKS[task_name] return run.Script( - path=str(CONTAINER_REPO_ROOT / "scripts/inference/text_generation.py"), + path=str(task_path), entrypoint="python", env={ "PYTHONPATH": f"{CONTAINER_REPO_ROOT}/src:{CONTAINER_REPO_ROOT}/3rdparty/Megatron-LM:$PYTHONPATH", @@ -218,28 +230,29 @@ def _raise_on_failed_tasks(experiment: run.Experiment) -> None: def parse_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]: - """Parse launch arguments and preserve all text-generation arguments.""" + """Parse launch arguments and preserve all selected-task arguments.""" return _build_parser().parse_known_args(argv) def main(argv: list[str] | None = None) -> None: - """Build and submit one Bridge text-generation experiment.""" + """Build and submit one Bridge inference experiment.""" args, inference_args = parse_args(argv) _validate_args(args) env_names = _parse_env(args.env) mounts = _parse_mounts(args.mount) executor = _build_executor(args, env_names, mounts) - task = _build_task(inference_args) + task = _build_task(args.task, inference_args) + task_path = CONTAINER_REPO_ROOT / INFERENCE_TASKS[args.task] logger.info( "Inference command: %s", - shlex.join(["python", str(CONTAINER_REPO_ROOT / "scripts/inference/text_generation.py"), *inference_args]), + shlex.join(["python", str(task_path), *inference_args]), ) logger.info("Forwarded environment variables: %s", ", ".join(env_names) or "none") logger.info("Container mounts: %s", ", ".join(mounts) or "none") with run.Experiment(args.experiment_name or "inference") as experiment: - experiment.add(task, executor=executor, name="text-generation") + experiment.add(task, executor=executor, name=args.task) if args.submission_dry_run: experiment.dryrun() return diff --git a/scripts/inference/vlm_generation.py b/scripts/inference/vlm_generation.py new file mode 100644 index 00000000000..f9634f16868 --- /dev/null +++ b/scripts/inference/vlm_generation.py @@ -0,0 +1,427 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Example: + # Vision-Language generation with image from URL: + uv run python scripts/inference/vlm_generation.py --hf_model_path="Qwen/Qwen2.5-VL-3B-Instruct" --image_path="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" --prompt="Describe this image." + + # Vision-Language generation with local image: + uv run python scripts/inference/vlm_generation.py --hf_model_path="Qwen/Qwen2.5-VL-3B-Instruct" --image_path="/path/to/image.jpg" --prompt="What do you see in this image?" + + # Text-only generation (no image): + uv run python scripts/inference/vlm_generation.py --hf_model_path="Qwen/Qwen2.5-VL-3B-Instruct" --prompt="Hello, how are you?" + + # Load from Megatron checkpoint: + uv run python scripts/inference/vlm_generation.py --hf_model_path="Qwen/Qwen2.5-VL-3B-Instruct" --megatron_model_path="/path/to/megatron/checkpoint" --image_path="/path/to/image.jpg" --prompt="Describe this image." +""" + +import argparse + +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.pipeline_parallel.schedules import get_forward_backward_func +from transformers import AutoConfig, AutoProcessor, AutoTokenizer +from vlm_generation_utils import ( + pad_input_ids_to_tp_multiple, + patch_kimi_vision_processor, + process_image_inputs, + process_multi_image_inputs, + process_video_inputs, + to_cuda, +) + +from megatron.bridge import AutoBridge +from megatron.bridge.models.hf_pretrained.utils import is_safe_repo +from megatron.bridge.utils.common_utils import get_last_rank, print_rank_0, print_rank_last + + +# --------------------------------------------------------------------------- +# Forward step +# --------------------------------------------------------------------------- + + +class SingleBatchIterator: + """Iterator that yields a single batch then stops. Required by + ``get_forward_backward_func``.""" + + def __init__( + self, + input_ids, + position_ids, + attention_mask, + pixel_values=None, + image_grid_thw=None, + image_sizes=None, + mm_token_type_ids=None, + pixel_values_videos=None, + video_grid_thw=None, + image_position_ids=None, + ): + self.batch = dict( + tokens=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + ) + if pixel_values is not None: + self.batch["pixel_values"] = pixel_values + if image_grid_thw is not None: + self.batch["image_grid_thw"] = image_grid_thw + if image_sizes is not None: + self.batch["image_sizes"] = image_sizes + if mm_token_type_ids is not None: + self.batch["mm_token_type_ids"] = mm_token_type_ids + if pixel_values_videos is not None: + self.batch["pixel_values_videos"] = pixel_values_videos + if video_grid_thw is not None: + self.batch["video_grid_thw"] = video_grid_thw + if image_position_ids is not None: + self.batch["image_position_ids"] = image_position_ids + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if self._yielded: + raise StopIteration + self._yielded = True + return self.batch + + +def vlm_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: + """Forward step for VLM generation.""" + batch = next(data_iterator) + forward_args = { + "input_ids": batch["tokens"], + "position_ids": batch["position_ids"], + "attention_mask": batch.get("attention_mask"), + } + for key in ( + "pixel_values", + "image_grid_thw", + "image_sizes", + "mm_token_type_ids", + "pixel_values_videos", + "video_grid_thw", + "image_position_ids", + ): + if key in batch: + forward_args[key] = batch[key] + + def loss_func(x, **kwargs): + return x + + model_output = model(**forward_args) + if isinstance(model_output, tuple): + output_tensor, _ = model_output + else: + output_tensor = model_output + return output_tensor, loss_func + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def _hf_revision_kwargs(revision: str | None) -> dict[str, str]: + """Build keyword arguments for revision-pinned Hugging Face loads.""" + return {"revision": revision} if revision is not None else {} + + +def main(args) -> None: + """Run VLM inference with HuggingFace or Megatron checkpoints.""" + tp = args.tp + pp = args.pp + ep = args.ep + etp = args.etp + + trust_remote = is_safe_repo( + trust_remote_code=args.trust_remote_code, + hf_path=args.hf_model_path, + ) + + # Detect model family for processor-specific handling + config = AutoConfig.from_pretrained( + args.hf_model_path, + trust_remote_code=trust_remote, + **_hf_revision_kwargs(args.hf_revision), + ) + model_type = getattr(config, "model_type", "") + is_kimi = "kimi" in model_type + image_token_id = getattr(config, "image_token_id", None) + if is_kimi and image_token_id is None: + image_token_id = 163605 + is_gemma4 = "gemma4" in model_type + is_minimax = model_type == "minimax_m3_vl" + is_mistral3 = model_type == "mistral3" + + # ------------------------------------------------------------------ + # Load model + # ------------------------------------------------------------------ + bridge = AutoBridge.from_hf_pretrained( + args.hf_model_path, + trust_remote_code=trust_remote, + **_hf_revision_kwargs(args.hf_revision), + ) + + if args.megatron_model_path: + print_rank_0(f"Loading Megatron model from: {args.megatron_model_path}") + model_provider = bridge.to_megatron_provider(load_weights=False) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.init_model_with_meta_device = True + if args.pp_layout: + model_provider.pipeline_model_parallel_layout = args.pp_layout + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0) + + mp_overrides = { + "tensor_model_parallel_size": tp, + "pipeline_model_parallel_size": pp, + "expert_model_parallel_size": ep, + "expert_tensor_parallel_size": etp, + "pipeline_dtype": torch.bfloat16, + } + if args.pp_layout: + mp_overrides["pipeline_model_parallel_layout"] = args.pp_layout + model = bridge.load_megatron_model( + args.megatron_model_path, + mp_overrides=mp_overrides, + wrap_with_ddp=False, + ) + else: + print_rank_0(f"Loading HuggingFace model from: {args.hf_model_path}") + model_provider = bridge.to_megatron_provider(load_weights=True) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.expert_tensor_parallel_size = etp + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0) + model = model_provider.provide_distributed_model(wrap_with_ddp=False) + + def _disable_mtp(m): + m.config.mtp_num_layers = None + inner = m.module if hasattr(m, "module") else m + lang = getattr(inner, "language_model", inner) + if hasattr(lang, "mtp_process"): + lang.mtp_process = False + + model = [m.cuda() for m in model] + for m in model: + m.eval() + _disable_mtp(m) + if hasattr(m, "config"): + m.config.grad_scale_func = None + + # ------------------------------------------------------------------ + # Tokenizer & processor + # ------------------------------------------------------------------ + tokenizer = AutoTokenizer.from_pretrained( + args.hf_model_path, + trust_remote_code=trust_remote, + **_hf_revision_kwargs(args.hf_revision), + ) + if is_kimi: + patch_kimi_vision_processor(args.hf_model_path) + processor = AutoProcessor.from_pretrained( + args.hf_model_path, + trust_remote_code=trust_remote, + **_hf_revision_kwargs(args.hf_revision), + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + pad_token_id = tokenizer.pad_token_id or 0 + + # ------------------------------------------------------------------ + # Process inputs + # ------------------------------------------------------------------ + pixel_values = image_grid_thw = image_sizes = mm_token_type_ids = image_position_ids = None + pixel_values_videos = video_grid_thw = None + + if args.video_path: + input_ids_raw, pixel_values_videos, video_grid_thw = process_video_inputs( + processor, args.video_path, args.prompt, fps=args.video_fps + ) + elif args.image_paths: + input_ids_raw, pixel_values, image_grid_thw = process_multi_image_inputs( + processor, args.image_paths, args.prompt + ) + else: + ( + input_ids_raw, + pixel_values, + image_grid_thw, + image_sizes, + mm_token_type_ids, + image_position_ids, + ) = process_image_inputs( + processor, + args.image_path, + args.prompt, + is_gemma4=is_gemma4, + is_kimi=is_kimi, + is_minimax=is_minimax, + is_mistral3=is_mistral3, + image_token_id=image_token_id, + ) + + input_ids_raw = input_ids_raw.cuda() + pixel_values = to_cuda(pixel_values) + image_grid_thw = to_cuda(image_grid_thw) + image_sizes = to_cuda(image_sizes) + mm_token_type_ids = to_cuda(mm_token_type_ids) + pixel_values_videos = to_cuda(pixel_values_videos) + video_grid_thw = to_cuda(video_grid_thw) + image_position_ids = to_cuda(image_position_ids) + + # ------------------------------------------------------------------ + # Greedy generation loop + # ------------------------------------------------------------------ + generated_ids = input_ids_raw.clone() + stop_tokens = [tokenizer.eos_token_id] + + for step in range(args.max_new_tokens): + with torch.no_grad(): + print_rank_0(f"Generation step {step}") + + real_seq_len = generated_ids.size(1) + input_ids = pad_input_ids_to_tp_multiple(generated_ids, tp, pad_token_id) + + mm_ids_padded = None + if mm_token_type_ids is not None: + mm_ids_padded = pad_input_ids_to_tp_multiple(mm_token_type_ids, tp, 0) + + position_ids = ( + torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device) + .unsqueeze(0) + .expand_as(input_ids) + ) + + fwd_bwd_function = get_forward_backward_func() + iterator = SingleBatchIterator( + input_ids, + position_ids, + None, + pixel_values, + image_grid_thw, + image_sizes, + mm_ids_padded, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + image_position_ids=image_position_ids, + ) + + output = fwd_bwd_function( + forward_step_func=vlm_forward_step, + data_iterator=iterator, + model=model, + num_microbatches=1, + forward_only=True, + seq_length=input_ids.size(1), + micro_batch_size=1, + collect_non_loss_data=True, + ) + if isinstance(output, list) and len(output) > 0: + output = output[0] + + if parallel_state.is_pipeline_last_stage(): + world_size = parallel_state.get_tensor_model_parallel_world_size() + gathered_tensors = [torch.zeros_like(output) for _ in range(world_size)] + dist.all_gather(gathered_tensors, output, group=parallel_state.get_tensor_model_parallel_group()) + output = torch.cat(gathered_tensors, dim=2) + + last_pos = real_seq_len - 1 + next_token_ids = torch.argmax(output[:, last_pos], dim=-1, keepdim=True) + + if step < 5: + print_rank_last( + f"Step {step}: output shape={output.shape}, " + f"real_seq_len={real_seq_len}, var={output.var():.4f}" + ) + logits = output[0, last_pos, :] + top5_vals, top5_ids = torch.topk(logits, 5) + top5_tokens = [tokenizer.decode([idx]) for idx in top5_ids] + print_rank_last(f"Top 5: {list(zip(top5_tokens, top5_vals.tolist()))}") + print_rank_last( + f"Selected: '{tokenizer.decode([next_token_ids.item()])}' (id={next_token_ids.item()})" + ) + else: + next_token_ids = torch.ones((1, 1), device=generated_ids.device, dtype=generated_ids.dtype) + + torch.distributed.broadcast(next_token_ids, get_last_rank()) + generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1) + + if mm_token_type_ids is not None: + mm_token_type_ids = torch.cat( + [mm_token_type_ids, torch.zeros_like(next_token_ids, dtype=mm_token_type_ids.dtype)], dim=-1 + ) + + if next_token_ids.item() in stop_tokens: + break + + generated_text = tokenizer.decode(list(generated_ids[0])) + print_rank_0("======== GENERATED TEXT OUTPUT ========") + if args.image_path: + print_rank_0(f"Image: {args.image_path}") + print_rank_0(f"Prompt: {args.prompt}") + print_rank_0(f"Generated: {generated_text}") + print_rank_0("=======================================") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="VLM Generation from HuggingFace Models") + parser.add_argument("--hf_model_path", type=str, required=True, help="Path to the HuggingFace VL model.") + parser.add_argument( + "--hf-revision", + dest="hf_revision", + default=None, + help="Hugging Face revision for reproducible model and processor loading.", + ) + parser.add_argument("--prompt", type=str, default="Describe this image.", help="Input prompt.") + parser.add_argument("--max_new_tokens", type=int, default=20, help="Maximum number of new tokens to generate.") + parser.add_argument("--tp", type=int, default=1, help="Tensor parallelism size") + parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size") + parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size") + parser.add_argument("--etp", type=int, default=1, help="Expert tensor parallelism size") + parser.add_argument( + "--pp_layout", type=str, default=None, help="Pipeline model parallel layout (e.g. 'Et*15|t*15|t*16|t*15L')" + ) + parser.add_argument("--megatron_model_path", type=str, default=None, help="Path to Megatron model checkpoint") + parser.add_argument("--image_path", type=str, default=None, help="Path or URL to a single image (optional).") + parser.add_argument( + "--image_paths", + type=str, + nargs="+", + default=None, + help="Paths to N image files in order (multi-image; Qwen-family only).", + ) + parser.add_argument("--video_path", type=str, default=None, help="Path to a video file (Qwen-family only).") + parser.add_argument( + "--video_fps", type=float, default=2.0, help="Frames per second to sample from the video (default: 2.0)." + ) + parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code for HF model loading") + args = parser.parse_args() + + main(args) + + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/scripts/inference/vlm_generation_utils.py b/scripts/inference/vlm_generation_utils.py new file mode 100644 index 00000000000..d86a666fccd --- /dev/null +++ b/scripts/inference/vlm_generation_utils.py @@ -0,0 +1,379 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared utilities for the canonical VLM generation entry point.""" + +import io +import logging +from typing import Optional + +import torch +from PIL import Image + +from megatron.bridge.utils.safe_url import is_safe_public_http_url, safe_url_open + + +logger = logging.getLogger(__name__) + + +try: + from qwen_vl_utils import process_vision_info + + _HAS_QWEN_VL_UTILS = True +except ImportError: + _HAS_QWEN_VL_UTILS = False + + +def patch_kimi_vision_processor(hf_model_path: str): + """Monkey-patch KimiK25VisionProcessor.from_dict to avoid duplicate keyword errors. + + The upstream from_dict passes both **config and **kwargs to cls(), which causes + 'got multiple values for keyword argument' when AutoProcessor injects '_from_auto'. + """ + try: + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + klass = get_class_from_dynamic_module( + "kimi_k25_vision_processing.KimiK25VisionProcessor", + hf_model_path, + ) + if klass is None or getattr(klass, "_from_dict_patched", False): + return + + @classmethod # type: ignore[misc] + def _patched_from_dict(cls, config_dict, **kwargs): + config = config_dict.copy() + for key in list(kwargs.keys()): + config.pop(key, None) + media_proc_cfg = config.pop("media_proc_cfg", {}) + return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) + + klass.from_dict = _patched_from_dict + klass._from_dict_patched = True + except Exception: + pass + + +def pre_expand_vision_tokens(input_ids, grid_thws, vision_token_id, spatial_merge_size=2): + """Pre-expand single vision placeholders (image or video) to N placeholders + matching vision feature count. + + With PP > 1 the pipeline schedule needs to know the actual sequence length + upfront. Dynamic expansion inside the model changes seq_length during + forward, causing send/recv shape mismatches. Pre-expanding makes the model + use the 1:1 replacement path (is_pre_expanded=True). + """ + if grid_thws is None: + return input_ids + + feature_counts = [] + for grid_thw in grid_thws: + t, h, w = grid_thw.tolist() + num_features = int(t * (h // spatial_merge_size) * (w // spatial_merge_size)) + feature_counts.append(num_features) + + expanded = [] + feat_idx = 0 + for token_id in input_ids[0]: + if token_id.item() == vision_token_id and feat_idx < len(feature_counts): + expanded.extend([vision_token_id] * feature_counts[feat_idx]) + feat_idx += 1 + else: + expanded.append(token_id.item()) + + return torch.tensor([expanded], dtype=input_ids.dtype, device=input_ids.device) + + +def pad_input_ids_to_tp_multiple(input_ids, tp_size: int, pad_token_id: int = 0): + """Pad input_ids so sequence length is divisible by tp_size. + + Needed for sequence-parallel, which is required for MoE models using TP + EP. + No-op when tp_size is 1 or the sequence is already aligned. + """ + if tp_size <= 1: + return input_ids + seq_len = input_ids.shape[1] + remainder = seq_len % tp_size + if remainder == 0: + return input_ids + pad_len = tp_size - remainder + padding = torch.full((input_ids.shape[0], pad_len), pad_token_id, dtype=input_ids.dtype, device=input_ids.device) + return torch.cat([input_ids, padding], dim=1) + + +def load_image(image_path: str) -> Image.Image: + """Load an image from URL or file path.""" + if image_path.startswith(("http://", "https://")): + is_safe, reason = is_safe_public_http_url(image_path) + if not is_safe: + raise ValueError(f"Refusing to fetch image URL ({reason}): {image_path}") + with safe_url_open(image_path) as resp: + return Image.open(io.BytesIO(resp.read())) + return Image.open(image_path) + + +def to_cuda(x): + """Move a tensor, list of tensors, or None to CUDA.""" + if x is None: + return None + if isinstance(x, (list, tuple)): + return [t.cuda() for t in x] + return x.cuda() + + +def process_multi_image_inputs(processor, image_paths: list[str], prompt: str): + """Process N ordered images + prompt into model inputs (Qwen-family). + + Returns: + (input_ids, pixel_values, image_grid_thw) + """ + if not _HAS_QWEN_VL_UTILS: + raise ImportError("qwen-vl-utils required: pip install qwen-vl-utils") + pils = [load_image(p).convert("RGB") for p in image_paths] + messages = [ + { + "role": "user", + "content": [{"type": "image", "image": p} for p in pils] + [{"type": "text", "text": prompt}], + } + ] + image_inputs, video_inputs = process_vision_info(messages) + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt") + return inputs.input_ids, inputs.get("pixel_values"), inputs.get("image_grid_thw") + + +def process_video_inputs(processor, video_path: str, prompt: str, *, fps: float = 2.0): + """Process a video + prompt into model inputs (Qwen-family). + + Frame decoding mirrors the Qwen3-VL training pipeline: fetch_video decodes at + ``fps``, then video_processor is called with do_sample_frames=False to use the + pre-decoded frames as-is. + + Returns: + (input_ids, pixel_values_videos, video_grid_thw) + """ + if not _HAS_QWEN_VL_UTILS: + raise ImportError("qwen-vl-utils required: pip install qwen-vl-utils") + from qwen_vl_utils import fetch_video + + frames = fetch_video({"video": video_path, "fps": fps}) + messages = [ + { + "role": "user", + "content": [{"type": "video"}, {"type": "text", "text": prompt}], + } + ] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text_inputs = processor(text=[text], padding=True, return_tensors="pt") + video_proc = processor.video_processor(videos=[frames], return_tensors="pt", do_sample_frames=False) + # processor(text=...) without videos produces a single <|video_pad|> placeholder. + # Pre-expand to match actual vision feature count so PP send/recv shapes are correct. + video_token_id = processor.tokenizer.convert_tokens_to_ids("<|video_pad|>") + input_ids = pre_expand_vision_tokens( + text_inputs["input_ids"], + video_proc["video_grid_thw"], + vision_token_id=video_token_id, + ) + return input_ids, video_proc.get("pixel_values_videos"), video_proc.get("video_grid_thw") + + +def process_image_inputs( + processor, + image_path: Optional[str], + prompt: str, + *, + is_gemma4: bool = False, + is_kimi: bool = False, + is_minimax: bool = False, + is_mistral3: bool = False, + image_token_id: Optional[int] = None, +): + """Process image + prompt into model inputs. + + Returns: + (input_ids, pixel_values, image_grid_thw, image_sizes, mm_token_type_ids, + image_position_ids) + + Fields not applicable to the current model are None. + """ + if is_gemma4: + return _process_gemma4_inputs(processor, image_path, prompt) + + if not image_path: + inputs = processor(text=[prompt], return_tensors="pt") + return inputs.input_ids, None, None, None, None, None + + if is_kimi: + return _process_kimi_inputs(processor, image_path, prompt, image_token_id) + + if is_minimax: + return _process_minimax_inputs(processor, image_path, prompt) + + return _process_default_inputs(processor, image_path, prompt, allow_raw_prompt=is_mistral3) + + +def _process_kimi_inputs(processor, image_path, prompt, image_token_id): + messages = [ + {"role": "system", "content": "You are Kimi, an AI assistant created by Moonshot AI."}, + { + "role": "user", + "content": [ + {"type": "image", "image_url": load_image(image_path)}, + {"type": "text", "text": prompt}, + ], + }, + ] + inputs = processor(messages=messages) + grid_thws = getattr(inputs, "grid_thws", None) + input_ids = pre_expand_vision_tokens(inputs.input_ids, grid_thws, image_token_id) + return input_ids, inputs.pixel_values, grid_thws, None, None, None + + +def _process_minimax_inputs(processor, image_path, prompt): + """Format MiniMax multimodal input with its tokenizer chat template.""" + tokenizer = getattr(processor, "tokenizer", None) + if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: + raise ValueError("The MiniMax tokenizer has no chat template") + + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": prompt}, + ], + } + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = processor( + text=[text], + images=[load_image(image_path).convert("RGB")], + padding=True, + return_tensors="pt", + ) + return ( + inputs.input_ids, + inputs.get("pixel_values"), + inputs.get("image_grid_thw"), + inputs.get("image_sizes"), + inputs.get("mm_token_type_ids"), + inputs.get("image_position_ids"), + ) + + +def _process_default_inputs(processor, image_path, prompt, *, allow_raw_prompt=False): + if getattr(processor, "chat_template", None) is None: + if not allow_raw_prompt: + raise ValueError("The processor has no chat template and no model-specific raw-prompt fallback") + image = load_image(image_path).convert("RGB") + image_token = getattr(processor, "image_token", None) or getattr( + processor.tokenizer, "image_token", "<|image|>" + ) + if image_token not in prompt: + prompt = f"{image_token}\n{prompt}" + inputs = processor( + text=[prompt], + images=[image], + padding=True, + return_tensors="pt", + ) + return ( + inputs.input_ids, + inputs.get("pixel_values"), + inputs.get("image_grid_thw"), + inputs.get("image_sizes"), + inputs.get("mm_token_type_ids"), + inputs.get("image_position_ids"), + ) + + if not _HAS_QWEN_VL_UTILS: + raise ImportError( + "qwen_vl_utils is required for chat-template VLM image processing. Install it with: pip install qwen-vl-utils" + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "image": image_path}, + {"type": "text", "text": prompt}, + ], + } + ] + image_inputs, video_inputs = process_vision_info(messages) + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + return ( + inputs.input_ids, + inputs.get("pixel_values"), + inputs.get("image_grid_thw"), + inputs.get("image_sizes"), + inputs.get("mm_token_type_ids"), + inputs.get("image_position_ids"), + ) + + +def _process_gemma4_inputs(processor, image_path, prompt): + """Process inputs for Gemma 4 (text-only or vision+text). + + Uses apply_chat_template for instruction-tuned models; falls back to raw + processor with manual image token injection for base models. + + Returns: + (input_ids, pixel_values, None, None, None, image_position_ids) + """ + if image_path: + image = load_image(image_path) + has_chat_template = ( + hasattr(processor, "apply_chat_template") and getattr(processor, "chat_template", None) is not None + ) + if has_chat_template: + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": prompt}, + ], + } + ] + inputs = processor.apply_chat_template( + messages, + tokenize=True, + return_dict=True, + return_tensors="pt", + add_generation_prompt=True, + ) + else: + image_token = getattr(processor.tokenizer, "image_token", "<|image|>") + if image_token not in prompt: + prompt = image_token + "\n" + prompt + inputs = processor(text=[prompt], images=[image], return_tensors="pt") + else: + inputs = processor.tokenizer(text=[prompt], return_tensors="pt") + + return ( + inputs["input_ids"], + inputs.get("pixel_values"), + None, + None, + None, + inputs.get("image_position_ids"), + ) diff --git a/skills/create-model-verification-card/SKILL.md b/skills/create-model-verification-card/SKILL.md index af01ff5f02a..5a03eeb6df0 100644 --- a/skills/create-model-verification-card/SKILL.md +++ b/skills/create-model-verification-card/SKILL.md @@ -264,12 +264,33 @@ model-verification workload: conversion, with portable node and GPU counts; - use `scripts/training/train.sh --nodes ... --gpus-per-node ...` for every training item; -- use `uv run python ...` only for inference helpers that do not yet have a - public Slurm executor; +- use `scripts/inference/infer.sh --nodes ... --gpus-per-node ... --task ...` + for Megatron inference and manual model comparison; select + `text-generation`, `vlm-generation`, or `model-comparison` explicitly; +- invoke public shell launchers directly from the card. Do not create Slurm + jobs by calling their Python setup modules, and do not wrap the launchers in + `uv run`, `python`, `srun`, or `sbatch`; the shell entry point owns its Python + environment and scheduler setup; +- use `uv run python ...` only for a local artifact verifier that has no public + shell wrapper, such as deterministic Hugging Face inference after export. It + must not be used as a substitute for an available shell launcher or to create + a cluster job; - use short, ignored repository-relative logical paths under `work/...`; prefer aliases such as `work/data/` and `work/cache/` over reproducing a physical storage hierarchy. +For example, a portable multi-node VLM verification starts with: + +```bash +./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 \ + --task vlm-generation \ + --hf_model_path / \ + --megatron_model_path work/model-verification//iter_0000000 \ + --image_path docs/images/tp1.png \ + --prompt "Describe this image." \ + --max_new_tokens 64 +``` + The public launchers may read their required generic Slurm configuration from the caller's environment. Do not include `srun`, `sbatch`, concrete account or partition values, container image arguments, `--mount`, `--env`, shell exports, @@ -564,10 +585,11 @@ result. Private executor configuration stays outside the card. expected to be lossless; otherwise state the numerical tolerance. Do not use `--detach` or a dry-run flag in a verified conversion command. - **Manual forward pass:** Compare Hugging Face and Megatron logits on the same - prompt with `examples/conversion/compare_hf_and_megatron/compare.py`. Record - whether the next token matches, the cosine similarity, and the maximum and - mean absolute logit differences. For new evidence, pass `--hf-revision` with - the exact `model.hf_revision` so the command itself is reproducibly pinned. + prompt with `scripts/inference/infer.sh --task model-comparison`, which routes + to `examples/conversion/compare_hf_and_megatron/compare.py`. Record whether + the next token matches, the cosine similarity, and the maximum and mean + absolute logit differences. For new evidence, pass `--hf-revision` with the + exact `model.hf_revision` so the command itself is reproducibly pinned. Historical evidence verified before 2026-07-20, when the helper gained explicit revision pinning, may remain verified without a rerun when its clean-run provenance is tied to the card's immutable `model.hf_revision`; @@ -581,10 +603,11 @@ result. Private executor configuration stays outside the card. generation. Choose a prompt whose tokenized length is divisible by TP so the helper does not append padding before selecting the compared next-token position. -- **Megatron inference:** Disable sampling, run one deterministic greedy - generation with an exact token count, and record the literal completion - including whitespace. A second replay may help diagnose nondeterminism, but - it is not required verification evidence. +- **Megatron inference:** Launch through `scripts/inference/infer.sh` with the + explicit `text-generation` or `vlm-generation` task. Disable sampling, run + one deterministic greedy generation with an exact token count, and record + the literal completion including whitespace. A second replay may help + diagnose nondeterminism, but it is not required verification evidence. - **Pretrain:** Use a bounded public dataset description and a stable schedule. Save a middle and final checkpoint when resume is in scope. For expensive workloads, a 100-step reference with checkpoints at steps 50 and 100 is a @@ -597,9 +620,12 @@ result. Private executor configuration stays outside the card. - **SFT export and inference:** Depend on `sft`, export its final full-model checkpoint to HF, reload the exported model with Transformers, and run one deterministic greedy generation. Store this item as an ordered `commands` - list containing exactly two strings: the synchronous Slurm export first and - the `uv run` HF inference second. Specify an exact new-token count and record - the literal completion, including whitespace, in `expected_result`. + list containing exactly two strings: the synchronous `convert.sh` Slurm + export first and the local `uv run` HF artifact verifier second. The second + command is an explicit exception because it validates the exported HF + artifact and does not create a scheduler job. Specify an exact new-token + count and record the literal completion, including whitespace, in + `expected_result`. - **Long-context SFT:** Verify sequence packing and CP together. Record CP only when its size is greater than one. - **Checkpoint resume:** Depend on `pretrain`; load its middle checkpoint @@ -737,7 +763,10 @@ an item verified merely to make validation pass. guarding on them. New evidence must pass the exact `model.hf_revision` through `--hf-revision`; retain older unpinned evidence only under the explicitly documented grandfathering rule above. -- Use `convert.sh --executor slurm` for conversion and `train.sh` for training. +- Use `convert.sh --executor slurm` for conversion, `train.sh` for training, + and `infer.sh --task ...` for Megatron inference and model comparison. Invoke + these shell launchers directly; never call their Python setup modules to + create jobs. - Keep private executor wiring out of commands: no mounts, environment forwarding, concrete accounts/partitions/images, or remote-launch setup. - Put every training result under its canonical public hardware key and include diff --git a/skills/create-model-verification-card/scripts/validate_card.py b/skills/create-model-verification-card/scripts/validate_card.py index cf0ba992144..e7a53fbff6c 100644 --- a/skills/create-model-verification-card/scripts/validate_card.py +++ b/skills/create-model-verification-card/scripts/validate_card.py @@ -1086,8 +1086,26 @@ def _validate_inference( command_tokens = shlex.split(command) except ValueError: command_tokens = [] - if command_tokens[:2] != ["uv", "run"]: - errors.append(f"{_pointer(*resolved_command_path)}: inference must use uv run") + uses_uv = command_tokens[:2] == ["uv", "run"] + uses_inference_launcher = command_tokens[:1] == ["./scripts/inference/infer.sh"] + if not uses_uv and not uses_inference_launcher: + errors.append( + f"{_pointer(*resolved_command_path)}: inference must use ./scripts/inference/infer.sh " + "or a local uv run helper" + ) + if uses_inference_launcher: + task_values = _argument_values(command, "--task") + if len(task_values) != 1 or task_values[0] not in { + "text-generation", + "vlm-generation", + "model-comparison", + }: + errors.append(f"{_pointer(*resolved_command_path)}: infer.sh must specify one supported --task") + for resource_flag in ("--nodes", "--gpus-per-node"): + if len(_argument_values(command, resource_flag)) != 1: + errors.append( + f"{_pointer(*resolved_command_path)}: infer.sh must specify {resource_flag} exactly once" + ) prompts = _argument_values(command, "--prompt") if len(prompts) != 1: errors.append(f"{_pointer(*resolved_command_path)}: specify --prompt exactly once") @@ -1182,10 +1200,23 @@ def _validate_manual_forward_pass( if status != "verified" or not isinstance(expected, str): return - prefix = ["uv", "run", "python", "-m", "torch.distributed.run"] - if tokens[:5] != prefix: - errors.append(f"{_pointer(*path, 'command')}: manual forward pass must use uv distributed run") - if "examples/conversion/compare_hf_and_megatron/compare.py" not in tokens: + legacy_prefix = ["uv", "run", "python", "-m", "torch.distributed.run"] + uses_legacy_helper = tokens[:5] == legacy_prefix + uses_inference_launcher = tokens[:1] == ["./scripts/inference/infer.sh"] + if not uses_inference_launcher and not uses_legacy_helper: + errors.append( + f"{_pointer(*path, 'command')}: manual forward pass must use ./scripts/inference/infer.sh " + "or the legacy uv distributed helper" + ) + if uses_inference_launcher: + if _argument_values(command, "--task") != ["model-comparison"]: + errors.append( + f"{_pointer(*path, 'command')}: infer.sh manual forward pass must use --task model-comparison" + ) + for resource_flag in ("--nodes", "--gpus-per-node"): + if len(_argument_values(command, resource_flag)) != 1: + errors.append(f"{_pointer(*path, 'command')}: infer.sh must specify {resource_flag} exactly once") + elif "examples/conversion/compare_hf_and_megatron/compare.py" not in tokens: errors.append(f"{_pointer(*path, 'command')}: use the HF/Megatron comparison helper") for argument in ("--hf_model_path", "--megatron_model_path", "--prompt"): if len(_argument_values(command, argument)) != 1: diff --git a/tests/unit_tests/scripts/inference/test_setup_inference.py b/tests/unit_tests/scripts/inference/test_setup_inference.py index 0ebc32e8f37..fc7b669d89d 100644 --- a/tests/unit_tests/scripts/inference/test_setup_inference.py +++ b/tests/unit_tests/scripts/inference/test_setup_inference.py @@ -120,6 +120,29 @@ def test_parser_forwards_model_checkpoint_prompt_and_engine_args(): ] +@pytest.mark.parametrize( + ("task_name", "expected_path"), + [ + ("text-generation", "/opt/Megatron-Bridge/scripts/inference/text_generation.py"), + ("vlm-generation", "/opt/Megatron-Bridge/scripts/inference/vlm_generation.py"), + ( + "model-comparison", + "/opt/Megatron-Bridge/examples/conversion/compare_hf_and_megatron/compare.py", + ), + ], +) +def test_parser_selects_repository_inference_task(task_name, expected_path): + module = _load_setup_inference_module() + scripts = [] + module.run.Script = lambda **kwargs: scripts.append(types.SimpleNamespace(**kwargs)) or scripts[-1] + + args, inference_args = module.parse_args(["--task", task_name, "--prompt", "hello"]) + module._build_task(args.task, inference_args) + + assert scripts[0].path == expected_path + assert scripts[0].args == ["--prompt", "hello"] + + def test_parser_consumes_repeatable_srun_args(): module = _load_setup_inference_module() @@ -274,7 +297,7 @@ def test_build_task_quotes_prompts_and_uses_existing_entrypoint(): module.run.Script = lambda **kwargs: scripts.append(types.SimpleNamespace(**kwargs)) or scripts[-1] sentinel = "benign; echo should-not-run" - module._build_task(["--prompt", sentinel]) + module._build_task("text-generation", ["--prompt", sentinel]) assert scripts[0].path == "/opt/Megatron-Bridge/scripts/inference/text_generation.py" assert scripts[0].entrypoint == "python" diff --git a/tests/unit_tests/scripts/test_validate_model_verification_card.py b/tests/unit_tests/scripts/test_validate_model_verification_card.py index 1d7df997954..018574e2ddf 100644 --- a/tests/unit_tests/scripts/test_validate_model_verification_card.py +++ b/tests/unit_tests/scripts/test_validate_model_verification_card.py @@ -57,6 +57,74 @@ def _fsdp_metrics(): } +def test_verified_inference_accepts_canonical_bash_launcher(): + module = _load_validator() + errors = [] + item = { + "expected_result": 'The exact 1-token result produced completion "ok".', + } + + module._validate_inference( + item, + item_name="inference", + status="verified", + errors=errors, + command_override=( + "./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 " + "--task vlm-generation --prompt hello --max_new_tokens 1" + ), + ) + + assert errors == [] + + +def test_verified_inference_bash_launcher_requires_task_and_resources(): + module = _load_validator() + errors = [] + item = { + "expected_result": 'The exact 1-token result produced completion "ok".', + } + + module._validate_inference( + item, + item_name="inference", + status="verified", + errors=errors, + command_override="./scripts/inference/infer.sh --prompt hello --max_new_tokens 1", + ) + + assert "/items/inference/command: infer.sh must specify one supported --task" in errors + assert "/items/inference/command: infer.sh must specify --nodes exactly once" in errors + assert "/items/inference/command: infer.sh must specify --gpus-per-node exactly once" in errors + + +def test_verified_manual_forward_pass_accepts_canonical_bash_launcher(): + module = _load_validator() + errors = [] + item = { + "status": "verified", + "command": ( + "./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 --task model-comparison " + "--hf_model_path org/model --hf-revision revision " + "--megatron_model_path work/model/iter_0000000 --prompt hello" + ), + "last_verified": "2026-07-31", + "expected_result": ( + "The next token matches. Cosine similarity is 0.99. " + "Maximum absolute logit difference is 0.1. Mean absolute logit difference is 0.01." + ), + } + + module._validate_manual_forward_pass( + item, + status="verified", + model_revision="revision", + errors=errors, + ) + + assert errors == [] + + def test_fsdp_index_mirrors_concrete_hardware_leaf(): module = _load_validator() verification_index, items, hardware_groups = _complete_index_inputs(module) diff --git a/tests/unit_tests/utils/test_vlm_generate_utils.py b/tests/unit_tests/utils/test_vlm_generate_utils.py index 4a7bc93823c..af1a63f380d 100644 --- a/tests/unit_tests/utils/test_vlm_generate_utils.py +++ b/tests/unit_tests/utils/test_vlm_generate_utils.py @@ -13,7 +13,7 @@ # limitations under the License. """Unit tests for VLM generation helpers -in ``examples.conversion.vlm_generate_utils``. +in ``scripts/inference/vlm_generation_utils.py``. Covers the ``ImportError`` fallback when ``qwen_vl_utils`` is unavailable and the success paths with mocked processors and ``qwen_vl_utils`` helpers. @@ -22,20 +22,23 @@ import importlib.util import os import sys +import types from unittest import mock import pytest import torch -# Load examples/conversion/vlm_generate_utils.py directly from its file path. -# This avoids ambiguity when another `examples` package (e.g. from Megatron-LM) -# shadows the local one on sys.path. +# Load scripts/inference/vlm_generation_utils.py directly from its file path. _REPO_ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) -_VLM_GEN_UTILS_PATH = os.path.join(_REPO_ROOT, "examples", "conversion", "vlm_generate_utils.py") +_VLM_GEN_UTILS_PATH = os.path.join(_REPO_ROOT, "scripts", "inference", "vlm_generation_utils.py") _spec = importlib.util.spec_from_file_location("vlm_generate_utils_under_test", _VLM_GEN_UTILS_PATH) vlm_generate_utils = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(vlm_generate_utils) +_safe_url_module = types.ModuleType("megatron.bridge.utils.safe_url") +_safe_url_module.is_safe_public_http_url = mock.MagicMock() +_safe_url_module.safe_url_open = mock.MagicMock() +with mock.patch.dict(sys.modules, {"megatron.bridge.utils.safe_url": _safe_url_module}): + _spec.loader.exec_module(vlm_generate_utils) @pytest.mark.unit @@ -200,10 +203,62 @@ def test_unrecognized_processor_without_chat_template_fails_clearly(self): """Do not apply the Ministral raw-prompt convention to other VLMs.""" proc = mock.MagicMock() proc.chat_template = None + proc.tokenizer.chat_template = None with pytest.raises(ValueError, match="no model-specific raw-prompt fallback"): vlm_generate_utils.process_image_inputs(proc, "/image.png", "describe") + def test_minimax_uses_tokenizer_chat_template(self): + """MiniMax must format the multimodal prompt with its tokenizer template.""" + inputs = mock.MagicMock() + inputs.input_ids = torch.tensor([[10, 20]]) + inputs.get.side_effect = lambda key: { + "pixel_values": "PIXELS", + "image_grid_thw": "GRID", + }.get(key) + + proc = mock.MagicMock() + proc.chat_template = None + proc.tokenizer.chat_template = "TOKENIZER_TEMPLATE" + proc.tokenizer.apply_chat_template.return_value = "TEMPLATED" + proc.return_value = inputs + image = mock.MagicMock() + image.convert.return_value = "RGB_IMAGE" + + with mock.patch.object(vlm_generate_utils, "load_image", return_value=image) as load_image_mock: + result = vlm_generate_utils.process_image_inputs( + proc, + "/image.png", + "describe", + is_minimax=True, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": "describe"}, + ], + } + ] + proc.tokenizer.apply_chat_template.assert_called_once_with( + messages, + tokenize=False, + add_generation_prompt=True, + ) + load_image_mock.assert_called_once_with("/image.png") + image.convert.assert_called_once_with("RGB") + proc.assert_called_once_with( + text=["TEMPLATED"], + images=["RGB_IMAGE"], + padding=True, + return_tensors="pt", + ) + assert torch.equal(result[0], torch.tensor([[10, 20]])) + assert result[1] == "PIXELS" + assert result[2] == "GRID" + def test_kimi_image_path_returns_full_six_field_contract(self): """Kimi image preprocessing must match the VLM generation unpack contract.""" inputs = mock.MagicMock() From 44f56685b686f773cce720bc8a8a45f9e1a66810 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 01:46:24 +0200 Subject: [PATCH 15/85] cp: `fix(peft): sync grouped expert LoRA over expert DP (5231)` into `r0.6.0` (#5269) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- src/megatron/bridge/peft/utils.py | 12 +++++-- tests/unit_tests/peft/test_utils.py | 56 +++++++++++++++++++---------- 2 files changed, 46 insertions(+), 22 deletions(-) diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py index 259ba5eefbc..dbed296c82b 100644 --- a/src/megatron/bridge/peft/utils.py +++ b/src/megatron/bridge/peft/utils.py @@ -1792,8 +1792,9 @@ def __init__( self.pg_collection = _get_pg_collection( pg_collection, model_parallel_config, - required_pgs=["ep", "expt_tp", "expt_dp"], + required_pgs=["tp", "ep", "expt_tp", "expt_dp"], ) + tensor_parallel_group = _get_tensor_parallel_group(self.pg_collection) self.expert_tp_group = _get_tensor_parallel_group(self.pg_collection, is_expert=True) self.ep_group = _get_process_group(self.pg_collection, "ep") self.expert_dp_group = _get_process_group(self.pg_collection, "expt_dp") @@ -1804,6 +1805,10 @@ def __init__( self.expert_tp_group, model_parallel_config.expert_tensor_parallel_size or 1, ) + tensor_parallel_size = _process_group_size( + tensor_parallel_group, + getattr(model_parallel_config, "tensor_model_parallel_size", 1) or 1, + ) linear_in_tp_axis = 2 if input_is_parallel else 1 linear_out_tp_axis = 1 @@ -1841,12 +1846,13 @@ def __init__( ParallelLinearAdapter._get_init_fn(self, column_init_method)(linear_in_weight) ParallelLinearAdapter._get_init_fn(self, row_init_method)(linear_out_weight) - expert_parallel = ( + use_expert_process_groups = ( _process_group_size( self.ep_group, model_parallel_config.expert_model_parallel_size or 1, ) > 1 + or expert_tp_size != tensor_parallel_size ) self._linear_in_tp_axis = linear_in_tp_axis self._linear_out_tp_axis = linear_out_tp_axis @@ -1856,7 +1862,7 @@ def __init__( (self.linear_in.weight, linear_in_tp_axis), (self.linear_out.weight, linear_out_tp_axis), ): - setattr(weight, "allreduce", not expert_parallel) + setattr(weight, "allreduce", not use_expert_process_groups) if tp_axis is not None: set_tensor_model_parallel_attributes(weight, True, tp_axis, 1) diff --git a/tests/unit_tests/peft/test_utils.py b/tests/unit_tests/peft/test_utils.py index 1b4dbfff8f2..c2be2839065 100644 --- a/tests/unit_tests/peft/test_utils.py +++ b/tests/unit_tests/peft/test_utils.py @@ -2226,11 +2226,32 @@ def test_grouped_expert_linear_fc1_factory_merge_restores_gate_up_order(self): ) torch.testing.assert_close(merged, expected) - @pytest.mark.parametrize(("ep_size", "expected_allreduce"), [(1, True), (2, False)]) - def test_grouped_expert_linear_adapter_allreduce_flag_tracks_expert_parallelism(self, ep_size, expected_allreduce): - """Per-expert grouped adapters should use expert-DP grad sync only when EP is enabled.""" + @pytest.mark.parametrize( + ("ep_size", "tp_size", "etp_size", "expected_allreduce"), + [ + pytest.param(1, 1, 1, True, id="ordinary-dp"), + pytest.param(2, 1, 1, False, id="expert-dp-with-ep"), + pytest.param(1, 2, 1, False, id="expert-dp-with-tp-greater-than-etp"), + pytest.param(1, 1, 2, False, id="expert-dp-with-etp-greater-than-tp"), + ], + ) + def test_grouped_expert_linear_adapter_allreduce_flag_tracks_expert_topology( + self, + ep_size, + tp_size, + etp_size, + expected_allreduce, + ): + """Per-expert grouped adapters should select DP from the full expert topology.""" config = MockModelParallelConfig() - config._pg_collection = make_mock_pg_collection(ep_size=ep_size, etp_size=1) + config.tensor_model_parallel_size = tp_size + config.expert_model_parallel_size = ep_size + config.expert_tensor_parallel_size = etp_size + config._pg_collection = make_mock_pg_collection( + tp_size=tp_size, + ep_size=ep_size, + etp_size=etp_size, + ) adapter = GroupedExpertLinearAdapter( in_features=2, out_features=2, @@ -2249,17 +2270,20 @@ def test_grouped_expert_linear_adapter_allreduce_flag_tracks_expert_parallelism( assert adapter.linear_in.weight.partition_dim == 1 assert adapter.linear_out.weight.partition_dim == 1 - def test_grouped_expert_linear_adapter_groups_as_expert_ddp_buffer_when_ep_enabled(self): - """Per-expert adapter params must sync on expert-DP, not dense DP. + def test_grouped_expert_linear_adapter_groups_as_expert_ddp_buffer_when_etp_differs(self): + """Per-expert adapter params must sync on expert-DP when ETP differs from TP. - EP plus DP replicates each local expert across expert-DP ranks. Marking - these params as expert-parallel keeps replicas for the same expert in - sync without mixing different EP-owned experts. + Even with EP=1, differing TP and ETP topologies require the expert-DP + buffer so replicas synchronize over the process group that owns them. """ from megatron.core.distributed.param_and_grad_buffer import group_params_for_buffers + from megatron.core.optimizer.param_layout import BufferKey config = MockModelParallelConfig() - config._pg_collection = make_mock_pg_collection(ep_size=8, etp_size=1) + config.tensor_model_parallel_size = 2 + config.expert_model_parallel_size = 1 + config.expert_tensor_parallel_size = 1 + config._pg_collection = make_mock_pg_collection(tp_size=2, ep_size=1, etp_size=1) adapter = GroupedExpertLinearAdapter( in_features=2, out_features=2, @@ -2271,18 +2295,12 @@ def test_grouped_expert_linear_adapter_groups_as_expert_ddp_buffer_when_ep_enabl model_parallel_config=config, ) - buffer_groups = group_params_for_buffers( - [adapter.linear_in.weight, adapter.linear_out.weight], - grad_reduce_in_fp32=False, - ) + buffer_groups = group_params_for_buffers([adapter.linear_in.weight], grad_reduce_in_fp32=False) assert len(buffer_groups) == 1 buffer_key, (params, _param_indices) = next(iter(buffer_groups.items())) - assert buffer_key.is_expert_parallel - assert [id(param) for param in params] == [ - id(adapter.linear_in.weight), - id(adapter.linear_out.weight), - ] + assert buffer_key == BufferKey(torch.float32, torch.float32, True) + assert params == [adapter.linear_in.weight] def test_grouped_expert_linear_sharded_state_dict_uses_expert_parallel_offsets(self): """Grouped-expert weights should shard only across expert EP/ETP and use expert-DP replica ids.""" From af2d5b4ed9fdae545b9d2f4c9e3f8e305a39a670 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 02:28:42 +0200 Subject: [PATCH 16/85] cp: `chore(deps): update mlflow version to 3.15.0 (5267)` into `r0.6.0` (#5275) Signed-off-by: Ajay Balasa Signed-off-by: NeMo Bot Co-authored-by: Ajay --- pyproject.toml | 4 ++-- uv.lock | 22 +++++++++++----------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index af10bc752ec..2c3894f59c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,7 +95,7 @@ dependencies = [ "flash-linear-attention", "timm", "open-clip-torch>=3.2.0", - "mlflow>=3.9.0", # To address CVE-2025-15031 + "mlflow>=3.15.0", # To address CVE-2025-15031 "comet-ml>=3.50.0", "torch>=2.6.0", "flashinfer-python==0.6.8.post1", # keep flashinfer-python and flashinfer-cubin in lockstep: @@ -130,7 +130,7 @@ override-dependencies = [ "torchvision; sys_platform == 'never'", "triton; sys_platform == 'never'", "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@e7c550c5f80636cf841a8204b1d6f85a5f3f28b7", - "mlflow>=3.9.0", # To address CVE-2025-15031 + "mlflow>=3.15.0", # To address CVE-2025-15031 "cachetools>=5.0.0", "cryptography>=43.0.0,<47", "nvidia-modelopt==0.44.0rc5", diff --git a/uv.lock b/uv.lock index 5ca47561861..1ad0214cd3f 100644 --- a/uv.lock +++ b/uv.lock @@ -16,7 +16,7 @@ overrides = [ { name = "imageio-ffmpeg", marker = "sys_platform == 'never'" }, { name = "langchain", specifier = ">=0.3.28" }, { name = "langchain-core", specifier = ">=0.3.80" }, - { name = "mlflow", specifier = ">=3.9.0" }, + { name = "mlflow", specifier = ">=3.15.0" }, { name = "nvidia-cudnn-frontend", specifier = "==1.26.0" }, { name = "nvidia-cutlass-dsl", extras = ["cu13"], marker = "sys_platform == 'never'" }, { name = "nvidia-modelopt", specifier = "==0.44.0rc5", index = "https://pypi.nvidia.com/" }, @@ -2198,7 +2198,7 @@ requires-dist = [ { name = "mamba-ssm", marker = "extra == 'ssm'" }, { name = "megatron-core", extras = ["dev", "mlm"], editable = "3rdparty/Megatron-LM" }, { name = "mistral-common", specifier = ">=1.10.0" }, - { name = "mlflow", specifier = ">=3.9.0" }, + { name = "mlflow", specifier = ">=3.15.0" }, { name = "nemo-run", marker = "extra == 'recipes'" }, { name = "nvdlfw-inspect", marker = "extra == 'tensor-inspect'", specifier = "==0.2.1" }, { name = "nvidia-resiliency-ext" }, @@ -2486,7 +2486,7 @@ wheels = [ [[package]] name = "mlflow" -version = "3.14.0" +version = "3.15.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -2510,14 +2510,14 @@ dependencies = [ { name = "sqlalchemy" }, { name = "waitress", marker = "sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d6/0b/3404a057daceffe9ce18cd08868648a1e9b817270177bdf8a764576b988b/mlflow-3.14.0.tar.gz", hash = "sha256:5a1f818fa003035c724162096ce3ded7bc7bc47a1cae595df6173961983f4718", size = 11792369, upload-time = "2026-06-17T07:57:44.712Z" } +sdist = { url = "https://files.pythonhosted.org/packages/af/9f/249b80849aae0919aee3f590003419c394f0efe92ae98546ea626be86e1e/mlflow-3.15.0.tar.gz", hash = "sha256:c53d7933391fefda2ac0f6b942d1744a02b5b65da74a93cc125247f4c4da57c9", size = 10395873, upload-time = "2026-07-31T07:06:00.211Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/de/b9/76dcdef7f7f856b36f18cfcd752c2717d9847812a0aaa36d50a7baed569d/mlflow-3.14.0-py3-none-any.whl", hash = "sha256:dbf77f7cdb5b5c0ec59b4671c61730b1b914b4dff7a2892e267a547cb5454f56", size = 12564161, upload-time = "2026-06-17T07:57:42.348Z" }, + { url = "https://files.pythonhosted.org/packages/0e/7a/3da40be70e7ca0c4d8c2f98edfcb275ef2898d8f4da565f55a9fff95fc78/mlflow-3.15.0-py3-none-any.whl", hash = "sha256:3ae54c7f91a6b98ae9360aaeda9b31eb7930571c2690491a7b550c84f795a709", size = 11188317, upload-time = "2026-07-31T07:05:57.492Z" }, ] [[package]] name = "mlflow-skinny" -version = "3.14.0" +version = "3.15.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, @@ -2541,14 +2541,14 @@ dependencies = [ { name = "typing-extensions" }, { name = "uvicorn" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e8/4f/a054cd8860590e4e942aee1aab3c94307878159f945fa844acc9ea787721/mlflow_skinny-3.14.0.tar.gz", hash = "sha256:e50f4506422c7737157ae6643c165122af7898345f2e828fa93c4f10128653cf", size = 2901772, upload-time = "2026-06-17T07:57:44.252Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/a1/04489b8f10ec642cc5d16d8bcf0e3626cbe6b8924940b510de9e3cb37ac7/mlflow_skinny-3.15.0.tar.gz", hash = "sha256:6faad0bb5b0b8adc08550c6cd2ab5a87f9e7d7fe8814278e22b79f9cac99d943", size = 3034359, upload-time = "2026-07-31T07:04:58.359Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/e7/b80f76ce689b9d6f21cdb84abb2b02148a1149e63a6428dd2c629cefd061/mlflow_skinny-3.14.0-py3-none-any.whl", hash = "sha256:a4880e086365871ef9d78e727a34ea5fb1ce615689579998d48e8c65ee1665a9", size = 3462788, upload-time = "2026-06-17T07:57:42.583Z" }, + { url = "https://files.pythonhosted.org/packages/28/ca/b0194962614876abb8b9a4d38815b4c6bf5c6f626b5dd3bcd321048755fc/mlflow_skinny-3.15.0-py3-none-any.whl", hash = "sha256:fc609ea5f7325a76d6f905489d8be73e85bcba0eef64144a763cf446de83192e", size = 3612199, upload-time = "2026-07-31T07:04:56.564Z" }, ] [[package]] name = "mlflow-tracing" -version = "3.14.0" +version = "3.15.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, @@ -2560,9 +2560,9 @@ dependencies = [ { name = "protobuf" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/04/34/ff5e72919b4eec8fe65e6fc843978a1a512194e6fafbd1761deca48269ad/mlflow_tracing-3.14.0.tar.gz", hash = "sha256:c2f701e001d35964f23fbbdfdda36c818a76c157b912ae83781199fd714be09a", size = 1429017, upload-time = "2026-06-17T07:58:00.647Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/47/8a3c4f56199aadce644256684060df09f0af6ecaec460860657f2cb13b10/mlflow_tracing-3.15.0.tar.gz", hash = "sha256:e4e73243eed04d0771954b6fa17014dd21dabc67af17d4937fac2b62ec476e95", size = 1500398, upload-time = "2026-07-31T07:05:54.445Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5f/4a/4658a9e514c8f079e40b608661844b9beb21c7530e6ee1e7f830cf81541e/mlflow_tracing-3.14.0-py3-none-any.whl", hash = "sha256:854488dd18068f15e2a56f1cc7b8868c611d09ea39068d0a691a3f07e0048cae", size = 1703863, upload-time = "2026-06-17T07:57:58.687Z" }, + { url = "https://files.pythonhosted.org/packages/76/3c/5839690c3fe68f579074080832cba413f988249ddd350dcf7dfc3be1b874/mlflow_tracing-3.15.0-py3-none-any.whl", hash = "sha256:3762f654a858474c7988c0a2a5246765b173fa2b6a2bb8f035c115db1a8e3a6a", size = 1784776, upload-time = "2026-07-31T07:05:52.937Z" }, ] [[package]] From a2fec25de20b9040e9e5502ffdab6c42cb1bf7db Mon Sep 17 00:00:00 2001 From: weijiac0619 Date: Mon, 3 Aug 2026 21:05:48 -0700 Subject: [PATCH 17/85] chore: bump mcore dev commit for DSv4 support (#5284) Signed-off-by: Weijia Chen --- .dev.commit | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.dev.commit b/.dev.commit index 0b03684ed74..3df9b5e7a84 100644 --- a/.dev.commit +++ b/.dev.commit @@ -1 +1 @@ -0e6ac576fba8995fca541ceb6af5f35eb0e4f16e +43e45e13074f36da68b34c7a189cd997a01d1b41 From d4bf3f1b5565bed5790a268d0b970d9340c61502 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 3 Aug 2026 21:20:49 -0700 Subject: [PATCH 18/85] cp: `fix(recipes): derive pretraining vocab from tokenizer (4996)` into `r0.6.0` (#5280) Signed-off-by: Chen Cui --- docs/nemo2-migration-guide.md | 47 ++++++++++++++---- scripts/performance/run_recipe.py | 17 +++++-- .../bridge/models/gemma/gemma3_provider.py | 2 +- .../bridge/models/gemma/gemma4_provider.py | 2 +- .../qwen_vl/modelling_qwen3_vl/model.py | 15 +++++- .../modelling_qwen3_vl/transformer_config.py | 2 + src/megatron/bridge/perf_recipes/_common.py | 9 +++- .../bridge/perf_recipes/qwen_vl/common.py | 14 ++++++ src/megatron/bridge/recipes/common.py | 1 + .../bridge/recipes/qwen_vl/h100/qwen35_vl.py | 16 ++----- .../bridge/recipes/qwen_vl/h100/qwen3_vl.py | 10 ++-- src/megatron/bridge/training/setup.py | 14 ++++-- .../bridge/training/tokenizers/config.py | 11 +++++ .../ckpts/qwen3_4b/test_qwen3_4b_ckpt.py | 8 +++- .../models/gemma/test_gemma3_provider.py | 5 +- .../models/gemma/test_gemma4_provider.py | 23 +++++++++ .../qwen_vl/modelling_qwen3_vl/test_model.py | 48 +++++++++++++++++++ .../recipes/qwen_vl/test_qwen35_vl_recipes.py | 4 +- .../recipes/qwen_vl/test_qwen3_vl_recipes.py | 4 +- .../recipes/test_all_perf_recipe_factories.py | 8 +++- .../recipes/test_all_recipe_factories.py | 6 +++ .../unit_tests/recipes/test_gemma4_recipe.py | 16 ++++--- .../recipes/test_perf_recipe_environment.py | 4 +- .../performance/test_recipe_environment.py | 38 +++++++++++++++ .../training/test_modelopt_setup.py | 4 +- tests/unit_tests/training/test_setup.py | 22 +++++++++ 26 files changed, 297 insertions(+), 53 deletions(-) diff --git a/docs/nemo2-migration-guide.md b/docs/nemo2-migration-guide.md index 76464e3f57c..dfd0f9ff49d 100644 --- a/docs/nemo2-migration-guide.md +++ b/docs/nemo2-migration-guide.md @@ -693,19 +693,25 @@ tokenizer_config = TokenizerConfig( #### Vocab Size Priority -In Megatron Bridge, vocabulary size can be specified in either the model provider or derived from the tokenizer. The priority order is: +In Megatron Bridge, vocabulary size can be specified in the model provider or derived from the runtime tokenizer. The priority order is: -1. **Model provider `vocab_size` is set**: Uses the model's vocab size - - Must be `>= tokenizer.vocab_size` (raises error if smaller) - - Sets `should_pad_vocab=False` (no automatic padding) - - Useful when you need a specific vocab size (e.g., for checkpoint compatibility) +1. **`TokenizerConfig.use_tokenizer_vocab_size=True`**: Uses the tokenizer's vocab size + - Overrides a preset model-provider `vocab_size`. + - Sets `should_pad_vocab=True` (enables padding for efficient parallelism). + - Intended for from-scratch pretraining, where the dataset tokenizer defines the vocabulary. + - This policy remains active during checkpoint loading; disable it when checkpoint compatibility requires the explicit model vocabulary. -2. **Model provider `vocab_size` is None**: Uses tokenizer's vocab size +2. **Model provider `vocab_size` is set**: Uses the model's vocab size + - Must be `>= tokenizer.vocab_size` (raises an error if smaller). + - Sets `should_pad_vocab=False` (no automatic padding). + - Useful when a specific vocabulary is required for model or checkpoint compatibility. + +3. **Model provider `vocab_size` is `None`**: Uses the tokenizer's vocab size - Automatically derived from `tokenizer.vocab_size` after building the tokenizer. - - Sets `should_pad_vocab=True` (enables padding for efficient parallelism) + - Sets `should_pad_vocab=True`. ```python -# Option 1: Let tokenizer determine vocab size +# Option 1: Let tokenizer determine vocab size when the model has no preset config = ConfigContainer( model=GPTModelProvider( # vocab_size not set - will use tokenizer's vocab size @@ -717,7 +723,19 @@ config = ConfigContainer( ), ) -# Option 2: Explicitly set vocab size in model +# Option 2: Override a preset model vocab for from-scratch pretraining +config = ConfigContainer( + model=GPTModelProvider( + vocab_size=128256, # Ignored whenever the flag is enabled + ), + tokenizer=TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model="my-org/my-pretraining-tokenizer", + use_tokenizer_vocab_size=True, + ), +) + +# Option 3: Explicitly set vocab size in model config = ConfigContainer( model=GPTModelProvider( vocab_size=128256, # Explicitly set (must be >= tokenizer vocab size) @@ -726,6 +744,17 @@ config = ConfigContainer( ) ``` +Pretraining recipes enable `use_tokenizer_vocab_size` by default. For a new run, use an empty checkpoint directory so the runtime tokenizer defines the model vocabulary. A checkpoint created by that policy can be resumed with the same tokenizer and recipe configuration. + +Checkpoints created before a recipe enabled `use_tokenizer_vocab_size` may have used the model provider's larger explicit vocabulary. To preserve their embedding and output tensor shapes, disable the new policy and retain the vocabulary used to create the checkpoint: + +```python +config.tokenizer.use_tokenizer_vocab_size = False +config.model.vocab_size = 128256 # The vocabulary used to create the checkpoint +``` + +Do not change this setting partway through a run. Switching vocabulary policies changes model tensor shapes and is not a checkpoint migration mechanism. + ### Parallelism Configuration Migration In NeMo 2.0, parallelism settings were configured on `MegatronStrategy`. In Megatron Bridge, these are set directly on the model provider: diff --git a/scripts/performance/run_recipe.py b/scripts/performance/run_recipe.py index 53344b071aa..1954962376a 100644 --- a/scripts/performance/run_recipe.py +++ b/scripts/performance/run_recipe.py @@ -129,18 +129,29 @@ def _apply_training_argparse_overrides(config, args): # Tokenizer configuration from megatron.bridge.training.config import TokenizerConfig + use_tokenizer_vocab_size = config.tokenizer.use_tokenizer_vocab_size if args.tokenizer_type == "NullTokenizer": - config.tokenizer = TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=args.vocab_size) + config.tokenizer = TokenizerConfig( + tokenizer_type="NullTokenizer", + vocab_size=args.vocab_size, + use_tokenizer_vocab_size=use_tokenizer_vocab_size, + ) elif args.tokenizer_type == "HuggingFaceTokenizer": if not args.tokenizer_model: raise ValueError("--tokenizer-model is required when using HuggingFaceTokenizer") tokenizer_model = args.tokenizer_model - config.tokenizer = TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model=tokenizer_model) + config.tokenizer = TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=tokenizer_model, + use_tokenizer_vocab_size=use_tokenizer_vocab_size, + ) elif args.tokenizer_type == "SentencePieceTokenizer": if not args.tokenizer_model: raise ValueError("--tokenizer-model is required for SentencePieceTokenizer") config.tokenizer = TokenizerConfig( - tokenizer_type="SentencePieceTokenizer", tokenizer_model=args.tokenizer_model + tokenizer_type="SentencePieceTokenizer", + tokenizer_model=args.tokenizer_model, + use_tokenizer_vocab_size=use_tokenizer_vocab_size, ) else: # Diffusion recipes (FLUX, WAN) keep their own dataset object (Wan/FluxDatasetConfig). diff --git a/src/megatron/bridge/models/gemma/gemma3_provider.py b/src/megatron/bridge/models/gemma/gemma3_provider.py index d8f3eb24327..ac59576ce27 100644 --- a/src/megatron/bridge/models/gemma/gemma3_provider.py +++ b/src/megatron/bridge/models/gemma/gemma3_provider.py @@ -119,7 +119,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> "MCoreG if hasattr(model, "embedding"): model.embedding = Gemma3LanguageModelEmbedding( config=self, - vocab_size=self.vocab_size, + vocab_size=model.vocab_size, max_sequence_length=self.seq_length, position_embedding_type=self.position_embedding_type, scatter_to_sequence_parallel=self.scatter_embedding_sequence_parallel, diff --git a/src/megatron/bridge/models/gemma/gemma4_provider.py b/src/megatron/bridge/models/gemma/gemma4_provider.py index 85759f30dd4..7b7b28fe0e2 100644 --- a/src/megatron/bridge/models/gemma/gemma4_provider.py +++ b/src/megatron/bridge/models/gemma/gemma4_provider.py @@ -345,7 +345,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> "MCoreG if hasattr(model, "embedding"): model.embedding = Gemma3LanguageModelEmbedding( config=self, - vocab_size=self.vocab_size, + vocab_size=model.vocab_size, max_sequence_length=self.seq_length, position_embedding_type=self.position_embedding_type, scatter_to_sequence_parallel=self.scatter_embedding_sequence_parallel, diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py index 7d1e1b651d5..7ac9ea1cb59 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py @@ -55,6 +55,7 @@ get_packed_seq_cp_partition_indices, get_packed_seq_q_cu_seqlens, ) +from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size def _is_mrope_position_ids(position_ids: torch.Tensor | None) -> bool: @@ -324,10 +325,22 @@ def __init__( pg_collection=pg_collection, ) if self.add_decoder: + assert language_transformer_config.vocab_size is not None, ( + "vocab_size must be configured before constructing the Qwen3-VL language model" + ) + if language_transformer_config.should_pad_vocab: + language_model_vocab_size = calculate_padded_vocab_size( + language_transformer_config.vocab_size, + language_transformer_config.make_vocab_size_divisible_by, + language_transformer_config.tensor_model_parallel_size, + ) + else: + language_model_vocab_size = language_transformer_config.vocab_size + self.language_model = Qwen3VLGPTModel( config=language_transformer_config, transformer_layer_spec=language_transformer_layer_spec, - vocab_size=language_transformer_config.vocab_size, + vocab_size=language_model_vocab_size, max_sequence_length=language_transformer_config.language_max_sequence_length, parallel_output=parallel_output, position_embedding_type="mrope", diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py index 41cbc08fa6e..c8f2bec15dd 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py @@ -28,6 +28,8 @@ class Qwen3VLTransformerConfig(TransformerConfig): """Configuration for Qwen3-VL transformer with vision and language components.""" vocab_size: int = 64000 + make_vocab_size_divisible_by: int = 128 + should_pad_vocab: bool = False language_max_sequence_length: int = 4096 patch_size: int = 16 diff --git a/src/megatron/bridge/perf_recipes/_common.py b/src/megatron/bridge/perf_recipes/_common.py index f21d54b18dc..fb512465358 100644 --- a/src/megatron/bridge/perf_recipes/_common.py +++ b/src/megatron/bridge/perf_recipes/_common.py @@ -33,8 +33,9 @@ def _benchmark_common(cfg: ConfigContainer, cross_entropy_impl: str = "te") -> N Intended for performance benchmark recipes only. Sets short training runs, disables checkpointing/eval, tunes scheduler, and enables perf-oriented kernels. - Must stay in sync with ``_set_common_perf_overrides`` in - ``scripts/performance/utils/overrides.py``. + This is the fixed-model-shape policy for flat performance recipes. + Canonical recipes launched with ``scripts/performance --use_recipes`` + retain their own tokenizer vocabulary policy. Individual recipes may override any of these after calling this function (e.g. Kimi K2 sets ``grad_reduce_in_fp32 = True``). @@ -44,6 +45,10 @@ def _benchmark_common(cfg: ConfigContainer, cross_entropy_impl: str = "te") -> N cfg.train.manual_gc = True cfg.train.manual_gc_interval = 100 + # Performance recipes benchmark a fixed model shape. Synthetic or runtime + # tokenizers must not resize the embedding and output layers during setup. + cfg.tokenizer.use_tokenizer_vocab_size = False + cfg.checkpoint.save = None cfg.logger.log_interval = 1 diff --git a/src/megatron/bridge/perf_recipes/qwen_vl/common.py b/src/megatron/bridge/perf_recipes/qwen_vl/common.py index 4948781827e..17f6d975761 100644 --- a/src/megatron/bridge/perf_recipes/qwen_vl/common.py +++ b/src/megatron/bridge/perf_recipes/qwen_vl/common.py @@ -32,11 +32,24 @@ from megatron.bridge.training.config import ConfigContainer +def _use_model_vocab_null_tokenizer(cfg: ConfigContainer) -> None: + """Use a model-sized synthetic tokenizer for Qwen-VL performance runs.""" + if cfg.model.vocab_size is None: + raise ValueError("Qwen-VL performance recipes require a model vocabulary size.") + cfg.tokenizer.tokenizer_type = "NullTokenizer" + cfg.tokenizer.tokenizer_model = None + cfg.tokenizer.vocab_size = cfg.model.vocab_size + # The synthetic tokenizer mirrors the fixed benchmark model shape; it does + # not define a new tokenizer-derived vocabulary for from-scratch training. + cfg.tokenizer.use_tokenizer_vocab_size = False + + def _qwen35_vl_common(cfg: ConfigContainer) -> None: """Apply VLM-specific performance benchmark settings for Qwen3.5-VL. Must be called before ``_benchmark_common`` and after setting precision. """ + _use_model_vocab_null_tokenizer(cfg) cfg.model.bias_activation_fusion = True cfg.model.recompute_granularity = None cfg.model.recompute_method = None @@ -82,6 +95,7 @@ def _qwen35_vl_post_clear_scope_with_overlap(cfg: ConfigContainer) -> None: def _finalize_qwen3_vl(cfg: ConfigContainer) -> None: """Apply Qwen3-VL perf defaults that must override generic benchmark defaults.""" + _use_model_vocab_null_tokenizer(cfg) # _benchmark_common sets apply_rope_fusion=True; Qwen3-VL asserts it must be False # (per-token absolute positional frequencies are incompatible with TE's fused RoPE). cfg.model.apply_rope_fusion = False diff --git a/src/megatron/bridge/recipes/common.py b/src/megatron/bridge/recipes/common.py index 809a4e765f3..6c620ad6319 100644 --- a/src/megatron/bridge/recipes/common.py +++ b/src/megatron/bridge/recipes/common.py @@ -111,6 +111,7 @@ def _pretrain_common() -> ConfigContainer: tokenizer=TokenizerConfig( tokenizer_type="HuggingFaceTokenizer", tokenizer_model=None, # Must be set by each recipe + use_tokenizer_vocab_size=True, ), # Checkpoint config checkpoint=CheckpointConfig( diff --git a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py index 144d0184f30..24c6ac1a580 100644 --- a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py +++ b/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py @@ -30,7 +30,6 @@ from megatron.bridge.recipes.utils.dataset_utils import default_peft_config from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing -from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE from megatron.bridge.training.config import ConfigContainer @@ -72,8 +71,7 @@ def qwen35_vl_9b_pretrain_4gpu_h100_bf16_mock_config() -> ConfigContainer: persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False @@ -121,8 +119,7 @@ def qwen35_vl_27b_pretrain_16gpu_h100_bf16_mock_config() -> ConfigContainer: persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False @@ -171,8 +168,7 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_h100_bf16_mock_config() -> ConfigContainer: persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False @@ -222,8 +218,7 @@ def qwen35_vl_122b_a10b_pretrain_128gpu_h100_bf16_mock_config() -> ConfigContain persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False @@ -274,8 +269,7 @@ def qwen35_vl_397b_a17b_pretrain_512gpu_h100_bf16_mock_config() -> ConfigContain persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False diff --git a/src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py index dc68dabad8e..74a8ab150f9 100644 --- a/src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py +++ b/src/megatron/bridge/recipes/qwen_vl/h100/qwen3_vl.py @@ -32,7 +32,6 @@ from megatron.bridge.recipes.utils.dataset_utils import default_peft_config from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing -from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend @@ -77,8 +76,7 @@ def qwen3_vl_8b_pretrain_4gpu_h100_bf16_mock_config() -> ConfigContainer: persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False @@ -127,8 +125,7 @@ def qwen3_vl_30b_a3b_pretrain_8gpu_h100_bf16_mock_config() -> ConfigContainer: persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False @@ -178,8 +175,7 @@ def qwen3_vl_235b_a22b_pretrain_256gpu_h100_bf16_mock_config() -> ConfigContaine persistent_workers=False, pad_to_max_length=True, ) - cfg.tokenizer.tokenizer_type = "NullTokenizer" - cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + cfg.tokenizer.tokenizer_model = hf_path cfg.train.eval_interval = 500 cfg.train.eval_iters = 32 cfg.ddp.overlap_grad_reduce = False diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index ebe5ed86d1e..30efc5da646 100644 --- a/src/megatron/bridge/training/setup.py +++ b/src/megatron/bridge/training/setup.py @@ -248,6 +248,7 @@ def setup( cfg.model.vocab_size, cfg.model.should_pad_vocab = _validate_and_set_vocab_size( model_vocab_size=cfg.model.vocab_size, tokenizer_vocab_size=tokenizer.vocab_size, + use_tokenizer_vocab_size=getattr(cfg.tokenizer, "use_tokenizer_vocab_size", False), ) if hasattr(cfg.dataset, "tokenizer"): @@ -630,12 +631,18 @@ def _apply_peft_transformation(peft, base_model: list[MegatronModule]) -> list[M return transformed_model -def _validate_and_set_vocab_size(model_vocab_size: Optional[int], tokenizer_vocab_size: int) -> tuple[int, bool]: +def _validate_and_set_vocab_size( + model_vocab_size: Optional[int], + tokenizer_vocab_size: int, + use_tokenizer_vocab_size: bool = False, +) -> tuple[int, bool]: """Validate and determine the correct vocab size for the model. Args: model_vocab_size: Vocab size set in model config (can be None) tokenizer_vocab_size: Unpadded tokenizer vocab size + use_tokenizer_vocab_size: Ignore a preset model vocabulary and derive it + from the tokenizer. Intended for from-scratch pretraining recipes. Returns: tuple[int, bool]: The validated unpadded vocab size and padding flag @@ -645,8 +652,9 @@ def _validate_and_set_vocab_size(model_vocab_size: Optional[int], tokenizer_voca Raises: ValueError: If model vocab size is invalid """ - if model_vocab_size is None: - # If model vocab size is not set, use the tokenizer's vocab size + if use_tokenizer_vocab_size or model_vocab_size is None: + # Use the tokenizer's vocab size when the model vocab is unset, or when + # use_tokenizer_vocab_size forces it for from-scratch pretraining. # Enable padding since this came from tokenizer return tokenizer_vocab_size, True elif model_vocab_size < tokenizer_vocab_size: diff --git a/src/megatron/bridge/training/tokenizers/config.py b/src/megatron/bridge/training/tokenizers/config.py index 03fab40c7af..22508f0a851 100644 --- a/src/megatron/bridge/training/tokenizers/config.py +++ b/src/megatron/bridge/training/tokenizers/config.py @@ -33,6 +33,17 @@ class TokenizerConfig(MTrainTokenizerConfig): rank: int = 0 """Distributed rank used by MCore tokenizer helper logging.""" + use_tokenizer_vocab_size: bool = False + """Use the runtime tokenizer vocabulary size for the model. + + Enable this for from-scratch pretraining, where the tokenizer selected for + the dataset defines the embedding and output vocabulary. Keep it disabled + when model or checkpoint compatibility requires an explicitly configured + model vocabulary size. This policy also applies during checkpoint loading; + disable it and configure the checkpoint's original model vocabulary when + resuming a run created with a different vocabulary policy. + """ + hf_tokenizer_kwargs: dict[str, Any] | None = field(default_factory=dict) """Additional keyword arguments to pass to HuggingFace AutoTokenizer.from_pretrained. diff --git a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py index 2e21c8507fe..24c1ae77686 100644 --- a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py +++ b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py @@ -31,6 +31,7 @@ MBRIDGE_CKPT = f"{BASE_DIR}/mbridge" MCORE_CKPT = f"{BASE_DIR}/mcore" TB_DIR = f"{BASE_DIR}/tb" +QWEN3_4B_VOCAB_SIZE = 151936 class TestQwen3Ckpt: @@ -42,6 +43,11 @@ def test_qwen3_4b_ckpt_mbridge(self): config = qwen3_4b_pretrain_config() + # Keep both checkpoint paths independent of the shared Hugging Face cache. + config.tokenizer.tokenizer_type = "NullTokenizer" + config.tokenizer.tokenizer_model = None + config.tokenizer.vocab_size = QWEN3_4B_VOCAB_SIZE + config.checkpoint.save = MBRIDGE_CKPT config.checkpoint.load = MCORE_CKPT if os.path.exists(MCORE_CKPT) else None config.checkpoint.load_optim = False @@ -127,7 +133,7 @@ def test_qwen3_4b_ckpt_mcore(self, monkeypatch): "--tokenizer-type", "NullTokenizer", "--vocab-size", - "151936", + str(QWEN3_4B_VOCAB_SIZE), "--train-iters", f"{train_iters}", "--save-interval", diff --git a/tests/unit_tests/models/gemma/test_gemma3_provider.py b/tests/unit_tests/models/gemma/test_gemma3_provider.py index cd727f4276c..845d6bbcaf8 100644 --- a/tests/unit_tests/models/gemma/test_gemma3_provider.py +++ b/tests/unit_tests/models/gemma/test_gemma3_provider.py @@ -89,6 +89,7 @@ def test_gemma3_provider_provide_method(self, mock_rotary_embedding, mock_langua # Mock the parent provide method mock_model = Mock() mock_model.embedding = Mock() + mock_model.vocab_size = 262272 mock_model.setup_embeddings_and_output_layer = Mock() provider = Gemma3ModelProvider( @@ -96,7 +97,7 @@ def test_gemma3_provider_provide_method(self, mock_rotary_embedding, mock_langua hidden_size=1152, num_attention_heads=4, kv_channels=256, - vocab_size=262144, + vocab_size=262145, seq_length=32768, ) @@ -109,7 +110,7 @@ def test_gemma3_provider_provide_method(self, mock_rotary_embedding, mock_langua # Verify that custom embedding was created mock_language_embedding.assert_called_once_with( config=provider, - vocab_size=provider.vocab_size, + vocab_size=mock_model.vocab_size, max_sequence_length=provider.seq_length, position_embedding_type=provider.position_embedding_type, scatter_to_sequence_parallel=provider.scatter_embedding_sequence_parallel, diff --git a/tests/unit_tests/models/gemma/test_gemma4_provider.py b/tests/unit_tests/models/gemma/test_gemma4_provider.py index df617211be9..6ceafb164a8 100644 --- a/tests/unit_tests/models/gemma/test_gemma4_provider.py +++ b/tests/unit_tests/models/gemma/test_gemma4_provider.py @@ -501,6 +501,29 @@ def test_provide_restores_dual_rotary_base(self, provider): mock_rotary.assert_called_once() mock_tied_kv.assert_called_once_with(mock_model, provider) + def test_provide_uses_padded_model_vocab_for_custom_embedding(self): + provider = Gemma4ModelProvider(vocab_size=262145) + mock_model = Mock(vocab_size=262272) + mock_model.embedding = Mock() + mock_model.setup_embeddings_and_output_layer = Mock() + + with ( + patch.object(GPTModelProvider, "provide", return_value=mock_model), + patch("megatron.bridge.models.gemma.gemma4_provider.Gemma3LanguageModelEmbedding") as mock_embedding, + patch("megatron.bridge.models.gemma.gemma4_provider.Gemma4RotaryEmbedding"), + patch("megatron.bridge.models.gemma.gemma4_provider._install_tied_kv"), + ): + provider.provide(pre_process=True, post_process=True) + + assert provider.vocab_size == 262145 + mock_embedding.assert_called_once_with( + config=provider, + vocab_size=262272, + max_sequence_length=provider.seq_length, + position_embedding_type=provider.position_embedding_type, + scatter_to_sequence_parallel=provider.scatter_embedding_sequence_parallel, + ) + def test_provide_restores_dual_rotary_base_on_error(self, provider): with patch.object(GPTModelProvider, "provide", side_effect=RuntimeError("boom")): with pytest.raises(RuntimeError, match="boom"): diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py index 344f5446dc2..44073371454 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py @@ -23,6 +23,7 @@ import os from dataclasses import replace from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -411,6 +412,53 @@ def test_shared_embedding_or_output_weight(self, hf_config): weight_no_decoder = model_no_decoder.shared_embedding_or_output_weight() assert weight_no_decoder is None + @pytest.mark.parametrize( + ("vocab_size", "should_pad_vocab", "expected_vocab_size"), + [ + (151669, True, 152064), + (248077, True, 248320), + (151936, False, 151936), + ], + ) + def test_language_model_honors_vocab_padding_policy( + self, + hf_config, + monkeypatch, + vocab_size, + should_pad_vocab, + expected_vocab_size, + ): + """Apply tokenizer-derived padding before constructing the Qwen language model.""" + self._setup_parallel_state(tp_size=1, ep_size=1, pp_size=1) + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + language_transformer_config = self.get_language_transformer_config(hf_config) + language_transformer_config.vocab_size = vocab_size + language_transformer_config.should_pad_vocab = should_pad_vocab + language_transformer_config.make_vocab_size_divisible_by = 128 + language_transformer_config.tensor_model_parallel_size = 4 + + language_model = Mock() + language_model.config.cuda_graph_impl = "none" + language_model.share_embeddings_and_output_weights = False + language_model_constructor = Mock(return_value=language_model) + monkeypatch.setattr( + "megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model.Qwen3VLGPTModel", + language_model_constructor, + ) + + Qwen3VLModel( + vision_transformer_config=self.get_vision_transformer_config(hf_config), + language_transformer_config=language_transformer_config, + language_transformer_layer_spec=self.get_language_model_layer_spec(), + pre_process=False, + post_process=True, + add_encoder=False, + add_decoder=True, + pg_collection=pg_collection, + ) + + assert language_model_constructor.call_args.kwargs["vocab_size"] == expected_vocab_size + @pytest.mark.timeout(50) def test_set_input_tensor(self, hf_config): """Test set_input_tensor method.""" diff --git a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py index 33116c86b1c..a20fde30ddb 100644 --- a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py +++ b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py @@ -715,7 +715,9 @@ def test_each_qwen35_vl_pretrain_mock_recipe_builds_config(recipe_func: Callable _assert_basic_config(cfg) - assert cfg.tokenizer.tokenizer_type == "NullTokenizer" + assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert cfg.tokenizer.tokenizer_model == cfg.dataset.hf_processor_path + assert cfg.tokenizer.use_tokenizer_vocab_size is True assert getattr(cfg.model, "tensor_model_parallel_size", 1) >= 1 assert getattr(cfg.model, "pipeline_model_parallel_size", 1) >= 1 diff --git a/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py b/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py index e9a295d54f6..84850aa733a 100644 --- a/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py +++ b/tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py @@ -600,7 +600,9 @@ def test_each_qwen3_vl_pretrain_mock_recipe_builds_config(recipe_func: Callable, _assert_basic_config(cfg) - assert cfg.tokenizer.tokenizer_type == "NullTokenizer" + assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert cfg.tokenizer.tokenizer_model == cfg.dataset.hf_processor_path + assert cfg.tokenizer.use_tokenizer_vocab_size is True assert getattr(cfg.model, "tensor_model_parallel_size", 1) >= 1 assert getattr(cfg.model, "pipeline_model_parallel_size", 1) >= 1 diff --git a/tests/unit_tests/recipes/test_all_perf_recipe_factories.py b/tests/unit_tests/recipes/test_all_perf_recipe_factories.py index 4589795a519..79e8b84cf6f 100644 --- a/tests/unit_tests/recipes/test_all_perf_recipe_factories.py +++ b/tests/unit_tests/recipes/test_all_perf_recipe_factories.py @@ -21,7 +21,7 @@ import pytest -from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.config import ConfigContainer, MockVLMSFTDatasetConfig from tests.unit_tests.recipes.recipe_test_utils import ( discover_recipe_factories, exported_recipe_factory_keys, @@ -84,3 +84,9 @@ def test_perf_recipe_factory_builds_config(recipe_factory: Callable[..., object] "dist", ): assert getattr(cfg, section) is not None + + assert cfg.tokenizer.use_tokenizer_vocab_size is False + + if "pretrain" in recipe_factory.__name__ and isinstance(cfg.dataset, MockVLMSFTDatasetConfig): + assert cfg.tokenizer.tokenizer_type == "NullTokenizer" + assert cfg.tokenizer.vocab_size == cfg.model.vocab_size diff --git a/tests/unit_tests/recipes/test_all_recipe_factories.py b/tests/unit_tests/recipes/test_all_recipe_factories.py index 64ef8a4a4b4..c3d32ff88e7 100644 --- a/tests/unit_tests/recipes/test_all_recipe_factories.py +++ b/tests/unit_tests/recipes/test_all_recipe_factories.py @@ -20,6 +20,7 @@ import pytest +from megatron.bridge.recipes.common import _pretrain_common from megatron.bridge.training.config import ConfigContainer from tests.unit_tests.recipes.recipe_test_utils import ( discover_recipe_factories, @@ -65,6 +66,11 @@ def test_all_recipe_factories_are_exported() -> None: assert getattr(_RECIPES_PACKAGE, factory.__name__) is factory +def test_common_pretrain_uses_runtime_tokenizer_vocabulary() -> None: + """From-scratch pretraining derives the model vocabulary from its tokenizer.""" + assert _pretrain_common().tokenizer.use_tokenizer_vocab_size is True + + @pytest.mark.parametrize("recipe_factory", _RUNNABLE_RECIPE_FACTORIES, ids=recipe_factory_id) def test_recipe_factory_builds_config(recipe_factory: Callable[..., object]) -> None: """Every supported recipe can be called with defaults without GPU or network access.""" diff --git a/tests/unit_tests/recipes/test_gemma4_recipe.py b/tests/unit_tests/recipes/test_gemma4_recipe.py index f4ba2c557bf..54439cbbb9c 100644 --- a/tests/unit_tests/recipes/test_gemma4_recipe.py +++ b/tests/unit_tests/recipes/test_gemma4_recipe.py @@ -57,26 +57,28 @@ def to_megatron_provider(self, load_weights=True): return self.provider -def _load_gemma4_recipe_module(): +def _load_gemma4_recipe_module(monkeypatch): """Load the Gemma4 recipe without importing the umbrella recipes package.""" bridge_root = Path(__file__).resolve().parents[3] recipes_root = bridge_root / "src" / "megatron" / "bridge" / "recipes" recipes_pkg = types.ModuleType("megatron.bridge.recipes") recipes_pkg.__path__ = [str(recipes_root)] - sys.modules.setdefault("megatron.bridge.recipes", recipes_pkg) + if "megatron.bridge.recipes" not in sys.modules: + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes", recipes_pkg) common_mod = types.ModuleType("megatron.bridge.recipes.common") common_mod._pretrain_common = _minimal_pretrain_common - sys.modules["megatron.bridge.recipes.common"] = common_mod + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes.common", common_mod) utils_pkg = types.ModuleType("megatron.bridge.recipes.utils") utils_pkg.__path__ = [str(recipes_root / "utils")] - sys.modules.setdefault("megatron.bridge.recipes.utils", utils_pkg) + if "megatron.bridge.recipes.utils" not in sys.modules: + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes.utils", utils_pkg) tokenizer_mod = types.ModuleType("megatron.bridge.recipes.utils.tokenizer_utils") tokenizer_mod.DEFAULT_NULL_TOKENIZER_VOCAB_SIZE = 32000 - sys.modules["megatron.bridge.recipes.utils.tokenizer_utils"] = tokenizer_mod + monkeypatch.setitem(sys.modules, "megatron.bridge.recipes.utils.tokenizer_utils", tokenizer_mod) recipe_path = recipes_root / "gemma" / "gemma4.py" spec = importlib.util.spec_from_file_location("_gemma4_recipe_under_test", recipe_path) @@ -124,8 +126,8 @@ def bridge_provider(hf_config_e4b): @pytest.fixture -def recipe_module(): - return _load_gemma4_recipe_module() +def recipe_module(monkeypatch): + return _load_gemma4_recipe_module(monkeypatch) @pytest.fixture diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index ebf5c178cdb..6f7c7771c81 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -110,9 +110,10 @@ def test_common_environment_defaults_are_small_and_universal(): assert COMMON_PERF_ENV_VARS == {"TORCH_NCCL_HIGH_PRIORITY": 1} -def test_benchmark_common_preserves_legacy_manual_gc_defaults(): +def test_benchmark_common_preserves_legacy_manual_gc_and_model_shape_defaults(): cfg = SimpleNamespace( train=SimpleNamespace(train_iters=0, eval_iters=1, manual_gc=False, manual_gc_interval=0), + tokenizer=SimpleNamespace(use_tokenizer_vocab_size=True), checkpoint=SimpleNamespace(save="checkpoint"), logger=SimpleNamespace(log_interval=10, tensorboard_dir="tensorboard"), ddp=SimpleNamespace(check_for_nan_in_grad=True, check_for_large_grads=True, grad_reduce_in_fp32=True), @@ -133,6 +134,7 @@ def test_benchmark_common_preserves_legacy_manual_gc_defaults(): assert cfg.train.manual_gc is True assert cfg.train.manual_gc_interval == 100 + assert cfg.tokenizer.use_tokenizer_vocab_size is False def test_every_flat_recipe_builder_declares_its_environment_inline(): diff --git a/tests/unit_tests/scripts/performance/test_recipe_environment.py b/tests/unit_tests/scripts/performance/test_recipe_environment.py index 4890a3b1ec7..312495514a0 100644 --- a/tests/unit_tests/scripts/performance/test_recipe_environment.py +++ b/tests/unit_tests/scripts/performance/test_recipe_environment.py @@ -983,6 +983,44 @@ def apply_hydra(config, _overrides): assert effective_recipe.model.moe_flex_dispatcher_backend is None +def test_training_tokenizer_override_preserves_recipe_vocab_policy(): + """Selecting a runtime tokenizer keeps the canonical pretraining vocabulary policy.""" + from argument_parser import parse_cli_args + + from megatron.bridge.recipes.common import _pretrain_common + + parser = parse_cli_args() + args, unknown = parser.parse_known_args( + [ + "--model_family_name", + "gpt", + "--model_recipe_name", + "vanilla_gpt", + "--num_gpus", + "1", + "--gpu", + "h100", + "--use_recipes", + "--max_steps", + "1000", + "--tokenizer_type", + "HuggingFaceTokenizer", + "--tokenizer_model", + "test-tokenizer", + ] + ) + assert unknown == [] + + recipe = _pretrain_common() + assert recipe.tokenizer.use_tokenizer_vocab_size is True + + updated = run_recipe._apply_training_argparse_overrides(recipe, args) + + assert updated.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert updated.tokenizer.tokenizer_model == "test-tokenizer" + assert updated.tokenizer.use_tokenizer_vocab_size is True + + def test_prepare_recipe_runs_base_override_and_finalize_stages(monkeypatch): """Recipe preparation has one visible path with exactly three config stages.""" args = SimpleNamespace( diff --git a/tests/unit_tests/training/test_modelopt_setup.py b/tests/unit_tests/training/test_modelopt_setup.py index 4c7358ee427..c8d0fc694f6 100644 --- a/tests/unit_tests/training/test_modelopt_setup.py +++ b/tests/unit_tests/training/test_modelopt_setup.py @@ -46,7 +46,7 @@ class StopAfterHooksRegistered(Exception): model=SimpleNamespace( fine_grained_activation_offloading=False, restore_modelopt_state=True, - vocab_size=32, + vocab_size=64, ), peft=None, profiling=SimpleNamespace(), @@ -96,6 +96,8 @@ class StopAfterHooksRegistered(Exception): with pytest.raises(StopAfterHooksRegistered): training_setup.setup(state, Mock()) + assert cfg.model.vocab_size == 64 + assert cfg.model.should_pad_vocab is False assert len(hooks) == 1 if resume_has_modelopt_state: hooks[0]([]) diff --git a/tests/unit_tests/training/test_setup.py b/tests/unit_tests/training/test_setup.py index 6f099271ff4..ad917295215 100644 --- a/tests/unit_tests/training/test_setup.py +++ b/tests/unit_tests/training/test_setup.py @@ -195,6 +195,28 @@ def test_vocab_size_larger_than_tokenizer_returns_same_value(self): assert vocab_size == 40960 assert should_pad_vocab is False + def test_pretraining_uses_tokenizer_vocab_over_larger_model_vocab(self): + """Tokenizer-derived pretraining vocab ignores the source model vocabulary.""" + vocab_size, should_pad_vocab = _validate_and_set_vocab_size( + model_vocab_size=248320, + tokenizer_vocab_size=32000, + use_tokenizer_vocab_size=True, + ) + + assert vocab_size == 32000 + assert should_pad_vocab is True + + def test_checkpoint_compatibility_override_preserves_model_vocab(self): + """Disabling the policy preserves the explicit vocabulary used by an existing checkpoint.""" + vocab_size, should_pad_vocab = _validate_and_set_vocab_size( + model_vocab_size=248320, + tokenizer_vocab_size=32000, + use_tokenizer_vocab_size=False, + ) + + assert vocab_size == 248320 + assert should_pad_vocab is False + def test_vocab_size_equal_to_tokenizer_returns_same_value(self): """Test that vocab_size equal to tokenizer returns the same value and disables padding.""" vocab_size, should_pad_vocab = _validate_and_set_vocab_size( From a91e2ab55304bd1f32775f16d0621e05f48749b3 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 06:23:38 +0200 Subject: [PATCH 19/85] cp: `fix(data): preserve VLM pixel defaults (5218)` into `r0.6.0` (#5271) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../bridge/data/energon/hf_task_encoder.py | 25 ++++++++++--------- .../data/energon/test_hf_task_encoder.py | 24 ++++++++++++++++++ 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/src/megatron/bridge/data/energon/hf_task_encoder.py b/src/megatron/bridge/data/energon/hf_task_encoder.py index 232b3e53602..9759a1c30bd 100644 --- a/src/megatron/bridge/data/energon/hf_task_encoder.py +++ b/src/megatron/bridge/data/energon/hf_task_encoder.py @@ -156,18 +156,19 @@ def collate_fn(self, examples: list[dict[str, Any]]) -> dict[str, Any]: The exact batch dictionary returned by the selected HF collate function for this processor type. """ - return self._collate_impl( - examples, - self.processor, - visual_keys=self.visual_keys, - min_pixels=self.min_pixels, - max_pixels=self.max_pixels, - sequence_length=self.seq_length, - pad_to_max_length=self.pad_to_max_length, - pad_to_multiple_of=self.pad_to_multiple_of, - enable_in_batch_packing=self.enable_in_batch_packing, - in_batch_packing_pad_to_multiple_of=self.in_batch_packing_pad_to_multiple_of, - ) + collate_kwargs: dict[str, Any] = { + "visual_keys": self.visual_keys, + "sequence_length": self.seq_length, + "pad_to_max_length": self.pad_to_max_length, + "pad_to_multiple_of": self.pad_to_multiple_of, + "enable_in_batch_packing": self.enable_in_batch_packing, + "in_batch_packing_pad_to_multiple_of": self.in_batch_packing_pad_to_multiple_of, + } + if self.min_pixels is not None: + collate_kwargs["min_pixels"] = self.min_pixels + if self.max_pixels is not None: + collate_kwargs["max_pixels"] = self.max_pixels + return self._collate_impl(examples, self.processor, **collate_kwargs) # ------------------------------------------------------------------ # batch diff --git a/tests/unit_tests/data/energon/test_hf_task_encoder.py b/tests/unit_tests/data/energon/test_hf_task_encoder.py index f778b1219e8..e61f3cde4ab 100644 --- a/tests/unit_tests/data/energon/test_hf_task_encoder.py +++ b/tests/unit_tests/data/energon/test_hf_task_encoder.py @@ -321,6 +321,30 @@ def _collate( self.assertEqual(seen_kwargs["min_pixels"], 16) self.assertEqual(seen_kwargs["max_pixels"], 128) + def test_collate_fn_preserves_model_pixel_defaults_when_unset(self): + seen_kwargs = {} + + def _collate( + examples, + processor, + *, + min_pixels=16, + max_pixels=128, + **kwargs, + ): + seen_kwargs.update(min_pixels=min_pixels, max_pixels=max_pixels) + return _make_collate_fn()(examples, processor, **kwargs) + + encoder = HFTaskEncoder( + processor=self.processor, + seq_length=128, + collate_fn=_collate, + ) + + encoder.collate_fn([{"conversation": [{"role": "user", "content": "manual"}]}]) + + self.assertEqual(seen_kwargs, {"min_pixels": 16, "max_pixels": 128}) + def test_collate_fallback_rejects_oversized_batches(self): processor = _make_processor() encoder = HFTaskEncoder( From 3626e28450d1e37f7421464a19d946c980dec480 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 06:24:06 +0200 Subject: [PATCH 20/85] cp: `fix(data): reject truncation without supervision (5217)` into `r0.6.0` (#5272) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../bridge/data/collators/sequence_padding.py | 14 ++++++++- .../data/collators/test_sequence.py | 30 +++++++++++++++++++ .../qwen_omni/test_qwen3_omni_collate.py | 2 +- 3 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/data/collators/sequence_padding.py b/src/megatron/bridge/data/collators/sequence_padding.py index d06a4120ada..a8f98a1fe4c 100644 --- a/src/megatron/bridge/data/collators/sequence_padding.py +++ b/src/megatron/bridge/data/collators/sequence_padding.py @@ -153,9 +153,21 @@ def pad_or_truncate_sequence_batch( if pad_to_max_length else min(sequence_length, _ceil_to_multiple(tokens.size(1), pad_to_multiple_of)) ) + loss_mask = batch.get("loss_mask") + prepared_loss_mask = _pad_or_truncate_2d(loss_mask, target_length, 0) + if tokens.size(1) > target_length and loss_mask is not None and prepared_loss_mask is not None: + had_supervision = torch.any(loss_mask != 0, dim=1) + has_supervision = torch.any(prepared_loss_mask != 0, dim=1) + emptied_rows = torch.nonzero(had_supervision & ~has_supervision, as_tuple=False).flatten().tolist() + if emptied_rows: + raise ValueError( + f"Truncating sequence batch from {tokens.size(1)} to {target_length} tokens removed every " + f"supervised token from rows {emptied_rows}. Increase sequence_length or filter overlength samples." + ) + _set_tokens(batch, token_key, _pad_or_truncate_2d(tokens, target_length, pad_token_id)) batch["labels"] = _pad_or_truncate_2d(batch.get("labels"), target_length, ignore_index) - batch["loss_mask"] = _pad_or_truncate_2d(batch.get("loss_mask"), target_length, 0) + batch["loss_mask"] = prepared_loss_mask batch["position_ids"] = _pad_or_truncate_position_ids(batch.get("position_ids"), target_length) batch["attention_mask"] = _pad_or_truncate_attention_mask(batch.get("attention_mask"), target_length) for key, tensor_pad_value in (sequence_tensor_pad_values or {}).items(): diff --git a/tests/unit_tests/data/collators/test_sequence.py b/tests/unit_tests/data/collators/test_sequence.py index 0d1e38ac52c..b0e63ea86c8 100644 --- a/tests/unit_tests/data/collators/test_sequence.py +++ b/tests/unit_tests/data/collators/test_sequence.py @@ -62,6 +62,36 @@ def test_prepare_sequence_batch_pads_to_model_length_when_required(): assert batch["attention_mask"].tolist() == [[1, 1, 1, 0, 0, 0]] +def test_prepare_sequence_batch_rejects_truncation_that_removes_all_supervision(): + batch = { + "input_ids": torch.arange(8).unsqueeze(0), + "labels": torch.arange(8).unsqueeze(0), + "loss_mask": torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0]]), + "position_ids": torch.arange(8).unsqueeze(0), + "attention_mask": torch.ones((1, 8), dtype=torch.long), + } + + with pytest.raises(ValueError, match=r"removed every supervised token from rows \[0\]"): + prepare_sequence_batch(batch, sequence_length=6, pad_to_multiple_of=1) + + assert batch["input_ids"].shape == (1, 8) + + +def test_prepare_sequence_batch_allows_truncation_that_preserves_supervision(): + batch = { + "input_ids": torch.arange(8).unsqueeze(0), + "labels": torch.arange(8).unsqueeze(0), + "loss_mask": torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0]]), + "position_ids": torch.arange(8).unsqueeze(0), + "attention_mask": torch.ones((1, 8), dtype=torch.long), + } + + prepare_sequence_batch(batch, sequence_length=6, pad_to_multiple_of=1) + + assert batch["input_ids"].shape == (1, 6) + assert batch["loss_mask"].sum().item() == 1.0 + + def test_prepare_sequence_batch_handles_rectangular_4d_attention_mask(): batch = { "input_ids": torch.tensor([[1, 2, 3]]), diff --git a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_collate.py b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_collate.py index 1d6c7bba76b..7e1a28ad656 100644 --- a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_collate.py +++ b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_collate.py @@ -134,7 +134,7 @@ def test_direct_sft_dataset_collates_qwen3_omni_media_and_loss_mask(): def test_qwen3_omni_collator_right_pads_before_truncating_mixed_lengths(): short_row = [103, 20, 105, 100, 30, 31, 105] - long_row = [103, *([20] * 70), 105, 100, 30, 31, 105] + long_row = [103, 20, 105, 100, *([30] * 70), 31, 105] processor = Qwen3OmniMoeProcessor([short_row, long_row]) examples = _examples() dataset = DirectSFTDataset( From fdffbe0985910953966da39a86d5d41749e41a83 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 06:24:31 +0200 Subject: [PATCH 21/85] cp: `fix(data): normalize GLM VLM assistant content (5216)` into `r0.6.0` (#5273) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../bridge/models/glm_vl/data/collate_fn.py | 42 +++++++++++++++++++ .../data/collators/test_model_collators.py | 42 +++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/src/megatron/bridge/models/glm_vl/data/collate_fn.py b/src/megatron/bridge/models/glm_vl/data/collate_fn.py index d017342c7b6..b881fb311a7 100644 --- a/src/megatron/bridge/models/glm_vl/data/collate_fn.py +++ b/src/megatron/bridge/models/glm_vl/data/collate_fn.py @@ -14,6 +14,7 @@ """GLM VL collator implementations.""" +from collections.abc import Mapping from typing import Any import torch @@ -42,6 +43,46 @@ GLM4V_NEXT_ROLE_MARKERS = ("<|system|>\n", "<|user|>\n", "<|observation|>\n") +def _normalize_glm4v_assistant_content(example: dict[str, Any]) -> dict[str, Any]: + """Flatten text-only structured assistant content for GLM chat templates.""" + conversation = example.get("conversation") + if not isinstance(conversation, list): + return example + + normalized_conversation = [] + changed = False + for turn in conversation: + if not isinstance(turn, Mapping) or turn.get("role") != "assistant": + normalized_conversation.append(turn) + continue + + parts = turn.get("content") + if not isinstance(parts, list) or not parts: + normalized_conversation.append(turn) + continue + + text_parts = [] + for part in parts: + if isinstance(part, Mapping) and part.get("type") == "text" and isinstance(part.get("text"), str): + text_parts.append(part["text"]) + else: + break + else: + normalized_turn = dict(turn) + normalized_turn["content"] = "".join(text_parts) + normalized_conversation.append(normalized_turn) + changed = True + continue + + normalized_conversation.append(turn) + + if not changed: + return example + normalized_example = dict(example) + normalized_example["conversation"] = normalized_conversation + return normalized_example + + def _glm4v_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundaryConfig: """Build GLM-4.5V role boundaries and exclude its empty thinking prefix.""" tokenizer = get_processor_tokenizer(processor) @@ -107,6 +148,7 @@ def glm4v_collate_fn( """ del visual_keys, min_pixels, max_pixels + examples = [_normalize_glm4v_assistant_content(example) for example in examples] skipped_tokens = extract_skipped_token_ids(processor) boundary_config = _glm4v_assistant_mask_boundary_config(processor) diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py index c0e3348d454..e24965d7790 100644 --- a/tests/unit_tests/data/collators/test_model_collators.py +++ b/tests/unit_tests/data/collators/test_model_collators.py @@ -904,6 +904,48 @@ def apply_chat_template(self, conversations, **kwargs): assert processor.tokenizer.padding_side == "left" +def test_glm4v_collate_flattens_structured_assistant_content(monkeypatch): + class _GlmProcessor: + class _Tokenizer: + padding_side = "left" + pad_token_id = 0 + + def __init__(self): + self.tokenizer = self._Tokenizer() + self.conversations = None + + def apply_chat_template(self, conversations, **kwargs): + self.conversations = conversations + return { + "input_ids": torch.tensor([[1, 2, 3]]), + "attention_mask": torch.ones((1, 3), dtype=torch.long), + "mm_token_type_ids": torch.zeros((1, 3), dtype=torch.long), + } + + monkeypatch.setattr( + glm_vl_collate, "extract_skipped_token_ids", lambda processor: torch.empty(0, dtype=torch.long) + ) + monkeypatch.setattr(glm_vl_collate, "_glm4v_assistant_mask_boundary_config", lambda processor: None) + monkeypatch.setattr( + glm_vl_collate, + "_build_glm4v_assistant_loss_mask", + lambda example, input_ids, *args, **kwargs: torch.ones_like(input_ids, dtype=torch.float32), + ) + assistant_content = [{"type": "text", "text": "A picture."}] + example = { + "conversation": [ + {"role": "user", "content": "Describe."}, + {"role": "assistant", "content": assistant_content}, + ] + } + processor = _GlmProcessor() + + glm_vl_collate.glm4v_collate_fn([example], processor) + + assert processor.conversations[0][-1]["content"] == "A picture." + assert example["conversation"][-1]["content"] == assistant_content + + @pytest.mark.parametrize( ("input_ids", "expected_mask"), [ From c5a6aaefa17371e2c701b17ea5c2be4179403681 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 06:24:54 +0200 Subject: [PATCH 22/85] cp: `fix(data): pad Qwen VLM token type IDs (5215)` into `r0.6.0` (#5274) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../bridge/models/qwen_vl/data/collate_fn.py | 1 + .../data/collators/test_model_collators.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/src/megatron/bridge/models/qwen_vl/data/collate_fn.py b/src/megatron/bridge/models/qwen_vl/data/collate_fn.py index 8dfd8d5c79f..23223d1bdbf 100644 --- a/src/megatron/bridge/models/qwen_vl/data/collate_fn.py +++ b/src/megatron/bridge/models/qwen_vl/data/collate_fn.py @@ -364,5 +364,6 @@ def pad_to(x, tgt_len, value): enable_in_batch_packing=enable_in_batch_packing, in_batch_packing_pad_to_multiple_of=in_batch_packing_pad_to_multiple_of, ignore_index=IGNORE_INDEX, + sequence_tensor_pad_values={"mm_token_type_ids": 0}, ) return batch diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py index e24965d7790..2c035fc5952 100644 --- a/tests/unit_tests/data/collators/test_model_collators.py +++ b/tests/unit_tests/data/collators/test_model_collators.py @@ -227,6 +227,31 @@ def test_qwen2_5_collate_fn_uses_shared_pixel_defaults(monkeypatch): assert proc.processor_kwargs[-1]["max_pixels"] == qwen_vl_collate.QWEN_VL_MAX_PIXELS +def test_qwen2_5_collate_fn_pads_mm_token_type_ids_with_sequence(monkeypatch): + monkeypatch.setattr(qwen_vl_collate, "HAVE_QWEN_VL_UTILS", True) + monkeypatch.setattr(qwen_vl_collate, "process_vision_info", lambda conv: (None, None)) + + class _TokenTypeProcessor(_DummyProcessor): + def __call__(self, *args, **kwargs): + batch = super().__call__(*args, **kwargs) + batch["mm_token_type_ids"] = torch.tensor([[0, 1, 0]]) + return batch + + examples = [ + {"conversation": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}, + ] + + batch = collate.qwen2_5_collate_fn( + examples, + _TokenTypeProcessor(), + sequence_length=8, + pad_to_multiple_of=4, + ) + + assert batch["input_ids"].shape == batch["mm_token_type_ids"].shape == (1, 4) + assert batch["mm_token_type_ids"].tolist() == [[0, 1, 0, 0]] + + def test_qwen2_audio_collate_fn_uses_audio_inputs_key(monkeypatch): """qwen2_audio_collate_fn should store Qwen2AudioInputs under 'audio_inputs', not 'visual_inputs'.""" From b75770454f04bb46fa0f5309144535e5fc5b9f41 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 06:25:20 +0200 Subject: [PATCH 23/85] cp: `perf(recipe): tune Nemotron 3.5 Nano H100 benchmark (5266)` into `r0.6.0` (#5279) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../nemotron-3.5-nano/card.yaml | 36 +++++++----- .../perf_recipes/nemotronh/h100/nemotronh.py | 29 ++++++++-- .../test_nemotron_3_5_nano_perf_recipes.py | 58 +++++++++++++++++-- 3 files changed, 98 insertions(+), 25 deletions(-) diff --git a/examples/model_verification_cards/nemotron-3.5-nano/card.yaml b/examples/model_verification_cards/nemotron-3.5-nano/card.yaml index 89a68a3ea3e..9d26eb038fd 100644 --- a/examples/model_verification_cards/nemotron-3.5-nano/card.yaml +++ b/examples/model_verification_cards/nemotron-3.5-nano/card.yaml @@ -9,8 +9,10 @@ summary: > performance results. The convergence pretraining recipes retain real tokenization, real data, natural routing, safety checks, and checkpointing while adopting the corresponding hardware performance recipe's execution - stack. Nemotron 3.5 Nano support verification covers conversion, inference, - and training on H100 and GB200. The + stack. The H100 performance leaf uses BF16 Adam moments in a benchmark-only + workload and is not convergence-comparable with FP32 optimizer-state runs. + Nemotron 3.5 Nano support verification covers conversion, inference, and + training on H100 and GB200. The model.hf_id/--hf-model and model.hf_revision/--hf-revision values are non-functional placeholders pending publication. verification_index: @@ -673,26 +675,28 @@ items: H100: status: verified precision: bf16 - bridge_commit: 9b0c2dff42e6175e91d34cd2a504529f803ebde9 # pragma: allowlist secret + bridge_commit: 573e088c9c6740082c39744e03dc5b009e730ed4 # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 --recipe nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config --max_steps 50 - last_verified: 2026-07-28 + last_verified: 2026-08-03 metrics: - initial_loss: 12.16155 - final_loss: 0.01361997 - last_10_steps_step_time_ms_avg: 20826.370 - last_10_steps_model_tflops_per_gpu_avg: 354.300 - peak_allocated_memory_gib: 71.949 - peak_reserved_memory_gib: 74.631 + initial_loss: 12.16197 + final_loss: 0.01217863 + last_10_steps_step_time_ms_avg: 18467.280 + last_10_steps_model_tflops_per_gpu_avg: 399.580 + peak_allocated_memory_gib: 72.066 + peak_reserved_memory_gib: 75.340 expected_result: > - On 16x H100, the exact mock-data BF16 performance recipe completes - exactly 50 steps with finite LM/MTP-1/MTP-2 losses and no skipped or - NaN iterations. Transformer Engine CUDA graphs cover attention and - Mamba scopes. Steps 41-50 average 20,826.370 ms and 354.300 model - TFLOP/s/GPU; peak rank-0 allocated/reserved memory is 71.949/74.631 - GiB. + On 16x H100, the exact mock-data, force-balanced BF16 performance + recipe completes exactly 50 steps with finite LM/MTP-1/MTP-2 losses + and no skipped or NaN iterations. Transformer Engine CUDA graphs cover + attention and Mamba scopes, while expert-FC1 activations are offloaded. + Steps 41-50 average 18,467.280 ms and 399.580 model TFLOP/s/GPU; peak + rank-0 allocated/reserved memory is 72.066/75.340 GiB. This + benchmark-only run stores Adam moments in BF16, so its losses are not + convergence-comparable with FP32 optimizer-state runs. GB200: status: verified diff --git a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py index 22fac9f9a3b..eefd9c07037 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py @@ -13,6 +13,8 @@ # limitations under the License. """H100 performance recipes for NemotronH and Nemotron 3.""" +import torch + from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.perf_recipes.nemotronh.common import ( ConfigContainer, @@ -173,7 +175,7 @@ def nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer: def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: - """Nemotron 3.5 Nano pretrain: 16× H100, BF16.""" + """Nemotron 3.5 Nano benchmark pretrain: 16× H100, BF16.""" cfg = nemotron_3_nano_pretrain_16gpu_h100_bf16_config() # Keep the benchmark workload aligned with the GB200 BF16 recipe. The # hardware recipes may tune execution-only knobs such as microbatch size, @@ -186,6 +188,24 @@ def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: cfg.model.mtp_loss_scaling_factor = 0.3 cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + + # This mock-data, force-balanced benchmark uses lower-precision Adam moments + # to make the measured selective-recompute and activation-offload stack fit. + # It is not convergence-equivalent to recipes with FP32 optimizer states. + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.exp_avg_dtype = torch.bfloat16 + cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 + + cfg.model.recompute_modules = ["moe_act", "layernorm"] + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["expert_fc1"] + cfg.model.activation_offload_fraction = 1.0 + cfg.model.delay_offload_until_cuda_graph = True + + cfg.model.moe_router_fusion = True + cfg.model.moe_permute_fusion_into_hybridep = True + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_flex_dispatcher_num_sms = 32 cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, @@ -194,11 +214,12 @@ def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "NCCL_NVLS_ENABLE": 0, "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 64, "NVLINK_DOMAIN_SIZE": 8, "USE_MNNVL": 0, - "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, - "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 10, + "NVTE_CPU_OFFLOAD_V1": 1, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 10, "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, } diff --git a/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py b/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py index 3675747c53a..5f474447b43 100644 --- a/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py +++ b/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py @@ -19,6 +19,7 @@ from inspect import signature import pytest +import torch from megatron.bridge.perf_recipes.nemotronh import ( nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config, @@ -173,11 +174,12 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy( recipe_factory: Callable[[], ConfigContainer], base_recipe_factory: Callable[[], ConfigContainer], ) -> None: - """Nemotron 3.5 variants inherit environment, loss normalization, and RNG policy.""" + """Nemotron 3.5 variants inherit shared loss normalization and RNG policy.""" cfg = recipe_factory() base_cfg = base_recipe_factory() - assert cfg.env_vars == base_cfg.env_vars + if recipe_factory is not nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config: + assert cfg.env_vars == base_cfg.env_vars assert cfg.model.calculate_per_token_loss == base_cfg.model.calculate_per_token_loss assert cfg.model.use_te_rng_tracker == base_cfg.model.use_te_rng_tracker assert cfg.tokenizer.tokenizer_model != base_cfg.tokenizer.tokenizer_model @@ -196,14 +198,49 @@ def test_h100_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer] assert cfg.model.recompute_granularity == "selective" assert cfg.model.seq_length == 8192 assert cfg.dataset.seq_length == 8192 - assert cfg.model.moe_hybridep_num_sms == 16 - assert cfg.optimizer.use_precision_aware_optimizer is False assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 8 assert cfg.env_vars["USE_MNNVL"] == 0 +def test_h100_bf16_perf_recipe_uses_measured_benchmark_tuning() -> None: + """The H100 BF16 benchmark preserves the measured memory and HybridEP tuning.""" + cfg = nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() + + assert cfg.optimizer.use_precision_aware_optimizer is True + assert cfg.optimizer.main_params_dtype == torch.float32 + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.optimizer.exp_avg_dtype == torch.bfloat16 + assert cfg.optimizer.exp_avg_sq_dtype == torch.bfloat16 + + assert cfg.model.recompute_modules == ["moe_act", "layernorm"] + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["expert_fc1"] + assert cfg.model.activation_offload_fraction == 1.0 + assert cfg.model.delay_offload_until_cuda_graph is True + + assert cfg.model.moe_router_fusion is True + assert cfg.model.moe_permute_fusion_into_hybridep is True + assert cfg.model.moe_hybridep_num_sms is None + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 64 + assert cfg.env_vars["NVTE_BWD_LAYERNORM_SM_MARGIN"] == 10 + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + assert cfg.env_vars["NVTE_FWD_LAYERNORM_SM_MARGIN"] == 10 + + +def test_h100_fp8_perf_recipe_retains_fp32_optimizer_state() -> None: + """The unrelated H100 FP8 benchmark retains its existing optimizer and dispatcher policy.""" + cfg = nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config() + + assert cfg.optimizer.use_precision_aware_optimizer is False + assert cfg.optimizer.exp_avg_dtype == torch.float32 + assert cfg.optimizer.exp_avg_sq_dtype == torch.float32 + assert cfg.model.moe_hybridep_num_sms == 16 + assert cfg.model.moe_flex_dispatcher_num_sms is None + + def test_bf16_perf_recipes_share_training_workload() -> None: - """H100 and GB200 BF16 recipes differ only in execution tuning.""" + """H100 and GB200 BF16 recipes otherwise share the same model workload.""" h100_cfg = nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() gb200_cfg = nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config() @@ -213,6 +250,17 @@ def test_bf16_perf_recipes_share_training_workload() -> None: gb200_cfg.model.cuda_graph_impl = h100_cfg.model.cuda_graph_impl gb200_cfg.model.cuda_graph_scope = h100_cfg.model.cuda_graph_scope gb200_cfg.model.cuda_graph_modules = h100_cfg.model.cuda_graph_modules + gb200_cfg.model.fine_grained_activation_offloading = h100_cfg.model.fine_grained_activation_offloading + gb200_cfg.model.offload_modules = h100_cfg.model.offload_modules + gb200_cfg.model.activation_offload_fraction = h100_cfg.model.activation_offload_fraction + gb200_cfg.model.delay_offload_until_cuda_graph = h100_cfg.model.delay_offload_until_cuda_graph + gb200_cfg.model.moe_router_fusion = h100_cfg.model.moe_router_fusion + gb200_cfg.model.moe_permute_fusion_into_hybridep = h100_cfg.model.moe_permute_fusion_into_hybridep + gb200_cfg.model.moe_hybridep_num_sms = h100_cfg.model.moe_hybridep_num_sms + gb200_cfg.model.moe_flex_dispatcher_num_sms = h100_cfg.model.moe_flex_dispatcher_num_sms + gb200_cfg.optimizer.use_precision_aware_optimizer = h100_cfg.optimizer.use_precision_aware_optimizer + gb200_cfg.optimizer.exp_avg_dtype = h100_cfg.optimizer.exp_avg_dtype + gb200_cfg.optimizer.exp_avg_sq_dtype = h100_cfg.optimizer.exp_avg_sq_dtype gb200_cfg.env_vars = h100_cfg.env_vars assert gb200_cfg == h100_cfg From df9cfb12042c8d4f613b097b5699da2128de6c8e Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 4 Aug 2026 08:24:50 +0200 Subject: [PATCH 24/85] cp: `fix(training): restore weak scaling for benchmark recipes (5222)` into `r0.6.0` (#5282) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- scripts/common/benchmark_parallelism.py | 143 ++++++++++++++ scripts/performance/run_script.py | 42 +++- scripts/performance/utils/overrides.py | 29 ++- scripts/performance/utils/utils.py | 75 +++++++- scripts/training/README.md | 18 +- scripts/training/run_recipe.py | 90 +++++++-- scripts/training/setup_experiment.py | 8 +- .../recipes/test_override_precedence.py | 120 +++++++++++- .../performance/test_recipe_environment.py | 179 ++++++++++++++++++ .../scripts/training/test_run_recipe.py | 164 +++++++++++++++- .../scripts/training/test_setup_experiment.py | 5 +- 11 files changed, 824 insertions(+), 49 deletions(-) create mode 100644 scripts/common/benchmark_parallelism.py diff --git a/scripts/common/benchmark_parallelism.py b/scripts/common/benchmark_parallelism.py new file mode 100644 index 00000000000..24fc1aec7c5 --- /dev/null +++ b/scripts/common/benchmark_parallelism.py @@ -0,0 +1,143 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Lightweight parallel-topology helpers shared by benchmark launchers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class ParallelTopology: + """Model and expert parallel degrees that constrain a benchmark world.""" + + tensor_parallel: int = 1 + pipeline_parallel: int = 1 + context_parallel: int = 1 + expert_parallel: int = 1 + expert_tensor_parallel: int | None = None + gtp_remat: int = 1 + expert_gtp_remat: int = 1 + + +def _first_attribute(config: Any, names: tuple[str, ...], default: Any) -> Any: + """Return the first exposed configuration attribute from ``names``.""" + for name in names: + if hasattr(config, name): + return getattr(config, name) + return default + + +def topology_from_config(config: Any) -> ParallelTopology: + """Read parallel degrees from a model config or lightweight workload config.""" + return ParallelTopology( + tensor_parallel=getattr(config, "tensor_model_parallel_size", 1), + pipeline_parallel=getattr(config, "pipeline_model_parallel_size", 1), + context_parallel=getattr(config, "context_parallel_size", 1), + expert_parallel=getattr(config, "expert_model_parallel_size", 1), + expert_tensor_parallel=getattr(config, "expert_tensor_parallel_size", None), + gtp_remat=_first_attribute(config, ("gtp_weight_remat_size", "gtp_remat_size"), 1), + expert_gtp_remat=_first_attribute( + config, + ("expert_gtp_weight_remat_size", "expert_gtp_remat_size"), + 1, + ), + ) + + +def _validate_positive_sizes(sizes: dict[str, int]) -> None: + """Require positive integer parallel degrees.""" + invalid_sizes = {name: size for name, size in sizes.items() if not isinstance(size, int) or size <= 0} + if invalid_sizes: + raise ValueError(f"Parallel sizes must be positive integers, got {invalid_sizes}.") + + +def data_parallel_size(*, num_gpus: int, topology: ParallelTopology) -> int: + """Validate dense and expert grids, then return the dense DP degree.""" + if not isinstance(num_gpus, int) or num_gpus <= 0: + raise ValueError(f"Number of GPUs must be a positive integer, got {num_gpus!r}.") + + expert_tensor_parallel = ( + topology.tensor_parallel if topology.expert_tensor_parallel is None else topology.expert_tensor_parallel + ) + sizes = { + "TP": topology.tensor_parallel, + "PP": topology.pipeline_parallel, + "CP": topology.context_parallel, + "EP": topology.expert_parallel, + "ETP": expert_tensor_parallel, + "GTP remat": topology.gtp_remat, + "expert GTP remat": topology.expert_gtp_remat, + } + _validate_positive_sizes(sizes) + + dense_factor = ( + topology.tensor_parallel * topology.pipeline_parallel * topology.context_parallel * topology.gtp_remat + ) + if num_gpus % dense_factor != 0: + if topology.gtp_remat == 1: + factor_name = "TP * PP * CP" + factor_values = f"{topology.tensor_parallel} * {topology.pipeline_parallel} * {topology.context_parallel}" + else: + factor_name = "TP * PP * CP * GTP remat" + factor_values = ( + f"{topology.tensor_parallel} * {topology.pipeline_parallel} * {topology.context_parallel} " + f"* {topology.gtp_remat}" + ) + raise ValueError( + f"Requested {num_gpus} GPUs is not divisible by {factor_name} " + f"({factor_values} = {dense_factor}). Override the parallel sizes for this GPU count." + ) + + expert_factor = ( + expert_tensor_parallel * topology.expert_parallel * topology.pipeline_parallel * topology.expert_gtp_remat + ) + if num_gpus % expert_factor != 0: + if topology.expert_gtp_remat == 1: + factor_name = "ETP * EP * PP" + factor_values = f"{expert_tensor_parallel} * {topology.expert_parallel} * {topology.pipeline_parallel}" + else: + factor_name = "ETP * EP * PP * expert GTP remat" + factor_values = ( + f"{expert_tensor_parallel} * {topology.expert_parallel} * {topology.pipeline_parallel} " + f"* {topology.expert_gtp_remat}" + ) + raise ValueError( + f"The expert-parallel grid is incompatible with requested {num_gpus} GPUs: the world size must be " + f"divisible by {factor_name} ({factor_values} = {expert_factor}). " + "Override the expert or pipeline parallel sizes for this GPU count." + ) + + return num_gpus // dense_factor + + +def weak_scaled_global_batch_size(*, base_gbs: int, base_data_parallel: int, data_parallel: int) -> int: + """Preserve samples per DP rank while requiring an integral GBS.""" + sizes = { + "canonical DP": base_data_parallel, + "requested DP": data_parallel, + "canonical GBS": base_gbs, + } + invalid_sizes = {name: size for name, size in sizes.items() if not isinstance(size, int) or size <= 0} + if invalid_sizes: + raise ValueError(f"Weak-scaling inputs must be positive integers, got {invalid_sizes}.") + + scaled_gbs_numerator = base_gbs * data_parallel + if scaled_gbs_numerator % base_data_parallel != 0: + raise ValueError( + f"Weak scaling GBS {base_gbs} from DP {base_data_parallel} to DP {data_parallel} does not produce " + "an integer global batch size. Pass --global_batch_size explicitly." + ) + return scaled_gbs_numerator // base_data_parallel diff --git a/scripts/performance/run_script.py b/scripts/performance/run_script.py index 90e2b632e3a..e0bc3852ca4 100644 --- a/scripts/performance/run_script.py +++ b/scripts/performance/run_script.py @@ -20,7 +20,7 @@ import sys from argument_parser import parse_cli_args -from utils.utils import get_perf_recipe_by_name +from utils.utils import PerfRecipeNotFoundError, get_perf_optimized_recipe, get_perf_recipe_by_name logger = logging.getLogger(__name__) @@ -96,14 +96,38 @@ def _apply_perf_recipe_overrides(recipe, cli_overrides: list[str], args): def _prepare_perf_recipe(args, cli_overrides: list[str]): """Build a flat performance recipe with all user overrides applied.""" - recipe = get_perf_recipe_by_name( - model_recipe_name=args.model_recipe_name, - task=args.task, - num_gpus=args.num_gpus, - gpu=args.gpu, - precision=args.compute_dtype, - config_variant=getattr(args, "config_variant", None), - ) + try: + recipe = get_perf_recipe_by_name( + model_recipe_name=args.model_recipe_name, + task=args.task, + num_gpus=args.num_gpus, + gpu=args.gpu, + precision=args.compute_dtype, + config_variant=getattr(args, "config_variant", None), + ) + except PerfRecipeNotFoundError: + recipe = get_perf_optimized_recipe( + model_family_name=args.model_family_name, + model_recipe_name=args.model_recipe_name, + train_task=args.task, + gpu=args.gpu, + compute_dtype=args.compute_dtype, + config_variant=getattr(args, "config_variant", None), + ) + recipe = _apply_perf_recipe_overrides(recipe, cli_overrides, args) + from utils.overrides import set_post_overrides + + return set_post_overrides( + recipe, + model_family_name=args.model_family_name, + model_recipe_name=args.model_recipe_name, + gpu=args.gpu, + num_gpus=args.num_gpus, + compute_dtype=args.compute_dtype, + task=args.task, + user_gbs=args.global_batch_size, + config_variant=getattr(args, "config_variant", None), + ) return _apply_perf_recipe_overrides(recipe, cli_overrides, args) diff --git a/scripts/performance/utils/overrides.py b/scripts/performance/utils/overrides.py index 2ddd07bf70d..ad43a79c184 100644 --- a/scripts/performance/utils/overrides.py +++ b/scripts/performance/utils/overrides.py @@ -44,7 +44,13 @@ create_rp2_dataset_config, create_squad_dataset_config, ) -from utils.utils import WorkloadBaseConfig, get_workload_base_config +from utils.utils import ( + WorkloadBaseConfig, + _validated_data_parallel_size, + _weak_scaled_global_batch_size, + get_workload_base_config, + topology_from_config, +) logger = logging.getLogger(__name__) @@ -665,7 +671,14 @@ def set_post_overrides( cp = recipe.model.context_parallel_size vp = recipe.model.virtual_pipeline_model_parallel_size or 1 - dp = int(num_gpus / (tp * pp * cp)) + base_dp = _validated_data_parallel_size( + num_gpus=workload_base_config.num_gpus, + topology=topology_from_config(workload_base_config), + ) + dp = _validated_data_parallel_size( + num_gpus=num_gpus, + topology=topology_from_config(recipe.model), + ) logger.info(f"DP: {dp}; TP: {tp}; PP: {pp}; CP: {cp}; VP: {vp}") # NOTE: overlap_param_gather_with_optimizer_step causes NaN grad norm for fp8_mx. Disabling it until the issue is resolved. if dp > 1 and pp > 1 and vp > 1 and compute_dtype not in ("fp8_mx", "nvfp4"): @@ -675,16 +688,20 @@ def set_post_overrides( if hasattr(recipe, "comm_overlap") and isinstance(recipe.comm_overlap, CommOverlapConfig): recipe.comm_overlap.overlap_param_gather_with_optimizer_step = True - default_num_gpus = workload_base_config.num_gpus if user_gbs is None: # Only auto-rescale if the recipe's current GBS still equals the # workload default — i.e., no one (Hydra or earlier step) has already # expressed an intent. Otherwise Hydra-set GBS gets silently stomped. - if recipe.train.global_batch_size == workload_base_config.global_batch_size and num_gpus != default_num_gpus: - new_gbs = int(workload_base_config.gbs_scaling_factor * num_gpus) + if recipe.train.global_batch_size == workload_base_config.global_batch_size and dp != base_dp: + new_gbs = _weak_scaled_global_batch_size( + base_gbs=workload_base_config.global_batch_size, + base_data_parallel=base_dp, + data_parallel=dp, + ) recipe.train.global_batch_size = new_gbs logger.info( - f"Scaled global batch size from {workload_base_config.global_batch_size} to {new_gbs} based on {num_gpus} GPUs." + f"Scaled global batch size from {workload_base_config.global_batch_size} to {new_gbs} " + f"based on DP {base_dp} -> {dp}." ) return recipe diff --git a/scripts/performance/utils/utils.py b/scripts/performance/utils/utils.py index 17715f3d8c9..55def68af35 100644 --- a/scripts/performance/utils/utils.py +++ b/scripts/performance/utils/utils.py @@ -31,6 +31,21 @@ logger = logging.getLogger(__name__) +_COMMON_SCRIPTS_ROOT = Path(__file__).resolve().parents[2] / "common" +if str(_COMMON_SCRIPTS_ROOT) not in sys.path: + sys.path.insert(0, str(_COMMON_SCRIPTS_ROOT)) + +from benchmark_parallelism import ( # noqa: E402 + data_parallel_size as _validated_data_parallel_size, +) +from benchmark_parallelism import ( + topology_from_config, +) +from benchmark_parallelism import ( + weak_scaled_global_batch_size as _weak_scaled_global_batch_size, +) + + # Default timeout for interactive config variant selection (in seconds) CONFIG_VARIANT_SELECTION_TIMEOUT = 15 _REPO_ROOT = Path(__file__).resolve().parents[3] @@ -72,6 +87,30 @@ } +class PerfRecipeNotFoundError(ValueError): + """Raised when an exact flat performance recipe name is unavailable.""" + + +def _data_parallel_size( + *, + num_gpus: int, + tensor_parallel: int, + pipeline_parallel: int, + context_parallel: int, + expert_parallel: int = 1, + expert_tensor_parallel: int | None = None, +) -> int: + """Return DP after validating the requested dense and expert grids.""" + topology = WorkloadBaseConfig( + tensor_model_parallel_size=tensor_parallel, + pipeline_model_parallel_size=pipeline_parallel, + context_parallel_size=context_parallel, + expert_model_parallel_size=expert_parallel, + expert_tensor_parallel_size=expert_tensor_parallel, + ) + return _validated_data_parallel_size(num_gpus=num_gpus, topology=topology_from_config(topology)) + + def configure_slurm_gpu_tuning( executor: Any, *, @@ -170,6 +209,8 @@ class WorkloadBaseConfig: virtual_pipeline_model_parallel_size: int | None = None expert_model_parallel_size: int = 1 expert_tensor_parallel_size: int | None = None + gtp_weight_remat_size: int = 1 + expert_gtp_weight_remat_size: int = 1 global_batch_size: int = 1 micro_batch_size: int = 1 @@ -499,7 +540,7 @@ def get_perf_recipe_by_name( recipe_fn = find_perf_recipe(name) if recipe_fn is None: searched_modules = ", ".join(perf_recipe_family_modules()) or "none" - raise ValueError(f"No perf recipe {name!r} found in perf recipe packages: {searched_modules}.") + raise PerfRecipeNotFoundError(f"No perf recipe {name!r} found in perf recipe packages: {searched_modules}.") return recipe_fn() @@ -603,6 +644,12 @@ def _workload_base_config_from_recipe(config, *, num_gpus: int) -> WorkloadBaseC virtual_pipeline_model_parallel_size=model.virtual_pipeline_model_parallel_size, expert_model_parallel_size=getattr(model, "expert_model_parallel_size", 1), expert_tensor_parallel_size=getattr(model, "expert_tensor_parallel_size", None), + gtp_weight_remat_size=getattr(model, "gtp_weight_remat_size", getattr(model, "gtp_remat_size", 1)), + expert_gtp_weight_remat_size=getattr( + model, + "expert_gtp_weight_remat_size", + getattr(model, "expert_gtp_remat_size", 1), + ), global_batch_size=train.global_batch_size, micro_batch_size=train.micro_batch_size, env_vars=dict(getattr(config, "env_vars", {})), @@ -701,10 +748,32 @@ def get_exp_name_config( ) mbs_size = args.micro_batch_size if args.micro_batch_size is not None else base_config.micro_batch_size + requested_topology = WorkloadBaseConfig( + tensor_model_parallel_size=tp_size, + pipeline_model_parallel_size=pp_size, + context_parallel_size=cp_size, + expert_model_parallel_size=ep_size, + expert_tensor_parallel_size=etp_size, + gtp_weight_remat_size=base_config.gtp_weight_remat_size, + expert_gtp_weight_remat_size=base_config.expert_gtp_weight_remat_size, + ) + base_dp = _validated_data_parallel_size( + num_gpus=base_config.num_gpus, + topology=topology_from_config(base_config), + ) + requested_dp = _validated_data_parallel_size( + num_gpus=num_gpus, + topology=topology_from_config(requested_topology), + ) + if args.global_batch_size is not None: gbs_size = args.global_batch_size - elif num_gpus != base_config.num_gpus: - gbs_size = int(base_config.gbs_scaling_factor * num_gpus) + elif requested_dp != base_dp: + gbs_size = _weak_scaled_global_batch_size( + base_gbs=base_config.global_batch_size, + base_data_parallel=base_dp, + data_parallel=requested_dp, + ) else: gbs_size = base_config.global_batch_size diff --git a/scripts/training/README.md b/scripts/training/README.md index 158c1c0c2e2..ce024dec311 100644 --- a/scripts/training/README.md +++ b/scripts/training/README.md @@ -52,9 +52,9 @@ requested adapter scheme. ### Benchmark recipe The training launcher can run exact exported recipes from `src/megatron/bridge/perf_recipes`, including text -pretraining, text SFT/PEFT, Qwen-VL pretraining, and Wan pretraining. The total allocation must match the GPU count -encoded by the recipe name; the user selects the node shape, and the Slurm partition must provide the requested -hardware: +pretraining, text SFT/PEFT, Qwen-VL pretraining, and Wan pretraining. The GPU count encoded by the recipe name is its +canonical tuned allocation. The user selects the node shape and may weak-scale to a different compatible world size; +the Slurm partition must provide the requested hardware: `benchmark` is the unified runner's user-facing term. The existing `perf_recipes` package and `scripts/performance/` compatibility paths retain their legacy names. @@ -75,6 +75,11 @@ configuration. Benchmark recipes retain their selected dataset type, so `--datas Recipe environment defaults are installed before the launcher enters the training stack; values explicitly set by the shell or Slurm environment retain precedence. +When the requested world size differs from the canonical count, the runner preserves the recipe's samples per GPU by +scaling `train.global_batch_size` proportionally. The scaled batch size must be an integer, and the requested world size +must be divisible by the resolved TP × PP × CP topology. Pass `--global_batch_size` or a trailing +`train.global_batch_size=...` override to choose an explicit batch size instead of automatic weak scaling. + The launcher does not infer offline mode, cluster-specific CPU/NUMA binding, Slurm segment sizing, or NCCL fabric settings from the recipe name. Supply those deployment settings explicitly through the target cluster's launcher configuration, repeated `--srun-arg` options, or exported values forwarded with `--env NAME`. @@ -279,9 +284,10 @@ uv run python scripts/training/run_recipe.py \ --dry-run logger.save_config_filepath=/tmp/config.yaml ``` -For a benchmark dry run, use the complete flat recipe name and omit `--dataset`. The rank-local dry run discovers the -recipe and validates the final config against the total GPU count encoded in that name without -requiring a live allocation; the submission dry run additionally validates `--nodes` and `--gpus-per-node`. +For a benchmark dry run, use the complete flat recipe name and omit `--dataset`. Without a distributed world-size +environment, the rank-local dry run validates the canonical allocation encoded in the recipe name. If `WORLD_SIZE` or +`SLURM_NTASKS` is set, it validates and weak-scales against that requested size instead. The submission dry run validates +the launcher arguments while leaving topology validation to the rank-local resolved config. ## Rank-local entry point diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index cdcb90e7fdb..bddb4e853e7 100755 --- a/scripts/training/run_recipe.py +++ b/scripts/training/run_recipe.py @@ -77,10 +77,22 @@ from typing import TYPE_CHECKING, Literal +logger = logging.getLogger(__name__) + + SCRIPT_DIR = Path(__file__).resolve().parent if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) - +COMMON_SCRIPT_DIR = SCRIPT_DIR.parent / "common" +if str(COMMON_SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(COMMON_SCRIPT_DIR)) + +from benchmark_parallelism import ( # noqa: E402 + ParallelTopology, + data_parallel_size, + topology_from_config, + weak_scaled_global_batch_size, +) from recipe_metadata import ( # noqa: E402 BenchmarkRecipeMetadata, infer_recipe_mode, @@ -304,21 +316,61 @@ def _current_world_size() -> int | None: return None -def _validate_benchmark_world_size(metadata: BenchmarkRecipeMetadata, *, dryrun: bool) -> None: - """Require the recipe's total GPU count for an executable run.""" - if dryrun: - return +def _benchmark_world_size(metadata: BenchmarkRecipeMetadata, *, dryrun: bool) -> int: + """Resolve the requested benchmark world size from the active launcher.""" world_size = _current_world_size() if world_size is None: + if dryrun: + return metadata.num_gpus raise ValueError( "Benchmark recipes require an existing distributed environment with the world size set by torchrun " "or Slurm." ) - if world_size != metadata.num_gpus: - raise ValueError( - f"Benchmark recipe requires exactly {metadata.num_gpus} GPUs, but the distributed world size is " - f"{world_size}. Select a recipe matching the allocation." - ) + return world_size + + +def _config_override_is_explicit(cli_overrides: list[str], field_name: str) -> bool: + """Return whether CLI overrides explicitly own a config field or its parent.""" + parent_name = field_name.split(".", maxsplit=1)[0] + for override in cli_overrides: + override_name = override.lstrip("+~").split("=", maxsplit=1)[0] + if override_name in {field_name, parent_name}: + return True + return False + + +def _apply_benchmark_weak_scaling( + recipe: ConfigContainer, + metadata: BenchmarkRecipeMetadata, + cli_overrides: list[str], + *, + canonical_topology: ParallelTopology, + canonical_global_batch_size: int, + world_size: int, +) -> ConfigContainer: + """Validate topology and preserve samples per DP rank.""" + model = getattr(recipe, "model", None) + canonical_dp = data_parallel_size(num_gpus=metadata.num_gpus, topology=canonical_topology) + requested_dp = data_parallel_size(num_gpus=world_size, topology=topology_from_config(model)) + + if _config_override_is_explicit(cli_overrides, "train.global_batch_size") or requested_dp == canonical_dp: + return recipe + + train = getattr(recipe, "train", None) + scaled_gbs = weak_scaled_global_batch_size( + base_gbs=canonical_global_batch_size, + base_data_parallel=canonical_dp, + data_parallel=requested_dp, + ) + train.global_batch_size = scaled_gbs + logger.info( + "Weak scaled benchmark global batch size from %d to %d for DP %d -> %d.", + canonical_global_batch_size, + scaled_gbs, + canonical_dp, + requested_dp, + ) + return recipe def _apply_benchmark_runtime_defaults( @@ -418,14 +470,28 @@ def main(argv: list[str] | None = None) -> None: recipe = _apply_benchmark_dataset_defaults(recipe, benchmark_metadata) recipe = _apply_dataset(recipe, args) recipe = apply_determinism(recipe, deterministic=args.deterministic) + benchmark_canonical_topology = None + benchmark_canonical_global_batch_size = None + if benchmark_metadata is not None: + benchmark_canonical_topology = topology_from_config(getattr(recipe, "model", None)) + benchmark_canonical_global_batch_size = getattr(getattr(recipe, "train", None), "global_batch_size", 1) recipe = apply_cli_overrides(recipe, cli_overrides) recipe = sync_model_pipeline_layout(recipe, cli_overrides=cli_overrides) + benchmark_world_size = None if benchmark_metadata is not None: recipe = _apply_benchmark_runtime_defaults(recipe, benchmark_metadata, cli_overrides) + benchmark_world_size = _benchmark_world_size(benchmark_metadata, dryrun=args.dryrun) + recipe = _apply_benchmark_weak_scaling( + recipe, + benchmark_metadata, + cli_overrides, + canonical_topology=benchmark_canonical_topology, + canonical_global_batch_size=benchmark_canonical_global_batch_size, + world_size=benchmark_world_size, + ) configuration_mode = _train_mode(args.mode) if benchmark_metadata is not None: - _validate_benchmark_world_size(benchmark_metadata, dryrun=args.dryrun) recipe = bootstrap_recipe_environment( recipe, script_path=str(Path(__file__).resolve()), @@ -449,7 +515,7 @@ def main(argv: list[str] | None = None) -> None: mode=execution_mode, step_func=forward_step, dryrun=args.dryrun, - dryrun_world_size=benchmark_metadata.num_gpus if benchmark_metadata is not None else None, + dryrun_world_size=benchmark_world_size, dump_environment=args.dump_env, ) diff --git a/scripts/training/setup_experiment.py b/scripts/training/setup_experiment.py index c5ef6d000a4..ab0a0620f56 100755 --- a/scripts/training/setup_experiment.py +++ b/scripts/training/setup_experiment.py @@ -199,9 +199,11 @@ def _validate_args( if benchmark_metadata is not None: requested_gpus = args.nodes * args.gpus_per_node if requested_gpus != benchmark_metadata.num_gpus: - raise ValueError( - f"Benchmark recipe requires exactly {benchmark_metadata.num_gpus} GPUs, but --nodes and " - f"--gpus-per-node request {requested_gpus}." + logger.info( + "Weak scaling benchmark recipe '%s' from %d to %d GPUs.", + benchmark_metadata.recipe_name, + benchmark_metadata.num_gpus, + requested_gpus, ) diff --git a/tests/unit_tests/recipes/test_override_precedence.py b/tests/unit_tests/recipes/test_override_precedence.py index 902a522644e..b88bf3a2c48 100644 --- a/tests/unit_tests/recipes/test_override_precedence.py +++ b/tests/unit_tests/recipes/test_override_precedence.py @@ -235,6 +235,121 @@ def test_E_argparse_recompute_num_layers_switches_to_full(self): class TestGbsPrecedence: + @staticmethod + def _nemotronh_64gpu_base_recipe(): + return SimpleNamespace( + optimizer=SimpleNamespace(optimizer="adam"), + model=SimpleNamespace( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + ), + train=SimpleNamespace(global_batch_size=192), + comm_overlap=None, + ) + + def test_nemotronh_b300_64gpu_recipe_weak_scales_to_8_gpus(self, monkeypatch): + """Nemotron-H preserves samples per GPU when its canonical recipe is reused.""" + from utils import overrides as override_utils + + base_config = SimpleNamespace( + num_gpus=64, + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=None, + gtp_weight_remat_size=1, + expert_gtp_weight_remat_size=1, + global_batch_size=192, + gbs_scaling_factor=3, + ) + monkeypatch.setattr(override_utils, "get_workload_base_config", lambda *_args, **_kwargs: base_config) + + recipe = override_utils.set_post_overrides( + self._nemotronh_64gpu_base_recipe(), + model_family_name="nemotronh", + model_recipe_name="nemotronh_56b", + gpu="b300", + num_gpus=8, + compute_dtype="fp8_cs", + task="pretrain", + ) + + assert recipe.train.global_batch_size == 24 + + def test_weak_scaling_uses_effective_data_parallel_ratio(self, monkeypatch): + """Parallelism overrides preserve samples per data-parallel rank.""" + from utils import overrides as override_utils + + base_config = SimpleNamespace( + num_gpus=64, + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=None, + global_batch_size=192, + ) + monkeypatch.setattr(override_utils, "get_workload_base_config", lambda *_args, **_kwargs: base_config) + recipe = self._nemotronh_64gpu_base_recipe() + recipe.model.tensor_model_parallel_size = 1 + + recipe = override_utils.set_post_overrides( + recipe, + model_family_name="nemotronh", + model_recipe_name="nemotronh_56b", + gpu="b300", + num_gpus=8, + compute_dtype="fp8_cs", + task="pretrain", + ) + + # Canonical DP is 32 and requested DP is 8, so 192 * 8 / 32 = 48. + assert recipe.train.global_batch_size == 48 + + def test_perf_fallback_rejects_world_size_incompatible_with_expert_grid(self, monkeypatch): + """The flat fallback rejects GPT-OSS layouts before MCore initialization.""" + from utils import overrides as override_utils + + base_config = SimpleNamespace( + num_gpus=64, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=4, + context_parallel_size=1, + expert_model_parallel_size=8, + expert_tensor_parallel_size=1, + gtp_weight_remat_size=1, + expert_gtp_weight_remat_size=1, + global_batch_size=1280, + ) + monkeypatch.setattr(override_utils, "get_workload_base_config", lambda *_args, **_kwargs: base_config) + recipe = SimpleNamespace( + optimizer=SimpleNamespace(optimizer="adam"), + model=SimpleNamespace( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=4, + context_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + expert_model_parallel_size=8, + expert_tensor_parallel_size=1, + ), + train=SimpleNamespace(global_batch_size=1280), + comm_overlap=None, + ) + + with pytest.raises(ValueError, match=r"expert.*8 GPUs.*ETP \* EP \* PP.*1 \* 8 \* 4 = 32"): + override_utils.set_post_overrides( + recipe, + model_family_name="gpt_oss", + model_recipe_name="gpt_oss_120b", + gpu="h100", + num_gpus=8, + compute_dtype="bf16", + task="pretrain", + ) + def test_F_autoscale_fires_when_no_one_sets(self): """When neither Hydra nor argparse sets GBS and num_gpus differs from the workload default, set_post_overrides should rescale. This is the @@ -242,7 +357,7 @@ def test_F_autoscale_fires_when_no_one_sets(self): recipe = _fresh_recipe() # Canonical flat workload default for GB200 is GBS=4096 at 256 GPUs. At 64 GPUs, # gbs_scaling_factor * 64 should be applied. - recipe = _apply(recipe, num_gpus=64) + recipe = _apply(recipe, args_overrides={"expert_model_parallel_size": 16}, num_gpus=64) assert recipe.train.global_batch_size != 4096, "GBS auto-scale did not fire for num_gpus=64 (default=256)" def test_G_hydra_overrides_autoscale(self): @@ -252,6 +367,7 @@ def test_G_hydra_overrides_autoscale(self): recipe = _apply( recipe, cli_overrides=["train.global_batch_size=128"], + args_overrides={"expert_model_parallel_size": 16}, num_gpus=64, ) assert recipe.train.global_batch_size == 128 @@ -261,7 +377,7 @@ def test_H_argparse_overrides_autoscale(self): recipe = _fresh_recipe() recipe = _apply( recipe, - args_overrides={"global_batch_size": 256}, + args_overrides={"global_batch_size": 256, "expert_model_parallel_size": 16}, num_gpus=64, ) assert recipe.train.global_batch_size == 256 diff --git a/tests/unit_tests/scripts/performance/test_recipe_environment.py b/tests/unit_tests/scripts/performance/test_recipe_environment.py index 312495514a0..4a1d560d8d7 100644 --- a/tests/unit_tests/scripts/performance/test_recipe_environment.py +++ b/tests/unit_tests/scripts/performance/test_recipe_environment.py @@ -267,6 +267,185 @@ def test_flat_environment_preparation_defaults_missing_recipe_environment(model_ assert "Performance might be degraded" in caplog.text +def test_missing_exact_gpu_recipe_uses_canonical_family_recipe(monkeypatch): + """Weak scaling falls back to the canonical recipe for the same workload.""" + canonical_name = "nemotronh_56b_pretrain_64gpu_b300_fp8cs_config" + args = SimpleNamespace( + model_family_name="nemotronh", + model_recipe_name="nemotronh_56b", + task="pretrain", + num_gpus=8, + gpu="b300", + compute_dtype="fp8_cs", + config_variant=None, + global_batch_size=None, + ) + canonical_recipe = SimpleNamespace(train=SimpleNamespace(global_batch_size=192)) + + def missing_exact(**_kwargs): + raise utils.PerfRecipeNotFoundError("missing exact 8-GPU recipe") + + monkeypatch.setattr(run_script, "get_perf_recipe_by_name", missing_exact) + monkeypatch.setattr(utils, "flat_perf_recipe_names", lambda: (canonical_name,)) + monkeypatch.setattr( + utils, + "find_perf_recipe", + lambda name: (lambda: canonical_recipe) if name == canonical_name else None, + ) + monkeypatch.setattr(run_script, "_apply_perf_recipe_overrides", lambda recipe, _overrides, _args: recipe) + + def apply_weak_scaling(recipe, **kwargs): + assert kwargs["num_gpus"] == 8 + recipe.train.global_batch_size = 24 + return recipe + + override_utils = types.ModuleType("utils.overrides") + override_utils.set_post_overrides = apply_weak_scaling + monkeypatch.setitem(sys.modules, "utils.overrides", override_utils) + + result = run_script._prepare_perf_recipe(args, []) + + assert result is canonical_recipe + assert result.train.global_batch_size == 24 + + +def test_exact_gpu_recipe_takes_precedence_over_canonical_fallback(monkeypatch): + """A tuned exact-count recipe must not be replaced by weak scaling.""" + exact_recipe = object() + args = SimpleNamespace( + model_family_name="test", + model_recipe_name="test_model", + task="pretrain", + num_gpus=8, + gpu="b300", + compute_dtype="bf16", + config_variant=None, + global_batch_size=None, + ) + monkeypatch.setattr(run_script, "get_perf_recipe_by_name", lambda **_kwargs: exact_recipe) + monkeypatch.setattr( + run_script, + "get_perf_optimized_recipe", + lambda **_kwargs: pytest.fail("canonical fallback must not run for an exact recipe"), + ) + monkeypatch.setattr(run_script, "_apply_perf_recipe_overrides", lambda recipe, _overrides, _args: recipe) + + result = run_script._prepare_perf_recipe(args, []) + + assert result is exact_recipe + + +def test_exact_recipe_construction_error_is_not_treated_as_missing(monkeypatch): + args = SimpleNamespace( + model_family_name="test", + model_recipe_name="test_model", + task="pretrain", + num_gpus=8, + gpu="b300", + compute_dtype="bf16", + config_variant=None, + global_batch_size=None, + ) + + def construction_error(**_kwargs): + raise RuntimeError("recipe builder failed") + + monkeypatch.setattr(run_script, "get_perf_recipe_by_name", construction_error) + monkeypatch.setattr( + run_script, + "get_perf_optimized_recipe", + lambda **_kwargs: pytest.fail("construction failures must not invoke canonical fallback"), + ) + + with pytest.raises(RuntimeError, match="recipe builder failed"): + run_script._prepare_perf_recipe(args, []) + + +def test_missing_perf_recipe_family_reports_requested_workload(monkeypatch): + args = SimpleNamespace( + model_family_name="nemotronh", + model_recipe_name="nemotronh_56b", + task="pretrain", + num_gpus=8, + gpu="b300", + compute_dtype="fp8_cs", + config_variant=None, + global_batch_size=None, + ) + + def missing_exact(**_kwargs): + raise utils.PerfRecipeNotFoundError("missing exact 8-GPU recipe") + + def missing_family(**_kwargs): + raise ValueError("No flat perf recipe found for nemotronh_56b/pretrain/b300/fp8_cs/default") + + monkeypatch.setattr(run_script, "get_perf_recipe_by_name", missing_exact) + monkeypatch.setattr(run_script, "get_perf_optimized_recipe", missing_family) + + with pytest.raises( + ValueError, + match=r"nemotronh_56b/pretrain/b300/fp8_cs/default", + ): + run_script._prepare_perf_recipe(args, []) + + +def test_weak_scaling_validates_world_size_and_integer_gbs(): + assert ( + utils._data_parallel_size( + num_gpus=8, + tensor_parallel=2, + pipeline_parallel=1, + context_parallel=1, + ) + == 4 + ) + assert utils._weak_scaled_global_batch_size(base_gbs=192, base_data_parallel=32, data_parallel=4) == 24 + + with pytest.raises(ValueError, match=r"7 GPUs.*TP \* PP \* CP.*2 \* 1 \* 1"): + utils._data_parallel_size( + num_gpus=7, + tensor_parallel=2, + pipeline_parallel=1, + context_parallel=1, + ) + with pytest.raises(ValueError, match=r"does not produce an integer global batch size"): + utils._weak_scaled_global_batch_size(base_gbs=190, base_data_parallel=32, data_parallel=4) + + +def test_exp_name_rejects_fractional_data_parallel_weak_scaling(monkeypatch): + base_config = utils.WorkloadBaseConfig( + num_gpus=64, + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=None, + global_batch_size=190, + ) + monkeypatch.setattr(utils, "get_workload_base_config", lambda *_args, **_kwargs: base_config) + args = SimpleNamespace( + num_gpus=8, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=None, + context_parallel_size=None, + virtual_pipeline_model_parallel_size=-1, + expert_model_parallel_size=None, + expert_tensor_parallel_size=None, + micro_batch_size=None, + global_batch_size=None, + ) + + with pytest.raises(ValueError, match="does not produce an integer global batch size"): + utils.get_exp_name_config( + args, + model_family_name="nemotronh", + model_recipe_name="nemotronh_56b", + gpu="b300", + compute_dtype="fp8_cs", + task="pretrain", + ) + + def test_flat_environment_preparation_applies_cli_overrides(monkeypatch): """Hydra env overrides must be reflected in the bootstrap recipe.""" args = SimpleNamespace( diff --git a/tests/unit_tests/scripts/training/test_run_recipe.py b/tests/unit_tests/scripts/training/test_run_recipe.py index 0533508f408..919b6bb036f 100644 --- a/tests/unit_tests/scripts/training/test_run_recipe.py +++ b/tests/unit_tests/scripts/training/test_run_recipe.py @@ -501,25 +501,179 @@ def test_benchmark_dry_run_accepts_config_overrides(monkeypatch): ) -def test_benchmark_recipe_rejects_noncanonical_world_size(monkeypatch): +def test_benchmark_recipe_weak_scales_noncanonical_world_size(monkeypatch): module, handles = _load_module() monkeypatch.setenv("WORLD_SIZE", "8") monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") + config = SimpleNamespace( + optimizer=SimpleNamespace(optimizer="adam", use_precision_aware_optimizer=False), + model=SimpleNamespace( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ), + train=SimpleNamespace(global_batch_size=192), + ) + handles.recipe_runner.load_recipe.return_value = config + + module.main( + [ + "--recipe", + "nemotronh_56b_pretrain_64gpu_b300_fp8cs_config", + "--mode", + "pretrain", + ] + ) + + assert config.train.global_batch_size == 24 + handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() + + +def test_benchmark_recipe_weak_scales_by_data_parallel_ratio_after_topology_override(monkeypatch): + module, handles = _load_module() + monkeypatch.setenv("WORLD_SIZE", "8") + config = SimpleNamespace( + optimizer=SimpleNamespace(optimizer="adam", use_precision_aware_optimizer=False), + model=SimpleNamespace( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=None, + ), + train=SimpleNamespace(global_batch_size=192), + ) + handles.recipe_runner.load_recipe.return_value = config + + def apply_overrides(recipe, overrides): + assert "model.tensor_model_parallel_size=1" in overrides + recipe.model.tensor_model_parallel_size = 1 + return recipe + + handles.recipe_runner.apply_cli_overrides.side_effect = apply_overrides + + module.main( + [ + "--recipe", + "nemotronh_56b_pretrain_64gpu_b300_fp8cs_config", + "--mode", + "pretrain", + "--tensor_model_parallel_size", + "1", + ] + ) + + # Canonical DP is 64 / TP2 = 32; requested DP is 8 / TP1 = 8. + assert config.train.global_batch_size == 48 + + +def test_benchmark_recipe_rejects_world_size_incompatible_with_expert_grid(monkeypatch): + module, handles = _load_module() + monkeypatch.setenv("WORLD_SIZE", "8") handles.recipe_runner.load_recipe.return_value = SimpleNamespace( - optimizer=SimpleNamespace(optimizer="adam", use_precision_aware_optimizer=False) + optimizer=SimpleNamespace(optimizer="adam", use_precision_aware_optimizer=False), + model=SimpleNamespace( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=4, + context_parallel_size=1, + expert_model_parallel_size=8, + expert_tensor_parallel_size=1, + ), + train=SimpleNamespace(global_batch_size=1280), ) - with pytest.raises(ValueError, match="requires exactly 16 GPUs"): + with pytest.raises(ValueError, match=r"expert.*8 GPUs.*ETP \* EP \* PP.*1 \* 8 \* 4 = 32"): module.main( [ "--recipe", - "qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config", + "gpt_oss_120b_pretrain_64gpu_h100_bf16_config", "--mode", "pretrain", ] ) - handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() + +def test_benchmark_recipe_explicit_global_batch_size_disables_weak_scaling(monkeypatch): + module, handles = _load_module() + monkeypatch.setenv("WORLD_SIZE", "8") + config = SimpleNamespace( + optimizer=SimpleNamespace(optimizer="adam", use_precision_aware_optimizer=False), + model=SimpleNamespace( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ), + train=SimpleNamespace(global_batch_size=192), + ) + handles.recipe_runner.load_recipe.return_value = config + + def apply_overrides(recipe, overrides): + assert "train.global_batch_size=7" in overrides + recipe.train.global_batch_size = 7 + return recipe + + handles.recipe_runner.apply_cli_overrides.side_effect = apply_overrides + + module.main( + [ + "--recipe", + "nemotronh_56b_pretrain_64gpu_b300_fp8cs_config", + "--mode", + "pretrain", + "--global_batch_size", + "7", + ] + ) + + assert config.train.global_batch_size == 7 + + +def test_benchmark_recipe_rejects_world_size_incompatible_with_model_parallelism(monkeypatch): + module, handles = _load_module() + monkeypatch.setenv("WORLD_SIZE", "7") + handles.recipe_runner.load_recipe.return_value = SimpleNamespace( + optimizer=SimpleNamespace(optimizer="adam", use_precision_aware_optimizer=False), + model=SimpleNamespace( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ), + train=SimpleNamespace(global_batch_size=192), + ) + + with pytest.raises(ValueError, match=r"7 GPUs.*TP \* PP \* CP.*2 \* 1 \* 1"): + module.main( + [ + "--recipe", + "nemotronh_56b_pretrain_64gpu_b300_fp8cs_config", + "--mode", + "pretrain", + ] + ) + + +def test_benchmark_recipe_rejects_fractional_weak_scaled_global_batch_size(monkeypatch): + module, handles = _load_module() + monkeypatch.setenv("WORLD_SIZE", "8") + handles.recipe_runner.load_recipe.return_value = SimpleNamespace( + optimizer=SimpleNamespace(optimizer="adam", use_precision_aware_optimizer=False), + model=SimpleNamespace( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ), + train=SimpleNamespace(global_batch_size=190), + ) + + with pytest.raises(ValueError, match="does not produce an integer global batch size"): + module.main( + [ + "--recipe", + "nemotronh_56b_pretrain_64gpu_b300_fp8cs_config", + "--mode", + "pretrain", + ] + ) def test_benchmark_recipe_accepts_user_selected_per_node_topology(monkeypatch): diff --git a/tests/unit_tests/scripts/training/test_setup_experiment.py b/tests/unit_tests/scripts/training/test_setup_experiment.py index e7b5b2d2e0b..eab74bab157 100644 --- a/tests/unit_tests/scripts/training/test_setup_experiment.py +++ b/tests/unit_tests/scripts/training/test_setup_experiment.py @@ -116,7 +116,7 @@ def test_library_resolved_recipe_does_not_enable_benchmark_executor(): assert module.selected_benchmark_recipe(training_args) is None -def test_benchmark_recipe_validates_total_submission_size(): +def test_benchmark_recipe_accepts_weak_scaling_submission_size(): module = _load_setup_experiment_module() args, training_args = module.parse_args( [ @@ -135,8 +135,7 @@ def test_benchmark_recipe_validates_total_submission_size(): ] ) - with pytest.raises(ValueError, match="requires exactly 16 GPUs"): - module._validate_args(args, module.selected_benchmark_recipe(training_args)) + module._validate_args(args, module.selected_benchmark_recipe(training_args)) def test_benchmark_recipe_accepts_user_selected_node_shape(): From 677077254f6ff563c114ff25d2f770027ec0be31 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Tue, 4 Aug 2026 13:43:19 -0700 Subject: [PATCH 25/85] chore(deps): update transformer-engine dependency to a new commit hash Updated the transformer-engine dependency in pyproject.toml and uv.lock to use the latest commit (cfda3834) for improved stability and features. Signed-off-by: Ajay Balasa --- pyproject.toml | 4 ++-- uv.lock | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2c3894f59c4..c71058a8bb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,7 @@ override-dependencies = [ "torch; sys_platform == 'never'", "torchvision; sys_platform == 'never'", "triton; sys_platform == 'never'", - "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@e7c550c5f80636cf841a8204b1d6f85a5f3f28b7", + "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@cfda38340769c432c12db3c0df574f8e6ce46c8c", "mlflow>=3.15.0", # To address CVE-2025-15031 "cachetools>=5.0.0", "cryptography>=43.0.0,<47", @@ -168,7 +168,7 @@ requires-dist = [] # above) changes — an out-of-date entry sends uv back to building TE. [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.18.0+e7c550c5" +version = "2.18.0+cfda3834" requires-dist = [ "einops", "importlib-metadata>=1.0", diff --git a/uv.lock b/uv.lock index 1ad0214cd3f..274b41ef2d4 100644 --- a/uv.lock +++ b/uv.lock @@ -24,7 +24,7 @@ overrides = [ { name = "tokenizers", specifier = ">=0.22.0,<0.23" }, { name = "torch", marker = "sys_platform == 'never'" }, { name = "torchvision", marker = "sys_platform == 'never'" }, - { name = "transformer-engine", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=e7c550c5f80636cf841a8204b1d6f85a5f3f28b7" }, + { name = "transformer-engine", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=cfda38340769c432c12db3c0df574f8e6ce46c8c" }, { name = "triton", marker = "sys_platform == 'never'" }, { name = "urllib3", specifier = ">=2.6.3" }, ] @@ -35,7 +35,7 @@ version = "1.0.0+b7643bd" [[manifest.dependency-metadata]] name = "transformer-engine" -version = "2.18.0+e7c550c5" +version = "2.18.0+cfda3834" requires-dist = ["einops", "importlib-metadata>=1.0", "nvdlfw-inspect", "nvidia-cudnn-frontend>=1.25.0", "onnx", "onnxscript", "packaging", "pydantic", "torch>=2.1"] provides-extras = ["core-cu13", "pytorch"] @@ -4863,8 +4863,8 @@ wheels = [ [[package]] name = "transformer-engine" -version = "2.18.0+e7c550c5" -source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=e7c550c5f80636cf841a8204b1d6f85a5f3f28b7#e7c550c5f80636cf841a8204b1d6f85a5f3f28b7" } +version = "2.18.0+cfda3834" +source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=cfda38340769c432c12db3c0df574f8e6ce46c8c#cfda38340769c432c12db3c0df574f8e6ce46c8c" } dependencies = [ { name = "einops" }, { name = "importlib-metadata" }, From 1f2d8e60cef293099ed3cf26a49edf38978d3903 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 5 Aug 2026 06:51:50 +0200 Subject: [PATCH 26/85] cp: `feat (recipe): add Nemotron 3 Ultra MXFP8 GB200 pretraining perf-recipe (5287)` into `r0.6.0` (#5309) Signed-off-by: Ravi Ghadia Signed-off-by: Malay Nagda Signed-off-by: NeMo Bot Co-authored-by: Ravi Ghadia <40660742+ghadiaravi13@users.noreply.github.com> Co-authored-by: Ravi Ghadia Co-authored-by: Malay Nagda --- .../bridge/perf_recipes/nemotronh/__init__.py | 1 + .../perf_recipes/nemotronh/gb200/nemotronh.py | 90 +++++++++++++++++++ .../recipes/test_perf_recipe_environment.py | 16 +++- 3 files changed, 106 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py index 05b9864b963..785860600bd 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py @@ -29,6 +29,7 @@ nemotron_3_super_pretrain_64gpu_gb200_bf16_config, nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config, nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config, + nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config, nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config, ) from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import ( diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index d7d6255c846..c9452d46839 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -19,12 +19,15 @@ ConfigContainer, _apply_nemotron_3_nano_perf_defaults, _apply_nemotron_3_super_perf_defaults, + _apply_nemotron_3_ultra_fsdp_hsdp, + _apply_nemotron_3_ultra_perf_defaults, _benchmark_common, _nemotron_3_super_nvfp4_precision, _perf_precision, load_quantization_recipe, nemotron_3_nano_pretrain_config, nemotron_3_super_pretrain_config, + nemotron_3_ultra_pretrain_config, nemotronh_56b_pretrain_config, ) from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules @@ -212,6 +215,93 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer: return cfg +def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: + """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP). + + TP2 + SP (due to smaller GB200 HBM) / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed + precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + + fine-grained activation offload of the expert MLP, MTP=2. + """ + + num_gpus = 256 + expert_model_parallel_size = 64 + global_batch_size = 256 + hybrid_ep_ranks_per_nvlink_domain = 64 + + cfg = nemotron_3_ultra_pretrain_config() + cfg.mixed_precision = _perf_precision("fp8_mx") + + _apply_nemotron_3_ultra_perf_defaults(cfg) + + # Apply HSDP / FSDP dtype overrides last so they win over the generic defaults. + _apply_nemotron_3_ultra_fsdp_hsdp(cfg, num_gpus=num_gpus) + + # Parallelism + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.sequence_parallel = True + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = None + cfg.model.seq_length = 8192 + + # Only tensors larger than 350M elements are offloaded, which + # approximates offloading the moe_act (pre-activation input) for seq 8192/2 (due to SP) / MBS 1. + cfg.model.min_offloaded_tensor_size = 350_000_000 + + # MXFP8 requires router padding for quantization. + cfg.model.moe_router_padding_for_quantization = True + + # GPU-count specific overrides of the canonical (256-GPU / EP64) defaults. + cfg.model.expert_model_parallel_size = expert_model_parallel_size + cfg.train.global_batch_size = global_batch_size + + # Fine-grained activation offloading. Requires NVTE_CPU_OFFLOAD_V1=1 in the + # launch environment (set in this recipe's env_vars). + # NOTE: also requires setting the min_offloaded_tensor_size to selectively offload moe_act of the fused_group_mlp, to avoid CPU OOM issues + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["fused_group_mlp"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 1 + + # Selective recompute of the MoE activation + # recomputes the activation output of the MoE expert MLP, while FC1 output (activation input) is saved and offloaded to cpu + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["moe_act"] + + # Keep process settings next to the recipe so users can see the exact benchmark environment. + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + # CUDA stream scheduling for this model and parallel layout. + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + # CUDA graph and allocator behavior for this recipe. + "NCCL_GRAPH_REGISTER": 0, + # TODO: graph_capture_record_stream_reuse might be potentially useful + # when enabling CG, because it allows the freed-up memory buffers of the offloaded tensors + # to be reused during the CG capture, thus keeping the peak memory usage lower. + "PYTORCH_CUDA_ALLOC_CONF": ("expandable_segments:True,graph_capture_record_stream_reuse:True"), + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + # NCCL user-buffer and launch settings. + "NCCL_NVLS_ENABLE": 0, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": hybrid_ep_ranks_per_nvlink_domain, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + # Transformer Engine overlap settings for this model. + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid + # offloading weights; + "NVTE_CPU_OFFLOAD_V1": 1, + # Enable TE's CuteDSL fused grouped MLP kernel (sm100+). Required by the + # op fuser + fused weighted squared-ReLU with moe_act activation recompute + # (ScaledSReLU(activation_recompute_in_mlp=True) only runs on this path). + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + } + return cfg + + def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: """Nemotron 3 Nano pretrain: 8× GB200, BF16.""" cfg = nemotron_3_nano_pretrain_config() diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index 6f7c7771c81..ffd7803808e 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -72,6 +72,15 @@ def _function(path: Path, function_name: str) -> ast.FunctionDef: def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int | float | bool]: """Read the literal env mapping written in a flat recipe builder.""" function = _function(path, function_name) + local_constants = { + node.targets[0].id: node.value.value + for node in function.body + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and isinstance(node.value, ast.Constant) + and isinstance(node.value.value, (str, int, float, bool)) + } assignments = [ node for node in function.body @@ -93,7 +102,12 @@ def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int assert isinstance(value, ast.Name) and value.id == "COMMON_PERF_ENV_VARS" common_expansions += 1 continue - result[ast.literal_eval(key)] = ast.literal_eval(value) + if isinstance(value, ast.Name): + assert value.id in local_constants + env_value = local_constants[value.id] + else: + env_value = ast.literal_eval(value) + result[ast.literal_eval(key)] = env_value assert common_expansions == 1 return result From 9d4eef48ed6255bad3a370a628f62e484c506e75 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 5 Aug 2026 06:52:02 +0200 Subject: [PATCH 27/85] cp: `fix(training): count expert parameter norms with TP and EP (4998)` into `r0.6.0` (#5177) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: Chen Cui Co-authored-by: meatybobby --- .../bridge/training/utils/train_utils.py | 34 +++- .../training/utils/test_train_utils.py | 184 ++++++++++++++++-- 2 files changed, 203 insertions(+), 15 deletions(-) diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index 404e40c1b24..d4a4518f4e6 100644 --- a/src/megatron/bridge/training/utils/train_utils.py +++ b/src/megatron/bridge/training/utils/train_utils.py @@ -325,23 +325,31 @@ def calc_params_l2_norm( params_data = [] moe_params_data = [] sharded_params_data = [] + sharded_moe_params_data = [] data_parallel_group = None + pg_collection = get_pg_collection(model) for model_chunk in model: for param in model_chunk.parameters(): data_parallel_group = get_data_parallel_group_if_dtensor(param, data_parallel_group) - is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param) + # MCore uses allreduce=False to mark parameters that use expert-parallel process groups. + uses_expert_parallel_groups = not getattr(param, "allreduce", True) + is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate( + param, + tp_group=pg_collection.tp, + expert_tp_group=pg_collection.expt_tp, + ) if not is_not_tp_duplicate: continue assert is_not_tp_duplicate - if not getattr(param, "allreduce", True): + if uses_expert_parallel_groups: assert param_is_not_shared(param) param = to_local_if_dtensor(param) if model_config.bf16: if not force_create_fp32_copy and hasattr(param, "main_param"): if getattr(param, "main_param_sharded", False): if param.main_param is not None: - sharded_params_data.append(param.main_param) + sharded_moe_params_data.append(param.main_param) else: main_param = param.main_param moe_params_data.append(main_param if main_param is not None else param.data.float()) @@ -401,7 +409,6 @@ def calc_params_l2_norm( sharded_norm_2 = torch.zeros((1,), dtype=torch.float32, device="cuda") # Sum over all DP groups, including CP since distributed optimizer state is # sharded jointly over DP+CP. - pg_collection = get_pg_collection(model) torch.distributed.all_reduce( sharded_norm_2, op=torch.distributed.ReduceOp.SUM, @@ -425,6 +432,25 @@ def calc_params_l2_norm( else: moe_norm_2 = torch.zeros_like(norm_2) + # Distributed optimizer shards expert main parameters over expert DP rather + # than regular DP, so reduce their contribution separately. + if len(sharded_moe_params_data) > 0: + sharded_moe_norm, _ = multi_tensor_applier( + multi_tensor_l2norm, + dummy_overflow_buf, + [sharded_moe_params_data], + False, # no per-parameter norm. + ) + sharded_moe_norm_2 = sharded_moe_norm * sharded_moe_norm + else: + sharded_moe_norm_2 = torch.zeros((1,), dtype=torch.float32, device="cuda") + torch.distributed.all_reduce( + sharded_moe_norm_2, + op=torch.distributed.ReduceOp.SUM, + group=pg_collection.expt_dp, + ) + moe_norm_2 += sharded_moe_norm_2 + # Reduce norm across model parallel groups (dense and expert). # Dense params should sum across all model-parallel GPUs (tensor + pipeline). dense_reduce_group = pg_collection.mp diff --git a/tests/unit_tests/training/utils/test_train_utils.py b/tests/unit_tests/training/utils/test_train_utils.py index 68de1bf9785..0cb17a5ac6a 100644 --- a/tests/unit_tests/training/utils/test_train_utils.py +++ b/tests/unit_tests/training/utils/test_train_utils.py @@ -2508,7 +2508,10 @@ class _PG: def __init__(self): # Minimal set of groups used by calc_params_l2_norm self.dp_cp = object() + self.expt_dp = object() + self.expt_tp = object() self.mp = object() + self.tp = object() self.tp_ep_pp = object() self.pp = object() @@ -2519,11 +2522,13 @@ def size(self_inner): self.dp = _DP() + pg_collection = _PG() monkeypatch.setattr( "megatron.bridge.training.utils.train_utils.get_pg_collection", - lambda model: _PG(), + lambda model: pg_collection, raising=True, ) + return pg_collection @pytest.fixture def simple_model(self): @@ -2730,8 +2735,9 @@ def test_bf16_mode_with_sharded_main_param( mock_is_not_tp_dup, mock_get_dp_group_if_dtensor, mock_model_config_bf16, + _patch_pg_collection, ): - """Test calc_params_l2_norm with sharded main params (distributed optimizer).""" + """Test dense sharded main params reduce over ordinary DP/CP.""" model = torch.nn.Linear(5, 5, bias=False, dtype=torch.bfloat16).cuda() # Setup mocks @@ -2748,10 +2754,70 @@ def test_bf16_mode_with_sharded_main_param( result = calc_params_l2_norm(model, mock_model_config_bf16, force_create_fp32_copy=False) - # Should use sharded params path and call all_reduce + dense_sharded_reduce = next( + call for call in mock_all_reduce.call_args_list if call.kwargs["group"] is _patch_pg_collection.dp_cp + ) + expert_sharded_reduce = next( + call for call in mock_all_reduce.call_args_list if call.kwargs["group"] is _patch_pg_collection.expt_dp + ) + assert dense_sharded_reduce.args[0].item() == pytest.approx(13.0) + assert expert_sharded_reduce.args[0].item() == pytest.approx(0.0) assert isinstance(result, float) assert result > 0 + def test_bf16_sharded_dense_and_expert_norm_uses_matching_dp_groups( + self, + monkeypatch, + mock_model_config_bf16, + _patch_pg_collection, + ): + """Dense and expert optimizer shards contribute through their matching DP groups.""" + + class _DenseAndExpertModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.dense = torch.nn.Parameter(torch.ones(1, dtype=torch.bfloat16, device="cuda")) + self.expert = torch.nn.Parameter(torch.ones(1, dtype=torch.bfloat16, device="cuda")) + self.expert.allreduce = False + + self.dense.main_param = torch.tensor([2.0], device="cuda") + self.dense.main_param_sharded = True + self.expert.main_param = torch.tensor([3.0], device="cuda") + self.expert.main_param_sharded = True + + def fake_all_reduce(tensor, op, group): + del op + if group is _patch_pg_collection.dp_cp: + tensor.add_(12.0) + elif group is _patch_pg_collection.expt_dp: + tensor.add_(16.0) + elif group is not _patch_pg_collection.mp: + raise AssertionError("unexpected reduction group") + + monkeypatch.setattr( + "megatron.bridge.training.utils.train_utils.get_data_parallel_group_if_dtensor", + lambda param, group: None, + ) + monkeypatch.setattr( + "megatron.bridge.training.utils.train_utils.param_is_not_tensor_parallel_duplicate", + lambda param, **kwargs: True, + ) + monkeypatch.setattr( + "megatron.bridge.training.utils.train_utils.to_local_if_dtensor", + lambda param: param, + ) + monkeypatch.setattr("torch.distributed.get_process_group_ranks", lambda group: [0]) + monkeypatch.setattr("torch.distributed.all_reduce", fake_all_reduce) + + actual_norm = calc_params_l2_norm( + _DenseAndExpertModel(), + mock_model_config_bf16, + force_create_fp32_copy=False, + ) + + # Dense: 2**2 local + 12 remote over DP/CP. Expert: 3**2 local + 16 remote over expert DP. + assert actual_norm == pytest.approx(math.sqrt(41.0)) + @mock.patch("megatron.bridge.training.utils.train_utils.get_data_parallel_group_if_dtensor") @mock.patch("megatron.bridge.training.utils.train_utils.param_is_not_tensor_parallel_duplicate") @mock.patch("megatron.bridge.training.utils.train_utils.to_local_if_dtensor") @@ -2874,13 +2940,14 @@ def test_shared_params( @mock.patch("megatron.bridge.training.utils.train_utils.get_data_parallel_group_if_dtensor") @mock.patch("megatron.bridge.training.utils.train_utils.param_is_not_tensor_parallel_duplicate") - def test_tp_duplicate_params( + def test_duplicate_filter_receives_tp_and_expert_tp_groups( self, mock_is_not_tp_dup, mock_get_dp_group_if_dtensor, mock_model_config_fp32, + _patch_pg_collection, ): - """Test calc_params_l2_norm skips TP duplicate parameters.""" + """Test duplicate filtering receives both tensor-parallel groups.""" model = torch.nn.Linear(5, 5, bias=False).cuda() # Setup mocks @@ -2901,6 +2968,97 @@ def test_tp_duplicate_params( # Should be 0 since all params are TP duplicates assert result == pytest.approx(0.0, abs=1e-5) + mock_is_not_tp_dup.assert_called_once_with( + model.weight, + tp_group=_patch_pg_collection.tp, + expert_tp_group=_patch_pg_collection.expt_tp, + ) + + def test_moe_param_norm_counts_logical_parameter_when_tp_ranks_differ( + self, + monkeypatch, + mock_model_config_fp32, + ): + """Expert parameters use the expert-TP rank when it differs from regular TP.""" + + class _RankGroup: + def __init__(self, rank): + self._rank = rank + + def rank(self): + return self._rank + + class _MixedModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.dense = torch.nn.Parameter(torch.ones(4, device="cuda")) + self.expert = torch.nn.Parameter(torch.ones(4, device="cuda")) + self.expert.allreduce = False + + regular_tp_group = _RankGroup(rank=1) + expert_tp_group = _RankGroup(rank=0) + reduce_group = _RankGroup(rank=0) + pg_collection = SimpleNamespace( + tp=regular_tp_group, + expt_tp=expert_tp_group, + dp_cp=reduce_group, + expt_dp=reduce_group, + mp=reduce_group, + tp_ep_pp=reduce_group, + ) + monkeypatch.setattr( + "megatron.bridge.training.utils.train_utils.get_pg_collection", + lambda model: pg_collection, + ) + + with ( + mock.patch( + "megatron.core.tensor_parallel.layers.get_tensor_model_parallel_rank", + return_value=regular_tp_group.rank(), + ), + mock.patch("torch.distributed.get_process_group_ranks", return_value=[0]), + mock.patch("torch.distributed.all_reduce"), + ): + actual_norm = calc_params_l2_norm(_MixedModel(), mock_model_config_fp32) + + assert actual_norm == pytest.approx(2.0) + + def test_real_mcore_duplicate_filter_honors_expert_tp_group(self): + """Guard the Bridge<->MCore contract for expert-parameter duplicate filtering. + + ``calc_params_l2_norm`` calls MCore's ``param_is_not_tensor_parallel_duplicate`` with an + ``expert_tp_group`` so that ``allreduce=False`` expert parameters de-duplicate over expert + tensor parallel instead of regular TP. Every other test in this class mocks that function, + so they would still pass against an MCore build that dropped the ``expert_tp_group`` kwarg -- + silently undercounting expert parameter norms. Exercise the real function here so a pinned + MCore that regresses below this contract fails loudly (``TypeError`` or wrong result) rather + than degrading the training-log norm in production. + """ + from megatron.core.tensor_parallel.layers import param_is_not_tensor_parallel_duplicate + + assert "expert_tp_group" in inspect.signature(param_is_not_tensor_parallel_duplicate).parameters + + class _RankGroup: + def __init__(self, rank): + self._rank = rank + + def rank(self): + return self._rank + + duplicate_regular_tp = _RankGroup(rank=1) # would be dropped as a regular-TP duplicate + primary_expert_tp = _RankGroup(rank=0) # but is the primary shard on expert TP + + expert_param = SimpleNamespace(allreduce=False) + dense_param = SimpleNamespace(allreduce=True) + + # Expert param must route through expert_tp_group (rank 0 -> kept), not regular TP (rank 1). + assert param_is_not_tensor_parallel_duplicate( + expert_param, tp_group=duplicate_regular_tp, expert_tp_group=primary_expert_tp + ) + # Dense param still de-duplicates against the regular TP group (rank 1 -> dropped). + assert not param_is_not_tensor_parallel_duplicate( + dense_param, tp_group=duplicate_regular_tp, expert_tp_group=primary_expert_tp + ) @mock.patch("megatron.bridge.training.utils.train_utils.calc_dtensor_params_l2_norm") def test_megatron_fsdp_path(self, mock_calc_dtensor_norm, mock_model_config_fp32): @@ -3225,12 +3383,9 @@ def test_moe_params_bf16_with_sharded_main_param( mock_is_not_tp_dup, mock_get_dp_group_if_dtensor, mock_model_config_bf16, + _patch_pg_collection, ): - """Test calc_params_l2_norm with MoE params using sharded main_param (distributed optimizer). - - When MoE params have main_param_sharded=True, they should be added to - sharded_params_data for proper all-reduce across DP groups. - """ + """Test expert sharded main params reduce over expert DP.""" model = torch.nn.Linear(5, 5, bias=False, dtype=torch.bfloat16).cuda() # Setup mocks @@ -3248,7 +3403,14 @@ def test_moe_params_bf16_with_sharded_main_param( result = calc_params_l2_norm(model, mock_model_config_bf16, force_create_fp32_copy=False) - # Should use sharded params path and call all_reduce + dense_sharded_reduce = next( + call for call in mock_all_reduce.call_args_list if call.kwargs["group"] is _patch_pg_collection.dp_cp + ) + expert_sharded_reduce = next( + call for call in mock_all_reduce.call_args_list if call.kwargs["group"] is _patch_pg_collection.expt_dp + ) + assert dense_sharded_reduce.args[0].item() == pytest.approx(0.0) + assert expert_sharded_reduce.args[0].item() == pytest.approx(13.0) assert isinstance(result, float) assert result > 0 From 9182eed122970b1bfcd19c744620f691f3999ad4 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 5 Aug 2026 06:52:17 +0200 Subject: [PATCH 28/85] cp: `fix(model): restore DeepSeek V3 sequence auxiliary loss (5165)` into `r0.6.0` (#5172) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: Chen Cui --- src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py | 3 +++ tests/unit_tests/models/deepseek/test_deepseek_bridges.py | 3 +-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py index 015930a6e62..dbb50f8617e 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py @@ -66,6 +66,9 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.moe_router_pre_softmax = True provider.moe_token_dispatcher_type = "alltoall" provider.moe_router_load_balancing_type = "seq_aux_loss" + # The released V3 configs omit this training-only hyperparameter, but V3 uses a + # small complementary sequence-wise loss alongside expert-bias load balancing. + provider.moe_aux_loss_coeff = getattr(hf_config, "aux_loss_alpha", 0.0001) provider.moe_shared_expert_overlap = True provider.moe_router_score_function = "sigmoid" provider.moe_router_enable_expert_bias = True diff --git a/tests/unit_tests/models/deepseek/test_deepseek_bridges.py b/tests/unit_tests/models/deepseek/test_deepseek_bridges.py index 6b55264753e..68385e3809b 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_bridges.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_bridges.py @@ -92,7 +92,6 @@ def ds_v3_config(self): "tie_word_embeddings": False, "topk_group": 4, "topk_method": "noaux_tc", - "aux_loss_alpha": 0.0001, "torch_dtype": "bfloat16", "transformers_version": "4.33.1", "use_cache": True, @@ -125,7 +124,7 @@ def test_provider_bridge_maps_config(self, mock_pretrained_v3): assert provider.vocab_size == mock_pretrained_v3.config.vocab_size assert provider.layernorm_epsilon == mock_pretrained_v3.config.rms_norm_eps assert provider.rotary_base == mock_pretrained_v3.config.rope_theta - assert provider.moe_aux_loss_coeff == mock_pretrained_v3.config.aux_loss_alpha + assert provider.moe_aux_loss_coeff == 0.0001 # dtype mapping assert provider.bf16 is True assert provider.params_dtype == torch.bfloat16 From e1f83675d84c014b1b423a6aee801a4e23e3f7a1 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 5 Aug 2026 06:53:27 +0200 Subject: [PATCH 29/85] cp: Nemotron 3.5 Lightning checkpoint and naming updates (#5293, #5303) into r0.6.0 (#5300) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- .../card.yaml | 258 +++++++++--------- scripts/performance/dump_perf_configs.py | 12 +- .../bridge/perf_recipes/nemotronh/__init__.py | 12 +- .../perf_recipes/nemotronh/gb200/nemotronh.py | 41 +-- .../perf_recipes/nemotronh/h100/nemotronh.py | 25 +- .../bridge/recipes/nemotronh/__init__.py | 32 +-- .../recipes/nemotronh/gb200/__init__.py | 12 +- .../nemotronh/gb200/nemotron_3_nano.py | 26 +- .../bridge/recipes/nemotronh/h100/__init__.py | 8 +- .../recipes/nemotronh/h100/nemotron_3_nano.py | 49 ++-- .../recipes/nemotronh/nemotron_3_nano.py | 18 +- .../recipes/nemotronh/test_nemotron_3_nano.py | 61 +++-- ...st_nemotron_3_5_lightning_perf_recipes.py} | 77 +++--- .../recipes/test_nemotronh_recipes.py | 73 ++--- 14 files changed, 379 insertions(+), 325 deletions(-) rename examples/model_verification_cards/{nemotron-3.5-nano => nemotron-3.5-lightning}/card.yaml (77%) rename tests/unit_tests/recipes/{test_nemotron_3_5_nano_perf_recipes.py => test_nemotron_3_5_lightning_perf_recipes.py} (81%) diff --git a/examples/model_verification_cards/nemotron-3.5-nano/card.yaml b/examples/model_verification_cards/nemotron-3.5-lightning/card.yaml similarity index 77% rename from examples/model_verification_cards/nemotron-3.5-nano/card.yaml rename to examples/model_verification_cards/nemotron-3.5-lightning/card.yaml index 9d26eb038fd..0616de653d7 100644 --- a/examples/model_verification_cards/nemotron-3.5-nano/card.yaml +++ b/examples/model_verification_cards/nemotron-3.5-lightning/card.yaml @@ -1,7 +1,7 @@ -# Agent-readable model verification card for Nemotron 3.5 Nano. +# Agent-readable model verification card for Nemotron 3.5 Lightning. # status: unverified | verified | unsupported | not_applicable -title: nemotron_3_5_nano +title: nemotron_3_5_lightning summary: > Performance scope: pretrain_performance.H100 and pretrain_performance.GB200 use tuned canonical performance recipes; timing and throughput metrics from @@ -11,10 +11,11 @@ summary: > while adopting the corresponding hardware performance recipe's execution stack. The H100 performance leaf uses BF16 Adam moments in a benchmark-only workload and is not convergence-comparable with FP32 optimizer-state runs. - Nemotron 3.5 Nano support verification covers conversion, inference, and - training on H100 and GB200. The - model.hf_id/--hf-model and model.hf_revision/--hf-revision values are - non-functional placeholders pending publication. + Nemotron 3.5 Lightning support verification covers conversion, inference, + and training on H100 and GB200. Recipe names, work paths, and this + verification-card path use the public Lightning name. The evidence used the 26.08 + release-candidate software stack; the environment records its corresponding + public 26.08 base container. verification_index: model_level: verified: @@ -35,12 +36,12 @@ verification_index: fsdp: GB200: verified model: - hf_id: nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - hf_revision: PUBLIC_HF_REVISION_PENDING + hf_id: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + hf_revision: b3caaabed0263651a17dc1f2d4ce97e794f76c44 # pragma: allowlist secret architecture: NemotronHForCausalLM min_transformers_version: "5.8.1" verification_environment: - base_container: nvcr.io/nvidian/nemo:26.08.rc2 + base_container: nvcr.io/nvidia/nemo:26.08 bridge_commit: 619a2df1da9517b4408ea475b428bf2aba618c8f # pragma: allowlist secret items: @@ -50,9 +51,9 @@ items: bridge_commit: 7edd7c7f68bcdd2b30cbaff426ece2ee0da9deb5 # pragma: allowlist secret command: > ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 - --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron-path work/model-verification/nemotron-3.5-nano/cpu-megatron + --hf-model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron-path work/model-verification/nemotron-3.5-lightning/cpu-megatron --torch-dtype bfloat16 last_verified: 2026-07-23 expected_result: > @@ -67,9 +68,9 @@ items: command: > ./scripts/conversion/convert.sh import --executor slurm --device gpu --nodes 1 --gpus-per-node 8 --tp 1 --pp 1 --ep 8 --etp 1 - --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron-path work/model-verification/nemotron-3.5-nano/imported-megatron + --hf-model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron-path work/model-verification/nemotron-3.5-lightning/imported-megatron --torch-dtype bfloat16 last_verified: 2026-07-23 expected_result: > @@ -84,10 +85,10 @@ items: bridge_commit: 7edd7c7f68bcdd2b30cbaff426ece2ee0da9deb5 # pragma: allowlist secret command: > ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 - --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron-path work/model-verification/nemotron-3.5-nano/cpu-megatron/iter_0000000 - --hf-path work/model-verification/nemotron-3.5-nano/cpu-hf-export + --hf-model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron-path work/model-verification/nemotron-3.5-lightning/cpu-megatron/iter_0000000 + --hf-path work/model-verification/nemotron-3.5-lightning/cpu-hf-export --torch-dtype bfloat16 last_verified: 2026-07-23 expected_result: > @@ -101,10 +102,10 @@ items: command: > ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 1 --gpus-per-node 8 --tp 1 --pp 1 --ep 8 --etp 1 - --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron-path work/model-verification/nemotron-3.5-nano/imported-megatron/iter_0000000 - --hf-path work/model-verification/nemotron-3.5-nano/hf-export + --hf-model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron-path work/model-verification/nemotron-3.5-lightning/imported-megatron/iter_0000000 + --hf-path work/model-verification/nemotron-3.5-lightning/hf-export --torch-dtype bfloat16 --distributed-save last_verified: 2026-07-23 expected_result: > @@ -118,9 +119,9 @@ items: command: > uv run python -m torch.distributed.run --standalone --nproc_per_node=8 examples/conversion/compare_hf_and_megatron/compare.py - --hf_model_path nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron_model_path work/model-verification/nemotron-3.5-nano/imported-megatron/iter_0000000 + --hf_model_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron_model_path work/model-verification/nemotron-3.5-lightning/imported-megatron/iter_0000000 --tp 1 --pp 1 --ep 8 --etp 1 --prompt "The capital of France is the city of" last_verified: 2026-07-23 @@ -136,9 +137,9 @@ items: command: > uv run python -m torch.distributed.run --standalone --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_text.py - --hf_model_path nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron_model_path work/model-verification/nemotron-3.5-nano/imported-megatron/iter_0000000 + --hf_model_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron_model_path work/model-verification/nemotron-3.5-lightning/imported-megatron/iter_0000000 --tp 1 --pp 1 --ep 8 --etp 1 --prompt "The capital of France is" --max_new_tokens 2 last_verified: 2026-07-23 @@ -159,16 +160,16 @@ items: scopes: [mamba] command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 - --recipe nemotron_3_5_nano_pretrain_config + --recipe nemotron_3_5_lightning_pretrain_config --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 20 --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' - dataset.path_to_cache=work/cache/nemotron-3.5-nano-h100/rp2 + dataset.path_to_cache=work/cache/nemotron-3.5-lightning-h100/rp2 rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=20 validation.eval_iters=0 validation.eval_interval=0 model.moe_router_force_load_balancing=false checkpoint.load=null - --save_dir work/model-verification/nemotron-3.5-nano-h100/pretrain-reference + --save_dir work/model-verification/nemotron-3.5-lightning-h100/pretrain-reference --save_interval 10 checkpoint.async_save=false logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 @@ -198,16 +199,16 @@ items: command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 --additional-slurm-params segment=2 - --recipe nemotron_3_5_nano_pretrain_8k_config + --recipe nemotron_3_5_lightning_pretrain_8k_config --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 100 --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' - dataset.path_to_cache=work/cache/nemotron-3.5-nano-gb200/rp2 + dataset.path_to_cache=work/cache/nemotron-3.5-lightning-gb200/rp2 rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 model.moe_router_force_load_balancing=false checkpoint.load=null checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-reference + --save_dir work/model-verification/nemotron-3.5-lightning-gb200/pretrain-reference logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 last_verified: 2026-07-28 @@ -245,16 +246,16 @@ items: command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 --additional-slurm-params segment=2 - --recipe nemotron_3_5_nano_pretrain_8k_fsdp_config + --recipe nemotron_3_5_lightning_pretrain_8k_fsdp_config --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 100 --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' - dataset.path_to_cache=work/cache/nemotron-3.5-nano-gb200/rp2 + dataset.path_to_cache=work/cache/nemotron-3.5-lightning-gb200/rp2 rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 model.moe_router_force_load_balancing=false checkpoint.load=null checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-fsdp + --save_dir work/model-verification/nemotron-3.5-lightning-gb200/pretrain-fsdp logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 last_verified: 2026-07-28 @@ -292,7 +293,7 @@ items: command: > ./scripts/training/train.sh --wait --nodes 2 --gpus-per-node 4 --additional-slurm-params segment=2 - --recipe nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config + --recipe nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config --mode pretrain --max_steps 20 last_verified: 2026-07-24 metrics: @@ -323,21 +324,21 @@ items: moe_dispatcher: deepep command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 - --recipe nemotron_3_5_nano_sft_openmathinstruct2_packed_config + --recipe nemotron_3_5_lightning_sft_openmathinstruct2_packed_config --mode sft --max_steps 100 - --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-h100/imported-megatron/iter_0000000 - --save_dir work/model-verification/nemotron-3.5-nano-h100/sft-checkpoints - last_verified: 2026-07-30 + --pretrained_checkpoint work/model-verification/nemotron-3.5-lightning-h100/imported-megatron/iter_0000000 + --save_dir work/model-verification/nemotron-3.5-lightning-h100/sft-checkpoints + last_verified: 2026-08-03 metrics: - initial_loss: 0.4407026 - final_loss: 0.2756689 - last_10_steps_step_time_ms_avg: 9856.900 - last_10_steps_model_tflops_per_gpu_avg: 79.470 + initial_loss: 0.4406566 + final_loss: 0.2757806 + last_10_steps_step_time_ms_avg: 9887.800 + last_10_steps_model_tflops_per_gpu_avg: 79.240 expected_result: > Packed OpenMathInstruct-2 full SFT completes 100 finite-loss steps with - active MTP gradients, LM loss 0.4407026 -> 0.2756689, finite MTP-1 - loss 0.7634735 -> 0.4076096, and no skipped or NaN iterations. Steps - 91-100 average 9,856.900 ms / 79.470 TFLOP/s/GPU, and the complete + active MTP gradients, LM loss 0.4406566 -> 0.2757806, finite MTP-1 + loss 0.7371141 -> 0.4133309, and no skipped or NaN iterations. Steps + 91-100 average 9,887.800 ms / 79.240 TFLOP/s/GPU, and the complete reloadable full-model checkpoint is written at step 100. GB200: @@ -350,25 +351,25 @@ items: command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 --additional-slurm-params segment=2 - --recipe nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config + --recipe nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config --mode sft --max_steps 100 - --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-gb200/imported-megatron/iter_0000000 - --save_dir work/model-verification/nemotron-3.5-nano-gb200/sft-checkpoints - last_verified: 2026-07-30 + --pretrained_checkpoint work/model-verification/nemotron-3.5-lightning-gb200/imported-megatron/iter_0000000 + --save_dir work/model-verification/nemotron-3.5-lightning-gb200/sft-checkpoints + last_verified: 2026-08-04 metrics: initial_loss: 0.4407797 - final_loss: 0.2760494 - last_10_steps_step_time_ms_avg: 6697.760 - last_10_steps_model_tflops_per_gpu_avg: 260.770 - peak_allocated_memory_gib: 101.640 - peak_reserved_memory_gib: 104.100 + final_loss: 0.2762039 + last_10_steps_step_time_ms_avg: 6795.830 + last_10_steps_model_tflops_per_gpu_avg: 256.990 + peak_allocated_memory_gib: 101.510 + peak_reserved_memory_gib: 103.960 expected_result: > On two complete four-GPU GB200 nodes in one segment-2 NVLink domain, packed OpenMathInstruct-2 full SFT completes exactly 100 finite-loss optimizer steps with active MTP gradients. LM loss is 0.4407797 -> - 0.2760494, MTP-1 loss is 0.7641921 -> 0.4087451, and MTP-2 loss is - 0.8651200 -> 0.4649736. Steps 91-100 average 6,697.760 ms / 260.770 - TFLOP/s/GPU with 101.640 GiB peak allocated and 104.100 GiB peak + 0.2762039, MTP-1 loss is 0.7373838 -> 0.4146959, and MTP-2 loss is + 0.8092746 -> 0.4744196. Steps 91-100 average 6,795.830 ms / 256.990 + TFLOP/s/GPU with 101.510 GiB peak allocated and 103.960 GiB peak reserved memory. No iteration is skipped or NaN, and the complete 12-file reloadable full-model checkpoint is written at step 100. @@ -381,18 +382,18 @@ items: - > ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 2 --gpus-per-node 8 --tp 2 --pp 1 --ep 8 --etp 1 - --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron-path work/model-verification/nemotron-3.5-nano-h100/sft-checkpoints/iter_0000100 - --hf-path work/model-verification/nemotron-3.5-nano-h100/sft-hf-export + --hf-model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron-path work/model-verification/nemotron-3.5-lightning-h100/sft-checkpoints/iter_0000100 + --hf-path work/model-verification/nemotron-3.5-lightning-h100/sft-hf-export --torch-dtype bfloat16 --export-weight-dtype bfloat16 - > uv run python skills/create-model-verification-card/scripts/verify_hf_inference.py - --hf-model work/model-verification/nemotron-3.5-nano-h100/sft-hf-export + --hf-model work/model-verification/nemotron-3.5-lightning-h100/sft-hf-export --prompt "Solve briefly: If 3x + 5 = 20, what is x?" --max-new-tokens 45 --chat-template --disable-thinking - last_verified: 2026-07-30 + last_verified: 2026-08-03 expected_result: "The step-100 SFT checkpoint exports 6,513 tensors, including all 270 MTP tensors, to native Hugging Face format and reloads as NemotronHForCausalLM. The deterministic verifier performs one greedy generation and produces exactly 45 new tokens. The literal completion is: \"To solve for \u0024x\u0024, follow these steps:\n\n1. Subtract 5 from both sides of the equation:\n\\[ 3x + 5 - 5 = 20 - 5 \\Rightarrow 3x\"." GB200: @@ -404,18 +405,18 @@ items: - > ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 2 --gpus-per-node 4 --tp 1 --pp 1 --ep 8 --etp 1 - --hf-model nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16 - --hf-revision PUBLIC_HF_REVISION_PENDING - --megatron-path work/model-verification/nemotron-3.5-nano-gb200/sft-checkpoints/iter_0000100 - --hf-path work/model-verification/nemotron-3.5-nano-gb200/sft-hf-export + --hf-model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --hf-revision b3caaabed0263651a17dc1f2d4ce97e794f76c44 + --megatron-path work/model-verification/nemotron-3.5-lightning-gb200/sft-checkpoints/iter_0000100 + --hf-path work/model-verification/nemotron-3.5-lightning-gb200/sft-hf-export --torch-dtype bfloat16 --export-weight-dtype bfloat16 - > uv run python skills/create-model-verification-card/scripts/verify_hf_inference.py - --hf-model work/model-verification/nemotron-3.5-nano-gb200/sft-hf-export + --hf-model work/model-verification/nemotron-3.5-lightning-gb200/sft-hf-export --prompt "Solve briefly: If 3x + 5 = 20, what is x?" --max-new-tokens 45 --chat-template --disable-thinking - last_verified: 2026-07-30 + last_verified: 2026-08-04 expected_result: "The step-100 SFT checkpoint exports 6,513 tensors, including all 270 MTP tensors, to native Hugging Face format and reloads as NemotronHForCausalLM. The deterministic verifier performs one greedy generation and produces exactly 45 new tokens. The literal completion is: \"To solve for \u0024x\u0024, follow these steps:\n\n1. Subtract 5 from both sides of the equation:\n\\[ 3x + 5 - 5 = 20 - 5 \\Rightarrow 3x\"." sft_long_context: @@ -428,12 +429,12 @@ items: moe_dispatcher: deepep command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 - --recipe nemotron_3_5_nano_sft_config + --recipe nemotron_3_5_lightning_sft_config --mode sft --dataset openmathinstruct2 --max_steps 100 --seq_length 32768 --lr 5e-6 --min_lr 0 --warmup_iters 10 --tensor_model_parallel_size 2 --context_parallel_size 2 --expert_tensor_parallel_size 1 --expert_model_parallel_size 8 - --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-h100/imported-megatron/iter_0000000 + --pretrained_checkpoint work/model-verification/nemotron-3.5-lightning-h100/imported-megatron/iter_0000000 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.enable_offline_packing=true 'dataset.offline_packing_specs={packed_sequence_size:32768,pad_seq_to_mult:4}' @@ -446,20 +447,20 @@ items: scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null checkpoint.save_optim=false checkpoint.save_rng=false checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-h100/long-sft-checkpoints + --save_dir work/model-verification/nemotron-3.5-lightning-h100/long-sft-checkpoints --save_interval 50 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-30 + last_verified: 2026-08-03 metrics: initial_loss: 0.4268321 - final_loss: 0.2649768 - last_10_steps_step_time_ms_avg: 40995.130 - last_10_steps_model_tflops_per_gpu_avg: 188.280 + final_loss: 0.2650192 + last_10_steps_step_time_ms_avg: 40863.010 + last_10_steps_model_tflops_per_gpu_avg: 188.950 expected_result: > Packed OpenMathInstruct-2 CP2 SFT completes exactly 100 finite-loss 32K - steps with LM loss 0.4268321 -> 0.2649768 and MTP-1 loss 0.7375677 -> - 0.4385137. All 100 steps have finite LM/MTP losses and no skipped or - NaN iterations. Steps 91-100 average 40,995.130 ms / 188.280 + steps with LM loss 0.4268321 -> 0.2650192 and MTP-1 loss 0.7322379 -> + 0.4439273. All 100 steps have finite LM/MTP losses and no skipped or + NaN iterations. Steps 91-100 average 40,863.010 ms / 188.950 TFLOP/s/GPU. Complete step-50 and step-100 checkpoints each contain 20 files, distributed metadata, saved run configuration and training state, with latest marker 100. @@ -473,12 +474,12 @@ items: moe_dispatcher: hybridep command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 - --recipe nemotron_3_5_nano_sft_config + --recipe nemotron_3_5_lightning_sft_config --mode sft --dataset openmathinstruct2 --max_steps 100 --seq_length 32768 --lr 5e-6 --min_lr 0 --warmup_iters 10 --context_parallel_size 2 --expert_tensor_parallel_size 1 --expert_model_parallel_size 8 - --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-gb200/imported-megatron/iter_0000000 + --pretrained_checkpoint work/model-verification/nemotron-3.5-lightning-gb200/imported-megatron/iter_0000000 model.moe_flex_dispatcher_backend=hybridep 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.enable_offline_packing=true @@ -492,22 +493,23 @@ items: scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null checkpoint.save_optim=false checkpoint.save_rng=false checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-gb200/long-sft-checkpoints + --save_dir work/model-verification/nemotron-3.5-lightning-gb200/long-sft-checkpoints --save_interval 100 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-30 + last_verified: 2026-08-04 metrics: initial_loss: 0.4268995 - final_loss: 0.2649511 - last_10_steps_step_time_ms_avg: 100358.390 - last_10_steps_model_tflops_per_gpu_avg: 153.810 + final_loss: 0.2649925 + last_10_steps_step_time_ms_avg: 129692.530 + last_10_steps_model_tflops_per_gpu_avg: 119.030 expected_result: > Packed OpenMathInstruct-2 CP2 SFT completes exactly 100 finite-loss 32K - steps with LM loss 0.4268995 -> 0.2649511 and MTP-1 loss 0.7376469 -> - 0.4384451. All 100 steps have finite LM/MTP losses and no skipped or - NaN iterations. The complete step-100 checkpoint has distributed - metadata, 12 files, a latest marker of 100, and all four metrics - recorded. + steps with LM loss 0.4268995 -> 0.2649925 and MTP-1 loss 0.7322624 -> + 0.4439029. All 100 steps have finite LM/MTP losses and no skipped or + NaN iterations. Steps 91-100 average 129,692.530 ms / 119.030 + TFLOP/s/GPU. The complete step-100 checkpoint has distributed metadata, + saved run configuration and training state, 12 files, and a latest + marker of 100. peft: H100: @@ -518,12 +520,12 @@ items: moe_dispatcher: deepep command: > ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 - --recipe nemotron_3_5_nano_peft_config + --recipe nemotron_3_5_lightning_peft_config --mode lora --dataset openmathinstruct2 --max_steps 100 --seq_length 4096 --lr 1e-4 --min_lr 0 --warmup_iters 10 --tensor_model_parallel_size 2 --expert_tensor_parallel_size 1 --expert_model_parallel_size 8 - --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-h100/imported-megatron/iter_0000000 + --pretrained_checkpoint work/model-verification/nemotron-3.5-lightning-h100/imported-megatron/iter_0000000 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.enable_offline_packing=true 'dataset.offline_packing_specs={packed_sequence_size:4096,pad_seq_to_mult:2}' @@ -531,19 +533,19 @@ items: ddp.overlap_param_gather=false optimizer.overlap_param_gather=false scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-h100/peft-checkpoints + --save_dir work/model-verification/nemotron-3.5-lightning-h100/peft-checkpoints --save_interval 100 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-30 + last_verified: 2026-08-03 metrics: - initial_loss: 0.4406879 - final_loss: 0.2856177 - last_10_steps_step_time_ms_avg: 16927.070 - last_10_steps_model_tflops_per_gpu_avg: 92.650 + initial_loss: 0.4406696 + final_loss: 0.2854756 + last_10_steps_step_time_ms_avg: 16900.360 + last_10_steps_model_tflops_per_gpu_avg: 92.890 expected_result: > Packed rank-32, alpha-32, zero-dropout LoRA completes 100 finite-loss - steps with adapters on the MTP layers, LM loss 0.4406879 -> 0.2856177, - finite MTP-1 loss 0.7634883 -> 0.4294895, and no skipped or NaN + steps with adapters on the MTP layers, LM loss 0.4406696 -> 0.2854756, + finite MTP-1 loss 0.7370828 -> 0.4362708, and no skipped or NaN iterations. The complete adapter checkpoint is written at step 100 and reloads over the pinned base checkpoint at step 100. @@ -555,12 +557,12 @@ items: moe_dispatcher: hybridep command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 - --recipe nemotron_3_5_nano_peft_config + --recipe nemotron_3_5_lightning_peft_config --mode lora --dataset openmathinstruct2 --max_steps 100 --seq_length 4096 --lr 1e-4 --min_lr 0 --warmup_iters 10 --tensor_model_parallel_size 2 --expert_tensor_parallel_size 1 --expert_model_parallel_size 8 - --pretrained_checkpoint work/model-verification/nemotron-3.5-nano-gb200/imported-megatron/iter_0000000 + --pretrained_checkpoint work/model-verification/nemotron-3.5-lightning-gb200/imported-megatron/iter_0000000 model.moe_flex_dispatcher_backend=hybridep 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.enable_offline_packing=true @@ -569,22 +571,22 @@ items: ddp.overlap_param_gather=false optimizer.overlap_param_gather=false scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-gb200/peft-checkpoints + --save_dir work/model-verification/nemotron-3.5-lightning-gb200/peft-checkpoints --save_interval 100 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-30 + last_verified: 2026-08-04 metrics: - initial_loss: 0.4405894 - final_loss: 0.2855471 - last_10_steps_step_time_ms_avg: 19163.540 - last_10_steps_model_tflops_per_gpu_avg: 82.080 + initial_loss: 0.4406088 + final_loss: 0.2855029 + last_10_steps_step_time_ms_avg: 19230.400 + last_10_steps_model_tflops_per_gpu_avg: 81.780 expected_result: > Packed rank-32, alpha-32, zero-dropout LoRA completes 100 finite-loss steps with 48 rank-local adapter attachment records covering MTP - attention, expert, and shared-expert layers, finite MTP losses, no - skipped or NaN iterations, and all four metrics recorded. The complete - adapter checkpoint is written at step 100 and reloads over the pinned - base checkpoint at step 100. + attention, expert, and shared-expert layers, LM loss 0.4406088 -> + 0.2855029, finite MTP-1 loss 0.7371238 -> 0.4362065, and no skipped or + NaN iterations. The complete adapter checkpoint is written at step 100 + and reloads over the pinned base checkpoint at step 100. checkpoint_resume: H100: @@ -594,17 +596,17 @@ items: depends_on: pretrain command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 - --recipe nemotron_3_5_nano_pretrain_config + --recipe nemotron_3_5_lightning_pretrain_config --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 20 --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' - dataset.path_to_cache=work/cache/nemotron-3.5-nano-h100/rp2 + dataset.path_to_cache=work/cache/nemotron-3.5-lightning-h100/rp2 rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=20 validation.eval_iters=0 validation.eval_interval=0 model.moe_router_force_load_balancing=false - --load_dir work/model-verification/nemotron-3.5-nano-h100/pretrain-reference + --load_dir work/model-verification/nemotron-3.5-lightning-h100/pretrain-reference checkpoint.ckpt_step=10 checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-h100/pretrain-resumed + --save_dir work/model-verification/nemotron-3.5-lightning-h100/pretrain-resumed --save_interval 10 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 last_verified: 2026-07-28 @@ -637,17 +639,17 @@ items: command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 --additional-slurm-params segment=2 - --recipe nemotron_3_5_nano_pretrain_8k_config + --recipe nemotron_3_5_lightning_pretrain_8k_config --mode pretrain --dataset megatron-indexed --seq_length 8192 --max_steps 100 --lr 3e-4 --min_lr 3e-5 --warmup_iters 8 'dataset.blend=[["work/data/rp2/perp_middle/middle_01_text_document"],null]' - dataset.path_to_cache=work/cache/nemotron-3.5-nano-gb200/rp2 + dataset.path_to_cache=work/cache/nemotron-3.5-lightning-gb200/rp2 rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 model.moe_router_force_load_balancing=false - --load_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-reference + --load_dir work/model-verification/nemotron-3.5-lightning-gb200/pretrain-reference checkpoint.ckpt_step=50 checkpoint.async_save=false - --save_dir work/model-verification/nemotron-3.5-nano-gb200/pretrain-resumed + --save_dir work/model-verification/nemotron-3.5-lightning-gb200/pretrain-resumed logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 last_verified: 2026-07-28 @@ -678,7 +680,7 @@ items: bridge_commit: 573e088c9c6740082c39744e03dc5b009e730ed4 # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 - --recipe nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config + --recipe nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config --max_steps 50 last_verified: 2026-08-03 metrics: @@ -704,7 +706,7 @@ items: bridge_commit: 9b0c2dff42e6175e91d34cd2a504529f803ebde9 # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 4 - --recipe nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config + --recipe nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config --max_steps 50 last_verified: 2026-07-28 metrics: diff --git a/scripts/performance/dump_perf_configs.py b/scripts/performance/dump_perf_configs.py index d5e6e3e7f70..ae8ee0768f5 100644 --- a/scripts/performance/dump_perf_configs.py +++ b/scripts/performance/dump_perf_configs.py @@ -276,12 +276,12 @@ ("nemotronh", "nemotron_3_nano", "pretrain", 8, "b200", "nvfp4"), ("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "bf16"), ("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "fp8_cs"), - # Nemotron 3.5 Nano - ("nemotronh", "nemotron_3_5_nano", "pretrain", 8, "gb200", "bf16"), - ("nemotronh", "nemotron_3_5_nano", "pretrain", 8, "gb200", "fp8_mx"), - ("nemotronh", "nemotron_3_5_nano", "pretrain", 8, "gb200", "nvfp4"), - ("nemotronh", "nemotron_3_5_nano", "pretrain", 16, "h100", "bf16"), - ("nemotronh", "nemotron_3_5_nano", "pretrain", 16, "h100", "fp8_cs"), + # Nemotron 3.5 Lightning + ("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "bf16"), + ("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "fp8_mx"), + ("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "nvfp4"), + ("nemotronh", "nemotron_3_5_lightning", "pretrain", 16, "h100", "bf16"), + ("nemotronh", "nemotron_3_5_lightning", "pretrain", 16, "h100", "fp8_cs"), # Nemotron 3 Super ("nemotronh", "nemotron_3_super", "pretrain", 64, "gb300", "bf16"), ("nemotronh", "nemotron_3_super", "pretrain", 64, "gb300", "fp8_mx"), diff --git a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py index 785860600bd..423b5eb03ab 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py @@ -19,10 +19,10 @@ nemotronh_56b_pretrain_64gpu_b300_fp8cs_config, ) from megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh import ( - nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config, nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config, @@ -45,8 +45,8 @@ nemotronh_56b_pretrain_256gpu_gb300_fp8cs_config, ) from megatron.bridge.perf_recipes.nemotronh.h100.nemotronh import ( - nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config, - nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config, nemotron_3_nano_pretrain_16gpu_h100_bf16_config, nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config, nemotronh_56b_pretrain_64gpu_h100_fp8cs_config, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index c9452d46839..fcb166b99f6 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -33,8 +33,9 @@ from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules -# Placeholder until the public Nemotron 3.5 Nano repository is released. -_NEMOTRON_3_5_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" +# Public Nemotron 3.5 Lightning checkpoint used by the Lightning recipe API. +_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret def nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config() -> ConfigContainer: @@ -457,16 +458,18 @@ def nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: - """Nemotron 3.5 Nano pretrain: 8× GB200, BF16.""" +def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: + """Nemotron 3.5 Lightning pretrain: 8× GB200, BF16.""" cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, @@ -486,16 +489,18 @@ def nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: - """Nemotron 3.5 Nano pretrain: 8× GB200, MXFP8.""" +def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: + """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8.""" cfg = nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config() cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, @@ -515,9 +520,9 @@ def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config() -> ConfigContainer: - """Nemotron 3.5 Nano pretrain: 8× GB200, MXFP8, Megatron FSDP.""" - cfg = nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config() +def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config() -> ConfigContainer: + """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8, Megatron FSDP.""" + cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() # FSDP reduces the model-state footprint enough to use the larger measured # microbatch. Megatron FSDP registers module hooks that Transformer Engine @@ -560,16 +565,18 @@ def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config() -> ConfigContainer return cfg -def nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer: - """Nemotron 3.5 Nano pretrain: 8× GB200, NVFP4.""" +def nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer: + """Nemotron 3.5 Lightning pretrain: 8× GB200, NVFP4.""" cfg = nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config() cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py index eefd9c07037..e63ac273b1c 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py @@ -27,8 +27,9 @@ from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules -# Placeholder until the public Nemotron 3.5 Nano repository is released. -_NEMOTRON_3_5_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" +# Public Nemotron 3.5 Lightning checkpoint used by the Lightning recipe API. +_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret def nemotronh_56b_pretrain_64gpu_h100_fp8cs_config() -> ConfigContainer: @@ -174,8 +175,8 @@ def nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: - """Nemotron 3.5 Nano benchmark pretrain: 16× H100, BF16.""" +def nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: + """Nemotron 3.5 Lightning benchmark pretrain: 16× H100, BF16.""" cfg = nemotron_3_nano_pretrain_16gpu_h100_bf16_config() # Keep the benchmark workload aligned with the GB200 BF16 recipe. The # hardware recipes may tune execution-only knobs such as microbatch size, @@ -186,8 +187,10 @@ def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} # This mock-data, force-balanced benchmark uses lower-precision Adam moments # to make the measured selective-recompute and activation-offload stack fit. @@ -226,16 +229,18 @@ def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer: - """Nemotron 3.5 Nano pretrain: 16× H100, FP8 current-scaling.""" +def nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer: + """Nemotron 3.5 Lightning pretrain: 16× H100, FP8 current-scaling.""" cfg = nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config() cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 1540e438a80..0c562920291 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -13,19 +13,19 @@ # limitations under the License. # Nemotron Nano v2 models -# Nemotron 3 and 3.5 Nano models +# Nemotron 3 Nano and Nemotron 3.5 Lightning models from megatron.bridge.recipes.nemotronh.gb200 import ( - nemotron_3_5_nano_pretrain_8k_config, - nemotron_3_5_nano_pretrain_8k_fsdp_config, - nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, + nemotron_3_5_lightning_pretrain_8k_config, + nemotron_3_5_lightning_pretrain_8k_fsdp_config, + nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, ) from megatron.bridge.recipes.nemotronh.nemotron_3_nano import ( - nemotron_3_5_nano_peft_config, - nemotron_3_5_nano_pretrain_config, - nemotron_3_5_nano_sft_config, - nemotron_3_5_nano_sft_openmathinstruct2_packed_config, + nemotron_3_5_lightning_peft_config, + nemotron_3_5_lightning_pretrain_config, + nemotron_3_5_lightning_sft_config, + nemotron_3_5_lightning_sft_openmathinstruct2_packed_config, nemotron_3_nano_peft_config, nemotron_3_nano_pretrain_config, nemotron_3_nano_sft_config, @@ -95,14 +95,14 @@ "nemotron_nano_12b_v2_sft_config", "nemotron_nano_9b_v2_peft_config", "nemotron_nano_12b_v2_peft_config", - # Nemotron 3 and 3.5 Nano models - "nemotron_3_5_nano_peft_config", - "nemotron_3_5_nano_pretrain_8k_config", - "nemotron_3_5_nano_pretrain_8k_fsdp_config", - "nemotron_3_5_nano_pretrain_config", - "nemotron_3_5_nano_sft_config", - "nemotron_3_5_nano_sft_openmathinstruct2_packed_config", - "nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config", + # Nemotron 3 Nano and Nemotron 3.5 Lightning models + "nemotron_3_5_lightning_peft_config", + "nemotron_3_5_lightning_pretrain_8k_config", + "nemotron_3_5_lightning_pretrain_8k_fsdp_config", + "nemotron_3_5_lightning_pretrain_config", + "nemotron_3_5_lightning_sft_config", + "nemotron_3_5_lightning_sft_openmathinstruct2_packed_config", + "nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_pretrain_config", "nemotron_3_nano_sft_config", "nemotron_3_nano_peft_config", diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py index 2fc637f14fb..fe7149c827e 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py @@ -13,18 +13,18 @@ # limitations under the License. from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano import ( - nemotron_3_5_nano_pretrain_8k_config, - nemotron_3_5_nano_pretrain_8k_fsdp_config, - nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, + nemotron_3_5_lightning_pretrain_8k_config, + nemotron_3_5_lightning_pretrain_8k_fsdp_config, + nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, ) __all__ = [ - "nemotron_3_5_nano_pretrain_8k_config", - "nemotron_3_5_nano_pretrain_8k_fsdp_config", - "nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config", + "nemotron_3_5_lightning_pretrain_8k_config", + "nemotron_3_5_lightning_pretrain_8k_fsdp_config", + "nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py index 136a7ded2f1..78283e737b3 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py @@ -19,8 +19,8 @@ from megatron.bridge import AutoBridge from megatron.bridge.recipes.common import _pretrain_common from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import ( - nemotron_3_5_nano_pretrain_config, - nemotron_3_5_nano_sft_openmathinstruct2_packed_config, + nemotron_3_5_lightning_pretrain_config, + nemotron_3_5_lightning_sft_openmathinstruct2_packed_config, ) from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS from megatron.bridge.training.comm_overlap import CommOverlapConfig @@ -157,9 +157,9 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_8k_config() -> ConfigContainer: - """Return the Nemotron 3.5 Nano BF16 8K pretraining config.""" - cfg = nemotron_3_5_nano_pretrain_config() +def nemotron_3_5_lightning_pretrain_8k_config() -> ConfigContainer: + """Return the Nemotron 3.5 Lightning BF16 8K pretraining config.""" + cfg = nemotron_3_5_lightning_pretrain_config() # Keep the H100 convergence contract intact and change only GB200 # execution/performance settings. @@ -180,9 +180,9 @@ def nemotron_3_5_nano_pretrain_8k_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_8k_fsdp_config() -> ConfigContainer: - """Return the Nemotron 3.5 Nano BF16 8K Megatron FSDP config.""" - cfg = nemotron_3_5_nano_pretrain_8k_config() +def nemotron_3_5_lightning_pretrain_8k_fsdp_config() -> ConfigContainer: + """Return the Nemotron 3.5 Lightning BF16 8K Megatron FSDP config.""" + cfg = nemotron_3_5_lightning_pretrain_8k_config() # Megatron FSDP module hooks are incompatible with Transformer Engine CUDA # graph capture. Keep all convergence settings inherited from the BF16 @@ -205,9 +205,9 @@ def nemotron_3_5_nano_pretrain_8k_fsdp_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContainer: +def nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config() -> ConfigContainer: """Return the optimized TP1 4K packed OpenMathInstruct-2 SFT config.""" - cfg = nemotron_3_5_nano_sft_openmathinstruct2_packed_config() + cfg = nemotron_3_5_lightning_sft_openmathinstruct2_packed_config() cfg.model.tensor_model_parallel_size = 1 cfg.model.sequence_parallel = False @@ -245,9 +245,9 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContain __all__ = [ - "nemotron_3_5_nano_pretrain_8k_config", - "nemotron_3_5_nano_pretrain_8k_fsdp_config", - "nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config", + "nemotron_3_5_lightning_pretrain_8k_config", + "nemotron_3_5_lightning_pretrain_8k_fsdp_config", + "nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py index 267abdcc506..8184f275429 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py @@ -25,10 +25,10 @@ "nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config", "nemotron_3_nano_4b_sft_8gpu_h100_bf16_32k_config", "nemotron_3_nano_4b_sft_8gpu_h100_bf16_config", - "nemotron_3_5_nano_peft_config", - "nemotron_3_5_nano_pretrain_config", - "nemotron_3_5_nano_sft_config", - "nemotron_3_5_nano_sft_openmathinstruct2_packed_config", + "nemotron_3_5_lightning_peft_config", + "nemotron_3_5_lightning_pretrain_config", + "nemotron_3_5_lightning_sft_config", + "nemotron_3_5_lightning_sft_openmathinstruct2_packed_config", "nemotron_3_nano_peft_8gpu_h100_bf16_config", "nemotron_3_nano_pretrain_8gpu_h100_bf16_config", "nemotron_3_nano_sft_8gpu_h100_bf16_config", diff --git a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py index 15651cef5ee..d3c23879f7a 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py @@ -32,8 +32,9 @@ _NEMOTRON_3_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" -# Placeholder until the public Nemotron 3.5 Nano repository is released. -_NEMOTRON_3_5_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" +# Public Nemotron 3.5 Lightning checkpoint used by the Lightning recipe API. +_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret _OPENMATHINSTRUCT2_REVISION = "469216e3f46f4dacf476b382e192485ea51a143e" # pragma: allowlist secret @@ -252,8 +253,8 @@ def nemotron_3_nano_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_pretrain_config() -> ConfigContainer: - """Return the Nemotron 3.5 Nano BF16 pretraining config.""" +def nemotron_3_5_lightning_pretrain_config() -> ConfigContainer: + """Return the Nemotron 3.5 Lightning BF16 pretraining config.""" cfg = nemotron_3_nano_pretrain_8gpu_h100_bf16_config() cfg.train.global_batch_size = 512 # Split the 8K sequence across two ranks so each MTP head materializes only @@ -268,8 +269,10 @@ def nemotron_3_5_nano_pretrain_config() -> ConfigContainer: cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} return cfg @@ -405,22 +408,24 @@ def nemotron_3_nano_sft_8gpu_h100_bf16_config() -> ConfigContainer: return cfg -def nemotron_3_5_nano_sft_config() -> ConfigContainer: - """Return a full SFT config for Nemotron 3.5 Nano.""" +def nemotron_3_5_lightning_sft_config() -> ConfigContainer: + """Return a full SFT config for Nemotron 3.5 Lightning.""" cfg = nemotron_3_nano_sft_8gpu_h100_bf16_config() cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} return cfg -def nemotron_3_5_nano_sft_openmathinstruct2_packed_config() -> ConfigContainer: +def nemotron_3_5_lightning_sft_openmathinstruct2_packed_config() -> ConfigContainer: """Return the verified 4K packed OpenMathInstruct-2 SFT config.""" - cfg = nemotron_3_5_nano_sft_config() + cfg = nemotron_3_5_lightning_sft_config() cfg.model.seq_length = 4096 cfg.model.tensor_model_parallel_size = 2 @@ -439,7 +444,7 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_config() -> ConfigContainer: ) cfg.dataset.hf_dataset.load_kwargs = {"revision": _OPENMATHINSTRUCT2_REVISION} if cfg.dataset.offline_packing_specs is not None: - cfg.dataset.offline_packing_specs.tokenizer_model_name = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.dataset.offline_packing_specs.tokenizer_model_name = _NEMOTRON_3_5_LIGHTNING_MODEL_ID cfg.train.train_iters = 100 cfg.train.global_batch_size = 128 @@ -639,24 +644,26 @@ def nemotron_3_nano_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") return cfg -def nemotron_3_5_nano_peft_config(peft_scheme: str | PEFT = "lora") -> ConfigContainer: - """Return a PEFT config for Nemotron 3.5 Nano.""" +def nemotron_3_5_lightning_peft_config(peft_scheme: str | PEFT = "lora") -> ConfigContainer: + """Return a PEFT config for Nemotron 3.5 Lightning.""" cfg = nemotron_3_nano_peft_8gpu_h100_bf16_config(peft_scheme) cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True cfg.model.keep_mtp_spec_in_bf16 = True cfg.model.mtp_loss_scaling_factor = 0.3 - cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} return cfg __all__ = [ - "nemotron_3_5_nano_peft_config", - "nemotron_3_5_nano_pretrain_config", - "nemotron_3_5_nano_sft_config", - "nemotron_3_5_nano_sft_openmathinstruct2_packed_config", + "nemotron_3_5_lightning_peft_config", + "nemotron_3_5_lightning_pretrain_config", + "nemotron_3_5_lightning_sft_config", + "nemotron_3_5_lightning_sft_openmathinstruct2_packed_config", "nemotron_3_nano_peft_8gpu_h100_bf16_config", "nemotron_3_nano_pretrain_8gpu_h100_bf16_config", "nemotron_3_nano_sft_8gpu_h100_bf16_config", diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py index a51dd69ec36..dd6a0a6af30 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py @@ -12,15 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. # ruff: noqa: F401 -"""Hardware-agnostic aliases for Nemotron 3 and 3.5 Nano recipes.""" +"""Hardware-agnostic aliases for Nemotron 3 Nano and Nemotron 3.5 Lightning recipes.""" from __future__ import annotations from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import ( - nemotron_3_5_nano_peft_config, - nemotron_3_5_nano_pretrain_config, - nemotron_3_5_nano_sft_config, - nemotron_3_5_nano_sft_openmathinstruct2_packed_config, + nemotron_3_5_lightning_peft_config, + nemotron_3_5_lightning_pretrain_config, + nemotron_3_5_lightning_sft_config, + nemotron_3_5_lightning_sft_openmathinstruct2_packed_config, ) from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import ( nemotron_3_nano_peft_8gpu_h100_bf16_config as nemotron_3_nano_peft_config, @@ -34,10 +34,10 @@ __all__ = [ - "nemotron_3_5_nano_peft_config", - "nemotron_3_5_nano_pretrain_config", - "nemotron_3_5_nano_sft_config", - "nemotron_3_5_nano_sft_openmathinstruct2_packed_config", + "nemotron_3_5_lightning_peft_config", + "nemotron_3_5_lightning_pretrain_config", + "nemotron_3_5_lightning_sft_config", + "nemotron_3_5_lightning_sft_openmathinstruct2_packed_config", "nemotron_3_nano_pretrain_config", "nemotron_3_nano_sft_config", "nemotron_3_nano_peft_config", diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py index f63d744e944..570504259fd 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py @@ -13,7 +13,7 @@ # limitations under the License. """ -Unit tests for Nemotron 3 and 3.5 Nano recipe configuration builders. +Unit tests for Nemotron 3 Nano and Nemotron 3.5 Lightning recipe configuration builders. Tests cover: - Separate Nemotron 3 and 3.5 pretrain configurations @@ -34,9 +34,9 @@ from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.recipes.nemotronh.h100 import nemotron_3_nano as recipe_module from megatron.bridge.recipes.nemotronh.nemotron_3_nano import ( - nemotron_3_5_nano_peft_config, - nemotron_3_5_nano_pretrain_config, - nemotron_3_5_nano_sft_config, + nemotron_3_5_lightning_peft_config, + nemotron_3_5_lightning_pretrain_config, + nemotron_3_5_lightning_sft_config, nemotron_3_nano_peft_config, nemotron_3_nano_pretrain_config, nemotron_3_nano_sft_config, @@ -95,9 +95,9 @@ def test_pretrain_config_default_parameters(self): assert config.model.mtp_hybrid_override_pattern is None def test_nemotron_3_5_pretrain_config(self): - """The Nemotron 3.5 recipe enables the repeated Nano MTP head.""" + """The Nemotron 3.5 Lightning recipe enables the repeated MTP head.""" base_config = nemotron_3_nano_pretrain_config() - config = nemotron_3_5_nano_pretrain_config() + config = nemotron_3_5_lightning_pretrain_config() assert config.model.mtp_num_layers == 2 assert config.model.mtp_hybrid_override_pattern == "*E" @@ -106,9 +106,13 @@ def test_nemotron_3_5_pretrain_config(self): assert config.model.mtp_loss_scaling_factor == 0.3 assert config.model.calculate_per_token_loss == base_config.model.calculate_per_token_loss assert config.model.use_te_rng_tracker == base_config.model.use_te_rng_tracker - assert recipe_module._NEMOTRON_3_5_NANO_MODEL_ID == ("nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16") - assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID - assert config.tokenizer.tokenizer_model == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID + assert recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID == ("nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16") + assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID + assert config.model.hf_model_revision == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert config.tokenizer.tokenizer_model == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID + assert config.tokenizer.hf_tokenizer_kwargs == { + "revision": recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + } assert config.train.global_batch_size == 512 assert config.model.context_parallel_size == 2 assert config.model.cp_comm_type == "p2p" @@ -118,7 +122,7 @@ def test_nemotron_3_5_pretrain_config(self): def test_pretrain_recipes_do_not_expose_mtp_flag(self): """Nemotron 3 and 3.5 pretraining use distinct parameterless factories.""" assert "enable_mtp" not in signature(nemotron_3_nano_pretrain_config).parameters - assert "enable_mtp" not in signature(nemotron_3_5_nano_pretrain_config).parameters + assert "enable_mtp" not in signature(nemotron_3_5_lightning_pretrain_config).parameters def test_pretrain_config_hybridep_enabled(self): """Test that HybridEP is enabled by default for MoE pretrain.""" @@ -263,8 +267,8 @@ def test_sft_config_with_pretrained_checkpoint(self): [ (nemotron_3_nano_sft_config, 0, recipe_module._NEMOTRON_3_NANO_MODEL_ID), (nemotron_3_nano_peft_config, 0, recipe_module._NEMOTRON_3_NANO_MODEL_ID), - (nemotron_3_5_nano_sft_config, 2, recipe_module._NEMOTRON_3_5_NANO_MODEL_ID), - (nemotron_3_5_nano_peft_config, 2, recipe_module._NEMOTRON_3_5_NANO_MODEL_ID), + (nemotron_3_5_lightning_sft_config, 2, recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID), + (nemotron_3_5_lightning_peft_config, 2, recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID), ], ) def test_finetuning_uses_nemotron_3_base_model( @@ -292,18 +296,23 @@ def test_finetuning_uses_nemotron_3_base_model( assert config.model.mtp_loss_scaling_factor == (0.3 if mtp_num_layers else 0.1) assert config.model.hf_model_id == tokenizer_model assert config.tokenizer.tokenizer_model == tokenizer_model + if mtp_num_layers: + assert config.model.hf_model_revision == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert config.tokenizer.hf_tokenizer_kwargs == { + "revision": recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + } @pytest.mark.parametrize( ("recipe_factory", "base_factory_name", "recipe_kwargs", "base_args"), [ ( - nemotron_3_5_nano_sft_config, + nemotron_3_5_lightning_sft_config, "nemotron_3_nano_sft_8gpu_h100_bf16_config", {}, (), ), ( - nemotron_3_5_nano_peft_config, + nemotron_3_5_lightning_peft_config, "nemotron_3_nano_peft_8gpu_h100_bf16_config", {"peft_scheme": "dora"}, ("dora",), @@ -332,15 +341,19 @@ def test_nemotron_3_5_finetuning_overrides_base_recipe( assert config.model.mtp_use_repeated_layer is True assert config.model.keep_mtp_spec_in_bf16 is True assert config.model.mtp_loss_scaling_factor == 0.3 - assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID - assert config.tokenizer.tokenizer_model == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID + assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID + assert config.model.hf_model_revision == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert config.tokenizer.tokenizer_model == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID + assert config.tokenizer.hf_tokenizer_kwargs == { + "revision": recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + } def test_finetuning_recipes_do_not_expose_model_id(self): """Model selection is fixed by separate Nemotron 3 and 3.5 factories.""" assert not signature(nemotron_3_nano_sft_config).parameters - assert not signature(nemotron_3_5_nano_sft_config).parameters + assert not signature(nemotron_3_5_lightning_sft_config).parameters assert set(signature(nemotron_3_nano_peft_config).parameters) == {"peft_scheme"} - assert set(signature(nemotron_3_5_nano_peft_config).parameters) == {"peft_scheme"} + assert set(signature(nemotron_3_5_lightning_peft_config).parameters) == {"peft_scheme"} def test_nemotron_3_5_openmath_sft_recipe_owns_verified_h100_defaults(self): """The dedicated H100 recipe makes the standard packed SFT command concise.""" @@ -350,7 +363,7 @@ def test_nemotron_3_5_openmath_sft_recipe_owns_verified_h100_defaults(self): bridge.to_megatron_provider.return_value = provider with patch.object(recipe_module.AutoBridge, "from_hf_pretrained", return_value=bridge) as from_hf: - config = recipe_module.nemotron_3_5_nano_sft_openmathinstruct2_packed_config() + config = recipe_module.nemotron_3_5_lightning_sft_openmathinstruct2_packed_config() from_hf.assert_called_once_with(recipe_module._NEMOTRON_3_NANO_MODEL_ID) assert config.model.mtp_num_layers == 2 @@ -358,7 +371,11 @@ def test_nemotron_3_5_openmath_sft_recipe_owns_verified_h100_defaults(self): assert config.model.mtp_use_repeated_layer is True assert config.model.keep_mtp_spec_in_bf16 is True assert config.model.mtp_loss_scaling_factor == 0.3 - assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID + assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID + assert config.model.hf_model_revision == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert config.tokenizer.hf_tokenizer_kwargs == { + "revision": recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + } assert config.model.tensor_model_parallel_size == 2 assert config.model.sequence_parallel is True assert config.model.expert_tensor_parallel_size == 1 @@ -371,7 +388,9 @@ def test_nemotron_3_5_openmath_sft_recipe_owns_verified_h100_defaults(self): assert config.dataset.enable_offline_packing is True assert config.dataset.offline_packing_specs.packed_sequence_size == 4096 assert config.dataset.offline_packing_specs.pad_seq_to_mult == 2 - assert config.dataset.offline_packing_specs.tokenizer_model_name == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID + assert ( + config.dataset.offline_packing_specs.tokenizer_model_name == recipe_module._NEMOTRON_3_5_LIGHTNING_MODEL_ID + ) assert config.train.train_iters == 100 assert config.train.global_batch_size == 128 assert config.train.micro_batch_size == 1 diff --git a/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py b/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py similarity index 81% rename from tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py rename to tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py index 5f474447b43..a2dd89ea9a1 100644 --- a/tests/unit_tests/recipes/test_nemotron_3_5_nano_perf_recipes.py +++ b/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for Nemotron 3.5 Nano performance recipe exports.""" +"""Unit tests for Nemotron 3.5 Lightning performance recipe exports.""" import importlib from collections.abc import Callable @@ -22,12 +22,12 @@ import torch from megatron.bridge.perf_recipes.nemotronh import ( - nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config, - nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config, - nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config, nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config, @@ -39,18 +39,19 @@ pytestmark = pytest.mark.unit -_NEMOTRON_3_5_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" +_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret _H100_RECIPES = ( - nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config, - nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config, ) _GB200_RECIPES = ( - nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config, - nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config, ) -_GB200_FSDP_RECIPES = (nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config,) +_GB200_FSDP_RECIPES = (nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,) _NEMOTRON_3_RECIPES = ( nemotron_3_nano_pretrain_16gpu_h100_bf16_config, nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config, @@ -60,23 +61,23 @@ ) _NEMOTRON_3_5_BASE_RECIPE_PAIRS = ( ( - nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config, nemotron_3_nano_pretrain_16gpu_h100_bf16_config, ), ( - nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config, + nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config, nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config, ), ( - nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, ), ( - nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config, nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config, ), ( - nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config, nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config, ), ) @@ -100,27 +101,27 @@ ("megatron.bridge.perf_recipes.nemotronh.vr200.nemotronh", "nemotron_3_nano_pretrain_8gpu_vr200_nvfp4_config"), ( "megatron.bridge.perf_recipes.nemotronh.h100.nemotronh", - "nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config", + "nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config", ), ( "megatron.bridge.perf_recipes.nemotronh.h100.nemotronh", - "nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config", + "nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config", ), ( "megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh", - "nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config", ), ( "megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh", - "nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config", + "nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config", ), ( "megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh", - "nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config", + "nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config", ), ( "megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh", - "nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config", + "nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config", ), ) @@ -161,8 +162,10 @@ def test_perf_recipes_enable_mtp(recipe_factory: Callable[[], ConfigContainer]) assert cfg.model.mtp_loss_scaling_factor == 0.3 assert cfg.model.moe_router_force_load_balancing is True assert cfg.model.moe_flex_dispatcher_backend == "hybridep" - assert cfg.model.hf_model_id == _NEMOTRON_3_5_NANO_MODEL_ID - assert cfg.tokenizer.tokenizer_model == _NEMOTRON_3_5_NANO_MODEL_ID + assert cfg.model.hf_model_id == _NEMOTRON_3_5_LIGHTNING_MODEL_ID + assert cfg.model.hf_model_revision == _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert cfg.tokenizer.tokenizer_model == _NEMOTRON_3_5_LIGHTNING_MODEL_ID + assert cfg.tokenizer.hf_tokenizer_kwargs == {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} @pytest.mark.parametrize( @@ -178,7 +181,7 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy( cfg = recipe_factory() base_cfg = base_recipe_factory() - if recipe_factory is not nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config: + if recipe_factory is not nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config: assert cfg.env_vars == base_cfg.env_vars assert cfg.model.calculate_per_token_loss == base_cfg.model.calculate_per_token_loss assert cfg.model.use_te_rng_tracker == base_cfg.model.use_te_rng_tracker @@ -187,11 +190,13 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy( @pytest.mark.parametrize("recipe_factory", _H100_RECIPES, ids=lambda recipe: recipe.__name__) def test_h100_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer]) -> None: - """H100 Nemotron 3.5 variants retain the existing Nano performance topology.""" + """H100 Nemotron 3.5 Lightning variants retain the established performance topology.""" cfg = recipe_factory() assert cfg.model.expert_model_parallel_size == 8 - expected_global_batch_size = 512 if recipe_factory is nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config else 1024 + expected_global_batch_size = ( + 512 if recipe_factory is nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config else 1024 + ) assert cfg.train.global_batch_size == expected_global_batch_size assert cfg.train.micro_batch_size == 1 assert cfg.model.context_parallel_size == 1 @@ -204,7 +209,7 @@ def test_h100_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer] def test_h100_bf16_perf_recipe_uses_measured_benchmark_tuning() -> None: """The H100 BF16 benchmark preserves the measured memory and HybridEP tuning.""" - cfg = nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() + cfg = nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config() assert cfg.optimizer.use_precision_aware_optimizer is True assert cfg.optimizer.main_params_dtype == torch.float32 @@ -230,7 +235,7 @@ def test_h100_bf16_perf_recipe_uses_measured_benchmark_tuning() -> None: def test_h100_fp8_perf_recipe_retains_fp32_optimizer_state() -> None: """The unrelated H100 FP8 benchmark retains its existing optimizer and dispatcher policy.""" - cfg = nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config() + cfg = nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config() assert cfg.optimizer.use_precision_aware_optimizer is False assert cfg.optimizer.exp_avg_dtype == torch.float32 @@ -241,8 +246,8 @@ def test_h100_fp8_perf_recipe_retains_fp32_optimizer_state() -> None: def test_bf16_perf_recipes_share_training_workload() -> None: """H100 and GB200 BF16 recipes otherwise share the same model workload.""" - h100_cfg = nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() - gb200_cfg = nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config() + h100_cfg = nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config() + gb200_cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() gb200_cfg.train.micro_batch_size = h100_cfg.train.micro_batch_size gb200_cfg.model.recompute_granularity = h100_cfg.model.recompute_granularity @@ -268,7 +273,7 @@ def test_bf16_perf_recipes_share_training_workload() -> None: @pytest.mark.parametrize("recipe_factory", _GB200_RECIPES, ids=lambda recipe: recipe.__name__) def test_gb200_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer]) -> None: - """GB200 Nemotron 3.5 variants retain the existing Nano performance topology.""" + """GB200 Nemotron 3.5 Lightning variants retain the established performance topology.""" cfg = recipe_factory() assert cfg.model.expert_model_parallel_size == 8 @@ -284,7 +289,7 @@ def test_gb200_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer def test_gb200_fsdp_perf_recipe_defaults() -> None: """The GB200 FSDP variant retains its measured 8-GPU performance settings.""" - cfg = nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config() + cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config() assert cfg.train.global_batch_size == 384 assert cfg.train.micro_batch_size == 3 diff --git a/tests/unit_tests/recipes/test_nemotronh_recipes.py b/tests/unit_tests/recipes/test_nemotronh_recipes.py index 2622512e199..4353b548f75 100644 --- a/tests/unit_tests/recipes/test_nemotronh_recipes.py +++ b/tests/unit_tests/recipes/test_nemotronh_recipes.py @@ -31,6 +31,7 @@ _nemotronh_module = importlib.import_module("megatron.bridge.recipes.nemotronh") +_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret _NEMOTRONH_RECIPE_FUNCS = [ getattr(_nemotronh_module, name) for name in getattr(_nemotronh_module, "__all__", []) @@ -38,11 +39,11 @@ ] -def test_nemotron_3_5_nano_library_recipe_names_are_hardware_agnostic(): +def test_nemotron_3_5_lightning_library_recipe_names_are_hardware_agnostic(): """Library names describe behavior and do not collide with performance recipes.""" perf_module = importlib.import_module("megatron.bridge.perf_recipes.nemotronh") - library_names = {name for name in _nemotronh_module.__all__ if name.startswith("nemotron_3_5_nano_")} - perf_names = {name for name in dir(perf_module) if name.startswith("nemotron_3_5_nano_")} + library_names = {name for name in _nemotronh_module.__all__ if name.startswith("nemotron_3_5_lightning_")} + perf_names = {name for name in dir(perf_module) if name.startswith("nemotron_3_5_lightning_")} assert all(re.search(r"(?:\d+gpu|h100|b200|b300|gb200|gb300)", name) is None for name in library_names) assert library_names.isdisjoint(perf_names) @@ -159,12 +160,12 @@ def test_nemotron_3_nano_gb200_defers_vocab_size_to_training_tokenizer(): assert cfg.model.vocab_size is None -def test_nemotron_3_5_nano_h100_convergence_recipe_uses_perf_execution_policy(): +def test_nemotron_3_5_lightning_h100_convergence_recipe_uses_perf_execution_policy(): """The H100 convergence recipe keeps safety checks while using the measured fast path.""" - from megatron.bridge.recipes.nemotronh import nemotron_3_5_nano_pretrain_config + from megatron.bridge.recipes.nemotronh import nemotron_3_5_lightning_pretrain_config from megatron.bridge.utils.cuda_graph import cuda_graph_module_names - cfg = nemotron_3_5_nano_pretrain_config() + cfg = nemotron_3_5_lightning_pretrain_config() assert cfg.model.seq_length == 8192 assert cfg.dataset.seq_length == 8192 @@ -200,18 +201,20 @@ def test_nemotron_3_5_nano_h100_convergence_recipe_uses_perf_execution_policy(): assert cfg.checkpoint.async_save is False assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert cfg.model.hf_model_id == "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" - assert cfg.tokenizer.tokenizer_model == "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" + assert cfg.model.hf_model_id == "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + assert cfg.model.hf_model_revision == _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert cfg.tokenizer.tokenizer_model == "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + assert cfg.tokenizer.hf_tokenizer_kwargs == {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 8 assert cfg.env_vars["USE_MNNVL"] == 0 -def test_nemotron_3_5_nano_8k_convergence_recipe_uses_perf_execution_policy(): +def test_nemotron_3_5_lightning_8k_convergence_recipe_uses_perf_execution_policy(): """The 8K convergence recipe keeps safety checks while using the measured fast path.""" - from megatron.bridge.recipes.nemotronh import nemotron_3_5_nano_pretrain_8k_config + from megatron.bridge.recipes.nemotronh import nemotron_3_5_lightning_pretrain_8k_config from megatron.bridge.utils.cuda_graph import cuda_graph_module_names - cfg = nemotron_3_5_nano_pretrain_8k_config() + cfg = nemotron_3_5_lightning_pretrain_8k_config() assert cfg.model.seq_length == 8192 assert cfg.dataset.seq_length == 8192 @@ -242,21 +245,23 @@ def test_nemotron_3_5_nano_8k_convergence_recipe_uses_perf_execution_policy(): assert cfg.checkpoint.async_save is False assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert cfg.model.hf_model_id == "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" - assert cfg.tokenizer.tokenizer_model == "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" + assert cfg.model.hf_model_id == "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + assert cfg.model.hf_model_revision == _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert cfg.tokenizer.tokenizer_model == "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + assert cfg.tokenizer.hf_tokenizer_kwargs == {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 assert cfg.env_vars["USE_MNNVL"] == 1 -def test_nemotron_3_5_nano_h100_and_gb200_differ_only_in_execution_policy(): +def test_nemotron_3_5_lightning_h100_and_gb200_differ_only_in_execution_policy(): """H100 and GB200 must share one training contract.""" from megatron.bridge.recipes.nemotronh import ( - nemotron_3_5_nano_pretrain_8k_config, - nemotron_3_5_nano_pretrain_config, + nemotron_3_5_lightning_pretrain_8k_config, + nemotron_3_5_lightning_pretrain_config, ) - h100 = nemotron_3_5_nano_pretrain_config() - gb200 = nemotron_3_5_nano_pretrain_8k_config() + h100 = nemotron_3_5_lightning_pretrain_config() + gb200 = nemotron_3_5_lightning_pretrain_8k_config() # Normalize the explicit GB200 execution/performance overrides. Any other # difference, including a convergence-setting drift, must fail this test. @@ -277,16 +282,16 @@ def test_nemotron_3_5_nano_h100_and_gb200_differ_only_in_execution_policy(): assert gb200 == h100 -def test_nemotron_3_5_nano_8k_fsdp_recipe_preserves_convergence_contract(): +def test_nemotron_3_5_lightning_8k_fsdp_recipe_preserves_convergence_contract(): """The BF16 FSDP recipe changes only FSDP-required execution settings.""" from megatron.bridge.recipes.nemotronh import ( - nemotron_3_5_nano_pretrain_8k_config, - nemotron_3_5_nano_pretrain_8k_fsdp_config, + nemotron_3_5_lightning_pretrain_8k_config, + nemotron_3_5_lightning_pretrain_8k_fsdp_config, ) from megatron.bridge.utils.cuda_graph import cuda_graph_module_names - reference = nemotron_3_5_nano_pretrain_8k_config() - cfg = nemotron_3_5_nano_pretrain_8k_fsdp_config() + reference = nemotron_3_5_lightning_pretrain_8k_config() + cfg = nemotron_3_5_lightning_pretrain_8k_fsdp_config() assert cfg.model.seq_length == reference.model.seq_length == 8192 assert cfg.dataset.seq_length == reference.dataset.seq_length == 8192 @@ -315,9 +320,9 @@ def test_nemotron_3_5_nano_8k_fsdp_recipe_preserves_convergence_contract(): assert cfg.checkpoint.ckpt_format == "fsdp_dtensor" -def test_nemotron_3_5_nano_openmath_sft_tp1_recipe_uses_tuned_defaults(): +def test_nemotron_3_5_lightning_openmath_sft_tp1_recipe_uses_tuned_defaults(): """The TP1 packed SFT recipe owns the measured topology and run contract.""" - cfg = _nemotronh_module.nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() + cfg = _nemotronh_module.nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config() assert cfg.model.seq_length == 4096 assert cfg.model.tensor_model_parallel_size == 1 @@ -329,7 +334,9 @@ def test_nemotron_3_5_nano_openmath_sft_tp1_recipe_uses_tuned_defaults(): assert cfg.model.moe_hybridep_num_sms == 32 assert cfg.model.recompute_granularity is None assert cfg.model.recompute_modules is None - assert cfg.model.hf_model_id == "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" + assert cfg.model.hf_model_id == "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + assert cfg.model.hf_model_revision == _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION + assert cfg.tokenizer.hf_tokenizer_kwargs == {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION} assert cfg.dataset.seq_length == 4096 assert cfg.dataset.hf_dataset.dataset_name == "openmathinstruct2" @@ -338,7 +345,9 @@ def test_nemotron_3_5_nano_openmath_sft_tp1_recipe_uses_tuned_defaults(): assert cfg.dataset.enable_offline_packing is True assert cfg.dataset.offline_packing_specs.packed_sequence_size == 4096 assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 2 - assert cfg.dataset.offline_packing_specs.tokenizer_model_name == "nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16" + assert cfg.dataset.offline_packing_specs.tokenizer_model_name == ( + "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + ) assert cfg.train.train_iters == 100 assert cfg.train.global_batch_size == 128 @@ -354,15 +363,15 @@ def test_nemotron_3_5_nano_openmath_sft_tp1_recipe_uses_tuned_defaults(): assert cfg.env_vars["USE_MNNVL"] == 1 -def test_nemotron_3_5_nano_h100_and_gb200_sft_differ_only_in_execution_policy(): +def test_nemotron_3_5_lightning_h100_and_gb200_sft_differ_only_in_execution_policy(): """H100 and GB200 packed SFT must share one training contract.""" from megatron.bridge.recipes.nemotronh import ( - nemotron_3_5_nano_sft_openmathinstruct2_packed_config, - nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, + nemotron_3_5_lightning_sft_openmathinstruct2_packed_config, + nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config, ) - h100 = nemotron_3_5_nano_sft_openmathinstruct2_packed_config() - gb200 = nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() + h100 = nemotron_3_5_lightning_sft_openmathinstruct2_packed_config() + gb200 = nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config() # Remove the explicit GB200 execution/performance model overrides. Any # other difference, including a convergence-setting drift, must fail. From 3bc4e4db7f7e86b56707cf20d503ae9aaef5a150 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 21:54:06 -0700 Subject: [PATCH 30/85] fix(training): backport HybridEP uneven-dispatch compatibility to r0.6.0 (#5288) Signed-off-by: Chen Cui Signed-off-by: Malay Nagda Co-authored-by: Malay Nagda --- docs/parallelisms.md | 34 +++--- docs/training/packed-sequences.md | 15 ++- .../bridge/models/transformer_config.py | 52 ++++++++ .../models/test_transformer_config.py | 111 ++++++++++++++++++ 4 files changed, 185 insertions(+), 27 deletions(-) diff --git a/docs/parallelisms.md b/docs/parallelisms.md index 660f9ebd17e..2aff5f6ea89 100644 --- a/docs/parallelisms.md +++ b/docs/parallelisms.md @@ -223,27 +223,23 @@ model_config = GPTModelProvider( # Apply HybridEP optimization apply_flex_dispatcher_backend(model_config, moe_flex_dispatcher_backend="hybridep") - -# Opt in when local dispatch token counts can differ across HybridEP ranks. -model_config.moe_hybridep_pad_uneven_dispatch_inputs = True ``` -`moe_hybridep_pad_uneven_dispatch_inputs` is an explicit opt-in. Enable it -when ranks in the HybridEP communication group can enter dispatch with different -token counts, such as dynamically packed THD batches. HybridEP then takes the -group-wide maximum token count, aligns it for the kernel, pads each rank before -dispatch, and removes the padding after combine. - -Leave the option disabled when every rank already supplies the same number of -tokens. Sequence packing alone does not imply that dispatcher inputs are uneven, -and enabling this path unnecessarily adds a MAX all-reduce and padding overhead. -The option does not enable HybridEP or sequence packing; configure those -separately. With `scripts/training/run_recipe.py`, the equivalent trailing CLI -override is: - -```bash -model.moe_hybridep_pad_uneven_dispatch_inputs=true -``` +Megatron Bridge automatically enables `moe_hybridep_pad_uneven_dispatch_inputs` +when an eager model config uses the HybridEP flex backend. HybridEP requires +equal per-rank dispatch shapes, but dynamically packed THD batches can produce +different local token counts. The safety path takes the group-wide maximum token +count, aligns it for the kernel, pads each rank before dispatch, and removes the +padding after combine. + +This adds a MAX all-reduce and temporary padding overhead, including when every +rank already supplies the same number of tokens. The safety path does not enable +HybridEP or sequence packing; configure those separately. + +CUDA-graph configs preserve their explicit padding setting because the uneven-input +path performs a host scalar synchronization that is not capture-safe. They must +provide equal per-rank dispatch shapes. Disable CUDA graphs when runtime token +counts can differ so Bridge can enable safe padding. **GPU Architecture Requirements:** diff --git a/docs/training/packed-sequences.md b/docs/training/packed-sequences.md index c0f09d433dd..8cefcd1b2ce 100644 --- a/docs/training/packed-sequences.md +++ b/docs/training/packed-sequences.md @@ -113,19 +113,18 @@ The durable constraints for packed sequences in Bridge are: evaluation CP constraints and `CP * TP` when sequence parallelism is enabled - for fine-tuning with CP enabled, per-token loss behavior and reduction settings matter -- HybridEP users whose local dispatch token counts differ across ranks must - explicitly set `model.moe_hybridep_pad_uneven_dispatch_inputs=True`; this - pads only to the group-wide aligned maximum before dispatch and trims the - padding after combine +- Megatron Bridge automatically enables safe uneven-input padding for eager + HybridEP configs; this pads only to the group-wide aligned maximum before + dispatch and trims the padding after combine - CUDA-graph-friendly packed metadata requires additional padding constraints Model-family support is not universal. Some families and recipe paths explicitly opt out of packed sequences or related packing modes. -The HybridEP option is not enabled automatically by in-batch packing. Some -packed datasets already produce equal dispatcher sizes, while variable-sized -inputs can also occur without sequence packing. Keeping the option explicit -avoids an unnecessary synchronization and padding cost for equal-sized inputs. +HybridEP CUDA-graph configs preserve their explicit uneven-input setting because +the safety path performs a host scalar synchronization that is not capture-safe. +They must provide equal per-rank dispatch shapes. Disable CUDA graphs when packed +runtime token counts can differ so Bridge can enable safe padding. ## Relationship to Long-Sequence Training diff --git a/src/megatron/bridge/models/transformer_config.py b/src/megatron/bridge/models/transformer_config.py index 8b1714048d6..2c8b5365830 100644 --- a/src/megatron/bridge/models/transformer_config.py +++ b/src/megatron/bridge/models/transformer_config.py @@ -82,6 +82,55 @@ def _set_moe_expert_tensor_parallel_default(config: MCoreTransformerConfig) -> N config.expert_tensor_parallel_size = 1 +def _enable_safe_hybridep_dispatch(config: MCoreTransformerConfig) -> None: + """Ensure eager HybridEP can dispatch different token counts across ranks. + + Bridge model configs are finalized before runtime batches expose whether their + THD token counts differ by rank. HybridEP requires equal dispatch shapes, so use + Megatron Core's padding path for eager Bridge-configured HybridEP dispatchers. + CUDA-graph configs retain their explicit setting because the padding path's host + scalar synchronization is not capture-safe; those configs require equal inputs. + """ + + def _uses_legacy_full_iteration(value: object) -> bool: + values = value if isinstance(value, list) else [value] + return any( + "full_iteration" in item.split(",") + if isinstance(item, str) + else getattr(item, "name", None) == "full_iteration" + for item in values + ) + + cuda_graph_impl = getattr(config, "cuda_graph_impl", "none") + cuda_graphs_enabled = ( + cuda_graph_impl not in (None, "none") + or getattr(config, "enable_cuda_graph", False) + or getattr(config, "external_cuda_graph", False) + or _uses_legacy_full_iteration(getattr(config, "cuda_graph_modules", None)) + or _uses_legacy_full_iteration(getattr(config, "cuda_graph_scope", None)) + ) + if ( + config.moe_token_dispatcher_type != "flex" + or config.moe_flex_dispatcher_backend != "hybridep" + or cuda_graphs_enabled + ): + return + + padding_fields = tuple( + field_name + for field_name in ( + "moe_hybridep_pad_uneven_dispatch_inputs", + "moe_hybridep_pad_variable_tokens", + ) + if hasattr(config, field_name) + ) + if not padding_fields: + raise AttributeError("Megatron Core TransformerConfig does not expose a HybridEP uneven-input padding field") + for padding_field in padding_fields: + if not getattr(config, padding_field): + setattr(config, padding_field, True) + + @dataclass class TransformerConfig(MCoreTransformerConfig): """Megatron Core TransformerConfig with deferred post-init. @@ -127,6 +176,7 @@ def finalize(self) -> None: if self.sequence_parallel and self.tensor_model_parallel_size <= 1: self.sequence_parallel = False _set_moe_expert_tensor_parallel_default(self) + _enable_safe_hybridep_dispatch(self) MCoreTransformerConfig.__post_init__(self) # In-batch packing produces variable-length packed sequences across microbatches, @@ -203,6 +253,7 @@ def finalize(self) -> None: if self.sequence_parallel and self.tensor_model_parallel_size <= 1: self.sequence_parallel = False _set_moe_expert_tensor_parallel_default(self) + _enable_safe_hybridep_dispatch(self) MCoreMLATransformerConfig.__post_init__(self) if getattr(self, "_enable_in_batch_packing", False) and self.pipeline_model_parallel_size > 1: @@ -257,6 +308,7 @@ def finalize(self) -> None: if self.sequence_parallel and self.tensor_model_parallel_size <= 1: self.sequence_parallel = False _set_moe_expert_tensor_parallel_default(self) + _enable_safe_hybridep_dispatch(self) MCoreHeterogeneousTransformerConfig.__post_init__(self) def get_config_for_layer(self, layer_number: int) -> MCoreTransformerConfig: diff --git a/tests/unit_tests/models/test_transformer_config.py b/tests/unit_tests/models/test_transformer_config.py index c8eda5b6c13..5aeb7f87d40 100644 --- a/tests/unit_tests/models/test_transformer_config.py +++ b/tests/unit_tests/models/test_transformer_config.py @@ -14,6 +14,7 @@ """Unit tests for megatron.bridge.models.transformer_config.""" +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -23,6 +24,7 @@ HeterogeneousTransformerConfig, MLATransformerConfig, TransformerConfig, + _enable_safe_hybridep_dispatch, _resolve_string_fields, ) @@ -43,6 +45,31 @@ def _make_config(**kwargs) -> TransformerConfig: return TransformerConfig(**defaults) +class TestEnableSafeHybridepDispatch: + """Tests for cross-branch Megatron Core HybridEP padding compatibility.""" + + def test_enables_dev_padding_field(self): + cfg = SimpleNamespace( + moe_token_dispatcher_type="flex", + moe_flex_dispatcher_backend="hybridep", + moe_hybridep_pad_variable_tokens=False, + cuda_graph_impl="none", + ) + + _enable_safe_hybridep_dispatch(cfg) + + assert cfg.moe_hybridep_pad_variable_tokens is True + + def test_cuda_graph_config_does_not_require_padding_field(self): + cfg = SimpleNamespace( + moe_token_dispatcher_type="flex", + moe_flex_dispatcher_backend="hybridep", + cuda_graph_impl="full_iteration", + ) + + _enable_safe_hybridep_dispatch(cfg) + + # --------------------------------------------------------------------------- # _resolve_string_fields # --------------------------------------------------------------------------- @@ -224,6 +251,61 @@ def test_dense_unset_expert_tensor_parallel_size_is_not_changed_by_bridge(self): cfg.finalize() assert cfg.expert_tensor_parallel_size is None + def test_hybridep_finalization_enables_uneven_dispatch_padding(self): + """HybridEP must safely handle different token counts on each EP rank.""" + cfg = _make_config( + num_moe_experts=8, + moe_token_dispatcher_type="flex", + moe_flex_dispatcher_backend="hybridep", + moe_hybridep_pad_uneven_dispatch_inputs=False, + ) + + with patch(_FINALIZE_PATCH): + cfg.finalize() + + assert cfg.moe_hybridep_pad_uneven_dispatch_inputs is True + + def test_non_hybridep_finalization_preserves_uneven_dispatch_padding(self): + """Other flex backends must retain their configured padding behavior.""" + cfg = _make_config( + num_moe_experts=8, + moe_token_dispatcher_type="flex", + moe_flex_dispatcher_backend="deepep", + moe_hybridep_pad_uneven_dispatch_inputs=False, + ) + + with patch(_FINALIZE_PATCH): + cfg.finalize() + + assert cfg.moe_hybridep_pad_uneven_dispatch_inputs is False + + @pytest.mark.parametrize( + "cuda_graph_settings", + [ + {"cuda_graph_impl": "full_iteration"}, + {"cuda_graph_impl": "transformer_engine"}, + {"cuda_graph_impl": "local"}, + {"enable_cuda_graph": True}, + {"external_cuda_graph": True}, + {"cuda_graph_modules": "full_iteration"}, + {"cuda_graph_scope": "full_iteration"}, + ], + ) + def test_hybridep_cuda_graph_finalization_preserves_padding_setting(self, cuda_graph_settings): + """CUDA-graph HybridEP configs must not gain a host scalar synchronization.""" + cfg = _make_config( + num_moe_experts=8, + moe_token_dispatcher_type="flex", + moe_flex_dispatcher_backend="hybridep", + moe_hybridep_pad_uneven_dispatch_inputs=False, + **cuda_graph_settings, + ) + + with patch(_FINALIZE_PATCH): + cfg.finalize() + + assert cfg.moe_hybridep_pad_uneven_dispatch_inputs is False + class TestMLATransformerConfigFinalize: """Tests for MLATransformerConfig.finalize().""" @@ -241,6 +323,22 @@ def test_moe_unset_expert_tensor_parallel_size_defaults_to_one(self): cfg.finalize() assert cfg.expert_tensor_parallel_size == 1 + def test_hybridep_finalization_enables_uneven_dispatch_padding(self): + cfg = MLATransformerConfig( + num_layers=2, + hidden_size=64, + num_attention_heads=4, + num_moe_experts=8, + moe_token_dispatcher_type="flex", + moe_flex_dispatcher_backend="hybridep", + moe_hybridep_pad_uneven_dispatch_inputs=False, + ) + + with patch(_MLA_FINALIZE_PATCH): + cfg.finalize() + + assert cfg.moe_hybridep_pad_uneven_dispatch_inputs is True + # --------------------------------------------------------------------------- # HeterogeneousTransformerConfig.finalize @@ -287,3 +385,16 @@ def test_sequence_parallel_preserved_when_tp_gt1(self): with patch(_HETERO_FINALIZE_PATCH): cfg.finalize() assert cfg.sequence_parallel is True + + def test_hybridep_finalization_enables_uneven_dispatch_padding(self): + cfg = self._make_hetero( + num_moe_experts=8, + moe_token_dispatcher_type="flex", + moe_flex_dispatcher_backend="hybridep", + moe_hybridep_pad_uneven_dispatch_inputs=False, + ) + + with patch(_HETERO_FINALIZE_PATCH): + cfg.finalize() + + assert cfg.moe_hybridep_pad_uneven_dispatch_inputs is True From 872ac76908904c0b07304c1f2f60101beea1d489 Mon Sep 17 00:00:00 2001 From: malay-nagda Date: Wed, 5 Aug 2026 15:16:20 +0530 Subject: [PATCH 31/85] perf(recipe): add DeepSeek V4 Flash GB300 FP8MX recipes (#5258) Signed-off-by: Malay Nagda --- .../bridge/perf_recipes/deepseek/__init__.py | 4 + .../deepseek/gb200/deepseek_v4.py | 129 ++++++++++++++++++ .../deepseek/gb300/deepseek_v4.py | 29 +++- .../test_deepseek_v4_flash_perf_recipes.py | 124 +++++++++++++++++ 4 files changed, 285 insertions(+), 1 deletion(-) create mode 100644 src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py create mode 100644 tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py diff --git a/src/megatron/bridge/perf_recipes/deepseek/__init__.py b/src/megatron/bridge/perf_recipes/deepseek/__init__.py index 4cf2f43e866..bac965a98cc 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/__init__.py +++ b/src/megatron/bridge/perf_recipes/deepseek/__init__.py @@ -20,6 +20,9 @@ deepseek_v3_pretrain_256gpu_gb200_fp8mx_partial_cg_dev_config, deepseek_v3_pretrain_256gpu_gb200_nvfp4_config, ) +from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config, +) from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v3 import ( deepseek_v3_pretrain_64gpu_gb300_bf16_fsdp_config, deepseek_v3_pretrain_64gpu_gb300_fp8mx_fsdp_config, @@ -32,6 +35,7 @@ deepseek_v3_pretrain_256gpu_gb300_nvfp4_config, ) from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config, deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config, ) from megatron.bridge.perf_recipes.deepseek.h100.deepseek_v3 import ( diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py new file mode 100644 index 00000000000..e7379dbb991 --- /dev/null +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GB200 performance recipes for DeepSeek V4.""" + +import torch + +from megatron.bridge.perf_recipes._common import _benchmark_common +from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, +) +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.utils.cuda_graph import set_full_iteration_cuda_graph + + +def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: + """DeepSeek V4 Flash pretrain: 128× GB200, MXFP8, full-iteration CUDA graph.""" + cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + # The Sheet lists VPP=4, but the measured PP=1 W&B run resolves VPP to None. + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 64 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.pipeline_model_parallel_layout = None + cfg.train.global_batch_size = 2048 + cfg.train.micro_batch_size = 1 + + cfg.model.attention_backend = "auto" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = 32 + cfg.model.moe_hybridep_num_sms_preprocessing = 108 + cfg.model.moe_router_fusion = True + cfg.model.moe_router_force_load_balancing = True + cfg.model.moe_router_load_balancing_type = "seq_aux_loss" + cfg.model.moe_aux_loss_coeff = 1.0e-4 + + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["mla_up_proj"] + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = [] + cfg.model.fine_grained_offloading_max_inflight_offloads = None + + _benchmark_common(cfg, cross_entropy_impl="native") + + set_full_iteration_cuda_graph(cfg.model) + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.use_te_rng_tracker = True + cfg.rng.te_rng_tracker = True + cfg.train.manual_gc_interval = 10 + + cfg.model.moe_expert_rank_capacity_factor = 1.5 + cfg.model.moe_pad_experts_for_cuda_graph_inference = True + cfg.model.moe_paged_stash = True + cfg.model.moe_paged_stash_page_size = 64 + cfg.model.moe_paged_stash_buffer_size_factor_cpu = 0.0 + cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.2 + cfg.model.moe_mlp_glu_interleave_size = 32 + cfg.model.use_transformer_engine_op_fuser = True + + cfg.model.csa_compress_rotary_base = 40_000 + cfg.model.rotary_scaling_factor = 4 + cfg.model.apply_dsa_kernel_fusion = True + cfg.model.dsa_indexer_loss_coeff = 0.01 + cfg.model.dsa_indexer_use_sparse_loss = True + cfg.model.quant_recipe = None + cfg.model.moe_router_padding_for_fp8 = False + cfg.model.moe_router_padding_for_quantization = True + + cfg.mixed_precision.fp8_param_gather = True + cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = True + cfg.mixed_precision.grad_reduce_in_fp32 = False + + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.bfloat16 + cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 + cfg.optimizer.optimizer_offload_fraction = 1.0 + cfg.optimizer.offload_optimizer_states = False + + cfg.ddp.overlap_param_gather = True + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.average_in_collective = False + cfg.comm_overlap.overlap_grad_reduce = True + cfg.comm_overlap.overlap_moe_expert_parallel_comm = False + cfg.comm_overlap.delay_wgrad_compute = False + + cfg.checkpoint.load_optim = False + cfg.checkpoint.load_rng = False + cfg.checkpoint.save_optim = False + + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NCCL_GRAPH_REGISTER": 0, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, + "NCCL_NVLS_ENABLE": 0, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + } + return cfg diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py index 07624d197d4..6c5d18c8314 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py @@ -11,11 +11,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""GB300 performance recipe for DeepSeek V4 Pro.""" +"""GB300 performance recipes for DeepSeek V4.""" import torch from megatron.bridge.perf_recipes._common import _benchmark_common +from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config, +) from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, @@ -23,6 +26,30 @@ from megatron.bridge.training.config import ConfigContainer +def deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() -> ConfigContainer: + """DeepSeek V4 Flash pretrain: 128× GB300, MXFP8.""" + cfg = deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NCCL_GRAPH_REGISTER": 0, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, + "NCCL_NVLS_ENABLE": 0, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + } + return cfg + + def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: """DeepSeek V4 Pro pretrain: 256× GB300, MXFP8, dev Megatron-Core required.""" cfg = deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() diff --git a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py new file mode 100644 index 00000000000..3a6670d6b88 --- /dev/null +++ b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regression tests for the DeepSeek V4 Flash performance recipes.""" + +import pytest +import torch + +from megatron.bridge.perf_recipes.deepseek import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config, + deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config, +) +from megatron.bridge.utils.cuda_graph import cuda_graph_module_names, is_full_iteration_cuda_graph +from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies + + +pytestmark = pytest.mark.unit + + +@pytest.fixture(autouse=True) +def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: + patch_recipe_construction_dependencies(monkeypatch) + + +def test_deepseek_v4_flash_128gpu_gb200_fp8mx_config() -> None: + cfg = deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() + + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.virtual_pipeline_model_parallel_size is None + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 64 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.sequence_parallel is False + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.train.global_batch_size == 2048 + assert cfg.train.micro_batch_size == 1 + + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_router_force_load_balancing is True + assert cfg.model.moe_hybridep_num_sms == 32 + assert cfg.model.moe_hybridep_num_sms_preprocessing == 108 + assert cfg.model.recompute_granularity == "selective" + assert cfg.model.recompute_modules == ["mla_up_proj"] + + assert is_full_iteration_cuda_graph(cfg.model) + assert cuda_graph_module_names(cfg.model) == [] + assert cfg.model.cuda_graph_warmup_steps == 3 + assert cfg.model.use_te_rng_tracker is True + assert cfg.rng.te_rng_tracker is True + assert cfg.rerun_state_machine.check_for_nan_in_loss is False + assert cfg.ddp.check_for_nan_in_grad is False + + assert cfg.model.moe_expert_rank_capacity_factor == 1.5 + assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True + assert cfg.model.moe_paged_stash is True + assert cfg.model.moe_paged_stash_page_size == 64 + assert cfg.model.moe_paged_stash_buffer_size_factor_cpu == 0.0 + assert cfg.model.moe_paged_stash_buffer_size_factor_cuda == 1.2 + assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert cfg.model.use_transformer_engine_op_fuser is True + + assert cfg.model.csa_compress_rotary_base == 40_000 + assert cfg.model.rotary_scaling_factor == 4 + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.dsa_indexer_loss_coeff == 0.01 + assert cfg.model.dsa_indexer_use_sparse_loss is True + assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.model.quant_recipe is None + assert cfg.model.moe_router_padding_for_fp8 is False + assert cfg.model.moe_router_padding_for_quantization is True + + assert cfg.mixed_precision.fp8_param_gather is True + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.optimizer.main_params_dtype == torch.float32 + assert cfg.optimizer.exp_avg_dtype == torch.bfloat16 + assert cfg.optimizer.exp_avg_sq_dtype == torch.bfloat16 + assert cfg.optimizer.optimizer_offload_fraction == 1.0 + assert cfg.optimizer.offload_optimizer_states is False + assert cfg.ddp.grad_reduce_in_fp32 is False + assert cfg.ddp.average_in_collective is False + assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is False + assert cfg.comm_overlap.delay_wgrad_compute is False + + assert cfg.checkpoint.load_optim is False + assert cfg.checkpoint.load_rng is False + assert cfg.checkpoint.save_optim is False + assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == ( + "expandable_segments:True,graph_capture_record_stream_reuse:True" + ) + assert cfg.env_vars["NCCL_GRAPH_REGISTER"] == 0 + assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 0 + assert cfg.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] == 32 + assert cfg.env_vars["NVTE_FWD_LAYERNORM_SM_MARGIN"] == 20 + assert cfg.env_vars["NVTE_BWD_LAYERNORM_SM_MARGIN"] == 20 + assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 + assert cfg.env_vars["USE_MNNVL"] == 1 + + +def test_deepseek_v4_flash_128gpu_gb300_fp8mx_config() -> None: + cfg = deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() + + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.virtual_pipeline_model_parallel_size is None + assert cfg.model.expert_model_parallel_size == 64 + assert cfg.train.global_batch_size == 2048 + assert cfg.train.micro_batch_size == 1 + assert cfg.model.recompute_modules == ["mla_up_proj"] + assert is_full_iteration_cuda_graph(cfg.model) From 59a658f4123787d04bc29477fc3105ae0ea5a3db Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 5 Aug 2026 13:36:59 +0200 Subject: [PATCH 32/85] cp: `build: bump nvidia-modelopt to 0.46.0rc1 (5331)` into `r0.6.0` (#5338) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: NeMo Bot Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c71058a8bb7..9601d3df079 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -133,7 +133,7 @@ override-dependencies = [ "mlflow>=3.15.0", # To address CVE-2025-15031 "cachetools>=5.0.0", "cryptography>=43.0.0,<47", - "nvidia-modelopt==0.44.0rc5", + "nvidia-modelopt==0.46.0rc1", "urllib3>=2.6.3", # To address CVE-2026-21441 "langchain>=0.3.28", # To address CVE-2025-65106 "langchain-core>=0.3.80", # To address CVE-2025-65106 diff --git a/uv.lock b/uv.lock index 274b41ef2d4..4ab0be37198 100644 --- a/uv.lock +++ b/uv.lock @@ -19,7 +19,7 @@ overrides = [ { name = "mlflow", specifier = ">=3.15.0" }, { name = "nvidia-cudnn-frontend", specifier = "==1.26.0" }, { name = "nvidia-cutlass-dsl", extras = ["cu13"], marker = "sys_platform == 'never'" }, - { name = "nvidia-modelopt", specifier = "==0.44.0rc5", index = "https://pypi.nvidia.com/" }, + { name = "nvidia-modelopt", specifier = "==0.46.0rc1", index = "https://pypi.nvidia.com/" }, { name = "onnx", marker = "python_full_version < '3.13'", specifier = ">=1.21.0" }, { name = "tokenizers", specifier = ">=0.22.0,<0.23" }, { name = "torch", marker = "sys_platform == 'never'" }, @@ -2935,7 +2935,7 @@ wheels = [ [[package]] name = "nvidia-modelopt" -version = "0.44.0rc5" +version = "0.46.0rc1" source = { registry = "https://pypi.nvidia.com/" } dependencies = [ { name = "ninja" }, @@ -2955,7 +2955,7 @@ dependencies = [ { name = "tqdm" }, ] wheels = [ - { url = "https://pypi.nvidia.com/nvidia-modelopt/nvidia_modelopt-0.44.0rc5-py3-none-any.whl", hash = "sha256:4934c50f1fb94f7b60a6b53f2bdbd470e68b00d1161206118b0be196cc7d9b22" }, + { url = "https://pypi.nvidia.com/nvidia-modelopt/nvidia_modelopt-0.46.0rc1-py3-none-any.whl", hash = "sha256:9eb5393812462674283d1919626d8ef23e17f15530b6626a6aaa32c5b810c5ad" }, ] [[package]] From 36950891c9665251bfb6fc12cace326770612382 Mon Sep 17 00:00:00 2001 From: Ajay Date: Wed, 5 Aug 2026 11:09:09 -0700 Subject: [PATCH 33/85] chore(docker): update Dockerfile and install script to remove unnecessary packages (#5330) Signed-off-by: Ajay Balasa --- docker/Dockerfile.fw_final | 5 +++-- docker/common/install_aws_ofi_nccl.sh | 8 ++++++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/docker/Dockerfile.fw_final b/docker/Dockerfile.fw_final index 74a3132f8d6..4a804544695 100644 --- a/docker/Dockerfile.fw_final +++ b/docker/Dockerfile.fw_final @@ -92,8 +92,9 @@ RUN pip install "fastapi==0.139.1" \ "mistune>=3.3.0" \ "GitPython>=3.1.55" && \ # Address CVE from transitive dependency - pip uninstall -y opencv-python-headless nvidia-dali-cuda130 - + # Remove the package pulled in by Humming kernels. + # Humming-backed vLLM quantization is unavailable without pyelftools. + pip uninstall -y opencv-python-headless nvidia-dali-cuda130 pyelftools # Remove PyTorch's bundled third-party onnx source tree (build artifact, not used at runtime). # The actual onnx Python package is managed above via pip/uv; this stale copy triggers CVE scanners. RUN rm -rf /opt/pytorch/pytorch/third_party/onnx diff --git a/docker/common/install_aws_ofi_nccl.sh b/docker/common/install_aws_ofi_nccl.sh index 49de30f388a..6304d451d84 100644 --- a/docker/common/install_aws_ofi_nccl.sh +++ b/docker/common/install_aws_ofi_nccl.sh @@ -42,8 +42,6 @@ SRC_DIR="$(mktemp -d)" apt-get update apt-get install -y --no-install-recommends \ git ca-certificates build-essential autoconf automake libtool libhwloc-dev -apt-get clean -rm -rf /var/lib/apt/lists/* git clone --depth 1 --branch "${AWS_OFI_NCCL_VER}" \ https://github.com/aws/aws-ofi-nccl.git "${SRC_DIR}" @@ -63,6 +61,12 @@ make -j"$(nproc)" make install popd +# The installed plugin links to libhwloc15 at runtime but does not need the +# development headers or libltdl development files after compilation. +apt-get purge -y libhwloc-dev libltdl-dev +apt-get clean +rm -rf /var/lib/apt/lists/* + # Revive the ldconfig entry the base image left behind so a bare `libnccl-net.so` # lookup can also resolve here; EKSEnvPlugin still pins the absolute path. echo "${PREFIX}/lib" > /etc/ld.so.conf.d/aws-ofi-nccl.conf From d51ba9f392df4aea5c3a55788a90a13416ccd813 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Wed, 5 Aug 2026 22:27:40 +0200 Subject: [PATCH 34/85] build: AUT-1353 bump transformer-engine to 2.17.1 (#5349) Signed-off-by: svcnemo-autobot --- pyproject.toml | 6 ++---- uv.lock | 13 +++++-------- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9601d3df079..c02830fe598 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,7 @@ override-dependencies = [ "torch; sys_platform == 'never'", "torchvision; sys_platform == 'never'", "triton; sys_platform == 'never'", - "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@cfda38340769c432c12db3c0df574f8e6ce46c8c", + "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@4329ff84bfbdaa778a33cba02a15fb0807c64689", "mlflow>=3.15.0", # To address CVE-2025-15031 "cachetools>=5.0.0", "cryptography>=43.0.0,<47", @@ -168,12 +168,11 @@ requires-dist = [] # above) changes — an out-of-date entry sends uv back to building TE. [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.18.0+cfda3834" +version = "2.17.1+4329ff84" requires-dist = [ "einops", "importlib-metadata>=1.0", "nvdlfw-inspect", - "nvidia-cudnn-frontend>=1.25.0", "onnx", "onnxscript", "packaging", @@ -242,7 +241,6 @@ build = [ "numpy", "ninja", "nvidia-mathdx", - "nvidia-cudnn-frontend>=1.25.0", ] # CVE-carrying WAN diffusion codecs, needed only by WAN video caching # (src/megatron/bridge/diffusion/models/wan/inference/utils.py). All three are suppressed diff --git a/uv.lock b/uv.lock index 4ab0be37198..374b85c5356 100644 --- a/uv.lock +++ b/uv.lock @@ -24,7 +24,7 @@ overrides = [ { name = "tokenizers", specifier = ">=0.22.0,<0.23" }, { name = "torch", marker = "sys_platform == 'never'" }, { name = "torchvision", marker = "sys_platform == 'never'" }, - { name = "transformer-engine", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=cfda38340769c432c12db3c0df574f8e6ce46c8c" }, + { name = "transformer-engine", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689" }, { name = "triton", marker = "sys_platform == 'never'" }, { name = "urllib3", specifier = ">=2.6.3" }, ] @@ -35,8 +35,8 @@ version = "1.0.0+b7643bd" [[manifest.dependency-metadata]] name = "transformer-engine" -version = "2.18.0+cfda3834" -requires-dist = ["einops", "importlib-metadata>=1.0", "nvdlfw-inspect", "nvidia-cudnn-frontend>=1.25.0", "onnx", "onnxscript", "packaging", "pydantic", "torch>=2.1"] +version = "2.17.1+4329ff84" +requires-dist = ["einops", "importlib-metadata>=1.0", "nvdlfw-inspect", "onnx", "onnxscript", "packaging", "pydantic", "torch>=2.1"] provides-extras = ["core-cu13", "pytorch"] [[package]] @@ -2142,7 +2142,6 @@ build = [ { name = "cython" }, { name = "ninja" }, { name = "numpy" }, - { name = "nvidia-cudnn-frontend" }, { name = "nvidia-mathdx" }, { name = "pybind11" }, { name = "setuptools" }, @@ -2228,7 +2227,6 @@ build = [ { name = "cython", specifier = ">=3.0.0" }, { name = "ninja" }, { name = "numpy" }, - { name = "nvidia-cudnn-frontend", specifier = ">=1.25.0" }, { name = "nvidia-mathdx" }, { name = "pybind11" }, { name = "setuptools" }, @@ -4863,13 +4861,12 @@ wheels = [ [[package]] name = "transformer-engine" -version = "2.18.0+cfda3834" -source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=cfda38340769c432c12db3c0df574f8e6ce46c8c#cfda38340769c432c12db3c0df574f8e6ce46c8c" } +version = "2.17.1+4329ff84" +source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689#4329ff84bfbdaa778a33cba02a15fb0807c64689" } dependencies = [ { name = "einops" }, { name = "importlib-metadata" }, { name = "nvdlfw-inspect" }, - { name = "nvidia-cudnn-frontend" }, { name = "onnx" }, { name = "onnxscript" }, { name = "packaging" }, From 6084efb44edab23d78c741ec85f79ee61f94216c Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Wed, 5 Aug 2026 22:48:42 +0200 Subject: [PATCH 35/85] build: AUT-1346 pin wandb to 0.27.2 (#5348) Signed-off-by: svcnemo-autobot --- pyproject.toml | 2 +- uv.lock | 27 ++++++++++++++------------- 2 files changed, 15 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c02830fe598..aa26cb2f82d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,7 +83,7 @@ dependencies = [ "tensorboard>=2.19.0", "typing-extensions", "rich", - "wandb>=0.25.0", + "wandb==0.27.2", "six>=1.17.0", "regex>=2024.11.6", "pyyaml>=6.0.2", diff --git a/uv.lock b/uv.lock index 374b85c5356..3248010ca16 100644 --- a/uv.lock +++ b/uv.lock @@ -2218,7 +2218,7 @@ requires-dist = [ { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'" }, { name = "transformers", specifier = ">=5.8,<=5.12.1" }, { name = "typing-extensions" }, - { name = "wandb", specifier = ">=0.25.0" }, + { name = "wandb", specifier = "==0.27.2" }, ] provides-extras = ["recipes", "parquet", "tensor-inspect", "te", "ssm", "audio"] @@ -5005,10 +5005,11 @@ wheels = [ [[package]] name = "wandb" -version = "0.28.1" +version = "0.27.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, + { name = "gitpython" }, { name = "packaging" }, { name = "platformdirs" }, { name = "protobuf" }, @@ -5018,17 +5019,17 @@ dependencies = [ { name = "sentry-sdk" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/92/fb/8d3f96a8b143060d6fa145462d0785981373e04694e4152555ccb5d23939/wandb-0.28.1.tar.gz", hash = "sha256:870ccb1a01238b0ac07c6fd96a0810a1f79090aba04ea29f4ee012ac8327705d", size = 40578119, upload-time = "2026-07-16T18:47:05.413Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/06/21/8df50164d07623cfcefec19bbf9327d9be84b637a827cea1f0c06db005fd/wandb-0.28.1-py3-none-macosx_12_0_arm64.whl", hash = "sha256:da909a76e65c64c0d93acc485d2a19f66e336f1e3f725f1c98a070883e084943", size = 24277925, upload-time = "2026-07-16T18:46:42.383Z" }, - { url = "https://files.pythonhosted.org/packages/8f/18/6c3da7e6cb215ad363324db8dc4d83b93626f5e339822b05b1c38a6097fd/wandb-0.28.1-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:3da3db219c54bfd1082c00e9061c8ea894ba43e42733b5af00bb10c09d7158fe", size = 25480852, upload-time = "2026-07-16T18:46:45.102Z" }, - { url = "https://files.pythonhosted.org/packages/e2/1a/d15bcfb4417fa69edcaa33db8ea012db733da1057e193b047e3f69fdd671/wandb-0.28.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:ae9ae6fb29e2e2b1d097ed8b75c0c0240c778c2a8cad1d996dee870a1e401c2c", size = 24832138, upload-time = "2026-07-16T18:46:47.433Z" }, - { url = "https://files.pythonhosted.org/packages/b3/da/49924c7df2952dfd82c86c3779c339c0c3d6f6439387c03d97d0470c3658/wandb-0.28.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:8cfb898b6a6c884d9c9294b02764e88bce65049f027a124d6bee53fe722469b6", size = 26486533, upload-time = "2026-07-16T18:46:49.839Z" }, - { url = "https://files.pythonhosted.org/packages/11/c0/06b23518e29690784f1b3081e39c7679ca076cb0af094cb9b4bb309150f5/wandb-0.28.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:cf2b1533945395e4fdbe6182b272bb0ca8a02c10b3086a395e2d57686ae3ed0d", size = 25022635, upload-time = "2026-07-16T18:46:52.376Z" }, - { url = "https://files.pythonhosted.org/packages/23/30/6de2f7995a8a6eecbd03d24c79a139a734c0168f5520cf4c7ccb43c1dbbc/wandb-0.28.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7233061080507a4b4098bed1ccb381ce6f890c60397cd4153d060285bcb267bd", size = 27008895, upload-time = "2026-07-16T18:46:55.025Z" }, - { url = "https://files.pythonhosted.org/packages/b2/83/49deab9447687625371435ca21b6da82f223f1c7d014d77386b9cb91833c/wandb-0.28.1-py3-none-win32.whl", hash = "sha256:4bc461cda3ce23a19d8df5e42981a664d95fa3231efb10fd1e85d9d4824c7d29", size = 24418398, upload-time = "2026-07-16T18:46:57.43Z" }, - { url = "https://files.pythonhosted.org/packages/bf/6f/ed6616b11ea15b8ceabedcaa567286c1c9ec65fa50563230a90bfb627cc5/wandb-0.28.1-py3-none-win_amd64.whl", hash = "sha256:d98a10370162b1e970850237114c56e9c4c58f3cb701e4b8cb38f36f6749fd52", size = 24418404, upload-time = "2026-07-16T18:47:00.427Z" }, - { url = "https://files.pythonhosted.org/packages/07/78/75b6827a6665337a715c5347c5edbd84eca660f7a0f48d8d6d24d1f66bee/wandb-0.28.1-py3-none-win_arm64.whl", hash = "sha256:4aa07f13dd3bcac2c0524c8d0f49f76e83ab5c1054fd09f3b1a436cfcde146a6", size = 22299006, upload-time = "2026-07-16T18:47:02.71Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/14/a2/53ca062f430178e3af48ebc137396481d0ee885fb94a554c0df464cd8afa/wandb-0.27.2.tar.gz", hash = "sha256:c81ff93ab63f4dabc5a27b90ac3d12310fbfa6a14ca99201626921c99b2845be", size = 40300451, upload-time = "2026-06-06T01:47:02.74Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/95/18d3625558667b459d91c19630f7cecfbc133f87f5b144a7fb755e473e8c/wandb-0.27.2-py3-none-macosx_12_0_arm64.whl", hash = "sha256:978400b3c4b7d97e927c32264453da5e4a0040a3468d5b77a00d9c480613f370", size = 23990048, upload-time = "2026-06-06T01:46:38.902Z" }, + { url = "https://files.pythonhosted.org/packages/43/14/72c26f67b0b6cb307cbb76659465c6ab7d99ea27c268d1b4f5aa82c4d8e5/wandb-0.27.2-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:de8099f02f540c743069617db7d034511a64c193748783aa6d2d98310918d170", size = 25165812, upload-time = "2026-06-06T01:46:42.068Z" }, + { url = "https://files.pythonhosted.org/packages/b5/a3/f9fe31ca72b4f5854d1e488403d6310783127a6b7e267c28577e9bd51b43/wandb-0.27.2-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:e0779592410215a2762063c3585d3dcad73c7dca9cb6d63c4dcc1588267c1392", size = 24554366, upload-time = "2026-06-06T01:46:44.57Z" }, + { url = "https://files.pythonhosted.org/packages/76/e2/7a5064aba235ddb855b8c2250e07e6187fcc8382332e237e545d4de094ee/wandb-0.27.2-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:55bfebf4d382116a8e9610848cadc0de50d406bacd3d0a390d12dabde196f009", size = 26380293, upload-time = "2026-06-06T01:46:47.487Z" }, + { url = "https://files.pythonhosted.org/packages/ed/4c/0c845edac5ff0fd0930e881bec2569f2e2af2a4fc873249855600546eee0/wandb-0.27.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:566aa2fcd67d2a23c08713da75e9daf82f30f7136af76763ef1d7db3d901d940", size = 24728823, upload-time = "2026-06-06T01:46:49.887Z" }, + { url = "https://files.pythonhosted.org/packages/e1/7a/a6f7a02a0e6bf73e163b61caca03aaba3452836a02dbe2b64f9e1a3c6afc/wandb-0.27.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:41158181fc5b691438b3d04fee0a8c061e3f1f407a3258096afbebfe1db24e72", size = 26691957, upload-time = "2026-06-06T01:46:52.384Z" }, + { url = "https://files.pythonhosted.org/packages/e0/9b/aa94eb8265b0c55dc6c3e435c11241b3f885c7a1720718046efd7cbd8361/wandb-0.27.2-py3-none-win32.whl", hash = "sha256:5c55fad8c7be9d345dcebdc9dc10f7d2ac5af5bede62acbcd79a412ccaf48c87", size = 24151396, upload-time = "2026-06-06T01:46:54.793Z" }, + { url = "https://files.pythonhosted.org/packages/bf/0b/9e442779f5f24baaca044daf7546a735f41a811886b84ae12740d51b7f9d/wandb-0.27.2-py3-none-win_amd64.whl", hash = "sha256:87204d4fe40fbd9a1fe89a05927ce4ddb8be34d8210045457819fb4a35e0bcea", size = 24151404, upload-time = "2026-06-06T01:46:56.994Z" }, + { url = "https://files.pythonhosted.org/packages/aa/3a/01ab3afc1f6f962df93db34be8183147c9821e4dce82dc03313fb8d08635/wandb-0.27.2-py3-none-win_arm64.whl", hash = "sha256:32ed7456f40443c971e95dd63704d840fce66c24f88049a9bda8a09dfe85effe", size = 22063373, upload-time = "2026-06-06T01:47:00.263Z" }, ] [[package]] From 917708706023676d0b4d29ac7217545f3c8a648f Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Thu, 6 Aug 2026 01:48:16 +0200 Subject: [PATCH 36/85] fix(ci): AUT-1359 skip humming module in import check (#5358) Signed-off-by: svcnemo-autobot --- docker/common/import_check_skip.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/docker/common/import_check_skip.txt b/docker/common/import_check_skip.txt index fbf2373f7ab..3058527a361 100644 --- a/docker/common/import_check_skip.txt +++ b/docker/common/import_check_skip.txt @@ -25,6 +25,7 @@ torch_tensorrt quack # imports nvidia-cutlass-dsl MLIR libs, which dlopen libcuda at import peft # transitively imports optional CUDA-backed modules that dlopen libcuda transformer_engine # dlopens libcuda at import +humming # Humming vLLM quantization kernels dlopen libcuda at import # Requires the libfuse system library (multi-storage-client FUSE backend; optional): mfusepy From 63e2b009d3089d4b863de1c2151f24f91bd07a38 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Thu, 6 Aug 2026 16:54:46 +0200 Subject: [PATCH 37/85] fix(model): AUT-1326 backport Gemma 3 VL fused attention to r0.6.0 (#5336) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: yaoyu-33 Signed-off-by: svcnemo-autobot Signed-off-by: oliver könig Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: oliver könig --- .../compare_hf_and_megatron/compare.py | 2 +- .../gemma-3-4b-it/card.yaml | 48 ++++++---- .../bridge/models/gemma/gemma3_provider.py | 30 ++++-- .../models/gemma_vl/gemma3_vl_provider.py | 8 +- .../models/gemma_vl/modeling_gemma3_vl.py | 50 +++++++--- .../recipes/gemma3_vl/h100/gemma3_vl.py | 12 +-- tests/functional_tests/fixture_utils.py | 39 ++++++++ .../active/L2_Launch_recipes_gemma_vl.sh | 8 -- .../h100/active/L2_Launch_recipes_gemma_vl.sh | 8 -- .../converter/test_generate_vlm_from_hf.py | 10 +- .../data/download_unit_tests_dataset.py | 96 ++++++++++++------- .../qwen_audio/test_qwen2_audio_generation.py | 9 +- .../qwen_vl/test_qwen3_vl_generation.py | 8 +- .../test_gemma3_vl_recipes_finetune.py | 26 ----- .../gemma_vl/test_gemma3_vl_modeling.py | 92 +++++++++++++++--- .../gemma_vl/test_gemma3_vl_provider.py | 68 +++++++++++-- .../recipes/test_gemma3_vl_recipes.py | 2 +- .../test_download_unit_tests_dataset.py | 70 ++++++++++++++ tests/unit_tests/test_fixture_utils.py | 32 +++++++ 19 files changed, 462 insertions(+), 156 deletions(-) create mode 100644 tests/functional_tests/fixture_utils.py create mode 100644 tests/unit_tests/test_download_unit_tests_dataset.py create mode 100644 tests/unit_tests/test_fixture_utils.py diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index 0ce406fa2d1..78acea49663 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -258,7 +258,7 @@ def is_vision_language_model( except Exception as e: print_rank_0(f"Warning: Could not determine model type from config: {e}") - # Fallback: check if qwen_vl_utils is available and model name contains vl indicators + # Fallback: check whether the model name contains common VL indicators. return any(indicator in model_path.lower() for indicator in ["vl", "vision"]) diff --git a/examples/model_verification_cards/gemma-3-4b-it/card.yaml b/examples/model_verification_cards/gemma-3-4b-it/card.yaml index d2f19aa0d8f..04fb8d09a03 100644 --- a/examples/model_verification_cards/gemma-3-4b-it/card.yaml +++ b/examples/model_verification_cards/gemma-3-4b-it/card.yaml @@ -11,9 +11,9 @@ summary: > vision namespace gate does not pass; SFT, post-SFT export/inference, and PEFT remain pending. Pretraining and checkpoint resume are unsupported because there is no public exact-model pretraining recipe. Long-context SFT - is unsupported because the required arbitrary multimodal mask currently - excludes context parallelism and in-batch THD packing. Existing toy-model - tests are not used as exact-model evidence. + is unsupported because the required dense local/global attention biases + currently exclude context parallelism and in-batch THD packing. Existing + toy-model tests are not used as exact-model evidence. verification_index: model_level: verified: @@ -35,11 +35,12 @@ model: min_transformers_version: "5.8.0" verification_environment: base_container: nvcr.io/nvidia/pytorch:26.04-py3 - bridge_commit: 2d0175b116af6cb785bb8efab7ee053435496266 # pragma: allowlist secret + bridge_commit: 330f87c0a391eb2770cbf2306352cc807211d7ba # pragma: allowlist secret items: hf_to_megatron_cpu: status: verified + bridge_commit: 2d0175b116af6cb785bb8efab7ee053435496266 # pragma: allowlist secret precision: bf16 command: > ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 @@ -56,6 +57,7 @@ items: hf_to_megatron_gpu: status: verified + bridge_commit: 2d185b9619c302281eb1a3219b3c7afaee940afc # pragma: allowlist secret precision: bf16 command: > ./scripts/conversion/convert.sh import --executor slurm --device gpu @@ -63,12 +65,13 @@ items: --hf-revision 093f9f388b31de276ce2de164bdc2081324b9767 --megatron-path work/model-verification/gemma-3-4b-it/imported-megatron --torch-dtype bfloat16 --tp 2 --pp 2 --ep 1 --etp 1 - last_verified: 2026-07-24 + last_verified: 2026-07-30 expected_result: > Exact-revision distributed BF16 import completed with TP2/PP2/EP1/ETP1. All 781 Bridge tensors loaded and all four ranks saved a reloadable checkpoint with the vision tower and multimodal projector on the owning - pipeline stage. + pipeline stage. Independent TP2/PP2 and Gloo reloads retained + FusedAttention and matched all 883 mapped tensors exactly. megatron_to_hf_cpu: status: unverified @@ -92,6 +95,7 @@ items: megatron_to_hf_gpu: status: verified + bridge_commit: 2d185b9619c302281eb1a3219b3c7afaee940afc # pragma: allowlist secret precision: bf16 command: > ./scripts/conversion/convert.sh export --executor slurm --device gpu @@ -101,7 +105,7 @@ items: --hf-path work/model-verification/gemma-3-4b-it/hf-export --torch-dtype bfloat16 --export-weight-dtype bfloat16 --distributed-save --tp 2 --pp 2 --ep 1 --etp 1 - last_verified: 2026-07-24 + last_verified: 2026-07-30 expected_result: > TP2/PP2 distributed export completed without a missing PP owner. Source and export contained the same 883 keys with exact shapes, dtypes, and @@ -123,16 +127,17 @@ items: --image_path work/model-verification/gemma-3-4b-it/candy.JPG --prompt "What animal is on the candy?" --tp 1 --pp 1 --ep 1 --etp 1 - last_verified: 2026-07-24 + last_verified: 2026-07-30 expected_result: > - The exact-revision image-and-text comparison used TP1 without padding. - Token match: true; Hugging Face and Megatron both produced next token ID - 22515 ("Based"). Cosine similarity: 0.990209, passing the required 0.99 - threshold. Maximum and mean absolute logit differences were 6.046631 - and 1.016037 respectively and are report-only diagnostics. + The exact-revision TP1 image-and-text comparison selected Transformer + Engine FusedAttention. Token match: true; Hugging Face and Megatron both + produced next token ID 22515 ("Based"). Cosine similarity: 0.990387, + passing the required 0.99 threshold; maximum and mean absolute logit + differences were 6.234131 and 1.007509 and are report-only diagnostics. inference: status: verified + bridge_commit: 2d185b9619c302281eb1a3219b3c7afaee940afc # pragma: allowlist secret precision: bf16 command: > uv run python -m torch.distributed.run --standalone --nproc_per_node=4 @@ -144,11 +149,11 @@ items: --prompt "What animal is on the candy?" --max_new_tokens 32 --tp 2 --pp 2 --ep 1 --etp 1 - last_verified: 2026-07-24 + last_verified: 2026-07-30 expected_result: |- - One exact-revision TP2/PP2 BF16 greedy multimodal run generated exactly - 32 tokens without reaching EOS. The public input was p-blog/candy.JPG at - documentation-images revision + One exact-revision TP2/PP2 BF16 greedy multimodal run selected + FusedAttention and generated exactly 32 tokens without reaching EOS. The + public input was p-blog/candy.JPG at documentation-images revision 44c3249461c5a5336585d1e31e0a4ce7fce06a5a with SHA-256 fc417c899e94f8df465b7541c5a70f0eebb85c414d06345f0b290c061eccc84c. The JSON-escaped literal completion was: "Based on the image, the animal @@ -224,10 +229,11 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - Gemma 3 VL image-bidirectional attention currently requires an - arbitrary unfused mask. The provider rejects context parallelism - greater than one, and the public recipes disable in-batch THD packing, - so there is no supported exact-model long-context SFT contract. + Gemma 3 VL image-bidirectional attention uses dense local/global + additive biases with FusedAttention. The provider rejects context + parallelism greater than one, and the public recipes disable in-batch + THD packing, so there is no supported exact-model long-context SFT + contract. peft: H100: diff --git a/src/megatron/bridge/models/gemma/gemma3_provider.py b/src/megatron/bridge/models/gemma/gemma3_provider.py index ac59576ce27..2ebfb369549 100644 --- a/src/megatron/bridge/models/gemma/gemma3_provider.py +++ b/src/megatron/bridge/models/gemma/gemma3_provider.py @@ -142,7 +142,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> "MCoreG def gemma3_layer_spec(config) -> ModuleSpec: """Gemma3 custom layer spec.""" - attn_mask_type = AttnMaskType.arbitrary if config.is_vision_language else AttnMaskType.causal + attn_mask_type = AttnMaskType.no_mask if config.is_vision_language else AttnMaskType.causal return ModuleSpec( module=TransformerLayer, submodules=TransformerLayerSubmodules( @@ -187,7 +187,7 @@ def forward( rotary_pos_cos: Optional[Tensor] = None, rotary_pos_sin: Optional[Tensor] = None, rotary_pos_cos_sin: Optional[Tuple[Tensor, Tensor]] = None, - attention_bias: Optional[Tensor] = None, + attention_bias: Optional[Tensor | tuple[Tensor, Tensor]] = None, packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[int] = None, *, @@ -199,8 +199,14 @@ def forward( if _is_local_attn_layer(self.layer_number, self.config.interleaved_attn_pattern): final_rotary_pos_emb = rotary_pos_emb[0] + # Gemma3VLModel supplies (local, global) biases because TE cannot + # combine a finite sliding window with post-scale bias. + if isinstance(attention_bias, tuple): + attention_bias = attention_bias[0] else: final_rotary_pos_emb = rotary_pos_emb[1] + if isinstance(attention_bias, tuple): + attention_bias = attention_bias[1] return super().forward( hidden_states=hidden_states, attention_mask=attention_mask, @@ -235,19 +241,23 @@ def __init__( # Overwrite config.window_size based on layer_number config = copy.deepcopy(config) if _is_local_attn_layer(layer_number, config.interleaved_attn_pattern): - # local attention, (q, k) - window_size = config.window_size - 1 - # The VL mask already blocks future text tokens. Keep the right - # window open so image tokens in the same block stay bidirectional. - right_window_size = window_size if config.is_vision_language else 0 - config.window_size = (window_size, right_window_size) + if config.is_vision_language: + # Transformer Engine cannot combine a finite sliding window + # with post-scale bias. The VL model encodes both constraints + # in its local-layer attention bias instead. + config.window_size = None + else: + # local attention, (q, k) + window_size = config.window_size - 1 + config.window_size = (window_size, 0) else: # global attention config.window_size = None - # The VL model calculates mask manually + # The VL model supplies the full causal/image visibility pattern as an + # attention bias so Transformer Engine can use FusedAttention. if config.is_vision_language: - attn_mask_type = AttnMaskType.arbitrary + attn_mask_type = AttnMaskType.no_mask super().__init__( config=config, diff --git a/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py b/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py index 0e63c088697..95490319c09 100644 --- a/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py +++ b/src/megatron/bridge/models/gemma_vl/gemma3_vl_provider.py @@ -31,7 +31,7 @@ class Gemma3VLModelProvider(Gemma3ModelProvider): # These are model invariants, not user-selectable settings. Marking them as # non-init also drops stale causal/flash values from older checkpoint configs. is_vision_language: bool = field(default=True, init=False) - attention_backend: AttnBackend = field(default=AttnBackend.unfused, init=False) + attention_backend: AttnBackend = field(default=AttnBackend.fused, init=False) # VL models shouldn't scatter embeddings across sequence parallel regions because # the vision embeddings are going to be inserted into the language embeddings. @@ -61,10 +61,10 @@ class Gemma3VLModelProvider(Gemma3ModelProvider): def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Gemma3VLModel: if not self.is_vision_language: raise ValueError("Gemma 3 VL requires is_vision_language=True.") - if self.attention_backend is not AttnBackend.unfused: - raise ValueError("Gemma 3 VL requires the unfused attention backend for arbitrary attention masks.") + if self.attention_backend not in (AttnBackend.auto, AttnBackend.fused): + raise ValueError("Gemma 3 VL requires automatic or fused attention for attention bias.") if self.context_parallel_size != 1: - raise ValueError("Gemma 3 VL does not support context parallelism with arbitrary attention masks.") + raise ValueError("Gemma 3 VL does not support context parallelism with attention bias.") model = Gemma3VLModel(self, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) diff --git a/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py b/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py index 3acdc912829..df2177718fe 100644 --- a/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py +++ b/src/megatron/bridge/models/gemma_vl/modeling_gemma3_vl.py @@ -161,18 +161,22 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # (B, T, D) -> (T, B, D) - # Compute attention mask on FULL sequence (before CP slicing) - # This is needed because image regions need bidirectional attention - attention_mask = self._compute_attention_mask(input_ids) - - # CP slicing: slice embeddings, labels, loss_mask, position_ids, and attention_mask + # Compute the attention bias on the full sequence. The zero/negative + # bias expresses causal text attention plus bidirectional image blocks + # while keeping Transformer Engine on its fused backend. + attention_bias_dtype = inputs_embeds.dtype if inputs_embeds is not None else self.config.params_dtype + attention_biases = self._compute_attention_biases(input_ids, dtype=attention_bias_dtype) + + # CP slicing: slice embeddings, labels, loss_mask, and position_ids. + # Context parallelism is rejected by the provider while attention bias + # slicing is unsupported. # This must happen AFTER vision-text merge so image token positions are correct - inputs_embeds, labels, loss_mask, position_ids, attention_mask = slice_batch_for_context_parallel( + inputs_embeds, labels, loss_mask, position_ids, _ = slice_batch_for_context_parallel( inputs_embeds=inputs_embeds, labels=labels, loss_mask=loss_mask, position_ids=position_ids, - attention_mask=attention_mask, + attention_mask=None, packed_seq_params=packed_seq_params, pg_collection=self.config._pg_collection, ) @@ -188,12 +192,13 @@ def forward( outputs = self.language_model.forward( input_ids=None, position_ids=position_ids, - attention_mask=attention_mask, + attention_mask=None, decoder_input=inputs_embeds, labels=labels, loss_mask=loss_mask, runtime_gather_output=runtime_gather_output, packed_seq_params=packed_seq_params, + extra_block_kwargs={"attention_bias": attention_biases}, ) # Return both outputs and the CP-sliced loss_mask for consistent loss computation return (outputs, loss_mask) @@ -229,14 +234,21 @@ def freeze(self, freeze_language_model: bool, freeze_vision_model: bool, freeze_ for param in module.parameters(): param.requires_grad = False - def _compute_attention_mask( + def _compute_attention_biases( self, input_ids: Optional[torch.Tensor], - ) -> Optional[torch.Tensor]: + *, + dtype: torch.dtype, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Build local/global additive biases for Gemma 3 VL attention.""" if input_ids is None: return None - batch_size, seq_len = input_ids.shape - causal_mask = torch.tril(torch.ones((batch_size, 1, seq_len, seq_len))).to(input_ids.device) + _, seq_len = input_ids.shape + causal_mask = torch.ones( + (1, 1, seq_len, seq_len), + dtype=torch.bool, + device=input_ids.device, + ).tril() image_mask = input_ids == self.config.image_token_id padded_mask = F.pad(image_mask, (1, 0), value=0) @@ -248,9 +260,17 @@ def _compute_attention_mask( kv_block_indices[:, None, :] == q_block_indices.unsqueeze(-1), q_block_indices.unsqueeze(-1) > 0, ) - # See te.DotProductAttention for the requirement of custom mask - attention_mask = ~torch.logical_or(causal_mask, bidirectional_mask.unsqueeze(1)) - return attention_mask + global_allowed_mask = torch.logical_or(causal_mask, bidirectional_mask.unsqueeze(1)) + global_attention_bias = torch.zeros(global_allowed_mask.shape, dtype=dtype, device=input_ids.device) + global_attention_bias.masked_fill_(~global_allowed_mask, torch.finfo(dtype).min) + + window_size = self.config.window_size - 1 + positions = torch.arange(seq_len, device=input_ids.device) + local_window_mask = (positions.unsqueeze(0) - positions.unsqueeze(1)).abs() <= window_size + local_allowed_mask = torch.logical_and(global_allowed_mask, local_window_mask.view(1, 1, seq_len, seq_len)) + local_attention_bias = torch.zeros_like(global_attention_bias) + local_attention_bias.masked_fill_(~local_allowed_mask, torch.finfo(dtype).min) + return local_attention_bias, global_attention_bias @dataclass diff --git a/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py b/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py index 81319727619..56637c446e9 100644 --- a/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py +++ b/src/megatron/bridge/recipes/gemma3_vl/h100/gemma3_vl.py @@ -70,7 +70,7 @@ def gemma3_vl_4b_sft_1gpu_h100_bf16_config() -> ConfigContainer: cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = AttnBackend.unfused + cfg.model.attention_backend = AttnBackend.auto cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -182,7 +182,7 @@ def gemma3_vl_12b_sft_4gpu_h100_bf16_config() -> ConfigContainer: cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = AttnBackend.unfused + cfg.model.attention_backend = AttnBackend.auto cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -294,7 +294,7 @@ def gemma3_vl_27b_sft_16gpu_h100_bf16_config() -> ConfigContainer: cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = AttnBackend.unfused + cfg.model.attention_backend = AttnBackend.auto cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -415,7 +415,7 @@ def gemma3_vl_4b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = AttnBackend.unfused + cfg.model.attention_backend = AttnBackend.auto cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -536,7 +536,7 @@ def gemma3_vl_12b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = AttnBackend.unfused + cfg.model.attention_backend = AttnBackend.auto cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" @@ -657,7 +657,7 @@ def gemma3_vl_27b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - cfg.model.cuda_graph_warmup_steps = 3 # Kernel selections - cfg.model.attention_backend = AttnBackend.unfused + cfg.model.attention_backend = AttnBackend.auto cfg.model.cross_entropy_loss_fusion = True cfg.model.cross_entropy_fusion_impl = "native" diff --git a/tests/functional_tests/fixture_utils.py b/tests/functional_tests/fixture_utils.py new file mode 100644 index 00000000000..3ddf122d263 --- /dev/null +++ b/tests/functional_tests/fixture_utils.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from pathlib import Path + + +DEFAULT_TEST_DATA_ROOT = Path("/home/TestData") +TEST_DATA_ROOT_ENV = "NEMO_TEST_DATA_ROOT" + + +def get_test_data_root() -> Path: + """Return the configured shared TestData root.""" + return Path(os.environ.get(TEST_DATA_ROOT_ENV) or DEFAULT_TEST_DATA_ROOT) + + +def resolve_test_data_file(relative_path: str | Path, fallback: str) -> str: + """Resolve a staged TestData file before using its egress fallback. + + Args: + relative_path: File path relative to the TestData root. + fallback: Value to use when the staged file is unavailable. + + Returns: + The staged file path when it exists, otherwise ``fallback``. + """ + staged_path = get_test_data_root() / relative_path + return str(staged_path) if staged_path.is_file() else fallback diff --git a/tests/functional_tests/launch_scripts/gb200/active/L2_Launch_recipes_gemma_vl.sh b/tests/functional_tests/launch_scripts/gb200/active/L2_Launch_recipes_gemma_vl.sh index 0a48cf63047..c2ab2536557 100755 --- a/tests/functional_tests/launch_scripts/gb200/active/L2_Launch_recipes_gemma_vl.sh +++ b/tests/functional_tests/launch_scripts/gb200/active/L2_Launch_recipes_gemma_vl.sh @@ -29,12 +29,4 @@ uv run python -m torch.distributed.run --nproc_per_node=1 --nnodes=1 \ -m "not pleasefixme" --tb=short -rA \ tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py -k "test_gemma3_vl_finetune_recipes" -# Test Gemma3-VL finetune packed recipes -uv run python -m torch.distributed.run --nproc_per_node=1 --nnodes=1 \ - -m coverage run --data-file=/opt/Megatron-Bridge/.coverage \ - --source=/opt/Megatron-Bridge/ --parallel-mode \ - -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x \ - -m "not pleasefixme" --tb=short -rA \ - tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py -k "test_gemma3_vl_finetune_packed_recipes" - coverage combine -q diff --git a/tests/functional_tests/launch_scripts/h100/active/L2_Launch_recipes_gemma_vl.sh b/tests/functional_tests/launch_scripts/h100/active/L2_Launch_recipes_gemma_vl.sh index 0a48cf63047..c2ab2536557 100755 --- a/tests/functional_tests/launch_scripts/h100/active/L2_Launch_recipes_gemma_vl.sh +++ b/tests/functional_tests/launch_scripts/h100/active/L2_Launch_recipes_gemma_vl.sh @@ -29,12 +29,4 @@ uv run python -m torch.distributed.run --nproc_per_node=1 --nnodes=1 \ -m "not pleasefixme" --tb=short -rA \ tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py -k "test_gemma3_vl_finetune_recipes" -# Test Gemma3-VL finetune packed recipes -uv run python -m torch.distributed.run --nproc_per_node=1 --nnodes=1 \ - -m coverage run --data-file=/opt/Megatron-Bridge/.coverage \ - --source=/opt/Megatron-Bridge/ --parallel-mode \ - -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x \ - -m "not pleasefixme" --tb=short -rA \ - tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py -k "test_gemma3_vl_finetune_packed_recipes" - coverage combine -q diff --git a/tests/functional_tests/test_groups/converter/test_generate_vlm_from_hf.py b/tests/functional_tests/test_groups/converter/test_generate_vlm_from_hf.py index 4ffe45a6c50..f6c724dcc89 100644 --- a/tests/functional_tests/test_groups/converter/test_generate_vlm_from_hf.py +++ b/tests/functional_tests/test_groups/converter/test_generate_vlm_from_hf.py @@ -22,13 +22,19 @@ import pytest +from tests.functional_tests.fixture_utils import resolve_test_data_file + class TestGenerateVLMFromHF: @pytest.mark.run_only_on("GPU") def test_generate_vlm(self): """ - Run distributed VLM generation on a small instruct model with an image URL. + Run distributed VLM generation on a small instruct model with an image fixture. """ + image_path = resolve_test_data_file( + "megatron_bridge/assets/demo.jpeg", + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", + ) cmd = [ "python", "-m", @@ -45,7 +51,7 @@ def test_generate_vlm(self): "--hf_model_path", "Qwen/Qwen2.5-VL-3B-Instruct", "--image_path", - "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", + image_path, "--prompt", "Describe this image.", "--tp", diff --git a/tests/functional_tests/test_groups/data/download_unit_tests_dataset.py b/tests/functional_tests/test_groups/data/download_unit_tests_dataset.py index c3dec1540d7..4f0be34455f 100644 --- a/tests/functional_tests/test_groups/data/download_unit_tests_dataset.py +++ b/tests/functional_tests/test_groups/data/download_unit_tests_dataset.py @@ -28,6 +28,54 @@ import requests from github import Github +from tests.functional_tests.fixture_utils import get_test_data_root + + +DEFAULT_REPO_NAME = "NVIDIA/Megatron-LM" +STAGED_RELEASE_ASSETS = ( + Path("megatron-lm/release-assets/v2.5/datasets.zip"), + Path("megatron-lm/release-assets/v2.5/tokenizers.zip"), +) + + +def extract_asset(asset_path: Path, assets_dir: Path) -> bool: + """Extract a release asset into the writable test data directory.""" + try: + print(f" Extracting {asset_path.name} to {assets_dir}...") + + if asset_path.name.endswith(".zip"): + with zipfile.ZipFile(asset_path, "r") as zip_ref: + zip_ref.extractall(assets_dir) + elif asset_path.name.endswith((".tar.gz", ".tgz")): + with tarfile.open(asset_path, "r:gz") as tar_ref: + tar_ref.extractall(assets_dir) + elif asset_path.name.endswith(".tar"): + with tarfile.open(asset_path, "r") as tar_ref: + tar_ref.extractall(assets_dir) + else: + print(f" Warning: Unknown file type for {asset_path.name}, skipping extraction") + return False + + print(f" Successfully extracted to {assets_dir}") + return True + + except Exception as e: + print(f" Error extracting {asset_path.name}: {e}") + return False + + +def extract_staged_release_assets(repo_name: str, assets_dir: Path) -> bool: + """Extract staged Megatron-LM v2.5 assets when all of them are available.""" + if repo_name != DEFAULT_REPO_NAME: + return False + + staged_assets = tuple(get_test_data_root() / relative_path for relative_path in STAGED_RELEASE_ASSETS) + if not all(asset_path.is_file() for asset_path in staged_assets): + return False + + print(f"Using staged release assets from {staged_assets[0].parent}") + return all(extract_asset(asset_path, assets_dir) for asset_path in staged_assets) + def download_and_extract_asset(asset_url: str, asset_name: str, assets_dir: Path) -> bool: """ @@ -41,45 +89,29 @@ def download_and_extract_asset(asset_url: str, asset_name: str, assets_dir: Path Returns: bool: True if successful, False otherwise """ + temp_file = assets_dir / asset_name try: # Download the asset print(f" Downloading {asset_name}...") - response = requests.get(asset_url, stream=True) + response = requests.get(asset_url, stream=True, timeout=60) response.raise_for_status() # Save to temporary file - temp_file = assets_dir / asset_name with open(temp_file, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) - print(f" Extracting {asset_name} to {assets_dir}...") - - # Extract based on file type - if asset_name.endswith(".zip"): - with zipfile.ZipFile(temp_file, "r") as zip_ref: - zip_ref.extractall(assets_dir) - elif asset_name.endswith((".tar.gz", ".tgz")): - with tarfile.open(temp_file, "r:gz") as tar_ref: - tar_ref.extractall(assets_dir) - elif asset_name.endswith(".tar"): - with tarfile.open(temp_file, "r") as tar_ref: - tar_ref.extractall(assets_dir) - else: - print(f" Warning: Unknown file type for {asset_name}, skipping extraction") - return False - - # Clean up temporary file - temp_file.unlink() - print(f" Successfully extracted to {assets_dir}") - return True + return extract_asset(temp_file, assets_dir) except Exception as e: print(f" Error downloading/extracting {asset_name}: {e}") return False + finally: + if temp_file.is_file(): + temp_file.unlink() -def get_oldest_release_and_assets(repo_name: str = "NVIDIA/Megatron-LM", assets_dir: str = "assets") -> None: +def get_oldest_release_and_assets(repo_name: str = DEFAULT_REPO_NAME, assets_dir: str = "assets") -> None: """ Fetch the oldest release of a GitHub repository and list its assets. @@ -88,12 +120,15 @@ def get_oldest_release_and_assets(repo_name: str = "NVIDIA/Megatron-LM", assets_ assets_dir: Directory to extract assets to """ try: - # Initialize GitHub client - token = os.getenv("GH_TOKEN", None) - if token is None: - raise ValueError("GH_TOKEN environment variable is not set") + assets_path = Path(assets_dir) + assets_path.mkdir(parents=True, exist_ok=True) + + if extract_staged_release_assets(repo_name, assets_path): + return - g = Github(login_or_token=token) + # Initialize an authenticated GitHub client when a token is available. + token = os.getenv("GH_TOKEN", None) + g = Github(login_or_token=token) if token else Github() # Get the repository repo = g.get_repo(repo_name) @@ -161,9 +196,6 @@ def get_oldest_release_and_assets(repo_name: str = "NVIDIA/Megatron-LM", assets_ print("-" * 80) print("Downloading and extracting assets...") - # Create assets directory - assets_path = Path(assets_dir) - assets_path.mkdir(parents=True, exist_ok=True) print(f"Created assets directory: {assets_path.absolute()}") successful_downloads = 0 @@ -182,7 +214,7 @@ def get_oldest_release_and_assets(repo_name: str = "NVIDIA/Megatron-LM", assets_ @click.command() -@click.option("--repo", default="NVIDIA/Megatron-LM", help="GitHub repository name (format: owner/repo)") +@click.option("--repo", default=DEFAULT_REPO_NAME, help="GitHub repository name (format: owner/repo)") @click.option("--assets-dir", default="assets", help="Directory to extract assets to") def main(repo, assets_dir): """Fetch the oldest release of a GitHub repository and download its assets.""" diff --git a/tests/functional_tests/test_groups/models/qwen_audio/test_qwen2_audio_generation.py b/tests/functional_tests/test_groups/models/qwen_audio/test_qwen2_audio_generation.py index de8012828ab..44dc6a324ec 100644 --- a/tests/functional_tests/test_groups/models/qwen_audio/test_qwen2_audio_generation.py +++ b/tests/functional_tests/test_groups/models/qwen_audio/test_qwen2_audio_generation.py @@ -34,6 +34,8 @@ import torch from transformers import AutoTokenizer, Qwen2AudioConfig, Qwen2AudioForConditionalGeneration +from tests.functional_tests.fixture_utils import resolve_test_data_file + HF_QWEN2_AUDIO_TOY_MODEL_CONFIG = { "architectures": ["Qwen2AudioForConditionalGeneration"], @@ -150,6 +152,11 @@ def test_qwen2_audio_generation(self, qwen2_audio_toy_model_path): Args: qwen2_audio_toy_model_path: Path to the toy Qwen2 Audio model (from fixture) """ + audio_path = resolve_test_data_file( + "megatron_bridge/assets/glass-breaking-151256.mp3", + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3", + ) + audio_option = "--audio_url" if audio_path.startswith(("http://", "https://")) else "--audio_path" cmd = [ sys.executable, "-m", @@ -157,7 +164,7 @@ def test_qwen2_audio_generation(self, qwen2_audio_toy_model_path): "--nproc_per_node=2", "examples/conversion/hf_to_megatron_generate_audio_lm.py", f"--hf_model_path={qwen2_audio_toy_model_path}", - "--audio_url=https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3", + f"{audio_option}={audio_path}", "--prompt=What's that sound?", "--tp=2", "--max_new_tokens=50", diff --git a/tests/functional_tests/test_groups/models/qwen_vl/test_qwen3_vl_generation.py b/tests/functional_tests/test_groups/models/qwen_vl/test_qwen3_vl_generation.py index f9851705d40..0ee02a7ec18 100644 --- a/tests/functional_tests/test_groups/models/qwen_vl/test_qwen3_vl_generation.py +++ b/tests/functional_tests/test_groups/models/qwen_vl/test_qwen3_vl_generation.py @@ -42,6 +42,8 @@ ) from transformers.models.qwen3_vl import Qwen3VLConfig +from tests.functional_tests.fixture_utils import resolve_test_data_file + HF_QWEN3_VL_TOY_MODEL_CONFIG = { "architectures": ["Qwen3VLForConditionalGeneration"], @@ -325,6 +327,10 @@ def test_qwen3_vl_8b_image_generation(self, qwen3_vl_toy_model_path): Args: qwen3_vl_toy_model_path: Path to the toy Qwen3 VL model (from fixture) """ + image_path = resolve_test_data_file( + "megatron_bridge/assets/demo.jpeg", + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", + ) cmd = [ "python", "-m", @@ -332,7 +338,7 @@ def test_qwen3_vl_8b_image_generation(self, qwen3_vl_toy_model_path): "--nproc_per_node=2", "examples/conversion/hf_to_megatron_generate_vlm.py", f"--hf_model_path={qwen3_vl_toy_model_path}", - "--image_path=https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", + f"--image_path={image_path}", "--prompt=Describe this image.", ] diff --git a/tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py b/tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py index b158ed576a8..c2b5c9dfffe 100644 --- a/tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py +++ b/tests/functional_tests/test_groups/recipes/test_gemma3_vl_recipes_finetune.py @@ -31,16 +31,6 @@ ), ] -GEMMA3_VL_FINETUNE_PACKED_RECIPES = [ - # Small model with packed sequences, only use 2 layers - ( - gemma3_vl_4b_sft_config, - "gemma3_vl_4b_sft_packed", - {"tensor_model_parallel_size": 1, "pipeline_model_parallel_size": 1, "num_layers": 2}, - {"enable_in_batch_packing": True}, - ), -] - class TestGemma3VLRecipes: """Test class for Gemma3-VL recipe functional tests.""" @@ -50,19 +40,3 @@ class TestGemma3VLRecipes: def test_gemma3_vl_finetune_recipes(self, config_func, recipe_name, model_overrides, tmp_path): """Functional test for Gemma3-VL recipes with appropriate parallelism configurations.""" run_pretrain_vl_recipe_test(config_func, recipe_name, tmp_path, model_overrides=model_overrides) - - @pytest.mark.run_only_on("GPU") - @pytest.mark.parametrize( - "config_func,recipe_name,model_overrides,dataset_overrides", GEMMA3_VL_FINETUNE_PACKED_RECIPES - ) - def test_gemma3_vl_finetune_packed_recipes( - self, config_func, recipe_name, model_overrides, dataset_overrides, tmp_path - ): - """Functional test for Gemma3-VL recipes with packed sequences enabled.""" - run_pretrain_vl_recipe_test( - config_func, - recipe_name, - tmp_path, - model_overrides=model_overrides, - dataset_overrides=dataset_overrides, - ) diff --git a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py index e87f08bdcb7..17f850a5b61 100644 --- a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py +++ b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_modeling.py @@ -12,29 +12,99 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import Mock +from types import SimpleNamespace +from unittest.mock import Mock, patch import torch +import torch.nn as nn from megatron.bridge.models.gemma_vl.modeling_gemma3_vl import Gemma3VLModel -def test_attention_mask_is_available_on_later_pipeline_stages(): - """Later PP stages must receive the same image-bidirectional mask.""" +def test_attention_bias_is_available_on_later_pipeline_stages(): + """Later PP stages must receive the same image-bidirectional bias.""" model = object.__new__(Gemma3VLModel) - model.config = Mock(image_token_id=99) + model.config = Mock(image_token_id=99, window_size=512) model.pre_process = False input_ids = torch.tensor([[1, 99, 99, 2]]) - mask = model._compute_attention_mask(input_ids) + biases = model._compute_attention_biases(input_ids, dtype=torch.bfloat16) - assert mask is not None - assert mask[0, 0, 1, 2].item() is False - assert mask[0, 0, 0, 3].item() is True + assert biases is not None + local_bias, global_bias = biases + assert local_bias.dtype is torch.bfloat16 + assert global_bias.dtype is torch.bfloat16 + assert local_bias[0, 0, 1, 2].item() == 0.0 + assert global_bias[0, 0, 1, 2].item() == 0.0 + assert local_bias[0, 0, 0, 3].item() == torch.finfo(torch.bfloat16).min + assert global_bias[0, 0, 0, 3].item() == torch.finfo(torch.bfloat16).min -def test_attention_mask_requires_input_ids(): - """Mask construction is skipped only when token IDs are unavailable.""" +def test_attention_bias_requires_input_ids(): + """Bias construction is skipped only when token IDs are unavailable.""" model = object.__new__(Gemma3VLModel) - assert model._compute_attention_mask(None) is None + assert model._compute_attention_biases(None, dtype=torch.bfloat16) is None + + +def test_attention_bias_keeps_separate_image_blocks_isolated(): + """Only tokens in the same contiguous image block can see each other bidirectionally.""" + model = object.__new__(Gemma3VLModel) + model.config = Mock(image_token_id=99, window_size=512) + input_ids = torch.tensor([[1, 99, 99, 2, 99, 99, 3]]) + + biases = model._compute_attention_biases(input_ids, dtype=torch.float32) + + assert biases is not None + local_bias, global_bias = biases + for bias in (local_bias, global_bias): + assert bias[0, 0, 1, 2].item() == 0.0 + assert bias[0, 0, 4, 5].item() == 0.0 + assert bias[0, 0, 1, 5].item() == torch.finfo(torch.float32).min + assert bias[0, 0, 3, 6].item() == torch.finfo(torch.float32).min + assert torch.count_nonzero(bias[0, 0].tril()).item() == 0 + + +def test_local_attention_bias_includes_sliding_window(): + """Local layers apply the sliding window in bias while global layers do not.""" + model = object.__new__(Gemma3VLModel) + model.config = Mock(image_token_id=99, window_size=3) + input_ids = torch.tensor([[1, 99, 99, 2, 3]]) + + biases = model._compute_attention_biases(input_ids, dtype=torch.float32) + + assert biases is not None + local_bias, global_bias = biases + assert local_bias[0, 0, 4, 0].item() == torch.finfo(torch.float32).min + assert global_bias[0, 0, 4, 0].item() == 0.0 + assert local_bias[0, 0, 1, 2].item() == 0.0 + + +def test_forward_passes_attention_bias_to_language_model(): + """The language model receives fused-attention bias rather than a custom mask.""" + model = object.__new__(Gemma3VLModel) + nn.Module.__init__(model) + model.config = SimpleNamespace( + image_token_id=99, + window_size=512, + params_dtype=torch.bfloat16, + sequence_parallel=False, + _pg_collection=None, + ) + model.pre_process = False + model.language_model = SimpleNamespace(forward=Mock(return_value=torch.tensor(1.0))) + input_ids = torch.tensor([[1, 99, 99, 2]]) + + with patch( + "megatron.bridge.models.gemma_vl.modeling_gemma3_vl.slice_batch_for_context_parallel", + return_value=(None, None, None, None, None), + ): + model(input_ids=input_ids) + + call_kwargs = model.language_model.forward.call_args.kwargs + assert call_kwargs["attention_mask"] is None + local_bias, global_bias = call_kwargs["extra_block_kwargs"]["attention_bias"] + assert local_bias[0, 0, 1, 2].item() == 0.0 + assert global_bias[0, 0, 1, 2].item() == 0.0 + assert local_bias[0, 0, 0, 3].item() == torch.finfo(torch.bfloat16).min + assert global_bias[0, 0, 0, 3].item() == torch.finfo(torch.bfloat16).min diff --git a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py index 726341b7c7b..c62472f0891 100644 --- a/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py +++ b/tests/unit_tests/models/gemma_vl/test_gemma3_vl_provider.py @@ -20,7 +20,11 @@ from megatron.core.transformer.enums import AttnBackend, AttnMaskType from transformers import SiglipVisionConfig -from megatron.bridge.models.gemma.gemma3_provider import Gemma3TEDotProductAttention, gemma3_layer_spec +from megatron.bridge.models.gemma.gemma3_provider import ( + Gemma3SelfAttention, + Gemma3TEDotProductAttention, + gemma3_layer_spec, +) from megatron.bridge.models.gemma_vl.gemma3_vl_provider import Gemma3VLModelProvider from megatron.bridge.models.gemma_vl.modeling_gemma3_vl import Gemma3VLMultimodalProjectorConfig from megatron.bridge.training.utils.config_utils import apply_run_config_backward_compat @@ -70,7 +74,7 @@ def test_gemma3_vl_vl_specific_defaults(self): # Check VL-specific defaults assert provider.is_vision_language is True - assert provider.attention_backend is AttnBackend.unfused + assert provider.attention_backend is AttnBackend.fused assert provider.scatter_embedding_sequence_parallel is False assert isinstance(provider.vision_config, SiglipVisionConfig) assert isinstance(provider.vision_projector_config, Gemma3VLMultimodalProjectorConfig) @@ -88,8 +92,8 @@ def test_gemma3_vl_vl_specific_defaults(self): assert provider.freeze_vision_model is False assert provider.freeze_vision_projection is False - def test_gemma3_vl_attention_consumes_custom_mask(self): - """Test that Gemma3 VL attention selects the arbitrary-mask path.""" + def test_gemma3_vl_attention_consumes_attention_bias(self): + """Test that Gemma3 VL attention selects the fused attention-bias path.""" provider = Gemma3VLModelProvider( num_layers=28, hidden_size=2560, @@ -98,7 +102,7 @@ def test_gemma3_vl_attention_consumes_custom_mask(self): ) layer_spec = gemma3_layer_spec(provider) - assert layer_spec.submodules.self_attention.params["attn_mask_type"] is AttnMaskType.arbitrary + assert layer_spec.submodules.self_attention.params["attn_mask_type"] is AttnMaskType.no_mask with patch.object(Gemma3TEDotProductAttention.__bases__[0], "__init__", return_value=None) as init: Gemma3TEDotProductAttention( @@ -108,8 +112,40 @@ def test_gemma3_vl_attention_consumes_custom_mask(self): attention_type="self", ) - assert init.call_args.kwargs["attn_mask_type"] is AttnMaskType.arbitrary - assert init.call_args.kwargs["config"].window_size == (1023, 1023) + assert init.call_args.kwargs["attn_mask_type"] is AttnMaskType.no_mask + assert init.call_args.kwargs["config"].window_size is None + + def test_gemma3_vl_attention_selects_per_layer_bias(self): + """Test that local and global layers receive their matching precomputed bias.""" + provider = Gemma3VLModelProvider( + num_layers=28, + hidden_size=2560, + num_attention_heads=10, + ) + attention = object.__new__(Gemma3SelfAttention) + object.__setattr__(attention, "config", provider) + local_bias = torch.tensor(1.0) + global_bias = torch.tensor(2.0) + rotary_pos_emb = torch.zeros(2, 1) + + with patch.object(Gemma3SelfAttention.__bases__[0], "forward", return_value=None) as forward: + object.__setattr__(attention, "layer_number", 1) + attention.forward( + hidden_states=torch.zeros(1), + attention_mask=None, + rotary_pos_emb=rotary_pos_emb, + attention_bias=(local_bias, global_bias), + ) + assert forward.call_args.kwargs["attention_bias"] is local_bias + + object.__setattr__(attention, "layer_number", 6) + attention.forward( + hidden_states=torch.zeros(1), + attention_mask=None, + rotary_pos_emb=rotary_pos_emb, + attention_bias=(local_bias, global_bias), + ) + assert forward.call_args.kwargs["attention_bias"] is global_bias def test_legacy_checkpoint_attention_settings_are_dropped(self): """Test that old checkpoint values cannot disable the VL mask.""" @@ -126,7 +162,7 @@ def test_legacy_checkpoint_attention_settings_are_dropped(self): provider = instantiate(sanitized_config) assert provider.is_vision_language is True - assert provider.attention_backend is AttnBackend.unfused + assert provider.attention_backend is AttnBackend.fused def test_gemma3_vl_custom_vision_config(self): """Test Gemma3VLModelProvider with custom vision configuration.""" @@ -177,6 +213,20 @@ def test_gemma3_vl_rejects_causal_override(self): with pytest.raises(ValueError, match="requires is_vision_language=True"): provider.provide() + def test_gemma3_vl_accepts_auto_attention_override(self): + """Test that automatic backend selection preserves the VL attention bias.""" + provider = Gemma3VLModelProvider( + num_layers=28, + hidden_size=2560, + num_attention_heads=10, + ) + provider.attention_backend = AttnBackend.auto + + with patch("megatron.bridge.models.gemma_vl.gemma3_vl_provider.Gemma3VLModel") as model: + provider.provide() + + model.assert_called_once() + def test_gemma3_vl_rejects_flash_attention_override(self): """Test that runtime overrides cannot select a backend that drops the VL mask.""" provider = Gemma3VLModelProvider( @@ -186,7 +236,7 @@ def test_gemma3_vl_rejects_flash_attention_override(self): ) provider.attention_backend = AttnBackend.flash - with pytest.raises(ValueError, match="requires the unfused attention backend"): + with pytest.raises(ValueError, match="requires automatic or fused attention"): provider.provide() def test_gemma3_vl_custom_vision_projector_config(self): diff --git a/tests/unit_tests/recipes/test_gemma3_vl_recipes.py b/tests/unit_tests/recipes/test_gemma3_vl_recipes.py index 7a6dfbdf945..fe2ba587d80 100644 --- a/tests/unit_tests/recipes/test_gemma3_vl_recipes.py +++ b/tests/unit_tests/recipes/test_gemma3_vl_recipes.py @@ -97,7 +97,7 @@ def _assert_basic_config(cfg): assert cfg.tokenizer is not None assert cfg.checkpoint is not None assert cfg.rng is not None - assert cfg.model.attention_backend is AttnBackend.unfused + assert cfg.model.attention_backend is AttnBackend.auto assert cfg.dataset.enable_in_batch_packing is False assert cfg.train.global_batch_size >= 1 diff --git a/tests/unit_tests/test_download_unit_tests_dataset.py b/tests/unit_tests/test_download_unit_tests_dataset.py new file mode 100644 index 00000000000..3f44884171e --- /dev/null +++ b/tests/unit_tests/test_download_unit_tests_dataset.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import MagicMock +from zipfile import ZipFile + +from tests.functional_tests import fixture_utils +from tests.functional_tests.test_groups.data import download_unit_tests_dataset + + +def test_get_oldest_release_prefers_staged_assets(monkeypatch, tmp_path): + staged_root = tmp_path / "staged" + output_root = tmp_path / "output" + for relative_path in download_unit_tests_dataset.STAGED_RELEASE_ASSETS: + staged_asset = staged_root / relative_path + staged_asset.parent.mkdir(parents=True, exist_ok=True) + with ZipFile(staged_asset, "w") as archive: + archive.writestr(f"{staged_asset.stem}/fixture.txt", staged_asset.name) + + monkeypatch.setenv(fixture_utils.TEST_DATA_ROOT_ENV, str(staged_root)) + github = MagicMock(side_effect=AssertionError("GitHub fallback should not be used")) + monkeypatch.setattr(download_unit_tests_dataset, "Github", github) + + download_unit_tests_dataset.get_oldest_release_and_assets(assets_dir=str(output_root)) + + assert (output_root / "datasets" / "fixture.txt").read_text() == "datasets.zip" + assert (output_root / "tokenizers" / "fixture.txt").read_text() == "tokenizers.zip" + github.assert_not_called() + + +def test_get_oldest_release_falls_back_without_github_token(monkeypatch, tmp_path): + release = SimpleNamespace( + tag_name="v2.5", + title="Megatron-LM v2.5", + created_at=0, + published_at=0, + draft=False, + prerelease=False, + html_url="https://github.com/NVIDIA/Megatron-LM/releases/tag/v2.5", + body="", + get_assets=lambda: [], + ) + repo = SimpleNamespace( + full_name="NVIDIA/Megatron-LM", + description="Megatron-LM", + html_url="https://github.com/NVIDIA/Megatron-LM", + get_releases=lambda: [release], + ) + github = MagicMock() + github.return_value.get_repo.return_value = repo + + monkeypatch.setenv(fixture_utils.TEST_DATA_ROOT_ENV, str(tmp_path / "missing")) + monkeypatch.delenv("GH_TOKEN", raising=False) + monkeypatch.setattr(download_unit_tests_dataset, "Github", github) + + download_unit_tests_dataset.get_oldest_release_and_assets(assets_dir=str(tmp_path / "output")) + + github.assert_called_once_with() diff --git a/tests/unit_tests/test_fixture_utils.py b/tests/unit_tests/test_fixture_utils.py new file mode 100644 index 00000000000..d43915f044b --- /dev/null +++ b/tests/unit_tests/test_fixture_utils.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tests.functional_tests.fixture_utils import TEST_DATA_ROOT_ENV, resolve_test_data_file + + +def test_resolve_test_data_file_prefers_staged_file(monkeypatch, tmp_path): + relative_path = "megatron_bridge/assets/demo.jpeg" + staged_path = tmp_path / relative_path + staged_path.parent.mkdir(parents=True) + staged_path.write_bytes(b"fixture") + monkeypatch.setenv(TEST_DATA_ROOT_ENV, str(tmp_path)) + + assert resolve_test_data_file(relative_path, "https://example.com/demo.jpeg") == str(staged_path) + + +def test_resolve_test_data_file_uses_fallback_when_staged_file_is_missing(monkeypatch, tmp_path): + fallback = "https://example.com/demo.jpeg" + monkeypatch.setenv(TEST_DATA_ROOT_ENV, str(tmp_path)) + + assert resolve_test_data_file("megatron_bridge/assets/demo.jpeg", fallback) == fallback From c395bf3f4b51bb1008a6380e5ec8bb2f6c362777 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 6 Aug 2026 14:21:47 -0700 Subject: [PATCH 38/85] fix(recipe): backport #5386 batch-size alignment to r0.6.0 (#5393) Signed-off-by: Chen Cui --- .../models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh | 1 + .../models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh | 1 + 2 files changed, 2 insertions(+) diff --git a/examples/models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh b/examples/models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh index 720ef00e4d2..258c434d3dd 100755 --- a/examples/models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh +++ b/examples/models/nemotron/nemotron_3_omni/slurm_peft_valor32k_avqa.sh @@ -133,6 +133,7 @@ CLI_OVERRIDES="\ logger.tensorboard_dir=$OUTPUT_DIR/tb_logs \ dataset.path=$ENERGON_PATH \ dataset.seq_length=$SEQ_LENGTH \ + dataset.micro_batch_size=$MICRO_BATCH_SIZE \ dataset.enable_in_batch_packing=$PACKED_SEQ \ model.seq_length=$SEQ_LENGTH \ model.tensor_model_parallel_size=$TP \ diff --git a/examples/models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh b/examples/models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh index dbdafd8a88d..6822e7e19d5 100755 --- a/examples/models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh +++ b/examples/models/nemotron/nemotron_3_omni/slurm_sft_valor32k_avqa.sh @@ -129,6 +129,7 @@ CLI_OVERRIDES="\ logger.tensorboard_dir=$OUTPUT_DIR/tb_logs \ dataset.path=$ENERGON_PATH \ dataset.seq_length=$SEQ_LENGTH \ + dataset.micro_batch_size=$MICRO_BATCH_SIZE \ dataset.enable_in_batch_packing=$PACKED_SEQ \ model.seq_length=$SEQ_LENGTH \ model.tensor_model_parallel_size=$TP \ From c6ed307b10c3021771be3ae656d8b63c716246c9 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 09:27:52 +0200 Subject: [PATCH 39/85] cp: `fix(perf): mount common performance scripts (5381)` into `r0.6.0` (#5404) Signed-off-by: Raghav Hrishikeshan Mukundan Signed-off-by: NeMo Bot Co-authored-by: Raghav Hrishikeshan Mukundan <102543536+rhmukundan@users.noreply.github.com> --- scripts/performance/setup_experiment.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/performance/setup_experiment.py b/scripts/performance/setup_experiment.py index fe0288d0158..aaa89b46ddc 100755 --- a/scripts/performance/setup_experiment.py +++ b/scripts/performance/setup_experiment.py @@ -63,6 +63,7 @@ SCRIPT_DIR = Path(__file__).parent.resolve() REPO_ROOT = SCRIPT_DIR.parents[1] BRIDGE_SOURCE_DIR = REPO_ROOT / "src" / "megatron" / "bridge" +COMMON_SCRIPTS_DIR = REPO_ROOT / "scripts" / "common" ENTRYPOINT_BOOTSTRAP = "bootstrap.py" logging.basicConfig(level=logging.DEBUG) @@ -641,6 +642,7 @@ def main( [ f"{run_script_path}:{run_script_path}", f"{SCRIPT_DIR}:{SCRIPT_DIR}", + f"{COMMON_SCRIPTS_DIR}:{COMMON_SCRIPTS_DIR}", _bridge_source_mount(in_container_script_dir), ] ) From 095e8cb08ecdc751c92b449c052871df4668d18d Mon Sep 17 00:00:00 2001 From: Ajay Date: Fri, 7 Aug 2026 14:09:37 -0700 Subject: [PATCH 40/85] chore(deps): update mlflow to 3.15.1 (#5419) Signed-off-by: svcnemo-autobot Co-authored-by: svcnemo-autobot --- pyproject.toml | 4 ++-- uv.lock | 22 +++++++++++----------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index aa26cb2f82d..5c3bfacfc0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,7 +95,7 @@ dependencies = [ "flash-linear-attention", "timm", "open-clip-torch>=3.2.0", - "mlflow>=3.15.0", # To address CVE-2025-15031 + "mlflow>=3.15.1", # To address CVE-2025-15031 "comet-ml>=3.50.0", "torch>=2.6.0", "flashinfer-python==0.6.8.post1", # keep flashinfer-python and flashinfer-cubin in lockstep: @@ -130,7 +130,7 @@ override-dependencies = [ "torchvision; sys_platform == 'never'", "triton; sys_platform == 'never'", "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@4329ff84bfbdaa778a33cba02a15fb0807c64689", - "mlflow>=3.15.0", # To address CVE-2025-15031 + "mlflow>=3.15.1", # To address CVE-2025-15031 "cachetools>=5.0.0", "cryptography>=43.0.0,<47", "nvidia-modelopt==0.46.0rc1", diff --git a/uv.lock b/uv.lock index 3248010ca16..926a1ad5380 100644 --- a/uv.lock +++ b/uv.lock @@ -16,7 +16,7 @@ overrides = [ { name = "imageio-ffmpeg", marker = "sys_platform == 'never'" }, { name = "langchain", specifier = ">=0.3.28" }, { name = "langchain-core", specifier = ">=0.3.80" }, - { name = "mlflow", specifier = ">=3.15.0" }, + { name = "mlflow", specifier = ">=3.15.1" }, { name = "nvidia-cudnn-frontend", specifier = "==1.26.0" }, { name = "nvidia-cutlass-dsl", extras = ["cu13"], marker = "sys_platform == 'never'" }, { name = "nvidia-modelopt", specifier = "==0.46.0rc1", index = "https://pypi.nvidia.com/" }, @@ -2197,7 +2197,7 @@ requires-dist = [ { name = "mamba-ssm", marker = "extra == 'ssm'" }, { name = "megatron-core", extras = ["dev", "mlm"], editable = "3rdparty/Megatron-LM" }, { name = "mistral-common", specifier = ">=1.10.0" }, - { name = "mlflow", specifier = ">=3.15.0" }, + { name = "mlflow", specifier = ">=3.15.1" }, { name = "nemo-run", marker = "extra == 'recipes'" }, { name = "nvdlfw-inspect", marker = "extra == 'tensor-inspect'", specifier = "==0.2.1" }, { name = "nvidia-resiliency-ext" }, @@ -2484,7 +2484,7 @@ wheels = [ [[package]] name = "mlflow" -version = "3.15.0" +version = "3.15.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -2508,14 +2508,14 @@ dependencies = [ { name = "sqlalchemy" }, { name = "waitress", marker = "sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/af/9f/249b80849aae0919aee3f590003419c394f0efe92ae98546ea626be86e1e/mlflow-3.15.0.tar.gz", hash = "sha256:c53d7933391fefda2ac0f6b942d1744a02b5b65da74a93cc125247f4c4da57c9", size = 10395873, upload-time = "2026-07-31T07:06:00.211Z" } +sdist = { url = "https://files.pythonhosted.org/packages/51/01/aff6c6e1ee571769688e926d2cf55ae9ddd300f95541905a0a82ea1599c7/mlflow-3.15.1.tar.gz", hash = "sha256:6177b02c442d7d6ab3ad752daa826db9cf993e255450fb4b00ea8b7b5539c01d", size = 10396741, upload-time = "2026-08-03T09:29:20.073Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0e/7a/3da40be70e7ca0c4d8c2f98edfcb275ef2898d8f4da565f55a9fff95fc78/mlflow-3.15.0-py3-none-any.whl", hash = "sha256:3ae54c7f91a6b98ae9360aaeda9b31eb7930571c2690491a7b550c84f795a709", size = 11188317, upload-time = "2026-07-31T07:05:57.492Z" }, + { url = "https://files.pythonhosted.org/packages/5a/ff/24e530a328b3b2c9d004e94cdd59ea5ec9c052c4935799f675d7bb1f1d2c/mlflow-3.15.1-py3-none-any.whl", hash = "sha256:c91f78d304d8ef825869ea07e638c73a2850660f0d9f098993bacecc359f8a75", size = 11189154, upload-time = "2026-08-03T09:29:17.718Z" }, ] [[package]] name = "mlflow-skinny" -version = "3.15.0" +version = "3.15.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, @@ -2539,14 +2539,14 @@ dependencies = [ { name = "typing-extensions" }, { name = "uvicorn" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/53/a1/04489b8f10ec642cc5d16d8bcf0e3626cbe6b8924940b510de9e3cb37ac7/mlflow_skinny-3.15.0.tar.gz", hash = "sha256:6faad0bb5b0b8adc08550c6cd2ab5a87f9e7d7fe8814278e22b79f9cac99d943", size = 3034359, upload-time = "2026-07-31T07:04:58.359Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/d7/f71b2607ef378b8481a1cea08463a6025b9ca831b4607afc16c75df57872/mlflow_skinny-3.15.1.tar.gz", hash = "sha256:8c53c85bf0fe6e9ba8aa72f6b8675c774d1ff55b00735a5cd33a510803905237", size = 3035284, upload-time = "2026-08-03T09:29:35.266Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/28/ca/b0194962614876abb8b9a4d38815b4c6bf5c6f626b5dd3bcd321048755fc/mlflow_skinny-3.15.0-py3-none-any.whl", hash = "sha256:fc609ea5f7325a76d6f905489d8be73e85bcba0eef64144a763cf446de83192e", size = 3612199, upload-time = "2026-07-31T07:04:56.564Z" }, + { url = "https://files.pythonhosted.org/packages/f3/9e/4e8138583321f6dfbaa35360f46056b4ecc06d505d130c03a5d81aec5620/mlflow_skinny-3.15.1-py3-none-any.whl", hash = "sha256:b62056806d0afc425e8948233a3646b0b3af504702ff3b334b6395f48b22b832", size = 3613028, upload-time = "2026-08-03T09:29:33.54Z" }, ] [[package]] name = "mlflow-tracing" -version = "3.15.0" +version = "3.15.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, @@ -2558,9 +2558,9 @@ dependencies = [ { name = "protobuf" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7b/47/8a3c4f56199aadce644256684060df09f0af6ecaec460860657f2cb13b10/mlflow_tracing-3.15.0.tar.gz", hash = "sha256:e4e73243eed04d0771954b6fa17014dd21dabc67af17d4937fac2b62ec476e95", size = 1500398, upload-time = "2026-07-31T07:05:54.445Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/a6/76964d8f265be984838a13b77656b047686c492b61b3f46ca2ef29ef7325/mlflow_tracing-3.15.1.tar.gz", hash = "sha256:ba1c4d873151c0ceeab37f53daf0997feea18a1c092358995072dc53e1944a29", size = 1501188, upload-time = "2026-08-03T09:26:15.541Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/76/3c/5839690c3fe68f579074080832cba413f988249ddd350dcf7dfc3be1b874/mlflow_tracing-3.15.0-py3-none-any.whl", hash = "sha256:3762f654a858474c7988c0a2a5246765b173fa2b6a2bb8f035c115db1a8e3a6a", size = 1784776, upload-time = "2026-07-31T07:05:52.937Z" }, + { url = "https://files.pythonhosted.org/packages/f1/c7/13db54bd1b6525a0f09d880260d9dcf8072ae89155dae7a74c9e75a1ba03/mlflow_tracing-3.15.1-py3-none-any.whl", hash = "sha256:3cdef1f1675fe2c7c9f409e132439d65b063e37455b338027bfb2c0f986bebca", size = 1785591, upload-time = "2026-08-03T09:26:13.511Z" }, ] [[package]] From 2f7c88869c101149c49acdbf6f2c060cd62b756b Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 23:38:20 +0200 Subject: [PATCH 41/85] cp: `feat(model): add GLM-5.2 verification card (5034)` into `r0.6.0` (#5301) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Hollow Man Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 Co-authored-by: yaoyu-33 --- README.md | 2 +- docs/models/README.md | 2 +- docs/models/glm/glm5.md | 6 +- docs/models/glm/index.md | 6 +- .../hf_to_megatron_generate_text.py | 57 +- .../model_verification_cards/glm5-2/card.yaml | 519 ++++++++++++++++++ examples/models/glm/glm5/README.md | 2 +- examples/models/glm/glm5/slurm_inference.sh | 4 +- scripts/inference/README.md | 5 +- scripts/inference/setup_inference.py | 18 +- .../create-model-verification-card/SKILL.md | 14 +- .../scripts/validate_card.py | 7 + .../bridge/models/glm_moe_dsa/glm5_bridge.py | 4 + src/megatron/bridge/recipes/glm/__init__.py | 10 + .../bridge/recipes/glm/gb200/__init__.py | 23 + src/megatron/bridge/recipes/glm/gb200/glm5.py | 397 ++++++++++++++ .../bridge/recipes/glm/h100/__init__.py | 5 + src/megatron/bridge/recipes/glm/h100/glm5.py | 368 +++++++++++++ .../glm_moe_dsa/test_glm5_conversion.py | 106 +++- .../test_hf_to_megatron_generate_text.py | 38 ++ .../models/glm_moe_dsa/test_glm5_bridge.py | 42 ++ tests/unit_tests/recipes/test_glm5_recipes.py | 304 ++++++++++ .../scripts/inference/test_setup_inference.py | 32 ++ .../test_validate_model_verification_card.py | 23 +- 24 files changed, 1958 insertions(+), 36 deletions(-) create mode 100644 examples/model_verification_cards/glm5-2/card.yaml create mode 100644 src/megatron/bridge/recipes/glm/gb200/__init__.py create mode 100644 src/megatron/bridge/recipes/glm/gb200/glm5.py create mode 100644 src/megatron/bridge/recipes/glm/h100/glm5.py create mode 100644 tests/unit_tests/recipes/test_glm5_recipes.py diff --git a/README.md b/README.md index ebd23b256f3..1bc063f7702 100644 --- a/README.md +++ b/README.md @@ -229,7 +229,7 @@ Megatron Bridge provides out-of-the-box bridges and training recipes for a wide | **Ernie** | [Ernie 4.5 MoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/ernie), [Ernie 4.5 VL MoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/ernie_vl) | | [**Falcon**](docs/models/falcon/index.md) | Falcon H1 | | [**Gemma**](docs/models/gemma/index.md) | Gemma / Gemma 2 (deprecated), Gemma 3, Gemma 3-VL, Gemma 4 (26B-A4B MoE / 31B dense), Gemma 4-VL (26B-A4B MoE) | -| [**GLM**](docs/models/glm/index.md) | GLM-4.5 / 4.7 / 4.7-Flash, GLM-4.5V, GLM-5 / 5.1 | +| [**GLM**](docs/models/glm/index.md) | GLM-4.5 / GLM-4.7 / GLM-4.7-Flash, GLM-4.5V, GLM-5 / GLM-5.1 / GLM-5.2 | | [**GPT-OSS**](docs/models/gpt_oss/index.md) | GPT-oss | | [**HY V3**](https://huggingface.co/tencent/Hy3-preview-Base) | Hy3 preview-Base (HF → Megatron checkpoint conversion) | | [**Kimi**](docs/models/kimi/index.md) | Kimi K2, Kimi-K2.5-VL | diff --git a/docs/models/README.md b/docs/models/README.md index 79a0d62a773..0d15a9ca220 100644 --- a/docs/models/README.md +++ b/docs/models/README.md @@ -19,7 +19,7 @@ Megatron Bridge conversion, training recipe links, and model-specific notes. | **DeepSeek** | [DeepSeek V2 (deprecated)](deepseek/deepseek-v2.md), [DeepSeek V3](deepseek/deepseek-v3.md), [DeepSeek V4](deepseek/deepseek-v4.md) | | **Falcon** | [Falcon H1](falcon/falcon-h1.md) | | **Gemma** | [Gemma (deprecated)](gemma/gemma.md), [Gemma 2 (deprecated)](gemma/gemma2.md), [Gemma 3](gemma/gemma3.md), [Gemma 3 VL](gemma/gemma3-vl.md), [Gemma 4 VL](gemma/gemma4-vl.md) | -| **GLM** | [GLM 4.5](glm/glm45.md), [GLM-4.5V](glm/glm-45v.md), [GLM-4.7 / 4.7-Flash](glm/glm47.md), [GLM-5 / 5.1](glm/glm5.md) | +| **GLM** | [GLM-4.5](glm/glm45.md), [GLM-4.5V](glm/glm-45v.md), [GLM-4.7 / GLM-4.7-Flash](glm/glm47.md), [GLM-5 / GLM-5.1 / GLM-5.2](glm/glm5.md) | | **GPT-OSS** | [GPT OSS](gpt_oss/gpt-oss.md) | | **Kimi** | [Kimi K2](kimi/kimi-k2.md), [Kimi-K2.5-VL](kimi/kimi-k25-vl.md) | | **Llama** | [Llama 2 (deprecated)](llama/llama2.md), [Llama 3](llama/llama3.md) | diff --git a/docs/models/glm/glm5.md b/docs/models/glm/glm5.md index 53deb7a56a2..6c5e09d5f7a 100644 --- a/docs/models/glm/glm5.md +++ b/docs/models/glm/glm5.md @@ -1,6 +1,6 @@ -# GLM-5 and GLM-5.1 +# GLM-5, GLM-5.1, and GLM-5.2 -[GLM-5](https://huggingface.co/zai-org/GLM-5) and [GLM-5.1](https://huggingface.co/zai-org/GLM-5.1) are large sparse MoE language models with Multi-Latent Attention and Dynamic Sparse Attention. Megatron Bridge supports both checkpoints through the shared `GLM5Bridge`. +[GLM-5](https://huggingface.co/zai-org/GLM-5), [GLM-5.1](https://huggingface.co/zai-org/GLM-5.1), and [GLM-5.2](https://huggingface.co/zai-org/GLM-5.2) are large sparse MoE language models with Multi-Latent Attention and Dynamic Sparse Attention. Megatron Bridge supports these checkpoints through the shared `GLM5Bridge`. ## Supported Variants @@ -8,6 +8,7 @@ |---------|-----------------|-------| | GLM-5 | `zai-org/GLM-5` | MoE + MLA + DSA architecture | | GLM-5.1 | `zai-org/GLM-5.1` | Same architecture and mapping shape as GLM-5 | +| GLM-5.2 | `zai-org/GLM-5.2` | Adds IndexShare-style DSA index reuse settings | ## Architecture Notes @@ -26,4 +27,3 @@ For conversion, inference, dependency notes, hardware requirements, and MCore pa - Bridge implementation: [`src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py) - Examples: [`examples/models/glm/glm5`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/glm/glm5) - diff --git a/docs/models/glm/index.md b/docs/models/glm/index.md index b1875ee200d..de4089eb82c 100644 --- a/docs/models/glm/index.md +++ b/docs/models/glm/index.md @@ -13,7 +13,7 @@ glm5.md | Variant | Guide | |---------|-------| -| GLM 4.5 | [glm45.md](glm45.md) | +| GLM-4.5 | [glm45.md](glm45.md) | | GLM-4.5V | [glm-45v.md](glm-45v.md) | -| GLM-4.7 / 4.7-Flash | [glm47.md](glm47.md) | -| GLM-5 / 5.1 | [glm5.md](glm5.md) | +| GLM-4.7 / GLM-4.7-Flash | [glm47.md](glm47.md) | +| GLM-5 / GLM-5.1 / GLM-5.2 | [glm5.md](glm5.md) | diff --git a/examples/conversion/hf_to_megatron_generate_text.py b/examples/conversion/hf_to_megatron_generate_text.py index f8da1f4ad21..b32470e6ba6 100644 --- a/examples/conversion/hf_to_megatron_generate_text.py +++ b/examples/conversion/hf_to_megatron_generate_text.py @@ -68,6 +68,20 @@ def _hf_revision_kwargs(revision: str | None) -> dict[str, str]: return {"revision": revision} if revision is not None else {} +def _build_inference_context( + input_ids: torch.Tensor, + *, + legacy_full_prefix: bool, +) -> StaticInferenceContext | None: + """Build a static inference context unless legacy full-prefix decoding is requested.""" + if legacy_full_prefix: + return None + return StaticInferenceContext( + max_batch_size=input_ids.size(0), + max_sequence_length=input_ids.size(1), + ) + + def text_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: """Forward step function for text generation. Required by the forward_backward_func function. @@ -189,7 +203,9 @@ def main(args) -> None: model_provider.pipeline_dtype = torch.bfloat16 # Read pipeline layout from checkpoint for PP > 1 - if pp > 1: + if args.pipeline_model_parallel_layout is not None: + model_provider.pipeline_model_parallel_layout = args.pipeline_model_parallel_layout + elif pp > 1: from pathlib import Path import yaml @@ -210,15 +226,19 @@ def main(args) -> None: model_provider.initialize_model_parallel(seed=0) # Load the Megatron model directly + mp_overrides = { + "tensor_model_parallel_size": tp, + "pipeline_model_parallel_size": pp, + "expert_model_parallel_size": ep, + "expert_tensor_parallel_size": etp, + "pipeline_dtype": torch.bfloat16, + } + if args.pipeline_model_parallel_layout is not None: + mp_overrides["pipeline_model_parallel_layout"] = args.pipeline_model_parallel_layout + model = bridge.load_megatron_model( args.megatron_model_path, - mp_overrides={ - "tensor_model_parallel_size": tp, - "pipeline_model_parallel_size": pp, - "expert_model_parallel_size": ep, - "expert_tensor_parallel_size": etp, - "pipeline_dtype": torch.bfloat16, - }, + mp_overrides=mp_overrides, wrap_with_ddp=False, ) @@ -239,6 +259,8 @@ def main(args) -> None: model_provider.expert_model_parallel_size = ep model_provider.expert_tensor_parallel_size = etp model_provider.pipeline_dtype = torch.bfloat16 + if args.pipeline_model_parallel_layout is not None: + model_provider.pipeline_model_parallel_layout = args.pipeline_model_parallel_layout # Once all overrides are set, finalize the model provider to ensure the post initialization logic is run model_provider.finalize() @@ -284,9 +306,9 @@ def main(args) -> None: print_rank_0(f"Generation step {step}") fwd_bwd_function = get_forward_backward_func() - inference_context = StaticInferenceContext( - max_batch_size=input_ids.size(0), - max_sequence_length=input_ids.size(1), + inference_context = _build_inference_context( + input_ids, + legacy_full_prefix=args.legacy_full_prefix, ) iterator = SingleBatchIterator(input_ids, position_ids, inference_context) @@ -375,6 +397,15 @@ def build_parser() -> argparse.ArgumentParser: default=20, help="Maximum number of new tokens to generate.", ) + parser.add_argument( + "--legacy-full-prefix", + action="store_true", + help=( + "Use non-cached autoregressive generation by disabling the MCore inference context and recomputing " + "the full accumulated prefix at every decoding step. This slower legacy path supports models such " + "as GLM-5 whose AbsorbedMLA attention does not yet support cached inference." + ), + ) parser.add_argument( "--apply-chat-template", action="store_true", @@ -388,6 +419,10 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--tp", type=int, default=1, help="Tensor parallelism size") parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size") + parser.add_argument( + "--pipeline-model-parallel-layout", + help="Optional pipeline layout used to keep model-specific layer groups within pipeline stages.", + ) parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size") parser.add_argument("--etp", type=int, default=1, help="Expert tensor parallelism size") parser.add_argument("--megatron_model_path", type=str, default=None, help="Path to the Megatron model checkpoint") diff --git a/examples/model_verification_cards/glm5-2/card.yaml b/examples/model_verification_cards/glm5-2/card.yaml new file mode 100644 index 00000000000..534828ad30c --- /dev/null +++ b/examples/model_verification_cards/glm5-2/card.yaml @@ -0,0 +1,519 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: glm5_2 +summary: > + Performance disclaimer: this model has not been performance-tuned; reported + timing and throughput metrics are sanity checks, not optimized performance + results. GLM-5.2 H100 support verification currently covers exact GPU + checkpoint round-trip on CPU and GPU, HF-to-Megatron forward equivalence, + bounded full SFT with export/reload + inference, 200K-context CP32 SFT, bounded PEFT, and bounded pretraining with + a reloadable middle checkpoint and verified continuation. GB200 support + verification currently covers bounded pretraining, bounded full SFT, 128K + packed long-context SFT, bounded PEFT, bounded checkpoint resume, and + deterministic legacy Megatron inference; post-SFT export/reload inference + remains pending. +verification_index: + model_level: + verified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - megatron_to_hf_cpu + - megatron_to_hf_gpu + - manual_forward_pass + - inference + training: + GB200: + verified: [pretrain, sft, sft_long_context, peft, checkpoint_resume] + unverified: [sft_export_inference] + H100: + verified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] +model: + hf_id: zai-org/GLM-5.2 + hf_revision: 4d67f66cc64d3219133b767c253b2ad1425c6c88 # pragma: allowlist secret + architecture: GlmMoeDsaForCausalLM + min_transformers_version: "5.12.1" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.04-py3 + bridge_commit: 8645cda8738cbb87b8b998ad2d09341eef2510c7 # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model zai-org/GLM-5.2 + --hf-revision 4d67f66cc64d3219133b767c253b2ad1425c6c88 + --megatron-path work/model-verification/glm5-2/cpu-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-22 + expected_result: > + Import exits successfully and creates a reloadable CPU checkpoint. Its + strict round-trip comparison matches all 59,585 BF16 tensors and + 753,329,940,480 elements bitwise against the immutable HF revision. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 13 --gpus-per-node 8 + --hf-model zai-org/GLM-5.2 + --hf-revision 4d67f66cc64d3219133b767c253b2ad1425c6c88 + --megatron-path work/model-verification/glm5-2/gpu-megatron + --torch-dtype bfloat16 --tp 1 --pp 13 --ep 8 --etp 1 + last_verified: 2026-07-22 + expected_result: > + Import exits successfully at TP1/PP13/EP8/ETP1 and creates a reloadable + iter_0000000 checkpoint. Its strict round-trip comparison matches all + 59,585 BF16 tensors and 753,329,940,480 elements bitwise against the + immutable HF revision. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model zai-org/GLM-5.2 + --hf-revision 4d67f66cc64d3219133b767c253b2ad1425c6c88 + --megatron-path work/model-verification/glm5-2/cpu-megatron + --hf-path work/model-verification/glm5-2/cpu-hf-export + --torch-dtype bfloat16 + last_verified: 2026-07-22 + expected_result: > + Export exits successfully; all 59,585 BF16 tensors and + 753,329,940,480 elements match the immutable HF revision bitwise. + Transformers 5.12.1 strictly reloads the output as + GlmMoeDsaForCausalLM with 743,377,000,704 parameters. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 13 --gpus-per-node 8 + --hf-model zai-org/GLM-5.2 + --hf-revision 4d67f66cc64d3219133b767c253b2ad1425c6c88 + --megatron-path work/model-verification/glm5-2/gpu-megatron + --hf-path work/model-verification/glm5-2/gpu-hf-export-exact + --torch-dtype bfloat16 --distributed-save + --tp 1 --pp 13 --ep 8 --etp 1 + last_verified: 2026-07-22 + expected_result: > + Distributed export exits successfully; all 59,585 BF16 tensors and + 753,329,940,480 elements match the immutable HF revision bitwise. + Transformers strictly reloads the output as GlmMoeDsaForCausalLM with + 743,377,000,704 parameters. + + manual_forward_pass: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=52 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path zai-org/GLM-5.2 + --hf-revision 4d67f66cc64d3219133b767c253b2ad1425c6c88 + --hf-logits-input work/model-verification/glm5-2/hf-forward-reference.safetensors + --megatron-forward-mode standard + --megatron_model_path work/model-verification/glm5-2/gpu-megatron + --prompt "The capital of France is" --tp 1 --pp 13 --ep 4 --etp 1 + last_verified: 2026-07-22 + expected_result: > + The pinned HF and Megatron next-token results match at token 12089 (`Paris`). + Cosine similarity is 0.999443, maximum absolute logit difference is + 0.500000, and mean absolute logit difference is 0.063958. + + inference: + status: verified + precision: bf16 + bridge_commit: cd09c801474e5573ffceccf236b48b543c7a0b4a # pragma: allowlist secret + command: > + ./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 + --task legacy-full-prefix-generation + --hf_model_path zai-org/GLM-5.2 + --hf-revision 4d67f66cc64d3219133b767c253b2ad1425c6c88 + --prompt "Write one concise sentence explaining why deterministic tests matter:" + --apply-chat-template --thinking-mode disabled --max_new_tokens 16 + --legacy-full-prefix --tp 1 --pp 4 --ep 8 --etp 1 + --pipeline-model-parallel-layout + "Etttttt|tttttttttttttttttttttttt|tttttttttttttttttttttttt|ttttttttttttttttttttttttL" + last_verified: 2026-08-02 + expected_result: |- + One deterministic BF16 run imports the pinned complete HF model directly + into Megatron across 32 ranks at TP1/PP4/EP8/ETP1. It exits successfully + after exactly 16 new greedy tokens without constructing an MCore inference + context. The legacy path recomputes the accumulated prefix at every step + because cached inference is not yet supported for AbsorbedMLA. The literal + completion was: + "1. **Analyze the Request:** + * **Topic:**" + + pretrain: + H100: + status: verified + precision: bf16 + enabled_features: + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 44 --gpus-per-node 8 + --recipe glm52_pretrain_416gpu_h100_bf16_config --mode pretrain + --dataset megatron-indexed --deterministic --seq_length 2048 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/wikitext103/wikitext103_glm52_text_document"],null]' + dataset.path_to_cache=work/cache/glm5-2/wikitext-pretrain + dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 + model.pipeline_model_parallel_size=11 + model.virtual_pipeline_model_parallel_size=null + model.microbatch_group_size_per_vp_stage=null + model.pipeline_model_parallel_layout='Etttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttt|ttttmL' + --save_dir work/model-verification/glm5-2/pretrain-reference/checkpoints + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-23 + metrics: + initial_loss: 13.15907 + final_loss: 6.385243 + last_10_steps_step_time_ms_avg: 31204.920 + last_10_steps_model_tflops_per_gpu_avg: 50.620 + expected_result: > + The uninterrupted 352-H100 TP1/PP11/CP1/EP32/ETP1 run completed 100 + real-data steps at GBS/MBS 1024/1 with cuDNN DSA, sparse indexer loss + coefficient 0.001, MTP1, and flex/DeepEP dispatch. Loss remained finite + from 13.15907 to 6.385243, final grad norm was 0.796, and every step + recorded zero skipped and NaN iterations. The final 10 steps averaged + 31,204.920 ms and 50.620 TFLOPS/GPU; peak observed GPU memory was + 80,999 MiB. The matching step-50 and step-100 checkpoints contain + model, optimizer, scheduler, data-order, and RNG state; the step-50 + checkpoint reloads successfully. + GB200: + status: verified + precision: bf16 + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 + --recipe glm52_pretrain_192gpu_gb200_bf16_config --mode pretrain + --dataset megatron-indexed --deterministic --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/glm5-2/pretrain-tulu3-text-document/tulu3_glm52_text_document_text_document"],null]' + dataset.path_to_cache=work/cache/glm5-2/pretrain-tulu3-gb200 + dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 + --save_dir work/model-verification/glm5-2/pretrain-gb200-cudnn/checkpoints + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-24 + metrics: + initial_loss: 13.19136 + final_loss: 8.171059 + last_10_steps_step_time_ms_avg: 60610.850 + last_10_steps_model_tflops_per_gpu_avg: 101.310 + expected_result: > + The 192-GB200 TP1/PP6/CP1/EP32/ETP1 run completed exactly 100 + real-data pretraining steps at GBS/MBS 1024/1 with cuDNN DSA, sparse + indexer loss coefficient 0.001, MTP1, and all-to-all expert dispatch. + Loss remained finite from 13.19136 to 8.171059, final grad norm was + 1.394, and every step recorded zero skipped and NaN iterations. The + final 10 steps averaged 60,610.850 ms and 101.310 TFLOPS/GPU. Step-50 + and step-100 checkpoints were written with model, optimizer, scheduler, + data-order, and RNG state for resume verification. These losses are not + directly comparable with the H100 pretrain result: H100 used Wikitext + at sequence length 2,048, while GB200 used Tulu 3 at sequence length + 4,096. + + sft: + H100: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 52 --gpus-per-node 8 + --recipe glm52_sft_416gpu_h100_bf16_config --mode sft + --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + dataset.hf_output_root=work/data/glm5-2/tulu3-full-sft-pad32 + model.context_parallel_size=2 + model.virtual_pipeline_model_parallel_size=null + model.pipeline_model_parallel_layout='Etttttt|tttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttt|tttt|tttt|tttt|ttttmL' + --save_dir work/model-verification/glm5-2/sft-functional/checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-22 + metrics: + initial_loss: 1.909875 + final_loss: 0.6634203 + last_10_steps_step_time_ms_avg: 4436.290 + last_10_steps_model_tflops_per_gpu_avg: 9.450 + expected_result: > + The immutable-revision 416-H100 run completes exactly 100 full-SFT + steps at TP1/PP13/CP2/EP32/ETP1, GBS/MBS 32/1, with cuDNN DSA, sparse + indexer loss coefficient 0.001, MTP1, and natural routing. Pad-32 + packing is 99.60% efficient and contains 4,097,215 supervised tokens. + Loss remains finite from 1.909875 to 0.6634203 with zero skipped or NaN + iterations; the final grad norm is 0.895 and the step-100 checkpoint + reloads for export. + GB200: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 + --recipe glm52_sft_192gpu_gb200_bf16_config --mode sft + --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 + --max_steps 100 + tokenizer.chat_template_path=work/templates/glm52_chat_template_generation.jinja + --save_dir work/model-verification/glm5-2/sft-gb200-cudnn/checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-24 + metrics: + initial_loss: 1.831100 + final_loss: 0.5953441 + last_10_steps_step_time_ms_avg: 68370.590 + last_10_steps_model_tflops_per_gpu_avg: 9.180 + expected_result: > + The immutable-revision 192-GB200 run completed exactly 100 full-SFT + steps at TP1/PP6/CP1/EP32/ETP1, GBS/MBS 32/1, with offline-packed + Tulu 3 data, cuDNN DSA, sparse indexer loss coefficient 0.001, MTP1, + and all-to-all expert dispatch. Loss remained finite from 1.831100 to + 0.5953441, final grad norm was 0.837, and every step recorded zero + skipped and NaN iterations. The step-100 checkpoint was written and is + available for a future GB200 export/reload gate. + + sft_export_inference: + H100: + status: verified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 52 --gpus-per-node 8 --hf-model zai-org/GLM-5.2 + --hf-revision 4d67f66cc64d3219133b767c253b2ad1425c6c88 + --megatron-path work/model-verification/glm5-2/sft-functional/checkpoints/iter_0000100 + --hf-path work/model-verification/glm5-2/sft-hf-export + --torch-dtype bfloat16 --distributed-save + --tp 1 --pp 13 --ep 32 --etp 1 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/glm5-2/sft-hf-export + --prompt "Write one concise sentence explaining why deterministic tests matter:" + --max-new-tokens 8 --chat-template --disable-thinking + --device cpu --dtype bfloat16 + last_verified: 2026-07-22 + expected_result: > + The step-100 full-SFT checkpoint exports successfully and strictly + reloads with Transformers. One greedy run produces exactly 8 tokens + with the completion "Deterministic tests matter because they provide + consistent". + GB200: + status: unverified + precision: bf16 + depends_on: sft + commands: null + last_verified: null + expected_result: > + Export and deterministic HF reload/inference from the GLM-5.2 SFT + checkpoint have not been run for this card. + + sft_long_context: + H100: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + context_parallel_size: 32 + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 76 --gpus-per-node 8 + --recipe glm52_sft_608gpu_h100_bf16_200k_config --mode sft + --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 + --max_steps 20 --seq_length 200000 + dataset.dataset_root=work/data/glm5-2/synthetic-200k-260 + train.empty_unused_memory_level=2 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: 2026-07-22 + metrics: + initial_loss: 8.283127 + final_loss: 2.557671 + last_10_steps_step_time_ms_avg: 36603.370 + last_10_steps_model_tflops_per_gpu_avg: 210.650 + expected_result: > + The 608-H100 TP1/PP19/CP32/EP32/ETP1 run completed exactly 20 optimizer + steps over 200,000-token packed rows with cuDNN DSA, sparse indexer loss + coefficient 0.001, and MTP1. Loss remained finite from 8.283127 to + 2.557671, final grad norm was 10838.703, and every step recorded zero + skipped and NaN iterations. The last 10 steps averaged 36,603.370 ms + and 210.650 TFLOPS/GPU; peak observed GPU memory was 81,005 MiB. + GB200: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + context_parallel_size: 32 + command: > + ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 + --recipe glm52_sft_192gpu_gb200_bf16_128k_config --mode sft + --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 + --max_steps 20 --seq_length 131072 --lr 1e-6 --min_lr 0 --warmup_iters 2 + tokenizer.chat_template_path=work/templates/glm52_chat_template_generation.jinja + checkpoint.load=null ddp.check_for_nan_in_grad=true + ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: 2026-07-24 + metrics: + initial_loss: 0.3235812 + final_loss: 0.1924341 + last_10_steps_step_time_ms_avg: 29269.730 + last_10_steps_model_tflops_per_gpu_avg: 236.510 + expected_result: > + The immutable-revision 192-GB200 run completed exactly 20 optimizer + steps over 131,072-token packed rows at TP1/PP6/CP32/EP32/ETP1, + GBS/MBS 8/1, with cuDNN DSA, sparse indexer loss coefficient 0.001, + MTP1, padded cu_seqlens, and all-to-all expert dispatch. Loss remained + finite from 0.3235812 to 0.1924341, final grad norm was 1448.863, and + every step recorded zero skipped and NaN iterations. The final 10 steps + averaged 29,269.730 ms and 236.510 TFLOPS/GPU. + + peft: + H100: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 26 --gpus-per-node 8 + --recipe glm52_peft_208gpu_h100_bf16_config --mode lora + --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + dataset.hf_output_root=work/data/glm5-2/tulu3-peft-pad4 + --save_dir work/model-verification/glm5-2/peft-lora/checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-22 + metrics: + initial_loss: 1.909421 + final_loss: 0.8835775 + last_10_steps_step_time_ms_avg: 4638.550 + last_10_steps_model_tflops_per_gpu_avg: 17.240 + expected_result: > + The immutable-revision 208-H100 run completes exactly 100 LoRA steps at + TP1/PP13/CP1/EP16/ETP1, GBS/MBS 32/1, with cuDNN DSA, sparse indexer + loss coefficient 0.001, MTP1, and natural routing. Pad-4 packing is + 99.45% efficient and contains 4,251,129 supervised tokens. Loss remains + finite from 1.909421 to 0.8835775 with zero skipped or NaN iterations. + The final grad norm is 0.101; the checkpoint contains exactly 790 BF16 + rank-8 adapter tensors and no base-model weights. + GB200: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 + --recipe glm52_peft_192gpu_gb200_bf16_config --mode lora + --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 + --max_steps 100 + tokenizer.chat_template_path=work/templates/glm52_chat_template_generation.jinja + --save_dir work/model-verification/glm5-2/peft-gb200-cudnn/checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-24 + metrics: + initial_loss: 1.831100 + final_loss: 0.8957523 + last_10_steps_step_time_ms_avg: 68588.270 + last_10_steps_model_tflops_per_gpu_avg: 9.330 + expected_result: > + The immutable-revision 192-GB200 run completed exactly 100 LoRA steps + at TP1/PP6/CP1/EP32/ETP1, GBS/MBS 32/1, with offline-packed Tulu 3 + data, cuDNN DSA, sparse indexer loss coefficient 0.001, MTP1, and + all-to-all expert dispatch. Loss remained finite from 1.831100 to + 0.8957523, final grad norm was 0.101, and every step recorded zero + skipped and NaN iterations. The step-100 adapter checkpoint was written + for reload validation. + + checkpoint_resume: + H100: + status: verified + precision: bf16 + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 44 --gpus-per-node 8 + --recipe glm52_pretrain_416gpu_h100_bf16_config --mode pretrain + --dataset megatron-indexed --deterministic --seq_length 2048 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/wikitext103/wikitext103_glm52_text_document"],null]' + dataset.path_to_cache=work/cache/glm5-2/wikitext-pretrain + dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 + model.pipeline_model_parallel_size=11 + model.virtual_pipeline_model_parallel_size=null + model.microbatch_group_size_per_vp_stage=null + model.pipeline_model_parallel_layout='Etttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttt|ttttmL' + --load_dir work/model-verification/glm5-2/pretrain-reference/checkpoints + --save_dir work/model-verification/glm5-2/pretrain-resume/checkpoints + --save_interval 0 checkpoint.ckpt_step=50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: 2026-07-23 + metrics: + initial_loss: 7.219822 + final_loss: 6.372624 + last_10_steps_step_time_ms_avg: 32002.290 + last_10_steps_model_tflops_per_gpu_avg: 49.350 + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: true + expected_result: > + The run reloads the corrected reference's step-50 model, optimizer, + scheduler, data-order, and RNG state and completes steps 51-100 with + zero skipped or NaN iterations. Step 51 matches exactly at 7.219822. + Step 100 is 6.372624 versus the uninterrupted reference's 6.385243, an + absolute difference of 0.012619 within the allowed 0.063853. Final + grad norm is 2.169; the last 10 steps average 32,002.290 ms and 49.350 + TFLOPS/GPU. Peak observed allocator memory is 65.230 GiB allocated and + 69.684 GiB reserved. + GB200: + status: verified + precision: bf16 + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 + --recipe glm52_pretrain_192gpu_gb200_bf16_config --mode pretrain + --dataset megatron-indexed --deterministic --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/glm5-2/pretrain-tulu3-text-document/tulu3_glm52_text_document_text_document"],null]' + dataset.path_to_cache=work/cache/glm5-2/pretrain-tulu3-gb200 + dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 + --load_dir work/model-verification/glm5-2/pretrain-gb200-cudnn/checkpoints + --save_dir work/model-verification/glm5-2/pretrain-resume-gb200-cudnn/checkpoints + --save_interval 0 checkpoint.ckpt_step=50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: 2026-07-24 + metrics: + initial_loss: 8.581669 + final_loss: 8.192354 + last_10_steps_step_time_ms_avg: 67700.250 + last_10_steps_model_tflops_per_gpu_avg: 90.710 + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: true + expected_result: > + The run reloads the GB200 reference's step-50 model, optimizer, + scheduler, data-order, and RNG state and completes steps 51-100 with + zero skipped or NaN iterations. Step 51 matches exactly at 8.581669. + Step 100 is 8.192354 versus the uninterrupted reference's 8.171059, an + absolute difference of 0.021295 within the allowed 0.081712. Final grad + norm is 1.840; the last 10 steps average 67,700.250 ms and 90.710 + TFLOPS/GPU. diff --git a/examples/models/glm/glm5/README.md b/examples/models/glm/glm5/README.md index 98e5e0d2040..4f74072e8cc 100644 --- a/examples/models/glm/glm5/README.md +++ b/examples/models/glm/glm5/README.md @@ -38,7 +38,7 @@ The PyPI source distribution is incomplete; install from the git repo. ## Inference (Megatron) -[slurm_inference.sh](slurm_inference.sh) loads the HF checkpoint, converts to Megatron in-memory, and runs greedy text generation with `TP=1, PP=2, EP=32` across 64 GPUs. TP remains disabled because AbsorbedMLA requires sequence parallelism when `TP > 1`. +[slurm_inference.sh](slurm_inference.sh) loads the HF checkpoint, converts to Megatron in-memory, and runs greedy text generation with `TP=1, PP=2, EP=32` across 64 GPUs. TP remains disabled because AbsorbedMLA requires sequence parallelism when `TP > 1`. The launcher passes `--legacy-full-prefix` to use non-cached autoregressive generation, which recomputes the full prefix at every decoding step. Cached MCore inference is not yet supported for AbsorbedMLA. ```bash sbatch examples/models/glm/glm5/slurm_inference.sh diff --git a/examples/models/glm/glm5/slurm_inference.sh b/examples/models/glm/glm5/slurm_inference.sh index 2a10adbbf52..05e34c3a9a7 100644 --- a/examples/models/glm/glm5/slurm_inference.sh +++ b/examples/models/glm/glm5/slurm_inference.sh @@ -27,7 +27,8 @@ # avoids the AbsorbedMLA sequence-parallel requirement for TP > 1. # # Megatron inference engine support is not yet available for this architecture. -# This launcher defaults to the older slow sanity-check generation path for now. +# This launcher passes --legacy-full-prefix to use non-cached autoregressive +# generation, which recomputes the full prefix at every decoding step. # # Loads the HF checkpoint, converts to Megatron in-memory, and runs greedy # text generation via examples/conversion/hf_to_megatron_generate_text.py. @@ -119,6 +120,7 @@ print(m.group(1) + m.group(2) if m else s.split('"'"','"'"')[0]) --hf_model_path "$HF_MODEL_PATH" \ --prompt "$PROMPT" \ --max_new_tokens "$MAX_NEW_TOKENS" \ + --legacy-full-prefix \ --tp "$TP" --ep "$EP" --pp "$PP" ' status=$? diff --git a/scripts/inference/README.md b/scripts/inference/README.md index e83f022e587..fef556e8972 100644 --- a/scripts/inference/README.md +++ b/scripts/inference/README.md @@ -26,7 +26,10 @@ are forwarded unchanged to the selected inference entry point. Text generation uses `text_generation.py`; `--task vlm-generation` selects `vlm_generation.py` for multimodal generation, and `--task model-comparison` selects `examples/conversion/compare_hf_and_megatron/compare.py` for forward-logit -parity. +parity. `--task legacy-full-prefix-generation` selects a slow, non-optimized +compatibility path and requires `--legacy-full-prefix`. It recomputes the +accumulated prefix for every decoding step for models such as GLM-5 whose +AbsorbedMLA attention does not yet support cached inference. ```bash ./scripts/inference/infer.sh \ diff --git a/scripts/inference/setup_inference.py b/scripts/inference/setup_inference.py index ad57f32ad63..32b524b92de 100755 --- a/scripts/inference/setup_inference.py +++ b/scripts/inference/setup_inference.py @@ -32,6 +32,7 @@ CONTAINER_REPO_ROOT = Path("/opt/Megatron-Bridge") INFERENCE_TASKS = { "text-generation": Path("scripts/inference/text_generation.py"), + "legacy-full-prefix-generation": Path("examples/conversion/hf_to_megatron_generate_text.py"), "vlm-generation": Path("scripts/inference/vlm_generation.py"), "model-comparison": Path("examples/conversion/compare_hf_and_megatron/compare.py"), } @@ -52,8 +53,11 @@ def _build_parser() -> argparse.ArgumentParser: --prompt "Megatron Bridge inference is" --max_new_tokens 32 Use --task vlm-generation for multimodal generation or --task model-comparison -for a one-step Hugging Face/Megatron comparison. Arguments not owned by this -launcher are forwarded unchanged to the selected repository entry point. +for a one-step Hugging Face/Megatron comparison. The +legacy-full-prefix-generation task is a slow, non-optimized compatibility path +for models that do not yet support cached inference and requires +--legacy-full-prefix. Arguments not owned by this launcher are forwarded +unchanged to the selected repository entry point. """, ) execution = parser.add_argument_group("Execution") @@ -174,6 +178,15 @@ def _validate_args(args: argparse.Namespace) -> None: raise ValueError("Slurm execution requires --container-image or CONTAINER_IMAGE.") +def _validate_task_args(task_name: str, inference_args: list[str]) -> None: + """Validate arguments owned by a selected inference task.""" + uses_legacy_full_prefix = "--legacy-full-prefix" in inference_args + if task_name == "legacy-full-prefix-generation" and not uses_legacy_full_prefix: + raise ValueError("--task legacy-full-prefix-generation requires --legacy-full-prefix.") + if task_name != "legacy-full-prefix-generation" and uses_legacy_full_prefix: + raise ValueError("--legacy-full-prefix requires --task legacy-full-prefix-generation.") + + def _build_executor(args: argparse.Namespace, env_names: list[str], mounts: list[str]) -> object: """Build the srun-native NeMo-Run Slurm executor.""" gpu_kwargs = {} if args.no_gpu_resource_request else {"gpus_per_node": args.gpus_per_node} @@ -238,6 +251,7 @@ def main(argv: list[str] | None = None) -> None: """Build and submit one Bridge inference experiment.""" args, inference_args = parse_args(argv) _validate_args(args) + _validate_task_args(args.task, inference_args) env_names = _parse_env(args.env) mounts = _parse_mounts(args.mount) executor = _build_executor(args, env_names, mounts) diff --git a/skills/create-model-verification-card/SKILL.md b/skills/create-model-verification-card/SKILL.md index 5a03eeb6df0..6d782f23729 100644 --- a/skills/create-model-verification-card/SKILL.md +++ b/skills/create-model-verification-card/SKILL.md @@ -266,7 +266,9 @@ model-verification workload: training item; - use `scripts/inference/infer.sh --nodes ... --gpus-per-node ... --task ...` for Megatron inference and manual model comparison; select - `text-generation`, `vlm-generation`, or `model-comparison` explicitly; + `text-generation`, `legacy-full-prefix-generation`, `vlm-generation`, or + `model-comparison` explicitly. The legacy full-prefix task is a slow, + non-optimized compatibility path and requires `--legacy-full-prefix`; - invoke public shell launchers directly from the card. Do not create Slurm jobs by calling their Python setup modules, and do not wrap the launchers in `uv run`, `python`, `srun`, or `sbatch`; the shell entry point owns its Python @@ -604,10 +606,12 @@ result. Private executor configuration stays outside the card. helper does not append padding before selecting the compared next-token position. - **Megatron inference:** Launch through `scripts/inference/infer.sh` with the - explicit `text-generation` or `vlm-generation` task. Disable sampling, run - one deterministic greedy generation with an exact token count, and record - the literal completion including whitespace. A second replay may help - diagnose nondeterminism, but it is not required verification evidence. + explicit `text-generation`, `legacy-full-prefix-generation`, or + `vlm-generation` task. Use the slow, non-optimized legacy task only when + cached inference is unsupported, and pass `--legacy-full-prefix`. Disable + sampling, run one deterministic greedy generation with an exact token count, + and record the literal completion including whitespace. A second replay may + help diagnose nondeterminism, but it is not required verification evidence. - **Pretrain:** Use a bounded public dataset description and a stable schedule. Save a middle and final checkpoint when resume is in scope. For expensive workloads, a 100-step reference with checkpoints at steps 50 and 100 is a diff --git a/skills/create-model-verification-card/scripts/validate_card.py b/skills/create-model-verification-card/scripts/validate_card.py index e7a53fbff6c..c6274fa8f1a 100644 --- a/skills/create-model-verification-card/scripts/validate_card.py +++ b/skills/create-model-verification-card/scripts/validate_card.py @@ -1097,15 +1097,22 @@ def _validate_inference( task_values = _argument_values(command, "--task") if len(task_values) != 1 or task_values[0] not in { "text-generation", + "legacy-full-prefix-generation", "vlm-generation", "model-comparison", }: errors.append(f"{_pointer(*resolved_command_path)}: infer.sh must specify one supported --task") + if task_values == ["legacy-full-prefix-generation"] and "--legacy-full-prefix" not in command_tokens: + errors.append( + f"{_pointer(*resolved_command_path)}: legacy-full-prefix-generation requires --legacy-full-prefix" + ) for resource_flag in ("--nodes", "--gpus-per-node"): if len(_argument_values(command, resource_flag)) != 1: errors.append( f"{_pointer(*resolved_command_path)}: infer.sh must specify {resource_flag} exactly once" ) + if any(option in command_tokens for option in ("--dry-run", "--submission-dry-run")): + errors.append(f"{_pointer(*resolved_command_path)}: verified inference command must submit the workload") prompts = _argument_values(command, "--prompt") if len(prompts) != 1: errors.append(f"{_pointer(*resolved_command_path)}: specify --prompt exactly once") diff --git a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py index 56ff48fb08c..3deba106c12 100644 --- a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py +++ b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py @@ -111,6 +111,10 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider hf_config.num_hidden_layers - hf_config.first_k_dense_replace ) provider.moe_shared_expert_intermediate_size = hf_config.moe_intermediate_size * hf_config.n_shared_experts + # GlmMoeDsaConfig may normalize qk_rope_head_dim to head_dim while + # loading GLM-5.2. Recover the RoPE width from the model's invariant: + # total QK width = non-RoPE width + RoPE width. + provider.qk_pos_emb_head_dim = hf_config.qk_head_dim - hf_config.qk_nope_head_dim # GLM5-specific: rotary_base is nested in rope_parameters provider.rotary_base = hf_config.rope_parameters["rope_theta"] diff --git a/src/megatron/bridge/recipes/glm/__init__.py b/src/megatron/bridge/recipes/glm/__init__.py index f003ab4de2b..6f59a74fd69 100644 --- a/src/megatron/bridge/recipes/glm/__init__.py +++ b/src/megatron/bridge/recipes/glm/__init__.py @@ -12,6 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .gb200.glm5 import ( + glm52_peft_192gpu_gb200_bf16_config, + glm52_pretrain_192gpu_gb200_bf16_config, + glm52_sft_192gpu_gb200_bf16_128k_config, + glm52_sft_192gpu_gb200_bf16_config, +) from .glm45 import ( glm45_355b_peft_config, glm45_355b_pretrain_config, @@ -29,4 +35,8 @@ "glm45_air_106b_sft_config", "glm45_355b_peft_config", "glm45_air_106b_peft_config", + "glm52_peft_192gpu_gb200_bf16_config", + "glm52_pretrain_192gpu_gb200_bf16_config", + "glm52_sft_192gpu_gb200_bf16_128k_config", + "glm52_sft_192gpu_gb200_bf16_config", ] diff --git a/src/megatron/bridge/recipes/glm/gb200/__init__.py b/src/megatron/bridge/recipes/glm/gb200/__init__.py new file mode 100644 index 00000000000..0700698697b --- /dev/null +++ b/src/megatron/bridge/recipes/glm/gb200/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.recipes.glm.gb200.glm5 import * # noqa: F403 + + +__all__ = [ + "glm52_peft_192gpu_gb200_bf16_config", + "glm52_pretrain_192gpu_gb200_bf16_config", + "glm52_sft_192gpu_gb200_bf16_128k_config", + "glm52_sft_192gpu_gb200_bf16_config", +] diff --git a/src/megatron/bridge/recipes/glm/gb200/glm5.py b/src/megatron/bridge/recipes/glm/gb200/glm5.py new file mode 100644 index 00000000000..5250452f199 --- /dev/null +++ b/src/megatron/bridge/recipes/glm/gb200/glm5.py @@ -0,0 +1,397 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GB200 recipes for GLM-5.2.""" + +from __future__ import annotations + +from megatron.bridge import AutoBridge +from megatron.bridge.data.builders import ChatSFTPreprocessingConfig +from megatron.bridge.peft.base import PEFT +from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common +from megatron.bridge.recipes.utils.dataset_utils import default_peft_config, default_tulu3_config +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed + + +_GLM52_MODEL_ID = "zai-org/GLM-5.2" +_GLM52_MODEL_REVISION = "4d67f66cc64d3219133b767c253b2ad1425c6c88" # pragma: allowlist secret +_TULU3_REVISION = "b14afda60f1bbebe55d5d2fa1e4df5042f97f8be" # pragma: allowlist secret + + +def glm52_pretrain_192gpu_gb200_bf16_config() -> ConfigContainer: + """GLM-5.2 bounded pretraining on 192 GB200 GPUs.""" + cfg = _pretrain_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 4096 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 6 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_layout = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.num_layers_in_first_pipeline_stage = 14 + cfg.model.num_layers_in_last_pipeline_stage = 16 + cfg.model.account_for_embedding_in_pipeline_split = False + cfg.model.account_for_loss_in_pipeline_split = False + cfg.model.microbatch_group_size_per_vp_stage = 6 + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.apply_rope_fusion = False + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + cfg.ddp.grad_reduce_in_fp32 = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + cfg.dataset.seq_length = 4096 + cfg.dataset.num_workers = 8 + cfg.dataset.random_seed = 1234 + cfg.rng.seed = 1234 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=40, + max_lr=3e-4, + min_lr=3e-5, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.optimizer.use_precision_aware_optimizer = True + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + + return cfg + + +def glm52_sft_192gpu_gb200_bf16_config() -> ConfigContainer: + """GLM-5.2 bounded full SFT on 192 GB200 GPUs.""" + cfg = _sft_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 2048 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 6 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_layout = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.num_layers_in_first_pipeline_stage = 14 + cfg.model.num_layers_in_last_pipeline_stage = 16 + cfg.model.account_for_embedding_in_pipeline_split = False + cfg.model.account_for_loss_in_pipeline_split = False + cfg.model.microbatch_group_size_per_vp_stage = 6 + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.apply_rope_fusion = False + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + cfg.dataset = default_tulu3_config( + seq_length=2048, + enable_offline_packing=True, + pad_seq_to_mult=32, + ) + cfg.dataset.hf_dataset.split = "train[:10000]" + cfg.dataset.hf_dataset.load_kwargs = {"revision": _TULU3_REVISION} + cfg.dataset.hf_output_root = "work/data/glm5-2/tulu3-full-sft-gb200" + cfg.dataset.hf_rewrite = False + cfg.dataset.hf_validation_proportion = None + cfg.dataset.max_train_samples = 10000 + cfg.dataset.seed = 1234 + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=10, + max_lr=5e-6, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.save_optim = False + cfg.checkpoint.save_rng = False + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + return cfg + + +def glm52_sft_192gpu_gb200_bf16_128k_config() -> ConfigContainer: + """GLM-5.2 128K packed SFT with context parallelism on 192 GB200 GPUs.""" + cfg = _sft_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 131072 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 6 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_layout = None + cfg.model.context_parallel_size = 32 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.num_layers_in_first_pipeline_stage = 14 + cfg.model.num_layers_in_last_pipeline_stage = 16 + cfg.model.account_for_embedding_in_pipeline_split = False + cfg.model.account_for_loss_in_pipeline_split = False + cfg.model.microbatch_group_size_per_vp_stage = None + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.apply_rope_fusion = False + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.train.train_iters = 20 + cfg.train.global_batch_size = 8 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + cfg.dataset.seq_length = 131072 + cfg.dataset.hf_dataset = None + cfg.dataset.dataset_root = "work/data/glm5-2/synthetic-long-context-gb200" + cfg.dataset.hf_output_root = None + cfg.dataset.hf_validation_proportion = None + cfg.dataset.seed = 1234 + cfg.dataset.preprocessing = ChatSFTPreprocessingConfig() + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + cfg.dataset.offline_packing_specs.packed_sequence_size = 131072 + cfg.dataset.offline_packing_specs.tokenizer_model_name = "glm5" + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 64 + cfg.dataset.offline_packing_specs.pad_cu_seqlens = True + cfg.dataset.dataset_kwargs = {"pad_to_max_length": True} + + cfg.rng.seed = 5678 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=2, + max_lr=1e-6, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 20 + cfg.checkpoint.save = None + cfg.checkpoint.load = None + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + } + return cfg + + +def glm52_peft_192gpu_gb200_bf16_config(peft_scheme: str | PEFT = "lora") -> ConfigContainer: + """GLM-5.2 bounded PEFT on 192 GB200 GPUs.""" + cfg = _peft_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 2048 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 6 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_layout = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.num_layers_in_first_pipeline_stage = 14 + cfg.model.num_layers_in_last_pipeline_stage = 16 + cfg.model.account_for_embedding_in_pipeline_split = False + cfg.model.account_for_loss_in_pipeline_split = False + cfg.model.microbatch_group_size_per_vp_stage = 6 + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.apply_rope_fusion = False + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + peft_cfg = default_peft_config(peft_scheme) + if isinstance(peft_scheme, str) and peft_scheme.lower() in {"lora", "dora"}: + peft_cfg.dim = 8 + peft_cfg.alpha = 16 + peft_cfg.dropout = 0.0 + peft_cfg.target_modules = [ + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_down_proj", + "linear_kv_up_proj", + "linear_proj", + ] + cfg.peft = peft_cfg + + cfg.dataset = default_tulu3_config( + seq_length=2048, + enable_offline_packing=True, + pad_seq_to_mult=4, + ) + cfg.dataset.hf_dataset.split = "train[:10000]" + cfg.dataset.hf_dataset.load_kwargs = {"revision": _TULU3_REVISION} + cfg.dataset.hf_output_root = "work/data/glm5-2/tulu3-peft-gb200" + cfg.dataset.hf_rewrite = False + cfg.dataset.hf_validation_proportion = None + cfg.dataset.max_train_samples = 10000 + cfg.dataset.seed = 1234 + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=10, + max_lr=1e-4, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + return cfg + + +__all__ = [ + "glm52_peft_192gpu_gb200_bf16_config", + "glm52_pretrain_192gpu_gb200_bf16_config", + "glm52_sft_192gpu_gb200_bf16_128k_config", + "glm52_sft_192gpu_gb200_bf16_config", +] diff --git a/src/megatron/bridge/recipes/glm/h100/__init__.py b/src/megatron/bridge/recipes/glm/h100/__init__.py index 46d25b842c4..ba0f432648a 100644 --- a/src/megatron/bridge/recipes/glm/h100/__init__.py +++ b/src/megatron/bridge/recipes/glm/h100/__init__.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from megatron.bridge.recipes.glm.h100.glm5 import * # noqa: F403 from megatron.bridge.recipes.glm.h100.glm45 import * # noqa: F403 @@ -22,4 +23,8 @@ "glm45_air_106b_peft_8gpu_h100_bf16_config", "glm45_air_106b_pretrain_32gpu_h100_bf16_config", "glm45_air_106b_sft_32gpu_h100_bf16_config", + "glm52_peft_208gpu_h100_bf16_config", + "glm52_pretrain_416gpu_h100_bf16_config", + "glm52_sft_416gpu_h100_bf16_config", + "glm52_sft_608gpu_h100_bf16_200k_config", ] diff --git a/src/megatron/bridge/recipes/glm/h100/glm5.py b/src/megatron/bridge/recipes/glm/h100/glm5.py new file mode 100644 index 00000000000..405a53a5ad5 --- /dev/null +++ b/src/megatron/bridge/recipes/glm/h100/glm5.py @@ -0,0 +1,368 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""H100 recipes for GLM-5.2.""" + +from __future__ import annotations + +from megatron.bridge import AutoBridge +from megatron.bridge.data.builders import ChatSFTPreprocessingConfig +from megatron.bridge.peft.base import PEFT +from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common +from megatron.bridge.recipes.utils.dataset_utils import default_peft_config, default_tulu3_config +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed + + +_GLM52_MODEL_ID = "zai-org/GLM-5.2" +_GLM52_MODEL_REVISION = "4d67f66cc64d3219133b767c253b2ad1425c6c88" # pragma: allowlist secret +_TULU3_REVISION = "b14afda60f1bbebe55d5d2fa1e4df5042f97f8be" # pragma: allowlist secret + +_GLM52_VPP2_LAYOUT = ( + "Et|t||tttt|tttt|tttt|tttt||tttt|tttt|tttt|tttt||tttt|tttt|tttt|tttt||tttt|tttt|tttt|tttt||tttt|tttt|ttttmL" +) +_GLM52_PP13_LAYOUT = "Etttttt|tttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttttttt|tttt|tttt|tttt|tttt|ttttmL" +_GLM52_PP19_200K_LAYOUT = ( + "Etttttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|tttt|ttttmL" +) + + +def glm52_pretrain_416gpu_h100_bf16_config() -> ConfigContainer: + """GLM-5.2 bounded pretraining on 416 H100 GPUs.""" + cfg = _pretrain_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 4096 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 13 + cfg.model.virtual_pipeline_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_layout = _GLM52_VPP2_LAYOUT + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.microbatch_group_size_per_vp_stage = 16 + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "deepep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + cfg.dataset.seq_length = 4096 + cfg.dataset.num_workers = 8 + cfg.dataset.random_seed = 1234 + cfg.rng.seed = 1234 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = False + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=40, + max_lr=3e-4, + min_lr=3e-5, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.optimizer.use_precision_aware_optimizer = True + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + return cfg + + +def glm52_sft_416gpu_h100_bf16_config() -> ConfigContainer: + """GLM-5.2 bounded full SFT on 416 H100 GPUs.""" + cfg = _sft_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 2048 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 13 + cfg.model.virtual_pipeline_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_layout = _GLM52_VPP2_LAYOUT + cfg.model.context_parallel_size = 16 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.microbatch_group_size_per_vp_stage = 16 + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "deepep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + cfg.dataset = default_tulu3_config( + seq_length=2048, + enable_offline_packing=True, + pad_seq_to_mult=32, + ) + cfg.dataset.hf_dataset.split = "train[:10000]" + cfg.dataset.hf_dataset.load_kwargs = {"revision": _TULU3_REVISION} + cfg.dataset.hf_output_root = "work/data/glm5-2/tulu3-full-sft-pad32" + cfg.dataset.hf_validation_proportion = None + cfg.dataset.max_train_samples = 10000 + cfg.dataset.seed = 1234 + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=10, + max_lr=5e-6, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.save_optim = False + cfg.checkpoint.save_rng = False + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + return cfg + + +def glm52_sft_608gpu_h100_bf16_200k_config() -> ConfigContainer: + """GLM-5.2 200K packed SFT with context parallelism on 608 H100 GPUs.""" + cfg = _sft_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 200000 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 19 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_layout = _GLM52_PP19_200K_LAYOUT + cfg.model.context_parallel_size = 32 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.microbatch_group_size_per_vp_stage = None + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "deepep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + cfg.dataset.seq_length = 200000 + cfg.dataset.hf_dataset = None + cfg.dataset.dataset_root = "work/data/glm5-2/synthetic-200k" + cfg.dataset.hf_output_root = None + cfg.dataset.hf_validation_proportion = None + cfg.dataset.seed = 1234 + cfg.dataset.preprocessing = ChatSFTPreprocessingConfig() + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + cfg.dataset.offline_packing_specs.packed_sequence_size = 200000 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 64 + cfg.rng.seed = 5678 + cfg.train.train_iters = 20 + cfg.train.global_batch_size = 13 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=2, + max_lr=1e-6, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 20 + cfg.checkpoint.save = None + cfg.checkpoint.load = None + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + return cfg + + +def glm52_peft_208gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> ConfigContainer: + """GLM-5.2 bounded LoRA SFT on 208 H100 GPUs.""" + cfg = _peft_common() + + cfg.model = AutoBridge.from_hf_pretrained(_GLM52_MODEL_ID, revision=_GLM52_MODEL_REVISION).to_megatron_provider( + load_weights=False + ) + cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} + + cfg.model.seq_length = 2048 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 13 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_layout = _GLM52_PP13_LAYOUT + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 16 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.microbatch_group_size_per_vp_stage = None + + cfg.model.dsa_kernel_backend = "cudnn" + cfg.model.mtp_num_layers = 1 + cfg.model.calculate_per_token_loss = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.gradient_accumulation_fusion = True + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "deepep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.persist_layer_norm = True + cfg.model.bias_dropout_fusion = True + cfg.model.bias_activation_fusion = True + cfg.model.recompute_granularity = "full" + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 + + cfg.mixed_precision = bf16_mixed() + cfg.ddp.average_in_collective = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 10 + cfg.validation.eval_iters = 0 + cfg.validation.eval_interval = 0 + cfg.logger.log_interval = 1 + + peft_cfg = default_peft_config(peft_scheme) + if isinstance(peft_scheme, str) and peft_scheme.lower() in {"lora", "dora"}: + peft_cfg.dim = 8 + peft_cfg.alpha = 16 + peft_cfg.dropout = 0.0 + peft_cfg.target_modules = [ + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_down_proj", + "linear_kv_up_proj", + "linear_proj", + ] + cfg.peft = peft_cfg + + cfg.dataset = default_tulu3_config( + seq_length=2048, + enable_offline_packing=True, + pad_seq_to_mult=4, + ) + cfg.dataset.hf_dataset.split = "train[:10000]" + cfg.dataset.hf_dataset.load_kwargs = {"revision": _TULU3_REVISION} + cfg.dataset.hf_output_root = "work/data/glm5-2/tulu3-peft-pad4" + cfg.dataset.hf_validation_proportion = None + cfg.dataset.max_train_samples = 10000 + cfg.dataset.seed = 1234 + cfg.dataset.do_validation = False + cfg.dataset.do_test = False + + cfg.rng.seed = 5678 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=10, + max_lr=1e-4, + min_lr=0.0, + ) + cfg.scheduler.lr_decay_iters = 100 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + return cfg + + +__all__ = [ + "glm52_peft_208gpu_h100_bf16_config", + "glm52_pretrain_416gpu_h100_bf16_config", + "glm52_sft_416gpu_h100_bf16_config", + "glm52_sft_608gpu_h100_bf16_200k_config", +] diff --git a/tests/functional_tests/test_groups/models/glm_moe_dsa/test_glm5_conversion.py b/tests/functional_tests/test_groups/models/glm_moe_dsa/test_glm5_conversion.py index b8bd1fe998f..ab16bab1b75 100644 --- a/tests/functional_tests/test_groups/models/glm_moe_dsa/test_glm5_conversion.py +++ b/tests/functional_tests/test_groups/models/glm_moe_dsa/test_glm5_conversion.py @@ -20,6 +20,7 @@ import pytest import torch from packaging.version import Version +from safetensors import safe_open from transformers import AutoConfig, AutoTokenizer from transformers import __version__ as TRANSFORMERS_VERSION @@ -44,6 +45,7 @@ "index_head_dim": 128, "index_n_heads": 8, "index_topk": 256, + "indexer_types": ["full", "full"], "indexer_rope_interleave": True, # ---- LoRA ranks ---- "q_lora_rank": 256, @@ -61,6 +63,7 @@ "scoring_func": "sigmoid", "topk_method": "noaux_tc", "mlp_layer_types": ["dense", "sparse"], + "layer_types": ["deepseek_sparse_attention", "deepseek_sparse_attention"], # ---- Position encoding ---- "max_position_embeddings": 8192, "rope_interleave": True, @@ -87,18 +90,31 @@ "transformers_version": "5.2.0.dev0", } +HF_GLM52_INDEXSHARE_TOY_MODEL_CONFIG = { + **HF_GLM5_TOY_MODEL_CONFIG, + "num_hidden_layers": 4, + "index_topk_freq": 4, + "index_skip_topk_offset": 3, + "indexer_types": ["full", "full", "full", "shared"], + "mlp_layer_types": ["dense", "sparse", "sparse", "sparse"], + "layer_types": ["deepseek_sparse_attention"] * 4, + # Transformers does not instantiate the appended MTP layer. MTP indexer + # names are covered separately by the bridge mapping-registry unit test. + "num_nextn_predict_layers": 0, +} + pytestmark = pytest.mark.skipif( Version(TRANSFORMERS_VERSION) < Version("5.2.0"), reason=f"GLM5 conversion tests require transformers>=5.2.0, found {TRANSFORMERS_VERSION}", ) -def _create_glm5_toy_model(model_dir: Path) -> None: +def _create_glm5_toy_model(model_dir: Path, model_config: dict[str, object]) -> None: model_dir.mkdir(parents=True, exist_ok=True) config = AutoConfig.from_pretrained("zai-org/GLM-5") - for key, value in HF_GLM5_TOY_MODEL_CONFIG.items(): + for key, value in model_config.items(): setattr(config, key, value) config.torch_dtype = torch.bfloat16 @@ -117,12 +133,25 @@ def _create_glm5_toy_model(model_dir: Path) -> None: model.save_pretrained(model_dir, safe_serialization=True) - config_to_save = HF_GLM5_TOY_MODEL_CONFIG.copy() + config_to_save = model_config.copy() config_path = model_dir / "config.json" with open(config_path, "w") as f: json.dump(config_to_save, f, indent=2) +def _safetensor_keys(model_dir: Path) -> set[str]: + index_path = model_dir / "model.safetensors.index.json" + if index_path.exists(): + with open(index_path) as f: + return set(json.load(f)["weight_map"]) + + keys: set[str] = set() + for checkpoint_path in model_dir.glob("*.safetensors"): + with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint: + keys.update(checkpoint.keys()) + return keys + + class TestGLM5Conversion: """ Test GLM-5 (MoE + MLA + DSA) model conversion with different parallelism configurations. @@ -135,7 +164,17 @@ def glm5_toy_model_path(self, tmp_path_factory): temp_dir = tmp_path_factory.mktemp("glm5_toy_model") model_dir = temp_dir / "glm5_toy" - _create_glm5_toy_model(model_dir) + _create_glm5_toy_model(model_dir, HF_GLM5_TOY_MODEL_CONFIG) + + return str(model_dir) + + @pytest.fixture(scope="class") + def glm52_indexshare_toy_model_path(self, tmp_path_factory): + """Create a GLM-5.2-style toy model with one shared-indexer layer.""" + temp_dir = tmp_path_factory.mktemp("glm52_indexshare_toy_model") + model_dir = temp_dir / "glm52_indexshare_toy" + + _create_glm5_toy_model(model_dir, HF_GLM52_INDEXSHARE_TOY_MODEL_CONFIG) return str(model_dir) @@ -191,8 +230,16 @@ def test_toy_model_creation(self, glm5_toy_model_path): assert hasattr(model.model, "layers") assert len(model.model.layers) == 2 - second_layer = model.model.layers[1] - assert hasattr(second_layer, "mlp") + assert hasattr(model.model.layers[1], "mlp") + + def test_glm52_indexshare_toy_model_creation(self, glm52_indexshare_toy_model_path): + """The HF fixture omits all five indexer tensors from its shared layer.""" + model_path = Path(glm52_indexshare_toy_model_path) + indexer_keys = {key for key in _safetensor_keys(model_path) if ".self_attn.indexer." in key} + indexer_layer_ids = {int(key.split(".")[2]) for key in indexer_keys} + + assert indexer_layer_ids == {0, 1, 2} + assert len(indexer_keys) == 15 @pytest.mark.run_only_on("GPU") @pytest.mark.parametrize( @@ -275,6 +322,53 @@ def test_glm5_conversion_parallelism(self, glm5_toy_model_path, tmp_path, tp, pp assert saved_config["num_experts_per_tok"] == HF_GLM5_TOY_MODEL_CONFIG["num_experts_per_tok"] assert saved_config["moe_intermediate_size"] == HF_GLM5_TOY_MODEL_CONFIG["moe_intermediate_size"] + @pytest.mark.run_only_on("GPU") + def test_glm52_indexshare_strict_roundtrip(self, glm52_indexshare_toy_model_path, tmp_path): + """Full layers round-trip exactly while shared layers omit indexer tensors.""" + test_output_dir = tmp_path / "glm52_indexshare" + test_output_dir.mkdir(exist_ok=True) + + repo_root = "/opt/Megatron-Bridge" + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + "--nproc_per_node=2", + "--nnodes=1", + "-m", + "coverage", + "run", + f"--data-file={repo_root}/.coverage", + f"--source={repo_root}/", + "--parallel-mode", + f"{repo_root}/examples/conversion/hf_megatron_roundtrip_multi_gpu.py", + "--hf-model-id", + glm52_indexshare_toy_model_path, + "--output-dir", + str(test_output_dir), + "--ep", + "2", + "--strict", + "--atol", + "0", + "--rtol", + "0", + ] + + result = subprocess.run(cmd, capture_output=True, text=True, cwd=repo_root) + if result.returncode != 0: + pytest.fail( + "GLM-5.2 IndexShare strict round-trip failed with return code " + f"{result.returncode}\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + + converted_model_dir = test_output_dir / Path(glm52_indexshare_toy_model_path).name + exported_indexer_keys = {key for key in _safetensor_keys(converted_model_dir) if ".self_attn.indexer." in key} + exported_indexer_layer_ids = {int(key.split(".")[2]) for key in exported_indexer_keys} + + assert exported_indexer_layer_ids == {0, 1, 2} + assert len(exported_indexer_keys) == 15 + @pytest.mark.run_only_on("GPU") def test_glm5_autoconfig_roundtrip(self, glm5_toy_model_path, tmp_path): from tests.functional_tests.utils import autoconfig_roundtrip diff --git a/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py b/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py index e45d3490037..7ca31071fd8 100644 --- a/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py +++ b/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py @@ -23,6 +23,7 @@ _SCRIPT = Path(__file__).parents[3] / "examples" / "conversion" / "hf_to_megatron_generate_text.py" _SCRIPT_GLOBALS = runpy.run_path(_SCRIPT) _build_parser = _SCRIPT_GLOBALS["build_parser"] +_build_inference_context = _SCRIPT_GLOBALS["_build_inference_context"] _decode_completion = _SCRIPT_GLOBALS["_decode_completion"] _hf_revision_kwargs = _SCRIPT_GLOBALS["_hf_revision_kwargs"] _maybe_gather_tensor_parallel_logits = _SCRIPT_GLOBALS["_maybe_gather_tensor_parallel_logits"] @@ -85,6 +86,43 @@ def test_hf_revision_is_parsed_and_forwarded() -> None: assert _hf_revision_kwargs(None) == {} +@pytest.mark.unit +def test_pipeline_model_parallel_layout_is_parsed() -> None: + layout = "Etttttt|ttttmL" + + args = _build_parser().parse_args(["--hf_model_path", "org/model", "--pipeline-model-parallel-layout", layout]) + + assert args.pipeline_model_parallel_layout == layout + + +@pytest.mark.unit +def test_legacy_full_prefix_disables_inference_context() -> None: + input_ids = torch.ones((2, 3), dtype=torch.long) + + default_args = _build_parser().parse_args(["--hf_model_path", "org/model"]) + legacy_args = _build_parser().parse_args(["--hf_model_path", "org/model", "--legacy-full-prefix"]) + + assert default_args.legacy_full_prefix is False + assert legacy_args.legacy_full_prefix is True + assert _build_inference_context(input_ids, legacy_full_prefix=True) is None + + +@pytest.mark.unit +def test_kv_cache_builds_static_inference_context() -> None: + input_ids = torch.ones((2, 3), dtype=torch.long) + inference_context = object() + context_factory = MagicMock(return_value=inference_context) + + with patch.dict( + _build_inference_context.__globals__, + {"StaticInferenceContext": context_factory}, + ): + result = _build_inference_context(input_ids, legacy_full_prefix=False) + + assert result is inference_context + context_factory.assert_called_once_with(max_batch_size=2, max_sequence_length=3) + + @pytest.mark.unit def test_text_forward_step_passes_static_context_and_gather_request() -> None: inference_context = object() diff --git a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py index 2f58aabfb68..d7848bdc8c0 100644 --- a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py +++ b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py @@ -26,6 +26,14 @@ pytestmark = pytest.mark.unit +_DSA_INDEXER_SUFFIXES = { + "linear_wq_b.weight": "wq_b.weight", + "linear_wk.weight": "wk.weight", + "k_norm.weight": "k_norm.weight", + "k_norm.bias": "k_norm.bias", + "linear_weights_proj.weight": "weights_proj.weight", +} + @pytest.fixture def glm5_bridge() -> GLM5Bridge: @@ -45,6 +53,11 @@ def _provider_from_hf_config(monkeypatch: pytest.MonkeyPatch, **config_overrides "num_hidden_layers": 78, "moe_intermediate_size": 2048, "n_shared_experts": 1, + "qk_head_dim": 256, + "qk_nope_head_dim": 192, + # Transformers currently normalizes this raw-config value from 64 to + # head_dim=192 for GLM-5.2. The bridge must not propagate that value. + "qk_rope_head_dim": 192, "rope_parameters": {"rope_theta": 1_000_000}, "index_head_dim": 128, "index_n_heads": 32, @@ -115,6 +128,7 @@ def test_provider_bridge_maps_dsa_architecture_from_hf_config( assert provider.experimental_attention_variant == "dsa" assert provider.cp_comm_type == "allgather" + assert provider.qk_pos_emb_head_dim == 64 assert provider.dsa_indexer_head_dim == 128 assert provider.dsa_indexer_n_heads == 32 assert provider.dsa_indexer_topk == 2048 @@ -202,6 +216,34 @@ def test_mapping_registry_includes_mtp_attention_and_dense_mlp_mappings( } +def test_mapping_registry_includes_dsa_indexer_mappings(glm5_bridge: GLM5Bridge) -> None: + """Every tensor on a full GLM-5.2 DSA indexer maps to its HF peer.""" + mappings = _mapping_by_megatron_param(glm5_bridge) + registry = glm5_bridge.mapping_registry() + + for megatron_suffix, hf_suffix in _DSA_INDEXER_SUFFIXES.items(): + megatron_param = f"decoder.layers.*.self_attention.core_attention.indexer.{megatron_suffix}" + mapping = mappings[megatron_param] + assert isinstance(mapping, AutoMapping) + assert mapping.hf_param == f"model.layers.*.self_attn.indexer.{hf_suffix}" + + resolved_mapping = registry.megatron_to_hf_lookup(megatron_param.replace("*", "2")) + assert resolved_mapping is not None + assert resolved_mapping.hf_param == f"model.layers.2.self_attn.indexer.{hf_suffix}" + + +@pytest.mark.parametrize("layer_prefix", ["transformer_layer", "mtp_model_layer"]) +def test_mapping_registry_includes_mtp_dsa_indexer_mappings(glm5_bridge: GLM5Bridge, layer_prefix: str) -> None: + """A full GLM-5.2 MTP layer maps every DSA indexer tensor.""" + mappings = _mapping_by_megatron_param(glm5_bridge) + + for megatron_suffix, hf_suffix in _DSA_INDEXER_SUFFIXES.items(): + megatron_param = f"mtp.layers.0.{layer_prefix}.self_attention.core_attention.indexer.{megatron_suffix}" + mapping = mappings[megatron_param] + assert isinstance(mapping, AutoMapping) + assert mapping.hf_param == f"model.layers.4.self_attn.indexer.{hf_suffix}" + + def test_mapping_registry_includes_mtp_standalone_weights(glm5_bridge: GLM5Bridge) -> None: """GLM-5 MTP-only weights map to the appended HF MTP layer.""" mappings = _mapping_by_megatron_param(glm5_bridge) diff --git a/tests/unit_tests/recipes/test_glm5_recipes.py b/tests/unit_tests/recipes/test_glm5_recipes.py new file mode 100644 index 00000000000..ded9ad3eb47 --- /dev/null +++ b/tests/unit_tests/recipes/test_glm5_recipes.py @@ -0,0 +1,304 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for GLM-5.2 recipes.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from megatron.bridge.recipes.glm import gb200 +from megatron.bridge.recipes.glm.gb200 import glm5 as gb200_glm5 +from megatron.bridge.recipes.glm.h100 import glm5 + + +pytestmark = pytest.mark.unit + + +class _FakeMegatronProvider(SimpleNamespace): + _KNOWN_FIELDS = { + "account_for_embedding_in_pipeline_split", + "account_for_loss_in_pipeline_split", + "apply_rope_fusion", + "bias_activation_fusion", + "bias_dropout_fusion", + "calculate_per_token_loss", + "context_parallel_size", + "cross_entropy_fusion_impl", + "cross_entropy_loss_fusion", + "dsa_indexer_loss_coeff", + "dsa_indexer_skip_topk_offset", + "dsa_indexer_topk_freq", + "dsa_indexer_use_sparse_loss", + "dsa_kernel_backend", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "gradient_accumulation_fusion", + "microbatch_group_size_per_vp_stage", + "moe_flex_dispatcher_backend", + "moe_router_force_load_balancing", + "moe_shared_expert_overlap", + "moe_token_dispatcher_type", + "mtp_num_layers", + "num_layers_in_first_pipeline_stage", + "num_layers_in_last_pipeline_stage", + "persist_layer_norm", + "pipeline_model_parallel_layout", + "pipeline_model_parallel_size", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "seq_length", + "sequence_parallel", + "tensor_model_parallel_size", + "virtual_pipeline_model_parallel_size", + } + + def __init__(self, **kwargs: object) -> None: + unknown_fields = set(kwargs) - self._KNOWN_FIELDS + if unknown_fields: + raise AssertionError(f"Unexpected GLM-5.2 model config fields: {sorted(unknown_fields)}") + super().__init__(**kwargs) + + def __setattr__(self, name: str, value: object) -> None: + if name not in self._KNOWN_FIELDS: + raise AssertionError(f"Unexpected GLM-5.2 model config field: {name}") + super().__setattr__(name, value) + + +class _FakeAutoBridge: + @classmethod + def from_hf_pretrained(cls, model_id: str, revision: str) -> "_FakeAutoBridge": + assert model_id == "zai-org/GLM-5.2" + assert len(revision) == 40 + return cls() + + def to_megatron_provider(self, load_weights: bool = False) -> _FakeMegatronProvider: + assert load_weights is False + return _FakeMegatronProvider( + dsa_indexer_loss_coeff=0.001, + dsa_indexer_use_sparse_loss=True, + dsa_indexer_topk_freq=4, + dsa_indexer_skip_topk_offset=3, + mtp_num_layers=None, + ) + + +@pytest.fixture(autouse=True) +def _patch_autobridge(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gb200_glm5, "AutoBridge", _FakeAutoBridge) + monkeypatch.setattr(glm5, "AutoBridge", _FakeAutoBridge) + + +@pytest.mark.parametrize( + ("recipe", "world_size", "pp", "cp", "ep", "gbs", "steps"), + [ + (glm5.glm52_pretrain_416gpu_h100_bf16_config, 416, 13, 1, 32, 1024, 100), + (glm5.glm52_sft_416gpu_h100_bf16_config, 416, 13, 16, 32, 32, 100), + (glm5.glm52_sft_608gpu_h100_bf16_200k_config, 608, 19, 32, 32, 13, 20), + (glm5.glm52_peft_208gpu_h100_bf16_config, 208, 13, 1, 16, 32, 100), + ], +) +def test_glm52_h100_recipe_topologies(recipe, world_size, pp, cp, ep, gbs, steps) -> None: + cfg = recipe() + + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == pp + assert cfg.model.context_parallel_size == cp + assert cfg.model.expert_model_parallel_size == ep + assert cfg.train.global_batch_size == gbs + assert cfg.train.micro_batch_size == 1 + assert cfg.train.train_iters == steps + assert cfg.model.dsa_kernel_backend == "cudnn" + assert cfg.model.mtp_num_layers == 1 + assert cfg.model.dsa_indexer_loss_coeff == 0.001 + assert cfg.model.dsa_indexer_use_sparse_loss is True + assert cfg.model.moe_router_force_load_balancing is False + assert world_size % (pp * cp) == 0 + assert world_size % (pp * ep) == 0 + + +def test_glm52_h100_200k_recipe_uses_packed_cp() -> None: + cfg = glm5.glm52_sft_608gpu_h100_bf16_200k_config() + + assert cfg.model.seq_length == 200000 + assert cfg.dataset.seq_length == 200000 + assert cfg.dataset.offline_packing_specs.packed_sequence_size == 200000 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 64 + assert cfg.model.virtual_pipeline_model_parallel_size is None + assert cfg.model.microbatch_group_size_per_vp_stage is None + assert cfg.model.pipeline_model_parallel_layout == glm5._GLM52_PP19_200K_LAYOUT + stages = cfg.model.pipeline_model_parallel_layout.split("|") + assert [stage.count("t") for stage in stages] == [6] + [4] * 18 + decoder_starts = [] + decoder_count = 0 + for stage in stages: + decoder_starts.append(decoder_count) + decoder_count += stage.count("t") + assert decoder_starts == [0, *range(6, 78, 4)] + assert cfg.dataset.dataset_root == "work/data/glm5-2/synthetic-200k" + assert cfg.dataset.hf_dataset is None + + +@pytest.mark.parametrize( + ("recipe", "gbs", "steps"), + [ + (gb200.glm52_pretrain_192gpu_gb200_bf16_config, 1024, 100), + (gb200.glm52_sft_192gpu_gb200_bf16_config, 32, 100), + (gb200.glm52_peft_192gpu_gb200_bf16_config, 32, 100), + ], +) +def test_glm52_gb200_recipe_topologies(recipe, gbs, steps) -> None: + cfg = recipe() + + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 6 + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 32 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.sequence_parallel is False + assert cfg.model.virtual_pipeline_model_parallel_size is None + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.model.num_layers_in_first_pipeline_stage == 14 + assert cfg.model.num_layers_in_last_pipeline_stage == 16 + assert cfg.model.microbatch_group_size_per_vp_stage == 6 + assert cfg.model.moe_token_dispatcher_type == "alltoall" + assert cfg.model.moe_flex_dispatcher_backend is None + assert cfg.model.apply_rope_fusion is False + assert cfg.train.global_batch_size == gbs + assert cfg.train.micro_batch_size == 1 + assert cfg.train.train_iters == steps + + +def test_glm52_gb200_128k_recipe_uses_packed_cp() -> None: + cfg = gb200.glm52_sft_192gpu_gb200_bf16_128k_config() + + assert cfg.model.seq_length == 131072 + assert cfg.model.pipeline_model_parallel_size == 6 + assert cfg.model.context_parallel_size == 32 + assert cfg.model.expert_model_parallel_size == 32 + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.model.num_layers_in_first_pipeline_stage == 14 + assert cfg.model.num_layers_in_last_pipeline_stage == 16 + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_shared_expert_overlap is False + assert cfg.model.apply_rope_fusion is False + assert cfg.train.global_batch_size == 8 + assert cfg.train.micro_batch_size == 1 + assert cfg.train.train_iters == 20 + assert cfg.dataset.seq_length == 131072 + assert cfg.dataset.dataset_root == "work/data/glm5-2/synthetic-long-context-gb200" + assert cfg.dataset.offline_packing_specs.packed_sequence_size == 131072 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 64 + assert cfg.dataset.offline_packing_specs.pad_cu_seqlens is True + assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 32 + assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 128 + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 + assert cfg.env_vars["USE_MNNVL"] == 1 + + +def test_glm52_gb200_recipes_do_not_depend_on_h100_recipes() -> None: + source = Path(gb200_glm5.__file__).read_text(encoding="utf-8") + + assert "megatron.bridge.recipes.glm.h100" not in source + + +@pytest.mark.parametrize( + "recipe", + [ + glm5.glm52_pretrain_416gpu_h100_bf16_config, + glm5.glm52_sft_416gpu_h100_bf16_config, + ], +) +def test_glm52_vpp2_default_layout_shape(recipe) -> None: + cfg = recipe() + + assert cfg.model.virtual_pipeline_model_parallel_size == 2 + assert cfg.model.pipeline_model_parallel_layout == glm5._GLM52_VPP2_LAYOUT + stages = cfg.model.pipeline_model_parallel_layout.split("|") + assert len(stages) == cfg.model.pipeline_model_parallel_size * cfg.model.virtual_pipeline_model_parallel_size + assert [stage.count("t") for stage in stages] == [ + 1, + 1, + 0, + 4, + 4, + 4, + 4, + 0, + 4, + 4, + 4, + 4, + 0, + 4, + 4, + 4, + 4, + 0, + 4, + 4, + 4, + 4, + 0, + 4, + 4, + 4, + ] + assert sum(stage.count("t") for stage in stages) == 78 + assert cfg.model.pipeline_model_parallel_layout.count("E") == 1 + assert cfg.model.pipeline_model_parallel_layout.count("m") == 1 + assert cfg.model.pipeline_model_parallel_layout.count("L") == 1 + + +def test_glm52_pretrain_uses_reference_gradient_path() -> None: + cfg = glm5.glm52_pretrain_416gpu_h100_bf16_config() + + assert cfg.mixed_precision.grad_reduce_in_fp32 is False + assert cfg.ddp.grad_reduce_in_fp32 is False + assert cfg.ddp.average_in_collective is False + assert cfg.optimizer.use_precision_aware_optimizer is True + + +def test_glm52_peft_targets_mla_attention_projections() -> None: + cfg = glm5.glm52_peft_208gpu_h100_bf16_config() + + assert cfg.peft.dim == 8 + assert cfg.peft.alpha == 16 + assert cfg.peft.dropout == 0.0 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert cfg.dataset.hf_dataset.split == "train[:10000]" + assert cfg.peft.target_modules == [ + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_down_proj", + "linear_kv_up_proj", + "linear_proj", + ] + + +def test_glm52_platform_recipes_are_exported() -> None: + import megatron.bridge.recipes as recipes + from megatron.bridge.recipes import glm as glm_recipes + from megatron.bridge.recipes.glm import h100 + + for recipe_name in gb200.__all__: + assert getattr(glm_recipes, recipe_name) is getattr(gb200, recipe_name) + assert getattr(recipes, recipe_name) is getattr(gb200, recipe_name) + assert recipe_name in glm_recipes.__all__ + for recipe_name in glm5.__all__: + assert getattr(h100, recipe_name) is getattr(glm5, recipe_name) + assert getattr(recipes, recipe_name) is getattr(glm5, recipe_name) + assert recipe_name in h100.__all__ diff --git a/tests/unit_tests/scripts/inference/test_setup_inference.py b/tests/unit_tests/scripts/inference/test_setup_inference.py index fc7b669d89d..cdcbebc80ce 100644 --- a/tests/unit_tests/scripts/inference/test_setup_inference.py +++ b/tests/unit_tests/scripts/inference/test_setup_inference.py @@ -124,6 +124,10 @@ def test_parser_forwards_model_checkpoint_prompt_and_engine_args(): ("task_name", "expected_path"), [ ("text-generation", "/opt/Megatron-Bridge/scripts/inference/text_generation.py"), + ( + "legacy-full-prefix-generation", + "/opt/Megatron-Bridge/examples/conversion/hf_to_megatron_generate_text.py", + ), ("vlm-generation", "/opt/Megatron-Bridge/scripts/inference/vlm_generation.py"), ( "model-comparison", @@ -186,6 +190,20 @@ def test_resource_validation_rejects_invalid_values(options, message): module._validate_args(args) +@pytest.mark.parametrize( + ("task_name", "inference_args", "message"), + [ + ("legacy-full-prefix-generation", [], "requires --legacy-full-prefix"), + ("text-generation", ["--legacy-full-prefix"], "requires --task legacy-full-prefix-generation"), + ], +) +def test_task_validation_rejects_inconsistent_legacy_full_prefix_usage(task_name, inference_args, message): + module = _load_setup_inference_module() + + with pytest.raises(ValueError, match=message): + module._validate_task_args(task_name, inference_args) + + def test_parse_env_deduplicates_names_and_rejects_values(monkeypatch): module = _load_setup_inference_module() monkeypatch.setenv("HF_TOKEN", "secret") @@ -307,6 +325,20 @@ def test_build_task_quotes_prompts_and_uses_existing_entrypoint(): assert scripts[0].args == ["--prompt", "'benign; echo should-not-run'"] +def test_build_task_preserves_legacy_full_prefix_argument(): + module = _load_setup_inference_module() + scripts = [] + module.run.Script = lambda **kwargs: scripts.append(types.SimpleNamespace(**kwargs)) or scripts[-1] + + module._build_task( + "legacy-full-prefix-generation", + ["--hf_model_path", "zai-org/GLM-5.2", "--legacy-full-prefix"], + ) + + assert scripts[0].path == "/opt/Megatron-Bridge/examples/conversion/hf_to_megatron_generate_text.py" + assert scripts[0].args == ["--hf_model_path", "zai-org/GLM-5.2", "--legacy-full-prefix"] + + @pytest.mark.parametrize( ("extra_options", "expected_run", "expected_dryrun"), [ diff --git a/tests/unit_tests/scripts/test_validate_model_verification_card.py b/tests/unit_tests/scripts/test_validate_model_verification_card.py index 018574e2ddf..2d8300fea03 100644 --- a/tests/unit_tests/scripts/test_validate_model_verification_card.py +++ b/tests/unit_tests/scripts/test_validate_model_verification_card.py @@ -71,13 +71,34 @@ def test_verified_inference_accepts_canonical_bash_launcher(): errors=errors, command_override=( "./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 " - "--task vlm-generation --prompt hello --max_new_tokens 1" + "--task legacy-full-prefix-generation --legacy-full-prefix --prompt hello --max_new_tokens 1" ), ) assert errors == [] +def test_verified_inference_legacy_task_requires_full_prefix_flag(): + module = _load_validator() + errors = [] + item = { + "expected_result": 'The exact 1-token result produced completion "ok".', + } + + module._validate_inference( + item, + item_name="inference", + status="verified", + errors=errors, + command_override=( + "./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 " + "--task legacy-full-prefix-generation --prompt hello --max_new_tokens 1" + ), + ) + + assert "/items/inference/command: legacy-full-prefix-generation requires --legacy-full-prefix" in errors + + def test_verified_inference_bash_launcher_requires_task_and_resources(): module = _load_validator() errors = [] From 9b5e30444b69c48b5221a6b51f5d1be2b2d3bb22 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 23:39:01 +0200 Subject: [PATCH 42/85] cp: `fix(ckpt): validate conversion world size against TP*PP, not TP*PP*EP (5302)` into `r0.6.0` (#5312) Signed-off-by: Pruthviraj Prakash Signed-off-by: NeMo Bot Co-authored-by: pruprakash --- scripts/conversion/README.md | 8 +++++--- scripts/conversion/setup_conversion.py | 6 +++--- .../conversion/launcher/test_setup_conversion.py | 14 ++++++++++---- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/scripts/conversion/README.md b/scripts/conversion/README.md index 4014628b415..1c9f237ce86 100644 --- a/scripts/conversion/README.md +++ b/scripts/conversion/README.md @@ -42,9 +42,11 @@ The GPU backend uses NeMo Run's torchrun launcher for local execution and srun-native tasks for Slurm. Users should not wrap the command in `torchrun`, `srun`, or `sbatch`. -The requested topology must satisfy -`nodes * gpus-per-node == TP * PP * EP`; conversion does not create redundant -data-parallel replicas. `ETP * EP * PP` must also divide the total GPU count. +The requested topology must satisfy `nodes * gpus-per-node % (TP * PP) == 0` and +`nodes * gpus-per-node % (ETP * EP * PP) == 0`. Expert parallelism is an +alternative slicing of the same ranks rather than an extra multiplicand, so it +is not part of the first product. Ranks left over after either split form +data-parallel replicas. ```bash export HF_TOKEN="$(<${HOME}/HF_TOKEN)" diff --git a/scripts/conversion/setup_conversion.py b/scripts/conversion/setup_conversion.py index 54d17010b63..feb7ae798a2 100755 --- a/scripts/conversion/setup_conversion.py +++ b/scripts/conversion/setup_conversion.py @@ -119,9 +119,9 @@ def _validate_args(args: argparse.Namespace) -> None: "metacharacters through NeMo Run 0.10; use shell-safe names and paths." ) world_size = args.nodes * args.gpus_per_node - model_parallel_size = args.tp * args.pp * args.ep - if world_size != model_parallel_size: - raise ValueError("nodes*gpus-per-node must equal TP*PP*EP.") + model_parallel_size = args.tp * args.pp + if world_size % model_parallel_size != 0: + raise ValueError("nodes*gpus-per-node must be divisible by TP*PP.") expert_model_parallel_size = args.etp * args.ep * args.pp if world_size % expert_model_parallel_size != 0: raise ValueError("nodes*gpus-per-node must be divisible by ETP*EP*PP.") diff --git a/tests/unit_tests/conversion/launcher/test_setup_conversion.py b/tests/unit_tests/conversion/launcher/test_setup_conversion.py index e9e92636c5f..d025ccc60b5 100644 --- a/tests/unit_tests/conversion/launcher/test_setup_conversion.py +++ b/tests/unit_tests/conversion/launcher/test_setup_conversion.py @@ -162,16 +162,22 @@ def test_gpu_backend_validates_parallelism_against_world_size(): module = _load_setup_conversion_module() args = _parse(module, "--device", "gpu", "--gpus-per-node", "4", "--tp", "3") - with pytest.raises(ValueError, match=r"nodes\*gpus-per-node must equal TP\*PP\*EP"): + with pytest.raises(ValueError, match=r"nodes\*gpus-per-node must be divisible by TP\*PP"): module._validate_args(args) -def test_gpu_backend_rejects_replicated_data_parallel_conversion(): +def test_gpu_backend_allows_replicated_data_parallel_conversion(): module = _load_setup_conversion_module() args = _parse(module, "--device", "gpu", "--gpus-per-node", "8", "--tp", "2", "--ep", "2") - with pytest.raises(ValueError, match=r"nodes\*gpus-per-node must equal TP\*PP\*EP"): - module._validate_args(args) + module._validate_args(args) + + +def test_gpu_backend_allows_tensor_parallelism_without_expert_parallelism(): + module = _load_setup_conversion_module() + args = _parse(module, "--device", "gpu", "--gpus-per-node", "8", "--tp", "4") + + module._validate_args(args) def test_roundtrip_accepts_exact_multinode_product_topology(): From 659077e31d2f032b021c1725aa8ef462adc1ae39 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 23:39:33 +0200 Subject: [PATCH 43/85] cp: `[docs] fix: point Gemma 4 VL inference to verification card (5311)` into `r0.6.0` (#5319) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- examples/models/gemma/gemma4_vl/README.md | 62 +++----------------- examples/models/gemma/gemma4_vl/inference.sh | 45 -------------- 2 files changed, 8 insertions(+), 99 deletions(-) delete mode 100644 examples/models/gemma/gemma4_vl/inference.sh diff --git a/examples/models/gemma/gemma4_vl/README.md b/examples/models/gemma/gemma4_vl/README.md index 441788391ed..40365461662 100644 --- a/examples/models/gemma/gemma4_vl/README.md +++ b/examples/models/gemma/gemma4_vl/README.md @@ -50,60 +50,14 @@ See the [conversion.sh](conversion.sh) script for more examples including multi- ## Inference -Gemma 4 uses VLM inference script (`hf_to_megatron_generate_vlm.py`) as other models. The script auto-detects the bridge type and switches to Gemma4-specific input preprocessing, attention mask handling, and stop tokens. - -### Text-only - -```bash -uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 \ - examples/conversion/hf_to_megatron_generate_vlm.py \ - --hf_model_path google/gemma-4-26B-A4B \ - --prompt "The capital of France is" \ - --max_new_tokens 20 \ - --tp 4 --pp 2 -``` - -### Vision + Text (HF weights) - -Use the instruction-tuned model (`-it`) for image+text queries — the base model has no chat template and requires manual image token injection. - -```bash -uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 \ - examples/conversion/hf_to_megatron_generate_vlm.py \ - --hf_model_path google/gemma-4-26B-A4B-it \ - --image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" \ - --prompt "What is shown in this image?" \ - --max_new_tokens 50 \ - --tp 4 --pp 2 -``` - -### Vision + Text (imported Megatron checkpoint) - -```bash -uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 \ - examples/conversion/hf_to_megatron_generate_vlm.py \ - --hf_model_path google/gemma-4-26B-A4B \ - --megatron_model_path ${WORKSPACE}/models/gemma-4-26B-A4B/iter_0000000 \ - --image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG" \ - --prompt "What animal is on the candy?" \ - --max_new_tokens 50 \ - --tp 4 --pp 2 -``` - -Note: when loading a Megatron checkpoint for VLM inference, use the base model (`gemma-4-26B-A4B`, not `-it`) as the `--hf_model_path` to match the checkpoint's tokenizer. - -See the [inference.sh](inference.sh) script for all three steps. - -**Expected output (cat image):** -``` -======== GENERATED TEXT OUTPUT ======== -Image: https://.../pipeline-cat-chonk.jpeg -Prompt: What is shown in this image? -New tokens: The image shows a large, fluffy orange and white cat sitting inside -what appears to be a wire cage or kennel. The cat looks quite large and appears -relaxed, with its paws tucked underneath its body. -======================================= -``` +Use the [Gemma 4 26B-A4B-it model verification +card](../../../model_verification_cards/gemma-4-26b-a4b-it/card.yaml) as the canonical source for supported inference. +The card records the immutable Hugging Face revision, verified Bridge commit, checkpoint prerequisite, parallelism +topology, input image, prompt, expected output, and last verification date. + +Run only an inference item whose card status is `verified`, and use its command without changing the recorded topology or +inputs when reproducing the verified result. Other combinations have not passed the same verification gate and should not +be inferred from older standalone examples. ## Finetune Recipes diff --git a/examples/models/gemma/gemma4_vl/inference.sh b/examples/models/gemma/gemma4_vl/inference.sh deleted file mode 100644 index d2b430ea853..00000000000 --- a/examples/models/gemma/gemma4_vl/inference.sh +++ /dev/null @@ -1,45 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Workspace directory for checkpoints and results -WORKSPACE=${WORKSPACE:-/workspace} - -# Inference with HuggingFace checkpoints (text only) -uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_vlm.py \ - --hf_model_path google/gemma-4-26B-A4B \ - --prompt "The capital of France is" \ - --max_new_tokens 20 \ - --tp 4 \ - --pp 2 - -# Inference with HuggingFace checkpoints (vision + text) -uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_vlm.py \ - --hf_model_path google/gemma-4-26B-A4B-it \ - --image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" \ - --prompt "What is shown in this image?" \ - --max_new_tokens 50 \ - --tp 4 \ - --pp 2 - -# Inference with imported Megatron checkpoints (IT model, VLM) -# conversion.sh imports both the base and IT models; step 3 uses the IT checkpoint. -uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_vlm.py \ - --hf_model_path google/gemma-4-26B-A4B-it \ - --megatron_model_path ${WORKSPACE}/models/gemma-4-26B-A4B-it/iter_0000000 \ - --image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG" \ - --prompt "What animal is on the candy?" \ - --max_new_tokens 50 \ - --tp 4 \ - --pp 2 From 3490b5952d0aa694195f3caf0c9177ffa75145dc Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 7 Aug 2026 14:40:10 -0700 Subject: [PATCH 44/85] fix(model): enable RADIO recompute for Nemotron Omni SFT (#5320) Signed-off-by: Chen Cui --- .../models/nemotron/nemotron_3_omni/README.md | 16 +- .../nemotron_3_omni/slurm_sft_cord_v2.sh | 26 ++- .../nemotron_omni/nemotron_omni_provider.py | 57 +++++- .../test_nemotron_omni_provider.py | 177 ++++++++++++++++++ 4 files changed, 260 insertions(+), 16 deletions(-) create mode 100644 tests/unit_tests/models/nemotron_omni/test_nemotron_omni_provider.py diff --git a/examples/models/nemotron/nemotron_3_omni/README.md b/examples/models/nemotron/nemotron_3_omni/README.md index e6751fc969b..784ffe94b85 100644 --- a/examples/models/nemotron/nemotron_3_omni/README.md +++ b/examples/models/nemotron/nemotron_3_omni/README.md @@ -189,8 +189,9 @@ bash examples/models/nemotron/nemotron_3_omni/inference.sh ## Training All training scripts use the Nemotron-3-Nano-Omni-30B-A3B-Reasoning -pretrained checkpoint and enable in-batch sequence packing via -`dataset.enable_in_batch_packing=True`. Default GPU layout per script: +pretrained checkpoint. Sequence packing is workload-specific: CORD Full SFT +is non-packed, while LoRA enables in-batch packing. Default GPU layout per +script: - **Full SFT** — 2 nodes / 16 GPUs (full optimizer state for ~33 B params) - **LoRA PEFT** — 1 node / 8 GPUs @@ -253,8 +254,15 @@ RADIO input contract. Recipe base: `nemotron_omni_cord_v2_*_config` in | Full SFT | [slurm_sft_cord_v2.sh](slurm_sft_cord_v2.sh) | `nemotron_omni_cord_v2_sft_config` | | LoRA | [slurm_peft_cord_v2.sh](slurm_peft_cord_v2.sh) | `nemotron_omni_cord_v2_peft_config` | -Parallelism (both): TP=2, EP=8, CP=1, MBS=2, GBS=16, packed sequences, -selective recompute. LoRA targets `linear_qkv`, `linear_proj`, `in_proj`, +Parallelism (both): TP=2, EP=8, CP=1, GBS=16, selective language recompute. +Full SFT uses non-packed MBS=1 so each CORD row remains within the 4096-token +limit after the legacy LLaVA path expands its image placeholder; the unchanged +global batch is accumulated over two microbatches. Full SFT also +enables per-layer full recompute for the unfrozen RADIO tower and sets +`model.radio_force_eval_mode=False` so its training-mode recompute guard +executes. It offloads half of the optimizer state to CPU with transfer overlap +to retain GPU memory headroom; CPE remains forced to eval mode. LoRA uses packed +MBS=2 and targets `linear_qkv`, `linear_proj`, `in_proj`, `out_proj` (LM attention + Mamba projections); vision / sound encoders + projections frozen. diff --git a/examples/models/nemotron/nemotron_3_omni/slurm_sft_cord_v2.sh b/examples/models/nemotron/nemotron_3_omni/slurm_sft_cord_v2.sh index 2ed24f08676..765ee895606 100755 --- a/examples/models/nemotron/nemotron_3_omni/slurm_sft_cord_v2.sh +++ b/examples/models/nemotron/nemotron_3_omni/slurm_sft_cord_v2.sh @@ -17,12 +17,11 @@ # Nemotron-3 Nano Omni - Full SFT on CORD-V2 (image+text) # # Recipe: nemotron_omni_cord_v2_sft_config (HF dataset backend, no shards needed) -# Default parallelism: TP=2, EP=8, CP=1, MBS=2, GBS=16, packed sequences, -# selective recompute +# Default parallelism: TP=2, EP=8, CP=1, MBS=1, GBS=16, selective recompute # Default layout: 2 nodes / 16 GPUs # -# Override TP/EP/CP/PACKED_SEQ via environment, e.g.: -# TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch slurm_sft_cord_v2.sh +# Override TP/EP/CP via environment, e.g.: +# TP=4 EP=4 CP=1 sbatch slurm_sft_cord_v2.sh # # Usage: # sbatch slurm_sft_cord_v2.sh @@ -53,16 +52,17 @@ PRETRAINED_CHECKPOINT=${WORKSPACE}/models/${MODEL_NAME} RECIPE=nemotron_omni_cord_v2_sft_config DATASET_NAME=cord_v2 -# Parallelism / batching (override via env: TP=4 EP=4 CP=1 PACKED_SEQ=false sbatch ...) +# Parallelism / batching (override via env: TP=4 EP=4 CP=1 sbatch ...) TP=${TP:-2} EP=${EP:-8} CP=${CP:-1} -PACKED_SEQ=${PACKED_SEQ:-true} SEQ_LENGTH=4096 TRAIN_ITERS=4000 GLOBAL_BATCH_SIZE=16 -MICRO_BATCH_SIZE=2 +# The legacy LLaVA path expands image placeholders during model merge, so one +# non-packed CORD row per microbatch stays within SEQ_LENGTH. +MICRO_BATCH_SIZE=1 EVAL_INTERVAL=50 EVAL_ITERS=10 SAVE_INTERVAL=500 @@ -105,7 +105,7 @@ echo "Nodes: ${SLURM_JOB_NUM_NODES:-N/A}" echo "GPUs per node: ${SLURM_GPUS_PER_NODE:-N/A}" echo "Recipe: $RECIPE" echo "Checkpoint: $PRETRAINED_CHECKPOINT" -echo "Parallelism: TP=$TP EP=$EP CP=$CP (packed=$PACKED_SEQ)" +echo "Parallelism: TP=$TP EP=$EP CP=$CP (packed=false)" echo "======================================" OUTPUT_DIR=${WORKSPACE}/results/${RECIPE}_sft @@ -125,16 +125,24 @@ CLI_OVERRIDES="\ model.sequence_parallel=True \ model.recompute_granularity=selective \ model.recompute_modules=[core_attn,mlp,layernorm,moe_act,moe] \ + model.recompute_vision=True \ + model.vision_recompute_granularity=full \ + model.vision_recompute_method=uniform \ + model.vision_recompute_num_layers=1 \ + model.radio_force_eval_mode=False \ model.freeze_language_model=False \ model.freeze_vision_model=False \ model.freeze_vision_projection=False \ model.freeze_sound_encoder=False \ model.freeze_sound_projection=False \ + optimizer.optimizer_cpu_offload=True \ + optimizer.optimizer_offload_fraction=0.5 \ + optimizer.overlap_cpu_optimizer_d2h_h2d=True \ train.train_iters=$TRAIN_ITERS \ train.global_batch_size=$GLOBAL_BATCH_SIZE \ train.micro_batch_size=$MICRO_BATCH_SIZE \ dataset.trust_remote_code=True \ - dataset.enable_in_batch_packing=$PACKED_SEQ \ + dataset.enable_in_batch_packing=False \ validation.eval_interval=$EVAL_INTERVAL \ validation.eval_iters=$EVAL_ITERS \ logger.log_interval=$LOG_INTERVAL \ diff --git a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py index be266f51009..60bd68094e5 100644 --- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py +++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_provider.py @@ -83,6 +83,10 @@ class NemotronVLModelProvider(HybridModelProvider, ABC): radio_interpolate_only_cpe: bool = True radio_cpe_aspect_ratio_select: bool = False radio_disable_cpe: bool = False + recompute_vision: bool = False + vision_recompute_granularity: Literal["full", "selective"] | None = None + vision_recompute_method: Literal["uniform", "block"] | None = None + vision_recompute_num_layers: int | None = None vision_proj_ffn_hidden_size: int = 20480 vision_class_token_len: Optional[int] = None @@ -93,13 +97,12 @@ class NemotronVLModelProvider(HybridModelProvider, ABC): def _build_vision_config(self, language_cfg): """Build RADIO ViT-H vision encoder config from a language config copy.""" + if self.recompute_vision and self.radio_force_eval_mode: + raise ValueError("Vision recompute requires radio_force_eval_mode=False.") vision_cfg = copy.deepcopy(language_cfg) vision_cfg.sequence_parallel = False vision_cfg.context_parallel_size = 1 vision_cfg.tp_comm_overlap = False - vision_cfg.recompute_granularity = None - vision_cfg.recompute_method = None - vision_cfg.recompute_num_layers = None vision_cfg.mtp_num_layers = None vision_cfg.num_layers = 32 vision_cfg.num_attention_heads = 16 @@ -117,6 +120,54 @@ def _build_vision_config(self, language_cfg): vision_cfg.normalization = "LayerNorm" vision_cfg.qk_layernorm = False vision_cfg.layernorm_epsilon = 1e-6 + if self.recompute_vision: + granularity = ( + self.vision_recompute_granularity + if self.vision_recompute_granularity is not None + else vision_cfg.recompute_granularity + ) + if granularity is None: + raise ValueError("Vision recompute requires an effective recompute granularity.") + if granularity not in {"full", "selective"}: + raise ValueError("Vision recompute granularity must be 'full' or 'selective'.") + vision_cfg.recompute_granularity = granularity + + if granularity == "selective": + if self.vision_recompute_method is not None or self.vision_recompute_num_layers is not None: + raise ValueError("Selective vision recompute does not use a method or layer count.") + vision_cfg.recompute_method = None + vision_cfg.recompute_num_layers = None + else: + method = ( + self.vision_recompute_method + if self.vision_recompute_method is not None + else vision_cfg.recompute_method + ) + num_layers = ( + self.vision_recompute_num_layers + if self.vision_recompute_num_layers is not None + else vision_cfg.recompute_num_layers + ) + if method is None: + raise ValueError("Full vision recompute requires a recompute method.") + if method not in {"uniform", "block"}: + raise ValueError("Full vision recompute method must be 'uniform' or 'block'.") + if num_layers is None: + raise ValueError("Full vision recompute requires a layer count.") + if ( + not isinstance(num_layers, int) + or isinstance(num_layers, bool) + or not 1 <= num_layers <= vision_cfg.num_layers + ): + raise ValueError( + f"Full vision recompute layer count must be an integer between 1 and {vision_cfg.num_layers}." + ) + vision_cfg.recompute_method = method + vision_cfg.recompute_num_layers = num_layers + else: + vision_cfg.recompute_granularity = None + vision_cfg.recompute_method = None + vision_cfg.recompute_num_layers = None if self.vision_class_token_len is not None: vision_cfg.class_token_len = self.vision_class_token_len return vision_cfg diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_provider.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_provider.py new file mode 100644 index 00000000000..698019c4dee --- /dev/null +++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_provider.py @@ -0,0 +1,177 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from megatron.bridge.models.nemotron_omni.nemotron_omni_provider import NemotronOmniModelProvider + + +pytestmark = pytest.mark.unit + + +def test_vision_recompute_is_disabled_by_default(): + provider = NemotronOmniModelProvider( + recompute_granularity="selective", + recompute_modules=["core_attn", "mlp"], + ) + + vision_config = provider._build_vision_config(provider) + + assert vision_config.recompute_granularity is None + assert vision_config.recompute_method is None + assert vision_config.recompute_num_layers is None + + +def test_vision_recompute_inherits_selective_model_config_when_enabled(): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + recompute_granularity="selective", + recompute_modules=["core_attn", "mlp"], + ) + + vision_config = provider._build_vision_config(provider) + + assert vision_config.recompute_granularity == "selective" + assert vision_config.recompute_modules == ["core_attn", "mlp"] + assert vision_config.recompute_num_layers is None + + +def test_vision_recompute_inherits_full_model_config_when_enabled(): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=4, + ) + + vision_config = provider._build_vision_config(provider) + + assert vision_config.recompute_granularity == "full" + assert vision_config.recompute_method == "uniform" + assert vision_config.recompute_num_layers == 4 + + +def test_vision_recompute_can_override_full_model_config_with_selective(): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=4, + vision_recompute_granularity="selective", + ) + + vision_config = provider._build_vision_config(provider) + + assert provider.recompute_granularity == "full" + assert vision_config.recompute_granularity == "selective" + assert vision_config.recompute_method is None + assert vision_config.recompute_num_layers is None + + +def test_vision_recompute_can_use_per_layer_full_recompute_independently(): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + recompute_granularity="selective", + recompute_modules=["core_attn", "mlp"], + vision_recompute_granularity="full", + vision_recompute_method="uniform", + vision_recompute_num_layers=1, + ) + + vision_config = provider._build_vision_config(provider) + + assert provider.recompute_granularity == "selective" + assert vision_config.recompute_granularity == "full" + assert vision_config.recompute_method == "uniform" + assert vision_config.recompute_num_layers == 1 + + +def test_vision_recompute_requires_effective_granularity(): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + ) + + with pytest.raises(ValueError, match="effective recompute granularity"): + provider._build_vision_config(provider) + + +@pytest.mark.parametrize( + ("method", "num_layers", "message"), + [ + (None, 1, "requires a recompute method"), + ("uniform", None, "requires a layer count"), + ], +) +def test_full_vision_recompute_requires_complete_policy(method, num_layers, message): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + recompute_granularity="selective", + vision_recompute_granularity="full", + vision_recompute_method=method, + vision_recompute_num_layers=num_layers, + ) + + with pytest.raises(ValueError, match=message): + provider._build_vision_config(provider) + + +def test_vision_recompute_rejects_invalid_granularity(): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + vision_recompute_granularity="typo", + ) + + with pytest.raises(ValueError, match="must be 'full' or 'selective'"): + provider._build_vision_config(provider) + + +def test_full_vision_recompute_rejects_invalid_method(): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + vision_recompute_granularity="full", + vision_recompute_method="typo", + vision_recompute_num_layers=1, + ) + + with pytest.raises(ValueError, match="must be 'uniform' or 'block'"): + provider._build_vision_config(provider) + + +@pytest.mark.parametrize("num_layers", [0, 33, 1.5, True]) +def test_full_vision_recompute_rejects_invalid_layer_count(num_layers): + provider = NemotronOmniModelProvider( + recompute_vision=True, + radio_force_eval_mode=False, + vision_recompute_granularity="full", + vision_recompute_method="uniform", + vision_recompute_num_layers=num_layers, + ) + + with pytest.raises(ValueError, match="integer between 1 and 32"): + provider._build_vision_config(provider) + + +def test_vision_recompute_rejects_forced_radio_eval_mode(): + provider = NemotronOmniModelProvider(recompute_vision=True) + + with pytest.raises(ValueError, match="radio_force_eval_mode=False"): + provider._build_vision_config(provider) From 9391fdf45b08593236131941224850bc3cf4e3b1 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 23:41:49 +0200 Subject: [PATCH 45/85] cp: `[model] fix: honor Gemma 4 image attention mode (5333)` into `r0.6.0` (#5350) Signed-off-by: Yu Yao Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../bridge/models/gemma_vl/modeling_gemma4_vl.py | 2 ++ .../models/gemma_vl/test_gemma4_vl_modeling.py | 12 +++++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/gemma_vl/modeling_gemma4_vl.py b/src/megatron/bridge/models/gemma_vl/modeling_gemma4_vl.py index 414f36c2e2b..c13fe53593d 100644 --- a/src/megatron/bridge/models/gemma_vl/modeling_gemma4_vl.py +++ b/src/megatron/bridge/models/gemma_vl/modeling_gemma4_vl.py @@ -371,6 +371,8 @@ def _compute_attention_mask(self, input_ids: torch.Tensor) -> Optional[torch.Ten causal_mask = torch.tril( torch.ones((batch_size, 1, seq_len, seq_len), dtype=torch.bool, device=input_ids.device) ) + if getattr(self.config.text_config, "use_bidirectional_attention", None) != "vision": + return ~causal_mask def _bidirectional_block_mask(token_mask: torch.Tensor) -> torch.Tensor: padded = F.pad(token_mask, (1, 0), value=0) diff --git a/tests/unit_tests/models/gemma_vl/test_gemma4_vl_modeling.py b/tests/unit_tests/models/gemma_vl/test_gemma4_vl_modeling.py index 4913348bf4d..cc663253ae6 100644 --- a/tests/unit_tests/models/gemma_vl/test_gemma4_vl_modeling.py +++ b/tests/unit_tests/models/gemma_vl/test_gemma4_vl_modeling.py @@ -36,12 +36,13 @@ # --------------------------------------------------------------------------- -def _make_model(image_token_id=IMAGE_TOKEN_ID): +def _make_model(image_token_id=IMAGE_TOKEN_ID, use_bidirectional_attention="vision"): """Build a Gemma4VLModel with all heavy dependencies mocked out.""" config = Mock() config.image_token_id = image_token_id config.vision_config = Mock() config.text_config = Mock() + config.text_config.use_bidirectional_attention = use_bidirectional_attention config.share_embeddings_and_output_weights = True config.sequence_parallel = False config._pg_collection = None @@ -109,6 +110,15 @@ def test_image_block_gets_bidirectional_attention(self): for j in range(2, 5): assert not mask[0, 0, i, j].item(), f"Image pos ({i},{j}) should be unblocked (bidirectional)" + def test_image_block_stays_causal_when_bidirectional_attention_is_disabled(self): + """Smaller Gemma 4 variants keep conventional causal attention for image tokens.""" + model = _make_model(use_bidirectional_attention=None) + input_ids = self._make_ids([self.IMAGE_TOKEN, self.IMAGE_TOKEN]) + + mask = model._compute_attention_mask(input_ids) + + assert mask[0, 0, 0, 1].item() is True + def test_text_after_image_cannot_attend_back_to_image_beyond_causal(self): """Text token after image block uses causal attention (cannot look into future).""" model = _make_model() From 95674527f11b903922cd11ea730974f67b0cdc5f Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 23:42:18 +0200 Subject: [PATCH 46/85] cp: `[ckpt] fix: support vocab padding with Megatron-FSDP TP (5353)` into `r0.6.0` (#5359) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- src/megatron/bridge/models/conversion/param_mapping.py | 6 +++++- .../test_groups/converter/test_hf_fsdp_conversion.py | 3 +-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py index fe8ee3a8557..2e626d97773 100644 --- a/src/megatron/bridge/models/conversion/param_mapping.py +++ b/src/megatron/bridge/models/conversion/param_mapping.py @@ -993,7 +993,11 @@ def hf_to_megatron( hf_weights = hf_weights.to(target_param.dtype) actual_dim0_size = hf_weights.shape[0] - expect_dim0_size = target_param.shape[0] * self.tp_size + # DTensor.shape is already the global shape across TP ranks, while + # a regular Megatron parameter stores only its local TP shard. + expect_dim0_size = target_param.shape[0] + if not isinstance(target_param, DTensor): + expect_dim0_size *= self.tp_size if actual_dim0_size != expect_dim0_size: assert self.megatron_param in {"embedding.word_embeddings.weight", "output_layer.weight"}, ( f"{hf_weights.shape=} {target_param.shape=} {self.tp_size=} {self.megatron_param=} {self.hf_param=}" diff --git a/tests/functional_tests/test_groups/converter/test_hf_fsdp_conversion.py b/tests/functional_tests/test_groups/converter/test_hf_fsdp_conversion.py index ae92cf319ff..dc4822c432b 100644 --- a/tests/functional_tests/test_groups/converter/test_hf_fsdp_conversion.py +++ b/tests/functional_tests/test_groups/converter/test_hf_fsdp_conversion.py @@ -101,9 +101,8 @@ def qwen3_moe_toy_model_path(self, tmp_path_factory): @pytest.mark.parametrize( "tp,ep,nproc,test_name", [ - # CI only covers TP=1 cases. TP>1 cases depend on Megatron-LM commit - # 8cbc45b6e (PR#3191, merged 2026-04-07). (1, 1, 1, "FSDP_base"), + (2, 1, 2, "FSDP_TP2"), ], ) def test_hf_fsdp_roundtrip(self, qwen3_moe_toy_model_path, tmp_path, tp, ep, nproc, test_name): From 29c30dd36cba018e8ea2672984b026215214303c Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 23:44:26 +0200 Subject: [PATCH 47/85] cp: `fix(ckpt): avoid roundtrip export barrier timeout (5377)` into `r0.6.0` (#5412) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../hf_megatron_roundtrip_multi_gpu.py | 126 +++++++++++++----- scripts/conversion/gpu_backend.py | 41 +++--- .../conversion/launcher/test_gpu_backend.py | 9 +- .../test_hf_megatron_roundtrip_multi_gpu.py | 98 ++++++++++++++ 4 files changed, 218 insertions(+), 56 deletions(-) diff --git a/examples/conversion/hf_megatron_roundtrip_multi_gpu.py b/examples/conversion/hf_megatron_roundtrip_multi_gpu.py index 23009f2f08e..09c445fc932 100644 --- a/examples/conversion/hf_megatron_roundtrip_multi_gpu.py +++ b/examples/conversion/hf_megatron_roundtrip_multi_gpu.py @@ -28,7 +28,9 @@ model back into the Hugging Face format. A new directory, named after the model, will be created for the converted model files. By default, this directory is created in the current working directory, but a different - parent directory can be specified via the `--output-dir` argument. + parent directory can be specified via the `--output-dir` argument. Multi-GPU + exports use distributed saving by default; pass `--no-distributed-save` to + select rank-0-only saving explicitly. 7. Optionally, the `save_megatron_model` method can be used to save the model in Megatron's native checkpoint format by specifying the `--megatron-save-path` argument. @@ -51,7 +53,6 @@ import torch from rich.console import Console -from rich.table import Table from megatron.bridge import AutoBridge from megatron.bridge.models.decorators import torchrun_main @@ -99,10 +100,59 @@ def _configure_slurm_distributed_environment() -> None: os.environ["MASTER_PORT"] = str(master_port) +def _synchronize_weight_verification(all_match: bool) -> bool: + """Return the rank-0 verification result consistently on every rank. + + Args: + all_match: Whether the locally verified weights match. Only rank 0's + value is authoritative. + + Returns: + Whether rank 0 observed any weight mismatch. + """ + mismatch = torch.tensor(not all_match, dtype=torch.int32, device=torch.cuda.current_device()) + torch.distributed.broadcast(mismatch, src=0) + return bool(mismatch.item()) + + +def _print_verification_results( + verified_count: int, + mismatch_count: int, + mismatch_samples: list[str], + fp8_skip_count: int, + fp8_skip_samples: list[str], +) -> None: + """Print a bounded rank-0-only verification summary. + + Args: + verified_count: Number of weights compared or skipped as lossy FP8. + mismatch_count: Number of weights that did not match. + mismatch_samples: Bounded sample of mismatched weight names. + fp8_skip_count: Number of lossy FP8 comparisons that were skipped. + fp8_skip_samples: Sample names and dtypes for skipped FP8 comparisons. + """ + compared_count = verified_count - fp8_skip_count + matched_count = compared_count - mismatch_count + color = "green" if mismatch_count == 0 else "red" + console.print(f"[{color}]Weight verification: {matched_count}/{compared_count} compared weights matched[/]") + if mismatch_count: + console.print(f"[red]{mismatch_count} weight mismatches (showing up to 20):[/red]") + for entry in mismatch_samples: + console.print(f" [red]{entry}[/red]") + if fp8_skip_count > 0: + console.print( + f"[yellow]WARNING: {fp8_skip_count} FP8 params skipped allclose (dequantisation is lossy):[/yellow]" + ) + for entry in fp8_skip_samples: + console.print(f" [yellow]{entry}[/yellow]") + if fp8_skip_count > len(fp8_skip_samples): + console.print(f" [yellow]... and {fp8_skip_count - len(fp8_skip_samples)} more[/yellow]") + + @torchrun_main def main( hf_model_id: str = HF_MODEL_ID, - output_dir: str = None, + output_dir: str | None = None, tp: int = 1, pp: int = 1, ep: int = 1, @@ -111,6 +161,7 @@ def main( megatron_load_path: str | None = None, trust_remote_code: bool | None = None, strict: bool = False, + distributed_save: bool = True, skip_save: bool = False, atol: float = 1e-1, rtol: float = 1e-5, @@ -180,15 +231,6 @@ def main( # Now we can check for rank is_rank_0 = torch.distributed.get_rank() == 0 - # Formatting - if is_rank_0: - table = Table(title="Hugging Face Weights Verification") - table.add_column("Weight Name", style="cyan") - table.add_column("Shape") - table.add_column("DType") - table.add_column("Device") - table.add_column("Matches Original", justify="center") - if is_rank_0: console.print(f"[yellow]Tensor parallel size: {model_provider.tensor_model_parallel_size}[/yellow]") console.print(f"[yellow]Pipeline parallel size: {model_provider.pipeline_model_parallel_size}[/yellow]") @@ -210,6 +252,9 @@ def main( # 3. Regular params (same dtype, not in ignore list) → direct allclose. all_match = True + verified_count = 0 + mismatch_count = 0 + mismatch_samples: list[str] = [] fp8_skip_count = 0 fp8_skip_samples: list[str] = [] for name, param in bridge.export_hf_weights(megatron_model, show_progress=False): @@ -248,26 +293,22 @@ def main( ) all_match = all_match and match - table.add_row( - name, - str(tuple(param.shape)), - str(param.dtype).replace("torch.", ""), - str(param.device), - "✅" if match else "❌", - ) - - if is_rank_0: - if fp8_skip_count > 0: - console.print( - f"[yellow]WARNING: {fp8_skip_count} FP8 params skipped allclose (dequantisation is lossy):[/yellow]" + verified_count += 1 + if not match: + mismatch_count += 1 + if len(mismatch_samples) < 20: + mismatch_samples.append(name) + + verification_failed = _synchronize_weight_verification(all_match) + if verification_failed: + if is_rank_0: + _print_verification_results( + verified_count, + mismatch_count, + mismatch_samples, + fp8_skip_count, + fp8_skip_samples, ) - for entry in fp8_skip_samples: - console.print(f" [yellow]{entry}[/yellow]") - if fp8_skip_count > len(fp8_skip_samples): - console.print(f" [yellow]... and {fp8_skip_count - len(fp8_skip_samples)} more[/yellow]") - console.print(table) - - if not all_match: raise ValueError("Weight mismatch detected") if skip_save: @@ -276,7 +317,12 @@ def main( else: if is_rank_0: console.print(f"Saving HF-ckpt in {save_path}...") - bridge.save_hf_pretrained(megatron_model, save_path, strict=strict) + bridge.save_hf_pretrained( + megatron_model, + save_path, + strict=strict, + distributed_save=distributed_save, + ) # Save in Megatron format if path is provided if megatron_save_path: @@ -284,6 +330,15 @@ def main( console.print(f"Saving Megatron checkpoint in {megatron_save_path}...") bridge.save_megatron_model(megatron_model, megatron_save_path) + if is_rank_0: + _print_verification_results( + verified_count, + mismatch_count, + mismatch_samples, + fp8_skip_count, + fp8_skip_samples, + ) + def _build_parser() -> argparse.ArgumentParser: """Build the standalone round-trip example argument parser.""" @@ -323,6 +378,12 @@ def _build_parser() -> argparse.ArgumentParser: ) strictness.add_argument("--not-strict", dest="strict", action="store_false", help=argparse.SUPPRESS) parser.set_defaults(strict=False) + parser.add_argument( + "--distributed-save", + action=argparse.BooleanOptionalAction, + default=True, + help="Let ranks save assigned Hugging Face shards independently (default: enabled)", + ) parser.add_argument( "--skip-save", action="store_true", help="Skip saving the model after comparison (verification only)" ) @@ -344,6 +405,7 @@ def _build_parser() -> argparse.ArgumentParser: args.megatron_load_path, args.trust_remote_code, strict=args.strict, + distributed_save=args.distributed_save, skip_save=args.skip_save, atol=args.atol, rtol=args.rtol, diff --git a/scripts/conversion/gpu_backend.py b/scripts/conversion/gpu_backend.py index 59dce2713f7..b650e95bbd1 100644 --- a/scripts/conversion/gpu_backend.py +++ b/scripts/conversion/gpu_backend.py @@ -21,7 +21,6 @@ import torch import yaml from rich.console import Console -from rich.table import Table from utils import parse_dtype, prepare_output_directory, validate_output_path from megatron.bridge import AutoBridge @@ -207,21 +206,15 @@ def _verify_roundtrip_weights(bridge: AutoBridge, megatron_model: list[torch.nn. """Verify exported Megatron weights against the original Hugging Face state.""" is_rank_0 = torch.distributed.get_rank() == 0 all_match = True + verified_count = 0 + mismatch_count = 0 + mismatch_samples: list[str] = [] fp8_skip_count = 0 fp8_skip_samples: list[str] = [] - table = None - if is_rank_0: - table = Table(title="Hugging Face Weights Verification") - table.add_column("Weight Name", style="cyan") - table.add_column("Shape") - table.add_column("DType") - table.add_column("Device") - table.add_column("Matches Original", justify="center") for name, exported in bridge.export_hf_weights(megatron_model, show_progress=False): if not is_rank_0: continue - assert table is not None original = bridge.hf_pretrained.state[name] match, skipped_fp8 = _roundtrip_weights_match(name, exported, original) if skipped_fp8: @@ -229,16 +222,24 @@ def _verify_roundtrip_weights(bridge: AutoBridge, megatron_model: list[torch.nn. if len(fp8_skip_samples) < 20: fp8_skip_samples.append(f"{name}: exported {exported.dtype} vs original {original.dtype}") all_match = all_match and match - table.add_row( - name, - str(tuple(exported.shape)), - str(exported.dtype).replace("torch.", ""), - str(exported.device), - "✅" if match else "❌", - ) + verified_count += 1 + if not match: + mismatch_count += 1 + if len(mismatch_samples) < 20: + mismatch_samples.append(name) + + mismatch = torch.tensor(not all_match, dtype=torch.int32, device=torch.cuda.current_device()) + torch.distributed.broadcast(mismatch, src=0) if is_rank_0: - assert table is not None + compared_count = verified_count - fp8_skip_count + matched_count = compared_count - mismatch_count + color = "green" if mismatch_count == 0 else "red" + _CONSOLE.print(f"[{color}]Weight verification: {matched_count}/{compared_count} compared weights matched[/]") + if mismatch_count: + _CONSOLE.print(f"[red]{mismatch_count} weight mismatches (showing up to 20):[/red]") + for entry in mismatch_samples: + _CONSOLE.print(f" [red]{entry}[/red]") if fp8_skip_count: _CONSOLE.print( f"[yellow]WARNING: {fp8_skip_count} FP8 params skipped allclose (dequantisation is lossy):[/yellow]" @@ -247,10 +248,6 @@ def _verify_roundtrip_weights(bridge: AutoBridge, megatron_model: list[torch.nn. _CONSOLE.print(f" [yellow]{entry}[/yellow]") if fp8_skip_count > len(fp8_skip_samples): _CONSOLE.print(f" [yellow]... and {fp8_skip_count - len(fp8_skip_samples)} more[/yellow]") - _CONSOLE.print(table) - - mismatch = torch.tensor(not all_match, dtype=torch.int32, device=torch.cuda.current_device()) - torch.distributed.broadcast(mismatch, src=0) if mismatch.item(): raise ValueError("Weight mismatch detected") diff --git a/tests/unit_tests/conversion/launcher/test_gpu_backend.py b/tests/unit_tests/conversion/launcher/test_gpu_backend.py index 776bf8903e1..632a591e909 100644 --- a/tests/unit_tests/conversion/launcher/test_gpu_backend.py +++ b/tests/unit_tests/conversion/launcher/test_gpu_backend.py @@ -417,6 +417,8 @@ def test_weight_comparison_casts_known_precision_parameters(self, cli): assert skipped_fp8 is False def test_verification_raises_on_weight_mismatch(self, cli, monkeypatch): + events = [] + class FakeBridge: hf_pretrained = types.SimpleNamespace(state={"weight": torch.tensor([2.0])}) @@ -424,9 +426,12 @@ def export_hf_weights(self, model, *, show_progress): yield "weight", torch.tensor([1.0]) monkeypatch.setattr(cli.torch.distributed, "get_rank", lambda: 0) - monkeypatch.setattr(cli.torch.distributed, "broadcast", lambda tensor, src: None) + monkeypatch.setattr(cli.torch.distributed, "broadcast", lambda tensor, src: events.append("broadcast")) monkeypatch.setattr(cli.torch.cuda, "current_device", lambda: "cpu") - monkeypatch.setattr(cli._CONSOLE, "print", lambda *args, **kwargs: None) + monkeypatch.setattr(cli._CONSOLE, "print", lambda *args, **kwargs: events.append("report")) with pytest.raises(ValueError, match="Weight mismatch detected"): cli._verify_roundtrip_weights(FakeBridge(), ["megatron-model"]) + + assert events[0] == "broadcast" + assert "report" in events diff --git a/tests/unit_tests/conversion/launcher/test_hf_megatron_roundtrip_multi_gpu.py b/tests/unit_tests/conversion/launcher/test_hf_megatron_roundtrip_multi_gpu.py index 6dfb9896b1e..0cacc81c667 100644 --- a/tests/unit_tests/conversion/launcher/test_hf_megatron_roundtrip_multi_gpu.py +++ b/tests/unit_tests/conversion/launcher/test_hf_megatron_roundtrip_multi_gpu.py @@ -16,8 +16,10 @@ import sys import types from pathlib import Path +from types import SimpleNamespace import pytest +import torch pytestmark = pytest.mark.unit @@ -95,3 +97,99 @@ def test_standalone_roundtrip_strict_export_is_opt_in(monkeypatch: pytest.Monkey with pytest.raises(SystemExit): module._build_parser().parse_args(["--strict", "--not-strict"]) + + +def test_standalone_roundtrip_defaults_to_distributed_export(monkeypatch: pytest.MonkeyPatch): + module = _load_roundtrip_module(monkeypatch) + + default_args = module._build_parser().parse_args([]) + serial_args = module._build_parser().parse_args(["--no-distributed-save"]) + + assert default_args.distributed_save is True + assert serial_args.distributed_save is False + + +def test_weight_verification_broadcasts_rank_zero_mismatch(monkeypatch: pytest.MonkeyPatch): + module = _load_roundtrip_module(monkeypatch) + mismatch = SimpleNamespace(value=0, item=lambda: mismatch.value) + tensor_calls = [] + + def fake_tensor(value, *, dtype, device): + tensor_calls.append((value, dtype, device)) + return mismatch + + def fake_broadcast(value, *, src): + assert value is mismatch + assert src == 0 + mismatch.value = 1 + + monkeypatch.setattr(module.torch, "tensor", fake_tensor) + monkeypatch.setattr(module.torch.cuda, "current_device", lambda: 3) + monkeypatch.setattr(module.torch.distributed, "broadcast", fake_broadcast) + + assert module._synchronize_weight_verification(all_match=True) is True + assert tensor_calls == [(False, torch.int32, 3)] + + +def test_roundtrip_synchronizes_result_and_saves_before_reporting_summary( + monkeypatch: pytest.MonkeyPatch, +): + """Ranks must synchronize and save before rank 0 reports verification.""" + module = _load_roundtrip_module(monkeypatch) + events = [] + + class FakeProvider: + tensor_model_parallel_size = 1 + pipeline_model_parallel_size = 1 + expert_model_parallel_size = 1 + expert_tensor_parallel_size = 1 + + def finalize(self): + return None + + def initialize_model_parallel(self, seed): + return None + + def provide_distributed_model(self, wrap_with_ddp): + return [object()] + + class FakeBridge: + hf_pretrained = SimpleNamespace(state={"weight": torch.ones(1)}) + + def to_megatron_provider(self, load_weights): + return FakeProvider() + + def export_hf_weights(self, model, show_progress): + yield "weight", torch.ones(1) + + def save_hf_pretrained(self, model, path, **kwargs): + events.append(("save", kwargs)) + + fake_bridge = FakeBridge() + monkeypatch.setattr( + module, + "AutoBridge", + SimpleNamespace(from_hf_pretrained=lambda *args, **kwargs: fake_bridge), + ) + monkeypatch.setattr(module.torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr( + module, + "_synchronize_weight_verification", + lambda all_match: events.append(("synchronize", all_match)), + raising=False, + ) + monkeypatch.setattr( + module, + "_print_verification_results", + lambda *args, **kwargs: events.append(("report", args)), + ) + monkeypatch.setenv("WORLD_SIZE", "2") + + module.main(hf_model_id="example/model", distributed_save=True) + + event_names = [event[0] for event in events] + assert "synchronize" in event_names + assert event_names.index("synchronize") < event_names.index("save") + assert event_names.index("save") < event_names.index("report") + save_kwargs = next(event[1] for event in events if event[0] == "save") + assert save_kwargs["distributed_save"] is True From 93715082e6696f8574389de2cae60975783145b0 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 7 Aug 2026 23:47:25 +0200 Subject: [PATCH 48/85] cp: `fix(peft): preserve grouped expert SwiGLU LoRA order (5376)` into `r0.6.0` (#5414) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- src/megatron/bridge/peft/utils.py | 12 --------- tests/unit_tests/peft/test_utils.py | 39 +++++++++++++++++++++-------- 2 files changed, 28 insertions(+), 23 deletions(-) diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py index dbed296c82b..f6f8d30ec8a 100644 --- a/src/megatron/bridge/peft/utils.py +++ b/src/megatron/bridge/peft/utils.py @@ -1621,18 +1621,6 @@ def sh_ten_build_fn(key: str, tensor: torch.Tensor, replica_id, flattened_range) ] def sh_ten_merge_fn(sub_state_dict): - if not singleton_local_shards and len(sub_state_dict) > 1: - # Dist checkpoint load reconstructs one local fused shard per expert-TP - # rank, so the incoming tensors look like [gate_0|up_0, gate_1|up_1, ...]. - # Restore the fused [gate_0, gate_1, ..., up_0, up_1, ...] layout before - # concatenating back along the SwiGLU axis. - gate_parts = [] - up_parts = [] - for tensor in sub_state_dict: - gate_part, up_part = torch.chunk(tensor, 2, dim=swiglu_shard_axis) - gate_parts.append(gate_part) - up_parts.append(up_part) - sub_state_dict = [*gate_parts, *up_parts] try: return torch.cat(sub_state_dict, dim=swiglu_shard_axis) except (RuntimeError, torch.cuda.OutOfMemoryError) as exc: diff --git a/tests/unit_tests/peft/test_utils.py b/tests/unit_tests/peft/test_utils.py index c2be2839065..fccce8118c0 100644 --- a/tests/unit_tests/peft/test_utils.py +++ b/tests/unit_tests/peft/test_utils.py @@ -2199,32 +2199,49 @@ def test_grouped_expert_linear_fc1_sharded_state_dict_preserves_expert_axis(self assert built[0].global_offset == (2, 0, 0) assert built[1].global_offset == (2, 2, 0) - def test_grouped_expert_linear_fc1_factory_merge_restores_gate_up_order(self): - """Grouped expert fc1 checkpoint reload should de-interleave gate/up expert-TP shards.""" + @pytest.mark.parametrize("etp_size", [1, 2]) + def test_grouped_expert_linear_fc1_factory_merge_preserves_gate_up_and_expert_order(self, etp_size): + """Grouped expert fc1 checkpoint reload should preserve exact adapter weight ordering.""" config = MockModelParallelConfig() config.gated_linear_unit = True - config._pg_collection = make_mock_pg_collection(ep_size=1, ep_rank=0, edp_rank=0, etp_size=2, etp_rank=0) + config._pg_collection = make_mock_pg_collection( + ep_size=1, + ep_rank=0, + edp_rank=0, + etp_size=etp_size, + etp_rank=0, + ) adapter = GroupedExpertLinearAdapter( in_features=2, out_features=8, dim=2, - num_local_experts=1, + num_local_experts=2, base_linear_name="decoder.layers.0.mlp.experts.linear_fc1", activation="identity", input_is_parallel=False, model_parallel_config=config, ) + expert = torch.arange(2).reshape(2, 1, 1) * 1000 + local_rows = adapter.linear_out.weight.shape[1] + gate_rows = local_rows // 2 + projection = torch.tensor([0] * gate_rows + [1] * gate_rows).reshape(1, local_rows, 1) * 100 + row = torch.arange(gate_rows).repeat(2).reshape(1, local_rows, 1) * 10 + column = torch.arange(2).reshape(1, 1, 2) + expected = (expert + projection + row + column).to(adapter.linear_out.weight) + with torch.no_grad(): + adapter.linear_out.weight.copy_(expected) + factory = adapter.sharded_state_dict("adapter.")["adapter.linear_out.weight"] + built = factory.build() - fused_tp0 = torch.tensor([[[1.0, 1.0], [1.0, 1.0], [2.0, 2.0], [2.0, 2.0]]]) - fused_tp1 = torch.tensor([[[3.0, 3.0], [3.0, 3.0], [4.0, 4.0], [4.0, 4.0]]]) + assert len(built) == 2 + torch.testing.assert_close(built[0].data, expected[:, :gate_rows], rtol=0, atol=0) + torch.testing.assert_close(built[1].data, expected[:, gate_rows:], rtol=0, atol=0) - merged = factory.merge_fn([fused_tp0, fused_tp1]) - expected = torch.tensor( - [[[1.0, 1.0], [1.0, 1.0], [3.0, 3.0], [3.0, 3.0], [2.0, 2.0], [2.0, 2.0], [4.0, 4.0], [4.0, 4.0]]] - ) - torch.testing.assert_close(merged, expected) + merged = factory.merge_fn([shard.data for shard in built]) + + torch.testing.assert_close(merged, expected, rtol=0, atol=0) @pytest.mark.parametrize( ("ep_size", "tp_size", "etp_size", "expected_allreduce"), From 8003b1174bb4443a59bf520a58c6441bab61a6c6 Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Sat, 8 Aug 2026 09:43:14 -0700 Subject: [PATCH 49/85] cp: `fix(training): include MTP stages in embedding groups (5380)` into `r0.6.0` (#5422) Signed-off-by: yaoyu-33 --- src/megatron/bridge/training/setup.py | 42 ++++++++++++ tests/unit_tests/training/test_setup.py | 86 +++++++++++++++++++++++++ 2 files changed, 128 insertions(+) diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index 30efc5da646..83a9e7f5e2d 100644 --- a/src/megatron/bridge/training/setup.py +++ b/src/megatron/bridge/training/setup.py @@ -28,11 +28,14 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.rerun_state_machine import RerunDataIterator from megatron.core.transformer import MegatronModule +from megatron.core.transformer.multi_token_prediction import get_mtp_ranks from megatron.training.models.base import ModelConfig from megatron.bridge.data.loaders import build_train_valid_test_datasets_for_num_epochs, setup_data_iterators from megatron.bridge.models.gpt.gpt_builder import GPTModelConfig +from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hybrid.hybrid_builder import HybridModelConfig +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.model_provider import ModelProviderMixin from megatron.bridge.models.transformer_config import TransformerConfig from megatron.bridge.training import fault_tolerance @@ -64,6 +67,43 @@ from megatron.bridge.utils.common_utils import get_rank_safe, print_rank_0 +def _get_embedding_ranks( + pp_ranks: list[int], + pipeline_model_parallel_size: int | None = None, + *, + model_config: GPTModelConfig | GPTModelProvider | HybridModelConfig | HybridModelProvider, +) -> list[int]: + """Get the embedding ranks for a Bridge language-model config.""" + # HyperCommGrid passes PP size as a second argument; MCore's MPU path does not. + del pipeline_model_parallel_size + + # Keep this rank construction aligned with pretrain_gpt.get_embedding_ranks in MCore. + embedding_ranks = [pp_ranks[0]] + if len(pp_ranks) > 1: + if model_config.share_embeddings_and_output_weights: + embedding_ranks.append(pp_ranks[-1]) + transformer_config = model_config.transformer if hasattr(model_config, "transformer") else model_config + mtp_ranks = get_mtp_ranks(pp_ranks, transformer_config) + embedding_ranks.extend(mtp_ranks) + embedding_ranks = list(set(embedding_ranks)) + embedding_ranks = sorted(embedding_ranks) + return embedding_ranks + + +def _resolve_embedding_ranks_fn( + model_config: object, + get_embedding_ranks: Callable[[list[int], Optional[int]], list[int]] | None, +) -> Callable[[list[int], Optional[int]], list[int]] | None: + """Use model-aware language-model embedding ranks unless the caller supplied an override.""" + if get_embedding_ranks is not None: + return get_embedding_ranks + + language_model_configs = (GPTModelConfig, GPTModelProvider, HybridModelConfig, HybridModelProvider) + if isinstance(model_config, language_model_configs): + return partial(_get_embedding_ranks, model_config=model_config) + return None + + class SetupOutput(NamedTuple): """Represents the output of the main setup function. @@ -200,6 +240,8 @@ def setup( set_level_for_all_loggers=cfg.logger.set_level_for_all_loggers, ) + get_embedding_ranks = _resolve_embedding_ranks_fn(cfg.model, get_embedding_ranks) + # pg_collection is returned from initialize_megatron: # - When use_decentralized_pg=True: uses HyperCommGrid to create local process groups # - When use_decentralized_pg=False: uses mpu's global parallel state diff --git a/tests/unit_tests/training/test_setup.py b/tests/unit_tests/training/test_setup.py index ad917295215..64612e09c7a 100644 --- a/tests/unit_tests/training/test_setup.py +++ b/tests/unit_tests/training/test_setup.py @@ -18,15 +18,19 @@ import pytest +import megatron.bridge.training.setup as training_setup from megatron.bridge.data.builders import GPTSFTDatasetConfig from megatron.bridge.models.gpt.gpt_builder import GPTModelConfig from megatron.bridge.models.gpt_provider import GPTModelProvider +from megatron.bridge.models.hybrid.hybrid_builder import HybridModelConfig +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.transformer_config import TransformerConfig from megatron.bridge.training.checkpointing import load_checkpoint from megatron.bridge.training.setup import ( _bind_dataset_provider_context, _build_distributed_model, _register_pre_wrap_hook, + _resolve_embedding_ranks_fn, _should_load_checkpoint, _update_model_config_funcs, _validate_and_set_vocab_size, @@ -48,6 +52,30 @@ def _make_gpt_model_config(**kwargs): return GPTModelConfig(**defaults) +def _make_mtp_model_config(model_config_cls, *, share_embeddings_and_output_weights=False): + transformer_kwargs = { + "hidden_size": 128, + "mtp_num_layers": 1, + "num_attention_heads": 1, + "num_layers": 1, + "pipeline_model_parallel_layout": "Et|m|L", + "pipeline_model_parallel_size": 3, + } + if model_config_cls in (GPTModelProvider, HybridModelProvider): + model_config = model_config_cls( + share_embeddings_and_output_weights=share_embeddings_and_output_weights, + **transformer_kwargs, + ) + else: + model_config = model_config_cls( + share_embeddings_and_output_weights=share_embeddings_and_output_weights, + transformer=_make_transformer(**transformer_kwargs), + vocab_size=128, + ) + model_config.finalize() + return model_config + + def _make_checkpoint_source_config( *, load, @@ -70,6 +98,64 @@ def _make_checkpoint_source_config( ) +@pytest.mark.parametrize( + ("share_embeddings_and_output_weights", "expected_ranks"), + [(False, [0, 1]), (True, [0, 1, 2])], +) +@pytest.mark.parametrize( + "model_config_cls", + [GPTModelConfig, GPTModelProvider, HybridModelConfig, HybridModelProvider], +) +def test_mtp_embedding_group_includes_standalone_mtp_stage( + model_config_cls, + share_embeddings_and_output_weights, + expected_ranks, +): + model_config = _make_mtp_model_config( + model_config_cls, + share_embeddings_and_output_weights=share_embeddings_and_output_weights, + ) + + get_embedding_ranks = _resolve_embedding_ranks_fn(model_config, None) + explicit_callback = Mock() + + assert get_embedding_ranks is not None + assert get_embedding_ranks([0, 1, 2]) == expected_ranks + assert _resolve_embedding_ranks_fn(model_config, explicit_callback) is explicit_callback + + +def test_setup_forwards_model_aware_embedding_ranks_to_parallel_initialization(): + model_config = _make_mtp_model_config(GPTModelProvider) + cfg = SimpleNamespace( + dist=SimpleNamespace(disable_jit_fuser=False, enable_megatron_core_experimental=False), + logger=SimpleNamespace( + filter_warnings=False, + logging_level="INFO", + modules_to_filter=[], + set_level_for_all_loggers=False, + ), + model=model_config, + ) + state = SimpleNamespace(cfg=cfg, initialize_async_checkpoint_worker=Mock()) + initialize_megatron = Mock(side_effect=RuntimeError("stop after process-group initialization")) + + with ( + patch.multiple( + training_setup, + initialize_megatron=initialize_megatron, + maybe_log_and_save_config=Mock(), + set_experimental_flag=Mock(), + setup_logging=Mock(), + ), + pytest.raises(RuntimeError, match="stop after process-group initialization"), + ): + training_setup.setup(state, Mock()) + + get_embedding_ranks = initialize_megatron.call_args.kwargs["get_embedding_ranks"] + assert get_embedding_ranks([0, 1, 2]) == [0, 1] + assert get_embedding_ranks([0, 1, 2], 3) == [0, 1] + + def test_gpt_sft_config_receives_tokenizer_through_builder_binding_without_mutation(): config = GPTSFTDatasetConfig(seq_length=128, dataset_root="/tmp/data") tokenizer = object() From 41d813549d3be2c0a4008087e245be1b3c997139 Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Sat, 8 Aug 2026 09:43:49 -0700 Subject: [PATCH 50/85] cp: `fix: repair release QA regressions (5351)` into `r0.6.0` (#5423) Signed-off-by: yaoyu-33 Signed-off-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- docs/index.md | 3 +++ docs/models/README.md | 6 ++++- docs/models/ernie/ernie45.md | 20 ++++++++++++++ docs/models/exaone/exaone.md | 18 +++++++++++++ docs/models/hy_v3/hy-v3.md | 27 +++++++++++++++++++ .../mfsdp/convert_checkpoints_fsdp.py | 4 +-- examples/models/exaone/exaone45/README.md | 18 +++++++------ examples/models/exaone/exaone45/conversion.sh | 5 ++-- examples/models/gemma/gemma4_vl/README.md | 22 +++------------ .../bridge/recipes/exaone/h100/exaone45.py | 1 + .../bridge/recipes/glm_vl/h100/glm_45v.py | 2 ++ 11 files changed, 95 insertions(+), 31 deletions(-) create mode 100644 docs/models/ernie/ernie45.md create mode 100644 docs/models/exaone/exaone.md create mode 100644 docs/models/hy_v3/hy-v3.md diff --git a/docs/index.md b/docs/index.md index 227196831ad..2b17a43a049 100644 --- a/docs/index.md +++ b/docs/index.md @@ -36,10 +36,13 @@ bridge-tech-details.md models/README.md models/bailing/index.md models/deepseek/index.md +models/ernie/ernie45.md +models/exaone/exaone.md models/falcon/index.md models/gemma/index.md models/glm/index.md models/gpt_oss/index.md +models/hy_v3/hy-v3.md models/kimi/index.md models/llama/index.md models/minimax/index.md diff --git a/docs/models/README.md b/docs/models/README.md index 0d15a9ca220..a666f9583c2 100644 --- a/docs/models/README.md +++ b/docs/models/README.md @@ -17,10 +17,13 @@ Megatron Bridge conversion, training recipe links, and model-specific notes. |----------------|---------------------| | **Bailing** | [Ling 2.0](bailing/ling-2.md) | | **DeepSeek** | [DeepSeek V2 (deprecated)](deepseek/deepseek-v2.md), [DeepSeek V3](deepseek/deepseek-v3.md), [DeepSeek V4](deepseek/deepseek-v4.md) | +| **ERNIE** | [ERNIE 4.5 and ERNIE 4.5 VL](ernie/ernie45.md) | +| **EXAONE** | [EXAONE 4, EXAONE 4.5 VL, and K-EXAONE](exaone/exaone.md) | | **Falcon** | [Falcon H1](falcon/falcon-h1.md) | | **Gemma** | [Gemma (deprecated)](gemma/gemma.md), [Gemma 2 (deprecated)](gemma/gemma2.md), [Gemma 3](gemma/gemma3.md), [Gemma 3 VL](gemma/gemma3-vl.md), [Gemma 4 VL](gemma/gemma4-vl.md) | | **GLM** | [GLM-4.5](glm/glm45.md), [GLM-4.5V](glm/glm-45v.md), [GLM-4.7 / GLM-4.7-Flash](glm/glm47.md), [GLM-5 / GLM-5.1 / GLM-5.2](glm/glm5.md) | | **GPT-OSS** | [GPT OSS](gpt_oss/gpt-oss.md) | +| **HY V3** | [HY V3](hy_v3/hy-v3.md) | | **Kimi** | [Kimi K2](kimi/kimi-k2.md), [Kimi-K2.5-VL](kimi/kimi-k25-vl.md) | | **Llama** | [Llama 2 (deprecated)](llama/llama2.md), [Llama 3](llama/llama3.md) | | **MiniMax** | [MiniMax-M2 / M2.5 / M2.7](minimax/minimax-m2.md), [MiniMax-M3](minimax/minimax-m3.md) | @@ -65,12 +68,13 @@ Each model documentation page typically includes: ### Decoder-Only and Hybrid Backbones -- Bailing, DeepSeek, Falcon, Gemma, GLM, GPT-OSS, Kimi, Llama, MiniMax, Mistral, Moonlight, Nemotron, OLMoE, Qwen, Sarvam, StepFun, and Xiaomi-MiMo +- Bailing, DeepSeek, ERNIE, EXAONE, Falcon, Gemma, GLM, GPT-OSS, HY V3, Kimi, Llama, MiniMax, Mistral, Moonlight, Nemotron, OLMoE, Qwen, Sarvam, StepFun, and Xiaomi-MiMo - MoE and hybrid variants including Bailing, DeepSeek, GLM, GPT-OSS, MiniMax, Nemotron-3, OLMoE, Qwen3-MoE, Qwen3-Next, and Sarvam ### Multimodal Variants - Gemma 3 VL and Gemma 4 VL +- ERNIE 4.5 VL and EXAONE 4.5 VL - GLM-4.5V - Kimi-K2.5-VL - Ministral 3 diff --git a/docs/models/ernie/ernie45.md b/docs/models/ernie/ernie45.md new file mode 100644 index 00000000000..bcd0ade881d --- /dev/null +++ b/docs/models/ernie/ernie45.md @@ -0,0 +1,20 @@ +# ERNIE 4.5 + +Megatron Bridge registers Hugging Face conversion support for the text-only +`baidu/ERNIE-4.5-0.3B-PT` MoE model and the +`baidu/ERNIE-4.5-VL-28B-A3B-Instruct` and `baidu/ERNIE-4.5-VL-28B-A3B-Thinking` +vision-language MoE models. + +Import a text checkpoint with the shared conversion CLI: + +```bash +./scripts/conversion/convert.sh import \ + --hf-model baidu/ERNIE-4.5-0.3B-PT \ + --megatron-path /workspace/models/ernie-4.5-0.3b \ + --trust-remote-code +``` + +For ERNIE 4.5 VL import, export, and inference commands, see the +[checked-in examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/vlm/ernie_vl). +The VL bridge supports its modality-specific dual-pool MoE layout; use the +parallelism settings documented with those examples. diff --git a/docs/models/exaone/exaone.md b/docs/models/exaone/exaone.md new file mode 100644 index 00000000000..a973fb140e1 --- /dev/null +++ b/docs/models/exaone/exaone.md @@ -0,0 +1,18 @@ +# EXAONE + +Megatron Bridge supports Hugging Face conversion for EXAONE 4.0 dense language +models, EXAONE 4.5 vision-language models, and K-EXAONE MoE language models. +Checked-in H100 recipes cover EXAONE 4.0 1.2B, EXAONE 4.5 VL 33B, +K-EXAONE-236B-A23B, and K-EXAONE 2.0 750B-A37B. + +Use the model-specific examples for commands and supported parallelism: + +- [EXAONE 4.0](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/exaone/exaone4) +- [EXAONE 4.5 VL](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/exaone/exaone45) +- [K-EXAONE MoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/exaone/exaone_moe) + +The EXAONE 4.0 and EXAONE 4.5 examples use the shared +`scripts/conversion/convert.sh import` entry point. The K-EXAONE examples use +their model-specific distributed round-trip workflows. K-EXAONE 2.0 also has a checked-in +[model verification card](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/model_verification_cards/k-exaone-2/card.yaml) +that records the current verification status of individual workflows. diff --git a/docs/models/hy_v3/hy-v3.md b/docs/models/hy_v3/hy-v3.md new file mode 100644 index 00000000000..aa4e9b166b6 --- /dev/null +++ b/docs/models/hy_v3/hy-v3.md @@ -0,0 +1,27 @@ +# HY V3 + +Megatron Bridge registers Hugging Face checkpoint conversion for the +`tencent/Hy3-preview-Base` MoE causal language model. The bridge maps HY V3's +GQA with QK layer normalization, dense-first layer schedule, sigmoid router, +shared experts, and grouped-GEMM routed experts to Megatron Core. + +Import the checkpoint with the shared conversion CLI: + +```bash +./scripts/conversion/convert.sh import \ + --hf-model tencent/Hy3-preview-Base \ + --megatron-path /workspace/models/hy3-preview-base \ + --trust-remote-code +``` + +Export a converted checkpoint back to Hugging Face format: + +```bash +./scripts/conversion/convert.sh export \ + --hf-model tencent/Hy3-preview-Base \ + --megatron-path /workspace/models/hy3-preview-base/iter_0000000 \ + --hf-path /workspace/models/hy3-preview-base-hf \ + --trust-remote-code +``` + +No library training recipe is currently exported for HY V3. diff --git a/examples/conversion/mfsdp/convert_checkpoints_fsdp.py b/examples/conversion/mfsdp/convert_checkpoints_fsdp.py index a86726d374b..0c4fd62b301 100644 --- a/examples/conversion/mfsdp/convert_checkpoints_fsdp.py +++ b/examples/conversion/mfsdp/convert_checkpoints_fsdp.py @@ -269,9 +269,9 @@ def _add_common_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--trust-remote-code", action="store_true", help="Allow custom model code execution") parser.add_argument( "--ckpt-format", - choices=["torch_dist", "fsdp_dtensor"], + choices=["fsdp_dtensor"], default="fsdp_dtensor", - help="Megatron checkpoint format to load/save (default: fsdp_dtensor for Megatron-FSDP)", + help="Megatron-FSDP checkpoint format (only fsdp_dtensor is supported)", ) diff --git a/examples/models/exaone/exaone45/README.md b/examples/models/exaone/exaone45/README.md index 44455bdf597..3ec1a5dc4a4 100644 --- a/examples/models/exaone/exaone45/README.md +++ b/examples/models/exaone/exaone45/README.md @@ -8,7 +8,7 @@ Scripts for [EXAONE-4.5-33B](https://huggingface.co/LGAI-EXAONE/EXAONE-4.5-33B), | Architecture | Vision-language model (`Exaone4_5_ForConditionalGeneration`) | | Default dtype | BF16 | | Default inference parallelism | `TP=4, PP=1, EP=1` | -| Sequence length in recipes | 8192 | +| Sequence length in recipes | 4096 | **Requirements:** use `--trust_remote_code` when loading the Hugging Face checkpoint. @@ -64,26 +64,28 @@ ${WORKSPACE}/models/EXAONE-4.5-33B Equivalent direct command: ```bash -uv run python examples/conversion/convert_checkpoints.py import \ +./scripts/conversion/convert.sh import \ --hf-model LGAI-EXAONE/EXAONE-4.5-33B \ - --megatron-path "${WORKSPACE}/models/EXAONE-4.5-33B" + --megatron-path "${WORKSPACE}/models/EXAONE-4.5-33B" \ + --trust-remote-code ``` ## Finetune Recipes Available recipes: -- `exaone45_vl_sft_config`: Full supervised fine-tuning for EXAONE 4.5 VL. -- `exaone45_vl_peft_config`: Parameter-efficient fine-tuning with LoRA or DoRA. +- `exaone45_vl_33b_sft_16gpu_h100_bf16_config`: Full supervised fine-tuning for EXAONE 4.5 VL. +- `exaone45_vl_33b_peft_4gpu_h100_bf16_config`: Parameter-efficient fine-tuning with LoRA or DoRA. Recipe defaults: | Recipe | TP | PP | Notes | |---|---:|---:|---| -| `exaone45_vl_sft_config` | 4 | 1 | Full SFT, sequence parallel enabled | -| `exaone45_vl_peft_config` | 1 | 1 | PEFT, sequence parallel disabled | +| `exaone45_vl_33b_sft_16gpu_h100_bf16_config` | 4 | 4 | Full SFT, sequence parallel disabled | +| `exaone45_vl_33b_peft_4gpu_h100_bf16_config` | 4 | 1 | PEFT, sequence parallel disabled | -The recipes freeze the vision encoder by default (`freeze_vision_model=True`) and keep the language model and vision projection trainable. +Both recipes train the vision encoder, language model, and vision projection. Their micro batch size is 1, so in-batch +packing is disabled. Before training, set the usual runtime environment variables as needed: diff --git a/examples/models/exaone/exaone45/conversion.sh b/examples/models/exaone/exaone45/conversion.sh index 39dbd7c0581..53ab8b741de 100755 --- a/examples/models/exaone/exaone45/conversion.sh +++ b/examples/models/exaone/exaone45/conversion.sh @@ -17,6 +17,7 @@ set -euo pipefail WORKSPACE=${WORKSPACE:-/workspace} -uv run python examples/conversion/convert_checkpoints.py import \ +./scripts/conversion/convert.sh import \ --hf-model LGAI-EXAONE/EXAONE-4.5-33B \ - --megatron-path "${WORKSPACE}/models/EXAONE-4.5-33B" + --megatron-path "${WORKSPACE}/models/EXAONE-4.5-33B" \ + --trust-remote-code diff --git a/examples/models/gemma/gemma4_vl/README.md b/examples/models/gemma/gemma4_vl/README.md index 40365461662..d82b5580c47 100644 --- a/examples/models/gemma/gemma4_vl/README.md +++ b/examples/models/gemma/gemma4_vl/README.md @@ -111,24 +111,10 @@ We provide a [Weights & Biases report](https://api.wandb.ai/links/nvidia-nemo-fw ## Evaluation -After training, use [eval_sft_cord_v2.py](eval_sft_cord_v2.py) to verify the fine-tuned checkpoint on CORD-v2. It feeds the full conversation (image + prompt + ground-truth response) through the model in a single forward pass and reports per-example cross-entropy loss, token accuracy, and GT vs. predicted text. - -Example invocation (single node, 8 GPUs). Replace `` with your training job ID: - -```bash -uv run python -m torch.distributed.run --nproc_per_node=8 \ - examples/models/gemma/gemma4_vl/eval_sft_cord_v2.py \ - --hf_model_path google/gemma-4-26B-A4B-it \ - --megatron_model_path ${WORKSPACE}/results/gemma4_vl_sft_tp2_pp1_ep8_ \ - --tp 2 --pp 1 --ep 4 \ - --num_examples 20 -``` - -For batch evaluation on Slurm, see [slurm_eval_sft.sh](slurm_eval_sft.sh). - -After 100 SFT iterations on CORD-v2, expected teacher-forced token accuracy is ~98%. - -> **These scripts run one sample at a time and are intended only as sanity checks of the trained checkpoint.** For production inference, re-export the checkpoint to HF format using the export step in [conversion.sh](conversion.sh) and run with vLLM. +Use the [Gemma 4 26B-A4B-it model verification +card](../../../model_verification_cards/gemma-4-26b-a4b-it/card.yaml) as the canonical source for checked-in training, +export, and deterministic inference evidence. Run only items marked `verified`; an item without a checked-in command is +not a supported evaluation workflow. ## LoRA Merge diff --git a/src/megatron/bridge/recipes/exaone/h100/exaone45.py b/src/megatron/bridge/recipes/exaone/h100/exaone45.py index 9e9a3e450c2..80000719625 100644 --- a/src/megatron/bridge/recipes/exaone/h100/exaone45.py +++ b/src/megatron/bridge/recipes/exaone/h100/exaone45.py @@ -62,6 +62,7 @@ def _apply_exaone45_common(cfg: ConfigContainer) -> None: cfg.dataset.seq_length = 4096 cfg.dataset.hf_processor_path = _HF_PATH + cfg.dataset.enable_in_batch_packing = False cfg.ddp.overlap_grad_reduce = False cfg.ddp.overlap_param_gather = False diff --git a/src/megatron/bridge/recipes/glm_vl/h100/glm_45v.py b/src/megatron/bridge/recipes/glm_vl/h100/glm_45v.py index ed15fc542c4..5a092a173c6 100644 --- a/src/megatron/bridge/recipes/glm_vl/h100/glm_45v.py +++ b/src/megatron/bridge/recipes/glm_vl/h100/glm_45v.py @@ -94,6 +94,7 @@ def glm_45v_sft_128gpu_h100_bf16_config() -> ConfigContainer: - Sequence length: 8192 """ cfg = _sft_common_vlm() + cfg.dataset.enable_in_batch_packing = False # Model configuration hf_path = "zai-org/GLM-4.5V" @@ -244,6 +245,7 @@ def glm_45v_peft_32gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con peft_scheme: PEFT scheme - "lora", "dora", or a custom PEFT instance. """ cfg = _peft_common_vlm() + cfg.dataset.enable_in_batch_packing = False # PEFT scheme if isinstance(peft_scheme, str) and peft_scheme.lower() in ["lora", "dora"]: From 324a8f711ff31aed60df85b0d3f337a7c1cedd99 Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Sat, 8 Aug 2026 09:43:58 -0700 Subject: [PATCH 51/85] cp: `[models, conversion] fix: verify Step-3.5 model workflows (5334)` into `r0.6.0` (#5424) Signed-off-by: yaoyu-33 --- .../step35-flash/card.yaml | 259 ++++++++++++++++++ .../bridge/models/conversion/model_bridge.py | 21 +- .../bridge/models/hf_pretrained/state.py | 15 +- .../bridge/models/stepfun/step35_bridge.py | 46 +++- .../models/stepfun/test_step35_bridge.py | 21 ++ .../models/test_hf_pretrained_state.py | 27 ++ tests/unit_tests/models/test_model_bridge.py | 59 ++++ 7 files changed, 439 insertions(+), 9 deletions(-) create mode 100644 examples/model_verification_cards/step35-flash/card.yaml diff --git a/examples/model_verification_cards/step35-flash/card.yaml b/examples/model_verification_cards/step35-flash/card.yaml new file mode 100644 index 00000000000..12ed26c46fa --- /dev/null +++ b/examples/model_verification_cards/step35-flash/card.yaml @@ -0,0 +1,259 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: step35_flash +summary: > + Performance disclaimer: this model has not been performance-tuned; reported + timing and throughput metrics are sanity checks, not optimized performance + results. This card intentionally limits its current verification scope to + model-level conversion, parity, and checkpoint-backed inference for the + pinned Step-3.5-Flash checkpoint. Training verification is deferred and no + training-support result is claimed by this card. Four model-level conversion + items are verified; inference is unverified because the source evidence used + a launcher task not shipped on r0.6.0, and manual HF/Megatron logit + correlation remains below the required gate. +verification_index: + model_level: + verified: + - hf_to_megatron_cpu + - hf_to_megatron_gpu + - megatron_to_hf_cpu + - megatron_to_hf_gpu + unverified: + - manual_forward_pass + - inference + training: + H100: + unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] +model: + hf_id: stepfun-ai/Step-3.5-Flash + hf_revision: ab446a3de5e171ea341227e24bb1f090e1b771f7 # pragma: allowlist secret + architecture: Step3p5ForCausalLM + min_transformers_version: "5.8.1" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.04-py3 + bridge_commit: d5c416c50e60a6fa2ddb4a840d2ff036a08ce864 # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + bridge_commit: dd0fe603c598d1e031ada0f1af4c5040b237cdef # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu + --nodes 1 --hf-model stepfun-ai/Step-3.5-Flash + --hf-revision ab446a3de5e171ea341227e24bb1f090e1b771f7 + --megatron-path work/model-verification/step35-flash/cpu-megatron + --torch-dtype bfloat16 --trust-remote-code + last_verified: 2026-08-05 + expected_result: > + CPU import completed all mappings for the immutable 44-shard source, + constructed 199,997,706,240 parameters, and created iter_0000000. The + checkpoint reloaded successfully for CPU export; its subsequent strict + audit covered all 804 source tensors. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + bridge_commit: dd0fe603c598d1e031ada0f1af4c5040b237cdef # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 --hf-model stepfun-ai/Step-3.5-Flash + --hf-revision ab446a3de5e171ea341227e24bb1f090e1b771f7 + --megatron-path work/model-verification/step35-flash/imported-megatron + --torch-dtype bfloat16 --tp 1 --pp 1 --ep 8 --etp 1 + --distributed-timeout-minutes 110 --trust-remote-code --low-memory-save + last_verified: 2026-08-05 + expected_result: > + Distributed import completed at TP1/PP1/EP8/ETP1, constructed + 33,525,780,480 parameters per rank, and created iter_0000000. The + persisted checkpoint then reloaded successfully for both distributed + export and checkpoint-backed generation. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu + --nodes 1 --hf-model stepfun-ai/Step-3.5-Flash + --hf-revision ab446a3de5e171ea341227e24bb1f090e1b771f7 + --megatron-path + work/model-verification/step35-flash/cpu-megatron/iter_0000000 + --hf-path work/model-verification/step35-flash/cpu-hf-export + --torch-dtype bfloat16 --trust-remote-code + last_verified: 2026-08-05 + expected_result: > + CPU export reloaded the imported checkpoint and wrote 44 indexed BF16 + shards with all 804 source keys, shapes, and dtypes. Of 398,768,626,944 + serialized bytes, 395,600,878,848 bytes across 801 tensors matched the + immutable source bitwise. The three published MTP output entries occupy + 3,167,748,096 bytes and were intentionally materialized from Megatron's + shared lm_head representation; each exported entry matched lm_head + bitwise and used independent safetensors storage. A Transformers reload + through the pinned checkpoint's metadata/RoPE compatibility view ignored + only the three published inference-unused MTP layers, predicted the same + next token as the source (' Paris'), and reached cosine similarity + 0.997260 with maximum/mean absolute logit differences 1.320312/0.209649. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + bridge_commit: 0b5cc0206e77b4d18139c2ff4fc0d353ea260c20 # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 --hf-model stepfun-ai/Step-3.5-Flash + --hf-revision ab446a3de5e171ea341227e24bb1f090e1b771f7 + --megatron-path + work/model-verification/step35-flash/imported-megatron/iter_0000000 + --hf-path work/model-verification/step35-flash/hf-export + --torch-dtype bfloat16 + --tp 1 --pp 1 --ep 8 --etp 1 --distributed-timeout-minutes 110 + --distributed-save --save-every-n-ranks 1 --no-progress + --trust-remote-code + last_verified: 2026-08-05 + expected_result: > + Distributed export reloaded the TP1/PP1/EP8/ETP1 checkpoint and wrote 44 + indexed BF16 shards with all 804 source keys, shapes, and dtypes. The + exact audit found 801 tensors and 395,600,878,848 bytes bit-identical to + the source. Megatron-Core shares the output head with MTP, so the three + source MTP output tensors were explicitly normalized to lm_head.weight; + all three exported copies matched that tensor bitwise and had independent + storage. The generic grouped-expert export path recovered from CUDA stack + pressure by moving only the affected merge to CPU. A Transformers reload + through the pinned checkpoint's metadata/RoPE compatibility view ignored + only the three published inference-unused MTP layers, predicted the same + next token as the source (' Paris'), and reached cosine similarity + 0.997260 with maximum/mean absolute logit differences 1.320312/0.209649. + + manual_forward_pass: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + The 2026-08-05 full-checkpoint attempt used the public model-comparison + launcher and memory-bounded pinned-HF reference logits for input IDs + [0, 671, 6102, 294, 8760, 344]. Hugging Face and Megatron both predicted + token 11111 (' Paris'), but cosine similarity was 0.987008, below the + required 0.99 gate; maximum and mean absolute logit differences were + 4.937500 and 1.064947. Exact weight audits passed, so this remains an + unverified runtime-architecture/configuration follow-up rather than a + conversion-key mismatch. A future rerun must preserve the next-token + match and reach cosine similarity of at least 0.99. + + inference: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Checkpoint-backed inference must be rerun through a public inference task + supported on r0.6.0. The source evidence used the main-only + legacy-full-prefix-generation task, so it is not transferred as verified + release evidence. + + pretrain: + H100: + status: unverified + precision: bf16 + enabled_features: + context_parallel_size: 8 + moe_dispatcher: deepep + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The public Step-3.5-Flash H100 recipe must complete a bounded 100-step + run with finite loss, no skipped or NaN iterations, all four metrics, + and a reloadable final checkpoint. Training is deferred by this card. + + sft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A pinned-data 100-step full-SFT run must finish with finite loss, no + skipped or NaN iterations, all four metrics, and a reloadable final + checkpoint. Training is deferred by this card. + + sft_export_inference: + H100: + status: unverified + precision: bf16 + depends_on: sft + commands: null + last_verified: null + expected_result: > + A verified SFT checkpoint must export to Hugging Face, strictly reload, + and produce deterministic checkpoint-backed inference. Training and + post-SFT export verification are deferred by this card. + + sft_long_context: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A dedicated packed long-context SFT run must complete 100 steps with + finite loss, no skipped or NaN iterations, all four metrics, and a + reloadable checkpoint. Training is deferred by this card. + + peft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A Step-3.5-Flash PEFT recipe with an audited adapter target set must + complete 100 steps with finite loss, all four metrics, and a reloadable + adapter checkpoint. Training is deferred by this card. + + checkpoint_resume: + H100: + status: unverified + precision: bf16 + depends_on: pretrain + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: null + expected_result: > + A direct continuation must restore model, optimizer, scheduler, + data-order, and RNG state, match declared sentinel losses, and save a + reloadable checkpoint to a distinct output root. Training is deferred + by this card. diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index 713ded62aa1..3514d765df6 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -15,6 +15,7 @@ import abc import contextlib import fnmatch +import gc import itertools import logging import math @@ -979,7 +980,23 @@ def _accumulate_grouped_export( if len(grouped_buffers[group_key]) != num_experts: continue - merged = torch.stack([grouped_buffers[group_key][i] for i in range(num_experts)], dim=0) + grouped_tensors = [grouped_buffers[group_key][i] for i in range(num_experts)] + try: + merged = torch.stack(grouped_tensors, dim=0) + except torch.OutOfMemoryError: + if not any(getattr(tensor, "is_cuda", False) for tensor in grouped_tensors): + raise + logger.warning( + "Grouped expert export ran out of CUDA memory while stacking %d experts; retrying on CPU.", + num_experts, + ) + cpu_tensors = [tensor.cpu() for tensor in grouped_tensors] + del grouped_tensors + del grouped_buffers[group_key] + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + merged = torch.stack(cpu_tensors, dim=0) if getattr(task.mapping, "transpose_on_export", False): if group_key in hf_state_dict: @@ -994,7 +1011,7 @@ def _accumulate_grouped_export( else: merged = merged.transpose(-1, -2).contiguous() - del grouped_buffers[group_key] + grouped_buffers.pop(group_key, None) merged_result[group_key] = merged return merged_result or None diff --git a/src/megatron/bridge/models/hf_pretrained/state.py b/src/megatron/bridge/models/hf_pretrained/state.py index d14adaa1c75..a6cc100dd64 100644 --- a/src/megatron/bridge/models/hf_pretrained/state.py +++ b/src/megatron/bridge/models/hf_pretrained/state.py @@ -80,6 +80,11 @@ def _resolve_output_shard_path(output_path: Path, filename: str) -> Path: return output_file_path +def _contiguous_safetensors(tensors: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Return tensors in the contiguous layout required by safetensors.""" + return {name: tensor if tensor.is_contiguous() else tensor.contiguous() for name, tensor in tensors.items()} + + class StateDict(Mapping[str, torch.Tensor]): """ A state dict accessor that provides a unified interface for querying model @@ -752,7 +757,7 @@ def save_generator( if not key_to_filename_map: buffered_tensors = dict(generator) if buffered_tensors: - save_file(buffered_tensors, output_path / "model.safetensors") + save_file(_contiguous_safetensors(buffered_tensors), output_path / "model.safetensors") return filename_to_keys_map = defaultdict(set) @@ -797,7 +802,7 @@ def save_generator( output_file_path = _resolve_output_shard_path(output_path, filename) output_file_path.parent.mkdir(parents=True, exist_ok=True) - save_file(tensors_to_save, output_file_path) + save_file(_contiguous_safetensors(tensors_to_save), output_file_path) total_saved_tensor_bytes += sum( tensor.numel() * tensor.element_size() for tensor in tensors_to_save.values() ) @@ -833,7 +838,7 @@ def save_generator( tensors_to_save = {key: buffered_tensors[key] for key in keys_for_file if key in buffered_tensors} output_file_path = _resolve_output_shard_path(output_path, filename) output_file_path.parent.mkdir(parents=True, exist_ok=True) - save_file(tensors_to_save, output_file_path) + save_file(_contiguous_safetensors(tensors_to_save), output_file_path) total_saved_tensor_bytes += sum( tensor.numel() * tensor.element_size() for tensor in tensors_to_save.values() ) @@ -959,7 +964,7 @@ def _save_generator_distributed( if is_saver_rank and saver_index == 0: buffered_tensors = dict(generator) if buffered_tensors: - save_file(buffered_tensors, output_path / "model.safetensors") + save_file(_contiguous_safetensors(buffered_tensors), output_path / "model.safetensors") else: for _ in generator: pass @@ -1033,7 +1038,7 @@ def _save_generator_distributed( continue output_file_path = _resolve_output_shard_path(output_path, fname) output_file_path.parent.mkdir(parents=True, exist_ok=True) - save_file(tensors_to_save, output_file_path) + save_file(_contiguous_safetensors(tensors_to_save), output_file_path) actually_saved_keys.update(tensors_to_save.keys()) local_saved_tensor_bytes += sum( tensor.numel() * tensor.element_size() for tensor in tensors_to_save.values() diff --git a/src/megatron/bridge/models/stepfun/step35_bridge.py b/src/megatron/bridge/models/stepfun/step35_bridge.py index b0785168911..ccd6d34aed7 100755 --- a/src/megatron/bridge/models/stepfun/step35_bridge.py +++ b/src/megatron/bridge/models/stepfun/step35_bridge.py @@ -14,7 +14,7 @@ import logging from functools import partial -from typing import Dict +from typing import Dict, Mapping import torch from megatron.core.models.gpt.gpt_layer_specs import ( @@ -26,7 +26,7 @@ from transformers import AutoConfig from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry -from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, WeightConversionTask from megatron.bridge.models.conversion.param_mapping import ( AutoMapping, GatedMLPMapping, @@ -214,6 +214,48 @@ class Step35Bridge(MegatronModelBridge): ("layer_types", "layer_types"), ] + def maybe_modify_converted_hf_weight( + self, + task: WeightConversionTask, + converted_weights_dict: Dict[str, torch.Tensor], + hf_state_dict: Mapping[str, torch.Tensor], + ) -> Dict[str, torch.Tensor]: + """Materialize HF MTP output entries from Megatron's shared output layer. + + The published checkpoint serializes a separate output tensor for each + MTP layer, although its HF inference implementation ignores those + layers. Megatron-Core instead defines MTP with a shared output layer. + Export the Megatron representation faithfully by materializing one + independent CPU tensor per HF entry from that shared output weight. + """ + converted_weights_dict = super().maybe_modify_converted_hf_weight( + task, + converted_weights_dict, + hf_state_dict, + ) + if self.hf_config is None or "lm_head.weight" not in converted_weights_dict: + return converted_weights_dict + + first_mtp_layer = self.hf_config.num_hidden_layers + num_mtp_layers = getattr(self.hf_config, "num_nextn_predict_layers", 0) + mtp_output_names = [ + f"model.layers.{first_mtp_layer + offset}.transformer.shared_head.output.weight" + for offset in range(num_mtp_layers) + ] + missing_output_names = [ + name for name in mtp_output_names if name in hf_state_dict and name not in converted_weights_dict + ] + if not missing_output_names: + return converted_weights_dict + + # Safetensors forbids shared storage. Keep distinct CPU storage for + # each entry without adding several gigabytes of peak CUDA memory. + output_weight = converted_weights_dict["lm_head.weight"].detach().cpu() + converted_weights_dict["lm_head.weight"] = output_weight + for name in missing_output_names: + converted_weights_dict[name] = output_weight.clone() + return converted_weights_dict + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider: """Convert HuggingFace Step3.5 config to GPTModelProvider. diff --git a/tests/unit_tests/models/stepfun/test_step35_bridge.py b/tests/unit_tests/models/stepfun/test_step35_bridge.py index 933014bfa60..4682a211ccc 100644 --- a/tests/unit_tests/models/stepfun/test_step35_bridge.py +++ b/tests/unit_tests/models/stepfun/test_step35_bridge.py @@ -675,6 +675,27 @@ def test_mtp_specific_norm_and_proj_present(self): assert "mtp.layers.0.final_layernorm.weight" in params +def test_mtp_output_entries_are_materialized_from_megatron_shared_head(): + bridge = Step35Bridge() + bridge.hf_config = SimpleNamespace(num_hidden_layers=45, num_nextn_predict_layers=3) + output_weight = torch.ones(8, 4) + expected_output_names = {f"model.layers.{layer}.transformer.shared_head.output.weight" for layer in range(45, 48)} + hf_state_dict = {name: torch.zeros_like(output_weight) for name in expected_output_names} + + converted = bridge.maybe_modify_converted_hf_weight( + SimpleNamespace(), + {"lm_head.weight": output_weight}, + hf_state_dict, + ) + + assert converted.keys() == {"lm_head.weight", *expected_output_names} + for name in expected_output_names: + torch.testing.assert_close(converted[name], output_weight) + assert not torch.equal(converted[name], hf_state_dict[name]) + all_output_names = {"lm_head.weight", *expected_output_names} + assert len({converted[name].data_ptr() for name in all_output_names}) == len(all_output_names) + + # --------------------------------------------------------------------------- # StackedExpertAutoMapping / StackedExpertGatedMLPMapping # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/models/test_hf_pretrained_state.py b/tests/unit_tests/models/test_hf_pretrained_state.py index 23a44fa13a6..70e52365f29 100644 --- a/tests/unit_tests/models/test_hf_pretrained_state.py +++ b/tests/unit_tests/models/test_hf_pretrained_state.py @@ -148,6 +148,33 @@ def gather_rank_zero(output: list[object | None], value: object) -> None: assert index_data["metadata"] == {"format": "pt", "total_size": expected_total_size} +@pytest.mark.parametrize("distributed_save", [False, True]) +def test_save_generator_materializes_noncontiguous_tensors(tmp_path, monkeypatch, distributed_save: bool) -> None: + shard = "model-00001-of-00001.safetensors" + _write_safetensors_index(tmp_path, {"model.weight": shard}) + source = SafeTensorsStateSource(tmp_path) + output_path = tmp_path / "output" + weight = torch.arange(12, dtype=torch.float32).reshape(3, 4).transpose(0, 1) + assert not weight.is_contiguous() + + if distributed_save: + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 1) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(torch.distributed, "barrier", lambda: None) + + def gather_rank_zero(output: list[object | None], value: object) -> None: + output[0] = value + + monkeypatch.setattr(torch.distributed, "all_gather_object", gather_rank_zero) + + source.save_generator(iter([("model.weight", weight)]), output_path, distributed_save=distributed_save) + + with safe_open(output_path / shard, framework="pt", device="cpu") as saved: + torch.testing.assert_close(saved.get_tensor("model.weight"), weight) + + def test_save_generator_writes_shard_as_soon_as_its_remaining_keys_arrive(tmp_path, monkeypatch) -> None: class _SubsetForbiddenSet(set[str]): def issubset(self, _other: Iterable[object]) -> bool: diff --git a/tests/unit_tests/models/test_model_bridge.py b/tests/unit_tests/models/test_model_bridge.py index 5ce38bcbf3a..9f598681c4a 100644 --- a/tests/unit_tests/models/test_model_bridge.py +++ b/tests/unit_tests/models/test_model_bridge.py @@ -285,6 +285,65 @@ def transform(name, tensor): ] +def test_grouped_export_retries_stack_on_cpu_after_cuda_oom(monkeypatch, caplog): + class FakeCudaTensor: + is_cuda = True + + def __init__(self, value): + self.value = value + + def cpu(self): + return torch.tensor([self.value]) + + original_stack = torch.stack + stack_devices = [] + + def stack_with_cuda_oom(tensors, dim=0): + stack_devices.append([type(tensor).__name__ for tensor in tensors]) + if isinstance(tensors[0], FakeCudaTensor): + raise torch.OutOfMemoryError("simulated grouped-export CUDA OOM") + return original_stack(tensors, dim=dim) + + monkeypatch.setattr(torch, "stack", stack_with_cuda_oom) + monkeypatch.setattr( + "megatron.bridge.models.conversion.model_bridge.parallel_state.get_expert_model_parallel_world_size", + lambda: 1, + ) + mapping = SimpleNamespace(is_grouped_export=True, ep_size=1) + model_config = SimpleNamespace(num_moe_experts=2) + buffers = {} + + first = MegatronModelBridge._accumulate_grouped_export( + None, + SimpleNamespace( + mapping=mapping, + param_name="decoder.layers.0.mlp.experts.linear_fc2.weight0", + ), + {"hf.grouped": FakeCudaTensor(1.0)}, + model_config, + buffers, + {}, + ) + with caplog.at_level("WARNING"): + second = MegatronModelBridge._accumulate_grouped_export( + None, + SimpleNamespace( + mapping=mapping, + param_name="decoder.layers.0.mlp.experts.linear_fc2.weight1", + ), + {"hf.grouped": FakeCudaTensor(2.0)}, + model_config, + buffers, + {}, + ) + + assert first is None + torch.testing.assert_close(second["hf.grouped"], torch.tensor([[1.0], [2.0]])) + assert stack_devices == [["FakeCudaTensor", "FakeCudaTensor"], ["Tensor", "Tensor"]] + assert "retrying on CPU" in caplog.text + assert buffers == {} + + def test_stream_weights_megatron_to_hf_finalizes_exported_tensors_before_cpu(monkeypatch): bridge = DummyBridge() events = [] From bad29070400f5f5409e055072bf203e2e1b003be Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:33:20 +0200 Subject: [PATCH 52/85] cp: `fix(training): surface distributed failures without cleanup hangs (5283)` into `r0.6.0` (#5327) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: Chen Cui --- .../gemma-4-26b-a4b-it/card.yaml | 120 +++++++++++------- .../recipes/gemma4_vl/h100/gemma4_vl.py | 35 +++-- src/megatron/bridge/training/pretrain.py | 45 ++++++- .../recipes/test_gemma4_vl_recipes.py | 51 ++++++++ tests/unit_tests/training/test_pretrain.py | 95 ++++++++++++-- 5 files changed, 273 insertions(+), 73 deletions(-) diff --git a/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml b/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml index a06a9db30fe..06b7e6da2b4 100644 --- a/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml +++ b/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml @@ -5,14 +5,17 @@ title: gemma_4_26b_a4b_it summary: > Performance disclaimer: this model has not been performance-tuned; reported timing and throughput metrics are sanity checks, not optimized performance - results. Gemma 4 26B-A4B-it BF16 CPU and GPU import, exact CPU and GPU - export, and deterministic multimodal inference are verified. Manual forward - comparison remains unverified because the recorded run compared a - tensor-parallel padding position rather than the prompt's next token. Full - SFT remains unverified because the default and selective-MLP configurations - exceed H100 memory, while full-layer recompute is too slow for the 10-step - verification window. Long-context SFT and PEFT are pending. The published - Bridge recipes do not currently provide VLM pretraining. + results. Gemma 4 26B-A4B-it BF16 CPU and GPU import and exact CPU and GPU + export are verified. Deterministic multimodal inference remains unverified + because the recorded run stopped after 29 of the requested 32 tokens. Manual + forward comparison remains unverified because the recorded run compared a + tensor-parallel padding position rather than the prompt's next token. + Decoder-focused SFT is verified on one 8-GPU H100 node with TP4/PP1/EP8/ETP1, + a frozen vision encoder, and selective core-attention and MoE-activation + recompute. The real CORD-v2 run completed 10 optimizer steps with stable + memory and saved iter_0000010. LoRA PEFT is also verified on four H100 GPUs + with TP2/PP1/EP4/ETP1 and a complete adapter checkpoint. Long-context SFT is + pending. The published Bridge recipes do not currently provide VLM pretraining. verification_index: model_level: verified: @@ -20,12 +23,13 @@ verification_index: - hf_to_megatron_gpu - megatron_to_hf_cpu - megatron_to_hf_gpu - - inference unverified: - manual_forward_pass + - inference training: H100: - unverified: [sft, sft_export_inference, sft_long_context, peft] + verified: [sft, peft] + unverified: [sft_export_inference, sft_long_context] unsupported: [pretrain, checkpoint_resume] model: hf_id: google/gemma-4-26B-A4B-it @@ -128,9 +132,8 @@ items: agreement and cosine similarity of at least 0.99. inference: - status: verified + status: unverified precision: bf16 - bridge_commit: c4250ac1382eccc8e2b58c2424b3db85415fbf6f # pragma: allowlist secret command: > uv run python -m torch.distributed.run --standalone --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_vlm.py @@ -141,17 +144,18 @@ items: --prompt "What animal is on the candy?" --max_new_tokens 32 --tp 2 --pp 1 --ep 4 --etp 1 - last_verified: 2026-07-23 + last_verified: null expected_result: > - One BF16 deterministic greedy run used an exact 32-token generation - limit and terminated on EOS after 29 generated tokens. The JSON-escaped - literal completion was "\u003c|channel\u003ethought\n\u003cchannel|\u003eThe + The 2026-07-23 BF16 deterministic greedy run stopped on EOS after 29 of + the requested 32 tokens, so it does not satisfy the exact-length gate. + Its JSON-escaped diagnostic completion was + "\u003c|channel\u003ethought\n\u003cchannel|\u003eThe animal on the candy is a turtle.\u003cturn|\u003e\u003cturn|\u003e\u003cturn|\u003e\u002f\u002fthought\n\u003cchannel|\u003eThe animal on the candy is a turtle.\u003cturn|\u003e\u003cturn|\u003e\u003cturn|\u003e\u003ceos\u003e". - The input was - p-blog/candy.JPG from huggingface/documentation-images at revision + Verification requires a fresh run that emits exactly 32 tokens. The + input was p-blog/candy.JPG from huggingface/documentation-images at revision 44c3249461c5a5336585d1e31e0a4ce7fce06a5a, with SHA-256 fc417c899e94f8df465b7541c5a70f0eebb85c414d06345f0b290c061eccc84c. @@ -174,9 +178,10 @@ items: sft: H100: - status: unverified + status: verified precision: bf16 enabled_features: {} + bridge_commit: 81e559ca4d579cf2d9ed5b8f3f48bef82e06451f # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 --recipe gemma4_vl_26b_sft_config --mode sft --dataset cord-v2 @@ -186,20 +191,29 @@ items: validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null --save_dir work/model-verification/gemma-4-26b-a4b-it/sft-checkpoints --save_interval 10 logger.log_interval=1 logger.log_throughput=true - last_verified: null + logger.save_config_filepath=work/model-verification/gemma-4-26b-a4b-it/sft-post-setup.yaml + last_verified: 2026-08-04 metrics: - initial_loss: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + initial_loss: 1.268018 + final_loss: 0.06543536 + last_10_steps_step_time_ms_avg: 18699.99 + last_10_steps_model_tflops_per_gpu_avg: 26.47 + peak_allocated_memory_gib: 70.624 + peak_reserved_memory_gib: 71.021 expected_result: > - On 2026-07-23 the default 8-GPU TP2/PP1/EP8/ETP1 run and a GBS4 - selective-MLP recompute run both failed before the first completed - optimizer step. Each rank had about 77.84 to 78.70 GiB in use and - needed an additional 1 to 2 GiB. Full 30-layer recompute crossed that - memory peak, but one GBS4 optimizer step did not finish during more - than nine observed minutes. No valid 10-step metrics or iter_0000010 - checkpoint were produced, so SFT remains unverified. + On 2026-08-04 the single-node TP4/PP1/EP8/ETP1 run loaded the immutable + BF16 checkpoint and real CORD-v2 revision, then completed exactly 10 + optimizer steps at MBS1 and GBS32. The persisted post-setup config + confirmed a frozen vision encoder, trainable projection and language + model, expandable CUDA allocator segments, and selective core_attn plus + moe_act recompute with CUDA graphs and CPU offload disabled. Loss moved + from 1.268018 to 0.06543536, final grad norm was 1.561, and every step + reported zero skipped and zero NaN iterations. Across all eight ranks, + allocated memory returned to 56.493 GiB after each step, while + cumulative peak allocated and reserved memory remained at or below + 70.624 and 71.021 GiB. The run saved a complete iter_0000010 torch_dist + checkpoint with eight rank shards, finalized metadata, training state, + and a latest-checkpoint tracker containing iteration 10. sft_export_inference: H100: @@ -209,12 +223,11 @@ items: commands: null last_verified: null expected_result: > - This item is blocked because SFT produced no verified iter_0000010 - checkpoint, so no export or post-SFT inference command is recorded. - Once that dependency passes, the checkpoint must export, strictly - reload as Gemma4ForConditionalGeneration, and complete one - deterministic greedy VLM generation with a recorded literal - completion. + SFT produced a verified iter_0000010 checkpoint, but export, strict + Hugging Face reload, and post-SFT inference have not been run. This + item requires the checkpoint to export, strictly reload as + Gemma4ForConditionalGeneration, and complete one deterministic greedy + VLM generation with a recorded literal completion. sft_long_context: H100: @@ -237,9 +250,10 @@ items: peft: H100: - status: unverified + status: verified precision: bf16 enabled_features: {} + bridge_commit: df1b978e751753707100417cb4e89645a23537f8 # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 1 --gpus-per-node 4 --recipe gemma4_vl_26b_peft_config --mode lora --dataset cord-v2 @@ -249,16 +263,30 @@ items: validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null --save_dir work/model-verification/gemma-4-26b-a4b-it/peft-checkpoints --save_interval 10 logger.log_interval=1 logger.log_throughput=true - last_verified: null + logger.save_config_filepath=work/model-verification/gemma-4-26b-a4b-it/peft-post-setup.yaml + last_verified: 2026-08-04 metrics: - initial_loss: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null + initial_loss: 1.288328 + final_loss: 0.1072377 + last_10_steps_step_time_ms_avg: 14163.25 + last_10_steps_model_tflops_per_gpu_avg: 65.01 + peak_allocated_memory_gib: 44.252 + peak_reserved_memory_gib: 45.637 expected_result: > - The 4-GPU TP2/PP1/EP4 LoRA run must complete exactly 10 optimizer steps - with finite loss, no skipped or NaN iterations, all four required - metrics, and a reloadable final adapter checkpoint. + On 2026-08-04 the 4-GPU TP2/PP1/EP4/ETP1 LoRA run loaded the immutable + BF16 checkpoint and real CORD-v2 revision, then completed exactly 10 + optimizer steps at MBS1 and GBS32. The persisted post-setup config + recorded rank-32 adapters with alpha 32 and zero dropout on linear_qkv, + linear_proj, linear_fc1, and linear_fc2. Runtime transformation froze + the base model and reported 22,272,000 trainable adapter parameters, + 0.30% of the local model shard. Loss moved from 1.288328 to 0.1072377, + final grad norm was 22.994, and every step reported zero skipped and + zero NaN iterations. Across all four ranks, allocated memory returned + to 14.582 GiB after each step, while cumulative peak allocated and + reserved memory remained at or below 44.252 and 45.637 GiB. The run + saved a complete iter_0000010 adapter checkpoint with four rank shards, + finalized metadata, training state, and a latest-checkpoint tracker + containing iteration 10. checkpoint_resume: all: diff --git a/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py b/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py index 051333ebb5a..aed09f0af6d 100644 --- a/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py +++ b/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py @@ -110,8 +110,9 @@ def gemma4_vl_26b_sft_8gpu_h100_bf16_config() -> ConfigContainer: """Return a full SFT config for Gemma 4 VL 26B-A4B (MoE VLM). Default configuration: 8 GPUs - - TP=2, PP=1, EP=8 (max EP with 16 GPUs at TP=2,PP=1; DP=8, EP divides DP) - - No activation recompute — EP=8 shards 87.5% of expert params per GPU + - TP=4, PP=1, EP=8, ETP=1 (dense DP=2, expert DP=1) + - Selective core-attention and MoE-activation recompute + - Frozen vision encoder with trainable projection and language model - LR=5e-5 (full SFT) - Sequence length: 4096 """ @@ -119,14 +120,26 @@ def gemma4_vl_26b_sft_8gpu_h100_bf16_config() -> ConfigContainer: _apply_gemma4_vl_common(cfg, _HF_PATH) - # Parallel settings — TP=2, PP=1, EP=8 on 2×8 GPUs - # DP = 16/(TP*PP) = 8; EP=8 shards experts across all DP ranks (1 expert replica) - cfg.model.tensor_model_parallel_size = 2 + # The dense TP and expert EP meshes overlap: PP * max(TP, EP * ETP) = 8 GPUs. + # Dense DP = 8 / (TP * PP) = 2; expert DP = 8 / (PP * EP * ETP) = 1. + cfg.model.tensor_model_parallel_size = 4 cfg.model.pipeline_model_parallel_size = 1 cfg.model.pipeline_dtype = None cfg.model.virtual_pipeline_model_parallel_size = None cfg.model.context_parallel_size = 1 - cfg.model.expert_model_parallel_size = 8 # override common EP=1 + cfg.model.expert_model_parallel_size = 8 + cfg.model.expert_tensor_parallel_size = 1 + + # Real CORD-v2 measurements require both modules to preserve H100 memory + # headroom after optimizer-state initialization. Full-layer recompute is not needed. + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["core_attn", "moe_act"] + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + + # Decoder-focused SFT keeps the vision encoder fixed while training the + # multimodal projection and language model. + cfg.model.freeze_vision_model = True # Reduce overhead to fit within 30-min wall time. # 40 iters × ~15s + 5 min init + 4 evals × 35s = ~20 min → 10 min for checkpoint save. @@ -170,7 +183,8 @@ def gemma4_vl_26b_peft_4gpu_h100_bf16_config( """Return a PEFT (LoRA/DoRA) config for Gemma 4 VL 26B-A4B (MoE VLM). Default configuration: 4 GPUs - - TP=4, PP=1, EP=4 (PEFT needs less memory, drop PP) + - TP=2, PP=1, EP=4, ETP=1 (dense DP=2, expert DP=1) + - MBS=1, GBS=32 - LR=2e-4 (PEFT) - Sequence length: 4096 @@ -194,7 +208,12 @@ def gemma4_vl_26b_peft_4gpu_h100_bf16_config( cfg.model.pipeline_dtype = None cfg.model.virtual_pipeline_model_parallel_size = None cfg.model.context_parallel_size = 1 - cfg.model.expert_model_parallel_size = 4 # override common EP=1 + cfg.model.expert_model_parallel_size = 4 + cfg.model.expert_tensor_parallel_size = 1 + + # Real CORD-v2 measurements fit with substantial activation-memory headroom + # at MBS=1, so recompute and LoRA sequence-parallel input re-gather are not needed. + cfg.train.micro_batch_size = 1 # Optimizer — higher LR for PEFT opt_cfg, scheduler_cfg = distributed_fused_adam_with_cosine_annealing( diff --git a/src/megatron/bridge/training/pretrain.py b/src/megatron/bridge/training/pretrain.py index 0b5fc59f36a..7c950d62513 100644 --- a/src/megatron/bridge/training/pretrain.py +++ b/src/megatron/bridge/training/pretrain.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +import os import torch.distributed as dist from nvidia_resiliency_ext.inprocess import CallWrapper @@ -196,9 +197,19 @@ def _pretrain( ) _finish_train(state, checkpoint_manager) - except BaseException: + except BaseException as error: if inprocess_call_wrapper is None: - _cleanup_after_pretrain_failure(state, should_destroy_process_group) + logger.exception( + "Pretrain failed on rank=%s with %s: %s. Bridge is aborting framework-owned distributed resources " + "because peer ranks may be blocked in collectives.", + _safe_distributed_rank(), + type(error).__name__, + error, + ) + try: + _cleanup_after_pretrain_failure(state, should_destroy_process_group) + except BaseException: + logger.exception("Unexpected failure while cleaning up after pretrain failure") raise _maybe_destroy_process_group(should_destroy_process_group) @@ -222,6 +233,16 @@ def _abort_async_checkpoint_worker(state: GlobalState) -> None: filesystem_async._results_queue = None +def _safe_distributed_rank() -> str: + """Return a rank identifier without obscuring an active training failure.""" + try: + if dist.is_initialized(): + return str(dist.get_rank()) + except Exception: + pass + return os.environ.get("RANK", "unknown") + + def _cleanup_after_pretrain_failure(state: GlobalState, should_destroy_process_group: bool) -> None: """Clean up framework-owned state after ordinary pretrain execution fails.""" try: @@ -235,19 +256,31 @@ def _cleanup_after_pretrain_failure(state: GlobalState, should_destroy_process_g logger.exception("Failed to destroy Megatron global state after pretrain failure") try: - _maybe_destroy_process_group(should_destroy_process_group, synchronize=False) + _maybe_destroy_process_group(should_destroy_process_group, synchronize=False, abort=True) except Exception: - logger.exception("Failed to destroy the process group after pretrain failure") + logger.exception("Failed to abort process groups after pretrain failure") -def _maybe_destroy_process_group(should_destroy: bool, *, synchronize: bool = True) -> None: - """Destroy the process group if it was created by this training session. +def _maybe_destroy_process_group( + should_destroy: bool, + *, + synchronize: bool = True, + abort: bool = False, +) -> None: + """Destroy or abort process groups created by this training session. Args: should_destroy: Whether the process group should be destroyed synchronize: Whether to synchronize ranks before destruction + abort: Whether to abort all process groups instead of waiting for their + outstanding work to finish """ if should_destroy and dist.is_initialized(): + if abort: + # Orderly shutdown waits for outstanding collectives, but failure + # cleanup may run while peer ranks are still blocked in those calls. + dist.distributed_c10d._abort_process_group() + return if synchronize: dist.barrier() dist.destroy_process_group() diff --git a/tests/unit_tests/recipes/test_gemma4_vl_recipes.py b/tests/unit_tests/recipes/test_gemma4_vl_recipes.py index 4c158d774b8..ee0fd95feff 100644 --- a/tests/unit_tests/recipes/test_gemma4_vl_recipes.py +++ b/tests/unit_tests/recipes/test_gemma4_vl_recipes.py @@ -49,3 +49,54 @@ def test_gemma4_vl_sft_uses_long_distributed_timeout(monkeypatch: pytest.MonkeyP cfg = _gemma4_vl_module.gemma4_vl_26b_sft_config() assert cfg.dist.distributed_timeout_minutes == 90 + + +def test_gemma4_vl_sft_uses_memory_stable_8gpu_contract(monkeypatch: pytest.MonkeyPatch): + """Full Gemma 4 VL SFT should use the measured single-node memory contract.""" + patch_recipe_module_global(monkeypatch, _gemma4_vl_module, "AutoBridge", _FakeAutoBridge) + + cfg = _gemma4_vl_module.gemma4_vl_26b_sft_config() + + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.sequence_parallel is True + assert cfg.model.recompute_granularity == "selective" + assert cfg.model.recompute_modules == ["core_attn", "moe_act"] + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None + assert cfg.model.cuda_graph_impl == "none" + assert cfg.model.freeze_vision_model is True + assert cfg.model.freeze_vision_projection is False + assert cfg.model.freeze_language_model is False + assert cfg.train.micro_batch_size == 1 + assert cfg.train.global_batch_size == 32 + assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + + +def test_gemma4_vl_peft_uses_memory_stable_4gpu_contract(monkeypatch: pytest.MonkeyPatch): + """Gemma 4 VL PEFT should use the measured single-node LoRA contract.""" + patch_recipe_module_global(monkeypatch, _gemma4_vl_module, "AutoBridge", _FakeAutoBridge) + + cfg = _gemma4_vl_module.gemma4_vl_26b_peft_config() + + assert cfg.model.tensor_model_parallel_size == 2 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 4 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.sequence_parallel is True + assert cfg.model.recompute_granularity is None + assert cfg.model.recompute_modules is None + assert cfg.model.cuda_graph_impl == "none" + assert cfg.train.micro_batch_size == 1 + assert cfg.train.global_batch_size == 32 + assert cfg.peft.target_modules == ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"] + assert cfg.peft.dim == 32 + assert cfg.peft.alpha == 32 + assert cfg.peft.dropout == 0.0 + assert cfg.peft.sequence_parallel_input_regather is False + assert cfg.peft.share_expert_adapters is True + assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" diff --git a/tests/unit_tests/training/test_pretrain.py b/tests/unit_tests/training/test_pretrain.py index 2b3edb485bf..2957244f6b2 100644 --- a/tests/unit_tests/training/test_pretrain.py +++ b/tests/unit_tests/training/test_pretrain.py @@ -14,6 +14,7 @@ """Unit tests for pretrain module process group cleanup.""" +import logging from unittest.mock import MagicMock, patch import pytest @@ -64,33 +65,97 @@ def test_no_destroy_when_neither_condition_met(self, mock_dist): mock_dist.barrier.assert_not_called() mock_dist.destroy_process_group.assert_not_called() + @patch("megatron.bridge.training.pretrain.dist") + def test_abort_does_not_wait_for_outstanding_collectives(self, mock_dist): + """Test failure cleanup aborts process groups without synchronization.""" + mock_dist.is_initialized.return_value = True + + _maybe_destroy_process_group(should_destroy=True, synchronize=False, abort=True) + + mock_dist.distributed_c10d._abort_process_group.assert_called_once_with() + mock_dist.barrier.assert_not_called() + mock_dist.destroy_process_group.assert_not_called() + class TestPretrainProcessGroupOwnership: """Test process group ownership across exceptional pretrain exits.""" - def test_framework_owned_process_group_is_destroyed_when_setup_raises(self): - """Test Bridge destroys a process group initialized during failed setup.""" + def test_framework_owned_failure_is_logged_before_cleanup_and_abort(self, caplog): + """Test Bridge logs the original failure before cleaning up framework-owned state.""" state = MagicMock() + async_queue = MagicMock() + state._async_calls_queue = async_queue + original_error = RuntimeError("setup failed after distributed initialization") + + def assert_failure_was_logged(): + failure_records = [ + record for record in caplog.records if record.getMessage().startswith("Pretrain failed") + ] + assert len(failure_records) == 1 with ( patch("megatron.bridge.training.pretrain.dist") as mock_dist, patch("megatron.bridge.training.pretrain.destroy_global_state") as mock_destroy_global_state, patch("megatron.bridge.training.pretrain.get_dataset_provider"), + patch("megatron.bridge.training.pretrain.setup", side_effect=original_error), + ): + mock_dist.is_initialized.side_effect = [False, True, True] + mock_dist.get_rank.return_value = 7 + async_queue.close.side_effect = lambda *, abort: assert_failure_was_logged() + mock_destroy_global_state.side_effect = assert_failure_was_logged + mock_dist.distributed_c10d._abort_process_group.side_effect = assert_failure_was_logged + + with ( + caplog.at_level(logging.ERROR, logger="megatron.bridge.training.pretrain"), + pytest.raises(RuntimeError) as exc_info, + ): + _pretrain(state, MagicMock()) + + assert exc_info.value is original_error + mock_destroy_global_state.assert_called_once_with() + mock_dist.barrier.assert_not_called() + mock_dist.distributed_c10d._abort_process_group.assert_called_once_with() + mock_dist.destroy_process_group.assert_not_called() + failure_record = next(record for record in caplog.records if record.getMessage().startswith("Pretrain failed")) + failure_message = failure_record.getMessage() + assert "rank=7" in failure_message + assert "RuntimeError: setup failed after distributed initialization" in failure_message + assert "aborting framework-owned distributed resources" in failure_message + assert "peer ranks may be blocked in collectives" in failure_message + assert failure_record.exc_info is not None + assert failure_record.exc_info[1] is original_error + assert "Traceback (most recent call last)" in caplog.text + assert "RuntimeError: setup failed after distributed initialization" in caplog.text + + def test_cleanup_failure_does_not_replace_original_failure(self, caplog): + """Test cleanup errors are logged while the original pretrain failure propagates.""" + state = MagicMock() + original_error = RuntimeError("original setup failure") + + with ( + patch("megatron.bridge.training.pretrain.dist") as mock_dist, patch( - "megatron.bridge.training.pretrain.setup", - side_effect=RuntimeError("setup failed after distributed initialization"), + "megatron.bridge.training.pretrain.destroy_global_state", + side_effect=RuntimeError("cleanup failure"), ), + patch("megatron.bridge.training.pretrain.get_dataset_provider"), + patch("megatron.bridge.training.pretrain.setup", side_effect=original_error), ): - mock_dist.is_initialized.side_effect = [False, True] + mock_dist.is_initialized.side_effect = [False, True, True] + mock_dist.get_rank.return_value = 2 - with pytest.raises(RuntimeError, match="setup failed after distributed initialization"): + with ( + caplog.at_level(logging.ERROR, logger="megatron.bridge.training.pretrain"), + pytest.raises(RuntimeError) as exc_info, + ): _pretrain(state, MagicMock()) - mock_destroy_global_state.assert_called_once_with() - mock_dist.barrier.assert_not_called() - mock_dist.destroy_process_group.assert_called_once_with() + assert exc_info.value is original_error + messages = [record.getMessage() for record in caplog.records] + assert messages[0].startswith("Pretrain failed on rank=2 with RuntimeError: original setup failure") + assert "Failed to destroy Megatron global state after pretrain failure" in messages - def test_framework_owned_async_worker_is_aborted_before_process_group_destroyed(self): + def test_framework_owned_async_worker_is_aborted_before_process_groups(self): """Test failed setup aborts its async checkpoint worker before distributed cleanup.""" state = MagicMock() async_queue = MagicMock() @@ -108,13 +173,15 @@ def setup_then_fail(*_args, **_kwargs): patch("megatron.bridge.training.pretrain.get_dataset_provider"), patch("megatron.bridge.training.pretrain.setup", side_effect=setup_then_fail), ): - mock_dist.is_initialized.side_effect = [False, True] - mock_dist.destroy_process_group.side_effect = lambda: cleanup_order.append(("process_group", None)) + mock_dist.is_initialized.side_effect = [False, True, True] + mock_dist.distributed_c10d._abort_process_group.side_effect = lambda: cleanup_order.append( + ("process_group", "abort") + ) with pytest.raises(RuntimeError, match="setup failed after async worker initialization"): _pretrain(state, MagicMock()) - assert cleanup_order == [("async_queue", True), ("process_group", None)] + assert cleanup_order == [("async_queue", True), ("process_group", "abort")] assert state._async_calls_queue is None def test_caller_owned_process_group_is_preserved_when_setup_raises(self): @@ -145,6 +212,7 @@ def test_inprocess_restart_wrapper_retains_cleanup_ownership_when_setup_raises(s patch("megatron.bridge.training.pretrain.destroy_global_state") as mock_destroy_global_state, patch("megatron.bridge.training.pretrain.get_dataset_provider"), patch("megatron.bridge.training.pretrain.setup", side_effect=RuntimeError("setup failed")), + patch("megatron.bridge.training.pretrain.logger.exception") as mock_logger_exception, ): mock_dist.is_initialized.return_value = False @@ -158,5 +226,6 @@ def test_inprocess_restart_wrapper_retains_cleanup_ownership_when_setup_raises(s mock_dist.PrefixStore.assert_called_once_with("1", store) mock_destroy_global_state.assert_not_called() + mock_logger_exception.assert_not_called() mock_dist.barrier.assert_not_called() mock_dist.destroy_process_group.assert_not_called() From d9121e7e051ee228f1cede4d39ea62a56c6672bd Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:33:43 +0200 Subject: [PATCH 53/85] cp: `fix(modelopt): handle sparse export task lists (5374)` into `r0.6.0` (#5446) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../models/conversion/modelopt_utils.py | 19 ++++--- tests/unit_tests/models/test_model_bridge.py | 56 +++++++++++++++++++ 2 files changed, 67 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/conversion/modelopt_utils.py b/src/megatron/bridge/models/conversion/modelopt_utils.py index e25c5c7bf5f..fa0da935dce 100644 --- a/src/megatron/bridge/models/conversion/modelopt_utils.py +++ b/src/megatron/bridge/models/conversion/modelopt_utils.py @@ -477,7 +477,7 @@ def _build_hf_modelopt_pre_ep_quant_metadata( def collect_modelopt_quant_metadata( - conversion_tasks: list[WeightConversionTask], + conversion_tasks: list[WeightConversionTask | None], ) -> dict[str, QuantMeta]: """Collect ModelOpt quantization metadata from conversion task modules.""" from modelopt.torch.export import quant_utils @@ -492,7 +492,7 @@ def collect_modelopt_quant_metadata( metadata: dict[str, QuantMeta] = {} quantizer_metadata: dict[tuple[int, int, str, int], QuantMeta | None] = {} for task in conversion_tasks: - if task.megatron_module is None or task.param_weight is None: + if task is None or task.megatron_module is None or task.param_weight is None: continue weight_quantizer, quant_module = find_modelopt_weight_quantizer_and_module( @@ -1018,7 +1018,7 @@ def export_and_finalize(hf_name: str, tensor: torch.Tensor) -> Iterable[HFWeight def build_modelopt_export_plan( - conversion_tasks: list[WeightConversionTask], + conversion_tasks: list[WeightConversionTask | None], *, model: list[torch.nn.Module], bridge: "MegatronModelBridge", @@ -1029,6 +1029,9 @@ def build_modelopt_export_plan( expected_qformat, export_weight = get_modelopt_quant_exporter(quant_mode) pg_collection = model_bridge_utils._get_pg_collection_from_model(model) local_metadata = collect_modelopt_quant_metadata(conversion_tasks) + # ``None`` slots represent global names without conversion tasks. They are + # not termination sentinels, so retain every concrete task after a hole. + concrete_tasks: list[WeightConversionTask] = [task for task in conversion_tasks if task is not None] if torch.distributed.is_initialized(): pp_group = model_bridge_utils._get_pp_group(model) ep_group = model_bridge_utils._get_ep_group(model) @@ -1049,7 +1052,7 @@ def build_modelopt_export_plan( candidate_mappings: dict[int, Any] = {} candidate_groups: dict[str, list[int]] = {} eligible_groups: dict[str, bool] = {} - for task_idx, task in enumerate(conversion_tasks): + for task_idx, task in enumerate(concrete_tasks): candidate = _modelopt_pre_ep_mapping(task.mapping, pg_collection) if can_export_before_ep else None if candidate is None: continue @@ -1078,7 +1081,7 @@ def build_modelopt_export_plan( if valid_expert_layout: try: for task_idx in task_indices: - candidate_task = conversion_tasks[task_idx] + candidate_task = concrete_tasks[task_idx] local_ids.append( extract_expert_number_from_param(candidate_task.param_name) % experts_per_rank ) @@ -1097,19 +1100,19 @@ def build_modelopt_export_plan( for group_key in all_group_keys } - mapped_tasks = list(conversion_tasks) + mapped_tasks = list(concrete_tasks) for group_key, task_indices in candidate_groups.items(): if not eligible_groups[group_key]: continue for task_idx in task_indices: mapped_tasks[task_idx] = replace( - conversion_tasks[task_idx], + concrete_tasks[task_idx], mapping=candidate_mappings[task_idx], ) regular_metadata_tasks: list[WeightConversionTask] = [] pre_ep_tasks: list[WeightConversionTask] = [] - for original_task, mapped_task in zip(conversion_tasks, mapped_tasks, strict=True): + for original_task, mapped_task in zip(concrete_tasks, mapped_tasks, strict=True): if getattr(mapped_task.mapping, "is_modelopt_pre_ep_export", False): pre_ep_tasks.append(mapped_task) else: diff --git a/tests/unit_tests/models/test_model_bridge.py b/tests/unit_tests/models/test_model_bridge.py index 9f598681c4a..669e9e37b56 100644 --- a/tests/unit_tests/models/test_model_bridge.py +++ b/tests/unit_tests/models/test_model_bridge.py @@ -18,7 +18,10 @@ import pytest import torch +from transformers import PretrainedConfig +from megatron.bridge.models.conversion import model_bridge as model_bridge_module +from megatron.bridge.models.conversion import modelopt_utils from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import HFWeightTuple, MegatronModelBridge, WeightConversionTask from megatron.bridge.models.conversion.param_mapping import AutoMapping @@ -32,6 +35,59 @@ def mapping_registry(self): # pragma: no cover - not used in tests return MegatronMappingRegistry() +def test_modelopt_plan_skips_unmapped_task_slot_and_keeps_later_task(monkeypatch): + first_name = "first.weight" + unmapped_name = "unmapped.amax" + last_name = "last.weight" + + class MappedBridge(DummyBridge): + def mapping_registry(self): + return MegatronMappingRegistry( + AutoMapping(first_name, "hf.first.weight"), + AutoMapping(last_name, "hf.last.weight"), + ) + + model = torch.nn.Module() + model.config = SimpleNamespace(share_embeddings_and_output_weights=False) + model.first = torch.nn.Linear(1, 1, bias=False) + model.unmapped = torch.nn.Module() + model.unmapped.register_buffer("amax", torch.ones(1)) + model.last = torch.nn.Linear(1, 1, bias=False) + bridge = MappedBridge() + global_names = [first_name, unmapped_name, last_name] + + monkeypatch.setattr(bridge, "_megatron_global_param_names_all_pp_ranks", lambda _model: global_names) + monkeypatch.setattr(bridge, "_share_embeddings_and_output_weights", lambda _config: False) + monkeypatch.setattr(model_bridge_module, "unwrap_model", lambda _model: [model]) + monkeypatch.setattr(model_bridge_module, "_get_pg_collection_from_model", lambda _model: None) + monkeypatch.setattr(model_bridge_module, "_get_pp_rank", lambda _model: 0) + monkeypatch.setattr( + model_bridge_module, + "_megatron_local_name_to_global", + lambda _models, _config, local_name, _vp_stage: local_name, + ) + + tasks = bridge.build_conversion_tasks(PretrainedConfig(), [model]) + + assert tasks[0].global_param_name == first_name + assert tasks[1] is None + assert tasks[2].global_param_name == last_name + + monkeypatch.setattr(modelopt_utils, "get_modelopt_quant_exporter", lambda _mode: ("unused", lambda *_args: ())) + monkeypatch.setattr(modelopt_utils, "get_pg_size", lambda _group: 1) + monkeypatch.setattr(modelopt_utils.model_bridge_utils, "_get_pg_collection_from_model", lambda _model: None) + + export_tasks = modelopt_utils.build_modelopt_export_plan( + tasks, + model=[model], + bridge=bridge, + quant_mode="nvfp4", + ignore_patterns=[], + ) + + assert [task.global_param_name for task in export_tasks] == [first_name, last_name] + + def test_hf_weight_tuple_iter_finalized_preserves_two_field_abi(): tensor = torch.ones(2) weight = HFWeightTuple("hf.weight", tensor) From 845bcf7822284dbe96e84dbbb8f5200fe12b770c Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:34:11 +0200 Subject: [PATCH 54/85] cp: `fix(ckpt): restrict dataloader state deserialization (5418)` into `r0.6.0` (#5437) Signed-off-by: Charlie Truong Signed-off-by: NeMo Bot Co-authored-by: Charlie Truong Co-authored-by: Claude Sonnet 4.6 (1M context) --- src/megatron/bridge/training/checkpointing.py | 3 +- src/megatron/bridge/utils/safe_pickle.py | 151 +++++++++++++++ .../unit_tests/training/test_checkpointing.py | 176 +++++++++++++++++- 3 files changed, 327 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 3daf8255782..781a06aafa8 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -90,6 +90,7 @@ ) from megatron.bridge.utils.import_utils import safe_import from megatron.bridge.utils.instantiate_utils import _validate_target_prefix +from megatron.bridge.utils.safe_pickle import energon_torch_load _, HAVE_RESIL = safe_import("nvidia_resiliency_ext.checkpointing") @@ -1818,7 +1819,7 @@ def maybe_load_dataloader_state( ) print_rank_0(f"restoring dataloader state at iteration {iteration} from {data_state_load_path}") - loaded = torch.load(data_state_load_path, map_location="cpu", weights_only=False) + loaded = energon_torch_load(data_state_load_path) iterable.restore_state(loaded["dataloader_state_dict"]) diff --git a/src/megatron/bridge/utils/safe_pickle.py b/src/megatron/bridge/utils/safe_pickle.py index 27d9398ad97..83805735b93 100644 --- a/src/megatron/bridge/utils/safe_pickle.py +++ b/src/megatron/bridge/utils/safe_pickle.py @@ -14,6 +14,7 @@ import io import pickle +import zipfile from types import MappingProxyType @@ -78,6 +79,8 @@ class _NumpyRestrictedUnpickler(pickle.Unpickler): "numpy.core.multiarray": frozenset({"_reconstruct", "scalar"}), # numpy ≥ 2.0 moved internals under ``numpy._core`` "numpy._core.multiarray": frozenset({"_reconstruct", "scalar"}), + # _codecs.encode is used by NumPy to encode raw array bytes into the pickle stream + "_codecs": frozenset({"encode"}), } ) @@ -90,6 +93,154 @@ def find_class(self, module: str, name: str) -> type: ) +class _EnergonUnpickler(_NumpyRestrictedUnpickler): + """Unpickler for Energon dataloader state files (``.pt``). + + Extends the NumPy-safe unpickler with the exact Energon dataclass types that Energon serialises + into dataloader checkpoint files. All other globals — including ``os``, ``subprocess``, and any + ``__reduce__`` payload callable outside this allowlist — are blocked, preventing arbitrary code + execution from attacker-controlled checkpoint files. + + Use via :func:`energon_torch_load` rather than instantiating directly. + """ + + _SAFE_MODULES: MappingProxyType = MappingProxyType( + { + **_NumpyRestrictedUnpickler._SAFE_MODULES, + # PyTorch tensor reconstruction — required for any .pt file containing tensors. + # These functions only rebuild tensor objects from pre-loaded storage; they do + # not execute arbitrary code. + "torch._utils": frozenset({"_rebuild_tensor_v2", "_rebuild_tensor"}), + # Energon dataloader state types — the explicit allowlist for this load site. + # If a real Energon checkpoint references a type not listed here the load will + # raise an UnpicklingError that names the missing ``module.name``; file a bug + # against Megatron Bridge so the allowlist can be extended. + "megatron.energon.state": frozenset({"FlexState"}), + "megatron.energon.rng": frozenset({"SystemRngState"}), + "megatron.energon.savable_loader": frozenset( + { + "SavableCheckpoint", + "SavableDataLoaderState", + "SavableDatasetCheckpoint", + "SavableDatasetState", + } + ), + "megatron.energon.flavors.webdataset.sample_loader": frozenset({"SliceState"}), + } + ) + + def find_class(self, module: str, name: str) -> type: + if module in self._SAFE_MODULES and name in self._SAFE_MODULES[module]: + return pickle.Unpickler.find_class(self, module, name) + raise pickle.UnpicklingError( + f"Restricted unpickler refused to load '{module}.{name}'. " + "This Energon checkpoint contains a type not in the dataloader-state allowlist. " + "Please file a bug against Megatron Bridge so it can be added." + ) + + +def energon_torch_load(path: str, *, map_location: str = "cpu") -> object: + """Load an Energon dataloader state ``.pt`` file through a restricted unpickler. + + Parses the torch zip format directly without calling ``torch.load``. Security is enforced + by :class:`_EnergonUnpickler`: any GLOBAL opcode whose ``(module, name)`` is not in the + explicit allowlist raises ``pickle.UnpicklingError``, blocking ``__reduce__``-based code + execution from attacker-controlled checkpoint files. + + ``torch.load(weights_only=True)`` is not used because PyTorch ≥ 2.13 restricts + SETITEM/SETITEMS to exact ``dict``, ``OrderedDict``, and ``Counter`` types, rejecting dict + subclasses such as Energon's ``FlexState`` — which is always present in real Energon + checkpoints (``SavableDatasetState.dataset_state`` is typed ``FlexState``, not Optional). + + ``torch.save`` writes a zip archive whose directory prefix is the file stem. This function + opens the zip, runs :class:`_EnergonUnpickler` on the pickle stream, and reconstructs + tensor storages from the raw blobs via ``persistent_load``. Storages are cached by key so + that tensors sharing a storage (views, slices) remain aliased after restore. + + Args: + path: Path to the ``.pt`` file written by + :func:`~megatron.bridge.training.checkpointing.maybe_save_dataloader_state`. + map_location: Device to map tensor storages to; defaults to ``"cpu"`` to avoid GPU + allocation during restore. + + Returns: + The deserialized object (a ``dict`` containing ``"dataloader_state_dict"``). + """ + import torch + + # PyTorch storage classes that torch.save may embed as GLOBAL opcodes in the persistent_id + # tuple. Enumerated explicitly rather than matched by suffix to avoid admitting future + # classes or monkey-patched attributes that happen to end with "Storage". + _TORCH_STORAGE_NAMES: frozenset[str] = frozenset( + { + "UntypedStorage", + "TypedStorage", + "FloatStorage", + "LongStorage", + "ByteStorage", + "HalfStorage", + "DoubleStorage", + "IntStorage", + "ShortStorage", + "CharStorage", + "BFloat16Storage", + "ComplexFloatStorage", + "ComplexDoubleStorage", + "QInt8Storage", + "QInt32Storage", + "QUInt8Storage", + "BoolStorage", + } + ) + + with zipfile.ZipFile(path) as zf: + names = zf.namelist() + + # torch.save uses the file stem as the archive prefix: "/data.pkl". + pkl_entry = next(n for n in names if n.endswith("/data.pkl")) + prefix = pkl_entry[: -len("/data.pkl")] + blob_prefix = f"{prefix}/data/" + + # Map storage key → raw bytes for every tensor storage blob. + blob_map = {n[len(blob_prefix) :]: zf.read(n) for n in names if n.startswith(blob_prefix)} + pkl_bytes = zf.read(pkl_entry) + + class _ZipLoader(_EnergonUnpickler): + def __init__(self, data: io.BytesIO) -> None: + super().__init__(data) + # Cache by storage key so tensors that share a storage remain aliased after restore, + # matching the behaviour of torch.load's own loaded_storages dict. + self._storage_cache: dict[str, object] = {} + + def find_class(self, module: str, name: str) -> type: + # torch.save embeds the storage class in the persistent_id tuple as a GLOBAL opcode. + # Storage classes hold raw bytes and are not executable; allow the known set here + # without adding them to the shared _SAFE_MODULES allowlist. + if module in ("torch", "torch.storage") and name in _TORCH_STORAGE_NAMES: + return pickle.Unpickler.find_class(self, module, name) + return super().find_class(module, name) + + def persistent_load(self, pid: tuple): + # torch.save persistent_id format (zip path): + # ('storage', storage_cls, key, location, nbytes) + _typename, storage_cls, key, _location, _nbytes = pid + key = key.decode() if isinstance(key, bytes) else key + if key in self._storage_cache: + return self._storage_cache[key] + raw = blob_map[key] + untyped = torch.frombuffer(bytearray(raw), dtype=torch.uint8).untyped_storage() + if map_location != _location: + untyped = untyped.to(torch.device(map_location)) + # _rebuild_tensor_v2 reads storage.dtype to determine the tensor dtype. + # UntypedStorage has no dtype attribute; wrap it with the typed storage class + # (e.g. torch.LongStorage) that torch.save recorded in the persistent_id. + result = untyped if storage_cls is torch.storage.UntypedStorage else storage_cls(wrap_storage=untyped) + self._storage_cache[key] = result + return result + + return _ZipLoader(io.BytesIO(pkl_bytes)).load() + + def safe_pickle_load(fp) -> object: """Deserialize from a file using a restricted unpickler that only allows safe types.""" return _RestrictedUnpickler(fp).load() diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 5ae9ed5c8de..7e88818781e 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -14,10 +14,12 @@ """Unit tests for megatron.bridge.training.checkpointing module.""" import os +import pickle import tempfile from pathlib import Path from unittest.mock import Mock, mock_open, patch +import numpy as np import pytest import torch from megatron.core.msc_utils import MultiStorageClientFeature @@ -73,6 +75,21 @@ class _DummyClass: _dummy_obj = _DummyClass() +def _write_dataloader_state_marker(path: str) -> None: + """Write a marker if an unsafe deserializer executes a test payload.""" + Path(path).write_text("dataloader state payload executed") + + +class _MaliciousDataloaderState: + """Pickle payload used to verify that dataloader restore blocks arbitrary globals.""" + + def __init__(self, marker_path: Path): + self.marker_path = marker_path + + def __reduce__(self): + return _write_dataloader_state_marker, (str(self.marker_path),) + + class TestCheckpointUtilities: """Test utility functions for checkpoint management.""" @@ -4797,7 +4814,7 @@ def test_existing_dir_missing_file_raises(self, tmp_path): with pytest.raises(RuntimeError, match="data-parallel size"): maybe_load_dataloader_state(train_iterator, 10, str(tmp_path), pg_collection=self._pg()) - @patch("megatron.bridge.training.checkpointing.torch.load") + @patch("megatron.bridge.training.checkpointing.energon_torch_load") def test_restores_on_every_cp_tp_pp_rank(self, mock_load, tmp_path): """Unlike save, restore must run on every cp/tp/pp because each rank pulls from its own data iterator. Keyed by the pure DP rank.""" @@ -4817,7 +4834,7 @@ def test_restores_on_every_cp_tp_pp_rank(self, mock_load, tmp_path): train_iterator.iterable.restore_state.assert_called_once_with({"dummy_energon_state": "xyz"}) - @patch("megatron.bridge.training.checkpointing.torch.load") + @patch("megatron.bridge.training.checkpointing.energon_torch_load") def test_restores_from_file(self, mock_load, tmp_path): """Happy path: the per-DP-rank file is loaded and restore_state is called with the saved dict.""" train_iterator = Mock() @@ -4831,6 +4848,161 @@ def test_restores_from_file(self, mock_load, tmp_path): train_iterator.iterable.restore_state.assert_called_once_with({"dummy_energon_state": "xyz"}) + def test_restores_energon_state_via_restricted_unpickler(self, tmp_path): + """Energon dataclass and NumPy state round-trips correctly through the restricted zip loader.""" + from megatron.energon.flavors.webdataset.sample_loader import SliceState + from megatron.energon.rng import SystemRngState + from megatron.energon.savable_loader import ( + SavableDataLoaderState, + SavableDatasetCheckpoint, + SavableDatasetState, + ) + from megatron.energon.state import FlexState + + train_iterator = Mock() + iter_dir = Path(get_checkpoint_name(str(tmp_path), 10)) + iter_dir.mkdir(parents=True) + state_path = iter_dir / "train_dataloader_dprank000.pt" + rng_state = SystemRngState( + torch=torch.arange(3), + numpy=np.arange(3, dtype=np.uint32), + random=(3, (1, 2, 3), None), + ) + loader_state = SavableDataLoaderState( + worker_states=[ + SavableDatasetCheckpoint( + state=SavableDatasetState( + rng=rng_state, + dataset_state=FlexState(active_slices=[SliceState(index=2, current=17)]), + sample_index=5, + ), + offset=1, + ) + ], + next_worker_id=0, + micro_batch_size=2, + ) + torch.save({"dataloader_state_dict": loader_state}, state_path) + + maybe_load_dataloader_state(train_iterator, 10, str(tmp_path), pg_collection=self._pg()) + + restored = train_iterator.iterable.restore_state.call_args.args[0] + assert isinstance(restored, SavableDataLoaderState) + assert restored.next_worker_id == loader_state.next_worker_id + assert restored.micro_batch_size == loader_state.micro_batch_size + restored_dataset_state = restored.worker_states[0].state + assert restored_dataset_state is not None + assert restored_dataset_state.sample_index == 5 + np.testing.assert_array_equal(restored_dataset_state.rng.numpy, rng_state.numpy) + assert restored_dataset_state.dataset_state["active_slices"] == [SliceState(index=2, current=17)] + + def test_rejects_reduce_payload_without_executing_it(self, tmp_path): + """A checkpoint cannot execute an arbitrary callable through pickle ``__reduce__``.""" + train_iterator = Mock() + iter_dir = Path(get_checkpoint_name(str(tmp_path), 10)) + iter_dir.mkdir(parents=True) + marker_path = tmp_path / "dataloader-state-payload-executed" + state_path = iter_dir / "train_dataloader_dprank000.pt" + torch.save({"dataloader_state_dict": _MaliciousDataloaderState(marker_path)}, state_path) + + assert not marker_path.exists() + with pytest.raises(pickle.UnpicklingError, match="Restricted unpickler refused to load"): + maybe_load_dataloader_state(train_iterator, 10, str(tmp_path), pg_collection=self._pg()) + + assert not marker_path.exists() + train_iterator.iterable.restore_state.assert_not_called() + + def test_restores_multi_worker_state_with_savable_checkpoint(self, tmp_path): + """Multi-worker state and SavableCheckpoint nested inside FlexState round-trip correctly. + + SavableCheckpoint (checkpoint_time + sample_index) is distinct from + SavableDatasetCheckpoint and can appear as a FlexState value when a + SavableDatasetWrapper is used inside a blending dataset. Two worker + states are used to confirm the loader handles the list correctly. + """ + from megatron.energon.flavors.webdataset.sample_loader import SliceState + from megatron.energon.rng import SystemRngState + from megatron.energon.savable_loader import ( + SavableCheckpoint, + SavableDataLoaderState, + SavableDatasetCheckpoint, + SavableDatasetState, + ) + from megatron.energon.state import FlexState + + def _rng(seed: int) -> SystemRngState: + return SystemRngState( + torch=torch.arange(seed, seed + 3), + numpy=np.arange(seed, seed + 3, dtype=np.uint32), + random=(seed, tuple(range(seed, seed + 3)), None), + ) + + # Worker 0: simple blending state with two active slices. + worker0 = SavableDatasetCheckpoint( + state=SavableDatasetState( + rng=_rng(0), + dataset_state=FlexState( + active_slices=[SliceState(index=0, current=11), SliceState(index=1, current=22)] + ), + sample_index=100, + ), + offset=0, + ) + # Worker 1: FlexState contains a nested SavableCheckpoint, simulating a + # SavableDatasetWrapper used as an inner dataset. + nested_ckpt = SavableCheckpoint( + state=SavableDatasetState( + rng=_rng(10), + dataset_state=FlexState(active_slices=[SliceState(index=5, current=99)]), + sample_index=50, + ), + checkpoint_time=1.0, + sample_index=50, + ) + worker1 = SavableDatasetCheckpoint( + state=SavableDatasetState( + rng=_rng(20), + dataset_state=FlexState(inner=nested_ckpt), + sample_index=200, + ), + offset=2, + ) + + loader_state = SavableDataLoaderState( + worker_states=[worker0, worker1], + next_worker_id=1, + micro_batch_size=8, + ) + + train_iterator = Mock() + iter_dir = Path(get_checkpoint_name(str(tmp_path), 10)) + iter_dir.mkdir(parents=True) + state_path = iter_dir / "train_dataloader_dprank000.pt" + torch.save({"dataloader_state_dict": loader_state}, state_path) + + maybe_load_dataloader_state(train_iterator, 10, str(tmp_path), pg_collection=self._pg()) + + restored = train_iterator.iterable.restore_state.call_args.args[0] + assert isinstance(restored, SavableDataLoaderState) + assert restored.next_worker_id == 1 + assert restored.micro_batch_size == 8 + assert len(restored.worker_states) == 2 + + r0 = restored.worker_states[0].state + assert r0.sample_index == 100 + assert r0.dataset_state["active_slices"] == [ + SliceState(index=0, current=11), + SliceState(index=1, current=22), + ] + np.testing.assert_array_equal(r0.rng.numpy, _rng(0).numpy) + + r1 = restored.worker_states[1].state + assert r1.sample_index == 200 + assert isinstance(r1.dataset_state["inner"], SavableCheckpoint) + assert r1.dataset_state["inner"].sample_index == 50 + assert r1.dataset_state["inner"].checkpoint_time == 1.0 + np.testing.assert_array_equal(r1.rng.numpy, _rng(20).numpy) + class TestMaybeSaveDataloaderState: """Tests for gating which rank writes the Energon dataloader stream-position state.""" From 02bd384ff69a8a86be7cbb91a7952f9588019aae Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:34:49 +0200 Subject: [PATCH 55/85] cp: `fix(recipes): use max-pool FSDP for Qwen3.5-VL (5402)` into `r0.6.0` (#5430) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py | 1 + tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py index 24c6ac1a580..0325632ef5d 100644 --- a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py +++ b/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py @@ -1011,6 +1011,7 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer: cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params" cfg.ddp.use_megatron_fsdp = True cfg.ddp.fsdp_double_buffer = True + cfg.ddp.megatron_fsdp_max_pool_double_buffer = True cfg.ddp.nccl_ub = False cfg.ddp.fsdp_db_use_persist_buf_on_alloc_fail = True cfg.ddp.overlap_grad_reduce = True diff --git a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py index a20fde30ddb..3982c201e8c 100644 --- a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py +++ b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py @@ -433,6 +433,7 @@ def test_qwen35_vl_35b_a3b_fsdp_sft_defaults(monkeypatch: pytest.MonkeyPatch): assert cfg.model.moe_router_fusion is True assert cfg.ddp.use_megatron_fsdp is True assert cfg.ddp.fsdp_double_buffer is True + assert cfg.ddp.megatron_fsdp_max_pool_double_buffer is True assert cfg.ddp.nccl_ub is False assert cfg.ddp.overlap_grad_reduce is True assert cfg.ddp.overlap_param_gather is True From cdafbebe2dd8cd3f38ce78367c76b8a993ba44b9 Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Sat, 8 Aug 2026 12:35:14 -0700 Subject: [PATCH 56/85] cp: `fix(recipe): verify GLM-5.2 RP2 pretraining and packed SFT (5344)` into `r0.6.0` (#5425) Signed-off-by: yaoyu-33 --- .../model_verification_cards/glm5-2/card.yaml | 188 +++++++++++------- src/megatron/bridge/data/builders/gpt_sft.py | 10 +- src/megatron/bridge/data/packing/config.py | 7 + src/megatron/bridge/data/packing/gpt_sft.py | 93 ++++----- src/megatron/bridge/recipes/glm/gb200/glm5.py | 72 +++++-- src/megatron/bridge/training/config.py | 29 +++ .../bridge/training/utils/packed_seq_utils.py | 18 +- .../data/builders/test_gpt_sft_builder.py | 21 ++ .../data/builders/test_gpt_sft_config.py | 24 ++- .../data/datasets/test_chat_template.py | 71 ++++++- tests/unit_tests/recipes/test_glm5_recipes.py | 71 +++++-- tests/unit_tests/training/test_config.py | 118 +++++++++++ tests/unit_tests/training/test_gpt_step.py | 20 +- .../training/utils/test_packed_seq_utils.py | 12 ++ 14 files changed, 568 insertions(+), 186 deletions(-) diff --git a/examples/model_verification_cards/glm5-2/card.yaml b/examples/model_verification_cards/glm5-2/card.yaml index 534828ad30c..a459be1b78d 100644 --- a/examples/model_verification_cards/glm5-2/card.yaml +++ b/examples/model_verification_cards/glm5-2/card.yaml @@ -7,13 +7,12 @@ summary: > timing and throughput metrics are sanity checks, not optimized performance results. GLM-5.2 H100 support verification currently covers exact GPU checkpoint round-trip on CPU and GPU, HF-to-Megatron forward equivalence, - bounded full SFT with export/reload - inference, 200K-context CP32 SFT, bounded PEFT, and bounded pretraining with - a reloadable middle checkpoint and verified continuation. GB200 support - verification currently covers bounded pretraining, bounded full SFT, 128K - packed long-context SFT, bounded PEFT, bounded checkpoint resume, and - deterministic legacy Megatron inference; post-SFT export/reload inference - remains pending. + bounded full SFT with export/reload inference, 200K-context CP32 SFT, + bounded PEFT, and bounded pretraining with a reloadable middle checkpoint + and verified continuation. GB200 support verification currently covers + bounded pretraining, bounded full SFT, 128K packed long-context SFT, bounded + PEFT, bounded checkpoint resume, and deterministic legacy Megatron + inference; post-SFT export/reload inference remains pending. verification_index: model_level: verified: @@ -149,6 +148,7 @@ items: completion was: "1. **Analyze the Request:** * **Topic:**" + It produced exactly 16 new tokens. pretrain: H100: @@ -176,50 +176,57 @@ items: initial_loss: 13.15907 final_loss: 6.385243 last_10_steps_step_time_ms_avg: 31204.920 - last_10_steps_model_tflops_per_gpu_avg: 50.620 + last_10_steps_model_tflops_per_gpu_avg: 54.595 expected_result: > The uninterrupted 352-H100 TP1/PP11/CP1/EP32/ETP1 run completed 100 real-data steps at GBS/MBS 1024/1 with cuDNN DSA, sparse indexer loss coefficient 0.001, MTP1, and flex/DeepEP dispatch. Loss remained finite from 13.15907 to 6.385243, final grad norm was 0.796, and every step recorded zero skipped and NaN iterations. The final 10 steps averaged - 31,204.920 ms and 50.620 TFLOPS/GPU; peak observed GPU memory was + 31,204.920 ms and 54.595 TFLOPS/GPU; peak observed GPU memory was 80,999 MiB. The matching step-50 and step-100 checkpoints contain model, optimizer, scheduler, data-order, and RNG state; the step-50 checkpoint reloads successfully. GB200: status: verified precision: bf16 + bridge_commit: e79f756196e0b97ee3ad94263a896c120b32c1e5 # pragma: allowlist secret enabled_features: {} command: > ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 --recipe glm52_pretrain_192gpu_gb200_bf16_config --mode pretrain --dataset megatron-indexed --deterministic --seq_length 4096 --max_steps 100 --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 - 'dataset.blend=[["work/data/glm5-2/pretrain-tulu3-text-document/tulu3_glm52_text_document_text_document"],null]' - dataset.path_to_cache=work/cache/glm5-2/pretrain-tulu3-gb200 - dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 - --save_dir work/model-verification/glm5-2/pretrain-gb200-cudnn/checkpoints + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/glm5-2/rp2-head-01-gb200 + dataset.random_seed=1234 dataset.num_workers=8 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer/tokenizer.model + tokenizer.use_tokenizer_vocab_size=false rng.seed=1234 + scheduler.lr_decay_iters=100 model.moe_router_force_load_balancing=false + checkpoint.load=null checkpoint.pretrained_checkpoint=null + checkpoint.save_optim=true checkpoint.save_rng=true + checkpoint.load_optim=true checkpoint.load_rng=true checkpoint.finetune=false + --save_dir work/model-verification/glm5-2/pretrain-rp2-gb200/checkpoints --save_interval 50 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-24 + last_verified: 2026-08-04 metrics: - initial_loss: 13.19136 - final_loss: 8.171059 - last_10_steps_step_time_ms_avg: 60610.850 - last_10_steps_model_tflops_per_gpu_avg: 101.310 + initial_loss: 13.168150 + final_loss: 6.523522 + last_10_steps_step_time_ms_avg: 64786.540 + last_10_steps_model_tflops_per_gpu_avg: 102.520 expected_result: > The 192-GB200 TP1/PP6/CP1/EP32/ETP1 run completed exactly 100 - real-data pretraining steps at GBS/MBS 1024/1 with cuDNN DSA, sparse + RP2 head_01 pretraining steps at GBS/MBS 1024/1 with cuDNN DSA, sparse indexer loss coefficient 0.001, MTP1, and all-to-all expert dispatch. - Loss remained finite from 13.19136 to 8.171059, final grad norm was - 1.394, and every step recorded zero skipped and NaN iterations. The - final 10 steps averaged 60,610.850 ms and 101.310 TFLOPS/GPU. Step-50 - and step-100 checkpoints were written with model, optimizer, scheduler, - data-order, and RNG state for resume verification. These losses are not - directly comparable with the H100 pretrain result: H100 used Wikitext - at sequence length 2,048, while GB200 used Tulu 3 at sequence length - 4,096. + The RP2 Llama SentencePiece tokenizer is used while retaining GLM-5.2's + configured vocabulary size, so this is support verification rather + than a cross-tokenizer convergence claim. Loss remained finite from + 13.168150 to 6.523522, final grad norm was 2.775, and every step + recorded zero skipped and NaN iterations. The final 10 steps averaged + 64,786.540 ms and 102.520 TFLOPS/GPU. Step-50 and step-100 checkpoints + contain model, optimizer, scheduler, data-order, and RNG state. sft: H100: @@ -246,7 +253,7 @@ items: initial_loss: 1.909875 final_loss: 0.6634203 last_10_steps_step_time_ms_avg: 4436.290 - last_10_steps_model_tflops_per_gpu_avg: 9.450 + last_10_steps_model_tflops_per_gpu_avg: 10.154 expected_result: > The immutable-revision 416-H100 run completes exactly 100 full-SFT steps at TP1/PP13/CP2/EP32/ETP1, GBS/MBS 32/1, with cuDNN DSA, sparse @@ -258,31 +265,44 @@ items: GB200: status: verified precision: bf16 + bridge_commit: cc190939aad0e1ee0edfb98c5455ae3fe833140b # pragma: allowlist secret enabled_features: sequence_packing: offline + context_parallel_size: 4 + moe_dispatcher: hybridep command: > ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 --recipe glm52_sft_192gpu_gb200_bf16_config --mode sft - --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 + --pretrained_checkpoint + work/model-verification/glm5-2/hf-4d67f66cc64d3219133b767c253b2ad1425c6c88 --max_steps 100 tokenizer.chat_template_path=work/templates/glm52_chat_template_generation.jinja --save_dir work/model-verification/glm5-2/sft-gb200-cudnn/checkpoints --save_interval 100 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-24 + last_verified: 2026-08-04 metrics: - initial_loss: 1.831100 - final_loss: 0.5953441 - last_10_steps_step_time_ms_avg: 68370.590 - last_10_steps_model_tflops_per_gpu_avg: 9.180 + initial_loss: 1.704979 + final_loss: 0.3682597 + last_10_steps_step_time_ms_avg: 8593.990 + last_10_steps_model_tflops_per_gpu_avg: 12.446 expected_result: > The immutable-revision 192-GB200 run completed exactly 100 full-SFT - steps at TP1/PP6/CP1/EP32/ETP1, GBS/MBS 32/1, with offline-packed - Tulu 3 data, cuDNN DSA, sparse indexer loss coefficient 0.001, MTP1, - and all-to-all expert dispatch. Loss remained finite from 1.831100 to - 0.5953441, final grad norm was 0.837, and every step recorded zero - skipped and NaN iterations. The step-100 checkpoint was written and is - available for a future GB200 export/reload gate. + steps at TP1/PP6/CP4/EP32/ETP1, DP8, and GBS/MBS 8/1 without gradient + accumulation. It used the pinned Tulu 3 train[:10000] split with the + GLM assistant-only chat template, 8,192-token offline packs, pad-8 + sequence alignment, an 8,184-token single-sequence cap, constant-shape + cu_seqlens, cuDNN DSA, sparse indexer loss coefficient 0.001, MTP1, + natural routing, and HybridEP. Packing is 99.72% efficient across 421 + rows with a 23.75 packing factor. The 100-step window contains + 6,553,600 token slots and 4,283,561 supervised tokens. Metadata and + Parquet SHA-256 values are cbd964c686b74f82804ceea2df75bb00fbe28648f0458a641b9db27977a55b1b + and b2c6f71a6b5add331750956622073cc8ced0bb6b2742e1d7a0e1a9f763cbaa0c. + Loss remained finite from 1.704979 to 0.3682597, final grad norm was + 0.873, and every step recorded zero skipped and NaN iterations. The + final 10 steps averaged 8,593.990 ms and 12.446 TFLOPS/GPU. A complete + step-100 distributed checkpoint was written for the future GB200 + export/reload gate. sft_export_inference: H100: @@ -342,43 +362,48 @@ items: initial_loss: 8.283127 final_loss: 2.557671 last_10_steps_step_time_ms_avg: 36603.370 - last_10_steps_model_tflops_per_gpu_avg: 210.650 + last_10_steps_model_tflops_per_gpu_avg: 39.648 expected_result: > The 608-H100 TP1/PP19/CP32/EP32/ETP1 run completed exactly 20 optimizer steps over 200,000-token packed rows with cuDNN DSA, sparse indexer loss coefficient 0.001, and MTP1. Loss remained finite from 8.283127 to 2.557671, final grad norm was 10838.703, and every step recorded zero skipped and NaN iterations. The last 10 steps averaged 36,603.370 ms - and 210.650 TFLOPS/GPU; peak observed GPU memory was 81,005 MiB. + and 39.648 TFLOPS/GPU; peak observed GPU memory was 81,005 MiB. GB200: status: verified precision: bf16 + bridge_commit: bdb637740154729a2409645703a0b3c78c38fa8c # pragma: allowlist secret enabled_features: sequence_packing: offline context_parallel_size: 32 + moe_dispatcher: hybridep command: > ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 --recipe glm52_sft_192gpu_gb200_bf16_128k_config --mode sft - --pretrained_checkpoint work/model-verification/glm5-2/gpu-megatron/iter_0000000 - --max_steps 20 --seq_length 131072 --lr 1e-6 --min_lr 0 --warmup_iters 2 + --pretrained_checkpoint + work/model-verification/glm5-2/hf-4d67f66cc64d3219133b767c253b2ad1425c6c88 + --max_steps 20 tokenizer.chat_template_path=work/templates/glm52_chat_template_generation.jinja - checkpoint.load=null ddp.check_for_nan_in_grad=true - ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true + dataset.dataset_root=work/data/glm5-2/synthetic-long-sft-128k + dataset.max_train_samples=1120 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-24 + last_verified: 2026-08-05 metrics: - initial_loss: 0.3235812 - final_loss: 0.1924341 - last_10_steps_step_time_ms_avg: 29269.730 - last_10_steps_model_tflops_per_gpu_avg: 236.510 + initial_loss: 0.4159648 + final_loss: 0.000511673 + last_10_steps_step_time_ms_avg: 143248.820 + last_10_steps_model_tflops_per_gpu_avg: 88.848 expected_result: > The immutable-revision 192-GB200 run completed exactly 20 optimizer - steps over 131,072-token packed rows at TP1/PP6/CP32/EP32/ETP1, - GBS/MBS 8/1, with cuDNN DSA, sparse indexer loss coefficient 0.001, - MTP1, padded cu_seqlens, and all-to-all expert dispatch. Loss remained - finite from 0.3235812 to 0.1924341, final grad norm was 1448.863, and - every step recorded zero skipped and NaN iterations. The final 10 steps - averaged 29,269.730 ms and 236.510 TFLOPS/GPU. + steps at TP1/PP6/CP32/EP32/ETP1, DP1, and GBS/MBS 56/1 with cuDNN DSA, + MTP1, HybridEP, and full activation recompute. The explicit PP6 decoder + layout [14, 16, 12, 12, 12, 12] completed without OOM using 32 + single-sequence, 131,072-token synthetic packs. Loss remained finite + from 0.4159648 to 0.000511673 with zero skipped or NaN iterations across + 146,786,080 supervised tokens. The final 10 steps averaged 143,248.820 + ms and 88.848 TFLOPS/GPU. This verifies long-context execution only; + the 8K pinned-Tulu run provides real-data convergence evidence. peft: H100: @@ -401,7 +426,7 @@ items: initial_loss: 1.909421 final_loss: 0.8835775 last_10_steps_step_time_ms_avg: 4638.550 - last_10_steps_model_tflops_per_gpu_avg: 17.240 + last_10_steps_model_tflops_per_gpu_avg: 19.423 expected_result: > The immutable-revision 208-H100 run completes exactly 100 LoRA steps at TP1/PP13/CP1/EP16/ETP1, GBS/MBS 32/1, with cuDNN DSA, sparse indexer @@ -429,7 +454,7 @@ items: initial_loss: 1.831100 final_loss: 0.8957523 last_10_steps_step_time_ms_avg: 68588.270 - last_10_steps_model_tflops_per_gpu_avg: 9.330 + last_10_steps_model_tflops_per_gpu_avg: 1.423 expected_result: > The immutable-revision 192-GB200 run completed exactly 100 LoRA steps at TP1/PP6/CP1/EP32/ETP1, GBS/MBS 32/1, with offline-packed Tulu 3 @@ -465,7 +490,7 @@ items: initial_loss: 7.219822 final_loss: 6.372624 last_10_steps_step_time_ms_avg: 32002.290 - last_10_steps_model_tflops_per_gpu_avg: 49.350 + last_10_steps_model_tflops_per_gpu_avg: 53.235 resume_comparison: reference_item: pretrain sentinel_steps: [51, 100] @@ -478,31 +503,39 @@ items: zero skipped or NaN iterations. Step 51 matches exactly at 7.219822. Step 100 is 6.372624 versus the uninterrupted reference's 6.385243, an absolute difference of 0.012619 within the allowed 0.063853. Final - grad norm is 2.169; the last 10 steps average 32,002.290 ms and 49.350 + grad norm is 2.169; the last 10 steps average 32,002.290 ms and 53.235 TFLOPS/GPU. Peak observed allocator memory is 65.230 GiB allocated and 69.684 GiB reserved. GB200: status: verified precision: bf16 depends_on: pretrain + bridge_commit: e79f756196e0b97ee3ad94263a896c120b32c1e5 # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 48 --gpus-per-node 4 --recipe glm52_pretrain_192gpu_gb200_bf16_config --mode pretrain --dataset megatron-indexed --deterministic --seq_length 4096 --max_steps 100 --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 - 'dataset.blend=[["work/data/glm5-2/pretrain-tulu3-text-document/tulu3_glm52_text_document_text_document"],null]' - dataset.path_to_cache=work/cache/glm5-2/pretrain-tulu3-gb200 - dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 - --load_dir work/model-verification/glm5-2/pretrain-gb200-cudnn/checkpoints - --save_dir work/model-verification/glm5-2/pretrain-resume-gb200-cudnn/checkpoints - --save_interval 0 checkpoint.ckpt_step=50 + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/glm5-2/rp2-head-01-gb200 + dataset.random_seed=1234 dataset.num_workers=8 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer/tokenizer.model + tokenizer.use_tokenizer_vocab_size=false rng.seed=1234 + scheduler.lr_decay_iters=100 model.moe_router_force_load_balancing=false + checkpoint.ckpt_step=50 checkpoint.pretrained_checkpoint=null + checkpoint.save_optim=true checkpoint.save_rng=true + checkpoint.load_optim=true checkpoint.load_rng=true checkpoint.finetune=false + --load_dir work/model-verification/glm5-2/pretrain-rp2-gb200/checkpoints + --save_dir work/model-verification/glm5-2/pretrain-rp2-resume-gb200/checkpoints + --save_interval 0 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-24 + last_verified: 2026-08-04 metrics: - initial_loss: 8.581669 - final_loss: 8.192354 - last_10_steps_step_time_ms_avg: 67700.250 - last_10_steps_model_tflops_per_gpu_avg: 90.710 + initial_loss: 7.233636 + final_loss: 6.503141 + last_10_steps_step_time_ms_avg: 60079.390 + last_10_steps_model_tflops_per_gpu_avg: 110.552 resume_comparison: reference_item: pretrain sentinel_steps: [51, 100] @@ -510,10 +543,11 @@ items: loss_absolute_tolerance: 1.0e-6 sentinels_match: true expected_result: > - The run reloads the GB200 reference's step-50 model, optimizer, + The run reloads the RP2 reference's step-50 model, optimizer, scheduler, data-order, and RNG state and completes steps 51-100 with - zero skipped or NaN iterations. Step 51 matches exactly at 8.581669. - Step 100 is 8.192354 versus the uninterrupted reference's 8.171059, an - absolute difference of 0.021295 within the allowed 0.081712. Final grad - norm is 1.840; the last 10 steps average 67,700.250 ms and 90.710 + zero skipped or NaN iterations. Step 51 is 7.233636 versus 7.233639 in + the uninterrupted reference, a relative difference of 4.1e-7. Step + 100 is 6.503141 versus 6.523522, a relative difference of 0.003124; + both satisfy the 1e-6 absolute plus 1% relative tolerance. Final grad + norm is 1.690; the last 10 steps average 60,079.390 ms and 110.552 TFLOPS/GPU. diff --git a/src/megatron/bridge/data/builders/gpt_sft.py b/src/megatron/bridge/data/builders/gpt_sft.py index 3a1e89a81fd..6f66f01b7c4 100644 --- a/src/megatron/bridge/data/builders/gpt_sft.py +++ b/src/megatron/bridge/data/builders/gpt_sft.py @@ -91,6 +91,11 @@ def _packing_fingerprint(config: "GPTSFTDatasetConfig", dataset_kwargs: dict[str "preprocessing": asdict(preprocessing), "dataset_kwargs": dataset_kwargs, } + if ( + config.offline_packing_specs is not None + and config.offline_packing_specs.max_single_sequence_length is not None + ): + packing_identity["max_single_sequence_length"] = config.offline_packing_specs.max_single_sequence_length return hashlib.sha256( json.dumps(packing_identity, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8") ).hexdigest()[:12] @@ -560,6 +565,9 @@ def __init__( self.packed_sequence_size = ( -1 if config.offline_packing_specs is None else config.offline_packing_specs.packed_sequence_size ) + self._max_single_sequence_length = ( + None if config.offline_packing_specs is None else config.offline_packing_specs.max_single_sequence_length + ) self.dataset_kwargs = normalize_gpt_sft_dataset_kwargs(config) self._pad_cu_seqlens = ( False if config.offline_packing_specs is None else config.offline_packing_specs.pad_cu_seqlens @@ -662,7 +670,7 @@ def _prepare_packed_split( output_metadata_path=self.pack_metadata, packed_sequence_size=self.packed_sequence_size, tokenizer=self.tokenizer, - max_seq_length=self.seq_length, + max_seq_length=self._max_single_sequence_length or self.seq_length, seed=self.seed, dataset_kwargs=self.dataset_kwargs, pad_seq_to_mult=self._pad_seq_to_mult, diff --git a/src/megatron/bridge/data/packing/config.py b/src/megatron/bridge/data/packing/config.py index c1150264be6..8fefec5a416 100644 --- a/src/megatron/bridge/data/packing/config.py +++ b/src/megatron/bridge/data/packing/config.py @@ -28,12 +28,14 @@ class PackedSequenceSpecs: """Settings and optional artifact paths for offline sequence packing.""" packed_sequence_size: int = -1 + max_single_sequence_length: int | None = None tokenizer_model_name: str | None = None num_tokenizer_workers: int = -1 packed_train_data_path: str | Path | None = None packed_val_data_path: str | Path | None = None packed_metadata_path: str | Path | None = None pad_cu_seqlens: bool = False + """Pad cumulative sequence boundaries for full-iteration, whole-layer, or attention-scoped CUDA graphs.""" pad_seq_to_mult: int | None = 1 def __post_init__(self) -> None: @@ -44,6 +46,11 @@ def __post_init__(self) -> None: self._validate_packed_path("packed_val_data_path", self.packed_val_data_path) if self.pad_seq_to_mult is not None and self.pad_seq_to_mult <= 0: raise ValueError("pad_seq_to_mult must be a positive integer when provided.") + if self.max_single_sequence_length is not None: + if self.max_single_sequence_length <= 0: + raise ValueError("max_single_sequence_length must be a positive integer when provided.") + if self.packed_sequence_size > 0 and self.max_single_sequence_length > self.packed_sequence_size: + raise ValueError("max_single_sequence_length cannot exceed packed_sequence_size.") def _validate_packed_path(self, attr_name: str, path_value: str | Path) -> None: """Validate an explicitly supplied packed artifact path.""" diff --git a/src/megatron/bridge/data/packing/gpt_sft.py b/src/megatron/bridge/data/packing/gpt_sft.py index 35495fb17a3..0aed186d8ad 100644 --- a/src/megatron/bridge/data/packing/gpt_sft.py +++ b/src/megatron/bridge/data/packing/gpt_sft.py @@ -61,8 +61,9 @@ def __init__( return_cu_seqlen: Whether to return `cu_seqlen` to pass to the model. Having `cu_seqlen` in the model input enables THD attention kernel, which is the correct format for training with packed sequence to prevent cross-sequence attention. This flag should be True unless you have a specific use case. - pad_seq_to_mult: The multiple used for padding sequences during packing. When > 1, cu_seqlens_unpadded - will be computed to support THD CP. When == 1 (no padding), cu_seqlens_unpadded is not computed. + pad_seq_to_mult: The multiple used for padding sequences during packing. When > 1, both logical and + physically padded cumulative sequence offsets are returned for THD CP. + pack_metadata_file_path: Path to the metadata JSON file used to stabilize cu_seqlens shapes. """ np.random.seed(kwargs.get("seed", 1234)) super().__init__(file_path, tokenizer, **kwargs) @@ -161,8 +162,8 @@ def collate_fn(self, batch): Returns: dict: A dictionary of batched tensors, including 'tokens', 'labels', - 'loss_mask', 'position_ids', and potentially 'cu_seqlens', - 'cu_seqlens_argmin', 'max_seqlen' if `return_cu_seqlen` is True. + 'loss_mask', 'position_ids', and the MCore/TE packed-sequence + fields if `return_cu_seqlen` is True. """ input_ids = [ np.concatenate( @@ -198,49 +199,60 @@ def collate_fn(self, batch): assert max_length <= self.max_seq_length position_ids: list[list[int]] = [] + # TE naming: cu_seqlens contains logical offsets, while cu_seqlens_padded + # contains physical offsets including per-sequence alignment padding. cu_seqlens: list[list[int]] = [] - # Only compute cu_seqlens_unpadded when pad_seq_to_mult > 1 (actual padding for CP) - cu_seqlens_unpadded: list[list[int]] | None = [] if self._pad_seq_to_mult > 1 else None + cu_seqlens_padded: list[list[int]] | None = [] if self._pad_seq_to_mult > 1 else None for item in batch: position_ids.append([]) cu_seqlens.append([0]) - if cu_seqlens_unpadded is not None: - cu_seqlens_unpadded.append([0]) + if cu_seqlens_padded is not None: + cu_seqlens_padded.append([0]) seqlens = np.array(item["seq_boundaries"][1:]) - np.array(item["seq_boundaries"][:-1]) for length in seqlens: # length minus 1 because input_ids is truncated by 1 for labels position_ids[-1].extend(list(range(length - 1))) - cu_seqlens[-1].append(cu_seqlens[-1][-1] + length - 1) + if cu_seqlens_padded is not None: + cu_seqlens_padded[-1].append(cu_seqlens_padded[-1][-1] + length - 1) + else: + cu_seqlens[-1].append(cu_seqlens[-1][-1] + length - 1) # the last seq needs to be the max seq len because rope and attn kernels expect no padding - assert cu_seqlens[-1][-1] <= max_length - - # since data is prepadded when cp_size > 1, there may be some extra padding at the end - # of the packed sequence. In this case, we need to add the max seq len to the end. - if cu_seqlens[-1][-1] != max_length: - cu_seqlens[-1].append(max_length) - - if cu_seqlens_unpadded is not None: + if cu_seqlens_padded is not None: + assert cu_seqlens_padded[-1][-1] <= max_length + if cu_seqlens_padded[-1][-1] != max_length: + cu_seqlens_padded[-1].append(max_length) for i in range(len(item["seq_boundaries"]) - 1): current_seq = item["input_ids"][item["seq_boundaries"][i] : item["seq_boundaries"][i + 1] - 1] - # Stop unpadded lengths at the last non-eos token so padding eos are excluded. + # Stop logical lengths at the last non-EOS token so alignment padding is excluded. current_seq_arr = np.array(current_seq) non_eos_positions = np.where(current_seq_arr != self.tokenizer.eos_id)[0] - seqlen_unpadded = non_eos_positions[-1] + 1 if non_eos_positions.size > 0 else 0 - cu_seqlens_unpadded[-1].append(cu_seqlens_unpadded[-1][-1] + seqlen_unpadded) + logical_seqlen = non_eos_positions[-1] + 1 if non_eos_positions.size > 0 else 0 + cu_seqlens[-1].append(cu_seqlens[-1][-1] + logical_seqlen) # if extra paddings are added in the packed sequence, they can't be counted as # actual tokens for training - if len(cu_seqlens[-1]) > len(cu_seqlens_unpadded[-1]): - cu_seqlens_unpadded[-1].append(cu_seqlens_unpadded[-1][-1]) + if len(cu_seqlens_padded[-1]) > len(cu_seqlens[-1]): + cu_seqlens[-1].append(cu_seqlens[-1][-1]) + else: + assert cu_seqlens[-1][-1] <= max_length + # Prepadded data may leave padding at the end of the packed sequence. + if cu_seqlens[-1][-1] != max_length: + cu_seqlens[-1].append(max_length) if self.pad_cu_seqlens: - # pad cu_seqlens to a constant shape with zero length sequences + # Pad both boundary representations to the same static shape with zero-length + # sequences. TE expects corresponding logical and physical boundary arrays. max_samples_per_bin = max(p["max_samples_per_bin"] for p in self.pack_metadata) # plus 2 since cu_seqlens additionally contains 0 and may append max_length - pad_num = max_samples_per_bin - len(cu_seqlens[-1]) + 2 - cu_seqlens[-1].extend([max_length] * pad_num) + if cu_seqlens_padded is not None: + pad_num = max_samples_per_bin - len(cu_seqlens_padded[-1]) + 2 + cu_seqlens_padded[-1].extend([max_length] * pad_num) + cu_seqlens[-1].extend([cu_seqlens[-1][-1]] * pad_num) + else: + pad_num = max_samples_per_bin - len(cu_seqlens[-1]) + 2 + cu_seqlens[-1].extend([max_length] * pad_num) assert len(input_ids[0]) == len(position_ids[0]), ( "Dataset problem: input_ids and position_ids lengths don't match" @@ -265,13 +277,10 @@ def collate_fn(self, batch): } if self.return_cu_seqlen: - cu_seqlens = self._collate_item( - cu_seqlens, max_length=max(len(length) for length in cu_seqlens) + 1, pad_id=-1 - ) - # Pre-generate `cu_seqlens_argmin` and `max_seqlen` as CPU tensor to avoid device-to-host copies. - cu_seqlens = torch.IntTensor(cu_seqlens) - cu_seqlens_argmin = torch.argmin(cu_seqlens, dim=1, keepdim=True) - seqlens = cu_seqlens[:, 1:] - cu_seqlens[:, :-1] + # Offline packed SFT requires micro-batch size 1, so no sentinel padding is + # needed to batch variable-length metadata. Emit MCore/TE field names directly. + seqlens = torch.IntTensor(cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens) + seqlens = seqlens[:, 1:] - seqlens[:, :-1] max_seqlen, _ = seqlens.max(dim=1, keepdim=True) if self.pad_cu_seqlens: @@ -286,25 +295,19 @@ def collate_fn(self, batch): safe_max_seqlen = max(dataset_max_seqlen, padding_gap) max_seqlen = torch.IntTensor([safe_max_seqlen] * len(cu_seqlens)) else: - seqlens = cu_seqlens[:, 1:] - cu_seqlens[:, :-1] max_seqlen, _ = seqlens.max(dim=1, keepdim=True) cu_seqlens_batch = { "attention_mask": None, # no attention mask is needed for packed seq - "cu_seqlens": torch.IntTensor(cu_seqlens), # cu_seqlens_q must be in dtype torch.int32 - "cu_seqlens_argmin": cu_seqlens_argmin, # only required for perf - "max_seqlen": max_seqlen, # only required for perf + "cu_seqlens_q": torch.IntTensor(cu_seqlens), + "max_seqlen_q": max_seqlen, + "max_seqlen_kv": max_seqlen, } + cu_seqlens_batch["cu_seqlens_kv"] = cu_seqlens_batch["cu_seqlens_q"] - # Only include cu_seqlens_unpadded when pad_seq_to_mult > 1 (actual CP padding) - if cu_seqlens_unpadded is not None: - cu_seqlens_unpadded = self._collate_item( - cu_seqlens_unpadded, max_length=max(len(length) for length in cu_seqlens_unpadded) + 1, pad_id=-1 - ) - cu_seqlens_unpadded = torch.IntTensor(cu_seqlens_unpadded) - cu_seqlens_unpadded_argmin = torch.argmin(cu_seqlens_unpadded, dim=1, keepdim=True) - cu_seqlens_batch["cu_seqlens_unpadded"] = cu_seqlens_unpadded - cu_seqlens_batch["cu_seqlens_unpadded_argmin"] = cu_seqlens_unpadded_argmin + if cu_seqlens_padded is not None: + cu_seqlens_batch["cu_seqlens_q_padded"] = torch.IntTensor(cu_seqlens_padded) + cu_seqlens_batch["cu_seqlens_kv_padded"] = cu_seqlens_batch["cu_seqlens_q_padded"] processed_batch.update(cu_seqlens_batch) else: diff --git a/src/megatron/bridge/recipes/glm/gb200/glm5.py b/src/megatron/bridge/recipes/glm/gb200/glm5.py index 5250452f199..9c071b6d7e7 100644 --- a/src/megatron/bridge/recipes/glm/gb200/glm5.py +++ b/src/megatron/bridge/recipes/glm/gb200/glm5.py @@ -29,6 +29,7 @@ _GLM52_MODEL_ID = "zai-org/GLM-5.2" _GLM52_MODEL_REVISION = "4d67f66cc64d3219133b767c253b2ad1425c6c88" # pragma: allowlist secret _TULU3_REVISION = "b14afda60f1bbebe55d5d2fa1e4df5042f97f8be" # pragma: allowlist secret +_GLM52_PP6_128K_LAYOUT = "|".join(("E" + "t" * 14, "t" * 16, "t" * 12, "t" * 12, "t" * 12, "t" * 12 + "mL")) def glm52_pretrain_192gpu_gb200_bf16_config() -> ConfigContainer: @@ -118,13 +119,12 @@ def glm52_sft_192gpu_gb200_bf16_config() -> ConfigContainer: ) cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} - - cfg.model.seq_length = 2048 + cfg.model.seq_length = 8192 cfg.model.tensor_model_parallel_size = 1 cfg.model.pipeline_model_parallel_size = 6 cfg.model.virtual_pipeline_model_parallel_size = None cfg.model.pipeline_model_parallel_layout = None - cfg.model.context_parallel_size = 1 + cfg.model.context_parallel_size = 4 cfg.model.expert_model_parallel_size = 32 cfg.model.expert_tensor_parallel_size = 1 cfg.model.sequence_parallel = False @@ -140,10 +140,13 @@ def glm52_sft_192gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.cross_entropy_loss_fusion = False cfg.model.cross_entropy_fusion_impl = "native" cfg.model.gradient_accumulation_fusion = True - cfg.model.moe_token_dispatcher_type = "alltoall" - cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.moe_shared_expert_overlap = False cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_permute_fusion = True + cfg.model.moe_grouped_gemm = True + cfg.model.deallocate_pipeline_outputs = True cfg.model.persist_layer_norm = True cfg.model.bias_dropout_fusion = True cfg.model.bias_activation_fusion = True @@ -165,24 +168,28 @@ def glm52_sft_192gpu_gb200_bf16_config() -> ConfigContainer: cfg.validation.eval_interval = 0 cfg.logger.log_interval = 1 + packing_alignment = 2 * cfg.model.context_parallel_size cfg.dataset = default_tulu3_config( - seq_length=2048, + seq_length=cfg.model.seq_length, enable_offline_packing=True, - pad_seq_to_mult=32, + pad_seq_to_mult=packing_alignment, ) cfg.dataset.hf_dataset.split = "train[:10000]" cfg.dataset.hf_dataset.load_kwargs = {"revision": _TULU3_REVISION} - cfg.dataset.hf_output_root = "work/data/glm5-2/tulu3-full-sft-gb200" + cfg.dataset.hf_output_root = "work/data/glm5-2/tulu3-full-sft-gb200-8k-v5" cfg.dataset.hf_rewrite = False cfg.dataset.hf_validation_proportion = None cfg.dataset.max_train_samples = 10000 cfg.dataset.seed = 1234 cfg.dataset.do_validation = False cfg.dataset.do_test = False + cfg.dataset.offline_packing_specs.max_single_sequence_length = cfg.model.seq_length - packing_alignment + # HybridEP needs a fixed token width; CUDA graphs are disabled, so cu_seqlens can remain dynamic. + cfg.dataset.dataset_kwargs = {"pad_to_max_length": True} cfg.rng.seed = 5678 cfg.train.train_iters = 100 - cfg.train.global_batch_size = 32 + cfg.train.global_batch_size = 8 cfg.optimizer, cfg.scheduler = distributed_fused_adam_with_cosine_annealing( lr_warmup_iters=10, max_lr=5e-6, @@ -193,7 +200,19 @@ def glm52_sft_192gpu_gb200_bf16_config() -> ConfigContainer: cfg.checkpoint.load = None cfg.checkpoint.save_optim = False cfg.checkpoint.save_rng = False - cfg.env_vars = {**COMMON_RECIPE_ENV_VARS} + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NCCL_GRAPH_REGISTER": 0, + "NCCL_NVLS_ENABLE": 0, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + "USE_MNNVL": 1, + } return cfg @@ -206,18 +225,17 @@ def glm52_sft_192gpu_gb200_bf16_128k_config() -> ConfigContainer: ) cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} - cfg.model.seq_length = 131072 cfg.model.tensor_model_parallel_size = 1 cfg.model.pipeline_model_parallel_size = 6 cfg.model.virtual_pipeline_model_parallel_size = None - cfg.model.pipeline_model_parallel_layout = None + cfg.model.pipeline_model_parallel_layout = _GLM52_PP6_128K_LAYOUT cfg.model.context_parallel_size = 32 cfg.model.expert_model_parallel_size = 32 cfg.model.expert_tensor_parallel_size = 1 cfg.model.sequence_parallel = False - cfg.model.num_layers_in_first_pipeline_stage = 14 - cfg.model.num_layers_in_last_pipeline_stage = 16 + cfg.model.num_layers_in_first_pipeline_stage = None + cfg.model.num_layers_in_last_pipeline_stage = None cfg.model.account_for_embedding_in_pipeline_split = False cfg.model.account_for_loss_in_pipeline_split = False cfg.model.microbatch_group_size_per_vp_stage = None @@ -232,6 +250,9 @@ def glm52_sft_192gpu_gb200_bf16_128k_config() -> ConfigContainer: cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.moe_shared_expert_overlap = False cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_permute_fusion = True + cfg.model.moe_grouped_gemm = True + cfg.model.deallocate_pipeline_outputs = True cfg.model.persist_layer_norm = True cfg.model.bias_dropout_fusion = True cfg.model.bias_activation_fusion = True @@ -250,24 +271,28 @@ def glm52_sft_192gpu_gb200_bf16_128k_config() -> ConfigContainer: cfg.train.manual_gc = True cfg.train.manual_gc_interval = 10 cfg.train.train_iters = 20 - cfg.train.global_batch_size = 8 + cfg.train.global_batch_size = 56 cfg.validation.eval_iters = 0 cfg.validation.eval_interval = 0 cfg.logger.log_interval = 1 - cfg.dataset.seq_length = 131072 + cfg.dataset = default_tulu3_config( + seq_length=131072, + enable_offline_packing=True, + pad_seq_to_mult=64, + ) cfg.dataset.hf_dataset = None - cfg.dataset.dataset_root = "work/data/glm5-2/synthetic-long-context-gb200" + cfg.dataset.dataset_root = "work/data/glm5-2/synthetic-long-sft-128k" cfg.dataset.hf_output_root = None + cfg.dataset.hf_rewrite = False cfg.dataset.hf_validation_proportion = None + cfg.dataset.max_train_samples = 1120 cfg.dataset.seed = 1234 cfg.dataset.preprocessing = ChatSFTPreprocessingConfig() cfg.dataset.do_validation = False cfg.dataset.do_test = False - cfg.dataset.offline_packing_specs.packed_sequence_size = 131072 cfg.dataset.offline_packing_specs.tokenizer_model_name = "glm5" - cfg.dataset.offline_packing_specs.pad_seq_to_mult = 64 - cfg.dataset.offline_packing_specs.pad_cu_seqlens = True + # HybridEP needs a fixed token width; CUDA graphs are disabled, so cu_seqlens can remain dynamic. cfg.dataset.dataset_kwargs = {"pad_to_max_length": True} cfg.rng.seed = 5678 @@ -281,9 +306,15 @@ def glm52_sft_192gpu_gb200_bf16_128k_config() -> ConfigContainer: cfg.checkpoint.load = None cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NCCL_GRAPH_REGISTER": 0, + "NCCL_NVLS_ENABLE": 0, "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, "NVLINK_DOMAIN_SIZE": 72, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "USE_MNNVL": 1, } return cfg @@ -298,7 +329,6 @@ def glm52_peft_192gpu_gb200_bf16_config(peft_scheme: str | PEFT = "lora") -> Con ) cfg.tokenizer.tokenizer_model = _GLM52_MODEL_ID cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _GLM52_MODEL_REVISION} - cfg.model.seq_length = 2048 cfg.model.tensor_model_parallel_size = 1 cfg.model.pipeline_model_parallel_size = 6 diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 03c470dff29..9fa1f89e471 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -79,6 +79,7 @@ warn_rank_0, ) from megatron.bridge.utils.cuda_graph import ( + cuda_graph_module_names, is_full_iteration_cuda_graph, validate_cuda_graph_configuration, ) @@ -1156,6 +1157,34 @@ def validate(self) -> None: if offline_packing_specs is not None and not enable_offline_packing: raise ValueError("enable_offline_packing must be True when offline_packing_specs is set.") + if offline_packing_specs is not None: + pad_cu_seqlens = offline_packing_specs.pad_cu_seqlens + dataset_kwargs = getattr(self.dataset, "dataset_kwargs", None) or {} + pad_to_max_length = ( + getattr(self.dataset, "pad_to_max_length", False) is True + or dataset_kwargs.get("pad_to_max_length", False) is True + ) + if pad_cu_seqlens and not pad_to_max_length: + raise ValueError("offline_packing_specs.pad_cu_seqlens=True requires dataset pad_to_max_length=True.") + + cuda_graph_impl = getattr(self.model, "cuda_graph_impl", "none") + cuda_graph_modules = cuda_graph_module_names(self.model) + is_full_iteration_graph = is_full_iteration_cuda_graph(self.model) + # Every training graph needs a static token width. Only graphs that include + # packed attention also consume cu_seqlens and require static boundary shapes. + uses_training_cuda_graphs = is_full_iteration_graph or cuda_graph_impl == "transformer_engine" + if uses_training_cuda_graphs and not pad_to_max_length: + raise ValueError("Offline packing with CUDA graphs requires dataset pad_to_max_length=True.") + + captures_packed_attention = is_full_iteration_graph or ( + cuda_graph_impl == "transformer_engine" and (not cuda_graph_modules or "attn" in cuda_graph_modules) + ) + if captures_packed_attention and not pad_cu_seqlens: + raise ValueError( + "Packed attention CUDA graphs require offline_packing_specs.pad_cu_seqlens=True. " + "Enable it for full-iteration, whole-layer, or attention-scoped capture." + ) + # Validate declarative SFT values before deriving runtime padding # multiples so normalization cannot hide an invalid user value. if isinstance( diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py index 25323af3773..acea8a31fe8 100644 --- a/src/megatron/bridge/training/utils/packed_seq_utils.py +++ b/src/megatron/bridge/training/utils/packed_seq_utils.py @@ -232,6 +232,18 @@ def _squeeze_metadata(value: PackedMetadataValue) -> PackedMetadataValue: return value.squeeze() +def _as_python_int(value: PackedMetadataValue, *, field_name: str) -> int | None: + """Normalize scalar packed-sequence metadata to MCore's integer contract.""" + value = _squeeze_metadata(value) + if value is None: + return None + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError(f"{field_name} must contain exactly one value, got shape {tuple(value.shape)}") + return int(value.item()) + return int(value) + + def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqParams: """Build packed sequence parameters from a batch dictionary. @@ -254,8 +266,8 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar if "cu_seqlens_q" in batch: cu_seqlens_q = _squeeze_metadata(batch["cu_seqlens_q"]) cu_seqlens_kv = _squeeze_metadata(batch.get("cu_seqlens_kv")) - max_seqlen_q = _squeeze_metadata(batch.get("max_seqlen_q")) - max_seqlen_kv = _squeeze_metadata(batch.get("max_seqlen_kv")) + max_seqlen_q = _as_python_int(batch.get("max_seqlen_q"), field_name="max_seqlen_q") + max_seqlen_kv = _as_python_int(batch.get("max_seqlen_kv"), field_name="max_seqlen_kv") return PackedSeqParams( cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv if cu_seqlens_kv is not None else cu_seqlens_q, @@ -296,7 +308,7 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar else: cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded)] - max_seqlen = batch["max_seqlen"].squeeze() if "max_seqlen" in batch else None + max_seqlen = _as_python_int(batch.get("max_seqlen"), field_name="max_seqlen") total_tokens = batch.get("total_tokens") # When cu_seqlens_unpadded is present (pad_seq_to_mult > 1), pass both unpadded and padded diff --git a/tests/unit_tests/data/builders/test_gpt_sft_builder.py b/tests/unit_tests/data/builders/test_gpt_sft_builder.py index bc1f9afcb21..087d95526b8 100644 --- a/tests/unit_tests/data/builders/test_gpt_sft_builder.py +++ b/tests/unit_tests/data/builders/test_gpt_sft_builder.py @@ -84,6 +84,27 @@ def test_default_pack_path_fingerprints_preprocessing(tmp_path): assert prompt_builder.default_pack_path != chat_builder.default_pack_path +def test_default_pack_path_fingerprints_max_single_sequence_length(tmp_path): + """Different single-sequence caps must not reuse the same packed artifact.""" + + def build(max_single_sequence_length: int) -> GPTSFTDatasetBuilder: + return GPTSFTDatasetBuilder( + config=GPTSFTDatasetConfig( + dataset_root=tmp_path, + seq_length=128, + enable_offline_packing=True, + offline_packing_specs=PackedSequenceSpecs( + packed_sequence_size=128, + max_single_sequence_length=max_single_sequence_length, + tokenizer_model_name="mock-tokenizer", + ), + ), + tokenizer=MagicMock(), + ) + + assert build(120).default_pack_path != build(112).default_pack_path + + def test_default_pack_path_is_stable_for_equivalent_non_hf_tokenizers(tmp_path): def build() -> GPTSFTDatasetBuilder: tokenizer = build_tokenizer(TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=128)) diff --git a/tests/unit_tests/data/builders/test_gpt_sft_config.py b/tests/unit_tests/data/builders/test_gpt_sft_config.py index cbdb824a297..b4399f48ac2 100644 --- a/tests/unit_tests/data/builders/test_gpt_sft_config.py +++ b/tests/unit_tests/data/builders/test_gpt_sft_config.py @@ -47,7 +47,11 @@ def _hf_config(tmp_path, **source_overrides): def test_config_round_trip_is_declarative_and_serializable(tmp_path): - specs = PackedSequenceSpecs(packed_sequence_size=128, pad_seq_to_mult=8) + specs = PackedSequenceSpecs( + packed_sequence_size=128, + max_single_sequence_length=120, + pad_seq_to_mult=8, + ) config = GPTSFTDatasetConfig( seq_length=128, hf_dataset=HFDatasetSourceConfig(dataset_name="squad"), @@ -67,10 +71,21 @@ def test_config_round_trip_is_declarative_and_serializable(tmp_path): assert isinstance(restored.hf_dataset, HFDatasetSourceConfig) assert restored.hf_dataset.dataset_name == "squad" assert restored.offline_packing_specs.packed_sequence_size == 128 + assert restored.offline_packing_specs.max_single_sequence_length == 120 assert isinstance(restored.preprocessing, PromptCompletionSFTPreprocessingConfig) assert "tokenizer" not in serialized +@pytest.mark.parametrize("max_single_sequence_length", [0, 129]) +def test_packed_specs_reject_invalid_max_single_sequence_length(max_single_sequence_length): + """The per-sequence cap must be positive and fit within the pack.""" + with pytest.raises(ValueError, match="max_single_sequence_length"): + PackedSequenceSpecs( + packed_sequence_size=128, + max_single_sequence_length=max_single_sequence_length, + ) + + @pytest.mark.parametrize( ("dataset_root", "hf_dataset"), [ @@ -556,7 +571,11 @@ def _build(path, **kwargs): def test_hf_rewrite_regenerates_existing_builder_managed_packed_data(monkeypatch, tmp_path): - specs = PackedSequenceSpecs(packed_sequence_size=128, tokenizer_model_name="test-tokenizer") + specs = PackedSequenceSpecs( + packed_sequence_size=128, + max_single_sequence_length=120, + tokenizer_model_name="test-tokenizer", + ) config = GPTSFTDatasetConfig( seq_length=128, hf_dataset=HFDatasetSourceConfig(dataset_name="squad"), @@ -582,6 +601,7 @@ def test_hf_rewrite_regenerates_existing_builder_managed_packed_data(monkeypatch builder.prepare_data() assert [call["input_path"].name for call in pack_calls] == ["training.jsonl", "validation.jsonl"] + assert all(call["max_seq_length"] == 120 for call in pack_calls) assert json.loads(builder.pack_metadata.read_text()) == [] diff --git a/tests/unit_tests/data/datasets/test_chat_template.py b/tests/unit_tests/data/datasets/test_chat_template.py index bc1994298c7..bd9e3707857 100644 --- a/tests/unit_tests/data/datasets/test_chat_template.py +++ b/tests/unit_tests/data/datasets/test_chat_template.py @@ -1001,7 +1001,7 @@ def _create_minimal_packed_dataset(tokenizer_eos_id: int = 0): dataset.pad_to_max_length = False dataset.max_seq_length = 32 dataset.pad_seq_length_to_mult = 1 - dataset._pad_seq_to_mult = 2 # Used in collate_fn for cu_seqlens_unpadded check (must be > 1 to compute) + dataset._pad_seq_to_mult = 2 # Emit distinct logical and physical cumulative offsets. dataset.ceil_to_power_2 = False dataset.tokenizer = SimpleNamespace(eos_id=tokenizer_eos_id) dataset.answer_only_loss = False @@ -1012,10 +1012,10 @@ def _create_minimal_packed_dataset(tokenizer_eos_id: int = 0): class TestEOSIndexFixInPackedDataset: - """Test EOS index fix for cu_seqlens_unpadded calculation.""" + """Test EOS index fix for logical cumulative sequence offsets.""" def test_eos_index_logic_uses_shape_check(self): - """Ensure cu_seqlens_unpadded handles sequences with <2 EOS tokens.""" + """Ensure logical offsets handle sequences with fewer than two EOS tokens.""" dataset = _create_minimal_packed_dataset() batch = [ { @@ -1026,12 +1026,61 @@ def test_eos_index_logic_uses_shape_check(self): ] processed = dataset.collate_fn(batch) - cu_unpadded = [val for val in processed["cu_seqlens_unpadded"][0].tolist() if val >= 0] + cu_logical = processed["cu_seqlens_q"][0].tolist() # Expect a single non-EOS token tracked without indexing errors. - assert cu_unpadded == [0, 1] + assert cu_logical == [0, 1] + assert processed["cu_seqlens_kv"].tolist() == processed["cu_seqlens_q"].tolist() assert processed["attention_mask"] is None + def test_static_metadata_keeps_logical_and_physical_boundaries_aligned(self): + """Static metadata pads logical and physical boundary arrays to the same shape.""" + dataset = _create_minimal_packed_dataset() + dataset.pad_to_max_length = True + dataset.max_seq_length = 8 + dataset.pad_cu_seqlens = True + dataset.pack_metadata = [ + { + "max_samples_per_bin": 3, + "dataset_max_seqlen": 8, + "min_packed_seqlen": 4, + } + ] + batch = [ + { + "input_ids": np.array([7, 0, 0, 0, 0], dtype=np.int64), + "seq_boundaries": [0, 5], + "loss_mask": np.ones(5, dtype=np.int64), + } + ] + + processed = dataset.collate_fn(batch) + + assert processed["cu_seqlens_q"].shape == processed["cu_seqlens_q_padded"].shape + assert processed["cu_seqlens_q"][0].tolist() == [0, 1, 1, 1, 1] + assert processed["cu_seqlens_kv"][0].tolist() == [0, 1, 1, 1, 1] + assert processed["cu_seqlens_q_padded"][0].tolist() == [0, 4, 8, 8, 8] + assert processed["cu_seqlens_kv_padded"][0].tolist() == [0, 4, 8, 8, 8] + + def test_without_alignment_padding_omits_physical_variants(self): + """Identical logical and physical layouts use only the faster logical TE fields.""" + dataset = _create_minimal_packed_dataset() + dataset._pad_seq_to_mult = 1 + batch = [ + { + "input_ids": np.array([7, 8, 9, 0, 0], dtype=np.int64), + "seq_boundaries": [0, 5], + "loss_mask": np.ones(5, dtype=np.int64), + } + ] + + processed = dataset.collate_fn(batch) + + assert processed["cu_seqlens_q"][0].tolist() == [0, 4] + assert processed["cu_seqlens_kv"][0].tolist() == [0, 4] + assert "cu_seqlens_q_padded" not in processed + assert "cu_seqlens_kv_padded" not in processed + class TestPackedChatDatasetIntegration: """Integration tests for packed chat datasets.""" @@ -1161,11 +1210,11 @@ def test_packed_dataset_preserves_chat_loss_mask(self): assert output_data[0]["loss_mask"] == expected_loss_mask -class TestCuSeqlensUnpaddedCalculation: - """Test cu_seqlens_unpadded calculation with EOS fix.""" +class TestLogicalCuSeqlensCalculation: + """Test logical cumulative sequence offsets with the EOS fix.""" - def test_cu_seqlens_unpadded_calculation_uses_correct_eos(self): - """Ensure cu_seqlens_unpadded honors the tokenizer's EOS id.""" + def test_logical_cu_seqlens_uses_correct_eos(self): + """Ensure logical cumulative offsets honor the tokenizer's EOS id.""" dataset = _create_minimal_packed_dataset(tokenizer_eos_id=999) batch = [ { @@ -1179,10 +1228,10 @@ def test_cu_seqlens_unpadded_calculation_uses_correct_eos(self): ] processed = dataset.collate_fn(batch) - cu_unpadded = [val for val in processed["cu_seqlens_unpadded"][0].tolist() if val >= 0] + cu_logical = processed["cu_seqlens_q"][0].tolist() # Each non-EOS token contributes exactly once despite EOS padding. - assert cu_unpadded == [0, 1, 2] + assert cu_logical == [0, 1, 2] class TestBackwardCompatibilityLossMask: diff --git a/tests/unit_tests/recipes/test_glm5_recipes.py b/tests/unit_tests/recipes/test_glm5_recipes.py index ded9ad3eb47..02a245b70ce 100644 --- a/tests/unit_tests/recipes/test_glm5_recipes.py +++ b/tests/unit_tests/recipes/test_glm5_recipes.py @@ -18,6 +18,7 @@ import pytest +from megatron.bridge.data.builders import ChatSFTPreprocessingConfig from megatron.bridge.recipes.glm import gb200 from megatron.bridge.recipes.glm.gb200 import glm5 as gb200_glm5 from megatron.bridge.recipes.glm.h100 import glm5 @@ -37,6 +38,7 @@ class _FakeMegatronProvider(SimpleNamespace): "context_parallel_size", "cross_entropy_fusion_impl", "cross_entropy_loss_fusion", + "deallocate_pipeline_outputs", "dsa_indexer_loss_coeff", "dsa_indexer_skip_topk_offset", "dsa_indexer_topk_freq", @@ -47,6 +49,8 @@ class _FakeMegatronProvider(SimpleNamespace): "gradient_accumulation_fusion", "microbatch_group_size_per_vp_stage", "moe_flex_dispatcher_backend", + "moe_grouped_gemm", + "moe_permute_fusion", "moe_router_force_load_balancing", "moe_shared_expert_overlap", "moe_token_dispatcher_type", @@ -152,19 +156,19 @@ def test_glm52_h100_200k_recipe_uses_packed_cp() -> None: @pytest.mark.parametrize( - ("recipe", "gbs", "steps"), + ("recipe", "cp", "gbs", "steps", "dispatcher", "backend"), [ - (gb200.glm52_pretrain_192gpu_gb200_bf16_config, 1024, 100), - (gb200.glm52_sft_192gpu_gb200_bf16_config, 32, 100), - (gb200.glm52_peft_192gpu_gb200_bf16_config, 32, 100), + (gb200.glm52_pretrain_192gpu_gb200_bf16_config, 1, 1024, 100, "alltoall", None), + (gb200.glm52_sft_192gpu_gb200_bf16_config, 4, 8, 100, "flex", "hybridep"), + (gb200.glm52_peft_192gpu_gb200_bf16_config, 1, 32, 100, "alltoall", None), ], ) -def test_glm52_gb200_recipe_topologies(recipe, gbs, steps) -> None: +def test_glm52_gb200_recipe_topologies(recipe, cp, gbs, steps, dispatcher, backend) -> None: cfg = recipe() assert cfg.model.tensor_model_parallel_size == 1 assert cfg.model.pipeline_model_parallel_size == 6 - assert cfg.model.context_parallel_size == 1 + assert cfg.model.context_parallel_size == cp assert cfg.model.expert_model_parallel_size == 32 assert cfg.model.expert_tensor_parallel_size == 1 assert cfg.model.sequence_parallel is False @@ -173,12 +177,14 @@ def test_glm52_gb200_recipe_topologies(recipe, gbs, steps) -> None: assert cfg.model.num_layers_in_first_pipeline_stage == 14 assert cfg.model.num_layers_in_last_pipeline_stage == 16 assert cfg.model.microbatch_group_size_per_vp_stage == 6 - assert cfg.model.moe_token_dispatcher_type == "alltoall" - assert cfg.model.moe_flex_dispatcher_backend is None + assert cfg.model.moe_token_dispatcher_type == dispatcher + assert cfg.model.moe_flex_dispatcher_backend == backend assert cfg.model.apply_rope_fusion is False assert cfg.train.global_batch_size == gbs assert cfg.train.micro_batch_size == 1 assert cfg.train.train_iters == steps + data_parallel_size = 192 // (cfg.model.pipeline_model_parallel_size * cfg.model.context_parallel_size) + assert cfg.train.global_batch_size % data_parallel_size == 0 def test_glm52_gb200_128k_recipe_uses_packed_cp() -> None: @@ -188,27 +194,64 @@ def test_glm52_gb200_128k_recipe_uses_packed_cp() -> None: assert cfg.model.pipeline_model_parallel_size == 6 assert cfg.model.context_parallel_size == 32 assert cfg.model.expert_model_parallel_size == 32 - assert cfg.model.pipeline_model_parallel_layout is None - assert cfg.model.num_layers_in_first_pipeline_stage == 14 - assert cfg.model.num_layers_in_last_pipeline_stage == 16 + assert cfg.model.pipeline_model_parallel_layout == gb200_glm5._GLM52_PP6_128K_LAYOUT + assert cfg.model.num_layers_in_first_pipeline_stage is None + assert cfg.model.num_layers_in_last_pipeline_stage is None + stages = cfg.model.pipeline_model_parallel_layout.split("|") + assert [stage.count("t") for stage in stages] == [14, 16, 12, 12, 12, 12] + decoder_starts = [] + decoder_count = 0 + for stage in stages: + decoder_starts.append(decoder_count) + decoder_count += stage.count("t") + assert decoder_starts == [0, 14, 30, 42, 54, 66] + assert all((start + 1 - 3) % 4 == 0 for start in decoder_starts[1:]) assert cfg.model.moe_token_dispatcher_type == "flex" assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.moe_shared_expert_overlap is False assert cfg.model.apply_rope_fusion is False - assert cfg.train.global_batch_size == 8 + assert cfg.train.global_batch_size == 56 assert cfg.train.micro_batch_size == 1 assert cfg.train.train_iters == 20 assert cfg.dataset.seq_length == 131072 - assert cfg.dataset.dataset_root == "work/data/glm5-2/synthetic-long-context-gb200" + assert cfg.dataset.hf_dataset is None + assert cfg.dataset.dataset_root == "work/data/glm5-2/synthetic-long-sft-128k" + assert cfg.dataset.hf_output_root is None + assert cfg.dataset.max_train_samples == 1120 + assert isinstance(cfg.dataset.preprocessing, ChatSFTPreprocessingConfig) assert cfg.dataset.offline_packing_specs.packed_sequence_size == 131072 assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 64 - assert cfg.dataset.offline_packing_specs.pad_cu_seqlens is True + assert cfg.dataset.offline_packing_specs.pad_cu_seqlens is False + assert cfg.dataset.dataset_kwargs == {"pad_to_max_length": True} assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 32 assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 128 assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 assert cfg.env_vars["USE_MNNVL"] == 1 +def test_glm52_gb200_sft_uses_8k_packed_tulu3() -> None: + cfg = gb200.glm52_sft_192gpu_gb200_bf16_config() + + assert cfg.model.seq_length == 8192 + assert cfg.model.context_parallel_size == 4 + assert cfg.dataset.seq_length == 8192 + assert cfg.train.global_batch_size == 8 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.hf_dataset.split == "train[:10000]" + assert cfg.dataset.hf_dataset.load_kwargs == {"revision": gb200_glm5._TULU3_REVISION} + assert cfg.dataset.hf_output_root == "work/data/glm5-2/tulu3-full-sft-gb200-8k-v5" + assert cfg.dataset.offline_packing_specs.packed_sequence_size == 8192 + assert ( + cfg.dataset.offline_packing_specs.max_single_sequence_length + == cfg.model.seq_length - cfg.dataset.offline_packing_specs.pad_seq_to_mult + ) + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 2 * cfg.model.context_parallel_size + assert cfg.dataset.offline_packing_specs.pad_cu_seqlens is False + assert cfg.dataset.dataset_kwargs == {"pad_to_max_length": True} + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + + def test_glm52_gb200_recipes_do_not_depend_on_h100_recipes() -> None: source = Path(gb200_glm5.__file__).read_text(encoding="utf-8") diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index 54f3d58f63f..d65a5187956 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -1425,6 +1425,124 @@ def test_offline_and_in_batch_packing_are_mutually_exclusive(self, monkeypatch): finally: restore_get_world_size_safe(og_ws, cfg_mod) + def test_pad_cu_seqlens_requires_fixed_token_width(self, monkeypatch): + """Test static packed boundaries also require a fixed packed-token width.""" + from megatron.bridge.data.packing import PackedSequenceSpecs + + dataset_cfg = create_test_gpt_sft_dataset_config(sequence_length=512) + dataset_cfg.enable_offline_packing = True + dataset_cfg.offline_packing_specs = PackedSequenceSpecs( + packed_sequence_size=512, + pad_cu_seqlens=True, + ) + container, og_ws, cfg_mod = create_test_config_container( + world_size_override=1, + model_config=create_test_gpt_config(), + dataset_config_override=dataset_cfg, + ) + + try: + with pytest.raises(ValueError, match="pad_cu_seqlens=True requires dataset pad_to_max_length=True"): + container.validate() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + @pytest.mark.parametrize("graph_modules", [[], ["attn"], ["attn", "mlp"]]) + def test_packed_attention_cuda_graph_requires_padded_cu_seqlens(self, graph_modules, monkeypatch): + """Test whole-layer and attention-scoped graphs require static packed boundaries.""" + from megatron.bridge.data.packing import PackedSequenceSpecs + + model_cfg = create_test_gpt_config( + cuda_graph_impl="transformer_engine", + use_te_rng_tracker=True, + ) + set_cuda_graph_modules(model_cfg, graph_modules) + dataset_cfg = create_test_gpt_sft_dataset_config(sequence_length=512) + dataset_cfg.enable_offline_packing = True + dataset_cfg.offline_packing_specs = PackedSequenceSpecs(packed_sequence_size=512) + dataset_cfg.dataset_kwargs = {"pad_to_max_length": True} + container, og_ws, cfg_mod = create_test_config_container( + world_size_override=1, + model_config=model_cfg, + dataset_config_override=dataset_cfg, + ) + + try: + with pytest.raises(ValueError, match="Packed attention CUDA graphs require.*pad_cu_seqlens=True"): + container.validate() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + def test_mlp_only_cuda_graph_does_not_require_padded_cu_seqlens(self, monkeypatch): + """Test an MLP-only graph does not capture packed attention metadata.""" + from megatron.bridge.data.packing import PackedSequenceSpecs + + model_cfg = create_test_gpt_config( + cuda_graph_impl="transformer_engine", + use_te_rng_tracker=True, + ) + set_cuda_graph_modules(model_cfg, ["mlp"]) + dataset_cfg = create_test_gpt_sft_dataset_config(sequence_length=512) + dataset_cfg.enable_offline_packing = True + dataset_cfg.offline_packing_specs = PackedSequenceSpecs(packed_sequence_size=512) + dataset_cfg.dataset_kwargs = {"pad_to_max_length": True} + container, og_ws, cfg_mod = create_test_config_container( + world_size_override=1, + model_config=model_cfg, + dataset_config_override=dataset_cfg, + ) + + try: + container.validate() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + def test_mlp_only_cuda_graph_still_requires_fixed_token_width(self, monkeypatch): + """Test every CUDA graph over offline-packed tokens requires a static token shape.""" + from megatron.bridge.data.packing import PackedSequenceSpecs + + model_cfg = create_test_gpt_config( + cuda_graph_impl="transformer_engine", + use_te_rng_tracker=True, + ) + set_cuda_graph_modules(model_cfg, ["mlp"]) + dataset_cfg = create_test_gpt_sft_dataset_config(sequence_length=512) + dataset_cfg.enable_offline_packing = True + dataset_cfg.offline_packing_specs = PackedSequenceSpecs(packed_sequence_size=512) + container, og_ws, cfg_mod = create_test_config_container( + world_size_override=1, + model_config=model_cfg, + dataset_config_override=dataset_cfg, + ) + + try: + with pytest.raises(ValueError, match="Offline packing with CUDA graphs requires.*pad_to_max_length=True"): + container.validate() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + def test_full_iteration_cuda_graph_requires_padded_cu_seqlens_for_offline_packing(self, monkeypatch): + """Test full-iteration graphs require static packed attention metadata.""" + from megatron.bridge.data.packing import PackedSequenceSpecs + + model_cfg = create_test_gpt_config(use_te_rng_tracker=True) + set_full_iteration_cuda_graph(model_cfg) + dataset_cfg = create_test_gpt_sft_dataset_config(sequence_length=512) + dataset_cfg.enable_offline_packing = True + dataset_cfg.offline_packing_specs = PackedSequenceSpecs(packed_sequence_size=512) + dataset_cfg.dataset_kwargs = {"pad_to_max_length": True} + container, og_ws, cfg_mod = create_test_config_container( + world_size_override=1, + model_config=model_cfg, + dataset_config_override=dataset_cfg, + ) + + try: + with pytest.raises(ValueError, match="Packed attention CUDA graphs require.*pad_cu_seqlens=True"): + container.validate() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + @pytest.mark.parametrize( "seq_length, context_parallel_size, expect_assertion_error", [ diff --git a/tests/unit_tests/training/test_gpt_step.py b/tests/unit_tests/training/test_gpt_step.py index d5b38f347c0..4bc31133168 100644 --- a/tests/unit_tests/training/test_gpt_step.py +++ b/tests/unit_tests/training/test_gpt_step.py @@ -747,10 +747,9 @@ def test_basic_packed_seq_params_with_max_seqlen(self): assert torch.equal(result.cu_seqlens_q, expected_cu_seqlens) assert torch.equal(result.cu_seqlens_kv, expected_cu_seqlens) - # Verify max_seqlen was squeezed - expected_max_seqlen = torch.tensor(15, dtype=torch.int32) - assert torch.equal(result.max_seqlen_q, expected_max_seqlen) - assert torch.equal(result.max_seqlen_kv, expected_max_seqlen) + # Verify max_seqlen was normalized to MCore's Python-int contract + assert result.max_seqlen_q == 15 + assert result.max_seqlen_kv == 15 # Verify qkv_format is correct assert result.qkv_format == "thd" @@ -797,9 +796,8 @@ def test_packed_seq_params_with_cu_seqlens_argmin(self): assert torch.equal(result.cu_seqlens_kv, expected_cu_seqlens) # Verify max_seqlen was processed correctly - expected_max_seqlen = torch.tensor(18, dtype=torch.int32) - assert torch.equal(result.max_seqlen_q, expected_max_seqlen) - assert torch.equal(result.max_seqlen_kv, expected_max_seqlen) + assert result.max_seqlen_q == 18 + assert result.max_seqlen_kv == 18 def test_packed_seq_params_with_cu_seqlens_argmin_zero(self): """Test edge case when cu_seqlens_argmin is 0.""" @@ -829,8 +827,7 @@ def test_packed_seq_params_batch_dimension_removal(self): expected_cu_seqlens = torch.tensor([0, 6, 12], dtype=torch.int32) assert torch.equal(result.cu_seqlens_q, expected_cu_seqlens) - expected_max_seqlen = torch.tensor(20, dtype=torch.int32) - assert torch.equal(result.max_seqlen_q, expected_max_seqlen) + assert result.max_seqlen_q == 20 def test_packed_seq_params_with_different_dtypes(self): """Test functionality with different tensor dtypes.""" @@ -845,8 +842,7 @@ def test_packed_seq_params_with_different_dtypes(self): expected_cu_seqlens = torch.tensor([0, 10, 20], dtype=torch.int64) assert torch.equal(result.cu_seqlens_q, expected_cu_seqlens) - expected_max_seqlen = torch.tensor(25, dtype=torch.int64) - assert torch.equal(result.max_seqlen_q, expected_max_seqlen) + assert result.max_seqlen_q == 25 def test_packed_seq_params_all_fields_match(self): """Test that cu_seqlens_q/kv and max_seqlen_q/kv are identical.""" @@ -859,7 +855,7 @@ def test_packed_seq_params_all_fields_match(self): # Verify that q and kv parameters are identical (as expected for this function) assert torch.equal(result.cu_seqlens_q, result.cu_seqlens_kv) - assert torch.equal(result.max_seqlen_q, result.max_seqlen_kv) + assert result.max_seqlen_q == result.max_seqlen_kv def test_packed_seq_params_with_cu_seqlens_unpadded(self): """Test functionality with cu_seqlens_unpadded for THD CP support.""" diff --git a/tests/unit_tests/training/utils/test_packed_seq_utils.py b/tests/unit_tests/training/utils/test_packed_seq_utils.py index 21581bc25fd..0cde368a3f4 100644 --- a/tests/unit_tests/training/utils/test_packed_seq_utils.py +++ b/tests/unit_tests/training/utils/test_packed_seq_utils.py @@ -215,8 +215,20 @@ def test_current_mcore_metadata_fields(self): torch.testing.assert_close(result.cu_seqlens_kv_padded, batch["cu_seqlens_kv_padded"]) assert result.max_seqlen_q == 128 assert result.max_seqlen_kv == 128 + assert isinstance(result.max_seqlen_q, int) + assert isinstance(result.max_seqlen_kv, int) assert result.qkv_format == "thd" + def test_current_mcore_metadata_rejects_non_scalar_max_seqlen(self): + """Test that malformed max-sequence metadata fails before reaching MCore.""" + batch = { + "cu_seqlens_q": torch.IntTensor([0, 64, 128]), + "max_seqlen_q": torch.IntTensor([64, 64]), + } + + with pytest.raises(ValueError, match="max_seqlen_q must contain exactly one value"): + get_packed_seq_params(batch) + def test_without_cu_seqlens_unpadded(self): """Test get_packed_seq_params when cu_seqlens_unpadded is NOT present. From bb94c3f489a36aaab768dfd4b54b7fefe3384ad9 Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Sat, 8 Aug 2026 12:35:33 -0700 Subject: [PATCH 57/85] cp: `[examples] refactor: direct GLM inference to model card (5346)` into `r0.6.0` (#5427) Signed-off-by: yaoyu-33 --- docs/models/glm/glm5.md | 2 +- examples/models/glm/glm5/README.md | 35 ++---- examples/models/glm/glm5/slurm_inference.sh | 131 -------------------- 3 files changed, 11 insertions(+), 157 deletions(-) delete mode 100644 examples/models/glm/glm5/slurm_inference.sh diff --git a/docs/models/glm/glm5.md b/docs/models/glm/glm5.md index 6c5e09d5f7a..cc845b7bb8c 100644 --- a/docs/models/glm/glm5.md +++ b/docs/models/glm/glm5.md @@ -21,7 +21,7 @@ ## Examples -For conversion, inference, dependency notes, hardware requirements, and MCore patch requirements, see the [GLM-5 examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/glm/glm5/README.md). +For the pinned checkpoint revision, tested topology, commands, and expected results for verified GLM-5.2 workflows, see the [GLM-5.2 model verification card](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/model_verification_cards/glm5-2/card.yaml). For the GLM-5 conversion wrapper, dependency notes, and architecture constraints, see the [GLM-5 examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/glm/glm5/README.md). ## Related Implementation diff --git a/examples/models/glm/glm5/README.md b/examples/models/glm/glm5/README.md index 4f74072e8cc..70eeec45354 100644 --- a/examples/models/glm/glm5/README.md +++ b/examples/models/glm/glm5/README.md @@ -1,8 +1,8 @@ -# GLM-5 / GLM-5.1 Examples +# GLM-5 Family Examples -Scripts for the GLM-5 family — [GLM-5](https://huggingface.co/zai-org/GLM-5) (`zai-org/GLM-5`) and [GLM-5.1](https://huggingface.co/zai-org/GLM-5.1) (`zai-org/GLM-5.1`) — large sparse MoE models with Multi-Latent Attention (MLA) and Dynamic Sparse Attention (DSA). +Examples for the GLM-5 family — [GLM-5](https://huggingface.co/zai-org/GLM-5) (`zai-org/GLM-5`), [GLM-5.1](https://huggingface.co/zai-org/GLM-5.1) (`zai-org/GLM-5.1`), and [GLM-5.2](https://huggingface.co/zai-org/GLM-5.2) (`zai-org/GLM-5.2`) — large sparse MoE models with Multi-Latent Attention (MLA) and Dynamic Sparse Attention (DSA). -GLM-5 and GLM-5.1 share the `GlmMoeDsaForCausalLM` architecture and identical MoE / MLA / DSA dimensions, so the same `GLM5Bridge` handles both. To run the GLM-5.1 checkpoint, replace `zai-org/GLM-5` with `zai-org/GLM-5.1` in the relevant example script. +All three checkpoints use the `GlmMoeDsaForCausalLM` architecture and are handled by `GLM5Bridge`. GLM-5 and GLM-5.1 have identical MoE / MLA / DSA dimensions, while GLM-5.2 adds IndexShare-style DSA index reuse settings. | Property | Value | |---|---| @@ -18,12 +18,11 @@ GLM-5 and GLM-5.1 share the `GlmMoeDsaForCausalLM` architecture and identical Mo ## Hardware Requirements -Full-model conversion and inference in BF16 requires **at least 8 nodes (64 GPUs × 80 GB)**. Key constraints: +The included GLM-5 round-trip conversion example uses **8 nodes (64 GPUs × 80 GB)**. See the [GLM-5.2 model verification card](../../../model_verification_cards/glm5-2/card.yaml) for the hardware and parallelism used by each verified GLM-5.2 workflow. Key constraints include: - EP must divide 256 (number of routed experts). Valid: 1, 2, 4, 8, 16, 32, 64, 128, 256. - TP does **not** reduce expert memory — increase EP instead. -- Minimum recommended: `TP=1, PP=2, EP=32` (64 GPUs, 8 nodes). PP splits the 78 transformer layers evenly, with 39 layers per stage, and EP places 8 routed experts per GPU. -- `TP=1, PP=1, EP=64` works for conversion but may cause empty-dispatch issues during autoregressive inference with single-token batches. Prefer `PP=2, EP=32` for inference on 64 GPUs. +- The conversion wrapper uses `TP=1, PP=2, EP=32` on 64 GPUs. PP splits the 78 transformer layers evenly, with 39 layers per stage, and EP places 8 routed experts per GPU. ### Pre-requisites @@ -38,22 +37,9 @@ The PyPI source distribution is incomplete; install from the git repo. ## Inference (Megatron) -[slurm_inference.sh](slurm_inference.sh) loads the HF checkpoint, converts to Megatron in-memory, and runs greedy text generation with `TP=1, PP=2, EP=32` across 64 GPUs. TP remains disabled because AbsorbedMLA requires sequence parallelism when `TP > 1`. The launcher passes `--legacy-full-prefix` to use non-cached autoregressive generation, which recomputes the full prefix at every decoding step. Cached MCore inference is not yet supported for AbsorbedMLA. +Use the verified `inference` item in the [GLM-5.2 model verification card](../../../model_verification_cards/glm5-2/card.yaml). The card is the canonical source for the pinned checkpoint revision, tested topology, `scripts/inference/infer.sh` command, and expected result. -```bash -sbatch examples/models/glm/glm5/slurm_inference.sh -``` - -### Expected output - -``` -======== GENERATED TEXT OUTPUT ======== -Prompt: What is artificial intelligence? -Generated: What is artificial intelligence? Artificial intelligence (AI) is a field of -computer Science and Engineering that deals with the creation of intelligent -machines, which are used in different areas such... -======================================= -``` +The verified command selects the `legacy-full-prefix-generation` compatibility task. It recomputes the accumulated prefix at every decoding step because cached inference is not yet supported for AbsorbedMLA. ## Checkpoint Conversion (Round-Trip) @@ -71,12 +57,11 @@ bash examples/models/glm/glm5/slurm_conversion.sh The script uses 8 nodes (64 GPUs) with `TP=1`, `PP=2`, and `EP=32`. > **Note:** The round-trip verification step (comparing ~63K weight tensors on rank 0) -> may hit Lustre I/O contention at this model scale. The HF→Megatron conversion -> itself is validated by the successful inference above. +> may hit shared-filesystem I/O contention at this model scale. -## Script Configuration +## Conversion Script Configuration -Both scripts resolve the HF model from the local cache to avoid `snapshot_download` race conditions with 64 concurrent processes. Set these environment variables before submitting: +Set these environment variables before submitting the round-trip conversion wrapper: | Variable | Description | |---|---| diff --git a/examples/models/glm/glm5/slurm_inference.sh b/examples/models/glm/glm5/slurm_inference.sh deleted file mode 100644 index 05e34c3a9a7..00000000000 --- a/examples/models/glm/glm5/slurm_inference.sh +++ /dev/null @@ -1,131 +0,0 @@ -#!/bin/bash -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# ============================================================================== -# GLM-5 / GLM-5.1 Multi-Node Inference (Slurm) -# -# GLM-5 and GLM-5.1 share the same GlmMoeDsaForCausalLM architecture -# (MoE + MLA + DSA: 256 routed experts, top-8, ~800B+ params, BF16), -# so this script handles both. Set MODEL_NAME=GLM-5.1 to run the 5.1 checkpoint. -# -# Full model requires 8 nodes (64 GPUs) minimum. -# TP does NOT reduce expert memory — increase EP instead. -# Recommended: TP=1, EP=32, PP=2 (64 GPUs, 8 nodes). -# PP=2 splits the 78 transformer layers evenly (39 per stage), while TP=1 -# avoids the AbsorbedMLA sequence-parallel requirement for TP > 1. -# -# Megatron inference engine support is not yet available for this architecture. -# This launcher passes --legacy-full-prefix to use non-cached autoregressive -# generation, which recomputes the full prefix at every decoding step. -# -# Loads the HF checkpoint, converts to Megatron in-memory, and runs greedy -# text generation via examples/conversion/hf_to_megatron_generate_text.py. -# -# Requirements: transformers >= 5.2.0 -# -# Usage: -# sbatch examples/models/glm/glm5/slurm_inference.sh # GLM-5 -# MODEL_NAME=GLM-5.1 sbatch examples/models/glm/glm5/slurm_inference.sh -# ============================================================================== - -#SBATCH --job-name=glm5-infer -#SBATCH --nodes=8 -#SBATCH --ntasks-per-node=1 -#SBATCH --gpus-per-node=8 -#SBATCH --time=0:30:00 -#SBATCH --account=${SLURM_ACCOUNT:-your_account} -#SBATCH --partition=batch -#SBATCH --output=logs/glm5_inference_%j.log -#SBATCH --exclusive - -# ── Container ──────────────────────────────────────────────────────────── -CONTAINER_IMAGE="${CONTAINER_IMAGE:?Set CONTAINER_IMAGE to your .sqsh container path}" -CONTAINER_MOUNTS="/lustre:/lustre" -WORKDIR="/opt/Megatron-Bridge" -BRIDGE_PATH="${BRIDGE_PATH:?Set BRIDGE_PATH to your Megatron-Bridge checkout on shared storage}" - -# ── Tokens / Caches ────────────────────────────────────────────────────── -export HF_TOKEN="${HF_TOKEN:?Set HF_TOKEN for gated model access}" -export HF_HOME="${HF_HOME:?Set HF_HOME to your HuggingFace cache directory}" -export UV_CACHE_DIR="${UV_CACHE_DIR:-/tmp/uv_cache}" - -# ── Model / Parallelism ────────────────────────────────────────────────── -# MODEL_NAME selects between GLM-5 and GLM-5.1 (same architecture). -MODEL_NAME="${MODEL_NAME:-GLM-5}" -# Use the direct local snapshot path to avoid 64 processes calling -# snapshot_download simultaneously (causes Lustre race conditions). -HF_MODEL_PATH="${HF_HOME}/hub/models--zai-org--${MODEL_NAME}/snapshots/$(ls ${HF_HOME}/hub/models--zai-org--${MODEL_NAME}/snapshots/ | head -1)" -TP="${TP:-1}" -EP="${EP:-32}" -PP="${PP:-2}" - -PROMPT="${PROMPT:-What is artificial intelligence?}" -MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-100}" - -# ── Environment ─────────────────────────────────────────────────────────── -export TORCH_NCCL_AVOID_RECORD_STREAMS=1 -export NCCL_NVLS_ENABLE=0 -export NCCL_DEBUG=WARN -export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -export HF_HUB_OFFLINE=1 -export NCCL_TIMEOUT=1800000 -export WORKDIR HF_MODEL_PATH TP EP PP PROMPT MAX_NEW_TOKENS - -# ============================================================================== -# Job Execution -# ============================================================================== - -echo "======================================" -echo "${MODEL_NAME} Inference" -echo "Job: $SLURM_JOB_ID | Nodes: $SLURM_JOB_NUM_NODES" -echo "TP=$TP PP=$PP EP=$EP (Total GPUs: $((SLURM_JOB_NUM_NODES * 8)))" -echo "======================================" - -mkdir -p logs - -srun --mpi=pmix \ - --export=HF_TOKEN,HF_HOME,UV_CACHE_DIR,NCCL_DEBUG,TORCH_NCCL_AVOID_RECORD_STREAMS,NCCL_NVLS_ENABLE,PYTORCH_CUDA_ALLOC_CONF,HF_HUB_OFFLINE,NCCL_TIMEOUT,WORKDIR,HF_MODEL_PATH,TP,EP,PP,PROMPT,MAX_NEW_TOKENS \ - --container-image="$CONTAINER_IMAGE" \ - --container-mounts="${BRIDGE_PATH}:${WORKDIR},${CONTAINER_MOUNTS}" \ - --container-env=HF_TOKEN,HF_HOME,UV_CACHE_DIR,NCCL_DEBUG,TORCH_NCCL_AVOID_RECORD_STREAMS,NCCL_NVLS_ENABLE,PYTORCH_CUDA_ALLOC_CONF,HF_HUB_OFFLINE,NCCL_TIMEOUT,WORKDIR,HF_MODEL_PATH,TP,EP,PP,PROMPT,MAX_NEW_TOKENS \ - --no-container-mount-home \ - bash -c ' - MASTER_ADDR=$(python3 -c " -import re, os -s = os.environ.get('"'"'SLURM_NODELIST'"'"', '"'"''"'"') -m = re.match(r'"'"'([\w-]+)\[(\d+)'"'"', s) -print(m.group(1) + m.group(2) if m else s.split('"'"','"'"')[0]) -") - cd "$WORKDIR" - export PYTHONPATH="$WORKDIR/.venv/lib/python3.12/site-packages:${PYTHONPATH:-}" - uv run --no-sync python -m torch.distributed.run \ - --nproc_per_node=8 \ - --nnodes="$SLURM_JOB_NUM_NODES" \ - --node_rank="$SLURM_PROCID" \ - --master_addr="$MASTER_ADDR" \ - --master_port=29500 \ - examples/conversion/hf_to_megatron_generate_text.py \ - --hf_model_path "$HF_MODEL_PATH" \ - --prompt "$PROMPT" \ - --max_new_tokens "$MAX_NEW_TOKENS" \ - --legacy-full-prefix \ - --tp "$TP" --ep "$EP" --pp "$PP" - ' -status=$? - -echo "======================================" -echo "Inference completed (exit $status)" -echo "======================================" -exit "$status" From f01b1f38e8d2e3d5036c55a0111659d5826be1bf Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:35:43 +0200 Subject: [PATCH 58/85] cp: `docs(models): link MiMo and StepFun examples (5432)` into `r0.6.0` (#5439) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- docs/models/mimo/mimo.md | 1 + docs/models/stepfun/step35.md | 1 + 2 files changed, 2 insertions(+) diff --git a/docs/models/mimo/mimo.md b/docs/models/mimo/mimo.md index f912bb01873..546056c5b9a 100644 --- a/docs/models/mimo/mimo.md +++ b/docs/models/mimo/mimo.md @@ -18,6 +18,7 @@ Megatron Bridge supports Hugging Face checkpoints using the `MiMoForCausalLM` ar Checkpoint conversion and text-generation examples for MiMo causal language models live under [`examples/models/mimo`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/6a706b1c3afab4d21a6f5dd88aa6a75296d33fe2/examples/models/mimo/README.md). +For MiMo-V2-Flash conversion and inference examples, see the [MiMo-V2-Flash examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/mimo_v2_flash/README.md). Heterogeneous multimodal training examples live separately under [`examples/megatron_mimo`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/megatron_mimo). diff --git a/docs/models/stepfun/step35.md b/docs/models/stepfun/step35.md index 9a4ef742a38..c867c6c5910 100644 --- a/docs/models/stepfun/step35.md +++ b/docs/models/stepfun/step35.md @@ -17,6 +17,7 @@ ## Examples For conversion, inference, pretraining, Slurm launch scripts, and MCore branch notes, see the [Step-3.5-Flash examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/stepfun/step35/README.md). +For Step 3.7 Flash conversion and training examples, see the [Step 3.7 Flash examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/stepfun/step37/README.md). ## Related Implementation From 3264eda8b35a0aecfa7c4fba154307dcb9a5c73d Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:35:51 +0200 Subject: [PATCH 59/85] cp: `docs(recipe): retire stale Gemma 4 VL SFT launchers (5433)` into `r0.6.0` (#5435) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- examples/models/gemma/gemma4_vl/README.md | 30 ++-- examples/models/gemma/gemma4_vl/sft.sh | 71 ---------- examples/models/gemma/gemma4_vl/slurm_sft.sh | 138 ------------------- 3 files changed, 15 insertions(+), 224 deletions(-) delete mode 100644 examples/models/gemma/gemma4_vl/sft.sh delete mode 100644 examples/models/gemma/gemma4_vl/slurm_sft.sh diff --git a/examples/models/gemma/gemma4_vl/README.md b/examples/models/gemma/gemma4_vl/README.md index d82b5580c47..ec8040a2127 100644 --- a/examples/models/gemma/gemma4_vl/README.md +++ b/examples/models/gemma/gemma4_vl/README.md @@ -73,16 +73,14 @@ Before training, ensure the following environment variables are set: ### Supervised Fine-Tuning (SFT) -For single-node interactive runs, see [sft.sh](sft.sh). +Use the verified `sft.H100` command in the [Gemma 4 26B-A4B-it model verification +card](../../../model_verification_cards/gemma-4-26b-a4b-it/card.yaml). It launches the public +`scripts/training/train.sh` entry point and records the exact checkpoint, dataset revision, topology, memory settings, +metrics, and expected result. -For multi-node Slurm jobs, see [slurm_sft.sh](slurm_sft.sh). Default configuration: TP=2, PP=1, EP=8 on 2 nodes (16 GPUs). - -```bash -# Override defaults via environment variables -PRETRAINED_CHECKPOINT=${WORKSPACE}/models/gemma-4-26B-A4B \ -TP=2 PP=1 EP=8 \ -sbatch --nodes=2 slurm_sft.sh -``` +The verified recipe owns its TP4/PP1/EP8/ETP1 topology and selective `core_attn` plus `moe_act` recompute settings. Do +not replace them with manual TP, PP, EP, or ETP overrides when reproducing the card result. The former model-specific +SFT launchers were removed because their hard-coded topologies diverged from the verified recipe. ### Parameter-Efficient Fine-Tuning (PEFT) with LoRA @@ -94,13 +92,15 @@ For multi-node Slurm jobs, see [slurm_peft.sh](slurm_peft.sh). Default configura sbatch slurm_peft.sh ``` -### Recommended Configurations +### Verified Configurations + +| Mode | Verified hardware | Topology | Global Batch Size | Learning Rate | Notes | +|------|-------------------|----------|-------------------|---------------|-------| +| Full SFT | 8 H100 GPUs | TP4/PP1/EP8/ETP1 | 32 | 5e-5 | Freeze vision; selective `core_attn` and `moe_act` recompute | +| LoRA | 4 H100 GPUs | TP2/PP1/EP4/ETP1 | 32 | 2e-4 | Train language-model adapters; no recompute | -| Mode | TP | PP | EP | Nodes | Global Batch Size | Learning Rate | Notes | -|------|----|----|----|----|-------------------|---------------|-------| -| Full SFT | 2 | 1 | 8 | 2 | 32 | 5e-5 | Max EP=DP=8; vision unfrozen; no activation recompute | -| Full SFT | 4 | 2 | 1 | 1 | 32 | 5e-5 | `recompute_granularity="selective"`; freeze vision | -| LoRA | 2 | 1 | 4 | 1 | 32 | 2e-4 | EP=4 required (see note above) | +The model verification card is authoritative for the commands, immutable inputs, Bridge commit, metrics, and current +verification status of these configurations. > **Note:** Do not use `recompute_granularity="full"`. Megatron's `CheckpointFunction` does not support non-tensor (tuple) arguments, causing a `TypeError` at runtime. Use `"selective"` instead. diff --git a/examples/models/gemma/gemma4_vl/sft.sh b/examples/models/gemma/gemma4_vl/sft.sh deleted file mode 100644 index f049613681a..00000000000 --- a/examples/models/gemma/gemma4_vl/sft.sh +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Workspace directory for checkpoints and results -WORKSPACE=${WORKSPACE:-/workspace} - -# Before training, make sure to set WANDB_API_KEY or disable wandb logging -# export WANDB_API_KEY= -# export WANDB_MODE=disabled - -# Common configurations -PRETRAINED_CHECKPOINT=${WORKSPACE}/models/gemma-4-26B-A4B -MODEL_NAME=gemma4_vl_26b -DATASET_NAME=cord_v2 -SEQ_LENGTH=4096 -TRAIN_ITERS=50 -GLOBAL_BATCH_SIZE=32 -MICRO_BATCH_SIZE=1 -EVAL_ITERS=10 -LR=0.00005 -MIN_LR=0.000005 -LR_WARMUP_ITERS=10 -LOG_INTERVAL=1 -WANDB_PROJECT=megatron-bridge-${DATASET_NAME} - -# TP/PP combinations: "TP,PP" -# These configs target a single 8-GPU node (DP=1 so EP must also be 1). -# With EP=1 the 128 MoE experts are NOT sharded across GPUs, so activation -# recompute and a frozen vision model are required to avoid OOM. -PARALLELISM_CONFIGS=("4,2" "8,1") - -for config in "${PARALLELISM_CONFIGS[@]}"; do - IFS=',' read -r TP PP <<< "$config" - - echo "Running full finetuning with TP=$TP, PP=$PP" - uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 scripts/training/run_recipe.py \ - --recipe ${MODEL_NAME}_sft_config \ - --step_func vlm_step \ - checkpoint.pretrained_checkpoint=$PRETRAINED_CHECKPOINT \ - model.seq_length=$SEQ_LENGTH \ - train.train_iters=$TRAIN_ITERS \ - train.global_batch_size=$GLOBAL_BATCH_SIZE \ - train.micro_batch_size=$MICRO_BATCH_SIZE \ - validation.eval_iters=$EVAL_ITERS \ - optimizer.lr=$LR \ - optimizer.min_lr=$MIN_LR \ - scheduler.lr_warmup_iters=$LR_WARMUP_ITERS \ - checkpoint.save=${WORKSPACE}/results/${MODEL_NAME}_sft_tp${TP}_pp${PP} \ - logger.log_interval=$LOG_INTERVAL \ - logger.wandb_project=$WANDB_PROJECT \ - logger.wandb_exp_name=${MODEL_NAME}_${DATASET_NAME}_sft_tp${TP}_pp${PP} \ - dataset.source.dataset_name=${DATASET_NAME} \ - dataset.seq_length=$SEQ_LENGTH \ - model.tensor_model_parallel_size=$TP \ - model.pipeline_model_parallel_size=$PP \ - model.expert_model_parallel_size=1 \ - model.recompute_granularity=selective \ - model.freeze_vision_model=True -done diff --git a/examples/models/gemma/gemma4_vl/slurm_sft.sh b/examples/models/gemma/gemma4_vl/slurm_sft.sh deleted file mode 100644 index f4e10dd5e4d..00000000000 --- a/examples/models/gemma/gemma4_vl/slurm_sft.sh +++ /dev/null @@ -1,138 +0,0 @@ -#!/bin/bash -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# Gemma 4 VL 26B-A4B SFT fine-tuning. -# -# Usage (single node, 8 GPUs): -# sbatch slurm_sft.sh -# -# Usage (multi-node, e.g. 2 nodes with EP=8): -# NUM_NODES=2 TP=2 PP=1 EP=8 sbatch --nodes=2 slurm_sft.sh - -#SBATCH --job-name=gemma4vl-sft -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=8 -#SBATCH --gpus-per-node=8 -#SBATCH --time=00:30:00 -#SBATCH --partition=gpu -#SBATCH --account=my_account -#SBATCH --output=gemma4_vl_sft_%j.out -#SBATCH --error=gemma4_vl_sft_%j.err -#SBATCH --exclusive - -# ============================================================================== -# CONFIGURATION -# ============================================================================== -WORKSPACE=${WORKSPACE:-/workspace} - -PRETRAINED_CHECKPOINT=${PRETRAINED_CHECKPOINT:-${WORKSPACE}/models/gemma-4-26B-A4B} -RECIPE=${RECIPE:-gemma4_vl_26b_sft_config} -DATASET_NAME=${DATASET_NAME:-cord_v2} -SEQ_LENGTH=${SEQ_LENGTH:-4096} -TRAIN_ITERS=${TRAIN_ITERS:-40} # 40 iters fits 30-min wall time with checkpoint save -GBS=${GBS:-32} -MBS=${MBS:-1} -EVAL_ITERS=${EVAL_ITERS:-5} # 5 eval iters (was 10) reduces eval cost ~36s vs 72s -LR=${LR:-0.00005} -MIN_LR=${MIN_LR:-0.000005} -LR_WARMUP_ITERS=${LR_WARMUP_ITERS:-10} -TP=${TP:-2} -PP=${PP:-1} -EP=${EP:-8} # Max EP with 16 GPUs, TP=2, PP=1: DP=8, EP=8 (EP must divide DP) -SAVE_DIR=${SAVE_DIR:-${WORKSPACE}/results/gemma4_vl_sft_tp${TP}_pp${PP}_ep${EP}_${SLURM_JOB_ID}} - -# Container image (required) -CONTAINER_IMAGE="${CONTAINER_IMAGE:-}" -# CONTAINER_IMAGE="/path/to/container.sqsh" - -# Container mounts (optional) -CONTAINER_MOUNTS="${CONTAINER_MOUNTS:-}" -# CONTAINER_MOUNTS="/data:/data,/workspace:/workspace" - -# ============================================================================== -# Environment -# ============================================================================== - -export TORCH_NCCL_AVOID_RECORD_STREAMS=1 -export NCCL_NVLS_ENABLE=0 -# Reduce fragmentation — helps when reserved-but-unallocated memory is ~1-2 GiB -export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" - -# export UV_CACHE_DIR="/path/to/shared/uv_cache" -# export HF_HOME="/path/to/shared/HF_HOME" -# export HF_TOKEN="hf_your_token_here" -# export WANDB_API_KEY="your_wandb_key_here" -# export WANDB_MODE=disabled - -export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -1) -export MASTER_PORT=${MASTER_PORT:-29501} - -# ============================================================================== -# Job -# ============================================================================== - -echo "======================================" -echo "Gemma 4 VL 26B-A4B SFT" -echo "======================================" -echo "Job ID: $SLURM_JOB_ID" -echo "Nodes: $SLURM_JOB_NUM_NODES" -echo "Total tasks: $SLURM_NTASKS" -echo "Recipe: $RECIPE" -echo "TP=$TP PP=$PP EP=$EP" -echo "Checkpoint: $PRETRAINED_CHECKPOINT" -echo "Save: $SAVE_DIR" -echo "Master: $MASTER_ADDR:$MASTER_PORT" -echo "======================================" - -CMD="cd /opt/Megatron-Bridge && " -CMD="${CMD}export RANK=\$SLURM_PROCID LOCAL_RANK=\$SLURM_LOCALID WORLD_SIZE=\$SLURM_NTASKS && " -CMD="${CMD}uv run --no-sync python scripts/training/run_recipe.py" -CMD="${CMD} --recipe ${RECIPE}" -CMD="${CMD} --mode sft" -CMD="${CMD} --step_func vlm_step" -CMD="${CMD} checkpoint.pretrained_checkpoint=${PRETRAINED_CHECKPOINT}" -CMD="${CMD} model.tensor_model_parallel_size=${TP}" -CMD="${CMD} model.pipeline_model_parallel_size=${PP}" -CMD="${CMD} model.expert_model_parallel_size=${EP}" -CMD="${CMD} model.seq_length=${SEQ_LENGTH}" -CMD="${CMD} train.train_iters=${TRAIN_ITERS}" -CMD="${CMD} train.global_batch_size=${GBS}" -CMD="${CMD} train.micro_batch_size=${MBS}" -CMD="${CMD} validation.eval_iters=${EVAL_ITERS}" -CMD="${CMD} optimizer.lr=${LR}" -CMD="${CMD} optimizer.min_lr=${MIN_LR}" -CMD="${CMD} scheduler.lr_warmup_iters=${LR_WARMUP_ITERS}" -CMD="${CMD} checkpoint.save=${SAVE_DIR}" -CMD="${CMD} dataset.source.dataset_name=${DATASET_NAME}" -CMD="${CMD} dataset.seq_length=${SEQ_LENGTH}" - -echo "Running training..." - -if [ -z "$CONTAINER_IMAGE" ]; then - echo "ERROR: CONTAINER_IMAGE must be set. Please specify a valid container image." - exit 1 -fi - -SRUN_CMD="srun --mpi=pmix --container-image=${CONTAINER_IMAGE} --no-container-mount-home" - -if [ -n "$CONTAINER_MOUNTS" ]; then - SRUN_CMD="${SRUN_CMD} --container-mounts=${CONTAINER_MOUNTS}" -fi - -$SRUN_CMD bash -c "$CMD" - -echo "======================================" -echo "Training job finished. EXIT=$?" -echo "======================================" From 75e6be5874d08afdb4d29ee0bb192afd2d9ae2e3 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:36:01 +0200 Subject: [PATCH 60/85] cp: `fix(peft): synchronize shared expert LoRA across EP (5396)` into `r0.6.0` (#5428) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- src/megatron/bridge/peft/utils.py | 32 +++ .../test_shared_expert_lora_distributed.py | 184 ++++++++++++++++++ tests/unit_tests/peft/test_utils.py | 83 +++++++- 3 files changed, 298 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/peft/test_shared_expert_lora_distributed.py diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py index f6f8d30ec8a..b170bc1fa44 100644 --- a/src/megatron/bridge/peft/utils.py +++ b/src/megatron/bridge/peft/utils.py @@ -1067,6 +1067,7 @@ def __init__( self.half() if self._uses_grouped_expert_sharding(): + self._synchronize_shared_expert_parameters() self._register_shared_expert_grad_sync_hooks() # revert config change in case it is read elsewhere @@ -1281,6 +1282,37 @@ def _allreduce_shared_expert_grad(self, grad: torch.Tensor) -> torch.Tensor: torch.distributed.all_reduce(grad, group=self.ep_group) return grad + def _synchronize_shared_expert_parameters(self) -> None: + """Broadcast shared expert adapter parameters from EP group rank zero.""" + + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return + if self.ep_group is None or _process_group_size(self.ep_group) <= 1: + return + + weights = [ + weight + for module in (self.linear_in, self.linear_out) + if isinstance(weight := getattr(module, "weight", None), torch.Tensor) + ] + if not weights: + return + + src_rank = torch.distributed.get_global_rank(self.ep_group, 0) + with torch.no_grad(): + for weight in weights: + if weight.is_meta: + raise RuntimeError( + "Shared expert adapter parameters must be materialized before EP synchronization" + ) + if weight.is_cuda or torch.distributed.get_backend(self.ep_group) != "nccl": + torch.distributed.broadcast(weight, src=src_rank, group=self.ep_group) + continue + + staged_weight = weight.to(torch.device("cuda", torch.cuda.current_device())) + torch.distributed.broadcast(staged_weight, src=src_rank, group=self.ep_group) + weight.copy_(staged_weight.cpu()) + def _register_shared_expert_grad_sync_hooks(self) -> None: """Keep shared grouped-expert adapters synchronized across EP ranks.""" diff --git a/tests/unit_tests/peft/test_shared_expert_lora_distributed.py b/tests/unit_tests/peft/test_shared_expert_lora_distributed.py new file mode 100644 index 00000000000..552da6913f5 --- /dev/null +++ b/tests/unit_tests/peft/test_shared_expert_lora_distributed.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Two-rank parity tests for shared expert LoRA parameters. + +Run with: +uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest \ + tests/unit_tests/peft/test_shared_expert_lora_distributed.py +""" + +import os +from collections.abc import Iterator +from contextlib import contextmanager + +import megatron.core.parallel_state as parallel_state +import pytest +import torch +import torch.distributed as dist +from megatron.core.model_parallel_config import ModelParallelConfig +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + +from megatron.bridge.peft.utils import ParallelLinearAdapter + + +_EP_SIZE = 2 +_ADAPTER_CASES = ( + ("decoder.layers.0.mlp.experts.linear_fc1", False), + ("decoder.layers.0.mlp.experts.linear_fc2", True), +) + + +@contextmanager +def _distributed_ep() -> Iterator[ProcessGroupCollection]: + """Initialize the minimum two-rank EP topology used by these tests.""" + owns_process_group = not dist.is_initialized() + owns_model_parallel = not parallel_state.model_parallel_is_initialized() + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + + if owns_process_group: + dist.init_process_group(backend="nccl") + if owns_model_parallel: + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=_EP_SIZE, + expert_tensor_parallel_size=1, + ) + model_parallel_cuda_manual_seed(2026, force_reset_rng=True) + + try: + yield ProcessGroupCollection.use_mpu_process_groups() + finally: + if owns_model_parallel and parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + if owns_process_group and dist.is_initialized(): + dist.destroy_process_group() + + +def _make_adapter( + pg_collection: ProcessGroupCollection, + *, + base_linear_name: str, + input_is_parallel: bool, + use_cpu_initialization: bool = False, +) -> ParallelLinearAdapter: + """Construct a real MCore-backed shared expert adapter.""" + config = ModelParallelConfig( + tensor_model_parallel_size=1, + expert_model_parallel_size=_EP_SIZE, + expert_tensor_parallel_size=1, + params_dtype=torch.float32, + gradient_accumulation_fusion=False, + use_cpu_initialization=use_cpu_initialization, + ) + return ParallelLinearAdapter( + in_features=8, + out_features=8, + dim=4, + base_linear_name=base_linear_name, + activation="identity", + input_is_parallel=input_is_parallel, + is_expert=True, + model_parallel_config=config, + pg_collection=pg_collection, + ) + + +def _assert_identical_across_ep(tensor: torch.Tensor, ep_group: object) -> None: + """Require bitwise equality for one logical tensor on every EP rank.""" + collective_tensor = tensor if tensor.is_cuda else tensor.cuda() + gathered = [torch.empty_like(collective_tensor) for _ in range(_EP_SIZE)] + dist.all_gather(gathered, collective_tensor, group=ep_group) + for replica in gathered[1:]: + torch.testing.assert_close(replica, gathered[0], rtol=0, atol=0) + + +def _set_identical_nonzero_weights(adapter: ParallelLinearAdapter) -> None: + """Make gradients for both LoRA projections nonzero on the first backward.""" + with torch.no_grad(): + for index, parameter in enumerate(adapter.parameters(), start=1): + values = torch.arange(1, parameter.numel() + 1, device=parameter.device, dtype=parameter.dtype) + parameter.copy_(values.reshape_as(parameter) * (0.01 * index)) + + +@pytest.fixture(scope="module") +def ep_pg_collection() -> Iterator[ProcessGroupCollection]: + """Provide one shared two-rank EP topology for this test module.""" + if int(os.environ.get("WORLD_SIZE", "1")) != _EP_SIZE: + pytest.skip("requires a two-rank torch.distributed launch") + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + with _distributed_ep() as pg_collection: + yield pg_collection + + +@pytest.mark.gpu +def test_shared_expert_lora_initialization_is_identical_across_ep( + ep_pg_collection: ProcessGroupCollection, +) -> None: + """Every shared expert LoRA parameter must start bitwise identical across EP.""" + for base_linear_name, input_is_parallel in _ADAPTER_CASES: + adapter = _make_adapter( + ep_pg_collection, + base_linear_name=base_linear_name, + input_is_parallel=input_is_parallel, + ) + for parameter in adapter.parameters(): + _assert_identical_across_ep(parameter.detach(), ep_pg_collection.ep) + + +@pytest.mark.gpu +def test_shared_expert_lora_gradients_are_identical_across_ep( + ep_pg_collection: ProcessGroupCollection, +) -> None: + """EP gradient synchronization must produce identical nonzero gradients.""" + ep_rank = dist.get_rank(group=ep_pg_collection.ep) + for base_linear_name, input_is_parallel in _ADAPTER_CASES: + adapter = _make_adapter( + ep_pg_collection, + base_linear_name=base_linear_name, + input_is_parallel=input_is_parallel, + ) + _set_identical_nonzero_weights(adapter) + inputs = torch.arange(1, 25, device="cuda", dtype=torch.float32).reshape(3, 8) + ep_rank + adapter(inputs).square().sum().backward() + + for parameter in adapter.parameters(): + assert parameter.grad is not None + assert torch.count_nonzero(parameter.grad) > 0 + _assert_identical_across_ep(parameter.grad, ep_pg_collection.ep) + + +@pytest.mark.gpu +def test_shared_expert_lora_cpu_initialization_is_identical_across_ep( + ep_pg_collection: ProcessGroupCollection, +) -> None: + """CPU initialization must be synchronized when EP ranks use different seeds.""" + ep_rank = dist.get_rank(group=ep_pg_collection.ep) + for base_linear_name, input_is_parallel in _ADAPTER_CASES: + torch.manual_seed(2026 + ep_rank) + adapter = _make_adapter( + ep_pg_collection, + base_linear_name=base_linear_name, + input_is_parallel=input_is_parallel, + use_cpu_initialization=True, + ) + for parameter in adapter.parameters(): + assert parameter.device.type == "cpu" + _assert_identical_across_ep(parameter.detach(), ep_pg_collection.ep) diff --git a/tests/unit_tests/peft/test_utils.py b/tests/unit_tests/peft/test_utils.py index fccce8118c0..3ef310340a8 100644 --- a/tests/unit_tests/peft/test_utils.py +++ b/tests/unit_tests/peft/test_utils.py @@ -21,7 +21,7 @@ import math from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import Mock, PropertyMock, patch import pytest import torch @@ -1010,6 +1010,87 @@ def test_parallel_linear_adapter_grouped_expert_shared_adapter_syncs_grad_across torch.testing.assert_close(grad, torch.ones_like(grad)) assert call.kwargs["group"] is ep_group + @patch("megatron.bridge.peft.utils.ColumnParallelLinear") + @patch("megatron.bridge.peft.utils.RowParallelLinear") + def test_parallel_linear_adapter_grouped_expert_shared_adapter_syncs_init_across_ep( + self, mock_row_linear, mock_col_linear, mock_config + ): + """Shared grouped-expert adapters must broadcast every weight across EP.""" + mock_linear_in = Mock() + mock_linear_out = Mock() + mock_linear_in.weight = nn.Parameter(torch.ones(2, 2)) + mock_linear_out.weight = nn.Parameter(torch.ones(2, 2)) + mock_col_linear.side_effect = [mock_linear_in, mock_linear_out] + mock_config._pg_collection = make_mock_pg_collection(ep_size=2) + ep_group = mock_config._pg_collection.ep + + with ( + patch("torch.distributed.is_available", return_value=True), + patch("torch.distributed.is_initialized", return_value=True), + patch("torch.distributed.get_global_rank", return_value=7) as mock_get_global_rank, + patch("torch.distributed.broadcast") as mock_broadcast, + patch.object(torch.Tensor, "is_cuda", new_callable=PropertyMock, return_value=True), + ): + ParallelLinearAdapter( + in_features=2, + out_features=2, + dim=2, + base_linear_name="decoder.layers.0.mlp.experts.linear_fc2", + is_expert=True, + model_parallel_config=mock_config, + ) + + mock_get_global_rank.assert_called_once_with(ep_group, 0) + assert mock_broadcast.call_count == 2 + for call, expected_weight in zip( + mock_broadcast.call_args_list, + (mock_linear_in.weight, mock_linear_out.weight), + strict=True, + ): + assert call.args[0] is expected_weight + assert call.kwargs == {"src": 7, "group": ep_group} + + @patch("megatron.bridge.peft.utils.ColumnParallelLinear") + @patch("megatron.bridge.peft.utils.RowParallelLinear") + def test_parallel_linear_adapter_grouped_expert_shared_adapter_syncs_cpu_init_across_ep( + self, mock_row_linear, mock_col_linear, mock_config + ): + """CPU-initialized shared adapters must stage both weights through NCCL.""" + mock_linear_in = Mock() + mock_linear_out = Mock() + mock_linear_in.weight = nn.Parameter(torch.ones(2, 2)) + mock_linear_out.weight = nn.Parameter(torch.ones(2, 2)) + mock_col_linear.side_effect = [mock_linear_in, mock_linear_out] + mock_config._pg_collection = make_mock_pg_collection(ep_size=2) + ep_group = mock_config._pg_collection.ep + staged_weights = (torch.full((2, 2), 3.0), torch.full((2, 2), 5.0)) + + with ( + patch("torch.distributed.is_available", return_value=True), + patch("torch.distributed.is_initialized", return_value=True), + patch("torch.distributed.get_global_rank", return_value=7), + patch("torch.distributed.get_backend", return_value="nccl"), + patch("torch.distributed.broadcast") as mock_broadcast, + patch("torch.cuda.current_device", return_value=0), + patch.object(torch.Tensor, "to", side_effect=staged_weights) as mock_to, + ): + ParallelLinearAdapter( + in_features=2, + out_features=2, + dim=2, + base_linear_name="decoder.layers.0.mlp.experts.linear_fc2", + is_expert=True, + model_parallel_config=mock_config, + ) + + assert mock_to.call_count == 2 + assert mock_broadcast.call_count == 2 + for call, staged_weight in zip(mock_broadcast.call_args_list, staged_weights, strict=True): + assert call.args[0] is staged_weight + assert call.kwargs == {"src": 7, "group": ep_group} + torch.testing.assert_close(mock_linear_in.weight, staged_weights[0]) + torch.testing.assert_close(mock_linear_out.weight, staged_weights[1]) + @patch("megatron.bridge.peft.utils.ColumnParallelLinear") @patch("megatron.bridge.peft.utils.RowParallelLinear") def test_parallel_linear_adapter_grouped_expert_swiglu_sharded_state_dict_uses_expert_axis( From 8fe66b5cfabb47c82f6f90e816c24632cc5e1246 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sat, 8 Aug 2026 21:36:10 +0200 Subject: [PATCH 61/85] cp: `fix(model): preserve Step3.7 text config fields (5417)` into `r0.6.0` (#5429) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../models/stepfun/configuration_step37.py | 3 + .../stepfun/test_configuration_step37.py | 107 ++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 tests/unit_tests/models/stepfun/test_configuration_step37.py diff --git a/src/megatron/bridge/models/stepfun/configuration_step37.py b/src/megatron/bridge/models/stepfun/configuration_step37.py index 440d15160a6..b3c881a2fe8 100644 --- a/src/megatron/bridge/models/stepfun/configuration_step37.py +++ b/src/megatron/bridge/models/stepfun/configuration_step37.py @@ -88,6 +88,9 @@ class Step37TextConfig(Step35Config): model_type = "step3p5" + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + class Step37Config(PretrainedConfig): """Top-level HF-style config for Step3.7 (the multimodal wrapper).""" diff --git a/tests/unit_tests/models/stepfun/test_configuration_step37.py b/tests/unit_tests/models/stepfun/test_configuration_step37.py new file mode 100644 index 00000000000..db45cdd17d6 --- /dev/null +++ b/tests/unit_tests/models/stepfun/test_configuration_step37.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest + +from megatron.bridge.models.hf_pretrained.safe_config_loader import safe_load_config_with_retry +from megatron.bridge.models.stepfun.step35_bridge import Step35Config +from megatron.bridge.models.stepfun.step37_bridge import Step37Config + + +pytestmark = pytest.mark.unit + + +def _write_config(path, config): + path.mkdir() + (path / "config.json").write_text(json.dumps(config)) + + +def test_step37_config_load_and_roundtrip_preserve_step35_text_fields(tmp_path): + step37_path = tmp_path / "step37" + step37_config = { + "architectures": ["Step3p7ForConditionalGeneration"], + "auto_map": { + "AutoConfig": "configuration_step3p7.Step3p7Config", + "AutoProcessor": "processing_step3.Step3VLProcessor", + "AutoModelForCausalLM": "modeling_step3p7.Step3p7ForConditionalGeneration", + }, + "model_type": "step3p7", + "image_token_id": 128001, + "understand_projector_stride": 2, + "projector_bias": False, + "vision_config": { + "model_type": "perception_encoder", + "image_size": 728, + "patch_size": 14, + "width": 1536, + "layers": 47, + "heads": 16, + }, + "text_config": { + "architectures": ["Step3p5ForCausalLM"], + "model_type": "step3p5", + "hidden_size": 4096, + "max_seq_len": 262144, + "max_position_embeddings": 262144, + "vocab_size": 128896, + "rope_scaling": { + "rope_type": "llama3", + "factor": 2.0, + "original_max_position_embeddings": 131072, + "low_freq_factor": 1.0, + "high_freq_factor": 32.0, + }, + "yarn_only_types": ["full_attention"], + }, + } + _write_config(step37_path, step37_config) + + loaded_step37 = safe_load_config_with_retry(step37_path, max_retries=0) + + assert isinstance(loaded_step37, Step37Config) + assert loaded_step37.auto_map["AutoConfig"] == "configuration_step3p7.Step3p7Config" + assert loaded_step37.max_position_embeddings == 262144 + assert loaded_step37.text_config.max_position_embeddings == 262144 + assert loaded_step37.text_config.yarn_only_types == ["full_attention"] + + roundtrip_path = tmp_path / "step37-roundtrip" + loaded_step37.save_pretrained(roundtrip_path) + reloaded_step37 = safe_load_config_with_retry(roundtrip_path, max_retries=0) + assert reloaded_step37.text_config.max_position_embeddings == 262144 + assert reloaded_step37.text_config.yarn_only_types == ["full_attention"] + loaded_dict = loaded_step37.to_dict() + reloaded_dict = reloaded_step37.to_dict() + loaded_dict.pop("_name_or_path") + reloaded_dict.pop("_name_or_path") + assert reloaded_dict == loaded_dict + + +def test_step35_config_load_preserves_inherited_fields(tmp_path): + step35_path = tmp_path / "step35" + step35_config = { + "architectures": ["Step3p5ForCausalLM"], + "model_type": "step3p5", + "hidden_size": 2048, + "max_position_embeddings": 131072, + "step35_only_flag": "preserved", + } + _write_config(step35_path, step35_config) + + loaded_step35 = safe_load_config_with_retry(step35_path, max_retries=0) + assert isinstance(loaded_step35, Step35Config) + assert loaded_step35.hidden_size == 2048 + assert loaded_step35.max_position_embeddings == 131072 + assert loaded_step35.step35_only_flag == "preserved" From 07d182e7e8df4421e445906b29e6b5f4a2f4a6f1 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Sun, 9 Aug 2026 02:55:44 +0200 Subject: [PATCH 62/85] cp: `[ckpt] fix: Deinterleave low-memory GLU checkpoints (5366)` into `r0.6.0` (#5370) Signed-off-by: Yu Yao Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- src/megatron/bridge/training/checkpointing.py | 39 +++++---- .../training/test_model_load_save.py | 81 +++++++++++++++++++ 2 files changed, 105 insertions(+), 15 deletions(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 781a06aafa8..091d755a970 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -1250,25 +1250,20 @@ def save_checkpoint( "converting to contiguous format for checkpoint" ) if routed_interleave_size is not None or shared_interleave_size is not None: - if len(model) == 1: - state_dict["model"] = _process_state_dict_for_model_glu_interleaving( - state_dict["model"], + use_megatron_fsdp = cfg.ddp is not None and cfg.ddp.use_megatron_fsdp + model_keys = ( + ["model"] + if "model" in state_dict + else sorted(key for key in state_dict if key.startswith("model") and key.removeprefix("model").isdigit()) + ) + for model_key in model_keys: + state_dict[model_key] = _process_state_dict_for_model_glu_interleaving( + state_dict[model_key], routed_interleave_size, shared_interleave_size, interleave=False, - use_megatron_fsdp=cfg.ddp.use_megatron_fsdp, + use_megatron_fsdp=use_megatron_fsdp, ) - else: - for i in range(len(model)): - model_key = "model%d" % i - if model_key in state_dict: - state_dict[model_key] = _process_state_dict_for_model_glu_interleaving( - state_dict[model_key], - routed_interleave_size, - shared_interleave_size, - interleave=False, - use_megatron_fsdp=cfg.ddp.use_megatron_fsdp, - ) # Apply PEFT filtering to preserve the existing adapter-only Megatron # checkpoint behavior. ``also_save_hf_checkpoint=True`` only adds the @@ -2472,6 +2467,20 @@ def _process_state_dict_for_glu_interleaving( num_keys_processed += 1 continue + if isinstance(value, list) and value and all(isinstance(shard, ShardedTensor) for shard in value): + if any(shard.data is None for shard in value): + processed_state_dict[key] = value + continue + shard_sizes = [shard.data.shape[0] for shard in value] + data = torch.cat([shard.data for shard in value], dim=0) + new_data = _apply_glu_interleave_to_tensor_data(data, interleave_size, interleave) + processed_state_dict[key] = [ + replace(shard, data=shard_data, local_shape=shard_data.shape) + for shard, shard_data in zip(value, new_data.split(shard_sizes, dim=0)) + ] + num_keys_processed += 1 + continue + if isinstance(value, ShardedObject): if not isinstance(value.data, torch.Tensor): processed_state_dict[key] = value diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py index ee93303cb51..11349fdc909 100644 --- a/tests/unit_tests/training/test_model_load_save.py +++ b/tests/unit_tests/training/test_model_load_save.py @@ -1057,6 +1057,87 @@ def finalize(self) -> None: callback_manager=None, ) + def test_low_memory_save_deinterleaves_expanded_glu_factory(self, tmp_path): + """Low-memory save must persist canonical contiguous SwiGLU weights.""" + from megatron.core.dist_checkpointing.mapping import ShardedTensor, ShardedTensorFactory + + from megatron.bridge.training.checkpointing import _interleave_glu_tensor + + interleave_size = 2 + key = "decoder.layers.0.mlp.experts.local_experts.0.linear_fc1.weight" + contiguous_weight = torch.arange(32, dtype=torch.float32).reshape(8, 4) + runtime_weight = _interleave_glu_tensor(contiguous_weight, interleave_size) + + module = torch.nn.Module() + module.register_parameter("linear_fc1_weight", torch.nn.Parameter(runtime_weight.clone())) + models = [module] + + provider = GPTModelProvider(num_layers=1, hidden_size=8, num_attention_heads=1) + provider.moe_mlp_glu_interleave_size = interleave_size + + def build_factory( + factory_key: str, + data: torch.Tensor, + replica_id: int, + flattened_range: slice | None, + ) -> list[ShardedTensor]: + assert flattened_range is None + return [ + ShardedTensor.from_rank_offsets(factory_key, chunk, replica_id=replica_id) + for chunk in torch.chunk(data, 2, dim=0) + ] + + factory = ShardedTensorFactory( + key=key, + data=module.linear_fc1_weight.data, + build_fn=build_factory, + merge_fn=lambda shards: torch.cat([shard.data for shard in shards], dim=0), + ) + generated_state = {"model": {key: factory}} + + pg_collection = Mock() + pg_collection.dp_cp = Mock() + pg_collection.tp.rank.return_value = 0 + pg_collection.tp.size.return_value = 1 + pg_collection.pp.rank.return_value = 0 + pg_collection.pp.size.return_value = 1 + + rerun_state_machine = Mock() + rerun_state_machine.state_dict.return_value = {} + captured_state = {} + + def capture_distributed_save(state_dict, *args, **kwargs): + captured_state.update(state_dict) + return None + + with ( + patch("megatron.bridge.training.model_load_save.get_model_config", return_value=provider), + patch("megatron.bridge.training.checkpointing.generate_state_dict", return_value=generated_state), + patch("megatron.bridge.training.checkpointing.get_rng_state", return_value=None), + patch("megatron.bridge.training.checkpointing.get_rerun_state_machine", return_value=rerun_state_machine), + patch("megatron.bridge.training.utils.pg_utils.get_pg_collection", return_value=pg_collection), + patch( + "megatron.bridge.training.checkpointing.dist_checkpointing.save", side_effect=capture_distributed_save + ), + patch("megatron.bridge.training.checkpointing.TorchDistSaveShardedStrategy", return_value=Mock()), + patch("megatron.bridge.training.checkpointing.FullyParallelSaveStrategyWrapper", return_value=Mock()), + patch("megatron.bridge.training.checkpointing.maybe_save_dataloader_state"), + patch("megatron.bridge.training.checkpointing.ensure_directory_exists"), + patch("megatron.bridge.training.checkpointing.fault_tolerance"), + patch("megatron.bridge.training.checkpointing.is_empty_async_queue", return_value=True), + patch("megatron.bridge.training.checkpointing.get_rank_safe", return_value=1), + patch("megatron.bridge.training.checkpointing.is_last_rank", return_value=False), + patch("megatron.bridge.training.checkpointing.print_rank_0"), + patch("torch.distributed.is_initialized", return_value=True), + patch("torch.distributed.get_rank", return_value=1), + patch("torch.distributed.barrier"), + ): + save_megatron_model(models, tmp_path, low_memory_save=True) + + serialized_shards = captured_state["model"][key] + serialized_weight = torch.cat([shard.data for shard in serialized_shards], dim=0) + assert torch.equal(serialized_weight, contiguous_weight) + class TestDtypeFromStr: """Test dtype_from_str function.""" From 9905b291c257370019fe3ad5f14dfb2c24c75ce2 Mon Sep 17 00:00:00 2001 From: Ajay Date: Mon, 10 Aug 2026 07:39:06 -0700 Subject: [PATCH 63/85] Update wandb dependency version to 0.27.2 (#5461) Signed-off-by: Ajay --- docker/common/fw_pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/common/fw_pyproject.toml b/docker/common/fw_pyproject.toml index ca9f9e83d9e..d70ef11915d 100644 --- a/docker/common/fw_pyproject.toml +++ b/docker/common/fw_pyproject.toml @@ -79,7 +79,7 @@ override-dependencies = [ "apache-tvm-ffi==0.1.11", "levenshtein; sys_platform == 'never'", "pytest==8.4.2", - "wandb>=0.25.0", + "wandb==0.27.2", "nvidia-lm-eval==25.11.1", "lilcom==1.8.1", "onnx==1.21.0", # TRT-LLM hard dependency From 13ed82034564f46e64feb63a2bc09eab0f9e29e6 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 10 Aug 2026 12:13:56 -0700 Subject: [PATCH 64/85] fix(model): backport Gemma4 PLE sequence parallel sharding to r0.6.0 (#5450) Signed-off-by: Chen Cui --- .../bridge/models/gemma/modeling_gemma4.py | 21 +- .../h100/active/L1_Launch_recipes_gemma.sh | 1 + .../models/gemma/test_gemma4_modeling.py | 67 ++++- ...emma4_ple_sequence_parallel_distributed.py | 273 ++++++++++++++++++ 4 files changed, 350 insertions(+), 12 deletions(-) create mode 100644 tests/unit_tests/models/gemma/test_gemma4_ple_sequence_parallel_distributed.py diff --git a/src/megatron/bridge/models/gemma/modeling_gemma4.py b/src/megatron/bridge/models/gemma/modeling_gemma4.py index 1c1dda9f3f4..43eb2b05b3a 100644 --- a/src/megatron/bridge/models/gemma/modeling_gemma4.py +++ b/src/megatron/bridge/models/gemma/modeling_gemma4.py @@ -894,12 +894,14 @@ def _attach_ple_modules( ple_vocab = provider.per_layer_embed_vocab_size if ple_dim <= 0 or ple_vocab <= 0: return + tp_group = model.pg_collection.tp model.per_layer_embedding = tp.VocabParallelEmbedding( ple_vocab, n_layers * ple_dim, config=config, init_method=config.init_method, + tp_group=tp_group, ) model.per_layer_model_proj = tp.ColumnParallelLinear( provider.hidden_size, @@ -908,6 +910,7 @@ def _attach_ple_modules( init_method=config.init_method, bias=False, gather_output=True, + tp_group=tp_group, ) model.per_layer_proj_norm = Gemma4RMSNorm(config, ple_dim, eps=provider.layernorm_epsilon) @@ -929,15 +932,23 @@ def _compute_per_layer_inputs( tok_emb = model.per_layer_embedding(input_ids) * (ple_dim**0.5) - if getattr(model.config, "sequence_parallel", False): - from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region as _scatter - - tok_emb = _scatter(tok_emb.transpose(0, 1)).transpose(0, 1) + sequence_parallel = getattr(model.config, "sequence_parallel", False) + if sequence_parallel: + tok_emb = tensor_parallel.scatter_to_sequence_parallel_region( + tok_emb.transpose(0, 1), group=model.pg_collection.tp + ).transpose(0, 1) s_local: int = tok_emb.shape[1] tok_emb = tok_emb.view(b, s_local, n_layers, ple_dim) - mdl_proj, _ = model.per_layer_model_proj(decoder_input.transpose(0, 1)) + if sequence_parallel: + # ColumnParallelLinear's sequence-parallel contract gathers and reduce-scatters + # dimension 0, so keep the input sequence-major through the projection. + mdl_proj, _ = model.per_layer_model_proj(decoder_input) + mdl_proj = tensor_parallel.scatter_to_sequence_parallel_region(mdl_proj, group=model.pg_collection.tp) + mdl_proj = mdl_proj.transpose(0, 1).contiguous() + else: + mdl_proj, _ = model.per_layer_model_proj(decoder_input.transpose(0, 1)) mdl_proj = mdl_proj * (model.config.hidden_size**-0.5) mdl_proj = mdl_proj.view(b, s_local, n_layers, ple_dim) mdl_proj = model.per_layer_proj_norm(mdl_proj) diff --git a/tests/functional_tests/launch_scripts/h100/active/L1_Launch_recipes_gemma.sh b/tests/functional_tests/launch_scripts/h100/active/L1_Launch_recipes_gemma.sh index e27364766a8..5d2223b5330 100755 --- a/tests/functional_tests/launch_scripts/h100/active/L1_Launch_recipes_gemma.sh +++ b/tests/functional_tests/launch_scripts/h100/active/L1_Launch_recipes_gemma.sh @@ -21,4 +21,5 @@ export CUDA_VISIBLE_DEVICES="0,1" # This script tests recipe configurations with their default settings to ensure # they can run basic training without crashes uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/test_groups/recipes/test_gemma3_recipes_pretrain.py +uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/unit_tests/models/gemma/test_gemma4_ple_sequence_parallel_distributed.py coverage combine -q diff --git a/tests/unit_tests/models/gemma/test_gemma4_modeling.py b/tests/unit_tests/models/gemma/test_gemma4_modeling.py index 6b1318e5196..21bf8146cc1 100644 --- a/tests/unit_tests/models/gemma/test_gemma4_modeling.py +++ b/tests/unit_tests/models/gemma/test_gemma4_modeling.py @@ -1223,12 +1223,14 @@ def test_attach_ple_modules_installs_embedding_projection_and_norm(self, monkeyp calls = [] class FakeVocabParallelEmbedding: - def __init__(self, vocab_size, hidden_size, config, init_method): - calls.append(("embedding", vocab_size, hidden_size, config, init_method)) + def __init__(self, vocab_size, hidden_size, config, init_method, tp_group): + calls.append(("embedding", vocab_size, hidden_size, config, init_method, tp_group)) class FakeColumnParallelLinear: - def __init__(self, input_size, output_size, config, init_method, bias, gather_output): - calls.append(("projection", input_size, output_size, config, init_method, bias, gather_output)) + def __init__(self, input_size, output_size, config, init_method, bias, gather_output, tp_group): + calls.append( + ("projection", input_size, output_size, config, init_method, bias, gather_output, tp_group) + ) monkeypatch.setattr( "megatron.core.tensor_parallel.VocabParallelEmbedding", @@ -1238,7 +1240,8 @@ def __init__(self, input_size, output_size, config, init_method, bias, gather_ou "megatron.core.tensor_parallel.ColumnParallelLinear", FakeColumnParallelLinear, ) - model = SimpleNamespace() + tp_group = object() + model = SimpleNamespace(pg_collection=SimpleNamespace(tp=tp_group)) config = _config(init_method="init", layernorm_epsilon=1e-6) provider = SimpleNamespace( num_layers=3, @@ -1254,8 +1257,8 @@ def __init__(self, input_size, output_size, config, init_method, bias, gather_ou assert isinstance(model.per_layer_model_proj, FakeColumnParallelLinear) assert isinstance(model.per_layer_proj_norm, Gemma4RMSNorm) assert calls == [ - ("embedding", 128, 6, config, "init"), - ("projection", 4, 6, config, "init", False, True), + ("embedding", 128, 6, config, "init", tp_group), + ("projection", 4, 6, config, "init", False, True, tp_group), ] def test_compute_per_layer_inputs_combines_token_and_model_projections(self): @@ -1289,6 +1292,56 @@ def forward(self, hidden_states): expected_value = (4.0 * (4**-0.5) + (3**0.5)) * (2.0**-0.5) torch.testing.assert_close(out, torch.full_like(out, expected_value)) + def test_compute_per_layer_inputs_preserves_sequence_major_sp_projection(self, monkeypatch): + tp_group = object() + scatter_shapes = [] + + def fake_scatter(tensor, *, group): + assert group is tp_group + assert tensor.shape[0] % 2 == 0 + scatter_shapes.append(tensor.shape) + return tensor.chunk(2, dim=0)[0].contiguous() + + class FakeEmbedding(torch.nn.Module): + def forward(self, input_ids): + batch, seq = input_ids.shape + return torch.ones(batch, seq, 6) + + class FakeProjection(torch.nn.Module): + def forward(self, hidden_states): + # SP input is already sharded and must remain [sequence, batch, hidden]. + assert hidden_states.shape == (3, 2, 4) + # Emulate ColumnParallelLinear's internal sequence all-gather. + return torch.full((6, 2, 6), 4.0), None + + monkeypatch.setattr( + "megatron.core.tensor_parallel.scatter_to_sequence_parallel_region", + fake_scatter, + ) + model = SimpleNamespace( + config=SimpleNamespace( + per_layer_embed_dim=3, + num_layers=2, + hidden_size=4, + sequence_parallel=True, + ), + pg_collection=SimpleNamespace(tp=tp_group), + per_layer_embedding=FakeEmbedding(), + per_layer_model_proj=FakeProjection(), + per_layer_proj_norm=torch.nn.Identity(), + ) + + out = _compute_per_layer_inputs( + model, + input_ids=torch.ones(2, 6, dtype=torch.long), + decoder_input=torch.zeros(3, 2, 4), + ) + + assert scatter_shapes == [torch.Size((6, 2, 6)), torch.Size((6, 2, 6))] + assert out.shape == (2, 3, 2, 3) + expected_value = (4.0 * (4**-0.5) + (3**0.5)) * (2.0**-0.5) + torch.testing.assert_close(out, torch.full_like(out, expected_value)) + def test_compute_per_layer_inputs_returns_none_without_modules(self): model = SimpleNamespace(per_layer_embedding=None) diff --git a/tests/unit_tests/models/gemma/test_gemma4_ple_sequence_parallel_distributed.py b/tests/unit_tests/models/gemma/test_gemma4_ple_sequence_parallel_distributed.py new file mode 100644 index 00000000000..a99490b83ae --- /dev/null +++ b/tests/unit_tests/models/gemma/test_gemma4_ple_sequence_parallel_distributed.py @@ -0,0 +1,273 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Two-rank Gemma4 PLE sequence-parallel parity and backward regression test. + +Run with: +uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest \ + tests/unit_tests/models/gemma/test_gemma4_ple_sequence_parallel_distributed.py +""" + +import os +from types import SimpleNamespace + +import megatron.core.parallel_state as parallel_state +import pytest +import torch +import torch.distributed as dist +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel import ( + gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region, +) + +from megatron.bridge.models.gemma.gemma4_provider import Gemma4DenseProvider +from megatron.bridge.models.gemma.modeling_gemma4 import ( + _attach_ple_modules, + _compute_per_layer_inputs, +) + + +_TP_SIZE = 2 +_RTOL = 1e-5 +_ATOL = 1e-6 + + +def _full_weights( + *, vocab_size: int, packed_ple_size: int, hidden_size: int, ple_dim: int +) -> tuple[torch.Tensor, ...]: + embedding = torch.arange(vocab_size * packed_ple_size, device="cuda", dtype=torch.float32).reshape( + vocab_size, packed_ple_size + ) + projection = torch.arange(packed_ple_size * hidden_size, device="cuda", dtype=torch.float32).reshape( + packed_ple_size, hidden_size + ) + norm = torch.arange(ple_dim, device="cuda", dtype=torch.float32) / 16 + 0.75 + output = torch.arange(5 * packed_ple_size, device="cuda", dtype=torch.float32).reshape(5, packed_ple_size) + return ( + (embedding.remainder(17) - 8) / 64, + (projection.remainder(13) - 6) / 32, + norm, + (output.remainder(11) - 5) / 32, + ) + + +def _build_ple_model( + *, + tp_size: int, + sequence_parallel: bool, + tp_group: dist.ProcessGroup, + num_layers: int, + ple_dim: int, + hidden_size: int, + vocab_size: int, +) -> SimpleNamespace: + config = Gemma4DenseProvider( + num_layers=num_layers, + hidden_size=hidden_size, + ffn_hidden_size=16, + num_attention_heads=2, + num_query_groups=2, + kv_channels=4, + global_kv_channels=4, + per_layer_embed_vocab_size=vocab_size, + per_layer_embed_dim=ple_dim, + vocab_size=vocab_size, + tensor_model_parallel_size=tp_size, + sequence_parallel=sequence_parallel, + perform_initialization=False, + gradient_accumulation_fusion=False, + params_dtype=torch.float32, + bf16=False, + ) + model = SimpleNamespace(config=config, pg_collection=ProcessGroupCollection(tp=tp_group)) + _attach_ple_modules(model, config, config) + model.per_layer_proj_norm.cuda() + return model + + +def _gather_sequence(local_tensor: torch.Tensor, tp_group: dist.ProcessGroup) -> torch.Tensor: + gathered = gather_from_sequence_parallel_region( + local_tensor.transpose(0, 1).contiguous(), + tensor_parallel_output_grad=False, + group=tp_group, + ) + return gathered.transpose(0, 1).contiguous() + + +def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close(actual, expected, rtol=_RTOL, atol=_ATOL) + + +@pytest.mark.gpu +def test_gemma4_ple_tp2_sp_matches_tp1_forward_backward() -> None: + """TP2/SP PLE values and gradients must match TP1 on identical tensors.""" + if int(os.environ.get("WORLD_SIZE", "1")) != _TP_SIZE: + pytest.skip("requires a two-rank torch.distributed launch") + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + owns_process_group = not dist.is_initialized() + owns_model_parallel = not parallel_state.model_parallel_is_initialized() + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + previous_matmul_precision = torch.get_float32_matmul_precision() + # Avoid TF32 so the tolerance measures only FP32 collective/reduction ordering. + torch.set_float32_matmul_precision("highest") + + if owns_process_group: + dist.init_process_group(backend="nccl") + + try: + if owns_model_parallel: + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=_TP_SIZE, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ) + tp_group = ProcessGroupCollection.use_mpu_process_groups().tp + singleton_groups = [dist.new_group(ranks=[rank]) for rank in range(_TP_SIZE)] + reference_tp_group = singleton_groups[dist.get_rank()] + + num_layers = 3 + ple_dim = 4 + packed_ple_size = num_layers * ple_dim + hidden_size = 8 + vocab_size = 16 + model = _build_ple_model( + tp_size=_TP_SIZE, + sequence_parallel=True, + tp_group=tp_group, + num_layers=num_layers, + ple_dim=ple_dim, + hidden_size=hidden_size, + vocab_size=vocab_size, + ) + reference_model = _build_ple_model( + tp_size=1, + sequence_parallel=False, + tp_group=reference_tp_group, + num_layers=num_layers, + ple_dim=ple_dim, + hidden_size=hidden_size, + vocab_size=vocab_size, + ) + + full_embedding, full_projection, norm_weight, output_weight = _full_weights( + vocab_size=vocab_size, + packed_ple_size=packed_ple_size, + hidden_size=hidden_size, + ple_dim=ple_dim, + ) + with torch.no_grad(): + vocab_start = model.per_layer_embedding.vocab_start_index + vocab_end = model.per_layer_embedding.vocab_end_index + model.per_layer_embedding.weight.copy_(full_embedding[vocab_start:vocab_end]) + output_partition = model.per_layer_model_proj.output_size_per_partition + output_start = dist.get_rank(group=tp_group) * output_partition + output_end = output_start + output_partition + model.per_layer_model_proj.weight.copy_(full_projection[output_start:output_end]) + model.per_layer_proj_norm.weight.copy_(norm_weight) + reference_model.per_layer_embedding.weight.copy_(full_embedding) + reference_model.per_layer_model_proj.weight.copy_(full_projection) + reference_model.per_layer_proj_norm.weight.copy_(norm_weight) + + for batch_size, seq_length in ((1, 4), (2, 6)): + input_ids = torch.arange(batch_size * seq_length, device="cuda").reshape(batch_size, seq_length) + input_ids = input_ids.remainder(vocab_size) + decoder_values = torch.arange( + seq_length * batch_size * hidden_size, + device="cuda", + dtype=torch.float32, + ).reshape(seq_length, batch_size, hidden_size) + decoder_values = (decoder_values.remainder(19) - 9) / 32 + coefficients = torch.arange( + batch_size * seq_length * output_weight.shape[0], device="cuda", dtype=torch.float32 + ).reshape(batch_size, seq_length, output_weight.shape[0]) + coefficients = (coefficients.remainder(7) - 3) / 16 + + for parameter in ( + reference_model.per_layer_embedding.weight, + reference_model.per_layer_model_proj.weight, + reference_model.per_layer_proj_norm.weight, + ): + parameter.grad = None + output_ref = output_weight.detach().clone().requires_grad_(True) + decoder_ref = decoder_values.detach().clone().requires_grad_(True) + ple_ref = _compute_per_layer_inputs(reference_model, input_ids, decoder_ref) + logits_ref = torch.nn.functional.linear(ple_ref.flatten(-2), output_ref) + loss_ref = (logits_ref * coefficients).sum() + loss_ref.backward() + + for parameter in ( + model.per_layer_embedding.weight, + model.per_layer_model_proj.weight, + model.per_layer_proj_norm.weight, + ): + parameter.grad = None + decoder_dist = decoder_values.detach().clone().requires_grad_(True) + local_decoder = scatter_to_sequence_parallel_region(decoder_dist, group=tp_group) + output_dist = output_weight.detach().clone().requires_grad_(True) + ple_local = _compute_per_layer_inputs(model, input_ids, local_decoder) + logits_local = torch.nn.functional.linear(ple_local.flatten(-2), output_dist) + coefficients_local = scatter_to_sequence_parallel_region( + coefficients.transpose(0, 1).contiguous(), group=tp_group + ).transpose(0, 1) + loss_local = (logits_local * coefficients_local).sum() + + _assert_close(_gather_sequence(ple_local.detach(), tp_group), ple_ref.detach()) + _assert_close(_gather_sequence(logits_local.detach(), tp_group), logits_ref.detach()) + loss_dist = loss_local.detach().clone() + dist.all_reduce(loss_dist, group=tp_group) + _assert_close(loss_dist, loss_ref.detach()) + + loss_local.backward() + _assert_close(decoder_dist.grad, decoder_ref.grad) + _assert_close( + model.per_layer_embedding.weight.grad, + reference_model.per_layer_embedding.weight.grad[vocab_start:vocab_end], + ) + _assert_close( + model.per_layer_model_proj.weight.grad, + reference_model.per_layer_model_proj.weight.grad[output_start:output_end], + ) + + norm_grad = model.per_layer_proj_norm.weight.grad.detach().clone() + dist.all_reduce(norm_grad, group=tp_group) + _assert_close(norm_grad, reference_model.per_layer_proj_norm.weight.grad) + output_grad = output_dist.grad.detach().clone() + dist.all_reduce(output_grad, group=tp_group) + _assert_close(output_grad, output_ref.grad) + + with pytest.raises( + AssertionError, + match="First dimension of the tensor should be divisible by tensor parallel size", + ): + scatter_to_sequence_parallel_region(torch.zeros(5, 1, hidden_size, device="cuda"), group=tp_group) + with pytest.raises(AssertionError, match="9 is not divisible by 2"): + _build_ple_model( + tp_size=_TP_SIZE, + sequence_parallel=True, + tp_group=tp_group, + num_layers=3, + ple_dim=3, + hidden_size=hidden_size, + vocab_size=vocab_size, + ) + finally: + torch.set_float32_matmul_precision(previous_matmul_precision) + if owns_model_parallel and parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + if owns_process_group and dist.is_initialized(): + dist.destroy_process_group() From e78f969e418b473636ce694bb8a41f95e919aa5d Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:16:26 -0700 Subject: [PATCH 65/85] cp: `fix(data): collate ragged packed sequence metadata (5400)` into `r0.6.0` (#5426) Signed-off-by: yaoyu-33 --- src/megatron/bridge/data/packing/gpt_sft.py | 14 +++++++++-- .../unit_tests/data/datasets/test_gpt_sft.py | 25 +++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/data/packing/gpt_sft.py b/src/megatron/bridge/data/packing/gpt_sft.py index 0aed186d8ad..0a7fe548f78 100644 --- a/src/megatron/bridge/data/packing/gpt_sft.py +++ b/src/megatron/bridge/data/packing/gpt_sft.py @@ -277,8 +277,18 @@ def collate_fn(self, batch): } if self.return_cu_seqlen: - # Offline packed SFT requires micro-batch size 1, so no sentinel padding is - # needed to batch variable-length metadata. Emit MCore/TE field names directly. + # The DataLoader collates every row assigned to this data-parallel rank before + # the training loop splits that global batch into MBS1 microbatches. Extend the + # boundary arrays with zero-length sequences so ragged rows form tensors while + # retaining the direct MCore/TE cumulative-offset contract. + boundary_batches = [cu_seqlens] + if cu_seqlens_padded is not None: + boundary_batches.append(cu_seqlens_padded) + max_num_boundaries = max(len(boundaries) for rows in boundary_batches for boundaries in rows) + for rows in boundary_batches: + for boundaries in rows: + boundaries.extend([boundaries[-1]] * (max_num_boundaries - len(boundaries))) + seqlens = torch.IntTensor(cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens) seqlens = seqlens[:, 1:] - seqlens[:, :-1] max_seqlen, _ = seqlens.max(dim=1, keepdim=True) diff --git a/tests/unit_tests/data/datasets/test_gpt_sft.py b/tests/unit_tests/data/datasets/test_gpt_sft.py index 176ec68c7af..2fbe0e9539e 100755 --- a/tests/unit_tests/data/datasets/test_gpt_sft.py +++ b/tests/unit_tests/data/datasets/test_gpt_sft.py @@ -305,6 +305,31 @@ def test_collate_fn(self, tmp_path): ] dataset.collate_fn(batch) + def test_collate_fn_supports_different_sequence_counts(self, tmp_path): + """Packed rows with different logical sequence counts collate into one global batch.""" + dataset, _ = get_gpt_sft(tmp_path, dataset_type="packed") + dataset.max_seq_length = 8 + dataset.pad_seq_length_to_mult = 1 + batch = [ + { + "input_ids": np.array([10, 11, 12, 13, 20, 21, 22, 23]), + "seq_boundaries": [0, 4, 8], + "loss_mask": np.ones(8, dtype=np.int64), + }, + { + "input_ids": np.array([30, 31, 32, 33, 34, 35, 36, 37]), + "seq_boundaries": [0, 8], + "loss_mask": np.ones(8, dtype=np.int64), + }, + ] + + processed = dataset.collate_fn(batch) + + assert processed["cu_seqlens_q"].tolist() == [[0, 3, 6, 7], [0, 7, 7, 7]] + assert processed["cu_seqlens_kv"].tolist() == processed["cu_seqlens_q"].tolist() + assert processed["max_seqlen_q"].tolist() == [[3], [7]] + assert processed["max_seqlen_kv"].tolist() == [[3], [7]] + def test_utils_func_packed(self, tmp_path): dataset, _ = get_gpt_sft(tmp_path, dataset_type="packed") From 2bc595f12d067afb602a1dd6fcda8d41b31fc08a Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Mon, 10 Aug 2026 23:20:02 +0200 Subject: [PATCH 66/85] build: AUT-1430 upgrade GitPython to 3.1.58 (#5480) Signed-off-by: svcnemo-autobot --- docker/Dockerfile.fw_final | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docker/Dockerfile.fw_final b/docker/Dockerfile.fw_final index 4a804544695..918a4d4dd47 100644 --- a/docker/Dockerfile.fw_final +++ b/docker/Dockerfile.fw_final @@ -90,7 +90,7 @@ RUN pip install "fastapi==0.139.1" \ "jupyterlab>=4.5.10" \ "jupyter-server>=2.20.0" \ "mistune>=3.3.0" \ - "GitPython>=3.1.55" && \ + "GitPython>=3.1.58" && \ # Address CVE from transitive dependency # Remove the package pulled in by Humming kernels. # Humming-backed vLLM quantization is unavailable without pyelftools. diff --git a/uv.lock b/uv.lock index 926a1ad5380..4d4ebc0b337 100644 --- a/uv.lock +++ b/uv.lock @@ -1277,14 +1277,14 @@ wheels = [ [[package]] name = "gitpython" -version = "3.1.57" +version = "3.1.58" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gitdb" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ba/0d/132ed135c871b6bf91adf16a0e43797cd535b81d4973b5d09291c54fc5ee/gitpython-3.1.57.tar.gz", hash = "sha256:c493ec57c0ef6b19743798b6a5af859c71814b524e7e6f97baa2f8e658961488", size = 225898, upload-time = "2026-07-26T07:33:26.351Z" } +sdist = { url = "https://files.pythonhosted.org/packages/26/d6/5f358ff283325580c2003a6d953aea18cfe10ae87b46f5ebc80fa3a386dc/gitpython-3.1.58.tar.gz", hash = "sha256:621416df10ef3fd0e19fabf9172ddeed0fa704d353d04f194eec56a625a95b22", size = 228498, upload-time = "2026-08-04T15:05:49.47Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/41/6e/2139de986d9c7c3ac86f1f8be43858ce90bdfe2f7175e6c80c650ba15242/gitpython-3.1.57-py3-none-any.whl", hash = "sha256:4ccf7d73c10f5c9e76043fbb2675ac5a1b3ff5b41e648f56bcbed5f63792ecaf", size = 217151, upload-time = "2026-07-26T07:33:24.838Z" }, + { url = "https://files.pythonhosted.org/packages/ec/0c/9d8752098bc442f0726e64aa6135940b3a96809915d1aa4206c1bb97881d/gitpython-3.1.58-py3-none-any.whl", hash = "sha256:d331e722577f0fd7fc1f857419b3ecc07af66282b933d2a4d95f84a042fdd50f", size = 220183, upload-time = "2026-08-04T15:05:48.025Z" }, ] [[package]] From 802b7f636cae1add60242e00afa3d4a6b09a7133 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 00:00:22 +0200 Subject: [PATCH 67/85] cp: `fix(models): use HybridEP for MiniMax-M3 expert parallelism (5448)` into `r0.6.0` (#5473) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- .../minimax-m3/card.yaml | 53 +++++++++---------- scripts/inference/vlm_generation.py | 4 +- scripts/inference/vlm_generation_utils.py | 16 +++++- .../models/minimax_m3/minimax_m3_bridge.py | 5 +- .../bridge/recipes/minimax/h100/minimax_m3.py | 16 +++++- .../minimax_m3/test_minimax_m3_bridge.py | 9 +++- .../minimax/test_minimax_m3_recipes.py | 7 +++ .../utils/test_vlm_generate_utils.py | 12 +++++ 8 files changed, 89 insertions(+), 33 deletions(-) diff --git a/examples/model_verification_cards/minimax-m3/card.yaml b/examples/model_verification_cards/minimax-m3/card.yaml index 1e869a1d5cf..bb1c2edb009 100644 --- a/examples/model_verification_cards/minimax-m3/card.yaml +++ b/examples/model_verification_cards/minimax-m3/card.yaml @@ -9,10 +9,12 @@ summary: > vision-language checkpoint, including the language model, vision tower, projector paths, and lightning-indexer state. The pinned 32-H100 GPU import and GPU export completed their verification gates. A current clean 32-H100 - run also reloaded the persisted full-VLM checkpoint through the public - VLM generation entry point, processed the pinned repository image, generated a - coherent completion containing the exact image-only OCR target, and reached - EOS without token-0 or NUL degeneration. All 59 exported safetensors shards + release-candidate run also reloaded the persisted full-VLM checkpoint through + the public VLM generation entry point, processed the pinned repository image, + and generated a coherent 64-token bounded completion without token-0 or NUL + degeneration. It reached the configured bound before EOS or the image-only + OCR target; that limitation is recorded in the inference result. All 59 + exported safetensors shards are byte-for-byte identical to the pinned source, and Transformers 5.12.1 natively reloads the export without loading errors. CPU import and export remain unverified because the CPU-only attempt exited during CUDA-dependent @@ -138,7 +140,7 @@ items: inference: status: verified precision: bf16 - bridge_commit: 9bd762e26525ab40703f9dac0b8a935b2fc78c19 # pragma: allowlist secret + bridge_commit: 30fae2e8796a7c78f62c9c2cf3b1e249b58fd56d # pragma: allowlist secret command: > ./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 --task vlm-generation @@ -153,25 +155,26 @@ items: explain or reason." --max_new_tokens 64 --trust_remote_code - last_verified: 2026-07-31 + last_verified: 2026-08-08 expected_result: > - One deterministic greedy generation on 32 H100s reloaded the persisted - full-VLM checkpoint at TP1/PP1/EP32/ETP1, processed docs/images/tp1.png, - generated 50 tokens under a 64-token maximum, reached EOS naturally, and - exited successfully. The first five selected token IDs were 200059, 758, - 3100, 9210, and 640; the completion contained zero NUL bytes and did not - exhibit token-0/NUL degeneration. The exact literal completion, - represented with JSON Unicode escapes for the model tags and quotes, was - "\u003cmm:think\u003eThe user wants me to read the green caption near the - top of the image. The green text says \u0022Pre-LayerNorm\u0022 Transformer - Block. The exact quoted term would be - \u0022Pre-LayerNorm\u0022.\u003c\u002fmm:think\u003e\u0022Pre-LayerNorm\u0022". + The 32-H100 TP1/PP1/EP32 run generated exactly 64 tokens under the + 64-token maximum, with finite sampled logits and no token-0/NUL + degeneration. The literal completion was "\u003cmm:think\u003eThe user + wants me to read the green caption near the top of the image and respond + with only the exact quoted term, including capitalization and hyphen." + A matched HybridEP run + kept all observed logits finite, reached EOS, and reproduced the expected + completion. Bridge-created providers and H100 recipes use + flex/hybridep with 16 dispatcher SMs, unfused HybridEP permutation, and + 8 ranks per H100 NVLink domain. Older checkpoints may require the same + overrides after deserialization. Training remains unverified. pretrain: H100: status: unverified precision: bf16 - enabled_features: {} + enabled_features: + moe_dispatcher: hybridep command: null last_verified: null metrics: @@ -180,11 +183,8 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - The public 256-H100 BF16 language-backbone recipe must complete a - bounded 100-step run with finite loss, no skipped or NaN iterations, - all four metrics, a persisted post-setup configuration, and a - reloadable final checkpoint. This item does not claim vision or - projector training. + Complete a 100-step HybridEP language-backbone run with finite loss, + reported metrics, and a reloadable checkpoint. sft: H100: @@ -192,6 +192,7 @@ items: precision: bf16 enabled_features: sequence_packing: offline + moe_dispatcher: hybridep command: null last_verified: null metrics: @@ -200,10 +201,8 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - The public 128-H100 packed BF16 language-backbone recipe must complete a - pinned-data 100-step full-SFT run with finite loss, no skipped or NaN - iterations, all four metrics, and a reloadable final checkpoint. This - item does not claim vision or projector training. + Complete a 100-step packed HybridEP language-backbone SFT run with + finite loss, reported metrics, and a reloadable checkpoint. sft_export_inference: all: diff --git a/scripts/inference/vlm_generation.py b/scripts/inference/vlm_generation.py index f9634f16868..18b45b6a387 100644 --- a/scripts/inference/vlm_generation.py +++ b/scripts/inference/vlm_generation.py @@ -35,6 +35,7 @@ from megatron.core.pipeline_parallel.schedules import get_forward_backward_func from transformers import AutoConfig, AutoProcessor, AutoTokenizer from vlm_generation_utils import ( + decode_generated_tokens, pad_input_ids_to_tp_multiple, patch_kimi_vision_processor, process_image_inputs, @@ -296,6 +297,7 @@ def _disable_mtp(m): # ------------------------------------------------------------------ # Greedy generation loop # ------------------------------------------------------------------ + prompt_length = input_ids_raw.size(1) generated_ids = input_ids_raw.clone() stop_tokens = [tokenizer.eos_token_id] @@ -378,7 +380,7 @@ def _disable_mtp(m): if next_token_ids.item() in stop_tokens: break - generated_text = tokenizer.decode(list(generated_ids[0])) + generated_text = decode_generated_tokens(tokenizer, generated_ids, prompt_length) print_rank_0("======== GENERATED TEXT OUTPUT ========") if args.image_path: print_rank_0(f"Image: {args.image_path}") diff --git a/scripts/inference/vlm_generation_utils.py b/scripts/inference/vlm_generation_utils.py index d86a666fccd..64d67e24bf0 100644 --- a/scripts/inference/vlm_generation_utils.py +++ b/scripts/inference/vlm_generation_utils.py @@ -16,7 +16,7 @@ import io import logging -from typing import Optional +from typing import Any, Optional import torch from PIL import Image @@ -112,6 +112,20 @@ def pad_input_ids_to_tp_multiple(input_ids, tp_size: int, pad_token_id: int = 0) return torch.cat([input_ids, padding], dim=1) +def decode_generated_tokens(tokenizer: Any, generated_ids: torch.Tensor, prompt_length: int) -> str: + """Decode only token IDs produced after the input prompt. + + Args: + tokenizer: Tokenizer used to decode generated IDs. + generated_ids: Batched prompt and generated token IDs. + prompt_length: Number of prompt tokens at the start of ``generated_ids``. + + Returns: + Decoded generated text without the prompt or multimodal placeholders. + """ + return tokenizer.decode(generated_ids[0, prompt_length:].tolist()) + + def load_image(image_path: str) -> Image.Image: """Load an image from URL or file path.""" if image_path.startswith(("http://", "https://")): diff --git a/src/megatron/bridge/models/minimax_m3/minimax_m3_bridge.py b/src/megatron/bridge/models/minimax_m3/minimax_m3_bridge.py index aa7b355b8e3..f0768762152 100644 --- a/src/megatron/bridge/models/minimax_m3/minimax_m3_bridge.py +++ b/src/megatron/bridge/models/minimax_m3/minimax_m3_bridge.py @@ -346,7 +346,10 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MiniMaxM3VLModel # MoE settings — sigmoid routing with expert bias correction and # normalized top-k weights scaled by routed_scaling_factor (DeepSeek-V3 style) provider.moe_grouped_gemm = True - provider.moe_token_dispatcher_type = "alltoall" + provider.moe_token_dispatcher_type = "flex" + provider.moe_flex_dispatcher_backend = "hybridep" + provider.moe_flex_dispatcher_num_sms = 16 + provider.moe_permute_fusion_into_hybridep = False provider.moe_permute_fusion = True provider.moe_router_pre_softmax = False provider.moe_router_score_function = "sigmoid" diff --git a/src/megatron/bridge/recipes/minimax/h100/minimax_m3.py b/src/megatron/bridge/recipes/minimax/h100/minimax_m3.py index 55c57c65402..4bcc85d5a32 100644 --- a/src/megatron/bridge/recipes/minimax/h100/minimax_m3.py +++ b/src/megatron/bridge/recipes/minimax/h100/minimax_m3.py @@ -63,7 +63,10 @@ def minimax_m3_pretrain_256gpu_h100_bf16_config() -> ConfigContainer: cfg.model.transformer_impl = "transformer_engine" cfg.model.attention_backend = None - cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_permute_fusion_into_hybridep = False cfg.model.moe_grouped_gemm = True cfg.model.moe_permute_fusion = True cfg.model.moe_router_force_load_balancing = False @@ -130,6 +133,9 @@ def minimax_m3_pretrain_256gpu_h100_bf16_config() -> ConfigContainer: # Keep the complete process environment visible on the recipe. cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NVLINK_DOMAIN_SIZE": 8, + "USE_MNNVL": 0, } return cfg @@ -167,7 +173,10 @@ def minimax_m3_sft_128gpu_h100_bf16_config() -> ConfigContainer: cfg.model.transformer_impl = "transformer_engine" cfg.model.attention_backend = None - cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_permute_fusion_into_hybridep = False cfg.model.moe_grouped_gemm = True cfg.model.moe_permute_fusion = True cfg.model.moe_router_force_load_balancing = False @@ -198,6 +207,9 @@ def minimax_m3_sft_128gpu_h100_bf16_config() -> ConfigContainer: # Keep the complete process environment visible on the recipe. cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NVLINK_DOMAIN_SIZE": 8, + "USE_MNNVL": 0, } return cfg diff --git a/tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py b/tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py index f51805979a7..40b2929c440 100644 --- a/tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py +++ b/tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py @@ -230,6 +230,10 @@ def test_vlm_provider_can_recover_checkpoint_compatible_text_provider(self, mock assert text_provider.hidden_size == vlm_provider.hidden_size assert text_provider.num_layers == vlm_provider.num_layers assert text_provider.transformer_layer_spec == vlm_provider.transformer_layer_spec + assert text_provider.moe_token_dispatcher_type == "flex" + assert text_provider.moe_flex_dispatcher_backend == "hybridep" + assert text_provider.moe_flex_dispatcher_num_sms == 16 + assert text_provider.moe_permute_fusion_into_hybridep is False assert not hasattr(text_provider, "vision_config") def test_provider_bridge_uses_top_level_embedding_tie_contract(self): @@ -257,7 +261,10 @@ def test_provider_bridge_sets_moe_sigmoid_routing(self, mock_pretrained): assert provider.moe_router_pre_softmax is False assert provider.moe_router_score_function == "sigmoid" assert provider.moe_router_enable_expert_bias is True - assert provider.moe_token_dispatcher_type == "alltoall" + assert provider.moe_token_dispatcher_type == "flex" + assert provider.moe_flex_dispatcher_backend == "hybridep" + assert provider.moe_flex_dispatcher_num_sms == 16 + assert provider.moe_permute_fusion_into_hybridep is False assert provider.moe_router_load_balancing_type == "aux_loss" assert provider.moe_router_topk_scaling_factor == 2.0 assert provider.moe_shared_expert_intermediate_size == 32 diff --git a/tests/unit_tests/recipes/minimax/test_minimax_m3_recipes.py b/tests/unit_tests/recipes/minimax/test_minimax_m3_recipes.py index f37f9b5794e..ed64d8d01ae 100644 --- a/tests/unit_tests/recipes/minimax/test_minimax_m3_recipes.py +++ b/tests/unit_tests/recipes/minimax/test_minimax_m3_recipes.py @@ -52,3 +52,10 @@ def test_text_recipe_uses_text_only_provider(monkeypatch, recipe): cfg = recipe() assert isinstance(cfg.model, _FakeTextProvider) + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_flex_dispatcher_num_sms == 16 + assert cfg.model.moe_permute_fusion_into_hybridep is False + assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 8 + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 8 + assert cfg.env_vars["USE_MNNVL"] == 0 diff --git a/tests/unit_tests/utils/test_vlm_generate_utils.py b/tests/unit_tests/utils/test_vlm_generate_utils.py index af1a63f380d..806b23ae885 100644 --- a/tests/unit_tests/utils/test_vlm_generate_utils.py +++ b/tests/unit_tests/utils/test_vlm_generate_utils.py @@ -41,6 +41,18 @@ _spec.loader.exec_module(vlm_generate_utils) +@pytest.mark.unit +def test_decode_generated_tokens_excludes_prompt_and_uses_integer_ids(): + tokenizer = mock.MagicMock() + tokenizer.decode.return_value = "completion" + generated_ids = torch.tensor([[10, 11, 12, 20, 21]]) + + output = vlm_generate_utils.decode_generated_tokens(tokenizer, generated_ids, prompt_length=3) + + assert output == "completion" + tokenizer.decode.assert_called_once_with([20, 21]) + + @pytest.mark.unit class TestProcessMultiImageInputs: """Tests for ``process_multi_image_inputs``.""" From e6b2cbbfc35d48291686790df2f804bedd378e28 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 04:38:46 +0200 Subject: [PATCH 68/85] cp: `fix(diffusion): calculate WAN FLOPs from runtime geometry (5249)` into `r0.6.0` (#5477) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- .../diffusion/models/wan/wan_provider.py | 53 ++++++++ .../bridge/diffusion/models/wan/wan_step.py | 15 +++ src/megatron/bridge/training/train.py | 15 ++- .../bridge/training/utils/flop_utils.py | 125 +++++++++++++++--- .../diffusion/model/wan/test_wan_provider.py | 34 +++++ .../training/utils/test_flop_utils.py | 102 +++++++++++++- 6 files changed, 322 insertions(+), 22 deletions(-) diff --git a/src/megatron/bridge/diffusion/models/wan/wan_provider.py b/src/megatron/bridge/diffusion/models/wan/wan_provider.py index 6c21accb82e..ad1cbe5115d 100644 --- a/src/megatron/bridge/diffusion/models/wan/wan_provider.py +++ b/src/megatron/bridge/diffusion/models/wan/wan_provider.py @@ -72,6 +72,59 @@ class WanModelProvider(TransformerConfig, ModelProviderMixin[VisionModule]): # text_len: int = 512 text_dim: int = 4096 + def _get_num_floating_point_operations_with_runtime_stats( + self, + *, + batch_size: int, + seqlen_sum: int | None, + seqlen_squared_sum: int | None, + cross_seqlen_sum: int | None, + cross_seqlen_product_sum: int | None, + ) -> int: + """Estimate WAN training FLOPs from the executed latent and context lengths.""" + if batch_size <= 0: + return 0 + if any( + value is None for value in (seqlen_sum, seqlen_squared_sum, cross_seqlen_sum, cross_seqlen_product_sum) + ): + raise ValueError("WAN FLOP accounting requires runtime self- and cross-attention sequence metadata") + + query_tokens = int(seqlen_sum) + query_tokens_squared = int(seqlen_squared_sum) + context_tokens = int(cross_seqlen_sum) + query_context_tokens = int(cross_seqlen_product_sum) + + hidden_size = self.hidden_size + kv_channels = self.kv_channels or hidden_size // self.num_attention_heads + num_query_groups = self.num_attention_heads if self.num_query_groups is None else self.num_query_groups + query_projection_size = kv_channels * self.num_attention_heads + kv_projection_size = kv_channels * num_query_groups + + self_attention_flops = ( + 2 * query_tokens * hidden_size * (query_projection_size + 2 * kv_projection_size) + + 4 * query_tokens_squared * query_projection_size + + 2 * query_tokens * query_projection_size * hidden_size + ) + cross_attention_flops = ( + 2 * query_tokens * hidden_size * query_projection_size + + 4 * context_tokens * self.crossattn_emb_size * kv_projection_size + + 4 * query_context_tokens * query_projection_size + + 2 * query_tokens * query_projection_size * hidden_size + ) + mlp_projection_factor = 6 if self.gated_linear_unit else 4 + mlp_flops = mlp_projection_factor * query_tokens * hidden_size * self.ffn_hidden_size + transformer_flops = self.num_layers * (self_attention_flops + cross_attention_flops + mlp_flops) + + patch_volume = self.patch_temporal * self.patch_spatial**2 + patch_embedding_flops = 2 * query_tokens * self.in_channels * patch_volume * hidden_size + text_embedding_flops = ( + 2 * context_tokens * (self.text_dim * self.crossattn_emb_size + self.crossattn_emb_size**2) + ) + output_head_flops = 2 * query_tokens * hidden_size * self.out_channels * patch_volume + + # Backward costs approximately twice the forward matrix-multiplication work. + return 3 * (transformer_flops + patch_embedding_flops + text_embedding_flops + output_head_flops) + def provide(self, pre_process=None, post_process=None, vp_stage=None) -> WanModel: vp_size = self.virtual_pipeline_model_parallel_size if vp_size: diff --git a/src/megatron/bridge/diffusion/models/wan/wan_step.py b/src/megatron/bridge/diffusion/models/wan/wan_step.py index 93177307832..b705edd2f6a 100644 --- a/src/megatron/bridge/diffusion/models/wan/wan_step.py +++ b/src/megatron/bridge/diffusion/models/wan/wan_step.py @@ -28,6 +28,7 @@ ) from megatron.bridge.training.losses import masked_next_token_loss from megatron.bridge.training.state import GlobalState +from megatron.bridge.training.utils.flop_utils import accumulate_flops_metadata, get_model_chunk_vp_stage logger = logging.getLogger(__name__) @@ -149,6 +150,20 @@ def __call__( batch = wan_data_step(qkv_format, data_iterator) timers("batch-generator").stop() + packed_seq_params = batch.get("packed_seq_params") + if packed_seq_params is not None: + self_attention = packed_seq_params["self_attention"] + cross_attention = packed_seq_params["cross_attention"] + accumulate_flops_metadata( + state, + batch["video_latents"], + vp_stage=get_model_chunk_vp_stage(model), + cu_seqlens=self_attention.cu_seqlens_q_padded, + cu_seqlens_unpadded=self_attention.cu_seqlens_q, + cross_cu_seqlens=cross_attention.cu_seqlens_kv_padded, + cross_cu_seqlens_unpadded=cross_attention.cu_seqlens_kv, + ) + check_for_nan_in_loss = state.cfg.rerun_state_machine.check_for_nan_in_loss check_for_spiky_loss = state.cfg.rerun_state_machine.check_for_spiky_loss diff --git a/src/megatron/bridge/training/train.py b/src/megatron/bridge/training/train.py index e13a2c06469..4a90e617af8 100644 --- a/src/megatron/bridge/training/train.py +++ b/src/megatron/bridge/training/train.py @@ -467,6 +467,8 @@ def train( global_state._flops_seqlen_sum = 0 global_state._flops_seqlen_sq_sum = 0 global_state._flops_vision_patches = 0 + global_state._flops_cross_seqlen_sum = 0 + global_state._flops_cross_seqlen_product_sum = 0 global_state._flops_requires_global_reduce = False ( @@ -579,11 +581,20 @@ def train( # fixed-length stats from the local DP rank; THD batches request one exact SUM # all-reduce over the pure DP group because packed sub-sequence lengths can # differ by rank. - seqlen_sum, seqlen_squared_sum, num_vision_patches = flop_utils.resolve_global_flops_seqlen_stats( + ( + seqlen_sum, + seqlen_squared_sum, + num_vision_patches, + cross_seqlen_sum, + cross_seqlen_product_sum, + ) = flop_utils.resolve_global_flops_runtime_stats( global_state, data_parallel_size=dp_size, vp_size=config.model.virtual_pipeline_model_parallel_size, dp_group=pg_collection.dp, + include_cross_attention_stats=hasattr( + config.model, "_get_num_floating_point_operations_with_runtime_stats" + ), ) num_floating_point_operations_in_batch = flop_utils.num_floating_point_operations( config, @@ -591,6 +602,8 @@ def train( seqlen_sum=seqlen_sum, seqlen_squared_sum=seqlen_squared_sum, num_vision_patches=num_vision_patches, + cross_seqlen_sum=cross_seqlen_sum, + cross_seqlen_product_sum=cross_seqlen_product_sum, ) global_state.train_state.floating_point_operations_so_far += num_floating_point_operations_in_batch num_floating_point_operations_so_far = global_state.train_state.floating_point_operations_so_far diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index 71b884c1a85..54090384c82 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -85,19 +85,21 @@ def _accumulator_to_int(value) -> int: return 0 -def resolve_global_flops_seqlen_stats( +def resolve_global_flops_runtime_stats( state, *, data_parallel_size: int, vp_size: int | None = None, dp_group=None, -) -> tuple[int | None, int | None, int]: + include_cross_attention_stats: bool = False, +) -> tuple[int | None, int | None, int, int | None, int | None]: """Resolve data-parallel-global FLOPS sequence stats from per-rank accumulators. - Reads the three accumulators populated by the forward step + Reads the accumulators populated by the forward step (``_flops_seqlen_sum`` = Σ padded tokens, ``_flops_seqlen_sq_sum`` = Σᵢ sᵢ² - over real sub-sequences, ``_flops_vision_patches``) and reduces them to - global totals across the data-parallel group. + over real sub-sequences, ``_flops_vision_patches``, and optional cross-attention + key/value and query-key products) and reduces them to global totals across the + data-parallel group. Under variable-length (THD packed) training the per-rank ``Σᵢ sᵢ²`` can differ across DP ranks, so a single SUM all-reduce over ``dp_group`` is used @@ -116,16 +118,22 @@ def resolve_global_flops_seqlen_stats( dp_group: Data-parallel process group to SUM-reduce over. Must be the pure DP group (excluding CP) matching ``data_parallel_size`` — CP ranks share the same ``cu_seqlens`` and would double-count. + include_cross_attention_stats: Whether the collective includes the two + optional cross-attention values. This must be uniform across every + rank in ``dp_group``; callers should derive it from model capability, + not local batch contents. Returns: - ``(seqlen_sum, seqlen_squared_sum, num_vision_patches)``. The first two - are ``None`` when no accumulation happened, signalling the caller to fall - back to a fixed-length estimate. ``num_vision_patches`` is ``0`` when no - vision tokens were seen. + ``(seqlen_sum, seqlen_squared_sum, num_vision_patches, + cross_seqlen_sum, cross_seqlen_product_sum)``. Sequence values are ``None`` + when no corresponding accumulation happened. ``num_vision_patches`` is + ``0`` when no vision tokens were seen. """ local_seqlen_sum = _accumulator_to_int(getattr(state, "_flops_seqlen_sum", 0)) local_seqlen_sq_sum = _accumulator_to_int(getattr(state, "_flops_seqlen_sq_sum", 0)) local_vision_patches = _accumulator_to_int(getattr(state, "_flops_vision_patches", 0)) + local_cross_seqlen_sum = _accumulator_to_int(getattr(state, "_flops_cross_seqlen_sum", 0)) + local_cross_seqlen_product_sum = _accumulator_to_int(getattr(state, "_flops_cross_seqlen_product_sum", 0)) _ = vp_size use_all_reduce = ( @@ -137,22 +145,58 @@ def resolve_global_flops_seqlen_stats( ) if use_all_reduce: device = torch.cuda.current_device() if torch.cuda.is_available() else None - stats = torch.tensor( - [local_seqlen_sum, local_seqlen_sq_sum, local_vision_patches], - dtype=torch.long, - device=device, - ) + stats_values = [local_seqlen_sum, local_seqlen_sq_sum, local_vision_patches] + # Preserve the original three-int collective for every non-cross-attention + # caller. Only WAN currently records the two optional cross-attention terms. + if include_cross_attention_stats: + stats_values.extend([local_cross_seqlen_sum, local_cross_seqlen_product_sum]) + stats = torch.tensor(stats_values, dtype=torch.long, device=device) torch.distributed.all_reduce(stats, op=torch.distributed.ReduceOp.SUM, group=dp_group) - seqlen_sum, seqlen_squared_sum, num_vision_patches = (int(x) for x in stats.tolist()) + reduced_stats = [int(x) for x in stats.tolist()] + seqlen_sum, seqlen_squared_sum, num_vision_patches = reduced_stats[:3] + if include_cross_attention_stats: + cross_seqlen_sum, cross_seqlen_product_sum = reduced_stats[3:] + else: + cross_seqlen_sum = 0 + cross_seqlen_product_sum = 0 else: # No process group: extrapolate from the local rank (approximation). seqlen_sum = local_seqlen_sum * data_parallel_size seqlen_squared_sum = local_seqlen_sq_sum * data_parallel_size num_vision_patches = local_vision_patches * data_parallel_size + cross_seqlen_sum = local_cross_seqlen_sum * data_parallel_size + cross_seqlen_product_sum = local_cross_seqlen_product_sum * data_parallel_size if seqlen_sum <= 0: - return None, None, max(num_vision_patches, 0) - return seqlen_sum, seqlen_squared_sum, max(num_vision_patches, 0) + seqlen_sum = None + seqlen_squared_sum = None + if cross_seqlen_sum <= 0: + cross_seqlen_sum = None + cross_seqlen_product_sum = None + return ( + seqlen_sum, + seqlen_squared_sum, + max(num_vision_patches, 0), + cross_seqlen_sum, + cross_seqlen_product_sum, + ) + + +def resolve_global_flops_seqlen_stats( + state, + *, + data_parallel_size: int, + vp_size: int | None = None, + dp_group=None, +) -> tuple[int | None, int | None, int]: + """Resolve the original self-attention and vision FLOP statistics.""" + return resolve_global_flops_runtime_stats( + state, + data_parallel_size=data_parallel_size, + vp_size=vp_size, + dp_group=dp_group, + include_cross_attention_stats=False, + )[:3] def _add_flops_accumulator(state, name: str, delta) -> None: @@ -221,6 +265,8 @@ def accumulate_flops_metadata( cu_seqlens_argmin: torch.Tensor | None = None, cu_seqlens_unpadded: torch.Tensor | None = None, cu_seqlens_unpadded_argmin: torch.Tensor | None = None, + cross_cu_seqlens: torch.Tensor | None = None, + cross_cu_seqlens_unpadded: torch.Tensor | None = None, num_vision_patches: int | torch.Tensor | None = None, ) -> None: """Accumulate per-microbatch FLOPS metadata onto ``state``. @@ -230,7 +276,7 @@ def accumulate_flops_metadata( 0 contributes metadata so model chunking does not multiply the full-model FLOPS estimate. ``None`` and ``0`` both represent the primary/only chunk. - Writes three accumulators consumed by ``train.py`` at end of step: + Writes accumulators consumed by ``train.py`` at end of step: - ``_flops_seqlen_sum``: ``mbs * tokens.shape[1]`` (padded total tokens this microbatch contributes), or ``mbs * config_seq_len`` for dense @@ -247,6 +293,9 @@ def accumulate_flops_metadata( pre-fix value). ``dense_seq_len`` is ``config_seq_len`` when provided, otherwise ``tokens.shape[1]``. - ``_flops_vision_patches``: running total of ``num_vision_patches``. + - ``_flops_cross_seqlen_sum`` and ``_flops_cross_seqlen_product_sum``: + optional cross-attention Σᵢ kᵢ and Σᵢ qᵢkᵢ terms for model-specific + estimators such as WAN. ``num_vision_patches`` is the precomputed number of vision patches in this microbatch (drives the ViT term). It is kept model-agnostic on purpose: the @@ -283,6 +332,28 @@ def accumulate_flops_metadata( _add_flops_accumulator(state, "_flops_seqlen_sum", mbs * dense_seq_len) _add_flops_accumulator(state, "_flops_seqlen_sq_sum", mbs * dense_seq_len**2) + cross_sub_seq_lens = _real_subseq_lengths( + cross_cu_seqlens, + cu_seqlens_unpadded=cross_cu_seqlens_unpadded, + ) + if cross_sub_seq_lens is not None and cross_sub_seq_lens.numel() > 0: + if sub_seq_lens is None or sub_seq_lens.numel() != cross_sub_seq_lens.numel(): + raise ValueError("Cross-attention FLOP metadata requires matching query and key/value sequences") + setattr(state, "_flops_requires_global_reduce", True) + cross_padded_sub_seq_lens = _real_subseq_lengths(cross_cu_seqlens) + if cross_padded_sub_seq_lens is None or cross_padded_sub_seq_lens.numel() == 0: + cross_padded_sub_seq_lens = cross_sub_seq_lens + _add_flops_accumulator( + state, + "_flops_cross_seqlen_sum", + _scalar_sum_for_accumulator(cross_padded_sub_seq_lens), + ) + _add_flops_accumulator( + state, + "_flops_cross_seqlen_product_sum", + _scalar_sum_for_accumulator(sub_seq_lens * cross_sub_seq_lens), + ) + if num_vision_patches is not None: _add_flops_accumulator(state, "_flops_vision_patches", num_vision_patches) @@ -357,6 +428,8 @@ def num_floating_point_operations( seqlen_sum: int | None = None, seqlen_squared_sum: int | None = None, num_vision_patches: int = 0, + cross_seqlen_sum: int | None = None, + cross_seqlen_product_sum: int | None = None, ): """Return the number of floating point operations. @@ -374,7 +447,21 @@ def num_floating_point_operations( result matches the legacy constant-length estimate. num_vision_patches: Total number of vision patches in the batch (before spatial merge). Used to compute ViT encoder FLOPS. + cross_seqlen_sum: Sum of cross-attention key/value sequence lengths. + cross_seqlen_product_sum: Sum of per-sample query and key/value length products. """ + peft = getattr(cfg, "peft", None) + is_lora = isinstance(peft, LoRA) + runtime_flops_estimator = getattr(cfg.model, "_get_num_floating_point_operations_with_runtime_stats", None) + if runtime_flops_estimator is not None and not is_lora: + return runtime_flops_estimator( + batch_size=batch_size, + seqlen_sum=seqlen_sum, + seqlen_squared_sum=seqlen_squared_sum, + cross_seqlen_sum=cross_seqlen_sum, + cross_seqlen_product_sum=cross_seqlen_product_sum, + ) + # Compute effective sequence length from actual values or fall back to config. if seqlen_sum is not None and batch_size > 0: effective_seq_length = seqlen_sum / batch_size @@ -395,8 +482,6 @@ def num_floating_point_operations( else: core_attn_seq_factor = effective_seq_length - peft = getattr(cfg, "peft", None) - is_lora = isinstance(peft, LoRA) # If the model provider has a custom TFLOPS calculation method, use it (non-LoRA only). if not is_lora and hasattr(cfg.model, "_get_num_floating_point_operations"): return cfg.model._get_num_floating_point_operations(batch_size) diff --git a/tests/unit_tests/diffusion/model/wan/test_wan_provider.py b/tests/unit_tests/diffusion/model/wan/test_wan_provider.py index dafd2ce1fd4..e8d7151ccec 100644 --- a/tests/unit_tests/diffusion/model/wan/test_wan_provider.py +++ b/tests/unit_tests/diffusion/model/wan/test_wan_provider.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace + import pytest import torch import torch.nn as nn @@ -20,6 +22,38 @@ import megatron.bridge.diffusion.models.wan.wan_model as wan_model_module from megatron.bridge.diffusion.models.wan.wan_model import WanModel from megatron.bridge.diffusion.models.wan.wan_provider import WanModelProvider +from megatron.bridge.training.utils.flop_utils import num_floating_point_operations + + +def test_wan_provider_uses_runtime_geometry_for_flops(): + provider = WanModelProvider( + num_layers=2, + hidden_size=8, + ffn_hidden_size=16, + num_attention_heads=2, + crossattn_emb_size=6, + in_channels=4, + out_channels=4, + patch_spatial=2, + patch_temporal=1, + text_dim=12, + ) + config = SimpleNamespace(model=provider) + + runtime_stats = { + "batch_size": 2, + "seqlen_sum": 10, + "seqlen_squared_sum": 58, + "cross_seqlen_sum": 6, + "cross_seqlen_product_sum": 34, + } + provider.seq_length = 1024 + old_value = num_floating_point_operations(config, **runtime_stats) + provider.seq_length = 8192 + new_value = num_floating_point_operations(config, **runtime_stats) + + assert old_value == 120_624 + assert new_value == old_value @pytest.mark.parametrize( diff --git a/tests/unit_tests/training/utils/test_flop_utils.py b/tests/unit_tests/training/utils/test_flop_utils.py index a8029628179..04a1e6412ba 100644 --- a/tests/unit_tests/training/utils/test_flop_utils.py +++ b/tests/unit_tests/training/utils/test_flop_utils.py @@ -27,6 +27,7 @@ _packed_data_exists, accumulate_flops_metadata, num_floating_point_operations, + resolve_global_flops_runtime_stats, resolve_global_flops_seqlen_stats, vit_flops, ) @@ -2275,6 +2276,37 @@ def custom(batch_size): # Override must have been invoked twice with the right batch_size args. assert captured == [1, 4], f"Override call log mismatch: {captured}" + def test_runtime_stats_override_short_circuits(self): + model = MockModelConfig() + sentinel = 7_654_321 + captured = {} + + def custom(**kwargs): + captured.update(kwargs) + return sentinel + + model._get_num_floating_point_operations_with_runtime_stats = custom + cfg = MockConfigContainer(model=model) + + assert ( + num_floating_point_operations( + cfg, + batch_size=4, + seqlen_sum=100, + seqlen_squared_sum=2_500, + cross_seqlen_sum=20, + cross_seqlen_product_sum=500, + ) + == sentinel + ) + assert captured == { + "batch_size": 4, + "seqlen_sum": 100, + "seqlen_squared_sum": 2_500, + "cross_seqlen_sum": 20, + "cross_seqlen_product_sum": 500, + } + class _State: """Minimal stand-in for GlobalState — just an attribute bag.""" @@ -2376,6 +2408,41 @@ def test_accumulates_additively_across_microbatches(self): assert state._flops_seqlen_sum == 2 * 128 assert state._flops_seqlen_sq_sum == (32**2 + 96**2) + (64**2 + 64**2) + def test_cross_attention_metadata_uses_matching_query_and_key_lengths(self): + state = _State() + tokens = torch.zeros(1, 12) + query_cu_seqlens = torch.tensor([0, 4, 12]) + query_cu_seqlens_unpadded = torch.tensor([0, 3, 8]) + key_value_cu_seqlens = torch.tensor([0, 4, 12]) + key_value_cu_seqlens_unpadded = torch.tensor([0, 2, 9]) + + accumulate_flops_metadata( + state, + tokens, + cu_seqlens=query_cu_seqlens, + cu_seqlens_unpadded=query_cu_seqlens_unpadded, + cross_cu_seqlens=key_value_cu_seqlens, + cross_cu_seqlens_unpadded=key_value_cu_seqlens_unpadded, + ) + + assert state._flops_seqlen_sum == 12 + assert state._flops_seqlen_sq_sum == 3**2 + 5**2 + assert state._flops_cross_seqlen_sum == 4 + 8 + assert state._flops_cross_seqlen_product_sum == 3 * 2 + 5 * 7 + assert state._flops_requires_global_reduce + + def test_cross_attention_metadata_requires_matching_sequence_counts(self): + state = _State() + tokens = torch.zeros(1, 8) + + with pytest.raises(ValueError, match="matching query and key/value sequences"): + accumulate_flops_metadata( + state, + tokens, + cu_seqlens=torch.tensor([0, 3, 8]), + cross_cu_seqlens=torch.tensor([0, 9]), + ) + @pytest.mark.parametrize("vp_size", [1, 2, 10]) def test_vpp_accumulates_each_logical_microbatch_once(self, vp_size): # MCore's interleaved schedule calls forward_step once for every @@ -2499,12 +2566,16 @@ def test_extrapolates_local_by_dp_size(self): state._flops_seqlen_sum = 1000 state._flops_seqlen_sq_sum = 250_000 state._flops_vision_patches = 64 - seqlen_sum, seqlen_sq_sum, vision = resolve_global_flops_seqlen_stats( + state._flops_cross_seqlen_sum = 32 + state._flops_cross_seqlen_product_sum = 8_000 + seqlen_sum, seqlen_sq_sum, vision, cross_sum, cross_product_sum = resolve_global_flops_runtime_stats( state, data_parallel_size=4, dp_group=None ) assert seqlen_sum == 1000 * 4 assert seqlen_sq_sum == 250_000 * 4 assert vision == 64 * 4 + assert cross_sum == 32 * 4 + assert cross_product_sum == 8_000 * 4 def test_dp_size_one_returns_local(self): state = _State() @@ -2603,8 +2674,37 @@ def fake_all_reduce(stats, op=None, group=None): ) all_reduce.assert_called_once() + assert all_reduce.call_args.args[0].numel() == 3 assert (seqlen_sum, seqlen_sq_sum, vision) == (40, 400, 0) + def test_cross_capability_extends_all_reduce_when_local_stats_are_zero(self, monkeypatch): + state = _State() + state._flops_seqlen_sum = 10 + state._flops_seqlen_sq_sum = 100 + state._flops_cross_seqlen_sum = 0 + state._flops_cross_seqlen_product_sum = 0 + state._flops_requires_global_reduce = True + + def fake_all_reduce(stats, op=None, group=None): + stats.mul_(4) + + all_reduce = MagicMock(side_effect=fake_all_reduce) + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + stats = resolve_global_flops_runtime_stats( + state, + data_parallel_size=4, + dp_group=object(), + include_cross_attention_stats=True, + ) + + all_reduce.assert_called_once() + assert all_reduce.call_args.args[0].numel() == 5 + assert stats == (40, 400, 0, None, None) + @pytest.mark.unit class TestPackedDataExists: From 0ed053d1ef6e827b631cc506f960c1bb08607f2f Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 05:33:53 +0200 Subject: [PATCH 69/85] cp: `feat(data): infer assistant loss mask for gpt-oss Harmony chat template (5389)` into `r0.6.0` (#5447) Signed-off-by: Yu Yao Signed-off-by: yaoyu-33 Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: Chen Cui Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Maanu Grover --- .../bridge/data/conversation_processing.py | 290 +++++++- .../data/collators/test_model_collators.py | 24 +- tests/unit_tests/data/collators/test_sft.py | 4 + .../data/datasets/test_chat_template.py | 8 +- .../data/test_conversation_processing.py | 677 ++++++++++++++++++ 5 files changed, 994 insertions(+), 9 deletions(-) diff --git a/src/megatron/bridge/data/conversation_processing.py b/src/megatron/bridge/data/conversation_processing.py index 64e791f65df..47939425392 100644 --- a/src/megatron/bridge/data/conversation_processing.py +++ b/src/megatron/bridge/data/conversation_processing.py @@ -33,6 +33,10 @@ CHATML_ASSISTANT_ROLE = "assistant" _LEGACY_CHAT_ROLE_ALIASES = {"human": "user", "gpt": CHATML_ASSISTANT_ROLE} +# Upper bound on the rendered form of a non-tokenizable payload that the boundary-token scan +# will tokenize. Sized to admit payloads that print themselves in full — a few hundred frame +# reprs, or a float array below numpy's summarization threshold — and reject the rest. +_MAX_SCANNED_REPR_CHARS = 65536 @dataclass(frozen=True) @@ -980,6 +984,22 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar for role in ("system", "developer", "user", "tool") }, ), + # OpenAI Harmony (gpt-oss). One assistant turn renders as one or more + # channel segments, e.g. ``<|start|>assistant<|channel|>analysis<|message|>... + # <|end|><|start|>assistant<|channel|>final<|message|>...<|return|>``. The + # start marker stops at ``<|start|>assistant`` so the loss span begins at + # ``<|channel|>`` and covers every channel of the turn, analysis included. + # The template emits ``<|start|>`` and the role separately, so the required + # markers use the standalone structural specials rather than a concatenated + # ``<|start|>assistant``. End variants cover message (``<|end|>``), final + # generation (``<|return|>``), and tool call (``<|call|>``) terminators. + ( + ("<|start|>", "<|channel|>", "<|message|>"), + "<|start|>assistant", + "<|return|>", + ("<|end|>", "<|call|>"), + {role: f"<|start|>{role}" for role in ("system", "developer", "user", "tool")}, + ), ) for required_markers, assistant_start, assistant_end, assistant_end_fallbacks, other_role_starts in candidates: if not all(marker in template for marker in required_markers): @@ -1141,8 +1161,14 @@ def _align_assistant_mask_to_padded_ids( if source_ids == target_ids: return source_mask + if len(source_ids) > len(target_ids): + window_start = _rendered_window_start(source_ids, target_ids, tokenizer) + if window_start is None: + return None + return source_mask[window_start : window_start + len(target_ids)] + pad_token_id = getattr(tokenizer, "pad_token_id", None) - if pad_token_id is None or len(source_ids) > len(target_ids): + if pad_token_id is None: return None pad_len = len(target_ids) - len(source_ids) @@ -1155,6 +1181,120 @@ def _align_assistant_mask_to_padded_ids( return None +def _assistant_mask_from_conversation_turns( + example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]], + ids: Sequence[int], + template_owner: Any, + tokenizer: Any, + boundary_config: AssistantMaskBoundaryConfig, + *, + include_content: bool = True, +) -> torch.Tensor | None: + """Build a role-safe mask by locating turns through conversation-prefix renders.""" + conversation = _conversation_from_example(example_or_conversation) + loss_roles = set(boundary_config.loss_roles) + loss_turns = [(turn_index, turn) for turn_index, turn in enumerate(conversation) if turn.get("role") in loss_roles] + if not loss_turns: + return None + + tokenize_kwargs = { + "tokenize": True, + "add_generation_prompt": False, + "return_dict": True, + **chat_template_kwargs_from_example(example_or_conversation), + } + template = _get_chat_template(template_owner) + if template is not None and GENERATION_REGEX.search(template) is not None: + tokenize_kwargs["return_assistant_tokens_mask"] = True + try: + rendered = _apply_tokenized_chat_template(template_owner, conversation, tokenize_kwargs) + rendered_ids = _chat_template_input_ids(rendered) + except (AssertionError, AttributeError, KeyError, TypeError, ValueError): + return None + role_start_tokens = _token_map_from_boundary_config(boundary_config.role_start_tokens) + role_end_tokens = _token_map_from_boundary_config(boundary_config.role_end_tokens) + role_end_token_variants = { + role: [token_ids for raw_token_ids in raw_variants if (token_ids := _as_token_id_list(raw_token_ids))] + for role, raw_variants in boundary_config.role_end_token_variants.items() + } + include_start_roles = set(boundary_config.include_start_tokens_for_roles) + include_end_roles = set(boundary_config.include_end_tokens_for_roles) + rendered_mask = torch.zeros(len(rendered_ids), dtype=torch.float32) + + for turn_index, turn in loss_turns: + role = turn["role"] + start_tokens = role_start_tokens.get(role) + end_patterns = [role_end_tokens.get(role, []), *role_end_token_variants.get(role, [])] + if not start_tokens or not any(end_patterns): + return None + + try: + before_ids = ( + _chat_template_input_ids( + _apply_tokenized_chat_template(template_owner, conversation[:turn_index], tokenize_kwargs) + ) + if turn_index > 0 + else [] + ) + through_ids = ( + rendered_ids + if turn_index + 1 == len(conversation) + else _chat_template_input_ids( + _apply_tokenized_chat_template(template_owner, conversation[: turn_index + 1], tokenize_kwargs) + ) + ) + except (AssertionError, AttributeError, KeyError, TypeError, ValueError): + return None + + turn_search_start = _common_token_prefix_length(before_ids, rendered_ids) + turn_limit = _common_token_prefix_length(through_ids, rendered_ids) + if turn_search_start >= turn_limit: + return None + + role_start, content_start = find_token_span(rendered_ids, start_tokens, turn_search_start) + if role_start < 0 or content_start > turn_limit: + return None + + candidate_ends: list[tuple[int, int, int, list[int]]] = [] + for priority, pattern in enumerate(end_patterns): + if not pattern: + continue + search_start = content_start + while search_start < turn_limit: + end_start, after_end = find_token_span(rendered_ids, pattern, search_start) + if end_start < 0 or after_end > turn_limit: + break + candidate_ends.append((end_start, -priority, after_end, pattern)) + search_start = end_start + 1 + if not candidate_ends: + return None + content_end, _, after_end, _ = max(candidate_ends) + + loss_start = _trim_leading_token_sequences( + rendered_ids, + content_start, + content_end, + boundary_config.trim_leading_token_sequences, + ) + loss_start = _trim_leading_token_ids( + rendered_ids, + loss_start, + content_end, + _as_token_id_list(boundary_config.trim_leading_token_ids), + ) + if include_content: + rendered_mask[loss_start:content_end] = 1.0 + if role in include_start_roles: + rendered_mask[role_start : role_start + len(start_tokens)] = 1.0 + if role in include_end_roles: + rendered_mask[content_end:after_end] = 1.0 + + aligned_mask = _align_assistant_mask_to_padded_ids(rendered_ids, rendered_mask.tolist(), ids, tokenizer) + if aligned_mask is None: + return None + return torch.tensor(aligned_mask, dtype=torch.float32) + + def _token_map_from_boundary_config(token_map: Mapping[str, Sequence[int]] | None) -> dict[str, list[int]]: if token_map is None: return {} @@ -1163,6 +1303,87 @@ def _token_map_from_boundary_config(token_map: Mapping[str, Sequence[int]] | Non } +def _value_contains_token_sequence( + value: Any, + tokenizer: Any, + token_sequences: Sequence[Sequence[int]], +) -> bool | None: + """Check recursively tokenizable payload values for configured boundary sequences.""" + if isinstance(value, (bytes, bytearray, memoryview)): + # A raw buffer reaches the prompt only as an escaped repr literal, so escaping hides + # any marker inside it and this scan can never clear one. Stringifying a decoded image + # would also cost megabytes per example. + return None + if isinstance(value, str): + try: + token_ids = tokenize_text_without_special_tokens(tokenizer, value) + except (AssertionError, AttributeError, KeyError, TypeError, ValueError): + return None + return any(find_token_span(token_ids, token_sequence)[0] >= 0 for token_sequence in token_sequences) + if isinstance(value, Mapping): + results = [_value_contains_token_sequence(item, tokenizer, token_sequences) for item in value.values()] + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + results = [_value_contains_token_sequence(item, tokenizer, token_sequences) for item in value] + elif value is None or isinstance(value, (bool, int, float)): + return False + else: + # Leaves we cannot tokenize directly — PIL images, arrays, tensors, paths. Reporting + # unknown here would disable the boundary-config fallback for every multimodal example, + # so scan the text the payload renders as instead. Jinja emits an object through str() + # standalone but repr() inside a container, so both forms must look clean to report + # clean. This is a heuristic: a payload whose rendered text its own str()/repr() does + # not reveal can still slip through — an array that abbreviates its printout, or an + # opaque container whose repr escapes the newline inside a marker, the same escaping + # that makes raw buffers above undetectable. Reaching either needs a template that + # iterates or attribute-accesses the payload rather than rendering it whole. + try: + if callable(value) or (type(value).__str__ is object.__str__ and type(value).__repr__ is object.__repr__): + # Templates expand callables by introspection (HF renders `tools` entries through + # get_json_schema, pulling in names and docstrings), and an object with neither + # hook has no rendered text at all. Neither can be cleared by stringifying it. + return None + rendered = {str(value), repr(value)} + # str.__len__ rather than len(): __str__ may return a str subclass with its own. + oversized = any(str.__len__(form) > _MAX_SCANNED_REPR_CHARS for form in rendered) + except Exception: + # Rendering an arbitrary payload can raise anything, and callers rely on None to + # fail closed — tokenize_chat_example only recovers from ValueError, so anything + # escaping here kills the dataloader worker instead. + return None + if oversized: + # A payload that prints itself in full can run to megabytes; tokenizing that once + # per example would stall the dataloader. Treat it as unknown provenance instead. + return None + results = [_value_contains_token_sequence(form, tokenizer, token_sequences) for form in rendered] + if any(result is True for result in results): + return True + return None if any(result is None for result in results) else False + + +def _conversation_contains_boundary_tokens( + example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]], + tokenizer: Any, + boundary_config: AssistantMaskBoundaryConfig, +) -> bool | None: + """Return whether rendered conversation payloads contain structural token sequences.""" + token_sequences = [ + *(_token_map_from_boundary_config(boundary_config.role_start_tokens).values()), + *(_token_map_from_boundary_config(boundary_config.role_end_tokens).values()), + *( + token_ids + for raw_variants in boundary_config.role_end_token_variants.values() + for raw_token_ids in raw_variants + if (token_ids := _as_token_id_list(raw_token_ids)) + ), + ] + payloads: list[Any] = [ + {key: value for key, value in turn.items() if key != "role"} + for turn in _conversation_from_example(example_or_conversation) + ] + payloads.append(chat_template_kwargs_from_example(example_or_conversation)) + return _value_contains_token_sequence(payloads, tokenizer, token_sequences) + + def _trim_leading_token_ids( ids: Sequence[int], start: int, @@ -1216,6 +1437,7 @@ def _assistant_mask_from_boundary_config( boundary_config: AssistantMaskBoundaryConfig, *, include_content: bool = True, + conversation_roles: Sequence[str] | None = None, ) -> torch.Tensor | None: role_start_tokens = _token_map_from_boundary_config(boundary_config.role_start_tokens) role_end_tokens = _token_map_from_boundary_config(boundary_config.role_end_tokens) @@ -1253,6 +1475,11 @@ def _assistant_mask_from_boundary_config( if deduped_markers and marker[0] == deduped_markers[-1][0]: continue deduped_markers.append(marker) + if conversation_roles is not None: + expected_loss_roles = [role for role in conversation_roles if role in loss_roles] + observed_loss_roles = [role for _, role, _ in deduped_markers if role in loss_roles] + if observed_loss_roles != expected_loss_roles: + return None include_start_tokens_for_roles = set(boundary_config.include_start_tokens_for_roles) include_end_tokens_for_roles = set(boundary_config.include_end_tokens_for_roles) @@ -1269,11 +1496,17 @@ def _assistant_mask_from_boundary_config( candidate_end_spans = [] for priority, role_end_pattern in enumerate([role_end_tokens[role], *role_end_token_variants.get(role, [])]): - end_start, after_end = find_token_span(ids, role_end_pattern, content_start) - if end_start >= 0 and end_start < next_marker_start: + search_start = content_start + while search_start < next_marker_start: + end_start, after_end = find_token_span(ids, role_end_pattern, search_start) + if end_start < 0 or after_end > next_marker_start: + break candidate_end_spans.append((end_start, priority, after_end)) + search_start = end_start + 1 if not candidate_end_spans: continue + if conversation_roles is not None and len({end_start for end_start, _, _ in candidate_end_spans}) != 1: + return None end_start, _, after_end = min(candidate_end_spans) role_end_span = (end_start, after_end) @@ -1314,19 +1547,66 @@ def build_assistant_loss_mask( """ tokenizer = get_processor_tokenizer(processor) ids = _as_token_id_list(input_ids) + conversation = _conversation_from_example(example_or_conversation) + conversation_roles = [turn.get("role") for turn in conversation] if conversation else None mask = _assistant_mask_from_hf_chat_template(example_or_conversation, ids, processor, tokenizer) if boundary_config is not None: + boundary_scan_is_safe = ( + not conversation + or _conversation_contains_boundary_tokens(example_or_conversation, tokenizer, boundary_config) is False + ) if mask is None or float(mask.sum().item()) == 0.0: - boundary_mask = _assistant_mask_from_boundary_config(ids, boundary_config) + boundary_mask = _assistant_mask_from_conversation_turns( + example_or_conversation, + ids, + tokenizer, + tokenizer, + boundary_config, + ) + if boundary_mask is None and processor is not tokenizer: + boundary_mask = _assistant_mask_from_conversation_turns( + example_or_conversation, + ids, + processor, + tokenizer, + boundary_config, + ) + if boundary_mask is None and boundary_scan_is_safe: + boundary_mask = _assistant_mask_from_boundary_config( + ids, + boundary_config, + conversation_roles=conversation_roles, + ) + if boundary_mask is None and not boundary_scan_is_safe and mask is not None: + mask = None if mask is None or (boundary_mask is not None and float(boundary_mask.sum().item()) > 0.0): mask = boundary_mask else: - boundary_token_mask = _assistant_mask_from_boundary_config( + boundary_token_mask = _assistant_mask_from_conversation_turns( + example_or_conversation, ids, + tokenizer, + tokenizer, boundary_config, include_content=False, ) + if boundary_token_mask is None and processor is not tokenizer: + boundary_token_mask = _assistant_mask_from_conversation_turns( + example_or_conversation, + ids, + processor, + tokenizer, + boundary_config, + include_content=False, + ) + if boundary_token_mask is None and boundary_scan_is_safe: + boundary_token_mask = _assistant_mask_from_boundary_config( + ids, + boundary_config, + include_content=False, + conversation_roles=conversation_roles, + ) if boundary_token_mask is not None: mask = torch.maximum(mask, boundary_token_mask) if mask is None: diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py index 2c035fc5952..b1ad29d9d4d 100644 --- a/tests/unit_tests/data/collators/test_model_collators.py +++ b/tests/unit_tests/data/collators/test_model_collators.py @@ -76,7 +76,13 @@ def __call__(self, text, add_special_tokens=False): "<|im_end|>": [103], "<|im_end|>\n": [103, 104], } - return {"input_ids": mapping.get(text, [1])} + if text in mapping: + return {"input_ids": mapping[text]} + # Non-assistant role markers must tokenize to distinct ids so that content + # (default) never collides with a role-boundary token in the safety check. + if isinstance(text, str) and text.startswith("<|im_start|>"): + return {"input_ids": [199]} + return {"input_ids": [1]} def __init__(self): self.tokenizer = self._Tok() @@ -615,6 +621,8 @@ class _Tok: chat_template = "{% generation %}{{ messages }}{% endgeneration %}" def __call__(self, text, add_special_tokens=False): + if isinstance(text, str) and text.startswith("<|im_start|>"): + return {"input_ids": [199]} return {"input_ids": [1]} def __init__(self): @@ -694,7 +702,11 @@ def __call__(self, text, add_special_tokens=False): "<|im_end|>": [103], "<|im_end|>\n": [103, 104], } - return {"input_ids": mapping.get(text, [42])} + if text in mapping: + return {"input_ids": mapping[text]} + if isinstance(text, str) and text.startswith("<|im_start|>"): + return {"input_ids": [199]} + return {"input_ids": [42]} def __init__(self): self.tokenizer = self._Tok() @@ -735,6 +747,8 @@ class _Tok: chat_template = "{% generation %}{{ messages }}{% endgeneration %}" def __call__(self, text, add_special_tokens=False): + if isinstance(text, str) and text.startswith("<|im_start|>"): + return {"input_ids": [199]} return {"input_ids": [1]} def __init__(self): @@ -1848,7 +1862,11 @@ def __call__(self, texts, padding=True, truncation=True, return_tensors="pt", ** "<|im_end|>": [102], "<|im_end|>\n": [102, 103], } - return {"input_ids": marker_tokens.get(texts, [1])} + if texts in marker_tokens: + return {"input_ids": marker_tokens[texts]} + if texts.startswith("<|im_start|>"): + return {"input_ids": [199]} + return {"input_ids": [1]} self.tokenized_texts = list(texts) max_len = max(len(row) for row in self.tokenized_rows) out = torch.full((len(self.tokenized_rows), max_len), self.pad_token_id, dtype=torch.long) diff --git a/tests/unit_tests/data/collators/test_sft.py b/tests/unit_tests/data/collators/test_sft.py index 0f1fc7962aa..7c34250e688 100644 --- a/tests/unit_tests/data/collators/test_sft.py +++ b/tests/unit_tests/data/collators/test_sft.py @@ -130,6 +130,10 @@ def __call__( return {"input_ids": [102]} if text == "<|im_end|>\n": return {"input_ids": [102, 103]} + if text == "hi": + return {"input_ids": [10]} + if text == "answer": + return {"input_ids": [21, 22]} texts = text if isinstance(text, list) else [text] tokenized = [[100, 10, 102, 103, 101, 21, 22, 102, 103] for _ in texts] diff --git a/tests/unit_tests/data/datasets/test_chat_template.py b/tests/unit_tests/data/datasets/test_chat_template.py index bd9e3707857..d6ae8779f03 100644 --- a/tests/unit_tests/data/datasets/test_chat_template.py +++ b/tests/unit_tests/data/datasets/test_chat_template.py @@ -151,10 +151,13 @@ def __call__(self, text, add_special_tokens=False, **kwargs): mapping = { "<|im_start|>assistant\n": [101], "<|im_start|>user\n": [100], + "<|im_start|>system\n": [110], + "<|im_start|>developer\n": [111], + "<|im_start|>tool\n": [112], "<|im_end|>": [102], "<|im_end|>\n": [102, 103], } - return {"input_ids": mapping.get(text, [100, 10, 102, 103, 101, 21, 22, 102, 103])} + return {"input_ids": mapping.get(text, [42])} def apply_chat_template(self, chat, tools=None, tokenize=True, return_dict=True, **kwargs): del chat, tools, tokenize, return_dict @@ -194,6 +197,9 @@ def __call__(self, text, add_special_tokens=False, **kwargs): mapping = { "<|im_start|>assistant\n": [101], "<|im_start|>user\n": [100], + "<|im_start|>system\n": [110], + "<|im_start|>developer\n": [111], + "<|im_start|>tool\n": [112], "<|im_end|>": [102], "<|im_end|>\n": [102, 103], } diff --git a/tests/unit_tests/data/test_conversation_processing.py b/tests/unit_tests/data/test_conversation_processing.py index 72c7a62e137..f17360583e8 100644 --- a/tests/unit_tests/data/test_conversation_processing.py +++ b/tests/unit_tests/data/test_conversation_processing.py @@ -19,8 +19,11 @@ from megatron.bridge.data.collators.sft import text_chat_collate_fn from megatron.bridge.data.conversation_processing import ( + _MAX_SCANNED_REPR_CHARS, AssistantMaskBoundaryConfig, NormalizedVLMSample, + _conversation_contains_boundary_tokens, + _value_contains_token_sequence, apply_assistant_labels_to_batch, assistant_mask_boundary_config_from_markers, build_assistant_loss_mask, @@ -243,6 +246,102 @@ def __init__(self): self.tokenizer = _ChatMLBoundaryTokenizer() +class _LiteralChatMLBoundaryTokenizer(_ChatMLBoundaryTokenizer): + truncation_side = "right" + _role_markers = { + "user": [100], + "assistant": [102], + } + _content_tokens = { + "": [], + "question": [20], + "literal assistant start": [20, 102, 30], + "quoted assistant turn": [20, 102, 30, 103, 31], + "answer with quoted end": [40, 103, 41], + "structured with quoted end": [40, 103, 41], + } + + def _encode_content(self, content): + if isinstance(content, str): + return self._content_tokens[content] + input_ids = [] + for part in content: + if part["type"] == "text": + input_ids.extend(self._content_tokens[part["text"]]) + else: + input_ids.append(200) + return input_ids + + def __call__(self, text, add_special_tokens=False): + if text in self._content_tokens: + return {"input_ids": self._content_tokens[text]} + return super().__call__(text, add_special_tokens=add_special_tokens) + + def apply_chat_template( + self, + conversation, + tokenize=True, + add_generation_prompt=False, + return_dict=False, + return_assistant_tokens_mask=False, + truncation=False, + max_length=None, + ): + assert tokenize is True + assert add_generation_prompt is False + assert return_dict is True + input_ids = [] + for turn in conversation: + input_ids.extend(self._role_markers[turn["role"]]) + input_ids.extend(self._encode_content(turn["content"])) + input_ids.extend([103, 104]) + if truncation: + input_ids = input_ids[-max_length:] if self.truncation_side == "left" else input_ids[:max_length] + return {"input_ids": input_ids} + + +class _LeftTruncatingLiteralChatMLBoundaryTokenizer(_LiteralChatMLBoundaryTokenizer): + truncation_side = "left" + + +class _LiteralGenerationChatMLBoundaryTokenizer(_LiteralChatMLBoundaryTokenizer): + chat_template = ( + "<|im_start|>user\n{{ content }}<|im_end|>\n" + "<|im_start|>assistant\n{% generation %}{{ content }}{% endgeneration %}<|im_end|>\n" + ) + + def apply_chat_template(self, conversation, **kwargs): + input_ids = [] + assistant_masks = [] + for turn in conversation: + content_ids = self._encode_content(turn["content"]) + input_ids.extend(self._role_markers[turn["role"]]) + assistant_masks.append(0) + input_ids.extend(content_ids) + assistant_masks.extend([int(turn["role"] == "assistant")] * len(content_ids)) + input_ids.extend([103, 104]) + assistant_masks.extend([0, 0]) + if kwargs.get("truncation"): + max_length = kwargs["max_length"] + input_ids = input_ids[:max_length] + assistant_masks = assistant_masks[:max_length] + return {"input_ids": input_ids, "assistant_masks": assistant_masks} + + +class _PrefixUnsupportedLiteralChatMLBoundaryTokenizer(_LiteralChatMLBoundaryTokenizer): + def apply_chat_template(self, conversation, **kwargs): + if len(conversation) != 2: + raise ValueError("prefix rendering is unavailable") + return {"input_ids": [100, 20, 102, 30, 103, 104]} + + +class _PrefixUnsupportedLiteralGenerationChatMLBoundaryTokenizer(_LiteralGenerationChatMLBoundaryTokenizer): + def apply_chat_template(self, conversation, **kwargs): + if len(conversation) != 2: + raise ValueError("prefix rendering is unavailable") + return super().apply_chat_template(conversation, **kwargs) + + class _MoonlightBoundaryTokenizer(_Tokenizer): chat_template = ( "{%- for message in messages -%}" @@ -538,6 +637,463 @@ def test_infer_assistant_mask_boundary_config_from_chatml_template(): assert all(token_variants == [[103]] for token_variants in boundary_config.role_end_token_variants.values()) +def test_chatml_boundary_mask_does_not_treat_literal_control_markers_as_structure(): + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + }, + _LiteralChatMLBoundaryTokenizer(), + ) + + assert tokenized.input_ids.tolist() == [100, 20, 102, 30, 103, 31, 103, 104, 102, 40, 103, 41, 103, 104] + assert tokenized.assistant_mask.tolist() == [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + True, + True, + True, + True, + True, + ] + + +def test_chatml_boundary_mask_remains_role_safe_through_right_truncation(): + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + }, + _LiteralChatMLBoundaryTokenizer(), + max_length=8, + warn_on_all_masked=False, + ) + + assert tokenized.input_ids.tolist() == [100, 20, 102, 30, 103, 31, 103, 104] + assert tokenized.assistant_mask.tolist() == [False] * 8 + + +def test_direct_hf_chat_collation_does_not_train_truncated_user_marker_payload(): + batch = text_chat_collate_fn( + [ + { + "messages": [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + } + ], + _LiteralChatMLBoundaryTokenizer(), + sequence_length=8, + warn_on_all_masked=False, + ) + + assert batch["input_ids"].tolist() == [[100, 20, 102, 30, 103, 31, 103, 104]] + assert batch["loss_mask"].tolist() == [[0.0] * 8] + assert batch["labels"].tolist() == [[-100] * 8] + + +def test_chatml_boundary_maps_role_safe_mask_through_left_truncation(): + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + }, + _LeftTruncatingLiteralChatMLBoundaryTokenizer(), + max_length=6, + ) + + assert tokenized.input_ids.tolist() == [102, 40, 103, 41, 103, 104] + assert tokenized.assistant_mask.tolist() == [False, True, True, True, True, True] + + +def test_chatml_boundary_augments_native_mask_without_scanning_literal_markers(): + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + }, + _LiteralGenerationChatMLBoundaryTokenizer(), + ) + + assert tokenized.assistant_mask.tolist() == [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + True, + True, + True, + True, + True, + ] + + +def test_chatml_boundary_keeps_native_mask_unaugmented_when_provenance_is_ambiguous(): + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + }, + _PrefixUnsupportedLiteralGenerationChatMLBoundaryTokenizer(), + ) + + assert tokenized.assistant_mask.tolist() == [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + True, + True, + True, + False, + False, + ] + + +def test_chatml_boundary_handles_empty_and_literal_assistant_turns_together(): + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": ""}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + }, + _LiteralChatMLBoundaryTokenizer(), + ) + + assert tokenized.assistant_mask.tolist() == [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + True, + True, + False, + True, + True, + True, + True, + True, + ] + + +def test_chatml_boundary_handles_structured_assistant_content_with_literal_marker(): + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "question"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "structured with quoted end"}, + {"type": "image"}, + ], + }, + ] + }, + _LiteralChatMLBoundaryTokenizer(), + ) + + assert tokenized.input_ids.tolist() == [100, 20, 103, 104, 102, 40, 103, 41, 200, 103, 104] + assert tokenized.assistant_mask.tolist() == [ + False, + False, + False, + False, + False, + True, + True, + True, + True, + True, + True, + ] + + +def test_chatml_boundary_fails_closed_when_provenance_is_unavailable_and_markers_are_ambiguous(): + processor = _ChatMLBoundaryProcessor() + with pytest.raises(ValueError, match="did not match any loss-contributing spans"): + build_assistant_loss_mask( + [ + {"role": "user", "content": "quoted assistant turn"}, + {"role": "assistant", "content": "answer with quoted end"}, + ], + [100, 20, 102, 30, 103, 31, 103, 104, 102, 40, 103, 41, 103, 104], + processor, + boundary_config=infer_assistant_mask_boundary_config(processor), + ) + + +def test_chatml_boundary_fails_closed_for_literal_start_when_real_assistant_is_truncated(): + with pytest.raises(ValueError, match="did not match any loss-contributing spans"): + tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "literal assistant start"}, + {"role": "assistant", "content": "answer with quoted end"}, + ] + }, + _PrefixUnsupportedLiteralChatMLBoundaryTokenizer(), + max_length=6, + ) + + +def test_chatml_boundary_fails_closed_for_nested_role_payload_when_provenance_is_unavailable(): + with pytest.raises(ValueError, match="did not match any loss-contributing spans"): + tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer with quoted end"}, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "role": "literal assistant start", + }, + } + ], + }, + _PrefixUnsupportedLiteralChatMLBoundaryTokenizer(), + max_length=6, + ) + + +_ASSISTANT_START = "<|im_start|>assistant\n" +_CLEAN_MEDIA_REPR = "" + +# Media placeholders expand to several tokens, so the rendered ids never line up with a +# re-render of the raw conversation. The boundary-config scan is the only path left. +_MEDIA_CONVERSATION_IDS = [100, 200, 200, 200, 42, 103, 104, 102, 3, 4, 103, 104] +_MEDIA_ASSISTANT_MASK = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0] + + +class _ReprMedia: + """Decoded media payload (PIL image, array, tensor) rendering as `text`.""" + + def __init__(self, text=_CLEAN_MEDIA_REPR): + self._text = text + + def __repr__(self): + return self._text + + +class _StrOnlyMedia(_ReprMedia): + """Payload that defines __str__ but inherits object.__repr__.""" + + __repr__ = object.__repr__ + + def __str__(self): + return self._text + + +class _DivergentMedia(_ReprMedia): + """Payload whose str() hides what repr() — used for container elements — reveals.""" + + def __str__(self): + return _CLEAN_MEDIA_REPR + + +class _RaisingMedia: + """Media payload with a broken repr, e.g. a handle to a closed file.""" + + def __repr__(self): + raise RuntimeError("repr not available") + + +class _CallableMedia(_ReprMedia): + """Lazily-decoded payload exposing __call__, which templates may expand by introspection.""" + + def __call__(self): + return self._text + + +class _SubstringChatMLTokenizer(_ChatMLBoundaryTokenizer): + """Finds markers anywhere in the text, not only when they are the whole string.""" + + def __call__(self, text, add_special_tokens=False): + if _ASSISTANT_START in text and text != _ASSISTANT_START: + return {"input_ids": [42, 102, 42]} + return super().__call__(text, add_special_tokens=add_special_tokens) + + +def _media_conversation(*media, assistant_media=()): + user_content = [{"type": "image", "image": item} for item in media] + user_content.append({"type": "text", "text": "question"}) + assistant_content = [{"type": "image", "image": item} for item in assistant_media] + assistant_content.append({"type": "text", "text": "answer"}) + return [ + {"role": "user", "content": user_content}, + {"role": "assistant", "content": assistant_content}, + ] + + +def _boundary_scan(example, processor): + return _conversation_contains_boundary_tokens( + example, + processor.tokenizer, + infer_assistant_mask_boundary_config(processor), + ) + + +def test_chatml_boundary_scans_conversations_carrying_media_payloads(): + processor = _ChatMLBoundaryProcessor() + example = {"conversation": _media_conversation(_ReprMedia())} + + # False, not None: a media repr that carries no marker is known-clean provenance, which is + # what re-enables the boundary-config fallback. + assert _boundary_scan(example, processor) is False + + mask = build_assistant_loss_mask( + example, + _MEDIA_CONVERSATION_IDS, + processor, + boundary_config=infer_assistant_mask_boundary_config(processor), + ) + + assert mask.tolist() == _MEDIA_ASSISTANT_MASK + + +@pytest.mark.parametrize("media_field", ["media", "assistant_media"]) +def test_chatml_boundary_fails_closed_when_media_payload_renders_a_literal_marker(media_field): + processor = _ChatMLBoundaryProcessor() + marker_media = _ReprMedia(_ASSISTANT_START) + example = { + "conversation": ( + _media_conversation(marker_media) + if media_field == "media" + else _media_conversation(assistant_media=[marker_media]) + ) + } + + assert _boundary_scan(example, processor) is True + + with pytest.raises(ValueError, match="did not match any loss-contributing spans"): + build_assistant_loss_mask( + example, + _MEDIA_CONVERSATION_IDS, + processor, + boundary_config=infer_assistant_mask_boundary_config(processor), + ) + + +@pytest.mark.parametrize("sibling", [_ReprMedia(), object()], ids=["clean", "unknown"]) +def test_chatml_boundary_scan_reports_a_marker_beside_other_media(sibling): + # A detected marker outranks both a clean sibling and an unscannable one. + example = {"conversation": _media_conversation(sibling, _ReprMedia(_ASSISTANT_START))} + + assert _boundary_scan(example, _ChatMLBoundaryProcessor()) is True + + +@pytest.mark.parametrize( + "media", + [_RaisingMedia(), _ReprMedia("x" * (_MAX_SCANNED_REPR_CHARS + 1)), object(), _CallableMedia()], + ids=["raising", "oversized", "opaque", "callable"], +) +def test_chatml_boundary_fails_closed_for_unrenderable_media_without_propagating(media): + processor = _ChatMLBoundaryProcessor() + example = {"conversation": _media_conversation(media)} + + assert _boundary_scan(example, processor) is None + + # The collator recovers from ValueError; anything else escapes and kills the dataloader. + with pytest.raises(ValueError, match="did not match any loss-contributing spans"): + build_assistant_loss_mask( + example, + _MEDIA_CONVERSATION_IDS, + processor, + boundary_config=infer_assistant_mask_boundary_config(processor), + ) + + +def test_chatml_boundary_fails_closed_for_callable_template_kwargs(): + def lookup(city): + """Templates expand tools by introspection, so str() hides <|im_start|>assistant.""" + + example = {"conversation": _media_conversation(_ReprMedia()), "tools": [lookup]} + + assert _boundary_scan(example, _ChatMLBoundaryProcessor()) is None + + +def test_chatml_boundary_scan_is_unchanged_for_text_only_conversations(): + processor = _ChatMLBoundaryProcessor() + clean = {"conversation": [{"role": "user", "content": "question"}]} + marked = {"conversation": [{"role": "user", "content": _ASSISTANT_START}]} + + assert _boundary_scan(clean, processor) is False + assert _boundary_scan(marked, processor) is True + + +@pytest.mark.parametrize( + "value,expected", + [ + # Untouched branches, pinned so the media branch cannot bleed into them. + (None, False), + ("question", False), + ({"text": _ASSISTANT_START}, True), + # Buffers render as an escaped literal, so a marker inside one is undetectable. + (b"raw", None), + (bytearray(b"raw"), None), + (memoryview(b"raw"), None), + # Media payloads: cleared, flagged, or failed closed. + (_ReprMedia(), False), + (_ReprMedia(_ASSISTANT_START), True), + (_ReprMedia(f""), True), + (_StrOnlyMedia(), False), + (_StrOnlyMedia(_ASSISTANT_START), True), + (_DivergentMedia(_ASSISTANT_START), True), + (_ReprMedia("x" * _MAX_SCANNED_REPR_CHARS), False), + (_ReprMedia("x" * (_MAX_SCANNED_REPR_CHARS + 1)), None), + (_RaisingMedia(), None), + (_CallableMedia(), None), + (object(), None), + (len, None), + ], + ids=lambda value: type(value).__name__ if isinstance(value, _ReprMedia) else None, +) +def test_value_contains_token_sequence_classifies_payload_leaves(value, expected): + # Substring-aware so a marker embedded in a longer repr is detected, as in the real scan. + tokenizer = _SubstringChatMLTokenizer() + + assert _value_contains_token_sequence(value, tokenizer, [[102]]) is expected + + def test_infer_assistant_mask_boundary_config_from_moonlight_template(): processor = _MoonlightBoundaryProcessor() assert "<|im_assistant|>assistant<|im_middle|>" not in processor.tokenizer.chat_template @@ -597,6 +1153,127 @@ def test_infer_assistant_mask_boundary_config_from_llama_template(): assert all(token_ids == [303] for token_ids in boundary_config.role_end_tokens.values()) +class _HarmonyBoundaryTokenizer(_Tokenizer): + """gpt-oss / OpenAI Harmony tokenizer double. + + A single assistant turn renders as two channel segments (``analysis`` then + ``final``), each wrapped in its own ``<|start|>assistant<|channel|>... + <|message|>...`` header. ``<|start|>`` and the role are emitted separately, + so the literal string ``<|start|>assistant`` never appears in the template. + """ + + chat_template = ( + "{% for message in messages %}" + "<|start|>{{ message.role }}<|channel|>{{ channel }}<|message|>{{ message.content }}<|end|>" + "{% endfor %}<|return|>" + ) + + _role_ids = {"assistant": 210, "user": 211, "system": 212, "developer": 213, "tool": 214} + + def __call__(self, text, add_special_tokens=False): + mapping = { + "<|start|>assistant": [200, 210], + "<|start|>user": [200, 211], + "<|start|>system": [200, 212], + "<|start|>developer": [200, 213], + "<|start|>tool": [200, 214], + "<|return|>": [204], + "<|end|>": [203], + "<|call|>": [205], + "question": [20], + "reasoning": [30, 31], + "answer": [40, 41], + } + return {"input_ids": mapping.get(text, [42])} + + def apply_chat_template( + self, + conversation, + tokenize=True, + add_generation_prompt=False, + return_dict=False, + **kwargs, + ): + assert tokenize is True + assert add_generation_prompt is False + assert return_dict is True + input_ids = [] + for turn in conversation: + role_id = self._role_ids[turn["role"]] + content_ids = self(turn["content"])["input_ids"] + if turn["role"] == "assistant": + # Chain-of-thought in the analysis channel, then the user-facing final channel. + input_ids.extend([200, role_id, 201, 220, 202] + self("reasoning")["input_ids"] + [203]) + input_ids.extend([200, role_id, 201, 221, 202] + content_ids + [204]) + else: + input_ids.extend([200, role_id, 202] + content_ids + [203]) + return {"input_ids": input_ids} + + +class _HarmonyBoundaryProcessor(_Processor): + def __init__(self): + super().__init__() + self.tokenizer = _HarmonyBoundaryTokenizer() + + +def test_infer_assistant_mask_boundary_config_from_gpt_oss_harmony_template(): + processor = _HarmonyBoundaryProcessor() + # The concatenated assistant header is not a literal substring of the template; + # detection must rely on the standalone structural specials. + assert "<|start|>assistant" not in processor.tokenizer.chat_template + boundary_config = infer_assistant_mask_boundary_config(processor) + + assert boundary_config is not None + # Start markers stop at the role header so the loss span begins at <|channel|>, + # covering the analysis channel as well as the final channel. + assert boundary_config.role_start_tokens == { + "assistant": [200, 210], + "system": [200, 212], + "developer": [200, 213], + "user": [200, 211], + "tool": [200, 214], + } + assert all(token_ids == [204] for token_ids in boundary_config.role_end_tokens.values()) + # <|end|> (message) and <|call|> (tool call) terminators back up the <|return|> primary. + assert all(variants == [[203], [205]] for variants in boundary_config.role_end_token_variants.values()) + + +def test_gpt_oss_harmony_assistant_mask_covers_analysis_and_final_channels(): + processor = _HarmonyBoundaryProcessor() + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + }, + processor, + ) + + user_block = [200, 211, 202, 20, 203] # <|start|>user<|message|>question<|end|> + analysis_block = [200, 210, 201, 220, 202, 30, 31, 203] # ...<|channel|>analysis<|message|>reasoning<|end|> + final_block = [200, 210, 201, 221, 202, 40, 41, 204] # ...<|channel|>final<|message|>answer<|return|> + assert tokenized.input_ids.tolist() == user_block + analysis_block + final_block + + mask = tokenized.assistant_mask.tolist() + assert len(mask) == len(user_block + analysis_block + final_block) + + # The user turn and the assistant's priming role header (<|start|>assistant) carry no loss. + priming_header_len = 2 + assert not any(mask[: len(user_block) + priming_header_len]) + # Everything from the first <|channel|> through the final <|return|> is trained, so the + # analysis chain-of-thought and the final answer both contribute to the loss. + assert all(mask[len(user_block) + priming_header_len :]) + + # Pin the reasoning (analysis channel) and answer (final channel) token positions so a + # regression that drops the analysis channel from the loss fails loudly here. + reasoning_positions = [len(user_block) + 5, len(user_block) + 6] + answer_positions = [len(user_block) + len(analysis_block) + 5, len(user_block) + len(analysis_block) + 6] + assert tokenized.input_ids[reasoning_positions].tolist() == [30, 31] + assert tokenized.input_ids[answer_positions].tolist() == [40, 41] + assert all(mask[position] for position in reasoning_positions + answer_positions) + + @pytest.mark.parametrize("column", ["messages", "conversation", "conversations"]) def test_shared_chat_preprocessing_supports_all_declared_conversation_columns(column): turns = [ From 0494f0951fc28ccd23882d75027e8b218e820912 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 05:34:04 +0200 Subject: [PATCH 70/85] cp: `fix(scripts): preserve FHT and fetch missing MCore commits (5478)` into `r0.6.0` (#5481) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- pyproject.toml | 1 + scripts/switch_mcore.sh | 17 +++++++++++++++++ uv.lock | 2 ++ 3 files changed, 20 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 5c3bfacfc0f..395d08f7e05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,7 @@ dependencies = [ "torch>=2.6.0", "flashinfer-python==0.6.8.post1", # keep flashinfer-python and flashinfer-cubin in lockstep: "flashinfer-cubin==0.6.8.post1", # the runtime version-match check fails otherwise + "fast-hadamard-transform>=1.1.0", ] diff --git a/scripts/switch_mcore.sh b/scripts/switch_mcore.sh index b97e85a01d3..2ca0bbfb816 100755 --- a/scripts/switch_mcore.sh +++ b/scripts/switch_mcore.sh @@ -72,6 +72,22 @@ get_current_commit() { git -C "$SUBMODULE_PATH" rev-parse HEAD 2>/dev/null || echo "unknown" } +ensure_commit() { + local target_name="$1" + local target_commit="$2" + + if git -C "$SUBMODULE_PATH" cat-file -e "${target_commit}^{commit}" 2>/dev/null; then + return 0 + fi + + log "Pinned $target_name commit ${target_commit:0:12} is not available locally. Fetching from origin..." + git -C "$SUBMODULE_PATH" fetch --no-tags origin "$target_commit" \ + || err "Failed to fetch $target_name commit $target_commit from origin" + + git -C "$SUBMODULE_PATH" cat-file -e "${target_commit}^{commit}" 2>/dev/null \ + || err "Fetched from origin, but $target_name commit $target_commit is still unavailable" +} + do_status() { check_submodule local current @@ -106,6 +122,7 @@ do_switch() { log "Before: ${current:0:12}" log "Switching to $target_name commit: ${target_commit:0:12}" + ensure_commit "$target_name" "$target_commit" git -C "$SUBMODULE_PATH" checkout "$target_commit" || err "Failed to checkout $target_commit" log "After: $(get_current_commit | cut -c1-12)" diff --git a/uv.lock b/uv.lock index 4d4ebc0b337..4acf498cddc 100644 --- a/uv.lock +++ b/uv.lock @@ -2091,6 +2091,7 @@ dependencies = [ { name = "datasets" }, { name = "diffusers" }, { name = "einops" }, + { name = "fast-hadamard-transform" }, { name = "flash-linear-attention" }, { name = "flashinfer-cubin" }, { name = "flashinfer-python" }, @@ -2189,6 +2190,7 @@ requires-dist = [ { name = "datasets", specifier = ">=2.20.0" }, { name = "diffusers", specifier = ">=0.36.0" }, { name = "einops" }, + { name = "fast-hadamard-transform", specifier = ">=1.1.0" }, { name = "flash-linear-attention" }, { name = "flashinfer-cubin", specifier = "==0.6.8.post1" }, { name = "flashinfer-python", specifier = "==0.6.8.post1" }, From dd61c0c69721c9e1cb2183116cc21b8b221b4512 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 10 Aug 2026 20:43:37 -0700 Subject: [PATCH 71/85] cp: `fix(models): use HybridEP for Kimi-K2.5 inference (5421)` into `r0.6.0` (#5485) Signed-off-by: yaoyu-33 Signed-off-by: Chen Cui Co-authored-by: yaoyu-33 --- examples/models/kimi/kimi_k25_vl/README.md | 6 ++++++ examples/models/kimi/kimi_k25_vl/slurm_inference.sh | 3 +++ src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py | 5 ++++- tests/unit_tests/models/kimi_vl/test_kimi_k25_vl_bridge.py | 4 ++++ 4 files changed, 17 insertions(+), 1 deletion(-) diff --git a/examples/models/kimi/kimi_k25_vl/README.md b/examples/models/kimi/kimi_k25_vl/README.md index 055fec3fa57..f2f8d932473 100644 --- a/examples/models/kimi/kimi_k25_vl/README.md +++ b/examples/models/kimi/kimi_k25_vl/README.md @@ -78,6 +78,12 @@ See [slurm_inference.sh](slurm_inference.sh) for configuration details. Note: - `--trust_remote_code` is required for Kimi-K2.5 models. +- Bridge-created Kimi-K2.5 providers use the HybridEP dispatcher with 16 + communication SMs and unfused HybridEP permutation. On 8-GPU H100 nodes, set + `NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=8`, `NVLINK_DOMAIN_SIZE=8`, and + `USE_MNNVL=0`. A native Megatron checkpoint can retain the dispatcher + configuration with which it was saved, so loading an older checkpoint may + require applying the same overrides to the loaded model configuration. - Use `--pp_layout` to specify custom pipeline layouts (e.g. `--pp_layout "Et*15|t*15|t*16|t*15L"` for PP=4). - You can optionally pass `--megatron_model_path` to use a pre-converted diff --git a/examples/models/kimi/kimi_k25_vl/slurm_inference.sh b/examples/models/kimi/kimi_k25_vl/slurm_inference.sh index 16da7566604..fcb17ef896d 100644 --- a/examples/models/kimi/kimi_k25_vl/slurm_inference.sh +++ b/examples/models/kimi/kimi_k25_vl/slurm_inference.sh @@ -57,6 +57,9 @@ MAX_NEW_TOKENS=200 export TORCH_NCCL_AVOID_RECORD_STREAMS=1 export NCCL_NVLS_ENABLE=0 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=8 +export NVLINK_DOMAIN_SIZE=8 +export USE_MNNVL=0 echo "======================================" echo "Kimi-K2.5-VL Inference" diff --git a/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py b/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py index fa137521780..c5b5f9b0049 100644 --- a/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +++ b/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py @@ -86,7 +86,10 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> KimiK25VLModelPr # MoE settings provider.moe_grouped_gemm = True provider.moe_router_pre_softmax = True - provider.moe_token_dispatcher_type = "alltoall" + provider.moe_token_dispatcher_type = "flex" + provider.moe_flex_dispatcher_backend = "hybridep" + provider.moe_flex_dispatcher_num_sms = 16 + provider.moe_permute_fusion_into_hybridep = False provider.moe_router_load_balancing_type = "seq_aux_loss" provider.moe_shared_expert_overlap = True provider.moe_router_score_function = "sigmoid" diff --git a/tests/unit_tests/models/kimi_vl/test_kimi_k25_vl_bridge.py b/tests/unit_tests/models/kimi_vl/test_kimi_k25_vl_bridge.py index 9e45c3f52ac..1fff50a3185 100644 --- a/tests/unit_tests/models/kimi_vl/test_kimi_k25_vl_bridge.py +++ b/tests/unit_tests/models/kimi_vl/test_kimi_k25_vl_bridge.py @@ -175,6 +175,10 @@ def test_provider_bridge_moe_config(self, kimi_bridge, mock_hf_pretrained): assert provider.moe_router_topk == 8 assert provider.moe_router_score_function == "sigmoid" assert provider.moe_router_enable_expert_bias is True + assert provider.moe_token_dispatcher_type == "flex" + assert provider.moe_flex_dispatcher_backend == "hybridep" + assert provider.moe_flex_dispatcher_num_sms == 16 + assert provider.moe_permute_fusion_into_hybridep is False def test_provider_bridge_moe_layer_freq(self, kimi_bridge, mock_hf_pretrained): """Test MoE layer frequency computation.""" From e790259ea1ce3cd727739b4ebdf69cecf95ac7dc Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 06:29:10 +0200 Subject: [PATCH 72/85] cp: `fix(model): handle disabled GLM-5 MTP export (5445)` into `r0.6.0` (#5488) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: yaoyu-33 --- .../model_verification_cards/glm5/card.yaml | 232 ++++++++++++++++++ .../bridge/models/glm_moe_dsa/glm5_bridge.py | 7 +- .../models/glm_moe_dsa/test_glm5_bridge.py | 20 +- 3 files changed, 255 insertions(+), 4 deletions(-) create mode 100644 examples/model_verification_cards/glm5/card.yaml diff --git a/examples/model_verification_cards/glm5/card.yaml b/examples/model_verification_cards/glm5/card.yaml new file mode 100644 index 00000000000..8270b865b0f --- /dev/null +++ b/examples/model_verification_cards/glm5/card.yaml @@ -0,0 +1,232 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: glm5 +summary: > + Performance disclaimer: this model has not been performance-tuned; reported + timing and throughput metrics are sanity checks, not optimized performance + results. Model-level verification covers deterministic HybridEP inference + and exact distributed GPU import/export of the pinned GLM-5 Hugging Face + checkpoint's 78-layer inference graph. The checkpoint's appended MTP + auxiliary layer is intentionally disabled and excluded from those claims. + CPU conversion, Hugging Face/Megatron forward parity, training, checkpoint + resume, and post-SFT export/reload verification are deferred. +verification_index: + model_level: + verified: [hf_to_megatron_gpu, megatron_to_hf_gpu, inference] + unverified: + - hf_to_megatron_cpu + - megatron_to_hf_cpu + - manual_forward_pass + training: + H100: + unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] +model: + hf_id: zai-org/GLM-5 + hf_revision: 4e6698ba8e85059d749020e3c4d2123719f23926 # pragma: allowlist secret + architecture: GlmMoeDsaForCausalLM + min_transformers_version: "5.12.1" +verification_environment: + base_container: nvcr.io/nvidia/nemo:26.08.rc3 + bridge_commit: 5db9c262411cf4561d44cdf76a134b423fdaacc2 # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + A CPU import of the pinned Hugging Face revision must complete every + mapping and create a reloadable Megatron checkpoint. This workflow is + deferred by this card. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 4 --gpus-per-node 8 --hf-model zai-org/GLM-5 + --hf-revision 4e6698ba8e85059d749020e3c4d2123719f23926 + --megatron-path work/model-verification/glm5/gpu-megatron + --torch-dtype bfloat16 --tp 1 --pp 2 --ep 8 --etp 2 + --distributed-timeout-minutes 60 --low-memory-save + last_verified: 2026-08-10 + expected_result: > + The pinned 32-H100 import exits successfully at TP1/PP2/EP8/ETP2, + completes all 6,201 distributed mapping tasks, and persists a reloadable + iter_0000000 checkpoint. Reload plus exact export projection covers all + 59,079 tensors and 1,487,822,475,264 tensor-payload bytes in the 78-layer + inference graph with matching keys, shapes, dtypes, and values. The 791 + source tensors under model.layers.78 belong only to the intentionally + disabled appended MTP auxiliary layer and are outside this item. + + megatron_to_hf_cpu: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + A CPU export from a verified Megatron checkpoint must preserve every + mapped tensor and strictly reload in Transformers. This workflow is + deferred by this card. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + bridge_commit: 7e2181dcc5d09a1a859dfb0c655c9a1cd486e0ac # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 4 --gpus-per-node 8 --hf-model zai-org/GLM-5 + --hf-revision 4e6698ba8e85059d749020e3c4d2123719f23926 + --megatron-path work/model-verification/glm5/gpu-megatron/iter_0000000 + --hf-path work/model-verification/glm5/gpu-hf-export + --torch-dtype bfloat16 --tp 1 --pp 2 --ep 8 --etp 2 + --distributed-timeout-minutes 60 --distributed-save + --save-every-n-ranks 1 + last_verified: 2026-08-10 + expected_result: > + The 32-H100 distributed export exits successfully after all 6,201 + mapping tasks and writes 280 safetensors shards. Of those, 278 are + byte-for-byte identical to the pinned source shards; the two shards that + also contain excluded MTP keys match all 63 common tensors exactly. The + composite audit covers all 59,079 inference-graph tensors and + 1,487,822,475,264 tensor-payload bytes with zero key, shape, dtype, or + value mismatch. Transformers 5.12.1 strictly reloads the output as + GlmMoeDsaForCausalLM with 1,629 state tensors, 743,911,199,232 + parameters, and no loading discrepancies. + + manual_forward_pass: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + A pinned first-token Hugging Face versus Megatron comparison must match + the argmax token and meet the card's logit-correlation threshold. The + full Hugging Face reference pass is deferred by this card. + + inference: + status: verified + precision: bf16 + bridge_commit: b58afee863f40b64304373830ae79f3f61a1af19 # pragma: allowlist secret + command: > + ./scripts/inference/infer.sh --nodes 8 --gpus-per-node 8 + --task legacy-full-prefix-generation + --hf_model_path zai-org/GLM-5 + --hf-revision 4e6698ba8e85059d749020e3c4d2123719f23926 + --prompt "What is artificial intelligence?" --max_new_tokens 20 + --legacy-full-prefix --tp 1 --pp 2 --ep 32 --etp 1 + last_verified: 2026-08-10 + expected_result: |- + One deterministic BF16 run loads the pinned Hugging Face model directly + into Megatron across 64 H100 ranks at TP1/PP2/EP32/ETP1 and uses the + flex HybridEP dispatcher. The inference graph intentionally excludes the + checkpoint's appended MTP auxiliary layer. It exits successfully after + exactly 20 new greedy tokens. The literal completion, including its + leading space, was: + " How is artificial intelligence going to change our lives?" + + pretrain: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A public GLM-5 H100 recipe must complete a bounded 100-step run with + finite loss, no skipped or NaN iterations, all four metrics, and a + reloadable final checkpoint. Training is deferred by this card. + + sft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A pinned-data 100-step full-SFT run must finish with finite loss, no + skipped or NaN iterations, all four metrics, and a reloadable final + checkpoint. Training is deferred by this card. + + sft_export_inference: + H100: + status: unverified + precision: bf16 + depends_on: sft + commands: null + last_verified: null + expected_result: > + A verified SFT checkpoint must export to Hugging Face, strictly reload, + and produce deterministic checkpoint-backed inference. Training and + post-SFT export verification are deferred by this card. + + sft_long_context: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A dedicated packed long-context SFT run must complete a bounded run + with finite loss, no skipped or NaN iterations, all four metrics, and + a reloadable checkpoint. Training is deferred by this card. + + peft: + H100: + status: unverified + precision: bf16 + enabled_features: {} + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A GLM-5 PEFT recipe with an audited adapter target set must complete a + bounded run with finite loss, all four metrics, and a reloadable + adapter checkpoint. Training is deferred by this card. + + checkpoint_resume: + H100: + status: unverified + precision: bf16 + depends_on: pretrain + command: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: null + expected_result: > + A direct continuation must restore model, optimizer, scheduler, + data-order, and RNG state, match declared sentinel losses, and save a + reloadable checkpoint to a distinct output root. Training is deferred + by this card. diff --git a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py index 3deba106c12..780f3d7f64a 100644 --- a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py +++ b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py @@ -93,7 +93,10 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.moe_grouped_gemm = True provider.moe_router_pre_softmax = True - provider.moe_token_dispatcher_type = "alltoall" + provider.moe_token_dispatcher_type = "flex" + provider.moe_flex_dispatcher_backend = "hybridep" + provider.moe_flex_dispatcher_num_sms = 16 + provider.moe_permute_fusion_into_hybridep = False provider.moe_router_load_balancing_type = "seq_aux_loss" provider.moe_shared_expert_overlap = True provider.moe_router_score_function = "sigmoid" @@ -233,7 +236,7 @@ def mapping_registry(self) -> MegatronMappingRegistry: ) hf_config = self.hf_config - num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0) + num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0) or 0 num_transformer_layers = hf_config.num_hidden_layers for mtp_layer in range(num_mtp_layers): # MTP specific mappings diff --git a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py index d7848bdc8c0..0cf273827d3 100644 --- a/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py +++ b/tests/unit_tests/models/glm_moe_dsa/test_glm5_bridge.py @@ -141,6 +141,16 @@ def test_provider_bridge_maps_dsa_architecture_from_hf_config( assert provider.dsa_indexer_use_sparse_loss is True +def test_provider_bridge_uses_hybridep_dispatcher(monkeypatch: pytest.MonkeyPatch) -> None: + """GLM-5 avoids the affected grouped all-to-all transport path.""" + provider = _provider_from_hf_config(monkeypatch) + + assert provider.moe_token_dispatcher_type == "flex" + assert provider.moe_flex_dispatcher_backend == "hybridep" + assert provider.moe_flex_dispatcher_num_sms == 16 + assert provider.moe_permute_fusion_into_hybridep is False + + def test_mapping_registry_includes_grouped_and_local_expert_fc2_paths(glm5_bridge: GLM5Bridge) -> None: """GLM-5 MoE export supports both packed and local-expert down-projection names.""" mappings = _mapping_by_megatron_param(glm5_bridge) @@ -260,9 +270,15 @@ def test_mapping_registry_includes_mtp_standalone_weights(glm5_bridge: GLM5Bridg assert mapping.hf_param == hf_param -def test_mapping_registry_omits_mtp_mappings_without_nextn_layers() -> None: +@pytest.mark.parametrize("num_nextn_predict_layers", [0, None]) +def test_mapping_registry_omits_mtp_mappings_without_nextn_layers( + num_nextn_predict_layers: int | None, +) -> None: """No MTP mappings are registered when the HF config has no MTP layers.""" bridge = GLM5Bridge() - bridge.hf_config = SimpleNamespace(num_hidden_layers=4, num_nextn_predict_layers=0) + bridge.hf_config = SimpleNamespace( + num_hidden_layers=4, + num_nextn_predict_layers=num_nextn_predict_layers, + ) assert all(not mapping.megatron_param.startswith("mtp.") for mapping in bridge.mapping_registry()) From dd05e1af30c045ab7bce25072933ef12c2918b17 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 08:28:55 +0200 Subject: [PATCH 73/85] cp: `docs(ckpt): document roundtrip scaling limits (5491)` into `r0.6.0` (#5492) Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> --- scripts/conversion/gpu_backend.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/scripts/conversion/gpu_backend.py b/scripts/conversion/gpu_backend.py index b650e95bbd1..f5c6f129101 100644 --- a/scripts/conversion/gpu_backend.py +++ b/scripts/conversion/gpu_backend.py @@ -203,7 +203,14 @@ def _roundtrip_weights_match(name: str, exported: torch.Tensor, original: torch. def _verify_roundtrip_weights(bridge: AutoBridge, megatron_model: list[torch.nn.Module]) -> None: - """Verify exported Megatron weights against the original Hugging Face state.""" + """Exhaustively verify exported Megatron weights against the original Hugging Face state. + + Every rank participates in the collective Megatron-to-Hugging-Face export, but only rank 0 lazily reads each + original Hugging Face tensor and compares it serially. This does not materialize a complete Hugging Face model + on every rank. For very large checkpoints, the rank-0 work can create prolonged rank skew before the result + broadcast, exceed the process-group collective timeout, and incur substantial transient tensor memory and + storage I/O. + """ is_rank_0 = torch.distributed.get_rank() == 0 all_match = True verified_count = 0 @@ -421,6 +428,13 @@ def roundtrip_checkpoint( ) -> None: """Validate a Hugging Face to Megatron to Hugging Face round trip. + This workflow performs exhaustive equality validation and is not intended as the scalable conversion path for + very large checkpoints. When the goal is conversion rather than exhaustive validation, prefer separate + ``convert.sh import`` and ``convert.sh export --distributed-save`` workflows. If a full round trip is required, + provision sufficient time and memory and choose an appropriate ``--distributed-timeout-minutes`` value. A + longer timeout can accommodate expected rank skew, but does not reduce serial verification cost or memory and + I/O pressure. + Args: hf_model: Hugging Face model ID or local path. tp: Tensor parallelism size. From 3ee51738b77b7b1ec7b811064a95d86b7a1a4c25 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 09:36:56 +0200 Subject: [PATCH 74/85] cp: `docs(model-card): add Nemotron 3 Super verification card (5062)` into `r0.6.0` (#5494) Signed-off-by: yaoyu-33 Signed-off-by: yaoyu-33 Signed-off-by: NeMo Bot Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: yaoyu-33 --- .../nemotron-3-super-120b-a12b/card.yaml | 693 ++++++++++++++++++ scripts/training/run_recipe.py | 5 + .../bridge/perf_recipes/nemotronh/__init__.py | 1 + .../bridge/perf_recipes/nemotronh/common.py | 14 +- .../perf_recipes/nemotronh/h100/nemotronh.py | 37 + .../bridge/recipes/nemotronh/__init__.py | 2 + .../recipes/nemotronh/gb200/__init__.py | 4 + .../nemotronh/gb200/nemotron_3_super.py | 101 +++ .../bridge/recipes/nemotronh/h100/__init__.py | 6 +- .../nemotronh/h100/nemotron_3_super.py | 232 +++++- .../recipes/nemotronh/nemotron_3_super.py | 12 +- .../test_nemotronh_recipes_pretrain.py | 3 + .../nemotronh/test_nemotron_3_super.py | 134 +++- .../recipes/test_nemotronh_recipes.py | 154 +++- .../scripts/training/test_run_recipe.py | 11 + 15 files changed, 1350 insertions(+), 59 deletions(-) create mode 100644 examples/model_verification_cards/nemotron-3-super-120b-a12b/card.yaml create mode 100644 src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_super.py diff --git a/examples/model_verification_cards/nemotron-3-super-120b-a12b/card.yaml b/examples/model_verification_cards/nemotron-3-super-120b-a12b/card.yaml new file mode 100644 index 00000000000..4c43ba3693f --- /dev/null +++ b/examples/model_verification_cards/nemotron-3-super-120b-a12b/card.yaml @@ -0,0 +1,693 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: nemotron_3_super_120b_a12b +summary: > + Performance scope: pretrain_performance.H100 and pretrain_performance.GB200 + track tuned canonical 64-GPU performance recipes. Timing and throughput from + functional H100 and GB200 training are support-verification sanity checks + rather than optimized performance results. Hardware evidence is scoped + strictly by accelerator. Nemotron 3 Super 120B-A12B verification covers + conversion, inference, bounded real-data training, checkpoint resume, and + canonical H100 and GB200 benchmarks. +verification_index: + model_level: + verified: + - hf_to_megatron_gpu + - megatron_to_hf_gpu + - inference + unverified: + - hf_to_megatron_cpu + - megatron_to_hf_cpu + - manual_forward_pass + training: + H100: + verified: [pretrain, sft_long_context] + unverified: [sft, sft_export_inference, peft, checkpoint_resume] + GB200: + verified: [peft] + unverified: [pretrain, sft, sft_export_inference, sft_long_context, checkpoint_resume] + performance: + H100: verified + GB200: verified +model: + hf_id: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + hf_revision: d51eab0d1f979ebc26b546e634a04f450d99158e # pragma: allowlist secret + architecture: NemotronHForCausalLM + min_transformers_version: "5.8.1" +verification_environment: + base_container: nvcr.io/nvidia/nemo:26.06 + bridge_commit: ae2c24890ecd1cf1d535b9e4aa473af1bc94ab71 # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: unverified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + --hf-revision d51eab0d1f979ebc26b546e634a04f450d99158e + --megatron-path work/model-verification/nemotron-3-super-120b-a12b/cpu-megatron + --torch-dtype bfloat16 + last_verified: null + expected_result: > + The pinned CPU import must exit successfully, create iter_0000000, and + preserve the complete source checkpoint so the paired CPU export can audit + every supported tensor for exact keys, shapes, dtypes, and values. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + bridge_commit: 75f87e5894d38d01683dfcef8ce6092bf002783e # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + --hf-revision d51eab0d1f979ebc26b546e634a04f450d99158e + --megatron-path work/model-verification/nemotron-3-super-120b-a12b/imported-megatron + --torch-dtype bfloat16 --tp 1 --pp 1 --ep 8 --etp 1 + last_verified: 2026-07-23 + expected_result: > + The pinned eight-GPU BF16 import exits successfully and creates a + reloadable TP1/PP1/EP8/ETP1 iter_0000000 checkpoint. Its paired export + exactly reproduces all 42,683 source tensors and 123,611,033,088 + elements, including the 41 FP32 router correction-bias tensors. + + megatron_to_hf_cpu: + status: unverified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + --hf-revision d51eab0d1f979ebc26b546e634a04f450d99158e + --megatron-path work/model-verification/nemotron-3-super-120b-a12b/cpu-megatron/iter_0000000 + --hf-path work/model-verification/nemotron-3-super-120b-a12b/cpu-hf-export + --torch-dtype bfloat16 + last_verified: null + expected_result: > + Strict CPU export must exit successfully, reproduce every supported source + tensor with exact keys, shapes, dtypes, and values, and reload natively as + NemotronHForCausalLM. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + bridge_commit: 75f87e5894d38d01683dfcef8ce6092bf002783e # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + --hf-revision d51eab0d1f979ebc26b546e634a04f450d99158e + --megatron-path work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --hf-path work/model-verification/nemotron-3-super-120b-a12b/hf-export + --torch-dtype bfloat16 --distributed-save --tp 1 --pp 1 --ep 8 --etp 1 + last_verified: 2026-07-23 + expected_result: > + Strict distributed export exits successfully and writes 247,222,108,160 + bytes across all 42,683 tensors. Keys, shapes, dtypes, values, and + critical configuration match the pinned source with zero mismatches, + including 41 FP32 router correction-bias tensors, and Transformers + reloads the result natively as NemotronHForCausalLM with 88 layers. + + manual_forward_pass: + status: unverified + precision: bf16 + command: null + last_verified: null + expected_result: > + Verification requires a pinned-revision BF16 Hugging Face and Megatron + one-step comparison with matching next-token predictions and cosine + similarity of at least 0.99. Maximum and mean absolute logit differences + must be reported as diagnostics; the helper must first support sharding + the approximately 231-GiB Hugging Face checkpoint across H100 ranks. + + inference: + status: verified + precision: bf16 + bridge_commit: 75f87e5894d38d01683dfcef8ce6092bf002783e # pragma: allowlist secret + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/hf_to_megatron_generate_text.py + --hf_model_path nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + --hf-revision d51eab0d1f979ebc26b546e634a04f450d99158e + --megatron_model_path work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --tp 1 --pp 1 --ep 8 --etp 1 + --prompt "The capital of France is" --max_new_tokens 32 + last_verified: 2026-07-23 + expected_result: | + Two independent runs in BF16 reload the imported TP1/PP1/EP8/ETP1 + checkpoint, exit successfully after exactly 32 new tokens, and produce + byte-identical output with this literal completion: + ' Paris.", + "The Eiffel Tower is in Paris.", + "Berlin is the capital of Germany.", + "The Brandenburg Gate is in Berlin.", + ' + The token IDs are 6993, 32876, 1369, 1429, 1784, 1436, 4940, 1299, + 35549, 1395, 1294, 6993, 32876, 1369, 1429, 75613, 1395, 1278, 8961, + 1307, 10787, 32876, 1369, 1429, 1784, 59071, 42969, 1395, 1294, 10874, + 32876, and 1369. + + pretrain: + H100: + status: verified + precision: bf16 + bridge_commit: ffeb5a7ae817a5a6ccfa3850d7dee937faad7f0b # pragma: allowlist secret + enabled_features: + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --wait --nodes 8 --gpus-per-node 8 + --recipe nemotron_3_super_pretrain_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3-super-120b-a12b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2-tokenizer/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 + validation.eval_iters=0 validation.eval_interval=0 + checkpoint.ckpt_format=torch_dist checkpoint.async_save=false + checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/h100-pretrain-reference + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.log_device_memory_used=true + logger.tensorboard_dir=null + logger.save_config_filepath=work/model-verification/nemotron-3-super-120b-a12b/h100-pretrain-config/ConfigContainer.yaml + last_verified: 2026-08-03 + metrics: + initial_loss: 12.19922 + final_loss: 5.706429 + last_10_steps_step_time_ms_avg: 29368.690 + last_10_steps_model_tflops_per_gpu_avg: 236.330 + expected_result: > + On exactly 64 H100s, the BF16 support run completes 100 bounded RP2 + optimizer steps at TP1/PP2/CP1/EP32/ETP1, GBS/MBS 1280/1, natural + routing, HybridEP flex dispatch with 1.10x fixed expert capacity, eager + execution, selective recompute of layernorm, moe_act, moe, and + core_attn, no optimizer-state CPU offload, and ordered gradient + reduction. Loss remains finite from 12.19922 to 5.706429 with zero + skipped or NaN iterations. The final ten steps average 29368.690 ms and + 236.330 TFLOP/s/GPU. Peak device memory is 58.079 GiB allocated and + 70.336 GiB reserved. The resolved configuration persists, and complete + iter_0000050 and iter_0000100 checkpoints each contain 64 distributed + rank shards plus checkpoint metadata and training state. + + GB200: + status: unverified + precision: bf16 + enabled_features: {} + command: > + ./scripts/training/train.sh --wait --nodes 8 --gpus-per-node 2 + --recipe nemotron_3_super_pretrain_16gpu_h100_bf16_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + --seq_length 4096 model.cuda_graph_impl=none + --lr 4.5e-4 --min_lr 4.5e-6 --warmup_iters 10 + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3-super-120b-a12b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2-tokenizer/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + optimizer.optimizer_cpu_offload=true + optimizer.optimizer_offload_fraction=0.4 + optimizer.overlap_cpu_optimizer_d2h_h2d=false + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 checkpoint.ckpt_format=torch_dist + checkpoint.async_save=false checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/gb200-pretrain-reference + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + On 16 GB200s, the BF16 support run must complete exactly 100 bounded RP2 + optimizer steps at TP8/PP1/CP1/EP16/ETP1 and the recipe's cohort + GBS/MBS 1024/1. Full activation recompute, 40% optimizer CPU offload, + and eager execution keep the model and optimizer checkpoint within + memory. Loss must stay finite with no skipped or NaN iteration, and + complete iter_0000050 and iter_0000100 checkpoints must save. The + recipe moved from GBS 16 to the cohort GBS 1024, so the earlier metrics + no longer describe this command and it needs a rerun. + + sft: + H100: + status: unverified + precision: bf16 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --wait --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_super_sft_16gpu_h100_bf16_config + --mode sft --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 8192 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="2d59de1cbd51c0adf384eb906b766d1aee0e0517"' + dataset.hf_output_root=work/data/tulu3/nemotron-3-super-120b-a12b-pad8 + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:8192,pad_seq_to_mult:8}' + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 + optimizer.optimizer_cpu_offload=true + optimizer.optimizer_offload_fraction=0.75 + optimizer.overlap_cpu_optimizer_d2h_h2d=true + checkpoint.ckpt_format=torch_dist checkpoint.async_save=false + checkpoint.save_optim=false checkpoint.save_rng=false + scheduler.lr_decay_iters=100 validation.eval_iters=0 + validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/h100-sft + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The immutable-revision Tulu3 run must complete exactly 100 full-SFT steps + on 16 H100s at sequence length 8192, TP8/EP16, the recipe's cohort + GBS/MBS 32/1, and pad-8 offline packing. Loss must stay finite with no + skipped or NaN iteration, and the step-100 checkpoint must save + completely for export. The recipe moved from GBS 16 to the cohort + GBS 32, so the earlier metrics no longer describe this command and it + needs a rerun. + + GB200: + status: unverified + precision: bf16 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --wait --nodes 4 --gpus-per-node 4 + --recipe nemotron_3_super_sft_16gpu_h100_bf16_config + --mode sft --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 8192 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="2d59de1cbd51c0adf384eb906b766d1aee0e0517"' + dataset.hf_output_root=work/data/tulu3/nemotron-3-super-120b-a12b-pad8 + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:8192,pad_seq_to_mult:8}' + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 + optimizer.optimizer_cpu_offload=true + optimizer.optimizer_offload_fraction=0.75 + optimizer.overlap_cpu_optimizer_d2h_h2d=true + checkpoint.ckpt_format=torch_dist checkpoint.async_save=false + checkpoint.save_optim=false checkpoint.save_rng=false + scheduler.lr_decay_iters=100 validation.eval_iters=0 + validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/gb200-sft + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The immutable-revision Tulu3 run must complete exactly 100 full-SFT steps + on 16 GB200s at sequence length 8192, TP8/EP16, the recipe's cohort + GBS/MBS 32/1, and pad-8 offline packing. Loss must stay finite with no + skipped or NaN iteration, and the step-100 checkpoint must save + completely. The recipe moved from GBS 16 to the cohort GBS 32, so the + earlier metrics no longer describe this command and it needs a rerun. + + sft_export_inference: + H100: + status: unverified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 8 + --hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + --hf-revision d51eab0d1f979ebc26b546e634a04f450d99158e + --megatron-path work/model-verification/nemotron-3-super-120b-a12b/h100-sft/iter_0000100 + --hf-path work/model-verification/nemotron-3-super-120b-a12b/h100-sft-hf + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 8 --pp 1 --ep 16 --etp 1 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/nemotron-3-super-120b-a12b/h100-sft-hf + --prompt "Name the capital of France and explain its role in one sentence." + --max-new-tokens 45 --chat-template --disable-thinking + last_verified: null + expected_result: > + The H100 step-100 SFT checkpoint must export to indexed BF16 + safetensors, reload natively as NemotronHForCausalLM, and complete the + pinned deterministic 45-token greedy inference before promotion. + + GB200: + status: unverified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 4 --gpus-per-node 4 + --hf-model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 + --hf-revision d51eab0d1f979ebc26b546e634a04f450d99158e + --megatron-path work/model-verification/nemotron-3-super-120b-a12b/gb200-sft/iter_0000100 + --hf-path work/model-verification/nemotron-3-super-120b-a12b/gb200-sft-hf + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 8 --pp 1 --ep 16 --etp 1 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/nemotron-3-super-120b-a12b/gb200-sft-hf + --prompt "Name the capital of France and explain its role in one sentence." + --max-new-tokens 45 --chat-template --disable-thinking + last_verified: null + expected_result: > + The GB200 step-100 full-SFT checkpoint exports as 50 indexed BF16 + safetensor shards and reloads natively as NemotronHForCausalLM with 88 + layers. Greedy inference produces exactly 45 new tokens and the literal + completion "The capital of France is Paris, which serves as the + country's political, cultural, and economic center. It is home to many + important government institutions, such as the French National + Assembly and the Senate, as well as numerous museums". + + sft_long_context: + H100: + status: verified + precision: bf16 + bridge_commit: 44f3031a2e77ff455eb128efbaa19322ad6397a6 # pragma: allowlist secret + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --wait --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_super_sft_16gpu_h100_bf16_32k_config + --mode sft --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --max_steps 10 --seq_length 32768 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="2d59de1cbd51c0adf384eb906b766d1aee0e0517"' + dataset.hf_output_root=work/data/tulu3/nemotron-3-super-120b-a12b-long + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:32768,pad_seq_to_mult:4}' + model.attention_backend=auto model.cross_entropy_loss_fusion=false + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 + optimizer.optimizer_cpu_offload=true + optimizer.optimizer_offload_fraction=1.0 + optimizer.overlap_cpu_optimizer_d2h_h2d=false + scheduler.lr_decay_iters=100 validation.eval_iters=0 + validation.eval_interval=0 checkpoint.ckpt_format=torch_dist + checkpoint.async_save=false checkpoint.save_optim=false + checkpoint.save_rng=false checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/h100-long-sft + --save_interval 10 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-24 + metrics: + initial_loss: 1.168852 + final_loss: 1.008992 + last_10_steps_step_time_ms_avg: 132859.580 + last_10_steps_model_tflops_per_gpu_avg: 5.580 + expected_result: > + On 16 H100s, the dedicated TP1/PP8/CP2/EP2 recipe completes exactly + ten packed 32K SFT steps with pad-4 alignment, full recompute, and full + optimizer CPU offload. Loss is finite from 1.168852 to 1.008992 with no + skipped or NaN iterations. The ten-step metric window includes the + first-step cold-start cost by policy, the resolved configuration is + persisted, and the final model-only checkpoint saves successfully. + + GB200: + status: unverified + precision: bf16 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --wait --nodes 4 --gpus-per-node 4 + --recipe nemotron_3_super_sft_16gpu_h100_bf16_32k_config + --mode sft --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --max_steps 10 --seq_length 32768 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:32768,pad_seq_to_mult:4}' + checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/gb200-long-sft + --save_interval 10 logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-24 + metrics: + initial_loss: 7.022207 + final_loss: 6.406034 + last_10_steps_step_time_ms_avg: 21544.740 + last_10_steps_model_tflops_per_gpu_avg: 15.930 + expected_result: > + A dedicated GB200 run must complete at least ten packed 32K SFT steps + with CP2, finite loss, no skipped or NaN iterations, all four metrics, + a persisted resolved configuration, and a reloadable final checkpoint. + + peft: + H100: + status: unverified + precision: bf16 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --wait --nodes 2 --gpus-per-node 8 + --recipe nemotron_3_super_peft_16gpu_h100_bf16_config --mode lora + --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 8192 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="2d59de1cbd51c0adf384eb906b766d1aee0e0517"' + dataset.hf_output_root=work/data/tulu3/nemotron-3-super-120b-a12b-peft-pad8 + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:8192,pad_seq_to_mult:8}' + model.recompute_granularity=selective + 'model.recompute_modules=[core_attn]' + scheduler.lr_decay_iters=100 validation.eval_iters=0 + validation.eval_interval=0 checkpoint.ckpt_format=torch_dist + checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/h100-peft + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + H100 training has completed 100 finite LoRA steps and saved an adapter, + but this leaf remains unverified until the exported adapter reloads + natively over the pinned Hugging Face base and produces a nonzero, + finite logits effect. + + GB200: + status: verified + precision: bf16 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --wait --nodes 4 --gpus-per-node 4 + --recipe nemotron_3_super_peft_16gpu_h100_bf16_config --mode lora + --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-super-120b-a12b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 8192 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="2d59de1cbd51c0adf384eb906b766d1aee0e0517"' + dataset.hf_output_root=work/data/tulu3/nemotron-3-super-120b-a12b-peft-pad8 + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:8192,pad_seq_to_mult:8}' + model.recompute_granularity=selective + 'model.recompute_modules=[core_attn]' + scheduler.lr_decay_iters=100 validation.eval_iters=0 + validation.eval_interval=0 checkpoint.ckpt_format=torch_dist + checkpoint.async_save=false checkpoint.load=null + --save_dir work/model-verification/nemotron-3-super-120b-a12b/gb200-peft + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-24 + metrics: + initial_loss: 1.23861 + final_loss: 0.8878498 + last_10_steps_step_time_ms_avg: 12699.460 + last_10_steps_model_tflops_per_gpu_avg: 54.120 + expected_result: > + On 16 GB200s, the immutable-revision Tulu3 run completes exactly 100 + pad-8 packed LoRA steps with finite loss from 1.23861 to 0.8878498, no + skipped or NaN iterations, all four metrics, and a final adapter + checkpoint. The adapter contains only 82,304 LoRA tensors targeting + linear_qkv, linear_proj, linear_fc1, linear_fc2, in_proj, and out_proj. + Native PEFT reload over the pinned base succeeds; all logits are finite + and the adapter changes them by maximum absolute 3.1875 and mean + absolute 0.472428. + + checkpoint_resume: + H100: + status: unverified + precision: bf16 + depends_on: pretrain + command: > + ./scripts/training/train.sh --wait --nodes 8 --gpus-per-node 8 + --recipe nemotron_3_super_pretrain_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3-super-120b-a12b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2-tokenizer/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 + validation.eval_iters=0 validation.eval_interval=0 + checkpoint.ckpt_format=torch_dist checkpoint.async_save=false + --load_dir work/model-verification/nemotron-3-super-120b-a12b/h100-pretrain-reference + checkpoint.ckpt_step=50 + --save_dir work/model-verification/nemotron-3-super-120b-a12b/h100-pretrain-resumed + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.log_device_memory_used=true + logger.tensorboard_dir=null + logger.save_config_filepath=work/model-verification/nemotron-3-super-120b-a12b/h100-pretrain-config/ConfigContainer.yaml + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + A direct H100 continuation must restore model, optimizer, scheduler, + data-order, and RNG state from iter_0000050, execute exactly steps + 51-100 into a distinct output root, match both declared loss sentinels, + record all four metrics, and save a complete iter_0000100 checkpoint. + + GB200: + status: unverified + precision: bf16 + depends_on: pretrain + command: > + ./scripts/training/train.sh --wait --nodes 8 --gpus-per-node 2 + --recipe nemotron_3_super_pretrain_16gpu_h100_bf16_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + --seq_length 4096 model.cuda_graph_impl=none + --lr 4.5e-4 --min_lr 4.5e-6 --warmup_iters 10 + 'dataset.blend=[["work/data/rp2/head_01_text_document"],null]' + dataset.path_to_cache=work/cache/nemotron-3-super-120b-a12b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2-tokenizer/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + optimizer.optimizer_cpu_offload=true + optimizer.optimizer_offload_fraction=0.4 + optimizer.overlap_cpu_optimizer_d2h_h2d=false + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 checkpoint.ckpt_format=torch_dist + checkpoint.async_save=false + --load_dir work/model-verification/nemotron-3-super-120b-a12b/gb200-pretrain-reference + checkpoint.ckpt_step=50 + --save_dir work/model-verification/nemotron-3-super-120b-a12b/gb200-pretrain-resumed + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: true + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The direct GB200 continuation must restore full model, optimizer, + scheduler, data-order, and RNG state from iter_0000050 and execute + exactly steps 51-100 into a distinct output root. Each declared loss + sentinel must satisfy abs(resumed - reference) <= 1e-6 + 0.01 * + abs(reference) against the uninterrupted reference. Its pretrain + reference moved to the cohort GBS 1024, so both runs need redoing + before this item can be verified. + + pretrain_performance: + H100: + status: verified + precision: bf16 + bridge_commit: ffeb5a7ae817a5a6ccfa3850d7dee937faad7f0b # pragma: allowlist secret + command: > + ./scripts/training/train.sh --wait --nodes 8 --gpus-per-node 8 + --recipe nemotron_3_super_pretrain_64gpu_h100_bf16_config + --max_steps 50 logger.log_throughput=true + logger.log_device_memory_used=true + logger.save_config_filepath=work/model-verification/nemotron-3-super-120b-a12b/h100-performance/ConfigContainer.yaml + last_verified: 2026-08-03 + metrics: + initial_loss: 12.19134 + final_loss: 0.006422824 + last_10_steps_step_time_ms_avg: 27588.580 + last_10_steps_model_tflops_per_gpu_avg: 248.610 + expected_result: > + On exactly 64 H100s, the canonical BF16 performance recipe must + complete 50 mock-data steps at TP1/PP2/CP1/EP32/ETP1, GBS/MBS 1280/1, + forced routing, HybridEP flex dispatch with 1.10x fixed expert + capacity, eager execution, selective recompute of layernorm, moe_act, + moe, and core_attn, no optimizer-state CPU offload, and ordered + gradient reduction. All keyed rows have finite loss and throughput with + zero skipped or NaN + iterations. The final ten steps average 27588.580 ms and 248.610 + TFLOP/s/GPU, passing gates of at most 28000 ms and at least 240 + TFLOP/s/GPU. Peak device memory is 59.000 GiB allocated and 69.344 GiB + reserved, and the resolved configuration persists. Mock data and forced + routing make the loss a finiteness check rather than convergence + evidence. + + GB200: + status: verified + precision: bf16 + command: > + ./scripts/training/train.sh --wait --nodes 16 --gpus-per-node 4 + --recipe nemotron_3_super_pretrain_64gpu_gb200_bf16_config + --max_steps 50 logger.log_throughput=true + last_verified: 2026-07-24 + metrics: + initial_loss: 12.17778 + final_loss: 0.01120061 + last_10_steps_step_time_ms_avg: 9926.350 + last_10_steps_model_tflops_per_gpu_avg: 559.390 + expected_result: > + On exactly 64 GB200s, the canonical BF16 performance recipe completes + 50 mock-data steps at TP2/PP1/CP1/EP64, GBS/MBS 512/1, HybridEP flex + dispatch, and Transformer Engine CUDA graph scopes attn, mamba, + moe_router, and moe_preprocess. All 50 keyed rows have finite loss and + throughput with zero skipped or NaN iterations. The final ten steps + average 9926.350 ms and 559.390 TFLOP/s/GPU, passing gates of at most + 11000 ms and at least 500 TFLOP/s/GPU, and the resolved configuration + is persisted. Mock data and forced routing make the loss a finiteness + check rather than convergence evidence. These numbers are GB200-only + evidence and make no GB300 performance claim. diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index bddb4e853e7..0286ac863f3 100755 --- a/scripts/training/run_recipe.py +++ b/scripts/training/run_recipe.py @@ -81,6 +81,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parents[1] if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) COMMON_SCRIPT_DIR = SCRIPT_DIR.parent / "common" @@ -452,6 +453,10 @@ def parse_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[ def main(argv: list[str] | None = None) -> None: """Load, configure, and execute one library or benchmark recipe.""" + # NeMo-Run executes container tasks from its private /nemo_run/code + # directory. Resolve user-facing relative dataset, checkpoint, cache, and + # logger paths from the mounted Bridge checkout instead. + os.chdir(REPO_ROOT) logging.basicConfig(level=logging.INFO) args, cli_overrides = parse_args(argv) diff --git a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py index 423b5eb03ab..2606829e677 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py @@ -49,6 +49,7 @@ nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config, nemotron_3_nano_pretrain_16gpu_h100_bf16_config, nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config, + nemotron_3_super_pretrain_64gpu_h100_bf16_config, nemotronh_56b_pretrain_64gpu_h100_fp8cs_config, ) from megatron.bridge.perf_recipes.nemotronh.vr200.nemotronh import ( diff --git a/src/megatron/bridge/perf_recipes/nemotronh/common.py b/src/megatron/bridge/perf_recipes/nemotronh/common.py index 8baf41f4ced..c661af50e99 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/common.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/common.py @@ -20,8 +20,10 @@ from megatron.core.quantization.utils import load_quantization_recipe from megatron.bridge.perf_recipes._common import _benchmark_common, _perf_precision +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import ( + nemotron_3_super_pretrain_16gpu_h100_bf16_config as nemotron_3_super_pretrain_config, +) from megatron.bridge.recipes.nemotronh.nemotron_3_nano import nemotron_3_nano_pretrain_config -from megatron.bridge.recipes.nemotronh.nemotron_3_super import nemotron_3_super_pretrain_config from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import nemotron_3_ultra_pretrain_config from megatron.bridge.recipes.nemotronh.nemotronh import nemotronh_56b_pretrain_config from megatron.bridge.training.config import ConfigContainer @@ -73,6 +75,16 @@ def _apply_nemotron_3_super_perf_defaults(cfg: ConfigContainer) -> None: cfg.mixed_precision.grad_reduce_in_fp32 = False cfg.ddp.grad_reduce_in_fp32 = False + # The Nemotron 3 Super base recipe is memory-bounded for 16-GPU support + # runs. Benchmarks measure throughput on larger systems, so restore + # overlapped collectives and full-precision optimizer state. + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.optimizer.use_precision_aware_optimizer = False + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.model.moe_router_force_load_balancing = True cfg.checkpoint.async_save = False diff --git a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py index e63ac273b1c..2a6d5f1be7e 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/h100/nemotronh.py @@ -24,6 +24,7 @@ nemotron_3_nano_pretrain_config, nemotronh_56b_pretrain_config, ) +from megatron.bridge.recipes.nemotronh.nemotron_3_super import nemotron_3_super_pretrain_config from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules @@ -67,6 +68,42 @@ def nemotronh_56b_pretrain_64gpu_h100_fp8cs_config() -> ConfigContainer: return cfg +def nemotron_3_super_pretrain_64gpu_h100_bf16_config() -> ConfigContainer: + """Nemotron 3 Super pretrain: 64× H100, BF16.""" + cfg = nemotron_3_super_pretrain_config() + cfg.mixed_precision = _perf_precision("bf16") + + # The library recipe owns the measured execution layout. Performance mode + # adds only benchmark policy and forced routing on top of that shared base. + cfg.model.moe_router_force_load_balancing = True + _benchmark_common(cfg) + # `_benchmark_common` enables parameter prefetch globally. Restore the + # shared natural-routing-safe execution mapping owned by the library + # recipe; otherwise performance and convergence silently diverge here. + cfg.ddp.overlap_param_gather = False + cfg.model.moe_hybridep_num_sms = None + cfg.checkpoint.async_save = False + + # Keep the process settings aligned with the convergence recipe. In + # particular, do not silently change CUDA stream scheduling between the + # natural-routing and forced-routing measurements. + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NCCL_GRAPH_REGISTER": 0, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + "NCCL_NVLS_ENABLE": 0, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 64, + "NVLINK_DOMAIN_SIZE": 8, + "USE_MNNVL": 0, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + } + return cfg + + def nemotron_3_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: """Nemotron 3 Nano pretrain: 16× H100, BF16, recompute MoE+layernorm.""" cfg = nemotron_3_nano_pretrain_config() diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 0c562920291..3a1ee2f27f2 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -20,6 +20,7 @@ nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_super_pretrain_64gpu_gb200_bf16_config, ) from megatron.bridge.recipes.nemotronh.nemotron_3_nano import ( nemotron_3_5_lightning_peft_config, @@ -114,6 +115,7 @@ "nemotron_3_nano_4b_sft_32k_config", "nemotron_3_nano_4b_peft_config", # Nemotron 3 Super models + "nemotron_3_super_pretrain_64gpu_gb200_bf16_config", "nemotron_3_super_pretrain_config", "nemotron_3_super_sft_config", "nemotron_3_super_peft_config", diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py index fe7149c827e..273804ab605 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py @@ -19,6 +19,9 @@ nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, ) +from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_super import ( + nemotron_3_super_pretrain_64gpu_gb200_bf16_config, +) __all__ = [ @@ -27,4 +30,5 @@ "nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_super_pretrain_64gpu_gb200_bf16_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_super.py b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_super.py new file mode 100644 index 00000000000..7fdc37f38ec --- /dev/null +++ b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_super.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GB200 pretraining recipe for Nemotron 3 Super.""" + +import torch + +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import ( + nemotron_3_super_pretrain_16gpu_h100_bf16_config, +) +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed +from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules + + +def nemotron_3_super_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: + """Return the Nemotron 3 Super BF16 pretraining config for 64 GB200 GPUs. + + This is the convergence-oriented counterpart of the canonical 64-GPU + performance recipe. It uses the same GB200 parallel layout, HybridEP + dispatcher, CUDA graph scopes, and process environment while retaining + natural expert routing, runtime validation, and checkpointing. + + Returns: + GB200 BF16 pretraining configuration. + """ + cfg = nemotron_3_super_pretrain_16gpu_h100_bf16_config() + + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = False + + # The H100 base is memory-bounded for 16-GPU support runs. GB200 has the + # capacity for overlapped collectives and full-precision optimizer state. + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.optimizer.use_precision_aware_optimizer = False + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.context_parallel_size = 1 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.sequence_parallel = True + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.expert_model_parallel_size = 64 + cfg.train.global_batch_size = 512 + cfg.train.micro_batch_size = 1 + + # The H100 base carries the bounded 100-step verification schedule. This + # GB200 recipe is the convergence counterpart of the 64-GPU performance + # recipe and owns its own full-length schedule. + cfg.train.train_iters = 39735 + cfg.scheduler.lr_warmup_iters = 333 + cfg.scheduler.lr_decay_iters = None + + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_shared_expert_overlap = False + # Forced routing is benchmark-only and would change the training objective. + cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.moe_hybridep_num_sms = 32 + + cfg.model.cuda_graph_impl = "transformer_engine" + set_cuda_graph_modules(cfg.model, ["attn", "mamba", "moe_router", "moe_preprocess"]) + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.apply_rope_fusion = True + cfg.model.recompute_granularity = None + + cfg.optimizer.optimizer_cpu_offload = False + cfg.optimizer.optimizer_offload_fraction = 0.0 + + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + } + return cfg + + +__all__ = ["nemotron_3_super_pretrain_64gpu_gb200_bf16_config"] diff --git a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py index 8184f275429..faf98d5c6da 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py @@ -33,8 +33,10 @@ "nemotron_3_nano_pretrain_8gpu_h100_bf16_config", "nemotron_3_nano_sft_8gpu_h100_bf16_config", "nemotron_3_super_peft_1gpu_h100_bf16_config", - "nemotron_3_super_pretrain_8gpu_h100_bf16_config", - "nemotron_3_super_sft_8gpu_h100_bf16_config", + "nemotron_3_super_peft_16gpu_h100_bf16_config", + "nemotron_3_super_pretrain_16gpu_h100_bf16_config", + "nemotron_3_super_sft_16gpu_h100_bf16_32k_config", + "nemotron_3_super_sft_16gpu_h100_bf16_config", "nemotron_3_ultra_peft_32gpu_h100_bf16_openmathinstruct2_packed_config", "nemotron_3_ultra_pretrain_24gpu_h100_bf16_config", "nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config", diff --git a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_super.py b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_super.py index fecd4cb3e97..90f0c9a45b9 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_super.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_super.py @@ -21,17 +21,19 @@ from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common from megatron.bridge.recipes.utils.dataset_utils import default_peft_config from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed NEMOTRON_3_SUPER_HF_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16" -def nemotron_3_super_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: +def nemotron_3_super_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: """Return a pre-training config for Nemotron 3 Super (120B-A12B LatentMoE). This is a Latent MoE model with Multi-Token Prediction (MTP). Default parallelism: - - TP=4, PP=1, EP=8, SP=True + - TP=8, PP=1, EP=16, SP=True Returns: ConfigContainer: Pre-training configuration for Nemotron 3 Super. @@ -42,14 +44,14 @@ def nemotron_3_super_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: cfg.model = AutoBridge.from_hf_pretrained(NEMOTRON_3_SUPER_HF_MODEL_ID).to_megatron_provider(load_weights=False) # Parallelism Settings - cfg.model.tensor_model_parallel_size = 4 + cfg.model.tensor_model_parallel_size = 8 cfg.model.pipeline_model_parallel_size = 1 cfg.model.pipeline_dtype = torch.bfloat16 cfg.model.virtual_pipeline_model_parallel_size = None cfg.model.context_parallel_size = 1 cfg.model.sequence_parallel = True cfg.model.expert_tensor_parallel_size = 1 - cfg.model.expert_model_parallel_size = 8 + cfg.model.expert_model_parallel_size = 16 cfg.model.pipeline_model_parallel_layout = None cfg.model.seq_length = 8192 @@ -67,9 +69,10 @@ def nemotron_3_super_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: cfg.model.moe_shared_expert_overlap = False cfg.model.moe_flex_dispatcher_backend = "hybridep" - # Training Configuration - cfg.train.train_iters = 39735 - cfg.train.global_batch_size = 3072 + # Training Configuration. Bounded 100-step cohort schedule, matching the + # other H100 verification recipes (GBS 1024 / MBS 1 / 100 steps). + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 cfg.train.micro_batch_size = 1 cfg.train.manual_gc = False cfg.train.manual_gc_interval = 0 @@ -98,8 +101,11 @@ def nemotron_3_super_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: cfg.model.mtp_loss_scaling_factor = 0.3 cfg.model.mtp_use_repeated_layer = True - # Mixed Precision - cfg.mixed_precision = "nemotron_3_super_bf16_with_nvfp4_mixed" + # Mixed Precision. Gradients reduce in BF16 to bound the gradient buffer; + # grad_reduce_in_fp32 is a MixedPrecisionConfig field that setup() copies + # onto the DDP config, so it must be set here rather than on cfg.ddp alone. + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = False # Optimizer hyperparameters cfg.optimizer.lr = 4.5e-4 @@ -108,7 +114,8 @@ def nemotron_3_super_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: cfg.optimizer.adam_beta1 = 0.9 cfg.optimizer.adam_beta2 = 0.95 cfg.optimizer.adam_eps = 1e-8 - cfg.scheduler.lr_warmup_iters = 333 + cfg.scheduler.lr_warmup_iters = 40 + cfg.scheduler.lr_decay_iters = 100 cfg.scheduler.start_weight_decay = 0.1 cfg.scheduler.end_weight_decay = 0.1 cfg.scheduler.lr_decay_style = "WSD" @@ -119,13 +126,21 @@ def nemotron_3_super_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: cfg.checkpoint.dist_ckpt_strictness = "log_all" cfg.checkpoint.async_save = True - # DDP Configuration - cfg.ddp.overlap_grad_reduce = True - cfg.ddp.overlap_param_gather = True + # DDP Configuration. Overlap stays off so full-model checkpoint saves are + # ordered after all DDP collectives. + cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False cfg.ddp.check_for_nan_in_grad = True cfg.ddp.use_distributed_optimizer = True cfg.ddp.average_in_collective = False + # Reduced-memory optimizer state. + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.main_params_dtype = torch.float16 + cfg.optimizer.exp_avg_dtype = torch.bfloat16 + cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 + cfg.model.init_method_std = 0.014 cfg.model.apply_rope_fusion = False cfg.model.gradient_accumulation_fusion = True @@ -142,15 +157,105 @@ def nemotron_3_super_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: return cfg +def _nemotron_3_super_pretrain_64gpu_h100_bf16_config() -> ConfigContainer: + """Return the convergence-oriented Nemotron 3 Super config for 64 H100 GPUs. + + The public hardware-agnostic ``nemotron_3_super_pretrain_config`` alias + exposes this configuration. Its execution layout matches the canonical + H100 performance recipe while it retains natural routing, runtime checks, + and checkpointing for real-data training evidence. + + Returns: + BF16 pretraining configuration for 64 H100 GPUs. + """ + cfg = nemotron_3_super_pretrain_16gpu_h100_bf16_config() + + # The measured H100 layout avoids TP communication while using PP to fit + # the dense layers and EP to shard the experts across all 64 GPUs. + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 2 + cfg.model.context_parallel_size = 1 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_layout = None + cfg.model.sequence_parallel = False + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.expert_model_parallel_size = 32 + cfg.model.overlap_p2p_comm = False + cfg.model.batch_p2p_comm = True + + cfg.model.seq_length = 4096 + cfg.dataset.seq_length = 4096 + # Match the canonical GPT-OSS 120B H100 benchmark shape. DP=1 for + # TP1/PP2/EP32, so this yields 1,280 microbatches per optimizer step. + cfg.train.global_batch_size = 1280 + cfg.train.micro_batch_size = 1 + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.moe_deepep_num_sms = None + cfg.model.moe_hybridep_num_sms = None + # Keep HybridEP inputs static without its per-layer metadata all-reduce. + # At this shape, 1.10x capacity gives each expert 10% headroom over the + # 5,632-token mean. The fixed expert padding keeps downstream collective + # shapes consistent under natural routing without over-sizing activations. + cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = False + cfg.model.moe_expert_capacity_factor = 1.10 + cfg.model.moe_pad_expert_input_to_capacity = True + cfg.model.moe_shared_expert_overlap = False + # Forced routing changes the training objective and remains benchmark-only. + cfg.model.moe_router_force_load_balancing = False + + cfg.model.cuda_graph_impl = "none" + cfg.model.cuda_graph_scope = [] + cfg.model.apply_rope_fusion = True + cfg.model.use_te_rng_tracker = False + cfg.rng.te_rng_tracker = False + # Keep the provider-owned vocabulary shape identical between real-data + # convergence and mock-data performance runs. + cfg.tokenizer.use_tokenizer_vocab_size = False + cfg.model.recompute_granularity = "selective" + # Fixed-capacity dispatch keeps MoE replay collectives shape-stable while + # full MoE recompute provides the measured H100 activation headroom. + cfg.model.recompute_modules = ["layernorm", "moe_act", "moe", "core_attn"] + + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.main_grads_dtype = torch.bfloat16 + cfg.optimizer.main_params_dtype = torch.float16 + cfg.optimizer.exp_avg_dtype = torch.bfloat16 + cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 + cfg.optimizer.optimizer_cpu_offload = False + cfg.optimizer.optimizer_offload_fraction = 0.0 + cfg.optimizer.overlap_cpu_optimizer_d2h_h2d = False + + # Keep DDP collectives ordered with PP and expert communication. + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + # Retain enough watchdog headroom for first-step bring-up of the 120B model. + cfg.dist.distributed_timeout_minutes = 15 + + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 64, + "NVLINK_DOMAIN_SIZE": 8, + "USE_MNNVL": 0, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + } + return cfg + + # ============================================================================= # SFT Config # ============================================================================= -def nemotron_3_super_sft_8gpu_h100_bf16_config() -> ConfigContainer: +def nemotron_3_super_sft_16gpu_h100_bf16_config() -> ConfigContainer: """Return a full SFT config for Nemotron 3 Super (120B-A12B LatentMoE). - Default parallelism: TP=1, PP=1, EP=8, SP=True + Default parallelism: TP=8, PP=1, EP=16, SP=True Returns: ConfigContainer with all settings pre-configured for Nemotron 3 Super SFT. @@ -161,14 +266,14 @@ def nemotron_3_super_sft_8gpu_h100_bf16_config() -> ConfigContainer: cfg.model = AutoBridge.from_hf_pretrained(NEMOTRON_3_SUPER_HF_MODEL_ID).to_megatron_provider(load_weights=False) # Parallelism settings - cfg.model.tensor_model_parallel_size = 1 + cfg.model.tensor_model_parallel_size = 8 cfg.model.pipeline_model_parallel_size = 1 cfg.model.pipeline_dtype = torch.bfloat16 cfg.model.virtual_pipeline_model_parallel_size = None cfg.model.context_parallel_size = 1 cfg.model.sequence_parallel = True cfg.model.expert_tensor_parallel_size = 1 - cfg.model.expert_model_parallel_size = 8 + cfg.model.expert_model_parallel_size = 16 cfg.model.pipeline_model_parallel_layout = None cfg.model.seq_length = 2048 @@ -208,9 +313,27 @@ def nemotron_3_super_sft_8gpu_h100_bf16_config() -> ConfigContainer: cfg.scheduler.end_weight_decay = 0.1 cfg.scheduler.lr_decay_style = "cosine" + # Reduced-memory optimizer state. + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.main_params_dtype = torch.float16 + cfg.optimizer.exp_avg_dtype = torch.bfloat16 + cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 + # Tokenizer cfg.tokenizer.tokenizer_model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" + # Training Configuration. Bounded 100-step cohort schedule, matching the + # other H100 SFT verification recipes (GBS 32 / MBS 1 / 100 steps). + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.train.micro_batch_size = 1 + + # Mixed Precision. Gradients reduce in BF16 to bound the gradient buffer; + # grad_reduce_in_fp32 is a MixedPrecisionConfig field that setup() copies + # onto the DDP config, so it must be set here rather than on cfg.ddp alone. + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = False + # Checkpoint config overrides cfg.checkpoint.save_interval = 200 cfg.checkpoint.ckpt_format = "torch_dist" @@ -224,11 +347,12 @@ def nemotron_3_super_sft_8gpu_h100_bf16_config() -> ConfigContainer: # RNG config cfg.rng.seed = 1234 - # DDP config + # DDP config. Overlap stays off so full-model checkpoint saves are ordered + # after all DDP collectives. cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.grad_reduce_in_fp32 = True - cfg.ddp.overlap_grad_reduce = True - cfg.ddp.overlap_param_gather = True + cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False cfg.ddp.use_distributed_optimizer = True # Keep the complete process environment visible on the recipe. @@ -242,6 +366,46 @@ def nemotron_3_super_sft_8gpu_h100_bf16_config() -> ConfigContainer: return cfg +def nemotron_3_super_sft_16gpu_h100_bf16_32k_config() -> ConfigContainer: + """Return the 16-H100 BF16 32K full-SFT configuration. + + Long-context SFT uses a distinct PP8/CP2 layout rather than the default + SFT topology. + """ + cfg = nemotron_3_super_sft_16gpu_h100_bf16_config() + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 8 + # The final stage also owns the output projection and loss. Keep four + # decoder layers there and distribute 12 layers to each earlier stage. + cfg.model.num_layers_in_last_pipeline_stage = 4 + cfg.model.context_parallel_size = 2 + cfg.model.expert_model_parallel_size = 2 + cfg.model.seq_length = 32768 + cfg.model.cp_comm_type = "a2a" + # CP SFT needs per-token accounting because a CP rank can receive no + # supervised labels after masking. + cfg.model.cross_entropy_loss_fusion = False + cfg.model.calculate_per_token_loss = True + cfg.dataset.seq_length = 32768 + cfg.train.global_batch_size = 2 + cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False, batch_p2p_comm=False) + cfg.ddp.average_in_collective = False + # With uneven PP stages, asynchronous DDP collectives can race the + # cross-stage tied-embedding reduction and PP send/recv operations. + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + # Keep the complete process environment visible on the recipe and PP + # sends/receives ordered for the long-context PP8 schedule. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 1, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NVLINK_DOMAIN_SIZE": 8, + "USE_MNNVL": 0, + } + return cfg + + # ============================================================================= # PEFT Config # ============================================================================= @@ -364,8 +528,30 @@ def nemotron_3_super_peft_1gpu_h100_bf16_config( return cfg +def nemotron_3_super_peft_16gpu_h100_bf16_config( + peft_scheme: str | PEFT = "lora", +) -> ConfigContainer: + """Return the 16-H100 BF16 PEFT configuration.""" + cfg = nemotron_3_super_peft_1gpu_h100_bf16_config(peft_scheme) + cfg.model.tensor_model_parallel_size = 8 + cfg.model.expert_model_parallel_size = 16 + cfg.train.global_batch_size = 16 + # Keep checkpoint saves ordered after all DDP collectives. + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NVLINK_DOMAIN_SIZE": 8, + "USE_MNNVL": 0, + } + return cfg + + __all__ = [ "nemotron_3_super_peft_1gpu_h100_bf16_config", - "nemotron_3_super_pretrain_8gpu_h100_bf16_config", - "nemotron_3_super_sft_8gpu_h100_bf16_config", + "nemotron_3_super_peft_16gpu_h100_bf16_config", + "nemotron_3_super_pretrain_16gpu_h100_bf16_config", + "nemotron_3_super_sft_16gpu_h100_bf16_32k_config", + "nemotron_3_super_sft_16gpu_h100_bf16_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py index c50f6e5693b..a3355ed2726 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_super.py @@ -18,16 +18,20 @@ from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import ( NEMOTRON_3_SUPER_HF_MODEL_ID, + _nemotron_3_super_pretrain_64gpu_h100_bf16_config, ) from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import ( nemotron_3_super_peft_1gpu_h100_bf16_config as nemotron_3_super_peft_config, ) from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import ( - nemotron_3_super_pretrain_8gpu_h100_bf16_config as nemotron_3_super_pretrain_config, -) -from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import ( - nemotron_3_super_sft_8gpu_h100_bf16_config as nemotron_3_super_sft_config, + nemotron_3_super_sft_16gpu_h100_bf16_config as nemotron_3_super_sft_config, ) +from megatron.bridge.training.config import ConfigContainer + + +def nemotron_3_super_pretrain_config() -> ConfigContainer: + """Return the convergence-oriented 64-H100 Nemotron 3 Super pretraining config.""" + return _nemotron_3_super_pretrain_64gpu_h100_bf16_config() __all__ = [ diff --git a/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py b/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py index c07e5549285..6b29ed1cc13 100644 --- a/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py +++ b/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py @@ -76,6 +76,9 @@ def test_nemotron_3_nano_pretrain_recipes( "mtp_num_layers": 2, "mtp_hybrid_override_pattern": "*E", "moe_router_topk": 2, + # The production recipe uses HybridEP across 64 GPUs. Keep this + # two-GPU functional smoke on the topology-agnostic dispatcher. + "moe_token_dispatcher_type": "alltoall", # Keep this tiny MTP smoke out of MCore's MoE metric tracker path: # decoder and MTP routers initialize different layer-count views. "moe_aux_loss_coeff": 0.0, diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_super.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_super.py index 273ecb18541..23e5f04f8a5 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_super.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_super.py @@ -27,8 +27,15 @@ import tempfile import pytest +import torch from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import ( + nemotron_3_super_peft_16gpu_h100_bf16_config, + nemotron_3_super_pretrain_16gpu_h100_bf16_config, + nemotron_3_super_sft_16gpu_h100_bf16_32k_config, + nemotron_3_super_sft_16gpu_h100_bf16_config, +) from megatron.bridge.recipes.nemotronh.nemotron_3_super import ( nemotron_3_super_peft_config, nemotron_3_super_pretrain_config, @@ -53,35 +60,47 @@ def test_pretrain_config_default_parameters(self): assert isinstance(config.model, HybridModelProvider) # Check model configuration defaults - assert config.model.tensor_model_parallel_size == 4 - assert config.model.pipeline_model_parallel_size == 1 - assert config.model.sequence_parallel is True + assert config.model.tensor_model_parallel_size == 1 + assert config.model.pipeline_model_parallel_size == 2 + assert config.model.sequence_parallel is False # Check expert parallelism defaults assert config.model.expert_tensor_parallel_size == 1 - assert config.model.expert_model_parallel_size == 8 + assert config.model.expert_model_parallel_size == 32 # Check training configuration - assert config.train.train_iters == 39735 - assert config.train.global_batch_size == 3072 + assert config.train.train_iters == 100 + assert config.train.global_batch_size == 1280 assert config.train.micro_batch_size == 1 # Check dataset configuration - assert config.dataset.seq_length == 8192 + assert config.dataset.seq_length == 4096 # Check tokenizer (HuggingFace for this recipe) assert config.tokenizer.tokenizer_type == "HuggingFaceTokenizer" assert config.tokenizer.tokenizer_model == "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" # Check precision - assert config.mixed_precision == "nemotron_3_super_bf16_with_nvfp4_mixed" + assert config.mixed_precision.bf16 is True + assert config.mixed_precision.grad_reduce_in_fp32 is False + assert config.optimizer.main_grads_dtype == torch.bfloat16 + assert config.optimizer.optimizer_cpu_offload is False + assert config.optimizer.optimizer_offload_fraction == 0.0 + assert config.optimizer.overlap_cpu_optimizer_d2h_h2d is False + assert config.model.apply_rope_fusion is True + assert config.tokenizer.use_tokenizer_vocab_size is False def test_pretrain_config_moe_settings(self): """Test MoE settings for pretrain config.""" config = nemotron_3_super_pretrain_config() # Verify MoE settings - assert config.model.moe_token_dispatcher_type == "alltoall" + assert config.model.moe_token_dispatcher_type == "flex" + assert config.model.moe_flex_dispatcher_backend == "hybridep" + assert config.model.moe_hybridep_pad_uneven_dispatch_inputs is False + assert config.model.moe_expert_capacity_factor == 1.10 + assert config.model.moe_pad_expert_input_to_capacity is True + assert config.model.moe_router_force_load_balancing is False assert config.model.moe_shared_expert_overlap is False assert config.model.moe_grouped_gemm is True assert config.model.moe_permute_fusion is True @@ -109,7 +128,7 @@ def test_pretrain_config_optimizer_settings(self): assert config.optimizer.lr == 4.5e-4 assert config.optimizer.weight_decay == 0.1 assert config.optimizer.min_lr == 4.5e-6 - assert config.scheduler.lr_warmup_iters == 333 + assert config.scheduler.lr_warmup_iters == 40 def test_pretrain_config_checkpoint_settings(self): """Test checkpoint settings for pretrain config.""" @@ -133,12 +152,12 @@ def test_sft_config_defaults(self): assert isinstance(config.model, HybridModelProvider) # Check parallelism for full SFT - assert config.model.tensor_model_parallel_size == 1 + assert config.model.tensor_model_parallel_size == 8 assert config.model.pipeline_model_parallel_size == 1 assert config.model.sequence_parallel is True - # Check expert parallelism (EP=8 for full SFT) - assert config.model.expert_model_parallel_size == 8 + # Check expert parallelism (EP=16 for full SFT) + assert config.model.expert_model_parallel_size == 16 # No PEFT config for full SFT assert config.peft is None @@ -151,7 +170,8 @@ def test_sft_config_defaults(self): assert config.tokenizer.tokenizer_model == "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" # Check precision - assert config.mixed_precision == "bf16_mixed" + assert config.mixed_precision.bf16 is True + assert config.mixed_precision.grad_reduce_in_fp32 is False def test_sft_config_custom_parallelism(self): """Test SFT config with custom parallelism applied after creation.""" @@ -297,20 +317,20 @@ def test_config_container_structure(self, recipe_fn): assert config.mixed_precision is not None @pytest.mark.parametrize( - "recipe_fn", + ("recipe_fn", "overlap_grad_reduce", "overlap_param_gather"), [ - nemotron_3_super_pretrain_config, - nemotron_3_super_sft_config, - nemotron_3_super_peft_config, + (nemotron_3_super_pretrain_config, False, False), + (nemotron_3_super_sft_config, False, False), + (nemotron_3_super_peft_config, True, True), ], ) - def test_ddp_configuration(self, recipe_fn): + def test_ddp_configuration(self, recipe_fn, overlap_grad_reduce, overlap_param_gather): """Test distributed data parallel configuration.""" config = recipe_fn() assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True + assert config.ddp.overlap_grad_reduce is overlap_grad_reduce + assert config.ddp.overlap_param_gather is overlap_param_gather assert config.ddp.use_distributed_optimizer is True @pytest.mark.parametrize( @@ -332,3 +352,75 @@ def test_moe_model_configuration(self, recipe_fn): assert config.model.moe_router_topk == 22 assert config.model.moe_router_topk_scaling_factor == 5.0 assert config.model.moe_latent_size == 1024 + + +@pytest.mark.unit +class TestNemotron3Super16GpuH100: + """Test the 16-H100 support-verification recipes.""" + + @pytest.mark.parametrize( + ("recipe_fn", "expected_global_batch_size"), + [ + # Bounded cohort batch sizes shared with the other H100 verification + # recipes: pretrain GBS 1024, full SFT GBS 32. + (nemotron_3_super_pretrain_16gpu_h100_bf16_config, 1024), + (nemotron_3_super_sft_16gpu_h100_bf16_config, 32), + (nemotron_3_super_peft_16gpu_h100_bf16_config, 16), + ], + ) + def test_parallelism_batch_size_and_checkpoint_safe_ddp(self, recipe_fn, expected_global_batch_size): + """The verification recipes own their 16-GPU layout and checkpoint-safe DDP settings.""" + config = recipe_fn() + + assert config.model.tensor_model_parallel_size == 8 + assert config.model.pipeline_model_parallel_size == 1 + assert config.model.expert_model_parallel_size == 16 + assert config.train.global_batch_size == expected_global_batch_size + assert config.train.micro_batch_size == 1 + assert config.ddp.overlap_grad_reduce is False + assert config.ddp.overlap_param_gather is False + + @pytest.mark.parametrize( + "recipe_fn", + [ + nemotron_3_super_pretrain_16gpu_h100_bf16_config, + nemotron_3_super_sft_16gpu_h100_bf16_config, + nemotron_3_super_sft_16gpu_h100_bf16_32k_config, + ], + ) + def test_low_memory_optimizer_precision(self, recipe_fn): + """Full-model recipes use the explicit reduced-memory optimizer state.""" + config = recipe_fn() + + assert config.mixed_precision.grad_reduce_in_fp32 is False + assert config.ddp.grad_reduce_in_fp32 is False + assert config.ddp.overlap_grad_reduce is False + assert config.ddp.overlap_param_gather is False + assert config.optimizer.use_precision_aware_optimizer is True + assert config.optimizer.main_params_dtype == torch.float16 + assert config.optimizer.exp_avg_dtype == torch.bfloat16 + assert config.optimizer.exp_avg_sq_dtype == torch.bfloat16 + + def test_long_context_parallelism_batch_size_and_sequence_length(self): + """The 32K recipe owns its memory-bounded layout and batch size.""" + config = nemotron_3_super_sft_16gpu_h100_bf16_32k_config() + + assert config.model.tensor_model_parallel_size == 1 + assert config.model.pipeline_model_parallel_size == 8 + assert config.model.num_layers_in_last_pipeline_stage == 4 + assert config.model.context_parallel_size == 2 + assert config.model.expert_model_parallel_size == 2 + assert config.model.sequence_parallel is True + assert config.model.cp_comm_type == "a2a" + assert config.model.seq_length == 32768 + assert config.model.cross_entropy_loss_fusion is False + assert config.model.calculate_per_token_loss is True + assert config.dataset.seq_length == 32768 + assert config.train.global_batch_size == 2 + assert config.train.micro_batch_size == 1 + assert config.comm_overlap.tp_comm_overlap is False + assert config.comm_overlap.batch_p2p_comm is False + assert config.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] == 1 + assert config.ddp.average_in_collective is False + assert config.ddp.overlap_grad_reduce is False + assert config.ddp.overlap_param_gather is False diff --git a/tests/unit_tests/recipes/test_nemotronh_recipes.py b/tests/unit_tests/recipes/test_nemotronh_recipes.py index 4353b548f75..c3d564c8258 100644 --- a/tests/unit_tests/recipes/test_nemotronh_recipes.py +++ b/tests/unit_tests/recipes/test_nemotronh_recipes.py @@ -75,7 +75,9 @@ def _patch_hf_backed_recipe_providers(monkeypatch: pytest.MonkeyPatch) -> None: """Keep AutoBridge-backed recipe construction deterministic and offline.""" for module_name in ( "megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano", + "megatron.bridge.recipes.nemotronh.gb200.nemotron_3_super", "megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano", + "megatron.bridge.recipes.nemotronh.h100.nemotron_3_super", "megatron.bridge.recipes.nemotronh.nemotron_3_super", "megatron.bridge.recipes.nemotronh.nemotron_3_ultra", ): @@ -160,6 +162,142 @@ def test_nemotron_3_nano_gb200_defers_vocab_size_to_training_tokenizer(): assert cfg.model.vocab_size is None +def test_nemotron_3_super_64gpu_gb200_matches_benchmark_hardware_configuration(): + """The training recipe should share the tuned GB200 layout without benchmark-only behavior.""" + from megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh import ( + nemotron_3_super_pretrain_64gpu_gb200_bf16_config, + ) + + training_cfg = _nemotronh_module.nemotron_3_super_pretrain_64gpu_gb200_bf16_config() + benchmark_cfg = nemotron_3_super_pretrain_64gpu_gb200_bf16_config() + + for field_name in ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "context_parallel_size", + "expert_tensor_parallel_size", + "expert_model_parallel_size", + "moe_flex_dispatcher_backend", + "moe_hybridep_num_sms", + "moe_token_dispatcher_type", + "moe_shared_expert_overlap", + "cuda_graph_impl", + "apply_rope_fusion", + ): + assert getattr(training_cfg.model, field_name) == getattr(benchmark_cfg.model, field_name) + + training_scopes = [ + scope.name if hasattr(scope, "name") else scope for scope in training_cfg.model.cuda_graph_scope + ] + benchmark_scopes = [ + scope.name if hasattr(scope, "name") else scope for scope in benchmark_cfg.model.cuda_graph_scope + ] + assert training_scopes == benchmark_scopes + + assert training_cfg.train.global_batch_size == benchmark_cfg.train.global_batch_size == 512 + assert training_cfg.train.micro_batch_size == benchmark_cfg.train.micro_batch_size == 1 + assert training_cfg.env_vars == benchmark_cfg.env_vars + + assert training_cfg.model.moe_router_force_load_balancing is False + assert benchmark_cfg.model.moe_router_force_load_balancing is True + assert training_cfg.train.train_iters == 39735 + assert benchmark_cfg.train.train_iters == 50 + assert training_cfg.checkpoint.async_save is True + assert benchmark_cfg.checkpoint.async_save is False + assert training_cfg.ddp.check_for_nan_in_grad is True + assert benchmark_cfg.ddp.check_for_nan_in_grad is False + assert training_cfg.ddp.overlap_grad_reduce is True + assert training_cfg.ddp.overlap_param_gather is True + assert training_cfg.model.recompute_granularity is None + assert training_cfg.optimizer.optimizer_cpu_offload is False + assert training_cfg.optimizer.optimizer_offload_fraction == 0.0 + + # The GB200 recipe derives from the memory-bounded H100 support config, so it + # must restore overlapped collectives and full-precision optimizer state. + assert training_cfg.optimizer.use_precision_aware_optimizer is False + assert training_cfg.optimizer.main_params_dtype == torch.float32 + assert training_cfg.optimizer.exp_avg_dtype == torch.float32 + assert training_cfg.optimizer.exp_avg_sq_dtype == torch.float32 + + +def test_nemotron_3_super_64gpu_h100_matches_benchmark_execution_configuration(): + """The H100 convergence and benchmark recipes should share their execution mapping.""" + from megatron.bridge.perf_recipes.nemotronh.h100.nemotronh import ( + nemotron_3_super_pretrain_64gpu_h100_bf16_config, + ) + + training_cfg = _nemotronh_module.nemotron_3_super_pretrain_config() + benchmark_cfg = nemotron_3_super_pretrain_64gpu_h100_bf16_config() + + for field_name in ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "context_parallel_size", + "virtual_pipeline_model_parallel_size", + "expert_tensor_parallel_size", + "expert_model_parallel_size", + "sequence_parallel", + "overlap_p2p_comm", + "batch_p2p_comm", + "moe_token_dispatcher_type", + "moe_flex_dispatcher_backend", + "moe_flex_dispatcher_num_sms", + "moe_hybridep_num_sms", + "moe_hybridep_pad_uneven_dispatch_inputs", + "moe_shared_expert_overlap", + "cuda_graph_impl", + "apply_rope_fusion", + "recompute_granularity", + "recompute_modules", + ): + assert getattr(training_cfg.model, field_name) == getattr(benchmark_cfg.model, field_name) + + assert training_cfg.model.seq_length == benchmark_cfg.model.seq_length == 4096 + assert training_cfg.dataset.seq_length == benchmark_cfg.dataset.seq_length == 4096 + assert training_cfg.train.global_batch_size == benchmark_cfg.train.global_batch_size == 1280 + assert training_cfg.train.micro_batch_size == benchmark_cfg.train.micro_batch_size == 1 + assert training_cfg.ddp.overlap_grad_reduce == benchmark_cfg.ddp.overlap_grad_reduce is False + assert training_cfg.ddp.overlap_param_gather == benchmark_cfg.ddp.overlap_param_gather is False + assert training_cfg.dist.distributed_timeout_minutes == benchmark_cfg.dist.distributed_timeout_minutes == 15 + assert training_cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is False + assert training_cfg.model.moe_expert_capacity_factor == benchmark_cfg.model.moe_expert_capacity_factor == 1.10 + assert ( + training_cfg.model.moe_pad_expert_input_to_capacity + is benchmark_cfg.model.moe_pad_expert_input_to_capacity + is True + ) + assert training_cfg.model.recompute_modules == ["layernorm", "moe_act", "moe", "core_attn"] + assert training_cfg.model.apply_rope_fusion is benchmark_cfg.model.apply_rope_fusion is True + assert training_cfg.tokenizer.use_tokenizer_vocab_size is benchmark_cfg.tokenizer.use_tokenizer_vocab_size is False + assert training_cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 64 + assert training_cfg.optimizer.use_precision_aware_optimizer is True + assert benchmark_cfg.optimizer.use_precision_aware_optimizer is True + assert training_cfg.optimizer.main_grads_dtype == benchmark_cfg.optimizer.main_grads_dtype == torch.bfloat16 + assert training_cfg.optimizer.main_params_dtype == benchmark_cfg.optimizer.main_params_dtype == torch.float16 + assert training_cfg.optimizer.exp_avg_dtype == benchmark_cfg.optimizer.exp_avg_dtype == torch.bfloat16 + assert training_cfg.optimizer.exp_avg_sq_dtype == benchmark_cfg.optimizer.exp_avg_sq_dtype == torch.bfloat16 + assert training_cfg.optimizer.optimizer_cpu_offload is benchmark_cfg.optimizer.optimizer_cpu_offload is False + assert ( + training_cfg.optimizer.optimizer_offload_fraction == benchmark_cfg.optimizer.optimizer_offload_fraction == 0.0 + ) + assert ( + training_cfg.optimizer.overlap_cpu_optimizer_d2h_h2d + is benchmark_cfg.optimizer.overlap_cpu_optimizer_d2h_h2d + is False + ) + assert training_cfg.env_vars == benchmark_cfg.env_vars + assert training_cfg.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] == 32 + + assert training_cfg.model.moe_router_force_load_balancing is False + assert benchmark_cfg.model.moe_router_force_load_balancing is True + assert training_cfg.train.train_iters == 100 + assert benchmark_cfg.train.train_iters == 50 + assert training_cfg.ddp.check_for_nan_in_grad is True + assert benchmark_cfg.ddp.check_for_nan_in_grad is False + assert training_cfg.checkpoint.save is not None + assert benchmark_cfg.checkpoint.save is None + + def test_nemotron_3_5_lightning_h100_convergence_recipe_uses_perf_execution_policy(): """The H100 convergence recipe keeps safety checks while using the measured fast path.""" from megatron.bridge.recipes.nemotronh import nemotron_3_5_lightning_pretrain_config @@ -480,11 +618,11 @@ def test_nemotron_3_super_pretrain_defaults(): _assert_basic_config(cfg) - # Pretrain should use TP=4, PP=1 - assert cfg.model.tensor_model_parallel_size == 4 - assert cfg.model.pipeline_model_parallel_size == 1 - assert cfg.model.sequence_parallel is True - assert cfg.model.expert_model_parallel_size == 8 + # Pretrain uses the measured 64-H100 TP-free layout. + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 2 + assert cfg.model.sequence_parallel is False + assert cfg.model.expert_model_parallel_size == 32 def test_nemotron_3_super_peft_lora_defaults(): @@ -513,9 +651,9 @@ def test_nemotron_3_super_sft_defaults(): _assert_basic_config(cfg) - # For full SFT, should use TP=1, PP=1, EP=8 - assert cfg.model.tensor_model_parallel_size == 1 + # Full SFT uses the checkpoint-safe TP8/EP16 layout. + assert cfg.model.tensor_model_parallel_size == 8 assert cfg.model.pipeline_model_parallel_size == 1 assert cfg.model.sequence_parallel is True - assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.expert_model_parallel_size == 16 assert cfg.peft is None diff --git a/tests/unit_tests/scripts/training/test_run_recipe.py b/tests/unit_tests/scripts/training/test_run_recipe.py index 919b6bb036f..98d1596cb38 100644 --- a/tests/unit_tests/scripts/training/test_run_recipe.py +++ b/tests/unit_tests/scripts/training/test_run_recipe.py @@ -338,6 +338,17 @@ def test_full_recipe_auto_detects_benchmark_recipe(monkeypatch): ) +def test_main_resolves_relative_paths_from_repository_root(monkeypatch): + module, handles = _load_module() + handles.recipe_runner.load_recipe.return_value = SimpleNamespace() + change_directory = Mock() + monkeypatch.setattr(module.os, "chdir", change_directory) + + module.main(["--recipe", "gpt_oss_20b_pretrain_config", "--mode", "pretrain"]) + + change_directory.assert_called_once_with(module.REPO_ROOT) + + @pytest.mark.parametrize( "override", [ From 5796c06a1cf3a4cd9bfa848a419ae5e33fe32f7d Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Tue, 11 Aug 2026 20:00:10 +0200 Subject: [PATCH 75/85] cp: `docs: update performance summary for 26.08 (5496)` into `r0.6.0` (#5507) Signed-off-by: Malay Nagda Signed-off-by: NeMo Bot Co-authored-by: malay-nagda --- docs/performance-summary-archive.md | 50 +++++++++++++++++++++++++++++ docs/performance-summary.md | 48 ++++++++++++++++++--------- 2 files changed, 83 insertions(+), 15 deletions(-) diff --git a/docs/performance-summary-archive.md b/docs/performance-summary-archive.md index c5fd70620dd..0475abb51cc 100644 --- a/docs/performance-summary-archive.md +++ b/docs/performance-summary-archive.md @@ -39,6 +39,56 @@ The performance data includes: --- +## 26.06.01 NeMo Container + +### Pre-Training Performance + +#### Model: DeepSeekV3 + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 256 | MXFP8 | 4096 | 1 | 4096 | 1 | 2 | 1 | 8 | 32 | 6304 | 1640 | +| DGX-GB200 | 256 | MXFP8 | 4096 | 1 | 4096 | 1 | 4 | 1 | 4 | 64 | 4928 | 1280 | + +#### Model: GPT OSS 120B + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 64 | MXFP8 | 1280 | 4 | 4096 | 1 | 1 | 1 | n/a | 16 | 20288 | 661[^gpt-oss-note] | +| DGX-GB200 | 64 | MXFP8 | 1280 | 4 | 4096 | 1 | 1 | 1 | n/a | 64 | 18304 | 597[^gpt-oss-note] | + +#### Model: Qwen3_30B_a3B + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 8 | MXFP8 | 512 | 8 | 4096 | 1 | 1 | 1 | n/a | 8 | 44544 | 1029 | +| DGX-GB200 | 8 | MXFP8 | 512 | 4 | 4096 | 1 | 1 | 1 | n/a | 8 | 40960 | 937 | + +#### Model: Qwen3_235B_a22B + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 256 | MXFP8 | 8192 | 2 | 4096 | 1 | 4 | 1 | 12 | 32 | 9008 | 1333 | +| DGX-GB200 | 256 | MXFP8 | 8192 | 1 | 4096 | 1 | 8 | 1 | 3 | 32 | 7360 | 1089 | + +#### Model: Nemotron_3_Nano + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 8 | MXFP8 | 512 | 4 | 8192 | 1 | 1 | 1 | n/a | 8 | 40960 | 905 | +| DGX-GB200 | 8 | MXFP8 | 512 | 2 | 8192 | 1 | 1 | 1 | n/a | 8 | 34816 | 776 | + +#### Model: Nemotron_3_Super + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 64 | MXFP8 | 512 | 1 | 8192 | 1 | 1 | 1 | n/a | 64 | 9728 | 827 | +| DGX-GB300 | 64 | NVFP4 | 512 | 1 | 8192 | 1 | 1 | 1 | n/a | 64 | 9984 | 845 | +| DGX-GB200 | 64 | MXFP8 | 512 | 1 | 8192 | 2 | 1 | 1 | n/a | 64 | 7040 | 598 | +| DGX-GB200 | 64 | NVFP4 | 512 | 1 | 8192 | 2 | 1 | 1 | n/a | 64 | 7040 | 598 | + +[^gpt-oss-note]: Performance regression in GPT-OSS 120B pre-training is a [known issue](https://github.com/NVIDIA/cudnn-frontend/issues/256). It will be resolved with the next container release. + ## 26.06 NeMo Container ### Pre-Training Performance diff --git a/docs/performance-summary.md b/docs/performance-summary.md index 419427cd7b9..144e33e837a 100644 --- a/docs/performance-summary.md +++ b/docs/performance-summary.md @@ -34,7 +34,7 @@ The performance data includes: --- -## 26.06.01 NeMo Container +## 26.08 NeMo Container ### Pre-Training Performance @@ -42,48 +42,66 @@ The performance data includes: | System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | |--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| -| DGX-GB300 | 256 | MXFP8 | 4096 | 1 | 4096 | 1 | 2 | 1 | 8 | 32 | 6304 | 1640 | -| DGX-GB200 | 256 | MXFP8 | 4096 | 1 | 4096 | 1 | 4 | 1 | 4 | 64 | 4928 | 1280 | +| DGX-GB300 | 256 | MXFP8 | 4096 | 1 | 4096 | 1 | 2 | 1 | 8 | 32 | 6288 | 1635 | +| DGX-GB200 | 256 | MXFP8 | 4096 | 1 | 4096 | 1 | 4 | 1 | 4 | 64 | 4912 | 1276 | + +#### Model: DeepSeekV4 Flash + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 128 | MXFP8 | 2048 | 1 | 4096 | 1 | 1 | 1 | n/a | 64 | 8224 | 748 | +| DGX-GB200 | 128 | MXFP8 | 2048 | 1 | 4096 | 1 | 1 | 1 | n/a | 64 | 7776 | 706 | #### Model: GPT OSS 120B | System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | |--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| -| DGX-GB300 | 64 | MXFP8 | 1280 | 4 | 4096 | 1 | 1 | 1 | n/a | 16 | 20288 | 661[^gpt-oss-note] | -| DGX-GB200 | 64 | MXFP8 | 1280 | 4 | 4096 | 1 | 1 | 1 | n/a | 64 | 18304 | 597[^gpt-oss-note] | +| DGX-GB300 | 64 | MXFP8 | 1280 | 4 | 4096 | 1 | 1 | 1 | n/a | 16 | 33024 | 1077 | +| DGX-GB200 | 64 | MXFP8 | 1280 | 4 | 4096 | 1 | 1 | 1 | n/a | 64 | 28672 | 934 | #### Model: Qwen3_30B_a3B | System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | |--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| | DGX-GB300 | 8 | MXFP8 | 512 | 8 | 4096 | 1 | 1 | 1 | n/a | 8 | 44544 | 1029 | -| DGX-GB200 | 8 | MXFP8 | 512 | 4 | 4096 | 1 | 1 | 1 | n/a | 8 | 40960 | 937 | +| DGX-GB200 | 8 | MXFP8 | 512 | 4 | 4096 | 1 | 1 | 1 | n/a | 8 | 39936 | 923 | #### Model: Qwen3_235B_a22B | System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | |--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| -| DGX-GB300 | 256 | MXFP8 | 8192 | 2 | 4096 | 1 | 4 | 1 | 12 | 32 | 9008 | 1333 | -| DGX-GB200 | 256 | MXFP8 | 8192 | 1 | 4096 | 1 | 8 | 1 | 3 | 32 | 7360 | 1089 | +| DGX-GB300 | 256 | MXFP8 | 8192 | 2 | 4096 | 1 | 4 | 1 | 12 | 32 | 8816 | 1306 | +| DGX-GB200 | 256 | MXFP8 | 8192 | 1 | 4096 | 1 | 8 | 1 | 3 | 32 | 7280 | 1077 | #### Model: Nemotron_3_Nano | System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | |--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| -| DGX-GB300 | 8 | MXFP8 | 512 | 4 | 8192 | 1 | 1 | 1 | n/a | 8 | 40960 | 905 | -| DGX-GB200 | 8 | MXFP8 | 512 | 2 | 8192 | 1 | 1 | 1 | n/a | 8 | 34816 | 776 | +| DGX-GB300 | 8 | MXFP8 | 512 | 4 | 8192 | 1 | 1 | 1 | n/a | 8 | 39936 | 901 | + +#### Model: Nemotron_3_5_Lightning + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB200 | 8 | MXFP8 | 512 | 2 | 8192 | 1 | 1 | 1 | n/a | 8 | 28672 | 807 | #### Model: Nemotron_3_Super | System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | |--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| -| DGX-GB300 | 64 | MXFP8 | 512 | 1 | 8192 | 1 | 1 | 1 | n/a | 64 | 9728 | 827 | -| DGX-GB300 | 64 | NVFP4 | 512 | 1 | 8192 | 1 | 1 | 1 | n/a | 64 | 9984 | 845 | -| DGX-GB200 | 64 | MXFP8 | 512 | 1 | 8192 | 2 | 1 | 1 | n/a | 64 | 7040 | 598 | -| DGX-GB200 | 64 | NVFP4 | 512 | 1 | 8192 | 2 | 1 | 1 | n/a | 64 | 7040 | 598 | +| DGX-GB300 | 64 | MXFP8 | 512 | 1 | 8192 | 1 | 1 | 1 | n/a | 64 | 9856 | 831 | +| DGX-GB300 | 64 | NVFP4 | 512 | 1 | 8192 | 1 | 1 | 1 | n/a | 64 | 10240 | 864 | +| DGX-GB200 | 64 | MXFP8 | 512 | 1 | 8192 | 2 | 1 | 1 | n/a | 64 | 7040 | 600 | +| DGX-GB200 | 64 | NVFP4 | 512 | 1 | 8192 | 2 | 1 | 1 | n/a | 64 | 7168 | 610 | + +#### Model: Nemotron_3_Ultra + +| System | #-GPUs | Precision | GBS | MBS | Sequence Length | TP | PP | CP | VP | EP | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|--------|--------|-----------|-----|-----|-----------------|----|----|----|----|----|-----------------------|-------------------------| +| DGX-GB300 | 256 | MXFP8 | 256 | 1 | 8192 | 1 | 1 | 1 | n/a | 64 | 3552 | 1281 | +| DGX-GB200 | 256 | MXFP8 | 256 | 1 | 8192 | 2 | 1 | 1 | n/a | 64 | 2528 | 923 | [^moe-training-note]: In MoE training benchmarks, we force-balance the token distribution among experts and all benchmarks are token-dropless. -[^gpt-oss-note]: Performance regression in GPT-OSS 120B pre-training is a [known issue](https://github.com/NVIDIA/cudnn-frontend/issues/256). It will be resolved with the next container release. ## Archive From c2e93d863e13112b772cd8d617198fb7de0ae12d Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 12 Aug 2026 00:03:55 +0200 Subject: [PATCH 76/85] cp: `fix(deepseek-v4): use full-prefix inference (5509)` into `r0.6.0` (#5511) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- examples/models/deepseek_v4/inference.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/models/deepseek_v4/inference.sh b/examples/models/deepseek_v4/inference.sh index 0e660a54c64..cf4bc897ab6 100755 --- a/examples/models/deepseek_v4/inference.sh +++ b/examples/models/deepseek_v4/inference.sh @@ -46,6 +46,7 @@ if [[ -z "${PP:-}" ]]; then fi TP=1 +# DSv4 hybrid attention does not support inference contexts yet, so use full-prefix generation. # Inference directly from the HF checkpoint (Bridge dequantises in-flight). uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ examples/conversion/hf_to_megatron_generate_text.py \ @@ -53,6 +54,7 @@ uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ --prompt "${PROMPT}" \ --max_new_tokens 100 \ --tp ${TP} --pp ${PP} --ep ${EP} \ + --legacy-full-prefix \ --trust-remote-code # Inference from a previously-imported Megatron checkpoint (faster cold start). @@ -65,5 +67,6 @@ if [[ -d "${MEGATRON_DIR}/iter_0000000" ]]; then --prompt "${PROMPT}" \ --max_new_tokens 100 \ --tp ${TP} --pp ${PP} --ep ${EP} \ + --legacy-full-prefix \ --trust-remote-code fi From b88831815ffbe5116eed51054a7e1fb0ad948c11 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 12 Aug 2026 01:21:36 +0200 Subject: [PATCH 77/85] cp: `fix(recipe): prevent Qwen3.5-VL FSDP SFT OOM (5508)` into `r0.6.0` (#5515) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- .../bridge/recipes/qwen_vl/h100/qwen35_vl.py | 12 ++++++++---- .../recipes/qwen_vl/test_qwen35_vl_recipes.py | 6 ++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py index 0325632ef5d..5eb386c3122 100644 --- a/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py +++ b/src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py @@ -951,7 +951,9 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer: cfg.model.attention_backend = "auto" cfg.model.gradient_accumulation_fusion = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" + # Native fused cross entropy materializes a full FP32 vocabulary-sized + # softmax buffer. Use the TE kernel to keep the 248K-token loss bounded. + cfg.model.cross_entropy_fusion_impl = "te" # MoE settings cfg.model.moe_token_dispatcher_type = "alltoall" @@ -965,10 +967,12 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer: cfg.model.moe_router_padding_for_fp8 = False # Memory saving - cfg.model.recompute_granularity = None + # Full SFT retains FP32 optimizer state alongside the FSDP communication + # pools, so recompute each transformer layer to bound live activations. + cfg.model.recompute_granularity = "full" cfg.model.recompute_modules = None - cfg.model.recompute_method = None - cfg.model.recompute_num_layers = None + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 cfg.model.fine_grained_activation_offloading = False cfg.model.offload_modules = None diff --git a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py index 3982c201e8c..4c07b1ed6b8 100644 --- a/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py +++ b/tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py @@ -431,6 +431,12 @@ def test_qwen35_vl_35b_a3b_fsdp_sft_defaults(monkeypatch: pytest.MonkeyPatch): assert cfg.model.sequence_parallel is False assert cfg.model.moe_token_dispatcher_type == "alltoall" assert cfg.model.moe_router_fusion is True + assert cfg.model.cross_entropy_loss_fusion is True + assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.recompute_granularity == "full" + assert cfg.model.recompute_modules is None + assert cfg.model.recompute_method == "uniform" + assert cfg.model.recompute_num_layers == 1 assert cfg.ddp.use_megatron_fsdp is True assert cfg.ddp.fsdp_double_buffer is True assert cfg.ddp.megatron_fsdp_max_pool_double_buffer is True From 5b70a4c4bcbb883c1e901bbeaecd98246a381e68 Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:22:59 -0700 Subject: [PATCH 78/85] fix(recipe): recompute Gemma 4 VL SFT layers (#5506) Signed-off-by: yaoyu-33 --- .../gemma-4-26b-a4b-it/card.yaml | 49 ++++++++++--------- .../recipes/gemma4_vl/h100/gemma4_vl.py | 14 +++--- .../recipes/test_gemma4_vl_recipes.py | 8 +-- 3 files changed, 37 insertions(+), 34 deletions(-) diff --git a/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml b/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml index 06b7e6da2b4..9234f80a013 100644 --- a/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml +++ b/examples/model_verification_cards/gemma-4-26b-a4b-it/card.yaml @@ -11,11 +11,12 @@ summary: > forward comparison remains unverified because the recorded run compared a tensor-parallel padding position rather than the prompt's next token. Decoder-focused SFT is verified on one 8-GPU H100 node with TP4/PP1/EP8/ETP1, - a frozen vision encoder, and selective core-attention and MoE-activation - recompute. The real CORD-v2 run completed 10 optimizer steps with stable - memory and saved iter_0000010. LoRA PEFT is also verified on four H100 GPUs - with TP2/PP1/EP4/ETP1 and a complete adapter checkpoint. Long-context SFT is - pending. The published Bridge recipes do not currently provide VLM pretraining. + a frozen vision encoder, and full-layer uniform activation recompute with one + layer per recompute unit. The real CORD-v2 run completed 10 optimizer steps + with stable memory and saved iter_0000010. LoRA PEFT is also verified on four + H100 GPUs with TP2/PP1/EP4/ETP1 and a complete adapter checkpoint. + Long-context SFT is pending. The published Bridge recipes do not currently + provide VLM pretraining. verification_index: model_level: verified: @@ -181,7 +182,7 @@ items: status: verified precision: bf16 enabled_features: {} - bridge_commit: 81e559ca4d579cf2d9ed5b8f3f48bef82e06451f # pragma: allowlist secret + bridge_commit: f3130a0e37c78945ecf98b301dfc90505c3c305f # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 --recipe gemma4_vl_26b_sft_config --mode sft --dataset cord-v2 @@ -192,28 +193,30 @@ items: --save_dir work/model-verification/gemma-4-26b-a4b-it/sft-checkpoints --save_interval 10 logger.log_interval=1 logger.log_throughput=true logger.save_config_filepath=work/model-verification/gemma-4-26b-a4b-it/sft-post-setup.yaml - last_verified: 2026-08-04 + last_verified: 2026-08-11 metrics: - initial_loss: 1.268018 - final_loss: 0.06543536 - last_10_steps_step_time_ms_avg: 18699.99 - last_10_steps_model_tflops_per_gpu_avg: 26.47 - peak_allocated_memory_gib: 70.624 - peak_reserved_memory_gib: 71.021 + initial_loss: 1.282612 + final_loss: 0.09127883 + last_10_steps_step_time_ms_avg: 15936.25 + last_10_steps_model_tflops_per_gpu_avg: 32.20 + peak_allocated_memory_gib: 65.223 + peak_reserved_memory_gib: 72.356 expected_result: > - On 2026-08-04 the single-node TP4/PP1/EP8/ETP1 run loaded the immutable + On 2026-08-11 the single-node TP4/PP1/EP8/ETP1 run loaded the immutable BF16 checkpoint and real CORD-v2 revision, then completed exactly 10 optimizer steps at MBS1 and GBS32. The persisted post-setup config confirmed a frozen vision encoder, trainable projection and language - model, expandable CUDA allocator segments, and selective core_attn plus - moe_act recompute with CUDA graphs and CPU offload disabled. Loss moved - from 1.268018 to 0.06543536, final grad norm was 1.561, and every step - reported zero skipped and zero NaN iterations. Across all eight ranks, - allocated memory returned to 56.493 GiB after each step, while - cumulative peak allocated and reserved memory remained at or below - 70.624 and 71.021 GiB. The run saved a complete iter_0000010 torch_dist - checkpoint with eight rank shards, finalized metadata, training state, - and a latest-checkpoint tracker containing iteration 10. + model, expandable CUDA allocator segments, and full-layer uniform + recompute with one layer per recompute unit directly from the recipe + defaults, while CUDA graphs and CPU offload remained disabled. Loss + moved from 1.282612 to 0.09127883, final grad norm was 3.090, and all 10 + steps reported zero skipped and zero NaN iterations. After iteration 2, + logger-reported ranks returned to + 62.134-62.145 GiB allocated, while cumulative peak allocated and + reserved memory remained at or below 65.223 and 72.356 GiB. The run + saved a complete iter_0000010 torch_dist checkpoint with eight rank + shards, finalized metadata, run config, training state, and a + latest-checkpoint tracker containing iteration 10. sft_export_inference: H100: diff --git a/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py b/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py index aed09f0af6d..ca997f662ca 100644 --- a/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py +++ b/src/megatron/bridge/recipes/gemma4_vl/h100/gemma4_vl.py @@ -111,7 +111,7 @@ def gemma4_vl_26b_sft_8gpu_h100_bf16_config() -> ConfigContainer: Default configuration: 8 GPUs - TP=4, PP=1, EP=8, ETP=1 (dense DP=2, expert DP=1) - - Selective core-attention and MoE-activation recompute + - Full-layer uniform activation recompute - Frozen vision encoder with trainable projection and language model - LR=5e-5 (full SFT) - Sequence length: 4096 @@ -130,12 +130,12 @@ def gemma4_vl_26b_sft_8gpu_h100_bf16_config() -> ConfigContainer: cfg.model.expert_model_parallel_size = 8 cfg.model.expert_tensor_parallel_size = 1 - # Real CORD-v2 measurements require both modules to preserve H100 memory - # headroom after optimizer-state initialization. Full-layer recompute is not needed. - cfg.model.recompute_granularity = "selective" - cfg.model.recompute_modules = ["core_attn", "moe_act"] - cfg.model.recompute_method = None - cfg.model.recompute_num_layers = None + # Real-data SFT has batch-dependent activation and cross-entropy workspace peaks. + # Checkpoint each decoder layer to preserve H100 memory headroom after optimizer initialization. + cfg.model.recompute_granularity = "full" + cfg.model.recompute_modules = None + cfg.model.recompute_method = "uniform" + cfg.model.recompute_num_layers = 1 # Decoder-focused SFT keeps the vision encoder fixed while training the # multimodal projection and language model. diff --git a/tests/unit_tests/recipes/test_gemma4_vl_recipes.py b/tests/unit_tests/recipes/test_gemma4_vl_recipes.py index ee0fd95feff..cb5a1e9a7ac 100644 --- a/tests/unit_tests/recipes/test_gemma4_vl_recipes.py +++ b/tests/unit_tests/recipes/test_gemma4_vl_recipes.py @@ -63,10 +63,10 @@ def test_gemma4_vl_sft_uses_memory_stable_8gpu_contract(monkeypatch: pytest.Monk assert cfg.model.expert_model_parallel_size == 8 assert cfg.model.expert_tensor_parallel_size == 1 assert cfg.model.sequence_parallel is True - assert cfg.model.recompute_granularity == "selective" - assert cfg.model.recompute_modules == ["core_attn", "moe_act"] - assert cfg.model.recompute_method is None - assert cfg.model.recompute_num_layers is None + assert cfg.model.recompute_granularity == "full" + assert cfg.model.recompute_modules is None + assert cfg.model.recompute_method == "uniform" + assert cfg.model.recompute_num_layers == 1 assert cfg.model.cuda_graph_impl == "none" assert cfg.model.freeze_vision_model is True assert cfg.model.freeze_vision_projection is False From 7b5057e03224f08b122b132a7f8b1d361b039a7d Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Wed, 12 Aug 2026 01:23:37 +0200 Subject: [PATCH 79/85] cp: `docs(dsv4): preserve container packages during MCore sync (5514)` into `r0.6.0` (#5516) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- examples/models/deepseek_v4/README.md | 34 ++++++++++++++----- examples/models/deepseek_v4/slurm_pretrain.sh | 3 +- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index b6ba2cdc4d2..bcc1e675258 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -8,12 +8,35 @@ The bridge supports four published variants out of the same code path. The on-di The pretraining recipes were tested with Megatron-LM `dev` commit `35f36c7c9dba` plus PR [#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839) (`f04b762406f0` in the OCI test checkout). The Megatron-LM copy inside the current NeMo FW container is not expected to work for these recipes. +The NeMo Framework container uses one shared `/opt/venv` for several source +projects. A plain `uv sync` is exact by default and removes packages that are +not in the Megatron Bridge dependency graph, including container-provided +packages such as NeMo Run. When switching MCore in this container, preserve +those packages and then restore the aggregate container pins: + ```bash ./scripts/switch_mcore.sh dev -uv sync +uv sync --inexact +(cd /opt/NeMo-FW && uv sync --locked --all-groups --inexact) +export UV_NO_SYNC=1 +``` + +Keep `UV_NO_SYNC=1` set while running the examples. It makes `uv run` calls, +including those inside the example wrapper scripts, skip the implicit exact +sync. The unlocked dev sync updates `uv.lock`; do not commit that generated +change. From a clean checkout, return to the pinned main-branch submodule with: + +```bash +./scripts/switch_mcore.sh main +git restore uv.lock +uv sync --locked --inexact +(cd /opt/NeMo-FW && uv sync --locked --all-groups --inexact) +export UV_NO_SYNC=1 ``` -Use `./scripts/switch_mcore.sh main` and `uv sync --locked` to return to the pinned main-branch submodule. +In a standalone Megatron Bridge environment with its own virtual environment, +exact sync is appropriate: use `uv sync` after switching to dev and restore the +tracked lock file followed by `uv sync --locked` when switching back to main. The full-scale `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config` performance recipe preserves the stack validated by Megatron Bridge PR @@ -142,11 +165,6 @@ DSv4 currently requires **TP=1** because MLA tensor parallelism is not supported - **Fused mHC: use the unfused path for SFT.** The fused cuTile mHC kernel needs sm_100 (not H100) and is on by default in the bridge config; it works for import/inference on Blackwell, but **NaNs in SFT training** (see SFT Blockers). The SFT recipes therefore force `use_fused_mhc=False` — the validated unfused/reference path, which also runs on Hopper. -- **`fast_hadamard_transform` is required by the DSA attention variant.** `csa.py` and `dsa.py` import `hadamard_transform` from this package and hard-assert availability — there is no in-tree PyTorch fallback. Install from the Dao-AILab git repo (the PyPI source distribution is incomplete; see the sibling GLM-5 [README](../glm/glm5/README.md#pre-requisites) for the same dependency): - - ```bash - pip install --no-build-isolation \ - git+https://github.com/Dao-AILab/fast-hadamard-transform.git - ``` +- **`fast_hadamard_transform` is required by the DSA attention variant.** `csa.py` and `dsa.py` import `hadamard_transform` from this package and hard-assert availability — there is no in-tree PyTorch fallback. It is a required Megatron Bridge dependency, pinned to the complete Dao-AILab git source because the PyPI source distribution is incomplete, and is installed by the `uv sync` commands above. - **Logit parity is verified for Flash and Flash-Base** against the official inference stack at last-real-token logits. The remaining gap is structural, from different attention/HC kernel decompositions and accumulation precisions between MCore and official inference. diff --git a/examples/models/deepseek_v4/slurm_pretrain.sh b/examples/models/deepseek_v4/slurm_pretrain.sh index ddea2836ed5..b494cc55aa8 100755 --- a/examples/models/deepseek_v4/slurm_pretrain.sh +++ b/examples/models/deepseek_v4/slurm_pretrain.sh @@ -90,7 +90,8 @@ export NCCL_PXN_DISABLE=1 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True export CUDA_DEVICE_MAX_CONNECTIONS=1 -# Pre-sync once before submitting jobs: UV_CACHE_DIR=/path/to/cache uv sync +# Before submitting, follow README.md's MCore Checkout instructions. In a NeMo +# Framework container, use its two-stage inexact sync; do not run plain `uv sync`. # export UV_CACHE_DIR="/path/to/shared/uv_cache" # export HF_HOME="/path/to/shared/HF_HOME" # export HF_TOKEN="hf_your_token_here" From c93251151adeeadbae3ff2a2bf5ee7a1c34cff01 Mon Sep 17 00:00:00 2001 From: Ajay Date: Thu, 13 Aug 2026 12:45:37 -0700 Subject: [PATCH 80/85] Update nemo-evaluator path in fw_pyproject.toml (#5571) Signed-off-by: Ajay --- docker/common/fw_pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/common/fw_pyproject.toml b/docker/common/fw_pyproject.toml index d70ef11915d..4957e2b0809 100644 --- a/docker/common/fw_pyproject.toml +++ b/docker/common/fw_pyproject.toml @@ -41,7 +41,7 @@ test = ["coverage>=7.8.1"] [tool.uv.sources] "nemo-export-deploy" = { path = "/opt/Export-Deploy", editable = true } -"nemo-evaluator" = { path = "/opt/Evaluator", editable = true } +"nemo-evaluator" = { path = "/opt/Evaluator/packages/nemo-evaluator", editable = true } "nemo-run" = { path = "/opt/Run", editable = true } "lightning" = { path = "stubs/lightning" } From 8e43ce69d2f5dbe4b4ee7bcf809810c7962f5b29 Mon Sep 17 00:00:00 2001 From: Ajay Date: Thu, 13 Aug 2026 16:43:52 -0700 Subject: [PATCH 81/85] docs: Update 0.6.0 version entries for docs (#5567) Signed-off-by: Ajay --- docs/versions1.json | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/docs/versions1.json b/docs/versions1.json index 471dbd33202..fee0ec52657 100644 --- a/docs/versions1.json +++ b/docs/versions1.json @@ -4,11 +4,18 @@ "url": "https://docs.nvidia.com/nemo/megatron-bridge/nightly/" }, { + "name": "0.6.0 (latest) · 26.08", "version": "0.6.0", - "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.6.0/" + "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.6.0/", + "preferred": true + }, + { + "name": "0.5.0 · 26.06", + "version": "0.5.0", + "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.5.0/" }, { - "name": "0.4.2 (latest) · 26.04.01", + "name": "0.4.2 · 26.04.01", "version": "0.4.2", "url": "https://docs.nvidia.com/nemo/megatron-bridge/latest/", "preferred": true From cea3a730ae2ed9290b8aac54c323ca9e66181585 Mon Sep 17 00:00:00 2001 From: svcnvidia-nemo-ci Date: Fri, 14 Aug 2026 20:06:25 +0200 Subject: [PATCH 82/85] cp: `fix(deepseek): default DeepSeek V4 to HybridEP (5543)` into `r0.6.0` (#5582) Signed-off-by: Chen Cui Signed-off-by: NeMo Bot Co-authored-by: Chen Cui --- examples/models/deepseek_v4/README.md | 11 ++++++++ examples/models/deepseek_v4/inference.sh | 5 ++++ .../models/deepseek/deepseek_v4_bridge.py | 5 +++- .../deepseek/gb200/deepseek_v4.py | 3 ++- .../deepseek/gb300/deepseek_v4.py | 2 ++ .../recipes/deepseek/gb200/deepseek_v4.py | 3 ++- .../deepseek/test_deepseek_v4_bridge.py | 18 +++++++++++++ .../recipes/test_deepseek_recipes.py | 27 +++++++++++++++++++ .../test_deepseek_v4_flash_perf_recipes.py | 5 +++- .../recipes/test_r050_perf_recipe_ports.py | 2 ++ 10 files changed, 77 insertions(+), 4 deletions(-) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index bcc1e675258..e57dda2e711 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -61,6 +61,17 @@ performance features, so it is not a supported runtime for that recipe. - `slurm_pretrain.sh` runs the DeepSeek-V4-Flash pretraining recipes. - `slurm_sft.sh` runs DeepSeek-V4-Flash full SFT end to end (import, then fine-tune) on Hopper or Blackwell, with MTP on or off. +Bridge-created DeepSeek-V4 providers default to the `flex`/HybridEP MoE +dispatcher with 16 dispatcher SMs and unfused HybridEP permutation. Native +Megatron checkpoints retain the dispatcher configuration with which they were +saved, so older checkpoints may still use `alltoall`. HybridEP requires a +runtime containing the HybridEP package. It auto-detects accessible ranks and +fabric support when its topology environment variables are unset. To override +detection, set +`NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN`, `NVLINK_DOMAIN_SIZE`, and +`USE_MNNVL` to values that describe the actual deployment. Hardware recipes +retain their independently qualified dispatcher overrides. + Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each script's header comments for the expected environment variables and `#SBATCH` directives to edit before submitting. ## Pretraining Recipes diff --git a/examples/models/deepseek_v4/inference.sh b/examples/models/deepseek_v4/inference.sh index cf4bc897ab6..f1dbef30137 100755 --- a/examples/models/deepseek_v4/inference.sh +++ b/examples/models/deepseek_v4/inference.sh @@ -22,8 +22,13 @@ # EP: expert-parallel size (default: 4 for Flash, 8 for Pro) # PP: pipeline-parallel size (default: 1 for Flash, 4 for Pro) # PROMPT: prompt string (default: "Explain hyper-connections in transformer models.") +# NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: optional HybridEP topology override +# NVLINK_DOMAIN_SIZE: optional physical NVLink-domain size override +# USE_MNNVL: optional HybridEP fabric override (0 or 1) # # Defaults below are for GB200 (192 GB). For H100 (80 GB) configs, see README.md. +# HybridEP auto-detects accessible ranks and fabric support when the topology +# overrides are unset. If set, the values must describe the actual deployment. set -xeuo pipefail diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 29c4c4da82f..cdccfa4a44a 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -481,7 +481,10 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.gated_linear_unit = True provider.moe_grouped_gemm = True provider.moe_router_pre_softmax = False # V4 uses post-topk normalisation - provider.moe_token_dispatcher_type = "alltoall" + provider.moe_token_dispatcher_type = "flex" + provider.moe_flex_dispatcher_backend = "hybridep" + provider.moe_flex_dispatcher_num_sms = 16 + provider.moe_permute_fusion_into_hybridep = False provider.moe_router_load_balancing_type = "noaux_tc" provider.moe_shared_expert_overlap = True provider.moe_router_score_function = hf_config.scoring_func # "sqrtsoftplus" diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py index e7379dbb991..aa26ba7dbe8 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -44,7 +44,6 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_shared_expert_overlap = False - cfg.model.moe_hybridep_num_sms = 32 cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_router_fusion = True cfg.model.moe_router_force_load_balancing = True @@ -60,6 +59,8 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: cfg.model.fine_grained_offloading_max_inflight_offloads = None _benchmark_common(cfg, cross_entropy_impl="native") + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.moe_hybridep_num_sms = None set_full_iteration_cuda_graph(cfg.model) cfg.model.cuda_graph_warmup_steps = 3 diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py index 6c5d18c8314..8782cc5c61d 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py @@ -72,6 +72,8 @@ def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.model.recompute_modules = ["mla_up_proj", "mhc"] _benchmark_common(cfg, cross_entropy_impl="native") + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.moe_hybridep_num_sms = None cfg.model.cuda_graph_impl = "full_iteration" cfg.model.cuda_graph_scope = [] diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index aee81537eb4..677dabbe01a 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -80,7 +80,8 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_hybridep_num_sms = 16 + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_hybridep_num_sms = None cfg.model.moe_grouped_gemm = True cfg.model.moe_permute_fusion = True cfg.model.moe_aux_loss_coeff = 0.0 diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py index 68eab0f5721..f2223f9c1d9 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -427,6 +427,24 @@ def test_provider_bridge_forces_full_rotary_percent(self): assert out.csa_compress_rotary_base == 160000 +class TestDeepSeekV4MoEDispatcher: + """DeepSeek-V4 providers use the inference-validated HybridEP path.""" + + def test_provider_bridge_uses_hybridep(self): + hf_pretrained = MagicMock() + hf_pretrained.config = _deepseek_v4_hf_config() + provider = MagicMock() + + bridge = DeepSeekV4Bridge.__new__(DeepSeekV4Bridge) + with patch.object(MegatronModelBridge, "provider_bridge", return_value=provider): + out = bridge.provider_bridge(hf_pretrained) + + assert out.moe_token_dispatcher_type == "flex" + assert out.moe_flex_dispatcher_backend == "hybridep" + assert out.moe_flex_dispatcher_num_sms == 16 + assert out.moe_permute_fusion_into_hybridep is False + + class TestDeepSeekV4HardwareDefaults: """DSv4 Blackwell-only fused kernels must not default on for Hopper.""" diff --git a/tests/unit_tests/recipes/test_deepseek_recipes.py b/tests/unit_tests/recipes/test_deepseek_recipes.py index 365d8b16a07..e984f6ccfb7 100644 --- a/tests/unit_tests/recipes/test_deepseek_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_recipes.py @@ -426,6 +426,33 @@ def test_deepseek_v4_base_recipe_uses_blackwell_defaults(monkeypatch: pytest.Mon assert cfg.train.micro_batch_size == 1 +@pytest.mark.parametrize( + ("recipe_name", "expected_dispatcher"), + [ + ("deepseek_v4_flash_pretrain_config", "alltoall"), + ("deepseek_v4_flash_pretrain_mxfp8_config", "alltoall"), + ("deepseek_v4_flash_pretrain_muon_config", "alltoall"), + ("deepseek_v4_flash_sft_config", "alltoall"), + ("deepseek_v4_flash_no_mtp_sft_config", "alltoall"), + ("deepseek_v4_pro_pretrain_config", "alltoall"), + ("deepseek_v4_pro_pretrain_mxfp8_config", "alltoall"), + ("deepseek_v4_flash_pretrain_gb200_config", "flex"), + ("deepseek_v4_flash_pretrain_mxfp8_gb200_config", "flex"), + ("deepseek_v4_flash_pretrain_muon_gb200_config", "flex"), + ], +) +def test_deepseek_v4_recipes_keep_hardware_qualified_dispatcher( + recipe_name: str, expected_dispatcher: str, monkeypatch: pytest.MonkeyPatch +): + cfg = _build_deepseek_v4_recipe(recipe_name, monkeypatch) + + assert cfg.model.moe_token_dispatcher_type == expected_dispatcher + if expected_dispatcher == "flex": + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_flex_dispatcher_num_sms == 16 + assert cfg.model.moe_hybridep_num_sms is None + + @pytest.mark.parametrize( "recipe_name", [ diff --git a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py index 3a6670d6b88..fdcd31fea67 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py @@ -50,7 +50,8 @@ def test_deepseek_v4_flash_128gpu_gb200_fp8mx_config() -> None: assert cfg.model.moe_token_dispatcher_type == "flex" assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.moe_router_force_load_balancing is True - assert cfg.model.moe_hybridep_num_sms == 32 + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert cfg.model.moe_hybridep_num_sms is None assert cfg.model.moe_hybridep_num_sms_preprocessing == 108 assert cfg.model.recompute_granularity == "selective" assert cfg.model.recompute_modules == ["mla_up_proj"] @@ -121,4 +122,6 @@ def test_deepseek_v4_flash_128gpu_gb300_fp8mx_config() -> None: assert cfg.train.global_batch_size == 2048 assert cfg.train.micro_batch_size == 1 assert cfg.model.recompute_modules == ["mla_up_proj"] + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert cfg.model.moe_hybridep_num_sms is None assert is_full_iteration_cuda_graph(cfg.model) diff --git a/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py b/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py index fe3340fab00..51bd73c890a 100644 --- a/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py +++ b/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py @@ -154,6 +154,8 @@ def test_deepseek_v4_pro_gb300_matches_r050_performance_config() -> None: assert cfg.model.recompute_modules == ["mla_up_proj", "mhc"] assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert cfg.model.moe_hybridep_num_sms is None assert cfg.model.moe_shared_expert_overlap is False assert cfg.model.cuda_graph_impl == "full_iteration" From 0f113aba4bf625ab31c3a99f771ee24aa97d6dfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Mon, 17 Aug 2026 18:52:57 +0200 Subject: [PATCH 83/85] docs: fix preferred version and add 0.5.1 to the version picker (#5623) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- docs/versions1.json | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/versions1.json b/docs/versions1.json index fee0ec52657..4544b7fd574 100644 --- a/docs/versions1.json +++ b/docs/versions1.json @@ -9,6 +9,11 @@ "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.6.0/", "preferred": true }, + { + "name": "0.5.1 · 26.06.01", + "version": "0.5.1", + "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.5.1/" + }, { "name": "0.5.0 · 26.06", "version": "0.5.0", @@ -17,8 +22,7 @@ { "name": "0.4.2 · 26.04.01", "version": "0.4.2", - "url": "https://docs.nvidia.com/nemo/megatron-bridge/latest/", - "preferred": true + "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.4.2/" }, { "name": "0.4.0 · 26.04", From 51885cf132b2814188b6855c25a8588254274c2a Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 19 Aug 2026 18:50:28 +0000 Subject: [PATCH 84/85] =?UTF-8?q?beep=20boop=20=F0=9F=A4=96:=20Bumping=20M?= =?UTF-8?q?egatron-Bridge=20to=20v0.6.1=20[skip=20ci]?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: github-actions[bot] --- src/megatron/bridge/package_info.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/megatron/bridge/package_info.py b/src/megatron/bridge/package_info.py index 2e5611d84da..266cd97a815 100644 --- a/src/megatron/bridge/package_info.py +++ b/src/megatron/bridge/package_info.py @@ -15,7 +15,7 @@ MAJOR = 0 MINOR = 6 -PATCH = 0 +PATCH = 1 PRE_RELEASE = "" # Use the following formatting: (major, minor, patch, pre-release) From 53bea08c8417627d5e83af6c9c16163a47e3a8b1 Mon Sep 17 00:00:00 2001 From: Elia Palme Date: Tue, 1 Sep 2026 16:02:31 +0200 Subject: [PATCH 85/85] models: add Apertus / Apertus 1.5 bridges Upstream 0.6.x has no Apertus support. This adds: - models/apertus/: ApertusBridge for ApertusForCausalLM, plus MCoreXIELU (a Megatron-side wrapper of transformers' XIELUActivation that requires the CUDA xielu kernel and TP-sums its alpha grads) and get_apertus_decoder_block_spec (RMSNorm q/k-layernorm + MCoreXIELU MLP activation). - models/apertus1p5/: Apertus1p5Bridge for the multimodal Apertus1p5ForConditionalGeneration -- hyperparameters read from config.text_config, HF weight keys under model.language_model.*, pruned-LM-head handling (output_vocab_size < vocab_size), vision/audio tokenizer towers left unmapped (target is a text-only GPTModel). - models/__init__.py: import both (the @register_bridge decorator only fires on import) and add to __all__. - hf_pretrained/safe_config_loader.py: replace the filelock around AutoConfig.from_pretrained with contextlib.nullcontext -- fcntl.flock raises ENOLCK/ESTALE in-container on CSCS Lustre and the config is written once by rank 0 and read-only afterward. --- src/megatron/bridge/models/__init__.py | 4 + .../bridge/models/apertus/__init__.py | 3 + .../bridge/models/apertus/apertus_bridge.py | 169 ++++++++++ .../bridge/models/apertus1p5/__init__.py | 3 + .../models/apertus1p5/apertus1p5_bridge.py | 311 ++++++++++++++++++ .../hf_pretrained/safe_config_loader.py | 3 +- 6 files changed, 492 insertions(+), 1 deletion(-) create mode 100644 src/megatron/bridge/models/apertus/__init__.py create mode 100644 src/megatron/bridge/models/apertus/apertus_bridge.py create mode 100644 src/megatron/bridge/models/apertus1p5/__init__.py create mode 100644 src/megatron/bridge/models/apertus1p5/apertus1p5_bridge.py diff --git a/src/megatron/bridge/models/__init__.py b/src/megatron/bridge/models/__init__.py index c5ca28820af..eaad0cb1db9 100644 --- a/src/megatron/bridge/models/__init__.py +++ b/src/megatron/bridge/models/__init__.py @@ -13,6 +13,8 @@ # limitations under the License. # Import model providers for easy access +from megatron.bridge.models.apertus import ApertusBridge +from megatron.bridge.models.apertus1p5 import Apertus1p5Bridge from megatron.bridge.models.bailing import ( BailingMoeV2Bridge, ) @@ -206,6 +208,8 @@ "ReplicatedMapping", "RowParallelMapping", "AutoMapping", + "ApertusBridge", + "Apertus1p5Bridge", "BailingMoeV2Bridge", # DeepSeek Models "DeepSeekV2Bridge", diff --git a/src/megatron/bridge/models/apertus/__init__.py b/src/megatron/bridge/models/apertus/__init__.py new file mode 100644 index 00000000000..02bd836884b --- /dev/null +++ b/src/megatron/bridge/models/apertus/__init__.py @@ -0,0 +1,3 @@ +from megatron.bridge.models.apertus.apertus_bridge import ApertusBridge + +__all__ = ["ApertusBridge"] diff --git a/src/megatron/bridge/models/apertus/apertus_bridge.py b/src/megatron/bridge/models/apertus/apertus_bridge.py new file mode 100644 index 00000000000..2f5c207026e --- /dev/null +++ b/src/megatron/bridge/models/apertus/apertus_bridge.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +import torch +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.transformer.spec_utils import get_submodules +from megatron.core.utils import is_torch_min_version +from transformers import ApertusForCausalLM +from transformers.activations import XIELUActivation + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + ColumnParallelMapping, + QKVMapping, + ReplicatedMapping, +) +from megatron.bridge.models.conversion.utils import unwrap_model + +_ROPE_DEFAULTS = { + "rope_type": "llama3", + "original_max_position_embeddings": 8192, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, +} + + +class MCoreXIELU(XIELUActivation): + def __init__(self, *, config): + super().__init__(dtype=config.params_dtype, with_vector_loads=False) + if self._xielu_cuda_obj is None: + raise RuntimeError("CUDA xIELU is required. Install rubber-duck-debug/xielu.") + self.to( + device=torch.device("cpu") + if getattr(config, "use_cpu_initialization", False) + else torch.device("cuda", torch.cuda.current_device()) + ) + self.alpha_p.sum_gradients_across_tp_domain = True + self.alpha_n.sum_gradients_across_tp_domain = True + self._sync_runtime_scalars() + self.register_load_state_dict_post_hook(lambda m, _: m._sync_runtime_scalars()) + + @torch.no_grad() + def _sync_runtime_scalars(self): + """Refresh CUDA-kernel host caches after beta/eps buffers change.""" + if getattr(self.beta, "is_meta", False) or getattr(self.eps, "is_meta", False): + return + self._beta_scalar = float(self.beta.detach().cpu().float().item()) + self._eps_scalar = float(self.eps.detach().cpu().float().item()) + + +def get_apertus_decoder_block_spec(config, vp_stage=None, pp_rank=None): + assert is_torch_min_version("2.4.0a0"), "Torch RMSNorm requires PyTorch >= 2.4" + block_spec = get_gpt_decoder_block_spec( + config, use_transformer_engine=True, normalization="RMSNorm", vp_stage=vp_stage, pp_rank=pp_rank + ) + for layer in block_spec.layer_specs: + attn = layer.submodules.self_attention.submodules + attn.q_layernorm = attn.k_layernorm = (lambda config, hidden_size, eps=1e-5, **_: torch.nn.RMSNorm(hidden_size, eps=eps)) + get_submodules(layer.submodules.mlp).activation_func = MCoreXIELU + return block_spec + + +@MegatronModelBridge.register_bridge(source=ApertusForCausalLM, target=GPTModel, model_type="apertus") +class ApertusBridge(MegatronModelBridge): + @classmethod + def hf_to_megatron_activation(cls, hidden_act: str): + if hidden_act != "xielu": + return super().hf_to_megatron_activation(hidden_act) + return lambda _: (_ for _ in ()).throw(RuntimeError("expected MCoreXIELU")) + + def provider_bridge(self, hf_pretrained): + provider = super().provider_bridge(hf_pretrained) + config = hf_pretrained.config + rope = {**(getattr(config, "rope_scaling", None) or {}), **(getattr(config, "rope_parameters", None) or {})} + rope_type = rope.get("rope_type", rope.get("type", "llama3")) + factor = float(rope.get("factor", 1.0)) + theta = float(rope.get("rope_theta", getattr(config, "rope_theta", 10000.0))) + if config.hidden_act != "xielu": + raise ValueError(f"Expected hidden_act='xielu', got {config.hidden_act!r}") + if config.attention_bias: + raise ValueError("Apertus attention_bias=True is unsupported") + if rope_type != "llama3": + raise ValueError(f"Unsupported Apertus RoPE type: {rope_type!r}") + + provider.apertus_rope_scaling = { + "rope_type": rope_type, + "type": rope_type, + "factor": factor, + "original_max_position_embeddings": int(rope.get("original_max_position_embeddings", 8192)), + "low_freq_factor": float(rope.get("low_freq_factor", 1.0)), + "high_freq_factor": float(rope.get("high_freq_factor", 4.0)), + } + provider.normalization = "RMSNorm" + provider.qk_layernorm = True + provider.gated_linear_unit = False + provider.use_te_activation_func = False + provider.bias_activation_fusion = False + provider.add_bias_linear = False + provider.add_qkv_bias = False + provider.hidden_dropout = 0.0 + provider.rotary_interleaved = False + provider.position_embedding_type = "rope" + provider.rotary_base = theta + provider.rope_scaling = True + provider.rope_scaling_factor = factor + provider.transformer_layer_spec = get_apertus_decoder_block_spec + return provider + + def load_weights_hf_to_megatron(self, hf_pretrained, megatron_model, allowed_mismatched_params=None): + models = super().load_weights_hf_to_megatron( + hf_pretrained, megatron_model, allowed_mismatched_params=allowed_mismatched_params + ) + [ + m._sync_runtime_scalars() + for model in unwrap_model(models) + for m in model.modules() + if isinstance(m, MCoreXIELU) + ] + return models + + @classmethod + def megatron_to_hf_config(cls, provider): + config = super().megatron_to_hf_config(provider) + theta = float(provider.rotary_base) + rope = { + **(getattr(provider, "apertus_rope_scaling", None) or _ROPE_DEFAULTS), + "factor": float(provider.rope_scaling_factor), + } + rope["rope_type"] = rope["type"] = rope.get("rope_type", rope.get("type", "llama3")) + config.update( + hidden_act="xielu", + attention_bias=False, + rope_theta=theta, + rope_scaling=rope, + rope_parameters={**rope, "rope_theta": theta}, + ) + return config + + def mapping_registry(self) -> MegatronMappingRegistry: + L, H = "decoder.layers.*", "model.layers.*" + auto = { + "embedding.word_embeddings.weight": "model.embed_tokens.weight", + "output_layer.weight": "lm_head.weight", + "decoder.final_layernorm.weight": "model.norm.weight", + f"{L}.self_attention.linear_proj.weight": f"{H}.self_attn.o_proj.weight", + f"{L}.self_attention.q_layernorm.weight": f"{H}.self_attn.q_norm.weight", + f"{L}.self_attention.k_layernorm.weight": f"{H}.self_attn.k_norm.weight", + f"{L}.mlp.linear_fc2.weight": f"{H}.mlp.down_proj.weight", + } + repl = { + f"{L}.self_attention.linear_qkv.layer_norm_weight": f"{H}.attention_layernorm.weight", + f"{L}.mlp.linear_fc1.layer_norm_weight": f"{H}.feedforward_layernorm.weight", + **{f"{L}.mlp.activation_func.{n}": f"{H}.mlp.act_fn.{n}" for n in ("alpha_p", "alpha_n", "beta", "eps")}, + } + qkv = QKVMapping( + f"{L}.self_attention.linear_qkv.weight", + q=f"{H}.self_attn.q_proj.weight", + k=f"{H}.self_attn.k_proj.weight", + v=f"{H}.self_attn.v_proj.weight", + ) + qkv._tp_mapping = ColumnParallelMapping(qkv.megatron_param, qkv.megatron_param) + return MegatronMappingRegistry( + *(AutoMapping(m, h) for m, h in auto.items()), + *(ReplicatedMapping(m, h) for m, h in repl.items()), + qkv, + ColumnParallelMapping(f"{L}.mlp.linear_fc1.weight", f"{H}.mlp.up_proj.weight"), + ) diff --git a/src/megatron/bridge/models/apertus1p5/__init__.py b/src/megatron/bridge/models/apertus1p5/__init__.py new file mode 100644 index 00000000000..9d58fdd1e5b --- /dev/null +++ b/src/megatron/bridge/models/apertus1p5/__init__.py @@ -0,0 +1,3 @@ +from megatron.bridge.models.apertus1p5.apertus1p5_bridge import Apertus1p5Bridge + +__all__ = ["Apertus1p5Bridge"] diff --git a/src/megatron/bridge/models/apertus1p5/apertus1p5_bridge.py b/src/megatron/bridge/models/apertus1p5/apertus1p5_bridge.py new file mode 100644 index 00000000000..9cb764d926f --- /dev/null +++ b/src/megatron/bridge/models/apertus1p5/apertus1p5_bridge.py @@ -0,0 +1,311 @@ +# Megatron-Bridge bridge for Apertus 1.5 (`Apertus1p5ForConditionalGeneration`). +# +# wqwqazwsxedc/Megatron-Bridge (apertus branch) only ever added a bridge for the +# older, text-only `ApertusForCausalLM` (see apertus/apertus_bridge.py) -- there is +# no apertus1p5 bridge anywhere in that fork (confirmed against its git history: +# the entire branch is stock NVIDIA r0.5.0 plus one commit, "add apertus", which +# touches only apertus/{__init__.py,apertus_bridge.py} and a functional test). +# AutoBridge.from_hf_pretrained() therefore raises "Model architecture +# 'Apertus1p5ForConditionalGeneration' is not yet supported" (run 3149776) -- this +# is a real, total gap, not a registration bug fixable with a one-line patch. +# +# Feasibility: swiss-ai/transformers' Apertus1p5TextModel (the actual language +# backbone) is architecturally identical to plain Apertus -- same xielu MLP +# (up_proj/down_proj only, no gate), same q_norm/k_norm RMSNorm on attention, +# same attention_bias=False, same llama3-style rope_parameters dict shape +# (rope_type/rope_theta/factor/original_max_position_embeddings/low_freq_factor/ +# high_freq_factor) -- confirmed by reading modeling_apertus1p5.py at the exact +# commit this script pins (SWISS_AI_TRANSFORMERS_SHA). It only differs in two +# structural ways: +# 1. Hyperparameters live under `config.text_config` (an Apertus1p5TextConfig), +# not the top-level Apertus1p5Config -- the top level also carries +# vision_config/audio_config for the multimodal tokenizers this benchmark +# never uses. +# 2. The decoder stack sits one level deeper in the HF module tree: +# Apertus1p5ForConditionalGeneration.model.language_model.{embed_tokens, +# layers,norm} instead of plain Apertus's Apertus...ForCausalLM.model.*. +# `lm_head` itself stays at the top level, tied to +# model.language_model.embed_tokens.weight (per _tied_weights_keys) -- +# exactly analogous to plain Apertus's own lm_head/output_layer tie. +# +# This bridge is therefore adapted directly from apertus/apertus_bridge.py: same +# MCoreXIELU activation and get_apertus_decoder_block_spec (imported, not +# duplicated, so behavior stays identical to the one upstream commit that +# actually exists), same rope/attention_bias/hidden_act validation, only the HF +# config source (text_config instead of the top-level config) and the HF-side +# weight key prefix (model.language_model.* instead of model.*) change. +# +# Deliberately NOT mapped: vision_tower/audio_tower/vision_tokenizer/ +# audio_tokenizer weights. Megatron's target here is plain GPTModel (text-only -- +# see "want to take advantage of Megatron parallelism" in the dispatching +# conversation, not multimodal fidelity), which has no parameter slots for them +# anyway, and this benchmark (text-only GSM8K GRPO) never calls +# get_image_feature/get_audio_feature -- same reasoning already established for +# the SGLang side of this script (see apertus-benchmarks/patches/ +# sglang-apertus1p5-local-fixes.patch point 6 in CLAUDE.md: those weights never +# need to be numerically correct for this benchmark, only absent-is-fine). No +# strict/leftover-key check was found anywhere in model_bridge.py's weight-load +# path, so simply omitting these mappings is sufficient -- there is nothing to +# suppress. +# +# Pruned LM head: confirmed (run 3171151) that swiss-ai/Apertus-v1.5-70B does +# use a pruned head -- output_vocab_size=131072 vs the extended +# vocab_size=266752 that also covers the visual/audio token ranges. An +# earlier version of this bridge raised rather than risk mismapping shapes; +# now handled properly since this benchmark (text-only GSM8K GRPO) never +# produces or consumes a token id outside [0, output_vocab_size) anyway -- +# `output_vocab_size`'s own docstring in configuration_apertus1p5.py confirms +# ids `0..output_vocab_size - 1` are exactly the retained/generatable +# (non-multimodal) ones. Megatron-core's GPTModelProvider only has one +# `vocab_size` sizing both tables, so provider_bridge sets it to +# `output_vocab_size` (not the full extended vocab_size) whenever the two +# differ, and mapping_registry truncates the *input* embedding table +# (`model.language_model.embed_tokens.weight`, shape (266752, hidden) in the +# checkpoint) down to its first `output_vocab_size` rows to match -- the +# checkpoint's `lm_head.weight` is already exactly (131072, hidden) since it +# really is a physically pruned Linear layer (see +# Apertus1p5ForConditionalGeneration.__init__), so the output-side mapping +# needs no such transform. `_TruncatedVocabEmbeddingMapping` below also +# best-effort zero-pads back out to the full vocab_size on the reverse +# (megatron_to_hf) direction, for whatever HF-format-export path might call +# it -- those padded rows carry no real multimodal embedding, same caveat as +# megatron_to_hf_config below. +# +# megatron_to_hf_config below (used for full HF-format checkpoint export, not +# for the live NCCL weight sync this benchmark actually exercises during +# training) is a best-effort adaptation nesting the generic flat CONFIG_MAPPING +# output under "text_config" to match Apertus1p5Config's real shape -- lower +# confidence than the forward (provider_bridge/mapping_registry) path, since it +# is not needed to get past this benchmark's first training step and so was not +# a priority to verify. + +from __future__ import annotations + +import torch +from megatron.bridge.models.apertus.apertus_bridge import MCoreXIELU, get_apertus_decoder_block_spec +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + ColumnParallelMapping, + QKVMapping, + ReplicatedMapping, +) +from megatron.bridge.models.conversion.utils import unwrap_model +from megatron.core.models.gpt.gpt_model import GPTModel +from transformers import Apertus1p5ForConditionalGeneration + + +class _TruncatedVocabEmbeddingMapping(AutoMapping): + """AutoMapping variant for a pruned-LM-head checkpoint's input embedding table. + + HF's `embed_tokens.weight` covers the full extended vocab_size (text plus + multimodal token ids); Megatron-core's single `vocab_size` here is set to + the narrower `output_vocab_size` (see provider_bridge and the module + docstring's "Pruned LM head" section) so it matches the checkpoint's + already-pruned `lm_head.weight`. hf_to_megatron truncates the source + tensor to the first output_vocab_size rows before handing off to the + normal VocabParallelEmbedding sharding logic; megatron_to_hf best-effort + zero-pads back out to the full width (those rows carry no real multimodal + embedding -- acceptable since this bridge's HF-export path is already + lower-priority, see megatron_to_hf_config below). + """ + + def __init__(self, megatron_param, hf_param, output_vocab_size, full_vocab_size, permute_dims=None): + super().__init__(megatron_param, hf_param, permute_dims) + self._output_vocab_size = output_vocab_size + self._full_vocab_size = full_vocab_size + + def hf_to_megatron(self, hf_weights, megatron_module): + if hf_weights.shape[0] > self._output_vocab_size: + hf_weights = hf_weights[: self._output_vocab_size].contiguous() + return super().hf_to_megatron(hf_weights, megatron_module) + + def megatron_to_hf(self, megatron_weights, megatron_module): + result = super().megatron_to_hf(megatron_weights, megatron_module) + if not result: + return result + key = next(iter(result)) + value = result[key] + if value.shape[0] < self._full_vocab_size: + pad = value.new_zeros((self._full_vocab_size - value.shape[0], *value.shape[1:])) + value = torch.cat([value, pad], dim=0) + return {key: value} + + def resolve(self, captures): + resolved_megatron_param, resolved_hf_param = self._resolve_names(captures) + return type(self)( + resolved_megatron_param, + resolved_hf_param, + self._output_vocab_size, + self._full_vocab_size, + self.permute_dims, + ) + +_ROPE_DEFAULTS = { + "rope_type": "llama3", + "original_max_position_embeddings": 8192, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, +} + + +class _TextConfigOnlyShim: + """Minimal stand-in for an hf_pretrained wrapper, exposing only `.config`. + + MegatronModelBridge.provider_bridge() (the base implementation we delegate + to below) only ever reads `hf_pretrained.config` -- this lets us hand it + `Apertus1p5Config.text_config` directly instead of the top-level config + that also carries the unrelated vision_config/audio_config blocks. + """ + + def __init__(self, config): + self.config = config + + +@MegatronModelBridge.register_bridge(source=Apertus1p5ForConditionalGeneration, target=GPTModel, model_type="apertus1p5") +class Apertus1p5Bridge(MegatronModelBridge): + @classmethod + def hf_to_megatron_activation(cls, hidden_act: str): + if hidden_act != "xielu": + return super().hf_to_megatron_activation(hidden_act) + return lambda _: (_ for _ in ()).throw(RuntimeError("expected MCoreXIELU")) + + def provider_bridge(self, hf_pretrained): + text_config = hf_pretrained.config.text_config + provider = super().provider_bridge(_TextConfigOnlyShim(text_config)) + + output_vocab_size = getattr(text_config, "output_vocab_size", None) + if output_vocab_size is not None and output_vocab_size != text_config.vocab_size: + # Pruned LM head: size Megatron's single vocab_size to the narrower, + # physically-real output_vocab_size (matching the checkpoint's actual + # lm_head.weight shape) rather than the full extended vocab_size -- + # see the module docstring's "Pruned LM head" section. mapping_registry + # truncates the embedding table to match. + provider.vocab_size = output_vocab_size + provider.make_vocab_size_divisible_by = self.make_vocab_size_divisible_by(output_vocab_size) + + rope = { + **(getattr(text_config, "rope_scaling", None) or {}), + **(getattr(text_config, "rope_parameters", None) or {}), + } + rope_type = rope.get("rope_type", rope.get("type", "llama3")) + factor = float(rope.get("factor", 1.0)) + theta = float(rope.get("rope_theta", getattr(text_config, "rope_theta", 10000.0))) + if text_config.hidden_act != "xielu": + raise ValueError(f"Expected hidden_act='xielu', got {text_config.hidden_act!r}") + if text_config.attention_bias: + raise ValueError("Apertus1p5 attention_bias=True is unsupported") + if rope_type != "llama3": + raise ValueError(f"Unsupported Apertus1p5 RoPE type: {rope_type!r}") + + provider.apertus_rope_scaling = { + "rope_type": rope_type, + "type": rope_type, + "factor": factor, + "original_max_position_embeddings": int(rope.get("original_max_position_embeddings", 8192)), + "low_freq_factor": float(rope.get("low_freq_factor", 1.0)), + "high_freq_factor": float(rope.get("high_freq_factor", 4.0)), + } + provider.normalization = "RMSNorm" + provider.qk_layernorm = True + provider.gated_linear_unit = False + provider.use_te_activation_func = False + provider.bias_activation_fusion = False + provider.add_bias_linear = False + provider.add_qkv_bias = False + provider.hidden_dropout = 0.0 + provider.rotary_interleaved = False + provider.position_embedding_type = "rope" + provider.rotary_base = theta + provider.rope_scaling = True + provider.rope_scaling_factor = factor + provider.transformer_layer_spec = get_apertus_decoder_block_spec + return provider + + def load_weights_hf_to_megatron(self, hf_pretrained, megatron_model, allowed_mismatched_params=None): + models = super().load_weights_hf_to_megatron( + hf_pretrained, megatron_model, allowed_mismatched_params=allowed_mismatched_params + ) + [ + m._sync_runtime_scalars() + for model in unwrap_model(models) + for m in model.modules() + if isinstance(m, MCoreXIELU) + ] + return models + + @classmethod + def megatron_to_hf_config(cls, provider) -> dict: + text_config = super().megatron_to_hf_config(provider) + theta = float(provider.rotary_base) + rope = { + **(getattr(provider, "apertus_rope_scaling", None) or _ROPE_DEFAULTS), + "factor": float(provider.rope_scaling_factor), + } + rope["rope_type"] = rope["type"] = rope.get("rope_type", rope.get("type", "llama3")) + text_config.update( + hidden_act="xielu", + attention_bias=False, + rope_theta=theta, + rope_scaling=rope, + rope_parameters={**rope, "rope_theta": theta}, + ) + return { + "model_type": "apertus1p5", + "architectures": ["Apertus1p5ForConditionalGeneration"], + "text_config": text_config, + } + + def mapping_registry(self) -> MegatronMappingRegistry: + text_config = self.hf_config.text_config + output_vocab_size = getattr(text_config, "output_vocab_size", None) + full_vocab_size = text_config.vocab_size + + L, H = "decoder.layers.*", "model.language_model.layers.*" + if output_vocab_size is not None and output_vocab_size != full_vocab_size: + embedding_mapping = _TruncatedVocabEmbeddingMapping( + "embedding.word_embeddings.weight", + "model.language_model.embed_tokens.weight", + output_vocab_size, + full_vocab_size, + ) + else: + embedding_mapping = AutoMapping( + "embedding.word_embeddings.weight", "model.language_model.embed_tokens.weight" + ) + + auto = { + "output_layer.weight": "lm_head.weight", + "decoder.final_layernorm.weight": "model.language_model.norm.weight", + f"{L}.self_attention.linear_proj.weight": f"{H}.self_attn.o_proj.weight", + f"{L}.self_attention.q_layernorm.weight": f"{H}.self_attn.q_norm.weight", + f"{L}.self_attention.k_layernorm.weight": f"{H}.self_attn.k_norm.weight", + f"{L}.mlp.linear_fc2.weight": f"{H}.mlp.down_proj.weight", + } + # Pre-attn / pre-MLP RMSNorm weights + the learnable xIELU scalars, mapped + # replicated (not TP-sharded). This recipe's use_transformer_engine spec + # produces the fused `linear_qkv.layer_norm_weight` / `linear_fc1.layer_norm_weight` + # names, so those direct mappings are what load. Validated end-to-end in run + # 3247540 (46 steps, healthy GRPO curve) -- and confirmed against fork + # megatron-core 0.18 / megatron-bridge 0.5.0; switching these to AutoMapping + # on stock megatron-bridge 0.6.1 loaded them as ~0 instead (run 3243271). + repl = { + f"{L}.self_attention.linear_qkv.layer_norm_weight": f"{H}.attention_layernorm.weight", + f"{L}.mlp.linear_fc1.layer_norm_weight": f"{H}.feedforward_layernorm.weight", + **{f"{L}.mlp.activation_func.{n}": f"{H}.mlp.act_fn.{n}" for n in ("alpha_p", "alpha_n", "beta", "eps")}, + } + qkv = QKVMapping( + f"{L}.self_attention.linear_qkv.weight", + q=f"{H}.self_attn.q_proj.weight", + k=f"{H}.self_attn.k_proj.weight", + v=f"{H}.self_attn.v_proj.weight", + ) + qkv._tp_mapping = ColumnParallelMapping(qkv.megatron_param, qkv.megatron_param) + return MegatronMappingRegistry( + embedding_mapping, + *(AutoMapping(m, h) for m, h in auto.items()), + *(ReplicatedMapping(m, h) for m, h in repl.items()), + qkv, + ColumnParallelMapping(f"{L}.mlp.linear_fc1.weight", f"{H}.mlp.up_proj.weight"), + ) diff --git a/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py b/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py index 5eb61869758..83be15cec27 100644 --- a/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py +++ b/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py @@ -23,6 +23,7 @@ import hashlib import os import time +import contextlib from pathlib import Path from typing import Union @@ -97,7 +98,7 @@ def safe_load_config_with_retry( lock_file.parent.mkdir(parents=True, exist_ok=True) - with filelock.FileLock(str(lock_file) + ".lock", timeout=60): + with contextlib.nullcontext(): # filelock disabled: config is written once by rank 0, read-only after; fcntl.flock is unsupported in-container on CSCS Lustre return AutoConfig.from_pretrained(path, trust_remote_code=trust_remote_code, **kwargs) except Exception as e: