From 6ca1400e87a14c48b1d22699181cc42877321b6a Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 13 Aug 2026 03:13:41 -0700 Subject: [PATCH 01/35] feat(recipes): modernize DeepSeek V4 library configs Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 48 ++- scripts/training/run_recipe.py | 12 +- .../references/recipe-index.md | 2 + .../deepseek/gb200/deepseek_v4.py | 25 +- .../deepseek/gb300/deepseek_v4.py | 25 +- .../bridge/recipes/deepseek/__init__.py | 4 + .../bridge/recipes/deepseek/gb200/__init__.py | 2 + .../recipes/deepseek/gb200/deepseek_v4.py | 67 +++- .../bridge/recipes/deepseek/gb300/__init__.py | 2 + .../recipes/deepseek/gb300/deepseek_v4.py | 58 +++- .../test_deepseek_v4_hardware_recipes.py | 300 ++++++++++++++++++ .../scripts/training/test_run_recipe.py | 11 +- 12 files changed, 482 insertions(+), 74 deletions(-) create mode 100644 tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index bcc1e67525..f7cebe360d 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -6,7 +6,13 @@ The bridge supports four published variants out of the same code path. The on-di ## MCore Checkout -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. +DeepSeek V4 pretraining requires Megatron-LM `dev`; the Megatron-LM copy +pinned by the current Megatron Bridge `main` branch is not supported. Stateless +grouped-FP8 checkpoint saves additionally require Megatron-LM PR +[#6509](https://github.com/NVIDIA/Megatron-LM/pull/6509) until that fix merges, +including saves from the existing MXFP8 recipes. Earlier pretraining +verification used Megatron-LM `dev` commit `35f36c7c9dba` plus PR +[#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839). 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 @@ -38,9 +44,11 @@ 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 -[#4824](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4824): +The full-scale `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config` +library recipe and its corresponding performance recipe require the stack +validated by Megatron Bridge PR +[#4824](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4824). The performance +definition preserves its exact benchmark stack: `nvcr.io/nvidia/nemo:26.06.01` with Megatron-LM dev commit `9d46c924dce3818f2b5f894f7380712c780d1801` and the capability-check patch documented in that PR. The Megatron-LM commit pinned by the current @@ -58,22 +66,44 @@ performance features, so it is not a supported runtime for that recipe. - `conversion.sh` imports HF weights into Megatron Bridge and exports Megatron checkpoints back to HF format. - `inference.sh` runs text generation against an HF or Megatron checkpoint. -- `slurm_pretrain.sh` runs the DeepSeek-V4-Flash pretraining recipes. +- `slurm_pretrain.sh` runs the legacy 32-GPU 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. 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 -See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the Slurm launcher and [`deepseek_v4.py`](../../../src/megatron/bridge/recipes/deepseek/deepseek_v4.py) for recipe definitions. +See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py) and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 -- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config`: 256-GPU GB300 - performance configuration (requires the PR #4824 container and dev-MCore - stack described above) +- `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 + Adam MXFP8 with PP2/EP64/HybridEP +- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config`: 256-GPU GB300 + Adam MXFP8 with PP4/VPP4/EP64/HybridEP (requires the PR #4824 container and + dev-MCore stack described above) + +The `_library_config` suffix keeps these real-training definitions distinct +from the corresponding benchmark recipe families. They retain natural routing, +convergence batch sizes, correctness checks, and checkpoint behavior. + +The hardware-count-specific recipes are intentionally not accepted by +`slurm_pretrain.sh`: that legacy launcher overwrites the recipe's batch and +TP/PP/EP/CP settings. Launch them through a site multi-node runner that invokes +`scripts/training/run_recipe.py --recipe --mode pretrain` with +128 ranks for GB200 Flash or 256 ranks for GB300 Pro. Keep the topology and +global batch size owned by the recipe; pass only dataset, run-length, logging, +and checkpoint overrides. + +The GB200 Flash and GB300 Pro library bases enable their supported DSA and +Transformer Engine fused grouped-MLP paths. The new hardware-count-specific +variants additionally enable GLU interleaving; the existing recipe identities +retain their checkpoint tensor layout. The library recipes retain unlimited +natural-routing capacity. All variants retain their precision-specific +training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer +reuse, and reduced-precision gradient reduction remain disabled. `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index e3bd313e2a..654c7f0ae5 100755 --- a/scripts/training/run_recipe.py +++ b/scripts/training/run_recipe.py @@ -106,7 +106,6 @@ from recipe_runner import ( # noqa: E402 apply_cli_overrides, apply_determinism, - apply_runtime_environment, bootstrap_recipe_environment, load_forward_step, load_recipe, @@ -496,17 +495,16 @@ def main(argv: list[str] | None = None) -> None: world_size=benchmark_world_size, ) configuration_mode = _train_mode(args.mode) + recipe = bootstrap_recipe_environment( + recipe, + script_path=str(Path(__file__).resolve()), + argv=list(argv) if argv is not None else sys.argv[1:], + ) if benchmark_metadata is not None: - recipe = bootstrap_recipe_environment( - recipe, - script_path=str(Path(__file__).resolve()), - argv=list(argv) if argv is not None else sys.argv[1:], - ) execution_mode = "pretrain" step_mode = benchmark_metadata.task else: - recipe = apply_runtime_environment(recipe) execution_mode = configuration_mode step_mode = configuration_mode diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 5b58640d63..35a50a9b9e 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,6 +74,8 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 | | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | +| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 2 | 64 | 128 (GB200; unlimited capacity) | +| `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config` | Pretrain | 1 | 4 | 64 | 256 (GB300; unlimited capacity) | ### GLM-4.5 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 e7379dbb99..54af62b017 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -18,7 +18,7 @@ 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, + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as deepseek_v4_flash_128gpu_library_config, ) from megatron.bridge.training.config import ConfigContainer from megatron.bridge.utils.cuda_graph import set_full_iteration_cuda_graph @@ -26,35 +26,20 @@ 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 = deepseek_v4_flash_128gpu_library_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.train.manual_gc_interval = 10 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 @@ -99,8 +84,6 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: 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 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 6c5d18c831..8998333555 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py @@ -21,7 +21,7 @@ ) 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, + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config as deepseek_v4_pro_256gpu_library_config, ) from megatron.bridge.training.config import ConfigContainer @@ -52,24 +52,10 @@ def deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() -> ConfigContainer: 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() + cfg = deepseek_v4_pro_256gpu_library_config() - cfg.model.tensor_model_parallel_size = 1 - cfg.model.pipeline_model_parallel_size = 4 - cfg.model.virtual_pipeline_model_parallel_size = 4 - 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.train.global_batch_size = 4096 - cfg.train.micro_batch_size = 1 - - cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_token_dispatcher_type = "flex" - cfg.model.moe_shared_expert_overlap = False - cfg.model.pipeline_model_parallel_layout = "Et*4|(tttt|)*14tmL" - cfg.model.recompute_granularity = "selective" - cfg.model.recompute_modules = ["mla_up_proj", "mhc"] + cfg.train.manual_gc_interval = 100 _benchmark_common(cfg, cross_entropy_impl="native") @@ -104,11 +90,6 @@ def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.ddp.grad_reduce_in_fp32 = False - cfg.model.fine_grained_activation_offloading = True - cfg.model.offload_modules = ["core_attn", "attn_proj"] - cfg.model.fine_grained_offloading_max_inflight_offloads = 2 - cfg.comm_overlap.overlap_grad_reduce = True - cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 1cf2ac77dd..b9dcb8c0bf 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -50,10 +50,12 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, ) from .gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config, deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config, ) @@ -80,7 +82,9 @@ "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", + "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", + "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config", "set_deepseek_v4_pipeline_model_parallel_layout", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py index 696f6cae8b..7e93b92af4 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -17,10 +17,12 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, ) __all__ = [ + "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index aee81537eb..6341cb6a3e 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -71,7 +71,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: set_deepseek_v4_pipeline_model_parallel_layout(cfg.model) cfg.model.transformer_impl = "transformer_engine" - cfg.model.attention_backend = None + cfg.model.attention_backend = "auto" cfg.model.apply_dsa_kernel_fusion = True cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = use_fused_mhc @@ -83,17 +83,20 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_hybridep_num_sms = 16 cfg.model.moe_grouped_gemm = True cfg.model.moe_permute_fusion = True + cfg.model.moe_router_fusion = True cfg.model.moe_aux_loss_coeff = 0.0 cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc", "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 = None + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["core_attn", "attn_proj"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 2 cfg.model.cuda_graph_impl = "none" cfg.model.cuda_graph_scope = "full" cfg.model.cuda_graph_warmup_steps = 3 @@ -127,6 +130,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + cfg.comm_overlap.overlap_grad_reduce = True cfg.comm_overlap.delay_wgrad_compute = False cfg.comm_overlap.overlap_moe_expert_parallel_comm = False @@ -134,6 +138,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.ddp.use_megatron_fsdp = False cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + "NVTE_CPU_OFFLOAD_V1": 1, } return cfg @@ -145,10 +150,9 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: """ cfg = deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() - cfg.model.apply_dsa_kernel_fusion = False cfg.model.dsa_indexer_loss_coeff = 0.0 cfg.model.dsa_indexer_use_sparse_loss = False - cfg.model.recompute_modules = ["mla_up_proj"] + cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] opt_cfg, scheduler_cfg = distributed_fused_adam_with_cosine_annealing( lr_warmup_iters=2000, @@ -188,6 +192,57 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: return cfg +def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigContainer: + """Return the real-training DeepSeek V4 Flash config for 128 GB200 GPUs. + + This variant uses one PP stage per 64-rank NVLink domain while preserving + the 64-GPU library recipe's optimizer, routing, loss, precision, validation, + and checkpoint contracts. + """ + cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 2 + 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 + set_deepseek_v4_pipeline_model_parallel_layout(cfg.model) + cfg.train.global_batch_size = 256 + cfg.train.micro_batch_size = 1 + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + 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_mlp_glu_interleave_size = 32 + cfg.model.use_transformer_engine_op_fuser = True + cfg.model.recompute_modules = ["mhc", "mla_up_proj"] + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = [] + cfg.model.fine_grained_offloading_max_inflight_offloads = None + cfg.ddp.average_in_collective = False + + cfg.env_vars = { + **cfg.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, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + } + cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) + return cfg + + def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config() -> ConfigContainer: """DeepSeek-V4-Flash BF16 Muon pre-training on 64-GPU GB200 (PP=8, EP=8).""" cfg = deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() diff --git a/src/megatron/bridge/recipes/deepseek/gb300/__init__.py b/src/megatron/bridge/recipes/deepseek/gb300/__init__.py index e4b9a7e485..db38cda838 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/__init__.py @@ -17,10 +17,12 @@ from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config, deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config, ) __all__ = [ + "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py index db0dcba42d..ea263b2370 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py @@ -63,7 +63,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.transformer_impl = "transformer_engine" cfg.model.attention_backend = None - cfg.model.apply_dsa_kernel_fusion = False + cfg.model.apply_dsa_kernel_fusion = True cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = deepseek_v4_supports_blackwell_fused_kernels() cfg.model.dsa_indexer_loss_coeff = 0.0 @@ -72,15 +72,17 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.moe_token_dispatcher_type = "alltoall" cfg.model.moe_aux_loss_coeff = 0.0 cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc"] cfg.model.recompute_method = None cfg.model.recompute_num_layers = None - cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = None + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["core_attn", "attn_proj"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 2 cfg.model.cuda_graph_impl = "none" cfg.model.cuda_graph_scope = "full" cfg.model.cuda_graph_warmup_steps = 3 @@ -114,6 +116,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + cfg.comm_overlap.overlap_grad_reduce = True cfg.comm_overlap.delay_wgrad_compute = False cfg.comm_overlap.overlap_moe_expert_parallel_comm = False @@ -121,6 +124,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.ddp.use_megatron_fsdp = False cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + "NVTE_CPU_OFFLOAD_V1": 1, } return cfg @@ -141,7 +145,6 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.train.train_iters = 1_000_000 cfg.train.global_batch_size = 128 cfg.train.micro_batch_size = 1 - cfg.model.apply_dsa_kernel_fusion = False cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = deepseek_v4_supports_blackwell_fused_kernels() cfg.model.dsa_indexer_loss_coeff = 0.0 @@ -186,7 +189,50 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.model.moe_router_padding_for_fp8 = True cfg.model.mtp_eval_in_bf16 = True cfg.model.quant_recipe = _deepseek_v4_mxfp8_quant_recipe() + return cfg + + +def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config() -> ConfigContainer: + """Return the real-training DeepSeek V4 Pro config for 256 GB300 GPUs. + + This variant adopts the measured PP4/VPP4/EP64 execution topology while + preserving the 32-GPU library recipe's optimizer, routing, loss, precision, + validation, and checkpoint contracts. + """ + cfg = deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.virtual_pipeline_model_parallel_size = 4 + 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 = "Et*4|(tttt|)*14tmL" + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 1 + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = 32 + cfg.model.moe_mlp_glu_interleave_size = 32 + cfg.model.use_transformer_engine_op_fuser = True + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["mla_up_proj", "mhc"] cfg.env_vars = { - **COMMON_RECIPE_ENV_VARS, + **cfg.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_CPU_OFFLOAD_V1": 1, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, } return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py new file mode 100644 index 0000000000..ba4f6693f0 --- /dev/null +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -0,0 +1,300 @@ +# 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. + +"""Execution-parity and convergence guardrails for DeepSeek V4 hardware recipes.""" + +from collections.abc import Callable + +import pytest +import torch + +import megatron.bridge.recipes as recipes +from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config as flash_perf_config, +) +from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config as pro_perf_config, +) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, +) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config as flash_fp8_config, +) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as flash_library_config, +) +from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config as pro_bf16_base_config, +) +from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config as pro_fp8_config, +) +from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( + deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config as pro_library_config, +) +from megatron.bridge.training.config import ConfigContainer +from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies + + +pytestmark = pytest.mark.unit + +_SHARED_MODEL_EXECUTION_FIELDS = ( + "tensor_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "moe_token_dispatcher_type", + "moe_flex_dispatcher_backend", + "moe_shared_expert_overlap", + "moe_hybridep_num_sms", + "recompute_granularity", + "cross_entropy_fusion_impl", + "moe_pad_experts_for_cuda_graph_inference", + "moe_mlp_glu_interleave_size", + "use_transformer_engine_op_fuser", +) + +_SHARED_ENV_VARS = ( + "CUDA_DEVICE_MAX_CONNECTIONS", + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API", + "NVLINK_DOMAIN_SIZE", + "USE_MNNVL", + "NVTE_BWD_LAYERNORM_SM_MARGIN", + "NVTE_FWD_LAYERNORM_SM_MARGIN", + "NVTE_NORM_BWD_USE_CUDNN", + "NVTE_NORM_FWD_USE_CUDNN", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", +) + + +@pytest.fixture(autouse=True) +def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: + patch_recipe_construction_dependencies(monkeypatch) + + +@pytest.mark.parametrize( + ("library_factory", "perf_factory"), + [ + (flash_library_config, flash_perf_config), + (pro_library_config, pro_perf_config), + ], + ids=["flash-gb200", "pro-gb300"], +) +def test_high_scale_library_recipes_share_approved_execution_fields( + library_factory: Callable[[], ConfigContainer], + perf_factory: Callable[[], ConfigContainer], +) -> None: + library_cfg = library_factory() + perf_cfg = perf_factory() + + for field_name in _SHARED_MODEL_EXECUTION_FIELDS: + library_value = getattr(library_cfg.model, field_name, None) + perf_value = getattr(perf_cfg.model, field_name, None) + assert library_value == perf_value, field_name + + assert library_cfg.train.micro_batch_size == perf_cfg.train.micro_batch_size == 1 + assert library_cfg.comm_overlap.overlap_grad_reduce is perf_cfg.comm_overlap.overlap_grad_reduce is True + assert library_cfg.ddp.average_in_collective == perf_cfg.ddp.average_in_collective + for variable_name in _SHARED_ENV_VARS: + assert library_cfg.env_vars[variable_name] == perf_cfg.env_vars[variable_name], variable_name + + +@pytest.mark.parametrize( + "base_factory", + [flash_bf16_base_config, pro_bf16_base_config], + ids=["flash-bf16", "pro-bf16"], +) +def test_base_recipes_enable_precision_independent_fusions( + base_factory: Callable[[], ConfigContainer], +) -> None: + cfg = base_factory() + + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True + assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None + assert cfg.model.use_transformer_engine_op_fuser is False + assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.comm_overlap.overlap_grad_reduce is True + assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars + + assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None + assert getattr(cfg.model, "moe_paged_stash", False) is False + assert cfg.model.cuda_graph_impl == "none" + assert cfg.rerun_state_machine.check_for_nan_in_loss is True + assert cfg.ddp.check_for_nan_in_grad is True + assert not getattr(cfg.model, "fp8", False) + assert getattr(cfg.model, "quant_recipe", None) is None + assert cfg.optimizer.optimizer_offload_fraction == 0.0 + + +def test_flash_base_recipe_ports_flash_fusions() -> None: + cfg = flash_bf16_base_config() + + assert cfg.model.attention_backend == "auto" + assert cfg.model.moe_router_fusion is True + assert cfg.train.manual_gc_interval == 5 + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + + +def test_pro_base_recipe_ports_activation_offload() -> None: + cfg = pro_bf16_base_config() + + assert cfg.model.attention_backend is None + assert cfg.train.manual_gc_interval == 5 + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + + +@pytest.mark.parametrize( + "fp8_factory", + [flash_fp8_config, pro_fp8_config], + ids=["flash-mxfp8", "pro-mxfp8"], +) +def test_mxfp8_recipes_keep_training_precision_contract( + fp8_factory: Callable[[], ConfigContainer], +) -> None: + cfg = fp8_factory() + + assert cfg.model.quant_recipe is not None + assert cfg.model.moe_router_padding_for_fp8 is True + assert cfg.mixed_precision.fp8_param_gather is False + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + assert cfg.ddp.grad_reduce_in_fp32 is True + assert cfg.optimizer.main_grads_dtype == torch.float32 + + +def test_flash_mxfp8_recipe_uses_activation_offload_to_fit() -> None: + cfg = flash_fp8_config() + + assert cfg.model.recompute_modules == ["moe_act", "mhc", "mla_up_proj"] + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + + +def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: + cfg = flash_library_config() + + assert cfg.train.train_iters == 1_000_000 + assert cfg.train.global_batch_size == 256 + assert cfg.validation.eval_interval == 2000 + assert cfg.validation.eval_iters == 32 + assert cfg.model.moe_router_force_load_balancing is False + assert cfg.model.dsa_indexer_loss_coeff == 0.0 + assert cfg.model.dsa_indexer_use_sparse_loss is False + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.quant_recipe is not None + assert cfg.model.moe_router_padding_for_fp8 is True + assert cfg.mixed_precision.fp8_param_gather is False + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + assert cfg.ddp.grad_reduce_in_fp32 is True + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.ddp.check_for_nan_in_grad is True + assert cfg.model.cuda_graph_impl == "none" + assert cfg.model.use_transformer_engine_op_fuser is True + assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert cfg.model.pipeline_model_parallel_size == 2 + assert len(cfg.model.pipeline_model_parallel_layout) == 2 + assert cfg.model.pipeline_model_parallel_layout[0][0] == "embedding" + assert cfg.model.pipeline_model_parallel_layout[-1][-2:] == ["mtp", "loss"] + assert sum(stage.count("decoder") for stage in cfg.model.pipeline_model_parallel_layout) == cfg.model.num_layers + assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"] + assert cfg.model.fine_grained_activation_offloading is False + assert cfg.model.offload_modules == [] + assert cfg.model.fine_grained_offloading_max_inflight_offloads is None + assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None + assert getattr(cfg.model, "moe_paged_stash", False) is False + assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 + assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 + assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars + + perf_cfg = flash_perf_config() + assert perf_cfg.model.pipeline_model_parallel_size == 1 + assert perf_cfg.model.pipeline_model_parallel_layout is None + assert perf_cfg.train.global_batch_size == 2048 + assert perf_cfg.model.moe_router_force_load_balancing is True + assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 + assert perf_cfg.model.quant_recipe is None + assert perf_cfg.model.recompute_modules == ["mla_up_proj"] + assert perf_cfg.model.fine_grained_activation_offloading is False + assert perf_cfg.model.offload_modules == [] + assert perf_cfg.mixed_precision.fp8_param_gather is True + assert perf_cfg.ddp.grad_reduce_in_fp32 is False + assert perf_cfg.ddp.check_for_nan_in_grad is False + + +def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: + cfg = pro_library_config() + + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.virtual_pipeline_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_layout == "Et*4|(tttt|)*14tmL" + assert cfg.train.train_iters == 1_000_000 + assert cfg.train.global_batch_size == 128 + assert cfg.validation.eval_interval == 2000 + assert cfg.validation.eval_iters == 32 + assert cfg.model.moe_router_force_load_balancing is False + assert cfg.model.dsa_indexer_loss_coeff == 0.0 + assert cfg.model.dsa_indexer_use_sparse_loss is False + assert cfg.model.apply_dsa_kernel_fusion is True + assert cfg.model.quant_recipe is not None + assert cfg.model.moe_router_padding_for_fp8 is True + assert cfg.mixed_precision.fp8_param_gather is False + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + assert cfg.ddp.grad_reduce_in_fp32 is True + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.ddp.check_for_nan_in_grad is True + assert cfg.model.cuda_graph_impl == "none" + assert cfg.model.use_transformer_engine_op_fuser is True + assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None + assert getattr(cfg.model, "moe_paged_stash", False) is False + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 + assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 + + perf_cfg = pro_perf_config() + assert perf_cfg.model.pipeline_model_parallel_size == 4 + assert perf_cfg.model.virtual_pipeline_model_parallel_size == 4 + assert perf_cfg.model.pipeline_model_parallel_layout == "Et*4|(tttt|)*14tmL" + assert perf_cfg.train.global_batch_size == 4096 + assert perf_cfg.model.moe_router_force_load_balancing is True + assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 + assert perf_cfg.model.quant_recipe is None + assert perf_cfg.mixed_precision.fp8_param_gather is True + assert perf_cfg.optimizer.main_grads_dtype == torch.bfloat16 + assert perf_cfg.ddp.grad_reduce_in_fp32 is False + assert perf_cfg.ddp.check_for_nan_in_grad is False + + +def test_high_scale_deepseek_v4_recipes_are_exported() -> None: + assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config + assert recipes.deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config is pro_library_config diff --git a/tests/unit_tests/scripts/training/test_run_recipe.py b/tests/unit_tests/scripts/training/test_run_recipe.py index dc17ee22f2..1f8b540e2e 100644 --- a/tests/unit_tests/scripts/training/test_run_recipe.py +++ b/tests/unit_tests/scripts/training/test_run_recipe.py @@ -475,7 +475,7 @@ def test_benchmark_finetuning_recipes_use_unified_runner(monkeypatch, mode, task assert handles.recipe_runner.run_config.call_args.kwargs["mode"] == "pretrain" -def test_library_only_canonical_name_does_not_enable_benchmark_runtime(): +def test_library_only_canonical_name_bootstraps_recipe_environment(): module, handles = _load_module() handles.recipe_runner.load_recipe.return_value = SimpleNamespace() @@ -488,7 +488,10 @@ def test_library_only_canonical_name_does_not_enable_benchmark_runtime(): ] ) - handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() + handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() + bootstrap_call = handles.recipe_runner.bootstrap_recipe_environment.call_args + assert bootstrap_call.args == (handles.recipe_runner.load_recipe.return_value,) + assert bootstrap_call.kwargs["script_path"].endswith("scripts/training/run_recipe.py") handles.recipe_runner.load_forward_step.assert_called_once_with("llm_step", mode="pretrain") @@ -1444,4 +1447,6 @@ def test_config_container_overrides_are_forwarded_directly(): config, ["train.train_iters=3", "train.global_batch_size=8", "train.micro_batch_size=1"], ) - handles.recipe_runner.apply_runtime_environment.assert_called_once_with(config) + handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() + bootstrap_call = handles.recipe_runner.bootstrap_recipe_environment.call_args + assert bootstrap_call.args == (config,) From 76a84bbf0775fad0f8975599465701d5975795cb Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 14:00:56 -0700 Subject: [PATCH 02/35] feat(recipes): record DeepSeek V4 PP1 memory settings Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 6 +++-- .../references/recipe-index.md | 2 +- .../recipes/deepseek/gb200/deepseek_v4.py | 21 ++++++++++-------- .../test_deepseek_v4_hardware_recipes.py | 22 +++++++++---------- 4 files changed, 27 insertions(+), 24 deletions(-) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index f7cebe360d..236d5f73a4 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -80,7 +80,8 @@ Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 - Adam MXFP8 with PP2/EP64/HybridEP + Adam MXFP8 with PP1/EP64/DP2/HybridEP, selective recompute, attention + activation offload, and MXFP8 parameter gather/buffer reuse - `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config`: 256-GPU GB300 Adam MXFP8 with PP4/VPP4/EP64/HybridEP (requires the PR #4824 container and dev-MCore stack described above) @@ -103,7 +104,8 @@ variants additionally enable GLU interleaving; the existing recipe identities retain their checkpoint tensor layout. The library recipes retain unlimited natural-routing capacity. All variants retain their precision-specific training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer -reuse, and reduced-precision gradient reduction remain disabled. +reuse, and reduced-precision gradient reduction remain disabled except for the +validated FP8 parameter gather/buffer reuse in the high-scale GB200 recipe. `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 35a50a9b9e..a6f4d8dbca 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,7 +74,7 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 | | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | -| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 2 | 64 | 128 (GB200; unlimited capacity) | +| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) | | `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config` | Pretrain | 1 | 4 | 64 | 256 (GB300; unlimited capacity) | ### GLM-4.5 diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 6341cb6a3e..f672c5253b 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -195,14 +195,16 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigContainer: """Return the real-training DeepSeek V4 Flash config for 128 GB200 GPUs. - This variant uses one PP stage per 64-rank NVLink domain while preserving - the 64-GPU library recipe's optimizer, routing, loss, precision, validation, - and checkpoint contracts. + This variant uses PP1 with two 64-rank expert/data-parallel replicas. A + matched PP1 screen completed eight finite steps and reduced peak allocated + memory by 25.51 GiB with attention activation offload plus MXFP8 parameter + gather/storage. Expert capacity, paged stash, and CUDA graphs remain + disabled to preserve natural-routing training semantics. """ cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() cfg.model.tensor_model_parallel_size = 1 - cfg.model.pipeline_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.expert_model_parallel_size = 64 @@ -219,10 +221,12 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_mlp_glu_interleave_size = 32 cfg.model.use_transformer_engine_op_fuser = True - cfg.model.recompute_modules = ["mhc", "mla_up_proj"] - cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = [] - cfg.model.fine_grained_offloading_max_inflight_offloads = None + cfg.model.recompute_modules = ["moe", "mhc", "mla_up_proj", "layernorm"] + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["core_attn", "attn_proj"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 2 + cfg.mixed_precision.fp8_param_gather = True + cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = True cfg.ddp.average_in_collective = False cfg.env_vars = { @@ -239,7 +243,6 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, } - cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index ba4f6693f0..863a6dc346 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -207,8 +207,8 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.model.apply_dsa_kernel_fusion is True assert cfg.model.quant_recipe is not None assert cfg.model.moe_router_padding_for_fp8 is True - assert cfg.mixed_precision.fp8_param_gather is False - assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False + assert cfg.mixed_precision.fp8_param_gather is True + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True assert cfg.mixed_precision.grad_reduce_in_fp32 is True assert cfg.ddp.grad_reduce_in_fp32 is True assert cfg.optimizer.main_grads_dtype == torch.float32 @@ -216,21 +216,18 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.model.cuda_graph_impl == "none" assert cfg.model.use_transformer_engine_op_fuser is True assert cfg.model.moe_mlp_glu_interleave_size == 32 - assert cfg.model.pipeline_model_parallel_size == 2 - assert len(cfg.model.pipeline_model_parallel_layout) == 2 - assert cfg.model.pipeline_model_parallel_layout[0][0] == "embedding" - assert cfg.model.pipeline_model_parallel_layout[-1][-2:] == ["mtp", "loss"] - assert sum(stage.count("decoder") for stage in cfg.model.pipeline_model_parallel_layout) == cfg.model.num_layers - assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"] - assert cfg.model.fine_grained_activation_offloading is False - assert cfg.model.offload_modules == [] - assert cfg.model.fine_grained_offloading_max_inflight_offloads is None + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None assert getattr(cfg.model, "moe_paged_stash", False) is False assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 perf_cfg = flash_perf_config() assert perf_cfg.model.pipeline_model_parallel_size == 1 @@ -243,6 +240,7 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert perf_cfg.model.fine_grained_activation_offloading is False assert perf_cfg.model.offload_modules == [] assert perf_cfg.mixed_precision.fp8_param_gather is True + assert perf_cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True assert perf_cfg.ddp.grad_reduce_in_fp32 is False assert perf_cfg.ddp.check_for_nan_in_grad is False From 9f1fbe9b7a71cdd1ecf00c19726a4c4686e62f6e Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 14:20:02 -0700 Subject: [PATCH 03/35] fix(recipes): keep DeepSeek V4 perf configs unchanged Signed-off-by: Chen Cui --- .../deepseek/gb200/deepseek_v4.py | 25 ++++- .../deepseek/gb300/deepseek_v4.py | 25 ++++- .../test_deepseek_v4_hardware_recipes.py | 94 +------------------ 3 files changed, 44 insertions(+), 100 deletions(-) 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 54af62b017..e7379dbb99 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -18,7 +18,7 @@ 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_128gpu_gb200_fp8mx_library_config as deepseek_v4_flash_128gpu_library_config, + 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 @@ -26,20 +26,35 @@ 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_128gpu_library_config() + 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.manual_gc_interval = 10 + 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_modules = ["mla_up_proj"] + 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 @@ -84,6 +99,8 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: 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 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 8998333555..6c5d18c831 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py @@ -21,7 +21,7 @@ ) from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( - deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config as deepseek_v4_pro_256gpu_library_config, + deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, ) from megatron.bridge.training.config import ConfigContainer @@ -52,10 +52,24 @@ def deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() -> ConfigContainer: 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_256gpu_library_config() + cfg = deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.virtual_pipeline_model_parallel_size = 4 + 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.train.global_batch_size = 4096 - cfg.train.manual_gc_interval = 100 + cfg.train.micro_batch_size = 1 + + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_shared_expert_overlap = False + cfg.model.pipeline_model_parallel_layout = "Et*4|(tttt|)*14tmL" + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["mla_up_proj", "mhc"] _benchmark_common(cfg, cross_entropy_impl="native") @@ -90,6 +104,11 @@ def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.ddp.grad_reduce_in_fp32 = False + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["core_attn", "attn_proj"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 2 + cfg.comm_overlap.overlap_grad_reduce = True + cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 863a6dc346..a3aa82022b 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Execution-parity and convergence guardrails for DeepSeek V4 hardware recipes.""" +"""Execution and convergence guardrails for DeepSeek V4 hardware recipes.""" from collections.abc import Callable @@ -20,12 +20,6 @@ import torch import megatron.bridge.recipes as recipes -from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config as flash_perf_config, -) -from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v4 import ( - deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config as pro_perf_config, -) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, ) @@ -50,70 +44,12 @@ pytestmark = pytest.mark.unit -_SHARED_MODEL_EXECUTION_FIELDS = ( - "tensor_model_parallel_size", - "virtual_pipeline_model_parallel_size", - "context_parallel_size", - "expert_model_parallel_size", - "expert_tensor_parallel_size", - "sequence_parallel", - "moe_token_dispatcher_type", - "moe_flex_dispatcher_backend", - "moe_shared_expert_overlap", - "moe_hybridep_num_sms", - "recompute_granularity", - "cross_entropy_fusion_impl", - "moe_pad_experts_for_cuda_graph_inference", - "moe_mlp_glu_interleave_size", - "use_transformer_engine_op_fuser", -) - -_SHARED_ENV_VARS = ( - "CUDA_DEVICE_MAX_CONNECTIONS", - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API", - "NVLINK_DOMAIN_SIZE", - "USE_MNNVL", - "NVTE_BWD_LAYERNORM_SM_MARGIN", - "NVTE_FWD_LAYERNORM_SM_MARGIN", - "NVTE_NORM_BWD_USE_CUDNN", - "NVTE_NORM_FWD_USE_CUDNN", - "NVTE_ALLOW_NONDETERMINISTIC_ALGO", -) - @pytest.fixture(autouse=True) def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: patch_recipe_construction_dependencies(monkeypatch) -@pytest.mark.parametrize( - ("library_factory", "perf_factory"), - [ - (flash_library_config, flash_perf_config), - (pro_library_config, pro_perf_config), - ], - ids=["flash-gb200", "pro-gb300"], -) -def test_high_scale_library_recipes_share_approved_execution_fields( - library_factory: Callable[[], ConfigContainer], - perf_factory: Callable[[], ConfigContainer], -) -> None: - library_cfg = library_factory() - perf_cfg = perf_factory() - - for field_name in _SHARED_MODEL_EXECUTION_FIELDS: - library_value = getattr(library_cfg.model, field_name, None) - perf_value = getattr(perf_cfg.model, field_name, None) - assert library_value == perf_value, field_name - - assert library_cfg.train.micro_batch_size == perf_cfg.train.micro_batch_size == 1 - assert library_cfg.comm_overlap.overlap_grad_reduce is perf_cfg.comm_overlap.overlap_grad_reduce is True - assert library_cfg.ddp.average_in_collective == perf_cfg.ddp.average_in_collective - for variable_name in _SHARED_ENV_VARS: - assert library_cfg.env_vars[variable_name] == perf_cfg.env_vars[variable_name], variable_name - - @pytest.mark.parametrize( "base_factory", [flash_bf16_base_config, pro_bf16_base_config], @@ -229,21 +165,6 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 - perf_cfg = flash_perf_config() - assert perf_cfg.model.pipeline_model_parallel_size == 1 - assert perf_cfg.model.pipeline_model_parallel_layout is None - assert perf_cfg.train.global_batch_size == 2048 - assert perf_cfg.model.moe_router_force_load_balancing is True - assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 - assert perf_cfg.model.quant_recipe is None - assert perf_cfg.model.recompute_modules == ["mla_up_proj"] - assert perf_cfg.model.fine_grained_activation_offloading is False - assert perf_cfg.model.offload_modules == [] - assert perf_cfg.mixed_precision.fp8_param_gather is True - assert perf_cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True - assert perf_cfg.ddp.grad_reduce_in_fp32 is False - assert perf_cfg.ddp.check_for_nan_in_grad is False - def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: cfg = pro_library_config() @@ -279,19 +200,6 @@ def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - perf_cfg = pro_perf_config() - assert perf_cfg.model.pipeline_model_parallel_size == 4 - assert perf_cfg.model.virtual_pipeline_model_parallel_size == 4 - assert perf_cfg.model.pipeline_model_parallel_layout == "Et*4|(tttt|)*14tmL" - assert perf_cfg.train.global_batch_size == 4096 - assert perf_cfg.model.moe_router_force_load_balancing is True - assert perf_cfg.model.dsa_indexer_loss_coeff == 0.01 - assert perf_cfg.model.quant_recipe is None - assert perf_cfg.mixed_precision.fp8_param_gather is True - assert perf_cfg.optimizer.main_grads_dtype == torch.bfloat16 - assert perf_cfg.ddp.grad_reduce_in_fp32 is False - assert perf_cfg.ddp.check_for_nan_in_grad is False - def test_high_scale_deepseek_v4_recipes_are_exported() -> None: assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config From 1709cac34e515b8c41cba734ee54fafc50fca13f Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 14:30:37 -0700 Subject: [PATCH 04/35] fix(recipes): restore TE cross entropy Signed-off-by: Chen Cui --- src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py | 2 +- src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py | 2 +- tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index f672c5253b..0f634237b8 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -88,7 +88,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.cross_entropy_fusion_impl = "te" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] diff --git a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py index ea263b2370..c55b2fe13f 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py @@ -74,7 +74,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.cross_entropy_fusion_impl = "te" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc"] diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index a3aa82022b..7a9946c450 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -64,7 +64,7 @@ def test_base_recipes_enable_precision_independent_fusions( assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None assert cfg.model.use_transformer_engine_op_fuser is False - assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.model.cross_entropy_fusion_impl == "te" assert cfg.comm_overlap.overlap_grad_reduce is True assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars From 8f7017cad21a328f58932b8d0d3ce07c1a4d241b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 15:29:49 -0700 Subject: [PATCH 05/35] feat(recipes): add DeepSeek V4 GB200 packed SFT Signed-off-by: Chen Cui --- .../bridge/recipes/deepseek/__init__.py | 2 ++ .../bridge/recipes/deepseek/deepseek_v4.py | 3 ++ .../bridge/recipes/deepseek/gb200/__init__.py | 2 ++ .../recipes/deepseek/gb200/deepseek_v4.py | 34 ++++++++++++++++++- .../recipes/deepseek/gb300/deepseek_v4.py | 3 +- .../test_deepseek_v4_hardware_recipes.py | 25 +++++++++++++- 6 files changed, 66 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index b9dcb8c0bf..083772d64b 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -51,6 +51,7 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, + deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config, ) from .gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config, @@ -83,6 +84,7 @@ "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config", diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index 5ad648049b..43a255b8d0 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py @@ -93,6 +93,9 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: 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. + + For GB200-optimized training with HybridEP dispatcher and DSA kernel fusion, + use ``deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config`` instead. """ cfg = deepseek_v4_flash_sft_config() # DSv4 hybrid attention requires contiguous CP partition when CP > 1; diff --git a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py index 7e93b92af4..3e1c7c09d7 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -18,6 +18,7 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, + deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config, ) @@ -26,4 +27,5 @@ "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", + "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 0f634237b8..ce2e8b336b 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -88,7 +88,8 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] @@ -286,3 +287,34 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config() -> ConfigContainer cfg.mixed_precision = bf16_mixed() cfg.mixed_precision.grad_reduce_in_fp32 = True return cfg + + +def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigContainer: + """Return the GB200-optimized offline-packed OpenMath SFT config. + + This variant adds HybridEP dispatch, safe uneven-input padding, DSA indexer + training, and MoE fusions to the hardware-agnostic packed SFT recipe. It + preserves non-grouped expert GEMMs for checkpoint export compatibility. + """ + from megatron.bridge.recipes.deepseek.deepseek_v4 import ( + deepseek_v4_flash_sft_openmath_thinking_packed_config, + ) + + cfg = deepseek_v4_flash_sft_openmath_thinking_packed_config() + + 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.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_hybridep_num_sms = 16 + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = True + cfg.model.moe_grouped_gemm = False + cfg.model.moe_permute_fusion = True + cfg.model.moe_router_fusion = True + + # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. + cfg.model.cross_entropy_fusion_impl = "native" + return cfg diff --git a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py index c55b2fe13f..3301a3ae15 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py @@ -74,7 +74,8 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "te" + # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. + cfg.model.cross_entropy_fusion_impl = "native" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc"] diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 7a9946c450..c606777290 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -29,6 +29,9 @@ from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as flash_library_config, ) +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config as flash_packed_sft_config, +) from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config as pro_bf16_base_config, ) @@ -64,7 +67,7 @@ def test_base_recipes_enable_precision_independent_fusions( assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None assert cfg.model.use_transformer_engine_op_fuser is False - assert cfg.model.cross_entropy_fusion_impl == "te" + assert cfg.model.cross_entropy_fusion_impl == "native" assert cfg.comm_overlap.overlap_grad_reduce is True assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars @@ -130,6 +133,25 @@ def test_flash_mxfp8_recipe_uses_activation_offload_to_fit() -> None: assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 +def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: + cfg = flash_packed_sft_config() + + assert cfg.model.cp_partition_mode == "contiguous" + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 2 + 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.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_hybridep_num_sms == 16 + assert cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is True + assert cfg.model.moe_shared_expert_overlap is False + assert cfg.model.moe_permute_fusion is True + assert cfg.model.moe_router_fusion is True + assert cfg.model.moe_grouped_gemm is False + assert cfg.model.cross_entropy_fusion_impl == "native" + + def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: cfg = flash_library_config() @@ -203,4 +225,5 @@ def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: def test_high_scale_deepseek_v4_recipes_are_exported() -> None: assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config + assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config assert recipes.deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config is pro_library_config From b3d67c91007a84d14ccbbfbb12fa6d16320c4348 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 15:45:04 -0700 Subject: [PATCH 06/35] docs(cards): add DeepSeek V4 verification card Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 373 ++++++++++++++++++ 1 file changed, 373 insertions(+) create mode 100644 examples/model_verification_cards/deepseek-v4-flash/card.yaml diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml new file mode 100644 index 0000000000..1db3f8e8b0 --- /dev/null +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -0,0 +1,373 @@ +# Agent-readable model verification card. +# status: unverified | verified | unsupported | not_applicable + +title: deepseek_v4_flash +summary: > + Performance disclaimer: timing and throughput metrics from functional training items + are sanity checks, not optimized performance results. pretrain_performance.GB200 uses + the tuned canonical GB200 performance recipe. + + DeepSeek-V4-Flash verification covers GPU conversion and GB200 SFT at the item-pinned + Bridge revision. The current 128-GPU GB200 MXFP8 library pretraining recipe uses + PP1/EP64 with natural routing, selective recompute, attention activation offload, and + MXFP8 parameter gather and gradient-buffer reuse. It has completed an eight-step + screen, but full pretraining and direct checkpoint resume remain unverified until the + bounded 100-step runs described below complete. CPU conversion is unsupported because + the 285B model exceeds node CPU RAM. Standard Megatron KV-cache inference is unsupported + for DSV4 hybrid attention; inference uses the HF-native path after export. + + The GB200 packed SFT recipe keeps moe_grouped_gemm=False for export compatibility and + uses the canonical MCore uneven-dispatch padding field. + +verification_index: + model_level: + verified: + - hf_to_megatron_gpu + - megatron_to_hf_gpu + unsupported: + - hf_to_megatron_cpu + - megatron_to_hf_cpu + - inference + unverified: + - manual_forward_pass + training: + GB200: + verified: [sft, sft_long_context, sft_export_inference] + unverified: [pretrain, checkpoint_resume] + unsupported: [peft] + performance: + GB200: verified + +model: + hf_id: deepseek-ai/DeepSeek-V4-Flash + hf_revision: "60d8d70770c6776ff598c94bb586a859a38244f1" # pragma: allowlist secret + architecture: DeepseekV4ForCausalLM + min_transformers_version: "5.12.0" + +verification_environment: + base_container: nvcr.io/nvidia/nemo:26.08.rc4 + bridge_commit: 8f7017cad21a328f58932b8d0d3ce07c1a4d241b # pragma: allowlist secret + +items: + hf_to_megatron_cpu: + status: unsupported + precision: null + command: null + last_verified: null + expected_result: > + CPU import requires approximately 570 GB of CPU RAM for the full BF16 weight + materialisation (285B parameters x 2 bytes). Standard compute nodes do not + provide this headroom; the import consistently exits with OOM before completing. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 4 --ep 4 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/import-gpu + --torch-dtype bfloat16 --trust-remote-code + last_verified: 2026-08-04 + expected_result: > + The command exits successfully and creates iter_0000000. The checkpoint + is reloadable and paired GPU export produces correct HF weights. + + megatron_to_hf_cpu: + status: unsupported + precision: null + command: null + last_verified: null + expected_result: > + CPU export requires the full 570 GB BF16 weight set to reside in CPU RAM during + tensor merging; nodes consistently OOM before the export completes. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 4 --ep 8 --pp 4 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/import-gpu/iter_0000000 + --hf-path work/model-verification/dsv4-flash/export-gpu + --torch-dtype bfloat16 --export-weight-dtype bfloat16 --trust-remote-code + last_verified: 2026-08-05 + expected_result: > + The command exits successfully and writes 46 BF16 safetensors shards covering + all transformer-layer parameters. The export completes 4176/4176 weight + conversions with zero errors. + + manual_forward_pass: + status: unverified + precision: bf16 + command: > + ./scripts/inference/infer.sh --executor slurm --task model-comparison + --nodes 2 --gpus-per-node 4 --ep 4 + --hf_model_path deepseek-ai/DeepSeek-V4-Flash + --megatron_model_path work/model-verification/dsv4-flash/import-gpu/iter_0000000 + --pp 1 --ep 4 + --prompt "The capital of France is the city of" + --hf-revision 60d8d70770c6776ff598c94bb586a859a38244f1 + last_verified: null + expected_result: > + Blocked: loading the HF model and Megatron model simultaneously on 4 GB200 + GPUs (4 x 192 GB = 768 GB) exceeds the available GPU memory for the combined + 285 B parameter model. Requires 8 or more nodes for side-by-side comparison. + + inference: + status: unsupported + precision: null + command: null + last_verified: null + expected_result: > + Standard Megatron KV-cache autoregressive inference is blocked by an assertion + in DSv4HybridAttention.forward (inference_context must be None). DSV4 hybrid + attention with CSA/DSA indexers requires a dedicated inference path with + MLA-specific KV caching. Inference is supported through HF native pipeline + (deepseek-ai/DeepSeek-V4-Flash) after Megatron-to-HF export. + + pretrain: + GB200: + status: unverified + precision: fp8_mx + enabled_features: + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_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/rp2 dataset.num_workers=0 + dataset.random_seed=1234 rng.seed=1234 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + --save_dir work/model-verification/dsv4-flash/pretrain-ref + --save_interval 50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=120 + 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 current library recipe resolves to TP1/PP1/EP64/CP1/DP2 with GBS/MBS 256/1, + HybridEP natural routing, native fused cross entropy, selective recompute over + moe, mhc, mla_up_proj, and layernorm, attention activation offload, and MXFP8 + parameter gather and gradient-buffer reuse. Expert capacity, paged stash, and + CUDA graphs remain disabled. Verification requires 100 finite steps with no + skipped or NaN iterations and complete checkpoints at steps 50 and 100; the + existing eight-step screen is insufficient for verified status. + + sft: + GB200: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 + --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config + --step-func dsv4_step + --pretrained_checkpoint work/models/deepseek-v4-flash-megatron + --save_dir work/model-verification/dsv4-flash/sft-ref + --save_interval 100 --max_steps 100 + model.pipeline_model_parallel_size=4 + 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' + model.context_parallel_size=2 + model.expert_model_parallel_size=8 + model.moe_grouped_gemm=false + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 + dataset.seq_length=1024 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + rng.seed=5678 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + ddp.overlap_grad_reduce=false + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=180 + last_verified: 2026-08-07 + metrics: + initial_loss: 0.94383 + final_loss: 0.28370 + last_10_steps_step_time_ms_avg: 13806.500 + last_10_steps_model_tflops_per_gpu_avg: 13.620 + expected_result: > + On 64 GB200 (16 nodes x 4 GPUs), TP=1/PP=4/EP=8/CP=2/DP=1, GBS/MBS 128/1. + OpenMathInstruct-2 thinking data with offline packing, seq_length=1024. + Full recompute (granularity=full) is required to fit in memory at seq=1024 + with CP=2. HybridEP dispatcher and DSA kernel fusion are enabled via the + GB200-specific recipe. The run completes 100 steps with finite LM and MTP + losses, no skipped or NaN iterations, and a complete step-100 checkpoint. + moe_grouped_gemm=False is preserved for export compatibility. + dist.distributed_timeout_minutes must be 180 or above to allow offline data + packing (approximately 63 minutes) before training begins. + + sft_export_inference: + GB200: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 4 --ep 8 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/sft-ref/iter_0000100 + --hf-path work/model-verification/dsv4-flash/sft-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 --trust-remote-code + --not-strict + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/dsv4-flash/sft-export + --prompt "In one short sentence, explain why Paris is important to France." + --max-new-tokens 32 --chat-template --disable-thinking + last_verified: 2026-08-06 + expected_result: > + The exported BF16 SFT checkpoint reloads successfully as DeepseekV4ForCausalLM. + Two independent greedy runs produce byte-identical token IDs and exactly 32 new + tokens with this literal completion: "Paris is important to France because it is + the country's capital and a major center for politics, culture, and economy. It + is also a symbol of French history" + + sft_long_context: + GB200: + status: verified + precision: bf16 + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 + --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config + --step-func dsv4_step + --pretrained_checkpoint work/models/deepseek-v4-flash-megatron + --save_dir work/model-verification/dsv4-flash/sft-ref + --save_interval 100 --max_steps 100 + model.pipeline_model_parallel_size=4 + 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' + model.context_parallel_size=2 model.cp_partition_mode=contiguous + model.expert_model_parallel_size=8 + model.moe_grouped_gemm=false + model.recompute_granularity=full model.recompute_method=uniform + model.recompute_num_layers=1 + dataset.seq_length=1024 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + rng.seed=5678 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + ddp.overlap_grad_reduce=false + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=180 + last_verified: 2026-08-07 + metrics: + initial_loss: 0.94383 + final_loss: 0.28370 + last_10_steps_step_time_ms_avg: 13806.500 + last_10_steps_model_tflops_per_gpu_avg: 13.620 + expected_result: > + Same run as sft.GB200. CP=2 with contiguous partitioning and offline-packed + OpenMathInstruct-2 at seq_length=1024 demonstrates sequence packing and + context parallelism working together over 100 training steps with finite loss + and no skipped or NaN iterations. + + peft: + all: + status: unsupported + precision: null + command: null + last_verified: null + enabled_features: {} + 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: > + No LoRA / PEFT recipe exists for DeepSeek-V4 in the current recipe set. + + checkpoint_resume: + GB200: + status: unverified + precision: fp8_mx + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_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/rp2 dataset.num_workers=0 + dataset.random_seed=1234 rng.seed=1234 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false + --load_dir work/model-verification/dsv4-flash/pretrain-ref + --save_dir work/model-verification/dsv4-flash/pretrain-resumed + --save_interval 50 checkpoint.ckpt_step=50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=120 + 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: false + expected_result: > + Verification requires restoring optimizer, scheduler, data-order, and RNG state + from the current recipe's step-50 reference checkpoint, then running steps 51-100 + in a separate output root. Step-51 and step-100 losses must match the uninterrupted + reference within the declared 1% relative and 1e-6 absolute tolerances, and all + four metrics must be recorded. + + pretrain_performance: + GB200: + status: verified + precision: fp8_mx + bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config + --max_steps 50 + scheduler.lr_warmup_iters=5 scheduler.lr_decay_iters=50 + validation.eval_interval=0 validation.eval_iters=0 + checkpoint.exit_on_missing_checkpoint=false + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=120 + last_verified: 2026-08-06 + metrics: + initial_loss: 13.58750 + final_loss: 3.62683 + last_10_steps_step_time_ms_avg: 8134.600 + last_10_steps_model_tflops_per_gpu_avg: 731.400 + expected_result: > + On 128 GB200 (32 nodes x 4 GPUs) with contiguous NVLink-domain placement, + PP=1/EP=64/CP=1/TP=1/DP=2, GBS/MBS 2048/1, FP8-MX with HybridEP dispatcher + and full-iteration CUDA graphs. The 50-step run completes with finite losses, + no skipped or NaN iterations, and all four metrics recorded. First iteration + takes approximately 24 minutes for CUDA graph compilation and kernel warmup; + subsequent iterations run at approximately 8.1 seconds per step at steady state. From e459eb45473820d6b164a673d801f1022d17bdce Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 16:00:53 -0700 Subject: [PATCH 07/35] docs(cards): correct DeepSeek V4 evidence labels Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 81 +++++++------------ examples/models/deepseek_v4/README.md | 4 +- 2 files changed, 32 insertions(+), 53 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 1db3f8e8b0..70ff7d940a 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -32,8 +32,8 @@ verification_index: - manual_forward_pass training: GB200: - verified: [sft, sft_long_context, sft_export_inference] - unverified: [pretrain, checkpoint_resume] + verified: [sft, sft_export_inference] + unverified: [pretrain, sft_long_context, checkpoint_resume] unsupported: [peft] performance: GB200: verified @@ -158,13 +158,14 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - The current library recipe resolves to TP1/PP1/EP64/CP1/DP2 with GBS/MBS 256/1, - HybridEP natural routing, native fused cross entropy, selective recompute over - moe, mhc, mla_up_proj, and layernorm, attention activation offload, and MXFP8 - parameter gather and gradient-buffer reuse. Expert capacity, paged stash, and - CUDA graphs remain disabled. Verification requires 100 finite steps with no - skipped or NaN iterations and complete checkpoints at steps 50 and 100; the - existing eight-step screen is insufficient for verified status. + The current library recipe resolves to TP1/PP1/EP64/CP1 with dense DP128, + expert DP2, and GBS/MBS 256/1. It uses HybridEP natural routing, native fused + cross entropy, selective recompute over moe, mhc, mla_up_proj, and layernorm, + attention activation offload, and MXFP8 parameter gather and gradient-buffer + reuse. Expert capacity, paged stash, and CUDA graphs remain disabled. + Verification requires 100 finite steps with no skipped or NaN iterations and + complete checkpoints at steps 50 and 100; the existing eight-step screen is + insufficient for verified status. sft: GB200: @@ -205,8 +206,9 @@ items: last_10_steps_step_time_ms_avg: 13806.500 last_10_steps_model_tflops_per_gpu_avg: 13.620 expected_result: > - On 64 GB200 (16 nodes x 4 GPUs), TP=1/PP=4/EP=8/CP=2/DP=1, GBS/MBS 128/1. - OpenMathInstruct-2 thinking data with offline packing, seq_length=1024. + On 64 GB200 (16 nodes x 4 GPUs), TP1/PP4/EP8/CP2 with dense DP8 and expert + DP1, GBS/MBS 128/1. OpenMathInstruct-2 thinking data uses offline packing + at seq_length=1024. Full recompute (granularity=full) is required to fit in memory at seq=1024 with CP=2. HybridEP dispatcher and DSA kernel fusion are enabled via the GB200-specific recipe. The run completes 100 steps with finite LM and MTP @@ -246,47 +248,23 @@ items: sft_long_context: GB200: - status: verified + status: unverified precision: bf16 - bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret enabled_features: sequence_packing: offline context_parallel_size: 2 - command: > - ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 - --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config - --step-func dsv4_step - --pretrained_checkpoint work/models/deepseek-v4-flash-megatron - --save_dir work/model-verification/dsv4-flash/sft-ref - --save_interval 100 --max_steps 100 - model.pipeline_model_parallel_size=4 - 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' - model.context_parallel_size=2 model.cp_partition_mode=contiguous - model.expert_model_parallel_size=8 - model.moe_grouped_gemm=false - model.recompute_granularity=full model.recompute_method=uniform - model.recompute_num_layers=1 - dataset.seq_length=1024 - scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 - validation.eval_interval=0 validation.eval_iters=0 - rng.seed=5678 - ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true - rerun_state_machine.check_for_nan_in_loss=true - checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null - ddp.overlap_grad_reduce=false - logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - dist.distributed_timeout_minutes=180 - last_verified: 2026-08-07 + command: null + last_verified: null metrics: - initial_loss: 0.94383 - final_loss: 0.28370 - last_10_steps_step_time_ms_avg: 13806.500 - last_10_steps_model_tflops_per_gpu_avg: 13.620 + 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: > - Same run as sft.GB200. CP=2 with contiguous partitioning and offline-packed - OpenMathInstruct-2 at seq_length=1024 demonstrates sequence packing and - context parallelism working together over 100 training steps with finite loss - and no skipped or NaN iterations. + The imported seq_length=1024 SFT run verifies offline packing and CP=2 plumbing, + but it is not a long-context result. Verification requires a separately sized + long-context run with finite loss, no skipped or NaN iterations, and all four + metrics recorded over at least 100 optimizer steps. peft: all: @@ -366,8 +344,9 @@ items: last_10_steps_model_tflops_per_gpu_avg: 731.400 expected_result: > On 128 GB200 (32 nodes x 4 GPUs) with contiguous NVLink-domain placement, - PP=1/EP=64/CP=1/TP=1/DP=2, GBS/MBS 2048/1, FP8-MX with HybridEP dispatcher - and full-iteration CUDA graphs. The 50-step run completes with finite losses, - no skipped or NaN iterations, and all four metrics recorded. First iteration - takes approximately 24 minutes for CUDA graph compilation and kernel warmup; - subsequent iterations run at approximately 8.1 seconds per step at steady state. + PP1/EP64/CP1/TP1 with dense DP128 and expert DP2, GBS/MBS 2048/1, FP8-MX + with HybridEP dispatcher and full-iteration CUDA graphs. The 50-step run + completes with finite losses, no skipped or NaN iterations, and all four + metrics recorded. First iteration takes approximately 24 minutes for CUDA + graph compilation and kernel warmup; subsequent iterations run at + approximately 8.1 seconds per step at steady state. diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 236d5f73a4..9b5c87927f 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -80,8 +80,8 @@ Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 - Adam MXFP8 with PP1/EP64/DP2/HybridEP, selective recompute, attention - activation offload, and MXFP8 parameter gather/buffer reuse + Adam MXFP8 with PP1/EP64/dense-DP128/expert-DP2/HybridEP, selective recompute, + attention activation offload, and MXFP8 parameter gather/buffer reuse - `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config`: 256-GPU GB300 Adam MXFP8 with PP4/VPP4/EP64/HybridEP (requires the PR #4824 container and dev-MCore stack described above) From fe394e7656a107fe3f0ac2972f239ad16a2fca6b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 16:06:54 -0700 Subject: [PATCH 08/35] docs(recipes): clarify DeepSeek V4 validation status Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 9b5c87927f..14cc503e75 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -105,7 +105,8 @@ retain their checkpoint tensor layout. The library recipes retain unlimited natural-routing capacity. All variants retain their precision-specific training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer reuse, and reduced-precision gradient reduction remain disabled except for the -validated FP8 parameter gather/buffer reuse in the high-scale GB200 recipe. +provisional FP8 parameter gather/buffer reuse in the high-scale GB200 recipe; +that PP1 configuration still requires 100-step checkpoint and resume validation. `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. From c5dc1a5a61224b2de902aed94c49c75c0259ba91 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 17:31:48 -0700 Subject: [PATCH 09/35] feat(recipes): add DeepSeek V4 B200 config Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 16 +++-- .../references/recipe-index.md | 1 + .../bridge/recipes/deepseek/__init__.py | 6 ++ .../bridge/recipes/deepseek/b200/__init__.py | 21 ++++++ .../recipes/deepseek/b200/deepseek_v4.py | 69 +++++++++++++++++++ .../test_deepseek_v4_hardware_recipes.py | 47 +++++++++++++ 6 files changed, 154 insertions(+), 6 deletions(-) create mode 100644 src/megatron/bridge/recipes/deepseek/b200/__init__.py create mode 100644 src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 14cc503e75..f2f16509ed 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -73,12 +73,15 @@ Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each ## Pretraining Recipes -See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py) and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. +See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`b200`](../../../src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py), [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py), and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 +- `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config`: 64-GPU B200 + Adam MXFP8 with PP8/VPP2/EP8, all-to-all dispatch, selective recompute, and + attention activation offload plus MXFP8 parameter gather/buffer reuse - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 Adam MXFP8 with PP1/EP64/dense-DP128/expert-DP2/HybridEP, selective recompute, attention activation offload, and MXFP8 parameter gather/buffer reuse @@ -94,9 +97,9 @@ The hardware-count-specific recipes are intentionally not accepted by `slurm_pretrain.sh`: that legacy launcher overwrites the recipe's batch and TP/PP/EP/CP settings. Launch them through a site multi-node runner that invokes `scripts/training/run_recipe.py --recipe --mode pretrain` with -128 ranks for GB200 Flash or 256 ranks for GB300 Pro. Keep the topology and -global batch size owned by the recipe; pass only dataset, run-length, logging, -and checkpoint overrides. +64 ranks for B200 Flash, 128 ranks for GB200 Flash, or 256 ranks for GB300 Pro. +Keep the topology and global batch size owned by the recipe; pass only dataset, +run-length, logging, and checkpoint overrides. The GB200 Flash and GB300 Pro library bases enable their supported DSA and Transformer Engine fused grouped-MLP paths. The new hardware-count-specific @@ -105,8 +108,9 @@ retain their checkpoint tensor layout. The library recipes retain unlimited natural-routing capacity. All variants retain their precision-specific training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer reuse, and reduced-precision gradient reduction remain disabled except for the -provisional FP8 parameter gather/buffer reuse in the high-scale GB200 recipe; -that PP1 configuration still requires 100-step checkpoint and resume validation. +provisional FP8 parameter gather/buffer reuse in the high-scale B200 and GB200 +Flash recipes; both configurations still require 100-step checkpoint and resume +validation. `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index a6f4d8dbca..ee85c5d2bc 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,6 +74,7 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 | | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | +| `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config` | Pretrain | 1 | 8 | 8 | 64 (B200 NVL8; unlimited capacity) | | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) | | `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config` | Pretrain | 1 | 4 | 64 | 256 (GB300; unlimited capacity) | diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 083772d64b..6afaca1080 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -18,6 +18,11 @@ models (V2, V2-Lite, V3, V4). """ +# DeepSeek V4 B200 +from .b200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config, +) + # DeepSeek V2/V2-Lite from .deepseek_v2 import ( deepseek_v2_lite_pretrain_config, @@ -80,6 +85,7 @@ "deepseek_v4_flash_no_mtp_sft_config", "deepseek_v4_pro_pretrain_config", "deepseek_v4_pro_pretrain_mxfp8_config", + "deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", diff --git a/src/megatron/bridge/recipes/deepseek/b200/__init__.py b/src/megatron/bridge/recipes/deepseek/b200/__init__.py new file mode 100644 index 0000000000..d035278680 --- /dev/null +++ b/src/megatron/bridge/recipes/deepseek/b200/__init__.py @@ -0,0 +1,21 @@ +# 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. +"""B200 recipes for DeepSeek V4 Flash.""" + +from megatron.bridge.recipes.deepseek.b200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config, +) + + +__all__ = ["deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config"] diff --git a/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py new file mode 100644 index 0000000000..e4b7b21eb3 --- /dev/null +++ b/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py @@ -0,0 +1,69 @@ +# 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. +"""B200 NVL8 recipes for DeepSeek V4 Flash.""" + +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, +) +from megatron.bridge.training.config import ConfigContainer + + +_DSV4_FLASH_PP8_VP2_LAYOUT = "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL" +_FLEX_DISPATCHER_ENV_VARS = { + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API", + "NVLINK_DOMAIN_SIZE", + "USE_MNNVL", +} + + +def deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config() -> ConfigContainer: + """Return real-training DeepSeek V4 Flash for 64 B200 GPUs. + + PP8/VPP2 keeps each eight-rank expert group within one NVL8 system. The + recipe uses the portable all-to-all dispatcher and preserves natural, + unlimited-capacity routing without paged stash or CUDA graphs. + """ + cfg = deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 8 + cfg.model.virtual_pipeline_model_parallel_size = 2 + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 8 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.pipeline_model_parallel_layout = _DSV4_FLASH_PP8_VP2_LAYOUT + + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_flex_dispatcher_num_sms = None + cfg.model.moe_deepep_num_sms = None + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_hybridep_num_sms_preprocessing = None + cfg.model.moe_shared_expert_overlap = False + + cfg.model.recompute_modules = ["moe", "mhc", "mla_up_proj", "layernorm"] + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["core_attn", "attn_proj"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 2 + cfg.model.moe_pad_experts_for_cuda_graph_inference = False + cfg.model.cuda_graph_impl = "none" + cfg.model.cuda_graph_modules = [] + cfg.model.cuda_graph_scope = None + cfg.model.use_te_rng_tracker = False + cfg.rng.te_rng_tracker = False + + cfg.env_vars = {key: value for key, value in cfg.env_vars.items() if key not in _FLEX_DISPATCHER_ENV_VARS} + return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index c606777290..9220c5884d 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -18,8 +18,12 @@ import pytest import torch +from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout import megatron.bridge.recipes as recipes +from megatron.bridge.recipes.deepseek.b200.deepseek_v4 import ( + deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config as flash_b200_library_config, +) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, ) @@ -188,6 +192,48 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 +def test_flash_b200_nvl8_recipe_preserves_real_training_contract() -> None: + cfg = flash_b200_library_config() + + assert cfg.train.train_iters == 1_000_000 + assert cfg.train.global_batch_size == 256 + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 8 + assert cfg.model.virtual_pipeline_model_parallel_size == 2 + assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.pipeline_model_parallel_layout == ( + "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL" + ) + parsed_layout = PipelineParallelLayerLayout( + cfg.model.pipeline_model_parallel_layout, + pipeline_model_parallel_size=cfg.model.pipeline_model_parallel_size, + ) + parsed_layout.validate_layer_layout(num_layers=43, mtp_num_layers=1) + assert parsed_layout.virtual_pipeline_model_parallel_size == 2 + assert cfg.model.moe_token_dispatcher_type == "alltoall" + assert cfg.model.moe_flex_dispatcher_backend is None + assert cfg.model.moe_flex_dispatcher_num_sms is None + assert cfg.model.moe_router_force_load_balancing is False + assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 + assert cfg.mixed_precision.fp8_param_gather is True + assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True + assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None + assert getattr(cfg.model, "moe_paged_stash", False) is False + assert cfg.model.cuda_graph_impl == "none" + assert cfg.model.cuda_graph_modules == [] + assert cfg.model.cuda_graph_scope is None + assert cfg.rerun_state_machine.check_for_nan_in_loss is True + assert cfg.ddp.check_for_nan_in_grad is True + assert "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" not in cfg.env_vars + assert "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API" not in cfg.env_vars + assert "NVLINK_DOMAIN_SIZE" not in cfg.env_vars + assert "USE_MNNVL" not in cfg.env_vars + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + + def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: cfg = pro_library_config() @@ -224,6 +270,7 @@ def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: def test_high_scale_deepseek_v4_recipes_are_exported() -> None: + assert recipes.deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config is flash_b200_library_config assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config assert recipes.deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config is pro_library_config From 9d4a77bd0bb624c7632f78e179959df5aedf433b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 21:43:43 -0700 Subject: [PATCH 10/35] perf(recipes): optimize DeepSeek V4 GB200 SFT Signed-off-by: Chen Cui --- .../recipes/deepseek/gb200/deepseek_v4.py | 25 +++++++++++++++++-- .../test_deepseek_v4_hardware_recipes.py | 13 +++++++++- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index ce2e8b336b..fa670a74cb 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -293,8 +293,9 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigConta """Return the GB200-optimized offline-packed OpenMath SFT config. This variant adds HybridEP dispatch, safe uneven-input padding, DSA indexer - training, and MoE fusions to the hardware-agnostic packed SFT recipe. It - preserves non-grouped expert GEMMs for checkpoint export compatibility. + training, static packed-sequence shapes, and MoE fusions to the + hardware-agnostic packed SFT recipe. It preserves non-grouped expert GEMMs + for checkpoint export compatibility. """ from megatron.bridge.recipes.deepseek.deepseek_v4 import ( deepseek_v4_flash_sft_openmath_thinking_packed_config, @@ -315,6 +316,26 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigConta cfg.model.moe_permute_fusion = True cfg.model.moe_router_fusion = True + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["moe", "mhc", "mla_up_proj", "layernorm"] + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.model.calculate_per_token_loss = True + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["core_attn", "attn_proj"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 2 + + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 + cfg.dataset.offline_packing_specs.pad_cu_seqlens = True + cfg.dataset.dataset_kwargs = { + **(cfg.dataset.dataset_kwargs or {}), + "pad_to_max_length": True, + } + # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. cfg.model.cross_entropy_fusion_impl = "native" + cfg.env_vars = { + **cfg.env_vars, + "NVTE_CPU_OFFLOAD_V1": 1, + } return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 9220c5884d..535c8596dc 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -141,7 +141,9 @@ def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: cfg = flash_packed_sft_config() assert cfg.model.cp_partition_mode == "contiguous" - assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 2 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert cfg.dataset.offline_packing_specs.pad_cu_seqlens is True + assert cfg.dataset.dataset_kwargs == {"pad_to_max_length": True} 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 @@ -154,6 +156,15 @@ def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: assert cfg.model.moe_router_fusion is True assert cfg.model.moe_grouped_gemm is False assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.model.recompute_granularity == "selective" + assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None + assert cfg.model.calculate_per_token_loss is True + assert cfg.model.fine_grained_activation_offloading is True + assert cfg.model.offload_modules == ["core_attn", "attn_proj"] + assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: From 6d32594bbe19921543890429ca01a195674edff8 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 14 Aug 2026 21:54:55 -0700 Subject: [PATCH 11/35] docs(cards): update DeepSeek V4 verification Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 121 +++++++++--------- 1 file changed, 64 insertions(+), 57 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 70ff7d940a..a77559bdd4 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -5,19 +5,24 @@ title: deepseek_v4_flash summary: > Performance disclaimer: timing and throughput metrics from functional training items are sanity checks, not optimized performance results. pretrain_performance.GB200 uses - the tuned canonical GB200 performance recipe. + the current 128-GPU GB200 library recipe and remains unverified for natural routing. DeepSeek-V4-Flash verification covers GPU conversion and GB200 SFT at the item-pinned Bridge revision. The current 128-GPU GB200 MXFP8 library pretraining recipe uses PP1/EP64 with natural routing, selective recompute, attention activation offload, and MXFP8 parameter gather and gradient-buffer reuse. It has completed an eight-step screen, but full pretraining and direct checkpoint resume remain unverified until the - bounded 100-step runs described below complete. CPU conversion is unsupported because - the 285B model exceeds node CPU RAM. Standard Megatron KV-cache inference is unsupported - for DSV4 hybrid attention; inference uses the HF-native path after export. + bounded 100-step runs described below complete. A separate capacity-1.5 performance + screen averaged 206.56 TFLOP/s/GPU over steps 2-8, but dropped 66.16% of expert + assignments and 64.23% of routing probability mass; that result is benchmark-only and + is excluded from the natural-routing library recipe and verification metrics. CPU + conversion is unsupported because the 285B model exceeds node CPU RAM. Standard + Megatron KV-cache inference is unsupported for DSV4 hybrid attention; inference uses + the HF-native path after export. - The GB200 packed SFT recipe keeps moe_grouped_gemm=False for export compatibility and - uses the canonical MCore uneven-dispatch padding field. + The GB200 packed SFT recipe keeps moe_grouped_gemm=False for export compatibility, + uses static packed-sequence shapes, selective whole-MoE recompute, and attention + activation offload, and completed 100 steps at 20.360 TFLOP/s/GPU over steps 91-100. verification_index: model_level: @@ -32,11 +37,11 @@ verification_index: - manual_forward_pass training: GB200: - verified: [sft, sft_export_inference] - unverified: [pretrain, sft_long_context, checkpoint_resume] + verified: [sft] + unverified: [pretrain, sft_export_inference, sft_long_context, checkpoint_resume] unsupported: [peft] performance: - GB200: verified + GB200: unverified model: hf_id: deepseek-ai/DeepSeek-V4-Flash @@ -171,57 +176,56 @@ items: GB200: status: verified precision: bf16 - bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret + bridge_commit: 9d4a77bd0bb624c7632f78e179959df5aedf433b # pragma: allowlist secret enabled_features: sequence_packing: offline - context_parallel_size: 2 + moe_dispatcher: hybridep command: > - ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 + ./scripts/training/train.sh --nodes 8 --gpus-per-node 4 --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config - --step-func dsv4_step + --mode sft --step-func dsv4_step --pretrained_checkpoint work/models/deepseek-v4-flash-megatron --save_dir work/model-verification/dsv4-flash/sft-ref --save_interval 100 --max_steps 100 - model.pipeline_model_parallel_size=4 + --seq_length 1024 + --pipeline_model_parallel_size 4 --context_parallel_size 1 + --expert_model_parallel_size 8 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' - model.context_parallel_size=2 - model.expert_model_parallel_size=8 - model.moe_grouped_gemm=false - model.recompute_granularity=full model.recompute_method=uniform - model.recompute_num_layers=1 - dataset.seq_length=1024 scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 validation.eval_interval=0 validation.eval_iters=0 rng.seed=5678 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + checkpoint.load_optim=false checkpoint.load_rng=false + checkpoint.save_optim=false ddp.overlap_grad_reduce=false logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=180 - last_verified: 2026-08-07 + last_verified: 2026-08-14 metrics: - initial_loss: 0.94383 - final_loss: 0.28370 - last_10_steps_step_time_ms_avg: 13806.500 - last_10_steps_model_tflops_per_gpu_avg: 13.620 + initial_loss: 0.9432103 + final_loss: 0.2836438 + last_10_steps_step_time_ms_avg: 17551.490 + last_10_steps_model_tflops_per_gpu_avg: 20.360 + peak_allocated_memory_gib: 184.82 + peak_reserved_memory_gib: 187.64 expected_result: > - On 64 GB200 (16 nodes x 4 GPUs), TP1/PP4/EP8/CP2 with dense DP8 and expert + On 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert DP1, GBS/MBS 128/1. OpenMathInstruct-2 thinking data uses offline packing - at seq_length=1024. - Full recompute (granularity=full) is required to fit in memory at seq=1024 - with CP=2. HybridEP dispatcher and DSA kernel fusion are enabled via the - GB200-specific recipe. The run completes 100 steps with finite LM and MTP - losses, no skipped or NaN iterations, and a complete step-100 checkpoint. - moe_grouped_gemm=False is preserved for export compatibility. + at seq_length=1024 with fixed token and cumulative-boundary shapes. The + GB200 recipe uses selective recompute over moe, mhc, mla_up_proj, and + layernorm, attention activation offload, HybridEP dispatch, and DSA kernel + fusion. The run completes 100 steps with finite LM and MTP losses, no skipped + or NaN iterations, and a complete model checkpoint at step 100. + moe_grouped_gemm=False and native fused cross entropy are recipe defaults. dist.distributed_timeout_minutes must be 180 or above to allow offline data packing (approximately 63 minutes) before training begins. sft_export_inference: GB200: - status: verified + status: unverified precision: bf16 - bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret depends_on: sft commands: - > @@ -238,13 +242,11 @@ items: --hf-model work/model-verification/dsv4-flash/sft-export --prompt "In one short sentence, explain why Paris is important to France." --max-new-tokens 32 --chat-template --disable-thinking - last_verified: 2026-08-06 + last_verified: null expected_result: > - The exported BF16 SFT checkpoint reloads successfully as DeepseekV4ForCausalLM. - Two independent greedy runs produce byte-identical token IDs and exactly 32 new - tokens with this literal completion: "Paris is important to France because it is - the country's capital and a major center for politics, culture, and economy. It - is also a symbol of French history" + Verification requires exporting the current optimized SFT step-100 checkpoint, + reloading it successfully as DeepseekV4ForCausalLM, and producing byte-identical + token IDs across two deterministic greedy runs with the configured 32-token bound. sft_long_context: GB200: @@ -324,29 +326,34 @@ items: pretrain_performance: GB200: - status: verified + status: unverified precision: fp8_mx - bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 - --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config - --max_steps 50 - scheduler.lr_warmup_iters=5 scheduler.lr_decay_iters=50 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_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/rp2 dataset.num_workers=0 + dataset.random_seed=1234 rng.seed=1234 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 validation.eval_interval=0 validation.eval_iters=0 - checkpoint.exit_on_missing_checkpoint=false + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + checkpoint.save=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 dist.distributed_timeout_minutes=120 - last_verified: 2026-08-06 + last_verified: null metrics: - initial_loss: 13.58750 - final_loss: 3.62683 - last_10_steps_step_time_ms_avg: 8134.600 - last_10_steps_model_tflops_per_gpu_avg: 731.400 + 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 128 GB200 (32 nodes x 4 GPUs) with contiguous NVLink-domain placement, - PP1/EP64/CP1/TP1 with dense DP128 and expert DP2, GBS/MBS 2048/1, FP8-MX - with HybridEP dispatcher and full-iteration CUDA graphs. The 50-step run - completes with finite losses, no skipped or NaN iterations, and all four - metrics recorded. First iteration takes approximately 24 minutes for CUDA - graph compilation and kernel warmup; subsequent iterations run at - approximately 8.1 seconds per step at steady state. + TP1/PP1/EP64/CP1 with dense DP128 and expert DP2, GBS/MBS 256/1, FP8-MX, + natural routing, selective recompute, and attention activation offload. The + 100-step run must complete with finite losses, no skipped or NaN iterations, + and all four metrics recorded over the final ten steps. The separate + capacity-1.5 screen in the summary is not accepted for this item because it + drops substantial expert assignments and routing probability mass. From 829aa18ef1159aef63880ea0f7018c8597684bc6 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 20 Aug 2026 17:42:00 -0700 Subject: [PATCH 12/35] perf(recipes): enable grouped DeepSeek V4 SFT experts Signed-off-by: Chen Cui --- .../recipes/deepseek/gb200/deepseek_v4.py | 7 +++---- .../deepseek/test_deepseek_v4_conversion.py | 20 +++++++++++++++++++ .../test_deepseek_v4_hardware_recipes.py | 2 +- 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index fa670a74cb..0cad0f3efc 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -293,9 +293,8 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigConta """Return the GB200-optimized offline-packed OpenMath SFT config. This variant adds HybridEP dispatch, safe uneven-input padding, DSA indexer - training, static packed-sequence shapes, and MoE fusions to the - hardware-agnostic packed SFT recipe. It preserves non-grouped expert GEMMs - for checkpoint export compatibility. + training, static packed-sequence shapes, and grouped expert GEMMs to the + hardware-agnostic packed SFT recipe. """ from megatron.bridge.recipes.deepseek.deepseek_v4 import ( deepseek_v4_flash_sft_openmath_thinking_packed_config, @@ -312,7 +311,7 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigConta cfg.model.moe_hybridep_num_sms = 16 cfg.model.moe_shared_expert_overlap = False cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = True - cfg.model.moe_grouped_gemm = False + cfg.model.moe_grouped_gemm = True cfg.model.moe_permute_fusion = True cfg.model.moe_router_fusion = True diff --git a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py index f9d0a39213..bc0527e0e8 100644 --- a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py +++ b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py @@ -225,6 +225,7 @@ def deepseek_v4_toy_model_path(self, tmp_path_factory): def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): test_output_dir = tmp_path / "deepseek_v4_ep" test_output_dir.mkdir(exist_ok=True) + megatron_checkpoint_dir = test_output_dir / "megatron" cmd = [ "python", @@ -249,6 +250,8 @@ def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): "1", "--ep", "2", + "--megatron-save-path", + str(megatron_checkpoint_dir), ] result = subprocess.run( @@ -260,6 +263,23 @@ def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): print(f"STDERR: {result.stderr}") assert result.returncode == 0, f"DeepSeek V4 conversion failed with {result.returncode}" + reload_cmd = [ + *cmd, + "--megatron-load-path", + str(megatron_checkpoint_dir), + "--skip-save", + ] + reload_result = subprocess.run( + reload_cmd, capture_output=True, text=True, cwd=Path(__file__).parent.parent.parent.parent.parent.parent + ) + + if reload_result.returncode != 0: + print(f"RELOAD STDOUT: {reload_result.stdout}") + print(f"RELOAD STDERR: {reload_result.stderr}") + assert reload_result.returncode == 0, ( + f"DeepSeek V4 checkpoint reload conversion failed with {reload_result.returncode}" + ) + converted_dir = test_output_dir / Path(deepseek_v4_toy_model_path).name assert (converted_dir / "config.json").exists() assert list(converted_dir.glob("*.safetensors")) diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 535c8596dc..ca26b0b8ba 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -154,7 +154,7 @@ def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: assert cfg.model.moe_shared_expert_overlap is False assert cfg.model.moe_permute_fusion is True assert cfg.model.moe_router_fusion is True - assert cfg.model.moe_grouped_gemm is False + assert cfg.model.moe_grouped_gemm is True assert cfg.model.cross_entropy_fusion_impl == "native" assert cfg.model.recompute_granularity == "selective" assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] From 758bb98965d0b98883604173e4ba7cdb6851734b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 20 Aug 2026 20:27:54 -0700 Subject: [PATCH 13/35] fix(conversion): preserve DeepSeek V4 indexer exports Signed-off-by: Chen Cui --- .../bridge/models/deepseek/deepseek_v4_bridge.py | 5 +++-- .../models/deepseek/test_deepseek_v4_bridge.py | 11 +++++++---- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 29c4c4da82..4840f8b335 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -941,11 +941,10 @@ def maybe_modify_converted_hf_weight( ) -> Dict[str, torch.Tensor]: """Recreate DSv4 quantized weight/scale pairs expected by the source shard index. + Legacy indexer scorer names are restored before selecting the export dtype. When ``task.weight_dtype`` is set, skip requantization and return the weights unchanged — the generic export path casts the dtype. """ - if task.weight_dtype is not None: - return converted_weights_dict native_scorer_key = next( ( key @@ -959,6 +958,8 @@ def maybe_modify_converted_hf_weight( if legacy_key in hf_state_dict: converted_weights_dict = dict(converted_weights_dict) converted_weights_dict[legacy_key] = converted_weights_dict.pop(native_scorer_key) + if task.weight_dtype is not None: + return converted_weights_dict return quantization_utils.requantize_hf_weight_scale_pairs( converted_weights_dict, hf_state_dict, 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 68eab0f572..6a23f9497c 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -56,10 +56,10 @@ def _by_megatron(registry): return {m.megatron_param: m for m in registry.mappings} -def _dummy_task(): +def _dummy_task(*, weight_dtype=None): from megatron.bridge.models.conversion.model_bridge import WeightConversionTask - return WeightConversionTask(param_name="", global_param_name="", mapping=None) + return WeightConversionTask(param_name="", global_param_name="", mapping=None, weight_dtype=weight_dtype) def _deepseek_v4_hf_config(): @@ -151,13 +151,16 @@ def test_import_accepts_legacy_flat_indexer_weight_name(self): assert result is weight - def test_export_uses_legacy_flat_indexer_weight_name_when_source_requires_it(self): + @pytest.mark.parametrize("weight_dtype", [None, torch.bfloat16]) + def test_export_uses_legacy_flat_indexer_weight_name_when_source_requires_it(self, weight_dtype): bridge = DeepSeekV4Bridge() scorer_name = "layers.1.attn.indexer.scorer.weights_proj.weight" flat_name = "layers.1.attn.indexer.weights_proj.weight" weight = torch.randn(4, 4, dtype=torch.bfloat16) - result = bridge.maybe_modify_converted_hf_weight(_dummy_task(), {scorer_name: weight}, {flat_name: weight}) + result = bridge.maybe_modify_converted_hf_weight( + _dummy_task(weight_dtype=weight_dtype), {scorer_name: weight}, {flat_name: weight} + ) assert scorer_name not in result assert result[flat_name] is weight From efc326458083579bfca797a685a5d1bb982b740b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 20 Aug 2026 20:52:15 -0700 Subject: [PATCH 14/35] docs(cards): record grouped DeepSeek V4 SFT validation Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 52 +++++++++++-------- 1 file changed, 30 insertions(+), 22 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index a77559bdd4..ac95c95847 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -20,9 +20,11 @@ summary: > Megatron KV-cache inference is unsupported for DSV4 hybrid attention; inference uses the HF-native path after export. - The GB200 packed SFT recipe keeps moe_grouped_gemm=False for export compatibility, - uses static packed-sequence shapes, selective whole-MoE recompute, and attention - activation offload, and completed 100 steps at 20.360 TFLOP/s/GPU over steps 91-100. + The GB200 packed SFT recipe uses grouped GEMM, static packed-sequence shapes, + selective whole-MoE recompute, and attention activation offload. It completed + 100 steps at 39.740 TFLOP/s/GPU over steps 91-100, saved and reloaded its + step-100 checkpoint, and exported a complete BF16 checkpoint whose two strict + HF reloads produced identical deterministic output. verification_index: model_level: @@ -37,8 +39,8 @@ verification_index: - manual_forward_pass training: GB200: - verified: [sft] - unverified: [pretrain, sft_export_inference, sft_long_context, checkpoint_resume] + verified: [sft, sft_export_inference] + unverified: [pretrain, sft_long_context, checkpoint_resume] unsupported: [peft] performance: GB200: unverified @@ -50,8 +52,8 @@ model: min_transformers_version: "5.12.0" verification_environment: - base_container: nvcr.io/nvidia/nemo:26.08.rc4 - bridge_commit: 8f7017cad21a328f58932b8d0d3ce07c1a4d241b # pragma: allowlist secret + base_container: nvcr.io/nvidia/nemo:26.08.rc5 + bridge_commit: 758bb98965d0b98883604173e4ba7cdb6851734b # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -176,7 +178,7 @@ items: GB200: status: verified precision: bf16 - bridge_commit: 9d4a77bd0bb624c7632f78e179959df5aedf433b # pragma: allowlist secret + bridge_commit: 829aa18ef1159aef63880ea0f7018c8597684bc6 # pragma: allowlist secret enabled_features: sequence_packing: offline moe_dispatcher: hybridep @@ -202,14 +204,14 @@ items: ddp.overlap_grad_reduce=false logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=180 - last_verified: 2026-08-14 + last_verified: 2026-08-20 metrics: - initial_loss: 0.9432103 - final_loss: 0.2836438 - last_10_steps_step_time_ms_avg: 17551.490 - last_10_steps_model_tflops_per_gpu_avg: 20.360 - peak_allocated_memory_gib: 184.82 - peak_reserved_memory_gib: 187.64 + initial_loss: 0.9435847 + final_loss: 0.2834085 + last_10_steps_step_time_ms_avg: 9124.730 + last_10_steps_model_tflops_per_gpu_avg: 39.740 + peak_allocated_memory_gib: 184.96 + peak_reserved_memory_gib: 187.67 expected_result: > On 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert DP1, GBS/MBS 128/1. OpenMathInstruct-2 thinking data uses offline packing @@ -217,14 +219,15 @@ items: GB200 recipe uses selective recompute over moe, mhc, mla_up_proj, and layernorm, attention activation offload, HybridEP dispatch, and DSA kernel fusion. The run completes 100 steps with finite LM and MTP losses, no skipped - or NaN iterations, and a complete model checkpoint at step 100. - moe_grouped_gemm=False and native fused cross entropy are recipe defaults. + or NaN iterations, and a complete model checkpoint at step 100 that reloads + successfully for finite step 101. moe_grouped_gemm=True and native fused cross + entropy are recipe defaults. dist.distributed_timeout_minutes must be 180 or above to allow offline data packing (approximately 63 minutes) before training begins. sft_export_inference: GB200: - status: unverified + status: verified precision: bf16 depends_on: sft commands: @@ -242,11 +245,16 @@ items: --hf-model work/model-verification/dsv4-flash/sft-export --prompt "In one short sentence, explain why Paris is important to France." --max-new-tokens 32 --chat-template --disable-thinking - last_verified: null + last_verified: 2026-08-20 expected_result: > - Verification requires exporting the current optimized SFT step-100 checkpoint, - reloading it successfully as DeepseekV4ForCausalLM, and producing byte-identical - token IDs across two deterministic greedy runs with the configured 32-token bound. + The grouped-GEMM SFT step-100 checkpoint exports 35,020 BF16 weight tensors + covering every non-scale source key across 46 safetensors shards. Two independent + strict GPU-only DeepseekV4ForCausalLM reloads have no missing, unexpected, or + mismatched weights and produced exactly 30 generated tokens before stopping at + EOS within the configured 32-token maximum. Both deterministic greedy runs had + identical token IDs and this literal completion: "Paris is important to France + because it is the country's capital, a major economic and cultural center, and a + symbol of French history and identity." sft_long_context: GB200: From 07c5cc30d0510d437bfbbe7b12d9d98387295d89 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 20 Aug 2026 21:12:27 -0700 Subject: [PATCH 15/35] fix(recipes): preserve DeepSeek V4 training semantics Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 90 +++++-------- examples/models/deepseek_v4/README.md | 47 +++---- .../references/recipe-index.md | 2 - .../bridge/recipes/deepseek/__init__.py | 8 -- .../bridge/recipes/deepseek/b200/__init__.py | 21 ---- .../recipes/deepseek/b200/deepseek_v4.py | 69 ---------- .../recipes/deepseek/gb200/deepseek_v4.py | 11 +- .../bridge/recipes/deepseek/gb300/__init__.py | 2 - .../recipes/deepseek/gb300/deepseek_v4.py | 59 +-------- .../test_deepseek_v4_hardware_recipes.py | 119 ++---------------- 10 files changed, 72 insertions(+), 356 deletions(-) delete mode 100644 src/megatron/bridge/recipes/deepseek/b200/__init__.py delete mode 100644 src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index ac95c95847..e6dc66c260 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -5,10 +5,10 @@ title: deepseek_v4_flash summary: > Performance disclaimer: timing and throughput metrics from functional training items are sanity checks, not optimized performance results. pretrain_performance.GB200 uses - the current 128-GPU GB200 library recipe and remains unverified for natural routing. + the canonical 128-GPU GB200 benchmark recipe and remains unverified. - DeepSeek-V4-Flash verification covers GPU conversion and GB200 SFT at the item-pinned - Bridge revision. The current 128-GPU GB200 MXFP8 library pretraining recipe uses + DeepSeek-V4-Flash verification covers GPU import at the item-pinned Bridge revision. + The current 128-GPU GB200 MXFP8 library pretraining recipe uses PP1/EP64 with natural routing, selective recompute, attention activation offload, and MXFP8 parameter gather and gradient-buffer reuse. It has completed an eight-step screen, but full pretraining and direct checkpoint resume remain unverified until the @@ -21,26 +21,24 @@ summary: > the HF-native path after export. The GB200 packed SFT recipe uses grouped GEMM, static packed-sequence shapes, - selective whole-MoE recompute, and attention activation offload. It completed - 100 steps at 39.740 TFLOP/s/GPU over steps 91-100, saved and reloaded its - step-100 checkpoint, and exported a complete BF16 checkpoint whose two strict - HF reloads produced identical deterministic output. + selective whole-MoE recompute, and attention activation offload while preserving + the base SFT objective. Its final objective-preserving configuration and public + post-export inference workflow remain unverified. verification_index: model_level: verified: - hf_to_megatron_gpu - - megatron_to_hf_gpu unsupported: - hf_to_megatron_cpu - megatron_to_hf_cpu - inference unverified: - manual_forward_pass + - megatron_to_hf_gpu training: GB200: - verified: [sft, sft_export_inference] - unverified: [pretrain, sft_long_context, checkpoint_resume] + unverified: [pretrain, sft, sft_export_inference, sft_long_context, checkpoint_resume] unsupported: [peft] performance: GB200: unverified @@ -91,9 +89,8 @@ items: tensor merging; nodes consistently OOM before the export completes. megatron_to_hf_gpu: - status: verified + status: unverified precision: bf16 - bridge_commit: 3ecb452af619bef88083e10a7738362df8f5d0bf # pragma: allowlist secret command: > ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 2 --gpus-per-node 4 --ep 8 --pp 4 @@ -101,11 +98,11 @@ items: --megatron-path work/model-verification/dsv4-flash/import-gpu/iter_0000000 --hf-path work/model-verification/dsv4-flash/export-gpu --torch-dtype bfloat16 --export-weight-dtype bfloat16 --trust-remote-code - last_verified: 2026-08-05 + last_verified: null expected_result: > - The command exits successfully and writes 46 BF16 safetensors shards covering - all transformer-layer parameters. The export completes 4176/4176 weight - conversions with zero errors. + Verification requires rerunning the model-level export on the fixed indexer + mapping and proving exact non-scale source-key coverage, all expected BF16 + shards, and a strict Hugging Face reload. The historical run predates that fix. manual_forward_pass: status: unverified @@ -176,9 +173,8 @@ items: sft: GB200: - status: verified + status: unverified precision: bf16 - bridge_commit: 829aa18ef1159aef63880ea0f7018c8597684bc6 # pragma: allowlist secret enabled_features: sequence_packing: offline moe_dispatcher: hybridep @@ -204,21 +200,20 @@ items: ddp.overlap_grad_reduce=false logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=180 - last_verified: 2026-08-20 + last_verified: null metrics: - initial_loss: 0.9435847 - final_loss: 0.2834085 - last_10_steps_step_time_ms_avg: 9124.730 - last_10_steps_model_tflops_per_gpu_avg: 39.740 - peak_allocated_memory_gib: 184.96 - peak_reserved_memory_gib: 187.67 + 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 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert DP1, GBS/MBS 128/1. OpenMathInstruct-2 thinking data uses offline packing at seq_length=1024 with fixed token and cumulative-boundary shapes. The GB200 recipe uses selective recompute over moe, mhc, mla_up_proj, and layernorm, attention activation offload, HybridEP dispatch, and DSA kernel - fusion. The run completes 100 steps with finite LM and MTP losses, no skipped + fusion without enabling the optional DSA indexer loss. Verification requires + 100 steps with finite LM and MTP losses, no skipped or NaN iterations, and a complete model checkpoint at step 100 that reloads successfully for finite step 101. moe_grouped_gemm=True and native fused cross entropy are recipe defaults. @@ -227,7 +222,7 @@ items: sft_export_inference: GB200: - status: verified + status: unverified precision: bf16 depends_on: sft commands: @@ -245,16 +240,12 @@ items: --hf-model work/model-verification/dsv4-flash/sft-export --prompt "In one short sentence, explain why Paris is important to France." --max-new-tokens 32 --chat-template --disable-thinking - last_verified: 2026-08-20 + last_verified: null expected_result: > - The grouped-GEMM SFT step-100 checkpoint exports 35,020 BF16 weight tensors - covering every non-scale source key across 46 safetensors shards. Two independent - strict GPU-only DeepseekV4ForCausalLM reloads have no missing, unexpected, or - mismatched weights and produced exactly 30 generated tokens before stopping at - EOS within the configured 32-token maximum. Both deterministic greedy runs had - identical token IDs and this literal completion: "Paris is important to France - because it is the country's capital, a major economic and cultural center, and a - symbol of French history and identity." + Verification requires exact non-scale source-key coverage across all BF16 + shards, two independent strict GPU-only reloads, and identical bounded greedy + output. The maintained verifier must first support sharding this 285B model and + its BF16 autocast requirement; the current single-device command is insufficient. sft_long_context: GB200: @@ -338,19 +329,8 @@ items: precision: fp8_mx command: > ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 - --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_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/rp2 dataset.num_workers=0 - dataset.random_seed=1234 rng.seed=1234 - scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 - validation.eval_interval=0 validation.eval_iters=0 - checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null - checkpoint.save=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 - dist.distributed_timeout_minutes=120 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config + --mode pretrain --max_steps 50 last_verified: null metrics: initial_loss: null @@ -358,10 +338,8 @@ items: last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null expected_result: > - On 128 GB200 (32 nodes x 4 GPUs) with contiguous NVLink-domain placement, - TP1/PP1/EP64/CP1 with dense DP128 and expert DP2, GBS/MBS 256/1, FP8-MX, - natural routing, selective recompute, and attention activation offload. The - 100-step run must complete with finite losses, no skipped or NaN iterations, - and all four metrics recorded over the final ten steps. The separate - capacity-1.5 screen in the summary is not accepted for this item because it - drops substantial expert assignments and routing probability mass. + On 128 GB200, the exact mock-data benchmark recipe completes 50 steps with + finite losses and throughput. This benchmark uses forced routing, a static + expert-rank capacity, paged stash, full-iteration CUDA graphs, and a different + batch/objective contract, so its metrics must not be compared with natural-routing + library pretraining as convergence evidence. diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index f2f16509ed..99c889a30c 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -44,11 +44,9 @@ 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_library_config` -library recipe and its corresponding performance recipe require the stack -validated by Megatron Bridge PR -[#4824](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4824). The performance -definition preserves its exact benchmark stack: +The full-scale `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config` performance +recipe preserves the stack validated by Megatron Bridge PR +[#4824](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4824): `nvcr.io/nvidia/nemo:26.06.01` with Megatron-LM dev commit `9d46c924dce3818f2b5f894f7380712c780d1801` and the capability-check patch documented in that PR. The Megatron-LM commit pinned by the current @@ -73,44 +71,39 @@ Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each ## Pretraining Recipes -See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`b200`](../../../src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py), [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py), and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. +See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py) and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 -- `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config`: 64-GPU B200 - Adam MXFP8 with PP8/VPP2/EP8, all-to-all dispatch, selective recompute, and - attention activation offload plus MXFP8 parameter gather/buffer reuse - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 Adam MXFP8 with PP1/EP64/dense-DP128/expert-DP2/HybridEP, selective recompute, attention activation offload, and MXFP8 parameter gather/buffer reuse -- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config`: 256-GPU GB300 - Adam MXFP8 with PP4/VPP4/EP64/HybridEP (requires the PR #4824 container and - dev-MCore stack described above) +- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config`: 256-GPU GB300 + performance configuration (requires the PR #4824 container and dev-MCore + stack described above) -The `_library_config` suffix keeps these real-training definitions distinct -from the corresponding benchmark recipe families. They retain natural routing, -convergence batch sizes, correctness checks, and checkpoint behavior. +The `_library_config` suffix keeps the real-training GB200 definition distinct +from the corresponding benchmark recipe. It retains natural routing, the +convergence batch size, correctness checks, and checkpoint behavior. The hardware-count-specific recipes are intentionally not accepted by `slurm_pretrain.sh`: that legacy launcher overwrites the recipe's batch and TP/PP/EP/CP settings. Launch them through a site multi-node runner that invokes `scripts/training/run_recipe.py --recipe --mode pretrain` with -64 ranks for B200 Flash, 128 ranks for GB200 Flash, or 256 ranks for GB300 Pro. -Keep the topology and global batch size owned by the recipe; pass only dataset, -run-length, logging, and checkpoint overrides. - -The GB200 Flash and GB300 Pro library bases enable their supported DSA and -Transformer Engine fused grouped-MLP paths. The new hardware-count-specific -variants additionally enable GLU interleaving; the existing recipe identities -retain their checkpoint tensor layout. The library recipes retain unlimited -natural-routing capacity. All variants retain their precision-specific +128 ranks for GB200 Flash. Keep the topology and global batch size owned by the +recipe; pass only dataset, run-length, logging, and checkpoint overrides. + +The GB200 Flash library base enables its supported DSA and Transformer Engine +fused grouped-MLP paths. The new hardware-count-specific variant additionally +enables GLU interleaving; the existing recipe identities retain their checkpoint +tensor layout. The library recipe retains unlimited natural-routing capacity. +All variants retain their precision-specific training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer reuse, and reduced-precision gradient reduction remain disabled except for the -provisional FP8 parameter gather/buffer reuse in the high-scale B200 and GB200 -Flash recipes; both configurations still require 100-step checkpoint and resume -validation. +provisional FP8 parameter gather/buffer reuse in the high-scale GB200 recipe; +that PP1 configuration still requires 100-step checkpoint and resume validation. `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index ee85c5d2bc..ba5191e787 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,9 +74,7 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 | | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | -| `deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config` | Pretrain | 1 | 8 | 8 | 64 (B200 NVL8; unlimited capacity) | | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) | -| `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config` | Pretrain | 1 | 4 | 64 | 256 (GB300; unlimited capacity) | ### GLM-4.5 diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 6afaca1080..2acef0cedc 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -18,11 +18,6 @@ models (V2, V2-Lite, V3, V4). """ -# DeepSeek V4 B200 -from .b200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config, -) - # DeepSeek V2/V2-Lite from .deepseek_v2 import ( deepseek_v2_lite_pretrain_config, @@ -61,7 +56,6 @@ from .gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config, deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, - deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config, ) @@ -85,7 +79,6 @@ "deepseek_v4_flash_no_mtp_sft_config", "deepseek_v4_pro_pretrain_config", "deepseek_v4_pro_pretrain_mxfp8_config", - "deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", @@ -93,6 +86,5 @@ "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", - "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config", "set_deepseek_v4_pipeline_model_parallel_layout", ] diff --git a/src/megatron/bridge/recipes/deepseek/b200/__init__.py b/src/megatron/bridge/recipes/deepseek/b200/__init__.py deleted file mode 100644 index d035278680..0000000000 --- a/src/megatron/bridge/recipes/deepseek/b200/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# 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. -"""B200 recipes for DeepSeek V4 Flash.""" - -from megatron.bridge.recipes.deepseek.b200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config, -) - - -__all__ = ["deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config"] diff --git a/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py deleted file mode 100644 index e4b7b21eb3..0000000000 --- a/src/megatron/bridge/recipes/deepseek/b200/deepseek_v4.py +++ /dev/null @@ -1,69 +0,0 @@ -# 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. -"""B200 NVL8 recipes for DeepSeek V4 Flash.""" - -from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, -) -from megatron.bridge.training.config import ConfigContainer - - -_DSV4_FLASH_PP8_VP2_LAYOUT = "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL" -_FLEX_DISPATCHER_ENV_VARS = { - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API", - "NVLINK_DOMAIN_SIZE", - "USE_MNNVL", -} - - -def deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config() -> ConfigContainer: - """Return real-training DeepSeek V4 Flash for 64 B200 GPUs. - - PP8/VPP2 keeps each eight-rank expert group within one NVL8 system. The - recipe uses the portable all-to-all dispatcher and preserves natural, - unlimited-capacity routing without paged stash or CUDA graphs. - """ - cfg = deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() - - cfg.model.tensor_model_parallel_size = 1 - cfg.model.pipeline_model_parallel_size = 8 - cfg.model.virtual_pipeline_model_parallel_size = 2 - cfg.model.context_parallel_size = 1 - cfg.model.expert_model_parallel_size = 8 - cfg.model.expert_tensor_parallel_size = 1 - cfg.model.sequence_parallel = False - cfg.model.pipeline_model_parallel_layout = _DSV4_FLASH_PP8_VP2_LAYOUT - - cfg.model.moe_token_dispatcher_type = "alltoall" - cfg.model.moe_flex_dispatcher_backend = None - cfg.model.moe_flex_dispatcher_num_sms = None - cfg.model.moe_deepep_num_sms = None - cfg.model.moe_hybridep_num_sms = None - cfg.model.moe_hybridep_num_sms_preprocessing = None - cfg.model.moe_shared_expert_overlap = False - - cfg.model.recompute_modules = ["moe", "mhc", "mla_up_proj", "layernorm"] - cfg.model.fine_grained_activation_offloading = True - cfg.model.offload_modules = ["core_attn", "attn_proj"] - cfg.model.fine_grained_offloading_max_inflight_offloads = 2 - cfg.model.moe_pad_experts_for_cuda_graph_inference = False - cfg.model.cuda_graph_impl = "none" - cfg.model.cuda_graph_modules = [] - cfg.model.cuda_graph_scope = None - cfg.model.use_te_rng_tracker = False - cfg.rng.te_rng_tracker = False - - cfg.env_vars = {key: value for key, value in cfg.env_vars.items() if key not in _FLEX_DISPATCHER_ENV_VARS} - return cfg diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 0cad0f3efc..59aba598c3 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -218,7 +218,8 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.moe_shared_expert_overlap = False - cfg.model.moe_hybridep_num_sms = 32 + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.moe_hybridep_num_sms = None cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_mlp_glu_interleave_size = 32 cfg.model.use_transformer_engine_op_fuser = True @@ -292,8 +293,8 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config() -> ConfigContainer def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigContainer: """Return the GB200-optimized offline-packed OpenMath SFT config. - This variant adds HybridEP dispatch, safe uneven-input padding, DSA indexer - training, static packed-sequence shapes, and grouped expert GEMMs to the + This variant adds HybridEP dispatch, safe uneven-input padding, fused DSA + execution, static packed-sequence shapes, and grouped expert GEMMs to the hardware-agnostic packed SFT recipe. """ from megatron.bridge.recipes.deepseek.deepseek_v4 import ( @@ -303,12 +304,10 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigConta cfg = deepseek_v4_flash_sft_openmath_thinking_packed_config() 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.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_shared_expert_overlap = False cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = True cfg.model.moe_grouped_gemm = True diff --git a/src/megatron/bridge/recipes/deepseek/gb300/__init__.py b/src/megatron/bridge/recipes/deepseek/gb300/__init__.py index db38cda838..e4b9a7e485 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/__init__.py @@ -17,12 +17,10 @@ from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config, deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config, - deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config, ) __all__ = [ - "deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", ] diff --git a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py index 3301a3ae15..db0dcba42d 100644 --- a/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py @@ -63,7 +63,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.transformer_impl = "transformer_engine" cfg.model.attention_backend = None - cfg.model.apply_dsa_kernel_fusion = True + cfg.model.apply_dsa_kernel_fusion = False cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = deepseek_v4_supports_blackwell_fused_kernels() cfg.model.dsa_indexer_loss_coeff = 0.0 @@ -72,18 +72,15 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.model.moe_token_dispatcher_type = "alltoall" cfg.model.moe_aux_loss_coeff = 0.0 cfg.model.moe_router_force_load_balancing = False - cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.cross_entropy_fusion_impl = "te" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc"] cfg.model.recompute_method = None cfg.model.recompute_num_layers = None - cfg.model.fine_grained_activation_offloading = True - cfg.model.offload_modules = ["core_attn", "attn_proj"] - cfg.model.fine_grained_offloading_max_inflight_offloads = 2 + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None cfg.model.cuda_graph_impl = "none" cfg.model.cuda_graph_scope = "full" cfg.model.cuda_graph_warmup_steps = 3 @@ -117,7 +114,6 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.dist.enable_megatron_core_experimental = True cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) - cfg.comm_overlap.overlap_grad_reduce = True cfg.comm_overlap.delay_wgrad_compute = False cfg.comm_overlap.overlap_moe_expert_parallel_comm = False @@ -125,7 +121,6 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config() -> ConfigContainer: cfg.ddp.use_megatron_fsdp = False cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, - "NVTE_CPU_OFFLOAD_V1": 1, } return cfg @@ -146,6 +141,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.train.train_iters = 1_000_000 cfg.train.global_batch_size = 128 cfg.train.micro_batch_size = 1 + cfg.model.apply_dsa_kernel_fusion = False cfg.model.apply_rope_fusion = True cfg.model.use_fused_mhc = deepseek_v4_supports_blackwell_fused_kernels() cfg.model.dsa_indexer_loss_coeff = 0.0 @@ -190,50 +186,7 @@ def deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() -> ConfigContainer: cfg.model.moe_router_padding_for_fp8 = True cfg.model.mtp_eval_in_bf16 = True cfg.model.quant_recipe = _deepseek_v4_mxfp8_quant_recipe() - return cfg - - -def deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config() -> ConfigContainer: - """Return the real-training DeepSeek V4 Pro config for 256 GB300 GPUs. - - This variant adopts the measured PP4/VPP4/EP64 execution topology while - preserving the 32-GPU library recipe's optimizer, routing, loss, precision, - validation, and checkpoint contracts. - """ - cfg = deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config() - - cfg.model.tensor_model_parallel_size = 1 - cfg.model.pipeline_model_parallel_size = 4 - cfg.model.virtual_pipeline_model_parallel_size = 4 - 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 = "Et*4|(tttt|)*14tmL" - cfg.train.global_batch_size = 128 - cfg.train.micro_batch_size = 1 - - cfg.model.moe_token_dispatcher_type = "flex" - cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_shared_expert_overlap = False - cfg.model.moe_hybridep_num_sms = 32 - cfg.model.moe_mlp_glu_interleave_size = 32 - cfg.model.use_transformer_engine_op_fuser = True - cfg.model.recompute_granularity = "selective" - cfg.model.recompute_modules = ["mla_up_proj", "mhc"] cfg.env_vars = { - **cfg.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_CPU_OFFLOAD_V1": 1, - "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NVTE_NORM_BWD_USE_CUDNN": 1, - "NVTE_NORM_FWD_USE_CUDNN": 1, - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, - "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + **COMMON_RECIPE_ENV_VARS, } return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index ca26b0b8ba..d1d2ae1daa 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -18,12 +18,8 @@ import pytest import torch -from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout import megatron.bridge.recipes as recipes -from megatron.bridge.recipes.deepseek.b200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config as flash_b200_library_config, -) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, ) @@ -36,15 +32,9 @@ from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config as flash_packed_sft_config, ) -from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( - deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config as pro_bf16_base_config, -) from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config as pro_fp8_config, ) -from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( - deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config as pro_library_config, -) from megatron.bridge.training.config import ConfigContainer from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies @@ -57,15 +47,8 @@ def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: patch_recipe_construction_dependencies(monkeypatch) -@pytest.mark.parametrize( - "base_factory", - [flash_bf16_base_config, pro_bf16_base_config], - ids=["flash-bf16", "pro-bf16"], -) -def test_base_recipes_enable_precision_independent_fusions( - base_factory: Callable[[], ConfigContainer], -) -> None: - cfg = base_factory() +def test_flash_base_recipe_enables_precision_independent_fusions() -> None: + cfg = flash_bf16_base_config() assert cfg.model.apply_dsa_kernel_fusion is True assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True @@ -97,17 +80,6 @@ def test_flash_base_recipe_ports_flash_fusions() -> None: assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 -def test_pro_base_recipe_ports_activation_offload() -> None: - cfg = pro_bf16_base_config() - - assert cfg.model.attention_backend is None - assert cfg.train.manual_gc_interval == 5 - assert cfg.model.fine_grained_activation_offloading is True - assert cfg.model.offload_modules == ["core_attn", "attn_proj"] - assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 - assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 - - @pytest.mark.parametrize( "fp8_factory", [flash_fp8_config, pro_fp8_config], @@ -145,11 +117,11 @@ def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: assert cfg.dataset.offline_packing_specs.pad_cu_seqlens is True assert cfg.dataset.dataset_kwargs == {"pad_to_max_length": True} 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.dsa_indexer_loss_coeff == 0.0 + assert cfg.model.dsa_indexer_use_sparse_loss is False assert cfg.model.moe_token_dispatcher_type == "flex" assert cfg.model.moe_flex_dispatcher_backend == "hybridep" - assert cfg.model.moe_hybridep_num_sms == 16 + assert cfg.model.moe_flex_dispatcher_num_sms == 16 assert cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is True assert cfg.model.moe_shared_expert_overlap is False assert cfg.model.moe_permute_fusion is True @@ -189,6 +161,8 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.model.cuda_graph_impl == "none" assert cfg.model.use_transformer_engine_op_fuser is True assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert cfg.model.moe_hybridep_num_sms is None assert cfg.model.pipeline_model_parallel_size == 1 assert cfg.model.pipeline_model_parallel_layout is None assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] @@ -203,85 +177,6 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 -def test_flash_b200_nvl8_recipe_preserves_real_training_contract() -> None: - cfg = flash_b200_library_config() - - assert cfg.train.train_iters == 1_000_000 - assert cfg.train.global_batch_size == 256 - assert cfg.model.tensor_model_parallel_size == 1 - assert cfg.model.pipeline_model_parallel_size == 8 - assert cfg.model.virtual_pipeline_model_parallel_size == 2 - assert cfg.model.expert_model_parallel_size == 8 - assert cfg.model.pipeline_model_parallel_layout == ( - "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL" - ) - parsed_layout = PipelineParallelLayerLayout( - cfg.model.pipeline_model_parallel_layout, - pipeline_model_parallel_size=cfg.model.pipeline_model_parallel_size, - ) - parsed_layout.validate_layer_layout(num_layers=43, mtp_num_layers=1) - assert parsed_layout.virtual_pipeline_model_parallel_size == 2 - assert cfg.model.moe_token_dispatcher_type == "alltoall" - assert cfg.model.moe_flex_dispatcher_backend is None - assert cfg.model.moe_flex_dispatcher_num_sms is None - assert cfg.model.moe_router_force_load_balancing is False - assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] - assert cfg.model.fine_grained_activation_offloading is True - assert cfg.model.offload_modules == ["core_attn", "attn_proj"] - assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 - assert cfg.mixed_precision.fp8_param_gather is True - assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is True - assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None - assert getattr(cfg.model, "moe_paged_stash", False) is False - assert cfg.model.cuda_graph_impl == "none" - assert cfg.model.cuda_graph_modules == [] - assert cfg.model.cuda_graph_scope is None - assert cfg.rerun_state_machine.check_for_nan_in_loss is True - assert cfg.ddp.check_for_nan_in_grad is True - assert "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" not in cfg.env_vars - assert "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API" not in cfg.env_vars - assert "NVLINK_DOMAIN_SIZE" not in cfg.env_vars - assert "USE_MNNVL" not in cfg.env_vars - assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 - - -def test_pro_high_scale_recipe_preserves_real_training_contract() -> None: - cfg = pro_library_config() - - assert cfg.model.pipeline_model_parallel_size == 4 - assert cfg.model.virtual_pipeline_model_parallel_size == 4 - assert cfg.model.pipeline_model_parallel_layout == "Et*4|(tttt|)*14tmL" - assert cfg.train.train_iters == 1_000_000 - assert cfg.train.global_batch_size == 128 - assert cfg.validation.eval_interval == 2000 - assert cfg.validation.eval_iters == 32 - assert cfg.model.moe_router_force_load_balancing is False - assert cfg.model.dsa_indexer_loss_coeff == 0.0 - assert cfg.model.dsa_indexer_use_sparse_loss is False - assert cfg.model.apply_dsa_kernel_fusion is True - assert cfg.model.quant_recipe is not None - assert cfg.model.moe_router_padding_for_fp8 is True - assert cfg.mixed_precision.fp8_param_gather is False - assert cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag is False - assert cfg.mixed_precision.grad_reduce_in_fp32 is True - assert cfg.ddp.grad_reduce_in_fp32 is True - assert cfg.optimizer.main_grads_dtype == torch.float32 - assert cfg.ddp.check_for_nan_in_grad is True - assert cfg.model.cuda_graph_impl == "none" - assert cfg.model.use_transformer_engine_op_fuser is True - assert cfg.model.moe_mlp_glu_interleave_size == 32 - assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None - assert getattr(cfg.model, "moe_paged_stash", False) is False - assert cfg.model.fine_grained_activation_offloading is True - assert cfg.model.offload_modules == ["core_attn", "attn_proj"] - assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 - assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" - assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 - assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - - def test_high_scale_deepseek_v4_recipes_are_exported() -> None: - assert recipes.deepseek_v4_flash_pretrain_64gpu_b200_fp8mx_library_config is flash_b200_library_config assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config - assert recipes.deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_library_config is pro_library_config From aa87fb743234ac0cc1173ba38dc3d24c78f9b381 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 21 Aug 2026 00:20:00 -0700 Subject: [PATCH 16/35] docs(recipes): record DeepSeek V4 validation Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 70 +++++++++++-------- examples/models/deepseek_v4/README.md | 16 +++-- .../deepseek/test_deepseek_v4_conversion.py | 6 +- 3 files changed, 52 insertions(+), 40 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index e6dc66c260..d2a5349932 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -10,10 +10,11 @@ summary: > DeepSeek-V4-Flash verification covers GPU import at the item-pinned Bridge revision. The current 128-GPU GB200 MXFP8 library pretraining recipe uses PP1/EP64 with natural routing, selective recompute, attention activation offload, and - MXFP8 parameter gather and gradient-buffer reuse. It has completed an eight-step - screen, but full pretraining and direct checkpoint resume remain unverified until the - bounded 100-step runs described below complete. A separate capacity-1.5 performance - screen averaged 206.56 TFLOP/s/GPU over steps 2-8, but dropped 66.16% of expert + MXFP8 parameter gather and gradient-buffer reuse. It completed 100 finite steps, + saved a grouped-MXFP8 checkpoint, then restored model, optimizer, and RNG state in a + fresh process for finite steps 101-105. Formal checkpoint-resume verification remains + separate because it requires the step-50 sentinel comparison below. A capacity-1.5 + performance screen averaged 206.56 TFLOP/s/GPU over steps 2-8, but dropped 66.16% of expert assignments and 64.23% of routing probability mass; that result is benchmark-only and is excluded from the natural-routing library recipe and verification metrics. CPU conversion is unsupported because the 285B model exceeds node CPU RAM. Standard @@ -22,8 +23,9 @@ summary: > The GB200 packed SFT recipe uses grouped GEMM, static packed-sequence shapes, selective whole-MoE recompute, and attention activation offload while preserving - the base SFT objective. Its final objective-preserving configuration and public - post-export inference workflow remain unverified. + the base SFT objective. Its final objective-preserving configuration completed + 100 finite steps and a fresh-process checkpoint reload. The public post-export + inference workflow remains unverified. verification_index: model_level: @@ -38,7 +40,8 @@ verification_index: - megatron_to_hf_gpu training: GB200: - unverified: [pretrain, sft, sft_export_inference, sft_long_context, checkpoint_resume] + verified: [pretrain, sft] + unverified: [sft_export_inference, sft_long_context, checkpoint_resume] unsupported: [peft] performance: GB200: unverified @@ -51,7 +54,7 @@ model: verification_environment: base_container: nvcr.io/nvidia/nemo:26.08.rc5 - bridge_commit: 758bb98965d0b98883604173e4ba7cdb6851734b # pragma: allowlist secret + bridge_commit: 07c5cc30d0510d437bfbbe7b12d9d98387295d89 # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -76,8 +79,8 @@ items: --torch-dtype bfloat16 --trust-remote-code last_verified: 2026-08-04 expected_result: > - The command exits successfully and creates iter_0000000. The checkpoint - is reloadable and paired GPU export produces correct HF weights. + The command exits successfully and creates a reloadable iter_0000000 + Megatron checkpoint. megatron_to_hf_cpu: status: unsupported @@ -135,7 +138,7 @@ items: pretrain: GB200: - status: unverified + status: verified precision: fp8_mx enabled_features: moe_dispatcher: hybridep @@ -152,28 +155,28 @@ items: rerun_state_machine.check_for_nan_in_loss=true checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null --save_dir work/model-verification/dsv4-flash/pretrain-ref - --save_interval 50 + --save_interval 100 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 - last_verified: null + last_verified: 2026-08-21 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: 12.59201 + final_loss: 5.960634 + last_10_steps_step_time_ms_avg: 5713.47 + last_10_steps_model_tflops_per_gpu_avg: 131.36 + last_10_steps_tokens_per_second_per_gpu_avg: 1433.804676 expected_result: > The current library recipe resolves to TP1/PP1/EP64/CP1 with dense DP128, expert DP2, and GBS/MBS 256/1. It uses HybridEP natural routing, native fused cross entropy, selective recompute over moe, mhc, mla_up_proj, and layernorm, attention activation offload, and MXFP8 parameter gather and gradient-buffer reuse. Expert capacity, paged stash, and CUDA graphs remain disabled. - Verification requires 100 finite steps with no skipped or NaN iterations and - complete checkpoints at steps 50 and 100; the existing eight-step screen is - insufficient for verified status. + The run completed 100 finite LM/MTP-loss steps with no skipped or NaN + iterations and wrote a complete grouped-MXFP8 checkpoint at step 100. sft: GB200: - status: unverified + status: verified precision: bf16 enabled_features: sequence_packing: offline @@ -200,22 +203,25 @@ items: ddp.overlap_grad_reduce=false logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=180 - last_verified: null + last_verified: 2026-08-20 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: 0.9435847 + final_loss: 0.2834835 + last_10_steps_step_time_ms_avg: 8575.85 + last_10_steps_model_tflops_per_gpu_avg: 42.18 + last_10_steps_tokens_per_second_per_gpu_avg: 477.620294 + peak_allocated_memory_gib: 184.96 + peak_reserved_memory_gib: 187.04 expected_result: > On 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert DP1, GBS/MBS 128/1. OpenMathInstruct-2 thinking data uses offline packing at seq_length=1024 with fixed token and cumulative-boundary shapes. The GB200 recipe uses selective recompute over moe, mhc, mla_up_proj, and layernorm, attention activation offload, HybridEP dispatch, and DSA kernel - fusion without enabling the optional DSA indexer loss. Verification requires - 100 steps with finite LM and MTP losses, no skipped - or NaN iterations, and a complete model checkpoint at step 100 that reloads - successfully for finite step 101. moe_grouped_gemm=True and native fused cross + fusion without enabling the optional DSA indexer loss. The run completed + 100 steps with finite LM and MTP losses, no skipped or NaN iterations, and a + complete model checkpoint at step 100. A fresh process reloaded that checkpoint + and completed finite step 101. moe_grouped_gemm=True and native fused cross entropy are recipe defaults. dist.distributed_timeout_minutes must be 180 or above to allow offline data packing (approximately 63 minutes) before training begins. @@ -261,6 +267,7 @@ items: final_loss: null last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null + last_10_steps_tokens_per_second_per_gpu_avg: null expected_result: > The imported seq_length=1024 SFT run verifies offline packing and CP=2 plumbing, but it is not a long-context result. Verification requires a separately sized @@ -279,6 +286,7 @@ items: final_loss: null last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null + last_10_steps_tokens_per_second_per_gpu_avg: null expected_result: > No LoRA / PEFT recipe exists for DeepSeek-V4 in the current recipe set. @@ -310,6 +318,7 @@ items: final_loss: null last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null + last_10_steps_tokens_per_second_per_gpu_avg: null resume_comparison: reference_item: pretrain sentinel_steps: [51, 100] @@ -337,6 +346,7 @@ items: final_loss: null last_10_steps_step_time_ms_avg: null last_10_steps_model_tflops_per_gpu_avg: null + last_10_steps_tokens_per_second_per_gpu_avg: null expected_result: > On 128 GB200, the exact mock-data benchmark recipe completes 50 steps with finite losses and throughput. This benchmark uses forced routing, a static diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 99c889a30c..92e378d400 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -8,10 +8,13 @@ The bridge supports four published variants out of the same code path. The on-di DeepSeek V4 pretraining requires Megatron-LM `dev`; the Megatron-LM copy pinned by the current Megatron Bridge `main` branch is not supported. Stateless -grouped-FP8 checkpoint saves additionally require Megatron-LM PR -[#6509](https://github.com/NVIDIA/Megatron-LM/pull/6509) until that fix merges, -including saves from the existing MXFP8 recipes. Earlier pretraining -verification used Megatron-LM `dev` commit `35f36c7c9dba` plus PR +grouped-FP8 checkpoint saves additionally require the empty extra-state fix from +Megatron-LM PR [#5997](https://github.com/NVIDIA/Megatron-LM/pull/5997), +including saves from the existing MXFP8 recipes. The fix is available on +Megatron-LM `main`; until it reaches `dev`, apply the equivalent signed commit +`d05cb07ff618` from [#6509](https://github.com/NVIDIA/Megatron-LM/pull/6509) to +a compatible `dev` checkout. Earlier pretraining verification used Megatron-LM +`dev` commit `35f36c7c9dba` plus PR [#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839). The NeMo Framework container uses one shared `/opt/venv` for several source @@ -79,7 +82,7 @@ Available Blackwell pretraining recipes: - `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 - `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 Adam MXFP8 with PP1/EP64/dense-DP128/expert-DP2/HybridEP, selective recompute, - attention activation offload, and MXFP8 parameter gather/buffer reuse + grouped GEMM, attention activation offload, and MXFP8 parameter gather/buffer reuse - `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config`: 256-GPU GB300 performance configuration (requires the PR #4824 container and dev-MCore stack described above) @@ -103,7 +106,8 @@ All variants retain their precision-specific training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer reuse, and reduced-precision gradient reduction remain disabled except for the provisional FP8 parameter gather/buffer reuse in the high-scale GB200 recipe; -that PP1 configuration still requires 100-step checkpoint and resume validation. +that PP1 configuration has completed 100 finite steps, grouped-MXFP8 checkpoint +save, and model/optimizer/RNG reload through step 105. `slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. diff --git a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py index bc0527e0e8..adb5023637 100644 --- a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py +++ b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py @@ -273,11 +273,9 @@ def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): reload_cmd, capture_output=True, text=True, cwd=Path(__file__).parent.parent.parent.parent.parent.parent ) - if reload_result.returncode != 0: - print(f"RELOAD STDOUT: {reload_result.stdout}") - print(f"RELOAD STDERR: {reload_result.stderr}") assert reload_result.returncode == 0, ( - f"DeepSeek V4 checkpoint reload conversion failed with {reload_result.returncode}" + f"DeepSeek V4 checkpoint reload conversion failed with {reload_result.returncode}\n" + f"STDOUT:\n{reload_result.stdout}\nSTDERR:\n{reload_result.stderr}" ) converted_dir = test_output_dir / Path(deepseek_v4_toy_model_path).name From 5d3ed196f950529fc52e8fd4e48eb7d18a52a5ff Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 21 Aug 2026 15:19:27 -0700 Subject: [PATCH 17/35] fix(recipes): close DeepSeek V4 review gaps Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 75 +++++++++---------- examples/models/deepseek_v4/README.md | 13 ++-- .../references/recipe-index.md | 1 + .../recipes/deepseek/gb200/deepseek_v4.py | 18 ++--- .../test_deepseek_v4_hardware_recipes.py | 1 - 5 files changed, 52 insertions(+), 56 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index d2a5349932..527e99a954 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -17,7 +17,7 @@ summary: > performance screen averaged 206.56 TFLOP/s/GPU over steps 2-8, but dropped 66.16% of expert assignments and 64.23% of routing probability mass; that result is benchmark-only and is excluded from the natural-routing library recipe and verification metrics. CPU - conversion is unsupported because the 285B model exceeds node CPU RAM. Standard + conversion remains unverified because the available nodes lacked sufficient RAM. Standard Megatron KV-cache inference is unsupported for DSV4 hybrid attention; inference uses the HF-native path after export. @@ -32,11 +32,11 @@ verification_index: verified: - hf_to_megatron_gpu unsupported: - - hf_to_megatron_cpu - - megatron_to_hf_cpu - inference unverified: + - hf_to_megatron_cpu - manual_forward_pass + - megatron_to_hf_cpu - megatron_to_hf_gpu training: GB200: @@ -58,14 +58,19 @@ verification_environment: items: hf_to_megatron_cpu: - status: unsupported - precision: null - command: null + status: unverified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu + --nodes 1 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/import-cpu + --torch-dtype bfloat16 --trust-remote-code last_verified: null expected_result: > CPU import requires approximately 570 GB of CPU RAM for the full BF16 weight - materialisation (285B parameters x 2 bytes). Standard compute nodes do not - provide this headroom; the import consistently exits with OOM before completing. + materialisation (285B parameters x 2 bytes). Verification requires a high-memory + node with enough additional headroom for conversion workspace and strict checkpoint reload. hf_to_megatron_gpu: status: verified @@ -83,20 +88,26 @@ items: Megatron checkpoint. megatron_to_hf_cpu: - status: unsupported - precision: null - command: null + status: unverified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu + --nodes 1 + --hf-model deepseek-ai/DeepSeek-V4-Flash + --megatron-path work/model-verification/dsv4-flash/import-cpu/iter_0000000 + --hf-path work/model-verification/dsv4-flash/export-cpu + --torch-dtype bfloat16 --trust-remote-code last_verified: null expected_result: > - CPU export requires the full 570 GB BF16 weight set to reside in CPU RAM during - tensor merging; nodes consistently OOM before the export completes. + CPU export requires a high-memory node with enough headroom for the full BF16 + weight set and tensor-merging workspace, followed by strict Hugging Face reload. megatron_to_hf_gpu: status: unverified precision: bf16 command: > ./scripts/conversion/convert.sh export --executor slurm --device gpu - --nodes 2 --gpus-per-node 4 --ep 8 --pp 4 + --nodes 8 --gpus-per-node 4 --ep 8 --pp 4 --hf-model deepseek-ai/DeepSeek-V4-Flash --megatron-path work/model-verification/dsv4-flash/import-gpu/iter_0000000 --hf-path work/model-verification/dsv4-flash/export-gpu @@ -269,10 +280,10 @@ items: last_10_steps_model_tflops_per_gpu_avg: null last_10_steps_tokens_per_second_per_gpu_avg: null expected_result: > - The imported seq_length=1024 SFT run verifies offline packing and CP=2 plumbing, - but it is not a long-context result. Verification requires a separately sized - long-context run with finite loss, no skipped or NaN iterations, and all four - metrics recorded over at least 100 optimizer steps. + The verified seq_length=1024 SFT run established offline packing only and used + CP=1; it is not a long-context or context-parallel result. Verification requires + a separately sized CP=2 long-context run with finite loss, no skipped or NaN + iterations, and all four metrics recorded over at least 100 optimizer steps. peft: all: @@ -295,23 +306,7 @@ items: status: unverified precision: fp8_mx depends_on: pretrain - command: > - ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 - --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_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/rp2 dataset.num_workers=0 - dataset.random_seed=1234 rng.seed=1234 - scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 - validation.eval_interval=0 validation.eval_iters=0 - ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true - rerun_state_machine.check_for_nan_in_loss=true - checkpoint.exit_on_missing_checkpoint=false - --load_dir work/model-verification/dsv4-flash/pretrain-ref - --save_dir work/model-verification/dsv4-flash/pretrain-resumed - --save_interval 50 checkpoint.ckpt_step=50 - logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - dist.distributed_timeout_minutes=120 + command: null last_verified: null metrics: initial_loss: null @@ -326,11 +321,11 @@ items: loss_absolute_tolerance: 1.0e-6 sentinels_match: false expected_result: > - Verification requires restoring optimizer, scheduler, data-order, and RNG state - from the current recipe's step-50 reference checkpoint, then running steps 51-100 - in a separate output root. Step-51 and step-100 losses must match the uninterrupted - reference within the declared 1% relative and 1e-6 absolute tolerances, and all - four metrics must be recorded. + Verification requires a new uninterrupted reference run that saves a step-50 + checkpoint, followed by a fresh process that restores optimizer, scheduler, + data-order, and RNG state and runs steps 51-100 in a separate output root. + Step-51 and step-100 losses must match the uninterrupted reference within the + declared 1% relative and 1e-6 absolute tolerances, and all four metrics must be recorded. pretrain_performance: GB200: diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index d8393f7ea3..50da0b3f56 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -139,12 +139,13 @@ Full-parameter SFT of DeepSeek-V4-Flash. See [`slurm_sft.sh`](slurm_sft.sh) for |--------|-----|------------|----------| | `deepseek_v4_flash_sft_config` | on | fused (cuTile, sm_100) | Blackwell (Hopper: set `use_fused_mhc=False`) | | `deepseek_v4_flash_no_mtp_sft_config` | off | fused (cuTile, sm_100) | Blackwell (Hopper: set `use_fused_mhc=False`) | +| `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | on | fused (cuTile, sm_100) | GB200 (32 GPUs; offline-packed OpenMath) | -Both recipes enable fused mHC and fused rope (`use_fused_mhc=True`, `apply_rope_fusion=True`), matching the pretrain recipes. The historical "fused-kernel SFT NaN" reports are both resolved: the fused-mHC report was a confound, and the fused-rope NaN was a bridge config-mapping bug (`partial_rotary_factor` double-applied) fixed by `rotary_percent=1.0` in #4271 — with that fix, full-model (43-layer, real weights) SFT with rope fusion on 8×GB300 matches the unfused control's loss trajectory. The fused mHC cuTile kernel is **sm_100 (Blackwell)**; on Hopper set `use_fused_mhc=False`. The `no_mtp` variant drops the MTP layer and trims `csa_compress_ratios` back to `num_layers` (the bridge appends an MTP-layer ratio that `transformer_config` would otherwise reject). +All three recipes enable fused mHC and fused rope (`use_fused_mhc=True`, `apply_rope_fusion=True`), matching the pretrain recipes. The historical "fused-kernel SFT NaN" reports are both resolved: the fused-mHC report was a confound, and the fused-rope NaN was a bridge config-mapping bug (`partial_rotary_factor` double-applied) fixed by `rotary_percent=1.0` in #4271 — with that fix, full-model (43-layer, real weights) SFT with rope fusion on 8×GB300 matches the unfused control's loss trajectory. The fused mHC cuTile kernel is **sm_100 (Blackwell)**; on Hopper set `use_fused_mhc=False`. The `no_mtp` variant drops the MTP layer and trims `csa_compress_ratios` back to `num_layers` (the bridge appends an MTP-layer ratio that `transformer_config` would otherwise reject). > There are intentionally **no MXFP8 or Muon SFT recipes**: both were prototyped (mirroring the pretrain recipes) but fail in full-model DSv4-Flash SFT — see Blockers. SFT ships **Adam/bf16**; the pretrain MXFP8/Muon recipes are unaffected. -`slurm_sft.sh` selects the recipe from `MTP` (`on`|`off`); `HARDWARE` (`blackwell`|`hopper`) sets the node topology (on Hopper, also set `use_fused_mhc=False` — the fused mHC kernel is Blackwell-only). It runs the model at `TP=1, PP=4, EP=8` (32 GPUs), and: +`slurm_sft.sh` selects one of the legacy unpacked recipes from `MTP` (`on`|`off`); `HARDWARE` (`blackwell`|`hopper`) sets the node topology (on Hopper, also set `use_fused_mhc=False` — the fused mHC kernel is Blackwell-only). It runs the model at `TP=1, PP=4, EP=8` (32 GPUs), and: 1. Imports `deepseek-ai/DeepSeek-V4-Flash` into a Megatron checkpoint (skipped if already present). 2. Runs full SFT from that checkpoint via `scripts/training/run_recipe.py`. @@ -156,17 +157,17 @@ GPUs per node differ by hardware, so 32 GPUs means a different node count: | GB200 NVL | 4 | 8 | `--nodes=8 --gpus-per-node=4` (default) | | H100/H200 | 8 | 4 | `--nodes=4 --gpus-per-node=8` | -**Sequences are unpacked (SBHD).** The CSA/DSA indexer asserts `packed_seq_params is None` (`csa.py`), so packed/THD sequences are not yet supported on the sparse layers. The recipes ship an unpacked SQuAD config; do not set `dataset.enable_offline_packing=true`. Select another built-in source with `--dataset gsm8k` or use `--dataset local-jsonl dataset.dataset_root=` for JSONL data — both stay unpacked by default. +The legacy `slurm_sft.sh` recipes use unpacked SBHD sequences. The separate GB200 packed recipe is launched through `scripts/training/run_recipe.py`; it uses the Megatron-LM development branch's THD support, offline-packed OpenMath data, fixed shapes with padded `cu_seqlens`, HybridEP uneven padding, and grouped GEMM. It completed 100 finite steps plus a fresh checkpoint reload at sequence length 1024 and CP=1. Long-context and CP=2 remain unverified. **Eval sizing.** Each evaluation draws `validation.eval_iters × global_batch_size` samples; if that exceeds your validation/test split the eval hangs trying to form a batch. `slurm_sft.sh` defaults to a small `EVAL_ITERS=2` and `DO_TEST=false` (the end-of-run test eval is the usual culprit on small test sets) — raise them only when your splits are large enough. ## SFT Status, TODO & Blockers -Validated end to end on **8× GB300 (32 GPU, TP1/PP4/EP8)** with real DeepSeek-V4-Flash weights: HF→Megatron import (FP8/MXFP4→bf16) + full SFT, **MTP on and off**, `lm loss` decreasing with no NaN — at SBHD / bf16 / Adam / 4K. (Those end-to-end runs used `use_fused_mhc=False`; the recipes now default `use_fused_mhc=True`, verified clean by an isolated GB300 re-run and confirmed by a reviewer on GB200.) The items below are **not** implemented and are gated on upstream Megatron-Core (verified against the code and PRs as of 2026-06-02; tracking: [NVIDIA/Megatron-LM#4468](https://github.com/NVIDIA/Megatron-LM/issues/4468)): +Validated end to end on **8× GB300 (32 GPU, TP1/PP4/EP8)** with real DeepSeek-V4-Flash weights: HF→Megatron import (FP8/MXFP4→bf16) + full SFT, **MTP on and off**, `lm loss` decreasing with no NaN — at SBHD / bf16 / Adam / 4K. (Those end-to-end runs used `use_fused_mhc=False`; the recipes now default `use_fused_mhc=True`, verified clean by an isolated GB300 re-run and confirmed by a reviewer on GB200.) Current SFT capability status and remaining upstream blockers are tracked below (verified against the code and PRs as of 2026-06-02; tracking: [NVIDIA/Megatron-LM#4468](https://github.com/NVIDIA/Megatron-LM/issues/4468)): | Capability | Status | Gating | Notes | |------------|--------|--------|-------| -| Packed sequence (**THD**) for DSv4 attention | **Experimental** (unmerged) | mcore PR #5011 (open) | THD *is implemented* in #5011 — it removes the `packed_seq_params is None` asserts in `csa.py`/`dsa.py` and adds the THD index/topk + `cu_seqlens` path. It **grafts cleanly onto our pinned mcore** (0-conflict 3-way merge of #5011) and needs `nvidia-cudnn-frontend[cutedsl]>=1.24.0` installed **`--no-deps`** for the DSA THD kernels (the stock 26.04 container ships an older FE; a full install shadows the container CUDA and breaks TE). The shipped/supported path stays **SBHD**, but an experimental THD (packed-seq) SFT is now **self-validated on 8×GB300** (10-iter, `lm loss` 4.46→1.48, `mtp_1` 13.5→2.70, no NaN) — ahead of #5011 merging. Merged dispatcher-THD #4816 is MoE-side only and does not lift the attention block. | +| Packed sequence (**THD**) for DSv4 attention | Supported by the GB200 packed SFT recipe | Megatron-LM development branch | Verified for 100 finite steps and a fresh checkpoint reload at sequence length 1024 and CP=1; long-context and CP=2 remain separate. | | CUDA Graphs for DSv4 THD | TODO | (no PR) | Follows THD. | | Context parallel / long-context (≥64K) | TODO | draft PR #5087 (depends on #5011) | SFT runs CP=1; this is the Phase-3 long-context target. | | MXFP8 / Muon **SFT** | **Fails (upstream); no SFT recipe shipped** | fp8 numerics; Muon + expert-parallel | Both prototyped (mirroring the pretrain recipes) and tested at full scale on 8×GB300 with `use_fused_mhc=False`: **MXFP8 NaNs at iter-2** (fp8 × hash-MoE/ClampedSwiGLU numerics) and **Muon hits an iter-2 `AssertionError`** (Muon + EP-MoE grad bookkeeping not yet supported upstream). Removed from the shipped recipe set; the **pretrain** MXFP8/Muon recipes remain. SFT ships **Adam/bf16**. | @@ -208,7 +209,7 @@ DSv4 currently requires **TP=1** because MLA tensor parallelism is not supported - **MTP is disabled for inference** via `disable_mtp_for_inference()`. MTP weights are mapped end-to-end and loaded into the Megatron model. -- **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. +- **Fused mHC requires sm_100.** The Blackwell recipes use the fused cuTile kernel. Set `use_fused_mhc=False` on Hopper; the unfused path remains the portable fallback. - **`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. diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index ba5191e787..15fc5a7e55 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -75,6 +75,7 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) | +| `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | SFT | 1 | 4 | 8 | 32 (GB200; offline packed) | ### GLM-4.5 diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index dfc2d73d24..b6821c779f 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -41,9 +41,10 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: """DeepSeek-V4-Flash BF16 pre-training on 64-GPU GB200 (PP=8, EP=8). - All features validated in 1k-step convergence runs on OCI GB200 NVL72: - fused DSA indexer, DSA indexer loss, HybridEP dispatcher, GroupedGEMM, - permute fusion, selective recompute with mla_up_proj. + The core 64-GPU topology completed 1k-step convergence runs with fused DSA + indexer and loss, HybridEP, grouped GEMM, permute fusion, and selective + recompute. The current router fusion, native cross entropy, and activation + offload settings were exercised by the derived 128-GPU recipe's 100-step run. Note: set checkpoint.save_optim=False when using HybridEP. """ @@ -197,11 +198,11 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigContainer: """Return the real-training DeepSeek V4 Flash config for 128 GB200 GPUs. - This variant uses PP1 with two 64-rank expert/data-parallel replicas. A - matched PP1 screen completed eight finite steps and reduced peak allocated - memory by 25.51 GiB with attention activation offload plus MXFP8 parameter - gather/storage. Expert capacity, paged stash, and CUDA graphs remain - disabled to preserve natural-routing training semantics. + This variant uses PP1 with dense DP128, EP64, and expert DP2. Its natural-routing + validation completed 100 finite steps, saved a grouped-MXFP8 checkpoint, and + restored model, optimizer, and RNG state through step 105. Expert capacity, + paged stash, and CUDA graphs remain disabled to preserve natural-routing + training semantics. """ cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() @@ -220,7 +221,6 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.moe_shared_expert_overlap = False cfg.model.moe_flex_dispatcher_num_sms = 32 - cfg.model.moe_hybridep_num_sms = None cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_mlp_glu_interleave_size = 32 cfg.model.use_transformer_engine_op_fuser = True diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index d1d2ae1daa..7d22f4fb74 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -162,7 +162,6 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.model.use_transformer_engine_op_fuser is True assert cfg.model.moe_mlp_glu_interleave_size == 32 assert cfg.model.moe_flex_dispatcher_num_sms == 32 - assert cfg.model.moe_hybridep_num_sms is None assert cfg.model.pipeline_model_parallel_size == 1 assert cfg.model.pipeline_model_parallel_layout is None assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] From ca551e9cb91dd868ae85f6555b558c070cf4aa46 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 21 Aug 2026 15:50:27 -0700 Subject: [PATCH 18/35] docs(deepseek): refresh V4 example guidance Signed-off-by: Chen Cui --- examples/models/deepseek_v4/README.md | 324 ++++++++++++-------------- 1 file changed, 153 insertions(+), 171 deletions(-) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 50da0b3f56..be4a068236 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -1,216 +1,198 @@ # DeepSeek V4 -End-to-end conversion and inference scripts for the DeepSeek V4 family on Megatron Bridge. +Megatron Bridge supports checkpoint conversion and training for the DeepSeek V4 +family. See the +[DeepSeek-V4-Flash verification card](../../model_verification_cards/deepseek-v4-flash/card.yaml) +for the exact commands, revisions, metrics, and current verification status. -The bridge supports four published variants out of the same code path. The on-disk quantisation differs between post-trained (Flash, Pro) and pretrained-only (Flash-Base, Pro-Base) models — see [`docs/models/deepseek/deepseek-v4.md`](../../../docs/models/deepseek/deepseek-v4.md) for the per-variant scheme. +## Supported Variants -## MCore Checkout +All variants use the same bridge. Import dequantizes the published checkpoint to +the requested Megatron training dtype. -DeepSeek V4 pretraining requires Megatron-LM `dev`; the Megatron-LM copy -pinned by the current Megatron Bridge `main` branch is not supported. Stateless -grouped-FP8 checkpoint saves additionally require the empty extra-state fix from -Megatron-LM PR [#5997](https://github.com/NVIDIA/Megatron-LM/pull/5997), -including saves from the existing MXFP8 recipes. The fix is available on -Megatron-LM `main`; until it reaches `dev`, apply the equivalent signed commit -`d05cb07ff618` from [#6509](https://github.com/NVIDIA/Megatron-LM/pull/6509) to -a compatible `dev` checkout. Earlier pretraining verification used Megatron-LM -`dev` commit `35f36c7c9dba` plus PR -[#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839). +| Variant | Hugging Face model | Published weight format | +|---------|--------------------|-------------------------| +| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | FP8 attention and MXFP4 experts | +| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | FP8 with float32 scales | +| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | FP8 attention and MXFP4 experts | +| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | FP8 with float32 scales | -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: +## Runtime Requirements -```bash -./scripts/switch_mcore.sh dev -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: +DeepSeek V4 training requires a compatible Megatron-LM `dev` revision; the +Megatron-LM revision pinned for the Megatron Bridge `main` branch is not a +supported training runtime. Switch the submodule before syncing the environment: ```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 +./scripts/switch_mcore.sh dev +uv sync ``` -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 -[#4824](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4824): -`nvcr.io/nvidia/nemo:26.06.01` with Megatron-LM dev commit -`9d46c924dce3818f2b5f894f7380712c780d1801` and the capability-check patch -documented in that PR. The Megatron-LM commit pinned by the current -Megatron Bridge `main` branch does not provide the required DeepSeek V4 -performance features, so it is not a supported runtime for that recipe. - -| Variant | HF path | Quant scheme | Validation | -|---------|---------|--------------|------------| -| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | FP8 attn + MXFP4 experts | Verified on GB200, last-token logit cosine 0.96-0.99 (short prompts ~0.98, long prompts >1024 tokens ~0.96-0.99) vs official inference | -| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | uniform FP8 (F32 scales) | Verified on GB200, last-real-token logit cosine 0.9866-0.9930, mean 0.9907 vs official inference | -| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | FP8 attn + MXFP4 experts | Import, export, inference verified on GB200 (PP=4 EP=8) and H100 (PP=16 EP=8) | -| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | uniform FP8 (F32 scales) | Same bridge code as Pro; end-to-end untested | - -## Examples - -- `conversion.sh` imports HF weights into Megatron Bridge and exports Megatron checkpoints back to HF format. -- `inference.sh` runs text generation against an HF or Megatron checkpoint. -- `slurm_pretrain.sh` runs the legacy 32-GPU 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 - -See [`slurm_pretrain.sh`](slurm_pretrain.sh) for the legacy 32-GPU Slurm launcher. Hardware-specific recipe definitions live in the [`gb200`](../../../src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py) and [`gb300`](../../../src/megatron/bridge/recipes/deepseek/gb300/deepseek_v4.py) modules. - -Available Blackwell pretraining recipes: - -- `deepseek_v4_flash_pretrain_mxfp8_config`: Adam MXFP8 -- `deepseek_v4_flash_pretrain_muon_config`: Muon BF16 -- `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config`: 128-GPU GB200 - Adam MXFP8 with PP1/EP64/dense-DP128/expert-DP2/HybridEP, selective recompute, - grouped GEMM, attention activation offload, and MXFP8 parameter gather/buffer reuse -- `deepseek_v4_pro_pretrain_256gpu_gb300_fp8mx_config`: 256-GPU GB300 - performance configuration (requires the PR #4824 container and dev-MCore - stack described above) +Grouped-FP8 checkpoint saves also require the stateless grouped extra-state fix +merged in Megatron-LM +[#5997](https://github.com/NVIDIA/Megatron-LM/pull/5997). If the selected `dev` +revision does not contain that fix yet, apply the equivalent signed commit +`d05cb07ff618` from +[#6509](https://github.com/NVIDIA/Megatron-LM/pull/6509) to the compatible `dev` +checkout before enabling checkpoint saves. -The `_library_config` suffix keeps the real-training GB200 definition distinct -from the corresponding benchmark recipe. It retains natural routing, the -convergence batch size, correctness checks, and checkpoint behavior. +`fast-hadamard-transform` is required by DSA and is installed from the pinned +source dependency by `uv sync`. Run the examples in a CUDA-enabled Megatron +Bridge container; see the +[repository installation instructions](../../../README.md#-installation). -The hardware-count-specific recipes are intentionally not accepted by -`slurm_pretrain.sh`: that legacy launcher overwrites the recipe's batch and -TP/PP/EP/CP settings. Launch them through a site multi-node runner that invokes -`scripts/training/run_recipe.py --recipe --mode pretrain` with -128 ranks for GB200 Flash. Keep the topology and global batch size owned by the -recipe; pass only dataset, run-length, logging, and checkpoint overrides. +The Slurm launchers require an account, partition, and container image. Set +them before running any `--executor slurm` conversion or `train.sh` command: -The GB200 Flash library base enables its supported DSA and Transformer Engine -fused grouped-MLP paths. The new hardware-count-specific variant additionally -enables GLU interleaving; the existing recipe identities retain their checkpoint -tensor layout. The library recipe retains unlimited natural-routing capacity. -All variants retain their precision-specific -training policies: full-iteration CUDA graphs, FP8 parameter gather/buffer -reuse, and reduced-precision gradient reduction remain disabled except for the -provisional FP8 parameter gather/buffer reuse in the high-scale GB200 recipe; -that PP1 configuration has completed 100 finite steps, grouped-MXFP8 checkpoint -save, and model/optimizer/RNG reload through step 105. - -`slurm_pretrain.sh` is a GB200 launcher with `TP=1,PP=4,EP=8,CP=1` by default. Indexer loss are disabled for now and is planned for a follow-up. - -Before submitting, set `CONTAINER_IMAGE`. For DCLM, also set `DCLM_DATA_DIR` and `DCLM_CACHE`. Use `CONTAINER_MOUNTS` and `EXTRA_PYTHONPATH` for cluster-specific data, checkouts, and Python dependencies. +```bash +export SLURM_ACCOUNT=your-account +export SLURM_PARTITION=your-partition +export CONTAINER_IMAGE=your-container-image +``` -The bridge's `maybe_modify_loaded_hf_weight` hook dispatches dequantisation by tensor dtype: +Pass `--mount HOST` or `--mount HOST:CONTAINER` for each dataset, checkpoint, +or output path that is not already visible inside the container. -- `int8` -> MXFP4 packed nibbles -> `bfloat16` via the E2M1 lookup table and per-row 16-K-tile E8M0 scales -- `float8_e4m3fn` with companion `.scale` -> `bfloat16` via 128x128 block-scale expansion, handling both E8M0 and F32 scale dtypes +## Conversion -No external dequantisation script is required. +The shared conversion CLI performs dtype-aware FP8/MXFP4 dequantization during +import; no preprocessing script is required. A DeepSeek-V4-Flash GPU import is: -## SFT Recipes +```bash +./scripts/conversion/convert.sh import \ + --executor local --device gpu --gpus-per-node 4 \ + --hf-model deepseek-ai/DeepSeek-V4-Flash \ + --megatron-path work/models/deepseek-v4-flash \ + --tp 1 --pp 1 --ep 4 \ + --torch-dtype bfloat16 --trust-remote-code +``` -Full-parameter SFT of DeepSeek-V4-Flash. See [`slurm_sft.sh`](slurm_sft.sh) for the end-to-end launcher and [`deepseek_v4.py`](../../../src/megatron/bridge/recipes/deepseek/deepseek_v4.py) for the recipe definitions. +Export accepts a different parallel layout because distributed checkpoints can +be resharded while loading: -| Recipe | MTP | mHC kernel | Hardware | -|--------|-----|------------|----------| -| `deepseek_v4_flash_sft_config` | on | fused (cuTile, sm_100) | Blackwell (Hopper: set `use_fused_mhc=False`) | -| `deepseek_v4_flash_no_mtp_sft_config` | off | fused (cuTile, sm_100) | Blackwell (Hopper: set `use_fused_mhc=False`) | -| `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | on | fused (cuTile, sm_100) | GB200 (32 GPUs; offline-packed OpenMath) | +```bash +./scripts/conversion/convert.sh export \ + --executor slurm --device gpu --nodes 8 --gpus-per-node 4 \ + --hf-model deepseek-ai/DeepSeek-V4-Flash \ + --megatron-path work/models/deepseek-v4-flash/iter_0000000 \ + --hf-path work/models/deepseek-v4-flash-hf \ + --tp 1 --pp 4 --ep 8 \ + --torch-dtype bfloat16 --export-weight-dtype bfloat16 \ + --trust-remote-code +``` -All three recipes enable fused mHC and fused rope (`use_fused_mhc=True`, `apply_rope_fusion=True`), matching the pretrain recipes. The historical "fused-kernel SFT NaN" reports are both resolved: the fused-mHC report was a confound, and the fused-rope NaN was a bridge config-mapping bug (`partial_rotary_factor` double-applied) fixed by `rotary_percent=1.0` in #4271 — with that fix, full-model (43-layer, real weights) SFT with rope fusion on 8×GB300 matches the unfused control's loss trajectory. The fused mHC cuTile kernel is **sm_100 (Blackwell)**; on Hopper set `use_fused_mhc=False`. The `no_mtp` variant drops the MTP layer and trims `csa_compress_ratios` back to `num_layers` (the bridge appends an MTP-layer ratio that `transformer_config` would otherwise reject). +GPU import is verified. GPU export, CPU conversion, and manual HF/Megatron +forward correlation remain unverified at the revision recorded in the card. +Do not infer verification from the presence of a command. -> There are intentionally **no MXFP8 or Muon SFT recipes**: both were prototyped (mirroring the pretrain recipes) but fail in full-model DSv4-Flash SFT — see Blockers. SFT ships **Adam/bf16**; the pretrain MXFP8/Muon recipes are unaffected. +[`conversion.sh`](conversion.sh) is a convenience wrapper for import, export, +and optional round-trip checks. Set `MODEL_VARIANT`, `WORKSPACE`, and the +parallelism and executor variables documented in its header before running it. -`slurm_sft.sh` selects one of the legacy unpacked recipes from `MTP` (`on`|`off`); `HARDWARE` (`blackwell`|`hopper`) sets the node topology (on Hopper, also set `use_fused_mhc=False` — the fused mHC kernel is Blackwell-only). It runs the model at `TP=1, PP=4, EP=8` (32 GPUs), and: +## Pretraining -1. Imports `deepseek-ai/DeepSeek-V4-Flash` into a Megatron checkpoint (skipped if already present). -2. Runs full SFT from that checkpoint via `scripts/training/run_recipe.py`. +Use the public [`scripts/training/train.sh`](../../../scripts/training/train.sh) +launcher. Hardware-qualified library recipes are defined under +[`recipes/deepseek`](../../../src/megatron/bridge/recipes/deepseek): -GPUs per node differ by hardware, so 32 GPUs means a different node count: +| Hardware | Recipe | Precision and optimizer | +|----------|--------|-------------------------| +| H100, 32 GPUs | `deepseek_v4_flash_pretrain_32gpu_h100_bf16_config` | BF16 Adam | +| H100, 32 GPUs | `deepseek_v4_flash_pretrain_32gpu_h100_fp8mx_config` | MXFP8 Adam | +| H100, 32 GPUs | `deepseek_v4_flash_pretrain_32gpu_h100_bf16_muon_config` | BF16 Muon | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config` | BF16 Adam | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config` | MXFP8 Adam | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config` | BF16 Muon | +| GB200, 128 GPUs | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | MXFP8 Adam | +| GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config` | BF16 Adam | +| GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config` | MXFP8 Adam | -| Hardware | GPUs/node | Nodes for 32 GPUs | `#SBATCH` | -|----------|-----------|-------------------|-----------| -| GB200 NVL | 4 | 8 | `--nodes=8 --gpus-per-node=4` (default) | -| H100/H200 | 8 | 4 | `--nodes=4 --gpus-per-node=8` | +For example, a short generated-data run of the 128-GPU GB200 library recipe is: -The legacy `slurm_sft.sh` recipes use unpacked SBHD sequences. The separate GB200 packed recipe is launched through `scripts/training/run_recipe.py`; it uses the Megatron-LM development branch's THD support, offline-packed OpenMath data, fixed shapes with padded `cu_seqlens`, HybridEP uneven padding, and grouped GEMM. It completed 100 finite steps plus a fresh checkpoint reload at sequence length 1024 and CP=1. Long-context and CP=2 remain unverified. +```bash +./scripts/training/train.sh --nodes 32 --gpus-per-node 4 \ + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config \ + --mode pretrain --dataset mock --max_steps 10 +``` -**Eval sizing.** Each evaluation draws `validation.eval_iters × global_batch_size` samples; if that exceeds your validation/test split the eval hangs trying to form a batch. `slurm_sft.sh` defaults to a small `EVAL_ITERS=2` and `DO_TEST=false` (the end-of-run test eval is the usual culprit on small test sets) — raise them only when your splits are large enough. +That recipe owns its TP/PP/EP/CP topology, global batch size, natural-routing +semantics, recompute, offload, and checkpoint settings. Keep those defaults for +the qualified configuration and override only the dataset, run length, logging, +and output paths. Its 100-step real-data validation and checkpoint evidence are +recorded in the verification card. -## SFT Status, TODO & Blockers +Compatibility aliases such as `deepseek_v4_flash_pretrain_mxfp8_config` remain +exported, but new launches should use the hardware-qualified names above. +Canonical benchmark recipes under `src/megatron/bridge/perf_recipes/` are +performance references, not substitutes for the library recipes. -Validated end to end on **8× GB300 (32 GPU, TP1/PP4/EP8)** with real DeepSeek-V4-Flash weights: HF→Megatron import (FP8/MXFP4→bf16) + full SFT, **MTP on and off**, `lm loss` decreasing with no NaN — at SBHD / bf16 / Adam / 4K. (Those end-to-end runs used `use_fused_mhc=False`; the recipes now default `use_fused_mhc=True`, verified clean by an isolated GB300 re-run and confirmed by a reviewer on GB200.) Current SFT capability status and remaining upstream blockers are tracked below (verified against the code and PRs as of 2026-06-02; tracking: [NVIDIA/Megatron-LM#4468](https://github.com/NVIDIA/Megatron-LM/issues/4468)): +## Supervised Fine-Tuning -| Capability | Status | Gating | Notes | -|------------|--------|--------|-------| -| Packed sequence (**THD**) for DSv4 attention | Supported by the GB200 packed SFT recipe | Megatron-LM development branch | Verified for 100 finite steps and a fresh checkpoint reload at sequence length 1024 and CP=1; long-context and CP=2 remain separate. | -| CUDA Graphs for DSv4 THD | TODO | (no PR) | Follows THD. | -| Context parallel / long-context (≥64K) | TODO | draft PR #5087 (depends on #5011) | SFT runs CP=1; this is the Phase-3 long-context target. | -| MXFP8 / Muon **SFT** | **Fails (upstream); no SFT recipe shipped** | fp8 numerics; Muon + expert-parallel | Both prototyped (mirroring the pretrain recipes) and tested at full scale on 8×GB300 with `use_fused_mhc=False`: **MXFP8 NaNs at iter-2** (fp8 × hash-MoE/ClampedSwiGLU numerics) and **Muon hits an iter-2 `AssertionError`** (Muon + EP-MoE grad bookkeeping not yet supported upstream). Removed from the shipped recipe set; the **pretrain** MXFP8/Muon recipes remain. SFT ships **Adam/bf16**. | -| Full SFT on **H100-80GB** | needs **≥64 GPU** | hardware | At 32 GPU (TP1/PP4/EP8 ⇒ DP=1) the fp32 master-param buffer can't shard and OOMs. Use **≥64 H100** (PP8 ⇒ DP≥2), or H200-140GB / Blackwell, or PEFT/LoRA. | +DeepSeek-V4-Flash provides BF16 Adam full-parameter SFT recipes: -Already incorporated in the pinned mcore (no action): dense-loss + per-layer rope-type fix (#5018), CSA/HCA (#4458), Hash MoE/ClampedSwiGLU (#4481), MTP+mHC (#4518), fusion kernels (#4894). +| Recipe | Sequence format | MTP | Target | +|--------|-----------------|-----|--------| +| `deepseek_v4_flash_sft_config` | Unpacked SBHD | On | Hopper or Blackwell | +| `deepseek_v4_flash_no_mtp_sft_config` | Unpacked SBHD | Off | Hopper or Blackwell | +| `deepseek_v4_flash_sft_openmath_thinking_packed_config` | Offline-packed THD | On | Portable base | +| `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | Offline-packed THD | On | 32-GPU GB200 | -## Container Image +The recipes select fused mHC only when the runtime supports the Blackwell +kernel; Hopper uses the unfused fallback. The GB200 packed recipe additionally +enables HybridEP, uneven-dispatch padding, DSA fusion, grouped GEMM, selective +recompute, and attention activation offload. -Run inside a container that has the DSv4 prerequisites: Megatron-Bridge on a **`main2dev`** Megatron-LM commit (validated on `ed6b1f65502aec7f2fe27e14a1245c29e435c2a6`; has both `safe_get_world_size` and the DSv4 `csa.py`/`dsa.py`), the DSA dependency `fast_hadamard_transform`, and a pre-built `helpers_cpp`. The [NeMo Framework container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo) — or an image built from `docker/Dockerfile.ci` with the submodule checked out to a `main2dev` commit — provides this. Set `slurm_sft.sh`'s `CONTAINER_IMAGE` to the resulting `.sqsh`. +Launch the verified GB200 packed recipe from an imported BF16 checkpoint: -The first import downloads ~285B of HF weights, so point `HF_HOME` at shared scratch (e.g. `/home/scratch./HF_HOME`) before submitting so the download persists across jobs, and set `HF_TOKEN` if the repo is gated. +```bash +./scripts/training/train.sh --nodes 8 --gpus-per-node 4 \ + --recipe deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config \ + --mode sft --step-func dsv4_step \ + --pretrained_checkpoint work/models/deepseek-v4-flash \ + --save_dir work/results/deepseek-v4-flash-sft \ + --save_interval 100 --max_steps 100 --seq_length 1024 \ + dist.distributed_timeout_minutes=180 +``` -### Disk footprint +The card records 100 finite steps and a fresh-process checkpoint reload for +this configuration at CP=1. Long-context and CP=2 SFT, post-SFT export and HF +inference, and PEFT remain unverified or unsupported as indicated there. MXFP8 +and Muon SFT recipes are intentionally not shipped because full-model tests did +not establish a stable supported configuration. -The full bf16 model is ~570 GB, so plan storage before submitting: +## Legacy Slurm Templates -| Artifact | Size | Notes | -|----------|------|-------| -| HF download cache (`HF_HOME`) | ~150–200 GB | quantized (FP8/MXFP4) | -| Imported Megatron checkpoint (loaded by SFT) | ~570 GB | bf16, `${MEGATRON_DIR}/iter_0000000` | -| Each saved SFT checkpoint | ~570 GB | bf16 model only | +[`slurm_pretrain.sh`](slurm_pretrain.sh) and +[`slurm_sft.sh`](slurm_sft.sh) are customizable 32-GPU site templates retained +for compatibility. They call the lower-level `run_recipe.py` entry point and +do not reproduce the current hardware-qualified validation commands. +`slurm_sft.sh` uses the unpacked SFT recipes; it does not launch the GB200 packed +recipe. Prefer `scripts/training/train.sh` for new runs. -`slurm_sft.sh` does **not** save checkpoints by default (`SAVE_CKPT=0`) — to get training running you only need the base weights *loaded*, not written back — so the whole run fits in ~750 GB (cache + base). Never save optimizer state: distributed Adam for a 285B model is multi-TB (`save_optim=false` is fixed in the script). For a real fine-tune, set `SAVE_CKPT=1` (saving is capped to `KEEP_CKPTS=1` latest checkpoint) and put `HF_HOME` + `MEGATRON_DIR` on shared/project storage so a small (e.g. 2 TB) personal scratch only holds the output. +## Storage -## Parallelism Configurations +DeepSeek-V4-Flash materializes approximately 570 GB of BF16 model weights. +Plan persistent storage before importing or saving checkpoints: -DSv4 currently requires **TP=1** because MLA tensor parallelism is not supported alongside the DSv4 hybrid attention path. Scale via expert and pipeline parallelism instead. +| Artifact | Approximate size | +|----------|------------------| +| Quantized Hugging Face cache | 150-200 GB | +| Imported BF16 Megatron checkpoint | 570 GB | +| Each BF16 model-only SFT checkpoint | 570 GB | -| Model | TP | PP | EP | GPUs | GPU | Verified | -|-------|---:|---:|---:|-----:|-----|----------| -| DeepSeek-V4-Flash | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference | -| DeepSeek-V4-Flash | 1 | 1 | 16 | 16 | H100 80GB | Import, export, inference | -| DeepSeek-V4-Flash-Base | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference | -| DeepSeek-V4-Pro | 1 | 4 | 8 | 32 | GB200 192GB | Import, export, inference | -| DeepSeek-V4-Pro | 1 | 16 | 8 | 128 | H100 80GB | Import, export, inference | +Optimizer state can add multiple terabytes. Disable optimizer-state saves only +when the workflow does not require resumable training; otherwise provision the +required storage and validate checkpoint resume explicitly. ## Known Limitations -- **MTP is disabled for inference** via `disable_mtp_for_inference()`. MTP weights are mapped end-to-end and loaded into the Megatron model. - -- **Fused mHC requires sm_100.** The Blackwell recipes use the fused cuTile kernel. Set `use_fused_mhc=False` on Hopper; the unfused path remains the portable fallback. - -- **`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. +- DeepSeek V4 currently uses TP=1 with hybrid attention; scale with PP, EP, and + DP. Use only recipe-qualified topology changes. +- Standard Megatron KV-cache autoregressive inference is unsupported because + the hybrid-attention path does not accept an inference context. The + [`inference.sh`](inference.sh) example uses legacy full-prefix generation and + is not a verified KV-cache path. HF-native inference is the target after a + verified Megatron-to-HF export. +- Fused mHC requires Blackwell (`sm_100`). Hopper uses the unfused path. +- CPU import and export require enough RAM for the full BF16 model plus + conversion workspace and remain unverified in the current card. From 2859a3ee0f02743c54b390de91fee878c6de8909 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 21 Aug 2026 16:48:19 -0700 Subject: [PATCH 19/35] feat(inference): support sharded HF verification Signed-off-by: Chen Cui --- scripts/inference/setup_inference.py | 11 +- .../scripts/verify_hf_inference.py | 76 +++++++++++-- .../scripts/inference/test_setup_inference.py | 25 ++++- .../test_verify_hf_inference.py | 106 ++++++++++++++++++ 4 files changed, 206 insertions(+), 12 deletions(-) diff --git a/scripts/inference/setup_inference.py b/scripts/inference/setup_inference.py index 1a61ac6b9c..ba1f1ae41d 100755 --- a/scripts/inference/setup_inference.py +++ b/scripts/inference/setup_inference.py @@ -75,7 +75,12 @@ def _build_parser() -> argparse.ArgumentParser: "--gpus_per_node", type=int, dest="gpus_per_node", - help="GPUs and inference tasks per node.", + help="GPUs allocated per node.", + ) + execution.add_argument( + "--tasks-per-node", + type=int, + help="Inference tasks per node; defaults to one task per GPU.", ) execution.add_argument("--cpus-per-task", type=int, help="CPUs allocated to each inference task.") execution.add_argument("--mem", help="Optional Slurm memory request, such as 64G.") @@ -170,6 +175,8 @@ def _validate_args(args: argparse.Namespace) -> None: raise ValueError("--nodes must be at least 1.") if args.gpus_per_node is None or args.gpus_per_node < 1: raise ValueError("--gpus-per-node is required and must be at least 1.") + if args.tasks_per_node is not None and not 1 <= args.tasks_per_node <= args.gpus_per_node: + raise ValueError("--tasks-per-node must be between 1 and --gpus-per-node.") if args.cpus_per_task is not None and args.cpus_per_task < 1: raise ValueError("--cpus-per-task must be at least 1.") if any(not value.strip() for value in args.srun_args): @@ -201,7 +208,7 @@ def _build_executor(args: argparse.Namespace, env_names: list[str], mounts: list account=args.account, partition=args.partition, nodes=args.nodes, - ntasks_per_node=args.gpus_per_node, + ntasks_per_node=args.tasks_per_node or args.gpus_per_node, cpus_per_task=args.cpus_per_task, mem=args.mem, exclusive=True if args.exclusive else None, diff --git a/skills/create-model-verification-card/scripts/verify_hf_inference.py b/skills/create-model-verification-card/scripts/verify_hf_inference.py index d462eb47b6..40fd88f2a6 100644 --- a/skills/create-model-verification-card/scripts/verify_hf_inference.py +++ b/skills/create-model-verification-card/scripts/verify_hf_inference.py @@ -49,12 +49,27 @@ def _parse_args() -> argparse.Namespace: help="Pass enable_thinking=False to the tokenizer chat template.", ) parser.add_argument("--device", default="cuda", help="Torch device used for inference.") + parser.add_argument( + "--device-map", + choices=("auto", "balanced", "balanced_low_0", "sequential"), + help="Optional Hugging Face device-map strategy for sharded model loading.", + ) parser.add_argument( "--dtype", choices=("bfloat16", "float16", "float32"), default="bfloat16", help="Model loading dtype.", ) + parser.add_argument( + "--autocast", + action="store_true", + help="Run generation under autocast using the model loading dtype.", + ) + parser.add_argument( + "--require-gpu-only", + action="store_true", + help="Fail if any model shard is placed on CPU or disk.", + ) args = parser.parse_args() if args.max_new_tokens <= 0: parser.error("--max-new-tokens must be positive") @@ -122,6 +137,23 @@ def _validate_loading_info(loading_info: dict[str, Any]) -> None: raise RuntimeError(f"Exported Hugging Face checkpoint did not reload strictly: {details}") +def _validate_gpu_only_placement(model: Any) -> None: + """Require every model shard to be resident on CUDA devices.""" + device_map = getattr(model, "hf_device_map", None) + if device_map is None: + if getattr(model.device, "type", str(model.device).split(":", 1)[0]) != "cuda": + raise RuntimeError(f"Model is not GPU-only: device={model.device}") + return + + non_gpu = { + name: str(device) + for name, device in device_map.items() + if not isinstance(device, int) and str(device).split(":", 1)[0] != "cuda" + } + if non_gpu: + raise RuntimeError(f"Model has non-GPU placements: {non_gpu}") + + def _load_runtime(args: argparse.Namespace) -> tuple[Any, Any, Any]: """Load torch, the selected HF auto-model, and its tokenizer or processor.""" import torch @@ -137,27 +169,53 @@ def _load_runtime(args: argparse.Namespace) -> tuple[Any, Any, Any]: processor = AutoTokenizer.from_pretrained(args.hf_model, trust_remote_code=args.trust_remote_code) model_cls = AutoModelForCausalLM - model, loading_info = model_cls.from_pretrained( - args.hf_model, - dtype=dtype, - trust_remote_code=args.trust_remote_code, - output_loading_info=True, - ) + model_kwargs = { + "dtype": dtype, + "trust_remote_code": args.trust_remote_code, + "output_loading_info": True, + } + if args.device_map: + model_kwargs["device_map"] = args.device_map + model, loading_info = model_cls.from_pretrained(args.hf_model, **model_kwargs) _validate_loading_info(loading_info) - model = model.to(args.device).eval() + if not args.device_map: + model = model.to(args.device) + model = model.eval() + if args.require_gpu_only: + _validate_gpu_only_placement(model) + LOG.info("Strict HF reload complete (%d model modules)", sum(1 for _ in model.modules())) return torch, model, processor +def _model_input_device(model: Any) -> Any: + """Return the device holding the model's input embeddings.""" + get_input_embeddings = getattr(model, "get_input_embeddings", None) + if get_input_embeddings is not None: + embeddings = get_input_embeddings() + weight = getattr(embeddings, "weight", None) + if weight is not None and getattr(weight.device, "type", None) != "meta": + return weight.device + return model.device + + def main() -> int: """Run one bounded greedy generation and print its completion.""" args = _parse_args() torch, model, processor = _load_runtime(args) tokenizer = getattr(processor, "tokenizer", processor) - inputs = _prepare_inputs(processor, args).to(model.device) + input_device = _model_input_device(model) + inputs = _prepare_inputs(processor, args).to(input_device) prompt_length = inputs["input_ids"].shape[1] pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id - with torch.inference_mode(): + with ( + torch.inference_mode(), + torch.autocast( + device_type=torch.device(input_device).type, + dtype=getattr(torch, args.dtype), + enabled=args.autocast, + ), + ): output = model.generate( **inputs, do_sample=False, diff --git a/tests/unit_tests/scripts/inference/test_setup_inference.py b/tests/unit_tests/scripts/inference/test_setup_inference.py index 8e440f8ca6..84ee969c29 100644 --- a/tests/unit_tests/scripts/inference/test_setup_inference.py +++ b/tests/unit_tests/scripts/inference/test_setup_inference.py @@ -182,6 +182,8 @@ def test_submission_dry_run_aliases_are_consumed(submission_option): [ (("--nodes", "0"), "--nodes"), (("--gpus-per-node", "0"), "--gpus-per-node"), + (("--gpus-per-node", "4", "--tasks-per-node", "0"), "--tasks-per-node"), + (("--gpus-per-node", "4", "--tasks-per-node", "5"), "--tasks-per-node"), (("--cpus-per-task", "0"), "--cpus-per-task"), (("--srun-arg=",), "--srun-arg"), ], @@ -256,6 +258,8 @@ def __init__(self, **kwargs): _launcher_args( "--gpus-per-node", "4", + "--tasks-per-node", + "1", "--srun-arg=--mpi=pmix", "--srun-arg=--container-writable", ) @@ -263,7 +267,7 @@ def __init__(self, **kwargs): executor = module._build_executor(args, ["HF_TOKEN"], ["/host:/container"]) - assert executor.kwargs["ntasks_per_node"] == 4 + assert executor.kwargs["ntasks_per_node"] == 1 assert executor.kwargs["gpus_per_node"] == 4 assert executor.kwargs["exclusive"] is None assert "launcher" not in executor.kwargs @@ -275,6 +279,25 @@ def __init__(self, **kwargs): assert executor.env_vars == {} +def test_slurm_executor_defaults_to_one_task_per_gpu(tmp_path, monkeypatch): + module = _load_setup_inference_module() + + class _SlurmExecutor: + def __init__(self, **kwargs): + self.kwargs = kwargs + + module.run.Packager = lambda: "packager" + module.run.LocalTunnel = lambda **kwargs: types.SimpleNamespace(**kwargs) + module.run.SlurmExecutor = _SlurmExecutor + monkeypatch.setattr(module, "get_nemorun_home", lambda: str(tmp_path)) + args, _ = module.parse_args(_launcher_args("--gpus-per-node", "4")) + + executor = module._build_executor(args, [], []) + + assert executor.kwargs["ntasks_per_node"] == 4 + assert executor.kwargs["gpus_per_node"] == 4 + + def test_slurm_executor_can_request_exclusive_nodes(tmp_path, monkeypatch): module = _load_setup_inference_module() diff --git a/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py b/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py index f98a69b174..49f4ddae26 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_verify_hf_inference.py @@ -18,6 +18,7 @@ import importlib.util import sys +import types from pathlib import Path import pytest @@ -85,6 +86,78 @@ def generate(self, **kwargs): return torch.tensor([[1, 2, 3, 4]]) +def test_sharded_load_uses_device_map_without_moving_model(monkeypatch): + module = _load_module() + calls = [] + + class _LoadedModel: + hf_device_map = {"model.embed_tokens": 0, "model.layers": 1} + + @classmethod + def from_pretrained(cls, model_name, **kwargs): + calls.append((model_name, kwargs)) + return cls(), {key: [] for key in module._LOADING_ISSUE_KEYS} + + def to(self, device): + raise AssertionError(f"sharded model must not be moved to {device}") + + def eval(self): + return self + + def modules(self): + return iter((self,)) + + transformers = types.ModuleType("transformers") + transformers.AutoModelForCausalLM = _LoadedModel + transformers.AutoTokenizer = types.SimpleNamespace(from_pretrained=lambda *args, **kwargs: _Tokenizer()) + monkeypatch.setitem(sys.modules, "transformers", transformers) + args = types.SimpleNamespace( + image=None, + hf_model="model", + trust_remote_code=True, + dtype="bfloat16", + device_map="auto", + device="cuda", + require_gpu_only=True, + ) + + _, model, _ = module._load_runtime(args) + + assert isinstance(model, _LoadedModel) + assert calls == [ + ( + "model", + { + "dtype": torch.bfloat16, + "trust_remote_code": True, + "output_loading_info": True, + "device_map": "auto", + }, + ) + ] + + +def test_model_input_device_prefers_input_embeddings(): + module = _load_module() + model = types.SimpleNamespace( + device=torch.device("cuda:1"), + get_input_embeddings=lambda: types.SimpleNamespace(weight=torch.empty(1, device="meta")), + ) + + assert module._model_input_device(model) == torch.device("cuda:1") + + model.get_input_embeddings = lambda: types.SimpleNamespace(weight=torch.empty(1)) + assert module._model_input_device(model) == torch.device("cpu") + + +def test_gpu_only_placement_rejects_cpu_or_disk_shards(): + module = _load_module() + model = types.SimpleNamespace(hf_device_map={"layers.0": 0, "layers.1": "cpu", "layers.2": "disk"}) + + with pytest.raises(RuntimeError, match="non-GPU placements"): + module._validate_gpu_only_placement(model) + + def test_image_content_uses_processor_native_location_keys(): module = _load_module() @@ -220,3 +293,36 @@ def test_text_main_keeps_the_legacy_tokenizer_path(monkeypatch): assert tokenizer.prompt == "formatted text prompt" assert len(model.calls) == 1 assert "pixel_values" not in model.calls[0] + + +def test_text_main_enables_requested_autocast(monkeypatch): + module = _load_module() + tokenizer = _Tokenizer() + + class _AutocastModel(_Model): + def generate(self, **kwargs): + assert torch.is_autocast_enabled("cpu") + assert torch.get_autocast_dtype("cpu") == torch.bfloat16 + return super().generate(**kwargs) + + model = _AutocastModel() + monkeypatch.setattr(module, "_load_runtime", lambda args: (torch, model, tokenizer)) + monkeypatch.setattr( + sys, + "argv", + [ + "verify_hf_inference.py", + "--hf-model", + "model", + "--prompt", + "Describe Paris.", + "--max-new-tokens", + "2", + "--dtype", + "bfloat16", + "--autocast", + ], + ) + + assert module.main() == 0 + assert len(model.calls) == 1 From ca6a7a177b39846f92b8aab2a7ebf5a17c827e55 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 21 Aug 2026 16:57:16 -0700 Subject: [PATCH 20/35] docs(deepseek): verify GPU and SFT exports Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 58 ++++++++++++------- examples/models/deepseek_v4/README.md | 25 ++++---- .../test_validate_card.py | 2 + 3 files changed, 54 insertions(+), 31 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 527e99a954..243a8801ce 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -24,24 +24,26 @@ summary: > The GB200 packed SFT recipe uses grouped GEMM, static packed-sequence shapes, selective whole-MoE recompute, and attention activation offload while preserving the base SFT objective. Its final objective-preserving configuration completed - 100 finite steps and a fresh-process checkpoint reload. The public post-export - inference workflow remains unverified. + 100 finite steps and a fresh-process checkpoint reload. GPU export of both the + imported model and the step-100 SFT checkpoint produced exact BF16 inventories + and passed independent strict GPU-only HF reloads. The public post-SFT workflow + also produced identical bounded greedy completions from two independent reloads. verification_index: model_level: verified: - hf_to_megatron_gpu + - megatron_to_hf_gpu unsupported: - inference unverified: - hf_to_megatron_cpu - manual_forward_pass - megatron_to_hf_cpu - - megatron_to_hf_gpu training: GB200: - verified: [pretrain, sft] - unverified: [sft_export_inference, sft_long_context, checkpoint_resume] + verified: [pretrain, sft, sft_export_inference] + unverified: [sft_long_context, checkpoint_resume] unsupported: [peft] performance: GB200: unverified @@ -103,20 +105,28 @@ items: weight set and tensor-merging workspace, followed by strict Hugging Face reload. megatron_to_hf_gpu: - status: unverified + status: verified precision: bf16 + bridge_commit: 2859a3ee0f02743c54b390de91fee878c6de8909 # pragma: allowlist secret command: > ./scripts/conversion/convert.sh export --executor slurm --device gpu - --nodes 8 --gpus-per-node 4 --ep 8 --pp 4 + --nodes 2 --gpus-per-node 4 --tp 1 --pp 1 --ep 8 --etp 1 --hf-model deepseek-ai/DeepSeek-V4-Flash --megatron-path work/model-verification/dsv4-flash/import-gpu/iter_0000000 --hf-path work/model-verification/dsv4-flash/export-gpu --torch-dtype bfloat16 --export-weight-dtype bfloat16 --trust-remote-code - last_verified: null + --not-strict + last_verified: 2026-08-21 expected_result: > - Verification requires rerunning the model-level export on the fixed indexer - mapping and proving exact non-scale source-key coverage, all expected BF16 - shards, and a strict Hugging Face reload. The historical run predates that fix. + Two GB200 nodes produced 46 complete safetensors shards containing all 35,020 + expected non-scale source keys. BF16 export intentionally omitted 34,167 + quantization-scale companions; 35,017 exported tensors are BF16 and the three + I32 tid2eid routing tables exactly match their source values after the expected + integer cast. Two independent Transformers processes each strictly reloaded all + 1,500 model modules across four GPUs with no CPU or disk placement and produced + the same bounded greedy token. --not-strict permits only the intentional scale + omission; exact key, shard, dtype, routing-table value, and reload checks remain + required correctness gates. manual_forward_pass: status: unverified @@ -239,8 +249,9 @@ items: sft_export_inference: GB200: - status: unverified + status: verified precision: bf16 + bridge_commit: 2859a3ee0f02743c54b390de91fee878c6de8909 # pragma: allowlist secret depends_on: sft commands: - > @@ -252,17 +263,24 @@ items: --torch-dtype bfloat16 --export-weight-dtype bfloat16 --trust-remote-code --not-strict - > - uv run python - skills/create-model-verification-card/scripts/verify_hf_inference.py + ./scripts/inference/infer.sh --task hf-inference + --nodes 2 --gpus-per-node 4 --tasks-per-node 1 --hf-model work/model-verification/dsv4-flash/sft-export --prompt "In one short sentence, explain why Paris is important to France." - --max-new-tokens 32 --chat-template --disable-thinking - last_verified: null + --max-new-tokens 32 --chat-template --disable-thinking --trust-remote-code + --dtype bfloat16 --device-map auto --autocast --require-gpu-only + last_verified: 2026-08-21 expected_result: > - Verification requires exact non-scale source-key coverage across all BF16 - shards, two independent strict GPU-only reloads, and identical bounded greedy - output. The maintained verifier must first support sharding this 285B model and - its BF16 autocast requirement; the current single-device command is insufficient. + The objective-preserving step-100 SFT checkpoint exported successfully on two + GB200 nodes. Its 46 complete safetensors shards contain all 35,020 expected + non-scale source keys: 35,017 BF16 tensors and three value-matched I32 tid2eid + routing tables, with 34,167 quantization-scale companions intentionally omitted + for BF16. Two maintained-verifier processes independently reloaded all 1,500 model + modules across four GPUs each, reported no missing, unexpected, mismatched, or + error loading entries, used no CPU or disk placement, and generated exactly 32 + new tokens at the 32-token maximum with identical output. The literal completion was: "Paris is + important to France because it is the country's capital, a major economic and + cultural center, and a symbol of French history, art, and architecture." sft_long_context: GB200: diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index be4a068236..c06e4fc8d2 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -72,18 +72,20 @@ be resharded while loading: ```bash ./scripts/conversion/convert.sh export \ - --executor slurm --device gpu --nodes 8 --gpus-per-node 4 \ + --executor slurm --device gpu --nodes 2 --gpus-per-node 4 \ --hf-model deepseek-ai/DeepSeek-V4-Flash \ --megatron-path work/models/deepseek-v4-flash/iter_0000000 \ --hf-path work/models/deepseek-v4-flash-hf \ - --tp 1 --pp 4 --ep 8 \ + --tp 1 --pp 1 --ep 8 \ --torch-dtype bfloat16 --export-weight-dtype bfloat16 \ - --trust-remote-code + --trust-remote-code --not-strict ``` -GPU import is verified. GPU export, CPU conversion, and manual HF/Megatron -forward correlation remain unverified at the revision recorded in the card. -Do not infer verification from the presence of a command. +BF16 export intentionally omits the source checkpoint's quantization-scale +companions, so it uses `--not-strict`. The verification card pairs that option +with exact key, shard, dtype, integer-routing-table, and strict HF reload checks. +GPU import and GPU export are verified. CPU conversion and manual HF/Megatron +forward correlation remain unverified at the revisions recorded in the card. [`conversion.sh`](conversion.sh) is a convenience wrapper for import, export, and optional round-trip checks. Set `MODEL_VARIANT`, `WORKSPACE`, and the @@ -154,11 +156,11 @@ Launch the verified GB200 packed recipe from an imported BF16 checkpoint: dist.distributed_timeout_minutes=180 ``` -The card records 100 finite steps and a fresh-process checkpoint reload for -this configuration at CP=1. Long-context and CP=2 SFT, post-SFT export and HF -inference, and PEFT remain unverified or unsupported as indicated there. MXFP8 -and Muon SFT recipes are intentionally not shipped because full-model tests did -not establish a stable supported configuration. +The card records 100 finite steps, a fresh-process checkpoint reload, and +post-SFT GPU export with deterministic HF inference for this configuration at +CP=1. Long-context and CP=2 SFT remain unverified, and PEFT remains unsupported. +MXFP8 and Muon SFT recipes are intentionally not shipped because full-model +tests did not establish a stable supported configuration. ## Legacy Slurm Templates @@ -179,6 +181,7 @@ Plan persistent storage before importing or saving checkpoints: | Quantized Hugging Face cache | 150-200 GB | | Imported BF16 Megatron checkpoint | 570 GB | | Each BF16 model-only SFT checkpoint | 570 GB | +| Each exported BF16 HF checkpoint | 570 GB | Optimizer state can add multiple terabytes. Disable optimizer-state saves only when the workflow does not require resumable training; otherwise provision the diff --git a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py index 59461a1aa3..0fc21796b7 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py @@ -21,6 +21,8 @@ ("deepseek-v3", "pretrain_performance", "H100"): (4096, 16384, 1024), ("deepseek-v3", "pretrain_performance", "GB200"): (4096, 4096, 256), ("deepseek-v3", "pretrain_performance", "GB300"): (4096, 4096, 256), + ("deepseek-v4-flash", "pretrain", "GB200"): (4096, 256, 128), + ("deepseek-v4-flash", "sft", "GB200"): (1024, 128, 32), ("gemma-4-26b-a4b-it", "sft", "H100"): (4096, 32, 8), ("gemma-4-26b-a4b-it", "peft", "H100"): (4096, 32, 4), ("glm5-2", "pretrain", "H100"): (2048, 1024, 352), From d7d6fe71874d6477cb068d864da33fd4c154bfee Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 24 Aug 2026 15:27:06 -0700 Subject: [PATCH 21/35] feat(recipes): add DeepSeek V4 expert LoRA Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 17 +++++---- examples/models/deepseek_v4/README.md | 37 +++++++++++++++++-- .../references/recipe-index.md | 2 + .../bridge/recipes/deepseek/__init__.py | 4 ++ .../bridge/recipes/deepseek/deepseek_v4.py | 37 +++++++++++++++++++ .../bridge/recipes/deepseek/gb200/__init__.py | 2 + .../recipes/deepseek/gb200/deepseek_v4.py | 9 +++++ .../recipes/test_deepseek_recipes.py | 1 + .../test_deepseek_v4_hardware_recipes.py | 33 +++++++++++++++++ 9 files changed, 132 insertions(+), 10 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 243a8801ce..328074cb27 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -43,8 +43,7 @@ verification_index: training: GB200: verified: [pretrain, sft, sft_export_inference] - unverified: [sft_long_context, checkpoint_resume] - unsupported: [peft] + unverified: [sft_long_context, peft, checkpoint_resume] performance: GB200: unverified @@ -304,12 +303,14 @@ items: iterations, and all four metrics recorded over at least 100 optimizer steps. peft: - all: - status: unsupported - precision: null + GB200: + status: unverified + precision: bf16 command: null last_verified: null - enabled_features: {} + enabled_features: + sequence_packing: offline + moe_dispatcher: hybridep metrics: initial_loss: null final_loss: null @@ -317,7 +318,9 @@ items: last_10_steps_model_tflops_per_gpu_avg: null last_10_steps_tokens_per_second_per_gpu_avg: null expected_result: > - No LoRA / PEFT recipe exists for DeepSeek-V4 in the current recipe set. + The packed OpenMath LoRA recipe must complete 100 finite steps with + separate routed-expert adapters, zero skipped or NaN iterations, all five + metrics, and a reloadable adapter checkpoint. checkpoint_resume: GB200: diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index c06e4fc8d2..d88e1358ae 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -158,9 +158,40 @@ Launch the verified GB200 packed recipe from an imported BF16 checkpoint: The card records 100 finite steps, a fresh-process checkpoint reload, and post-SFT GPU export with deterministic HF inference for this configuration at -CP=1. Long-context and CP=2 SFT remain unverified, and PEFT remains unsupported. -MXFP8 and Muon SFT recipes are intentionally not shipped because full-model -tests did not establish a stable supported configuration. +CP=1. Long-context and CP=2 SFT remain unverified. MXFP8 and Muon SFT recipes +are intentionally not shipped because full-model tests did not establish a +stable supported configuration. + +## Parameter-Efficient Fine-Tuning + +DeepSeek-V4-Flash provides packed OpenMath LoRA recipes that preserve the SFT +data and objective contract: + +| Recipe | Target | +|--------|--------| +| `deepseek_v4_flash_peft_openmath_thinking_packed_config` | Portable base | +| `deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config` | 32-GPU GB200 | + +The LoRA rank and alpha are both 32. Adapters cover the MLA query down/up, KV, +and output projections plus shared and routed expert FC1/FC2 projections. +Routed experts use separate adapters (`share_expert_adapters=False`) rather +than sharing one adapter across the experts local to an EP rank. + +Launch the GB200 recipe from the same imported BF16 checkpoint and packed data +used by SFT: + +```bash +./scripts/training/train.sh --nodes 8 --gpus-per-node 4 \ + --recipe deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config \ + --mode lora --step-func dsv4_step \ + --pretrained_checkpoint work/models/deepseek-v4-flash \ + --save_dir work/results/deepseek-v4-flash-peft \ + --save_interval 100 --max_steps 100 --seq_length 1024 \ + dist.distributed_timeout_minutes=180 +``` + +The verification card records the current loss, checkpoint, and performance +status for this exact recipe. ## Legacy Slurm Templates diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 15fc5a7e55..5878eb74c1 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -76,6 +76,8 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) | | `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | SFT | 1 | 4 | 8 | 32 (GB200; offline packed) | +| `deepseek_v4_flash_peft_openmath_thinking_packed_config` | PEFT | 1 | 4 | 8 | 32 (offline packed) | +| `deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config` | PEFT | 1 | 4 | 8 | 32 (GB200; offline packed) | ### GLM-4.5 diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 2acef0cedc..986a27991a 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -34,6 +34,7 @@ # DeepSeek V4 from .deepseek_v4 import ( deepseek_v4_flash_no_mtp_sft_config, + deepseek_v4_flash_peft_openmath_thinking_packed_config, deepseek_v4_flash_pretrain_config, deepseek_v4_flash_pretrain_gb200_config, deepseek_v4_flash_pretrain_muon_config, @@ -47,6 +48,7 @@ set_deepseek_v4_pipeline_model_parallel_layout, ) from .gb200.deepseek_v4 import ( + deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, @@ -68,6 +70,8 @@ "deepseek_v3_pretrain_config_32nodes", "set_deepseek_v3_pipeline_model_parallel_layout", # DeepSeek V4 + "deepseek_v4_flash_peft_openmath_thinking_packed_config", + "deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config", "deepseek_v4_flash_pretrain_config", "deepseek_v4_flash_pretrain_gb200_config", "deepseek_v4_flash_pretrain_muon_config", diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index 43a255b8d0..603e60e33a 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py @@ -19,6 +19,7 @@ from megatron.bridge.models.deepseek.deepseek_v4_bridge import ( set_deepseek_v4_pipeline_model_parallel_layout, ) +from megatron.bridge.peft.lora import LoRA from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, @@ -66,6 +67,7 @@ __all__ = [ "deepseek_v4_flash_no_mtp_sft_config", + "deepseek_v4_flash_peft_openmath_thinking_packed_config", "deepseek_v4_flash_pretrain_config", "deepseek_v4_flash_pretrain_muon_config", "deepseek_v4_flash_pretrain_mxfp8_config", @@ -85,6 +87,29 @@ ] +_DEEPSEEK_V4_LORA_TARGET_MODULES = [ + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_proj", + "linear_proj", + "linear_fc1", + "linear_fc2", +] + + +def _apply_deepseek_v4_lora(cfg: ConfigContainer) -> None: + """Apply the DeepSeek V4 LoRA convergence contract to an SFT config.""" + cfg.peft = LoRA( + target_modules=list(_DEEPSEEK_V4_LORA_TARGET_MODULES), + dim=32, + alpha=32, + dropout=0.0, + share_expert_adapters=False, + ) + cfg.optimizer.lr = 1.0e-4 + cfg.optimizer.min_lr = 0.0 + + def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: """DSv4 Flash SFT on OpenMathInstruct-2 with thinking channel and offline-packed sequences. @@ -107,3 +132,15 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: pad_seq_to_mult=2 * cfg.model.context_parallel_size, ) return cfg + + +def deepseek_v4_flash_peft_openmath_thinking_packed_config() -> ConfigContainer: + """DSv4 Flash LoRA on packed OpenMathInstruct-2 thinking data. + + The attention targets follow the DeepSeek MLA projection layout. Both shared + and routed MLP projections are adapted; grouped routed experts use one adapter + per local expert to match the verl training layout. + """ + cfg = deepseek_v4_flash_sft_openmath_thinking_packed_config() + _apply_deepseek_v4_lora(cfg) + return cfg diff --git a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py index 3e1c7c09d7..b91ceffda7 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -14,6 +14,7 @@ """GB200 recipes for DeepSeek V4 Flash.""" from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, @@ -23,6 +24,7 @@ __all__ = [ + "deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config", "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index b6821c779f..dc94f11d39 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -338,3 +338,12 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigConta "NVTE_CPU_OFFLOAD_V1": 1, } return cfg + + +def deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config() -> ConfigContainer: + """Return the GB200-optimized packed OpenMath LoRA config.""" + from megatron.bridge.recipes.deepseek.deepseek_v4 import _apply_deepseek_v4_lora + + cfg = deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() + _apply_deepseek_v4_lora(cfg) + return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_recipes.py b/tests/unit_tests/recipes/test_deepseek_recipes.py index 6981105be0..f1b55e424a 100644 --- a/tests/unit_tests/recipes/test_deepseek_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_recipes.py @@ -46,6 +46,7 @@ { "deepseek_v3_pretrain_config", "deepseek_v3_pretrain_config_32nodes", + "deepseek_v4_flash_peft_openmath_thinking_packed_config", "deepseek_v4_flash_pretrain_config", "deepseek_v4_flash_pretrain_mxfp8_config", "deepseek_v4_flash_pretrain_muon_config", diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 7d22f4fb74..8b75d05241 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -20,6 +20,9 @@ import torch import megatron.bridge.recipes as recipes +from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( + deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config as flash_packed_peft_config, +) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config as flash_bf16_base_config, ) @@ -139,6 +142,35 @@ def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 +def test_flash_packed_peft_recipe_adapts_mla_and_unshared_experts() -> None: + cfg = flash_packed_peft_config() + + assert cfg.peft.target_modules == [ + "linear_q_down_proj", + "linear_q_up_proj", + "linear_kv_proj", + "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.share_expert_adapters is False + assert cfg.optimizer.lr == 1.0e-4 + assert cfg.optimizer.min_lr == 0.0 + + sft_cfg = flash_packed_sft_config() + assert cfg.dataset == sft_cfg.dataset + assert cfg.train == sft_cfg.train + assert cfg.rng == sft_cfg.rng + assert cfg.model.seq_length == sft_cfg.model.seq_length + assert cfg.model.pipeline_model_parallel_size == sft_cfg.model.pipeline_model_parallel_size + assert cfg.model.expert_model_parallel_size == sft_cfg.model.expert_model_parallel_size + assert cfg.model.moe_token_dispatcher_type == sft_cfg.model.moe_token_dispatcher_type + assert cfg.model.recompute_modules == sft_cfg.model.recompute_modules + + def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: cfg = flash_library_config() @@ -177,5 +209,6 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: def test_high_scale_deepseek_v4_recipes_are_exported() -> None: + assert recipes.deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config is flash_packed_peft_config assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config From 2920e647812d22ecfccaad1b5952e5593f6c0681 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 24 Aug 2026 15:53:33 -0700 Subject: [PATCH 22/35] fix(training): support MCore dev GDN APIs Signed-off-by: Chen Cui --- .../bridge/training/utils/flop_utils.py | 19 +++++++--- .../training/utils/test_flop_utils.py | 36 +++++++++++++++++++ 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index b499cdd034..7e9d86ddd0 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -12,13 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path +from typing import cast import torch -from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( - is_gated_delta_net_variant, -) +from megatron.core.models.gpt import experimental_attention_variant_module_specs from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, get_hybrid_layer_counts, @@ -33,6 +33,17 @@ _lora_seq_stats_cache: dict = {} +_mcore_is_gated_delta_net_variant = cast( + Callable[[str | None], bool] | None, + getattr(experimental_attention_variant_module_specs, "is_gated_delta_net_variant", None), +) + + +def _is_gated_delta_net_variant(experimental_attention_variant: str | None) -> bool: + """Recognize GDN variants across current main and older MCore dev branches.""" + if _mcore_is_gated_delta_net_variant is not None: + return _mcore_is_gated_delta_net_variant(experimental_attention_variant) + return experimental_attention_variant in {"gated_delta_net", "gdn", "gdn2"} @dataclass(frozen=True) @@ -1409,7 +1420,7 @@ def count_indexer_layers(layer_count: int) -> int: # Handle GDN (Gated DeltaNet) hybrid attention variants. MCore normalizes # the deprecated "gated_delta_net" alias to "gdn" during config finalization. - if is_gated_delta_net_variant(experimental_attention_variant): + if _is_gated_delta_net_variant(experimental_attention_variant): linear_attention_freq = cfg.model.linear_attention_freq decoder_num_layers = cfg.model.num_layers if linear_attention_freq is None: diff --git a/tests/unit_tests/training/utils/test_flop_utils.py b/tests/unit_tests/training/utils/test_flop_utils.py index 6af0cb17ee..6dc69077db 100644 --- a/tests/unit_tests/training/utils/test_flop_utils.py +++ b/tests/unit_tests/training/utils/test_flop_utils.py @@ -14,6 +14,8 @@ """Unit tests for flop_utils module.""" +import subprocess +import sys from dataclasses import dataclass, field from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -22,6 +24,7 @@ import torch from megatron.bridge.peft.lora import LoRA +from megatron.bridge.training.utils import flop_utils from megatron.bridge.training.utils.flop_utils import ( GlobalFlopsRuntimeStats, _lora_seq_stats_cache, @@ -711,6 +714,39 @@ def _qwen35_27b_config(self, **overrides): defaults.update(overrides) return MockModelConfig(**defaults) + def test_flop_utils_imports_without_mcore_gdn_helper(self) -> None: + code = """ +import importlib +import sys + +from megatron.core.models.gpt import experimental_attention_variant_module_specs +from megatron.bridge.training import utils + +import megatron.bridge.training.utils.flop_utils + +experimental_attention_variant_module_specs.__dict__.pop("is_gated_delta_net_variant", None) +del sys.modules["megatron.bridge.training.utils.flop_utils"] +del utils.flop_utils +flop_utils = importlib.import_module("megatron.bridge.training.utils.flop_utils") + +assert flop_utils._mcore_is_gated_delta_net_variant is None +assert flop_utils._is_gated_delta_net_variant("gdn") is True +""" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + + assert result.returncode == 0, result.stderr + + @pytest.mark.parametrize( + ("variant", "expected"), + [("gated_delta_net", True), ("gdn", True), ("gdn2", True), ("mamba", False), (None, False)], + ) + def test_gdn_variant_fallback_supports_older_mcore_dev( + self, monkeypatch: pytest.MonkeyPatch, variant: str | None, expected: bool + ) -> None: + monkeypatch.setattr(flop_utils, "_mcore_is_gated_delta_net_variant", None) + + assert flop_utils._is_gated_delta_net_variant(variant) is expected + def test_gdn_flops_differ_from_pure_attention(self): """GDN-enabled config should produce different FLOPs than pure-attention baseline.""" batch_size = 1 From 263f9f5172611eeb17f1e0abed572fc49fdd639c Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 24 Aug 2026 16:19:32 -0700 Subject: [PATCH 23/35] fix(recipes): disable activation offload for V4 PEFT Signed-off-by: Chen Cui --- src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py | 3 +++ tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index dc94f11d39..8a349a036c 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -346,4 +346,7 @@ def deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config() -> ConfigCont cfg = deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() _apply_deepseek_v4_lora(cfg) + # Unified recompute cannot replay offloaded activations for unshared grouped-expert adapters. + cfg.model.fine_grained_activation_offloading = False + cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 8b75d05241..e7757bd8df 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -159,6 +159,8 @@ def test_flash_packed_peft_recipe_adapts_mla_and_unshared_experts() -> None: assert cfg.peft.share_expert_adapters is False assert cfg.optimizer.lr == 1.0e-4 assert cfg.optimizer.min_lr == 0.0 + assert cfg.model.fine_grained_activation_offloading is False + assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars sft_cfg = flash_packed_sft_config() assert cfg.dataset == sft_cfg.dataset From 1a67d67f01e39eb2399fe59fb1f24b7f241a787c Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 24 Aug 2026 16:44:05 -0700 Subject: [PATCH 24/35] fix(recipes): disable recompute for V4 PEFT Signed-off-by: Chen Cui --- src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py | 8 +++++++- .../recipes/test_deepseek_v4_hardware_recipes.py | 6 +++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 8a349a036c..86091dd141 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -346,7 +346,13 @@ def deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config() -> ConfigCont cfg = deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() _apply_deepseek_v4_lora(cfg) - # Unified recompute cannot replay offloaded activations for unshared grouped-expert adapters. + # PEFT's smaller training-state footprint lets this adapter graph avoid MCore's + # unified recompute path, which cannot replay the frozen grouped-expert graph. + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index e7757bd8df..b3230732ff 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -159,7 +159,12 @@ def test_flash_packed_peft_recipe_adapts_mla_and_unshared_experts() -> None: assert cfg.peft.share_expert_adapters is False assert cfg.optimizer.lr == 1.0e-4 assert cfg.optimizer.min_lr == 0.0 + assert cfg.model.recompute_granularity is None + assert cfg.model.recompute_modules is None + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None assert cfg.model.fine_grained_activation_offloading is False + assert cfg.model.offload_modules is None assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars sft_cfg = flash_packed_sft_config() @@ -170,7 +175,6 @@ def test_flash_packed_peft_recipe_adapts_mla_and_unshared_experts() -> None: assert cfg.model.pipeline_model_parallel_size == sft_cfg.model.pipeline_model_parallel_size assert cfg.model.expert_model_parallel_size == sft_cfg.model.expert_model_parallel_size assert cfg.model.moe_token_dispatcher_type == sft_cfg.model.moe_token_dispatcher_type - assert cfg.model.recompute_modules == sft_cfg.model.recompute_modules def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: From 943d9b0bc275f6c60592efe44facbacc9a92002a Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 24 Aug 2026 17:48:03 -0700 Subject: [PATCH 25/35] docs(deepseek): verify V4 PEFT Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 67 +++++++++++++++---- .../test_validate_card.py | 1 + 2 files changed, 55 insertions(+), 13 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 328074cb27..fad8540d5c 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -29,6 +29,11 @@ summary: > and passed independent strict GPU-only HF reloads. The public post-SFT workflow also produced identical bounded greedy completions from two independent reloads. + The GB200 packed PEFT recipe applies rank-32 LoRA to the DeepSeek MLA projections + and to separate routed/shared expert FC1/FC2 projections. It completed 100 finite + steps, saved an adapter-only checkpoint, and reloaded that checkpoint in a fresh + process for finite step 101. + verification_index: model_level: verified: @@ -42,8 +47,8 @@ verification_index: - megatron_to_hf_cpu training: GB200: - verified: [pretrain, sft, sft_export_inference] - unverified: [sft_long_context, peft, checkpoint_resume] + verified: [pretrain, sft, sft_export_inference, peft] + unverified: [sft_long_context, checkpoint_resume] performance: GB200: unverified @@ -304,23 +309,59 @@ items: peft: GB200: - status: unverified + status: verified precision: bf16 - command: null - last_verified: null + bridge_commit: 1a67d67f01e39eb2399fe59fb1f24b7f241a787c # pragma: allowlist secret + command: > + ./scripts/training/train.sh --nodes 8 --gpus-per-node 4 + --recipe deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config + --mode lora --step-func dsv4_step + --pretrained_checkpoint work/models/deepseek-v4-flash-megatron + --save_dir work/model-verification/dsv4-flash/peft-ref + --save_interval 100 --max_steps 100 + --seq_length 1024 + --pipeline_model_parallel_size 4 --context_parallel_size 1 + --expert_model_parallel_size 8 + 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + rng.seed=5678 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + checkpoint.load_optim=false checkpoint.load_rng=false + checkpoint.save_optim=false + ddp.overlap_grad_reduce=false + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=180 + last_verified: 2026-08-24 enabled_features: sequence_packing: offline moe_dispatcher: hybridep metrics: - initial_loss: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null - last_10_steps_tokens_per_second_per_gpu_avg: null + initial_loss: 0.9435847 + final_loss: 0.2915190 + last_10_steps_step_time_ms_avg: 9903.50 + last_10_steps_model_tflops_per_gpu_avg: 38.13 + last_10_steps_tokens_per_second_per_gpu_avg: 413.591155 + peak_allocated_memory_gib: 53.712 + peak_reserved_memory_gib: 53.979 expected_result: > - The packed OpenMath LoRA recipe must complete 100 finite steps with - separate routed-expert adapters, zero skipped or NaN iterations, all five - metrics, and a reloadable adapter checkpoint. + On 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert + DP1, GBS/MBS 128/1. The recipe applies rank-32, alpha-32 LoRA with zero + dropout to linear_q_down_proj, linear_q_up_proj, linear_kv_proj, + linear_proj, linear_fc1, and linear_fc2. Routed experts use separate + per-expert adapters; shared experts are also adapted. The run preserves + the packed OpenMath thinking data, objective, natural routing, schedule, + and seed used by the verified SFT item while freezing the base model. + It disables activation recompute and offload because the reduced PEFT + training-state footprint fits with 53.712 GiB peak allocated memory. All + 100 steps completed with finite LM/MTP losses and zero skipped or NaN + iterations. Initial LM/MTP losses exactly matched SFT; final LM and MTP + losses were 2.84% and 2.81% above SFT, respectively. The 9.9 GiB + adapter-only checkpoint reloaded after the same base checkpoint in a + fresh process and completed finite step 101 with LM loss 0.2829354 and + MTP loss 0.02593753. checkpoint_resume: GB200: diff --git a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py index 0fc21796b7..0f237afed5 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py @@ -23,6 +23,7 @@ ("deepseek-v3", "pretrain_performance", "GB300"): (4096, 4096, 256), ("deepseek-v4-flash", "pretrain", "GB200"): (4096, 256, 128), ("deepseek-v4-flash", "sft", "GB200"): (1024, 128, 32), + ("deepseek-v4-flash", "peft", "GB200"): (1024, 128, 32), ("gemma-4-26b-a4b-it", "sft", "H100"): (4096, 32, 8), ("gemma-4-26b-a4b-it", "peft", "H100"): (4096, 32, 4), ("glm5-2", "pretrain", "H100"): (2048, 1024, 352), From db31f50a1f858102426e63bb56e79eae6c8eb5c3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 24 Aug 2026 17:56:48 -0700 Subject: [PATCH 26/35] fix(recipes): disable recompute for portable V4 PEFT Signed-off-by: Chen Cui --- .../model_verification_cards/deepseek-v4-flash/card.yaml | 9 +++++---- src/megatron/bridge/recipes/deepseek/deepseek_v4.py | 5 +++++ .../bridge/recipes/deepseek/gb200/deepseek_v4.py | 7 +------ tests/unit_tests/recipes/test_deepseek_recipes.py | 9 +++++++++ 4 files changed, 20 insertions(+), 10 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index fad8540d5c..98cec2c70c 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -352,13 +352,14 @@ items: dropout to linear_q_down_proj, linear_q_up_proj, linear_kv_proj, linear_proj, linear_fc1, and linear_fc2. Routed experts use separate per-expert adapters; shared experts are also adapted. The run preserves - the packed OpenMath thinking data, objective, natural routing, schedule, - and seed used by the verified SFT item while freezing the base model. + the packed OpenMath thinking data, objective, natural routing, seed, and + warmup/decay horizon used by the verified SFT item while freezing the base + model and using PEFT's own 1e-4 peak learning rate. It disables activation recompute and offload because the reduced PEFT training-state footprint fits with 53.712 GiB peak allocated memory. All 100 steps completed with finite LM/MTP losses and zero skipped or NaN - iterations. Initial LM/MTP losses exactly matched SFT; final LM and MTP - losses were 2.84% and 2.81% above SFT, respectively. The 9.9 GiB + iterations. For context, initial LM/MTP losses exactly matched SFT; final + LM and MTP losses were 2.84% and 2.81% above SFT, respectively. The 9.9 GiB adapter-only checkpoint reloaded after the same base checkpoint in a fresh process and completed finite step 101 with LM loss 0.2829354 and MTP loss 0.02593753. diff --git a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py index 603e60e33a..09f627521f 100644 --- a/src/megatron/bridge/recipes/deepseek/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/deepseek_v4.py @@ -108,6 +108,11 @@ def _apply_deepseek_v4_lora(cfg: ConfigContainer) -> None: ) cfg.optimizer.lr = 1.0e-4 cfg.optimizer.min_lr = 0.0 + # MCore's unified recompute path cannot replay the frozen grouped-expert graph. + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None def deepseek_v4_flash_sft_openmath_thinking_packed_config() -> ConfigContainer: diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 86091dd141..279909a4dd 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -346,12 +346,7 @@ def deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config() -> ConfigCont cfg = deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() _apply_deepseek_v4_lora(cfg) - # PEFT's smaller training-state footprint lets this adapter graph avoid MCore's - # unified recompute path, which cannot replay the frozen grouped-expert graph. - cfg.model.recompute_granularity = None - cfg.model.recompute_modules = None - cfg.model.recompute_method = None - cfg.model.recompute_num_layers = None + # PEFT's smaller training-state footprint also lets this graph avoid activation offload. cfg.model.fine_grained_activation_offloading = False cfg.model.offload_modules = None cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) diff --git a/tests/unit_tests/recipes/test_deepseek_recipes.py b/tests/unit_tests/recipes/test_deepseek_recipes.py index f1b55e424a..99a0f770cb 100644 --- a/tests/unit_tests/recipes/test_deepseek_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_recipes.py @@ -386,6 +386,15 @@ def _build_deepseek_v4_recipe(name: str, monkeypatch: pytest.MonkeyPatch): return getattr(mod, name)() +def test_deepseek_v4_portable_peft_recipe_disables_recompute(monkeypatch: pytest.MonkeyPatch): + cfg = _build_deepseek_v4_recipe("deepseek_v4_flash_peft_openmath_thinking_packed_config", monkeypatch) + + assert cfg.model.recompute_granularity is None + assert cfg.model.recompute_modules is None + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None + + def test_deepseek_v4_adam_mxfp8_recipe_uses_validated_optimizer_defaults(monkeypatch: pytest.MonkeyPatch): cfg = _build_deepseek_v4_recipe("deepseek_v4_flash_pretrain_mxfp8_config", monkeypatch) From 043d9cdaa1ac0cb52437d8d417a9f3150dfcae3b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 25 Aug 2026 00:20:51 -0700 Subject: [PATCH 27/35] perf(recipes): optimize DeepSeek V4 PEFT Signed-off-by: Chen Cui --- .../bridge/recipes/deepseek/gb200/deepseek_v4.py | 15 +++++++++++++++ .../recipes/test_deepseek_v4_hardware_recipes.py | 15 ++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 279909a4dd..2c076bcbeb 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -36,6 +36,7 @@ from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.mixed_precision import bf16_mixed, bf16_with_mxfp8_mixed +from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: @@ -350,4 +351,18 @@ def deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config() -> ConfigCont cfg.model.fine_grained_activation_offloading = False cfg.model.offload_modules = None cfg.env_vars.pop("NVTE_CPU_OFFLOAD_V1", None) + + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.cuda_graph_impl = "transformer_engine" + set_cuda_graph_modules(cfg.model, ["moe_router", "moe_preprocess"]) + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.use_te_rng_tracker = True + cfg.rng.te_rng_tracker = True + cfg.env_vars = { + **cfg.env_vars, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + } return cfg diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index b3230732ff..78c3cf5871 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -39,6 +39,7 @@ deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config as pro_fp8_config, ) from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.utils.cuda_graph import cuda_graph_module_names from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies @@ -165,12 +166,24 @@ def test_flash_packed_peft_recipe_adapts_mla_and_unshared_experts() -> None: assert cfg.model.recompute_num_layers is None assert cfg.model.fine_grained_activation_offloading is False assert cfg.model.offload_modules is None + assert cfg.model.moe_flex_dispatcher_num_sms == 32 + assert cfg.model.cuda_graph_impl == "transformer_engine" + assert cuda_graph_module_names(cfg.model) == ["moe_router", "moe_preprocess"] + 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.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 8 + 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 assert "NVTE_CPU_OFFLOAD_V1" not in cfg.env_vars sft_cfg = flash_packed_sft_config() assert cfg.dataset == sft_cfg.dataset assert cfg.train == sft_cfg.train - assert cfg.rng == sft_cfg.rng + assert cfg.rng.seed == sft_cfg.rng.seed + assert cfg.rng.inference_rng_tracker == sft_cfg.rng.inference_rng_tracker + assert cfg.rng.data_parallel_random_init == sft_cfg.rng.data_parallel_random_init assert cfg.model.seq_length == sft_cfg.model.seq_length assert cfg.model.pipeline_model_parallel_size == sft_cfg.model.pipeline_model_parallel_size assert cfg.model.expert_model_parallel_size == sft_cfg.model.expert_model_parallel_size From 27ecd139b052be1bd3ca8bbfccaed3cf569dba20 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 25 Aug 2026 01:48:04 -0700 Subject: [PATCH 28/35] docs(verification): record optimized DeepSeek V4 PEFT Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 40 +++++++++++-------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 98cec2c70c..d7fa825e2a 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -30,9 +30,10 @@ summary: > also produced identical bounded greedy completions from two independent reloads. The GB200 packed PEFT recipe applies rank-32 LoRA to the DeepSeek MLA projections - and to separate routed/shared expert FC1/FC2 projections. It completed 100 finite - steps, saved an adapter-only checkpoint, and reloaded that checkpoint in a fresh - process for finite step 101. + and to separate routed/shared expert FC1/FC2 projections. Its execution path uses + scoped Transformer Engine CUDA graphs and tuned HybridEP dispatch. It completed + 100 finite steps, saved an adapter-only checkpoint, and reloaded that checkpoint + in a fresh process for finite step 101. verification_index: model_level: @@ -311,7 +312,7 @@ items: GB200: status: verified precision: bf16 - bridge_commit: 1a67d67f01e39eb2399fe59fb1f24b7f241a787c # pragma: allowlist secret + bridge_commit: 043d9cdaa1ac0cb52437d8d417a9f3150dfcae3b # pragma: allowlist secret command: > ./scripts/training/train.sh --nodes 8 --gpus-per-node 4 --recipe deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config @@ -334,18 +335,21 @@ items: ddp.overlap_grad_reduce=false logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=180 - last_verified: 2026-08-24 + last_verified: 2026-08-25 enabled_features: sequence_packing: offline + cuda_graph: + implementation: transformer_engine + scopes: [moe_router, moe_preprocess] moe_dispatcher: hybridep metrics: initial_loss: 0.9435847 - final_loss: 0.2915190 - last_10_steps_step_time_ms_avg: 9903.50 - last_10_steps_model_tflops_per_gpu_avg: 38.13 - last_10_steps_tokens_per_second_per_gpu_avg: 413.591155 + final_loss: 0.2914867 + last_10_steps_step_time_ms_avg: 6944.33 + last_10_steps_model_tflops_per_gpu_avg: 52.00 + last_10_steps_tokens_per_second_per_gpu_avg: 589.833720 peak_allocated_memory_gib: 53.712 - peak_reserved_memory_gib: 53.979 + peak_reserved_memory_gib: 53.953 expected_result: > On 32 GB200 (8 nodes x 4 GPUs), TP1/PP4/EP8/CP1 with dense DP8 and expert DP1, GBS/MBS 128/1. The recipe applies rank-32, alpha-32 LoRA with zero @@ -356,13 +360,15 @@ items: warmup/decay horizon used by the verified SFT item while freezing the base model and using PEFT's own 1e-4 peak learning rate. It disables activation recompute and offload because the reduced PEFT - training-state footprint fits with 53.712 GiB peak allocated memory. All - 100 steps completed with finite LM/MTP losses and zero skipped or NaN - iterations. For context, initial LM/MTP losses exactly matched SFT; final - LM and MTP losses were 2.84% and 2.81% above SFT, respectively. The 9.9 GiB - adapter-only checkpoint reloaded after the same base checkpoint in a - fresh process and completed finite step 101 with LM loss 0.2829354 and - MTP loss 0.02593753. + training-state footprint fits with 53.712 GiB peak allocated memory. + Transformer Engine CUDA graphs cover moe_router and moe_preprocess, while + HybridEP uses 32 flex-dispatcher SMs, 8 ranks per NVLink domain, 128-token + combine chunks, a 72-GPU domain, and MNNVL. All 100 steps completed with + finite LM/MTP losses and zero skipped or NaN iterations. For context, + initial LM/MTP losses exactly matched SFT; final LM and MTP losses were + 2.82% and 2.84% above SFT, respectively. The adapter-only checkpoint + reloaded after the same base checkpoint in a fresh process and completed + finite step 101 with LM loss 0.2827368 and MTP loss 0.02594211. checkpoint_resume: GB200: From 3cced766300aa65efeac090e55a4eef19c797181 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 27 Aug 2026 12:51:55 -0700 Subject: [PATCH 29/35] perf(recipes): update DeepSeek V4 GB200 validation Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 61 +++++++++++-------- examples/models/deepseek_v4/README.md | 18 +++--- scripts/training/run_recipe.py | 23 +++++-- .../references/recipe-index.md | 2 +- .../recipes/deepseek/gb200/deepseek_v4.py | 47 +++++++------- .../test_deepseek_v4_hardware_recipes.py | 29 ++++++--- .../scripts/training/test_run_recipe.py | 38 ++++++++++-- .../test_validate_card.py | 1 - 8 files changed, 140 insertions(+), 79 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index a293b11d84..68410d996b 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -9,12 +9,12 @@ summary: > performance results. DeepSeek-V4-Flash verification covers GPU import at the item-pinned Bridge revision. - The current 128-GPU GB200 MXFP8 library pretraining recipe uses - PP1/EP64 with natural routing, selective recompute, attention activation offload, and - MXFP8 parameter gather and gradient-buffer reuse. It completed 100 finite steps, - saved a grouped-MXFP8 checkpoint, then restored model, optimizer, and RNG state in a - fresh process for finite steps 101-105. Formal checkpoint-resume verification remains - separate because it requires the step-50 sentinel comparison below. A capacity-1.5 + The latest 64-GPU GB200 MXFP8 library pretraining candidate uses + PP4/VPP4/EP16 with imported model weights, natural routing, grouped GEMM, TE fused + cross entropy, and selective recompute without activation offload. It completed + 100 finite steps and saved a model-only grouped-MXFP8 checkpoint at step 100. + Formal checkpoint-resume verification remains separate because it requires the + step-50 sentinel comparison below. A capacity-1.5 performance screen averaged 206.56 TFLOP/s/GPU over steps 2-8, but dropped 66.16% of expert assignments and 64.23% of routing probability mass; that result is benchmark-only and is excluded from the natural-routing library recipe and verification metrics. CPU @@ -50,8 +50,8 @@ verification_index: - megatron_to_hf_cpu training: GB200: - verified: [pretrain, sft, sft_long_context, sft_export_inference, peft] - unverified: [checkpoint_resume] + verified: [sft, sft_long_context, sft_export_inference, peft] + unverified: [pretrain, checkpoint_resume] performance: GB200: verified GB300: verified @@ -178,41 +178,49 @@ items: pretrain: GB200: - status: verified + status: unverified precision: fp8_mx enabled_features: moe_dispatcher: hybridep command: > - ./scripts/training/train.sh --nodes 32 --gpus-per-node 4 + ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_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/rp2 dataset.num_workers=0 + --pretrained_checkpoint work/model-verification/dsv4-flash/import-gpu/iter_0000000 + 'dataset.blend=[["work/data/the-pile/my-gpt3_08_text_document"],null]' + dataset.path_to_cache=work/cache/the-pile dataset.num_workers=0 dataset.random_seed=1234 rng.seed=1234 scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 validation.eval_interval=0 validation.eval_iters=0 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null + checkpoint.load_optim=false checkpoint.load_rng=false checkpoint.save_optim=false --save_dir work/model-verification/dsv4-flash/pretrain-ref --save_interval 100 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 - last_verified: 2026-08-21 + last_verified: 2026-08-27 metrics: - initial_loss: 12.59201 - final_loss: 5.960634 - last_10_steps_step_time_ms_avg: 5713.47 - last_10_steps_model_tflops_per_gpu_avg: 131.36 - last_10_steps_tokens_per_second_per_gpu_avg: 1433.804676 + initial_loss: 7.250136 + final_loss: 3.237480 + last_10_steps_step_time_ms_avg: 8482.070 + last_10_steps_model_tflops_per_gpu_avg: 175.950 + last_10_steps_tokens_per_second_per_gpu_avg: null + peak_allocated_memory_gib: 159.93 + peak_reserved_memory_gib: 162.06 expected_result: > - The current library recipe resolves to TP1/PP1/EP64/CP1 with dense DP128, - expert DP2, and GBS/MBS 256/1. It uses HybridEP natural routing, native fused - cross entropy, selective recompute over moe, mhc, mla_up_proj, and layernorm, - attention activation offload, and MXFP8 parameter gather and gradient-buffer - reuse. Expert capacity, paged stash, and CUDA graphs remain disabled. - The run completed 100 finite LM/MTP-loss steps with no skipped or NaN - iterations and wrote a complete grouped-MXFP8 checkpoint at step 100. + On 64 GB200 (16 nodes x 4 GPUs), TP1/PP4/VPP4/EP16/CP1 uses dense DP16, + expert DP1, and GBS/MBS 256/1. The command loads the imported model weights + while starting optimizer and RNG state fresh. It uses HybridEP natural routing, + grouped GEMM, TE fused cross entropy, selective recompute over mhc and + mla_up_proj, and MXFP8 parameter gather and gradient-buffer reuse without + activation offload. Expert capacity, paged stash, CUDA graphs, and forced + load balancing remain disabled. The run completed 100 finite LM/MTP-loss + steps with no skipped or NaN iterations and wrote a complete model-only + grouped-MXFP8 checkpoint at step 100. Promote this item to verified only + after repeating the command from a clean Bridge commit that can be pinned + in this leaf. sft: GB200: @@ -261,8 +269,7 @@ items: fusion without enabling the optional DSA indexer loss. The run completed 100 steps with finite LM and MTP losses, no skipped or NaN iterations, and a complete model checkpoint at step 100. A fresh process reloaded that checkpoint - and completed finite step 101. moe_grouped_gemm=True and native fused cross - entropy are recipe defaults. + and completed finite step 101. moe_grouped_gemm=True is a recipe default. dist.distributed_timeout_minutes must be 180 or above to allow offline data packing (approximately 63 minutes) before training begins. diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index f357bbec68..3ce31bd319 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -105,23 +105,25 @@ launcher. Hardware-qualified library recipes are defined under | GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config` | BF16 Adam | | GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config` | MXFP8 Adam | | GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config` | BF16 Muon | -| GB200, 128 GPUs | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | MXFP8 Adam | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | MXFP8 Adam | | GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config` | BF16 Adam | | GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config` | MXFP8 Adam | -For example, a short generated-data run of the 128-GPU GB200 library recipe is: +For example, a short generated-data run of the GB200 library recipe is: ```bash -./scripts/training/train.sh --nodes 32 --gpus-per-node 4 \ +./scripts/training/train.sh --nodes 16 --gpus-per-node 4 \ --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config \ --mode pretrain --dataset mock --max_steps 10 ``` -That recipe owns its TP/PP/EP/CP topology, global batch size, natural-routing -semantics, recompute, offload, and checkpoint settings. Keep those defaults for -the qualified configuration and override only the dataset, run length, logging, -and output paths. Its 100-step real-data validation and checkpoint evidence are -recorded in the verification card. +The historical `128gpu` name is retained for compatibility. The recipe owns its +TP1/PP4/VPP4/EP16/CP1 topology, global batch size, natural-routing semantics, +recompute, offload, and checkpoint settings. Its current qualification is at 64 +GPUs; larger data-parallel scales require separate performance validation. Keep +the recipe defaults and override only the dataset, run length, logging, pretrained +checkpoint, and output paths. The latest 100-step real-data candidate and its +checkpoint evidence are recorded in the verification card. Compatibility aliases such as `deepseek_v4_flash_pretrain_mxfp8_config` remain exported, but new launches should use the hardware-qualified names above. diff --git a/scripts/training/run_recipe.py b/scripts/training/run_recipe.py index 654c7f0ae5..559f86de22 100755 --- a/scripts/training/run_recipe.py +++ b/scripts/training/run_recipe.py @@ -106,6 +106,7 @@ from recipe_runner import ( # noqa: E402 apply_cli_overrides, apply_determinism, + apply_runtime_environment, bootstrap_recipe_environment, load_forward_step, load_recipe, @@ -131,6 +132,9 @@ TrainMode = Literal["pretrain", "finetune"] +IMPORT_TIME_RECIPE_ENV_VARS = frozenset({"NVTE_CPU_OFFLOAD_V1"}) + + COMMON_OVERRIDE_FIELDS = ( ("max_steps", "train.train_iters"), ("global_batch_size", "train.global_batch_size"), @@ -419,6 +423,12 @@ def _selected_recipe_name(args: argparse.Namespace) -> str: return args.recipe or f"{args.model}_{recipe_task(args.mode)}_config" +def _requires_recipe_environment_bootstrap(recipe: ConfigContainer) -> bool: + """Return whether a recipe declares an unset import-time environment variable.""" + recipe_env = getattr(recipe, "env_vars", None) or {} + return any(name in recipe_env and name not in os.environ for name in IMPORT_TIME_RECIPE_ENV_VARS) + + def _load_selected_recipe(args: argparse.Namespace) -> ConfigContainer: """Load the requested recipe by its complete name or model-derived library name.""" peft_scheme = args.mode if args.mode in {"lora", "dora"} else None @@ -495,11 +505,14 @@ def main(argv: list[str] | None = None) -> None: world_size=benchmark_world_size, ) configuration_mode = _train_mode(args.mode) - recipe = bootstrap_recipe_environment( - recipe, - script_path=str(Path(__file__).resolve()), - argv=list(argv) if argv is not None else sys.argv[1:], - ) + if benchmark_metadata is not None or _requires_recipe_environment_bootstrap(recipe): + recipe = bootstrap_recipe_environment( + recipe, + script_path=str(Path(__file__).resolve()), + argv=list(argv) if argv is not None else sys.argv[1:], + ) + else: + recipe = apply_runtime_environment(recipe) if benchmark_metadata is not None: execution_mode = "pretrain" diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 5878eb74c1..77add30b7d 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,7 +74,7 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 | | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | -| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 1 | 64 | 128 (GB200; unlimited capacity) | +| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 4 | 16 | 64 (GB200; VPP4; natural unlimited capacity) | | `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | SFT | 1 | 4 | 8 | 32 (GB200; offline packed) | | `deepseek_v4_flash_peft_openmath_thinking_packed_config` | PEFT | 1 | 4 | 8 | 32 (offline packed) | | `deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config` | PEFT | 1 | 4 | 8 | 32 (GB200; offline packed) | diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index f77e70dd0c..6334a77f1f 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -41,8 +41,8 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: The core 64-GPU topology completed 1k-step convergence runs with fused DSA indexer and loss, HybridEP, grouped GEMM, permute fusion, and selective - recompute. The current router fusion, native cross entropy, and activation - offload settings were exercised by the derived 128-GPU recipe's 100-step run. + recompute. The current router fusion was also exercised by the derived + library recipe's 64-GPU 100-step run. Note: set checkpoint.save_optim=False when using HybridEP. """ @@ -80,7 +80,6 @@ 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_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_router_fusion = True @@ -88,8 +87,7 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_router_force_load_balancing = False cfg.model.moe_pad_experts_for_cuda_graph_inference = True cfg.model.cross_entropy_loss_fusion = True - # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.cross_entropy_fusion_impl = "te" cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act", "mhc", "mla_up_proj"] @@ -194,24 +192,24 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigContainer: - """Return the real-training DeepSeek V4 Flash config for 128 GB200 GPUs. + """Return the scalable real-training DeepSeek V4 Flash GB200 config. - This variant uses PP1 with dense DP128, EP64, and expert DP2. Its natural-routing - validation completed 100 finite steps, saved a grouped-MXFP8 checkpoint, and - restored model, optimizer, and RNG state through step 105. Expert capacity, - paged stash, and CUDA graphs remain disabled to preserve natural-routing - training semantics. + The historical name is retained for compatibility. This variant uses PP4/VPP4, + EP16, and selective recompute without activation offload; its current + qualification is at 64 GPUs. Expert capacity, paged stash, CUDA graphs, and + forced load balancing remain disabled to preserve natural-routing training + semantics. """ cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() cfg.model.tensor_model_parallel_size = 1 - cfg.model.pipeline_model_parallel_size = 1 - cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.virtual_pipeline_model_parallel_size = 4 cfg.model.context_parallel_size = 1 - cfg.model.expert_model_parallel_size = 64 + cfg.model.expert_model_parallel_size = 16 cfg.model.expert_tensor_parallel_size = 1 cfg.model.sequence_parallel = False - set_deepseek_v4_pipeline_model_parallel_layout(cfg.model) + cfg.model.pipeline_model_parallel_layout = "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL" cfg.train.global_batch_size = 256 cfg.train.micro_batch_size = 1 @@ -222,10 +220,15 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont cfg.model.moe_hybridep_num_sms_preprocessing = 108 cfg.model.moe_mlp_glu_interleave_size = 32 cfg.model.use_transformer_engine_op_fuser = True - cfg.model.recompute_modules = ["moe", "mhc", "mla_up_proj", "layernorm"] - cfg.model.fine_grained_activation_offloading = True - cfg.model.offload_modules = ["core_attn", "attn_proj"] - cfg.model.fine_grained_offloading_max_inflight_offloads = 2 + cfg.model.recompute_modules = ["mhc", "mla_up_proj"] + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = [] + cfg.model.fine_grained_offloading_max_inflight_offloads = None + cfg.model.moe_pad_experts_for_cuda_graph_inference = False + cfg.model.cuda_graph_impl = "none" + set_cuda_graph_modules(cfg.model, []) + cfg.model.use_te_rng_tracker = False + cfg.rng.te_rng_tracker = False cfg.mixed_precision.fp8_param_gather = True cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = True cfg.ddp.average_in_collective = False @@ -233,10 +236,11 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont cfg.env_vars = { **cfg.env_vars, "CUDA_DEVICE_MAX_CONNECTIONS": 32, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 16, "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, "NVLINK_DOMAIN_SIZE": 72, "USE_MNNVL": 1, + "NVTE_CPU_OFFLOAD_V1": 0, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, @@ -329,8 +333,7 @@ def deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config() -> ConfigConta "pad_to_max_length": True, } - # MCore warns of TE CE stability issues, and matched GB200 testing found no throughput or memory benefit. - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.cross_entropy_fusion_impl = "te" cfg.env_vars = { **cfg.env_vars, "NVTE_CPU_OFFLOAD_V1": 1, diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index 78c3cf5871..c009b9b85c 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -58,7 +58,7 @@ def test_flash_base_recipe_enables_precision_independent_fusions() -> None: assert cfg.model.moe_pad_experts_for_cuda_graph_inference is True assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None assert cfg.model.use_transformer_engine_op_fuser is False - assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.model.cross_entropy_fusion_impl == "te" assert cfg.comm_overlap.overlap_grad_reduce is True assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars @@ -131,7 +131,7 @@ def test_flash_packed_sft_recipe_uses_gb200_training_contract() -> None: assert cfg.model.moe_permute_fusion is True assert cfg.model.moe_router_fusion is True assert cfg.model.moe_grouped_gemm is True - assert cfg.model.cross_entropy_fusion_impl == "native" + assert cfg.model.cross_entropy_fusion_impl == "te" assert cfg.model.recompute_granularity == "selective" assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] assert cfg.model.recompute_method is None @@ -210,21 +210,32 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: assert cfg.optimizer.main_grads_dtype == torch.float32 assert cfg.ddp.check_for_nan_in_grad is True assert cfg.model.cuda_graph_impl == "none" + assert cuda_graph_module_names(cfg.model) == [] + assert cfg.model.use_te_rng_tracker is False + assert cfg.rng.te_rng_tracker is False assert cfg.model.use_transformer_engine_op_fuser is True assert cfg.model.moe_mlp_glu_interleave_size == 32 assert cfg.model.moe_flex_dispatcher_num_sms == 32 - assert cfg.model.pipeline_model_parallel_size == 1 - assert cfg.model.pipeline_model_parallel_layout is None - assert cfg.model.recompute_modules == ["moe", "mhc", "mla_up_proj", "layernorm"] - assert cfg.model.fine_grained_activation_offloading is True - assert cfg.model.offload_modules == ["core_attn", "attn_proj"] - assert cfg.model.fine_grained_offloading_max_inflight_offloads == 2 + assert getattr(cfg.model, "moe_deepep_num_sms", None) is None + assert getattr(cfg.model, "moe_hybridep_num_sms", None) is None + assert cfg.checkpoint.pretrained_checkpoint is None + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.virtual_pipeline_model_parallel_size == 4 + assert cfg.model.expert_model_parallel_size == 16 + assert cfg.model.pipeline_model_parallel_layout == ( + "Et*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*3|t*2|t*2|t*2|t*2|t*2mL" + ) + assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"] + assert cfg.model.fine_grained_activation_offloading is False + assert cfg.model.offload_modules == [] + assert cfg.model.fine_grained_offloading_max_inflight_offloads is None assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None assert getattr(cfg.model, "moe_paged_stash", False) is False assert cfg.env_vars["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1 + assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 16 + assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 0 def test_high_scale_deepseek_v4_recipes_are_exported() -> None: diff --git a/tests/unit_tests/scripts/training/test_run_recipe.py b/tests/unit_tests/scripts/training/test_run_recipe.py index 1f8b540e2e..93e9c882f4 100644 --- a/tests/unit_tests/scripts/training/test_run_recipe.py +++ b/tests/unit_tests/scripts/training/test_run_recipe.py @@ -475,7 +475,7 @@ def test_benchmark_finetuning_recipes_use_unified_runner(monkeypatch, mode, task assert handles.recipe_runner.run_config.call_args.kwargs["mode"] == "pretrain" -def test_library_only_canonical_name_bootstraps_recipe_environment(): +def test_library_only_canonical_name_applies_runtime_environment_without_restart(): module, handles = _load_module() handles.recipe_runner.load_recipe.return_value = SimpleNamespace() @@ -488,11 +488,38 @@ def test_library_only_canonical_name_bootstraps_recipe_environment(): ] ) + handles.recipe_runner.apply_runtime_environment.assert_called_once_with( + handles.recipe_runner.load_recipe.return_value + ) + handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() + handles.recipe_runner.load_forward_step.assert_called_once_with("llm_step", mode="pretrain") + + +def test_library_recipe_bootstraps_unset_import_time_environment(monkeypatch): + module, handles = _load_module() + monkeypatch.delenv("NVTE_CPU_OFFLOAD_V1", raising=False) + config = SimpleNamespace(env_vars={"NVTE_CPU_OFFLOAD_V1": 1}) + handles.recipe_runner.load_recipe.return_value = config + + module.main(["--recipe", "deepseek_v4_flash_pretrain_config", "--mode", "pretrain"]) + handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() bootstrap_call = handles.recipe_runner.bootstrap_recipe_environment.call_args - assert bootstrap_call.args == (handles.recipe_runner.load_recipe.return_value,) + assert bootstrap_call.args == (config,) assert bootstrap_call.kwargs["script_path"].endswith("scripts/training/run_recipe.py") - handles.recipe_runner.load_forward_step.assert_called_once_with("llm_step", mode="pretrain") + handles.recipe_runner.apply_runtime_environment.assert_not_called() + + +def test_library_recipe_preserves_explicit_import_time_environment(monkeypatch): + module, handles = _load_module() + monkeypatch.setenv("NVTE_CPU_OFFLOAD_V1", "0") + config = SimpleNamespace(env_vars={"NVTE_CPU_OFFLOAD_V1": 1}) + handles.recipe_runner.load_recipe.return_value = config + + module.main(["--recipe", "deepseek_v4_flash_pretrain_config", "--mode", "pretrain"]) + + handles.recipe_runner.apply_runtime_environment.assert_called_once_with(config) + handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() def test_benchmark_dry_run_accepts_config_overrides(monkeypatch): @@ -1447,6 +1474,5 @@ def test_config_container_overrides_are_forwarded_directly(): config, ["train.train_iters=3", "train.global_batch_size=8", "train.micro_batch_size=1"], ) - handles.recipe_runner.bootstrap_recipe_environment.assert_called_once() - bootstrap_call = handles.recipe_runner.bootstrap_recipe_environment.call_args - assert bootstrap_call.args == (config,) + handles.recipe_runner.apply_runtime_environment.assert_called_once_with(config) + handles.recipe_runner.bootstrap_recipe_environment.assert_not_called() diff --git a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py index 22e5ce7ea6..d3b3f7be43 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py @@ -21,7 +21,6 @@ ("deepseek-v3", "pretrain_performance", "H100"): (4096, 16384, 1024), ("deepseek-v3", "pretrain_performance", "GB200"): (4096, 4096, 256), ("deepseek-v3", "pretrain_performance", "GB300"): (4096, 4096, 256), - ("deepseek-v4-flash", "pretrain", "GB200"): (4096, 256, 128), ("deepseek-v4-flash", "sft", "GB200"): (1024, 128, 32), ("deepseek-v4-flash", "peft", "GB200"): (1024, 128, 32), ("deepseek-v4-flash", "sft_long_context", "GB200"): (1024, 128, 64), From dff6445009ebd054028a883ea6d84f9a16ad6b9c Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 27 Aug 2026 14:52:20 -0700 Subject: [PATCH 30/35] build(mcore): pin dev for DeepSeek V4 pretraining Signed-off-by: Chen Cui --- .dev.commit | 2 +- .../scripts/test_docker_dependency_layers.sh | 8 +- 3rdparty/Megatron-LM | 2 +- pyproject.toml | 6 +- uv.lock | 187 ++++++++---------- 5 files changed, 87 insertions(+), 118 deletions(-) diff --git a/.dev.commit b/.dev.commit index 7bcf891a9f..a3512a4504 100644 --- a/.dev.commit +++ b/.dev.commit @@ -1 +1 @@ -2f2f8ebaef90f4d62a18a43808398223b1fc5809 +8a77f4a9b78cd58dee862f234cc5cc84e47311be diff --git a/.github/scripts/test_docker_dependency_layers.sh b/.github/scripts/test_docker_dependency_layers.sh index 49384a7836..10949032b6 100755 --- a/.github/scripts/test_docker_dependency_layers.sh +++ b/.github/scripts/test_docker_dependency_layers.sh @@ -242,14 +242,14 @@ if "$temporary_dir/revision-validator" "$revision_repo" not-a-full-sha; then exit 1 fi -if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq ' "megatron-core[te]",'; then - echo "Bridge's TE extra must enable the selected MCore ref's TE extra" >&2 +if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq ' "transformer-engine[pytorch,core_cu13]",'; then + echo "Bridge's TE extra must install TransformerEngine's PyTorch and CUDA 13 packages" >&2 exit 1 fi if ! grep -Fq \ - '{ name = "megatron-core", extras = ["te"], marker = "extra == '\''te'\''", editable = "3rdparty/Megatron-LM" }' \ + '{ name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == '\''te'\''" }' \ uv.lock; then - echo "Bridge's lock metadata must preserve the MCore TE extra" >&2 + echo "Bridge's lock metadata must preserve the direct TransformerEngine extras" >&2 exit 1 fi if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml; then diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 731b791469..8a77f4a9b7 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 731b791469004f8fdcb896e65d610d6b1b0ebb32 +Subproject commit 8a77f4a9b78cd58dee862f234cc5cc84e47311be diff --git a/pyproject.toml b/pyproject.toml index 7526fb9e7c..486528fb0d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -161,8 +161,8 @@ requires-dist = [] # Lane-specific metadata for VCS packages that otherwise compile CUDA extensions during lock. [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.17.1+4329ff84" -requires-dist = ["pydantic<2.14", "importlib-metadata>=1.0", "packaging", "torch>=2.1", "einops", "onnxscript", "onnx", "nvdlfw-inspect"] +version = "2.14.0+f031cf87" +requires-dist = ["einops", "importlib-metadata", "nvdlfw-inspect", "onnx", "onnxscript", "packaging", "pydantic", "torch"] [[tool.uv.dependency-metadata]] name = "fast-hadamard-transform" @@ -185,7 +185,7 @@ tensor-inspect = [ "nvdlfw-inspect==0.2.1", ] te = [ - "megatron-core[te]", + "transformer-engine[pytorch,core_cu13]", ] ssm = [ "mamba-ssm", diff --git a/uv.lock b/uv.lock index 9a169b1307..479923e544 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = "==3.12.*" [options] @@ -40,8 +40,8 @@ version = "1.0.0+b7643bd" [[manifest.dependency-metadata]] name = "transformer-engine" -version = "2.17.1+4329ff84" -requires-dist = ["pydantic<2.14", "importlib-metadata>=1.0", "packaging", "torch>=2.1", "einops", "onnxscript", "onnx", "nvdlfw-inspect"] +version = "2.14.0+f031cf87" +requires-dist = ["einops", "importlib-metadata", "nvdlfw-inspect", "onnx", "onnxscript", "packaging", "pydantic", "torch"] [[package]] name = "absl-py" @@ -964,8 +964,8 @@ wheels = [ [[package]] name = "emerging-optimizers" -version = "0.2.0" -source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0#1effa026ff096b7fa1063ca2fba19d98be6e6cdf" } +version = "0.3.0" +source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0#b309e2f01cda75dc96a6dc1a2355a7b3b64b5e16" } dependencies = [ { name = "absl-py" }, { name = "torch", marker = "sys_platform == 'never'" }, @@ -1098,6 +1098,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/68/3c/5819fb19dc071302ca818616a4e64d4454e1f1193929eb8365c8d38e9052/flash_linear_attention-0.5.1-py3-none-any.whl", hash = "sha256:9022862f0a238752372c81290694b8b2ed1cb2d13fc40d340ffeeb554d6bee5c", size = 403096, upload-time = "2026-06-18T18:17:14.025Z" }, ] +[package.optional-dependencies] +tilelang = [ + { name = "tilelang" }, +] + [[package]] name = "flash-mla" version = "1.0.0+b7643bd" @@ -1613,14 +1618,14 @@ wheels = [ [[package]] name = "importlib-metadata" -version = "9.0.0" +version = "8.6.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "zipp" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } +sdist = { url = "https://files.pythonhosted.org/packages/33/08/c1395a292bb23fd03bdf572a1357c5a733d3eecbab877641ceacab23db6e/importlib_metadata-8.6.1.tar.gz", hash = "sha256:310b41d755445d74569f993ccfc22838295d9fe005425094fad953d7f15c8580", size = 55767, upload-time = "2025-01-20T22:21:30.429Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, + { url = "https://files.pythonhosted.org/packages/79/9d/0fb148dc4d6fa4a7dd1d8378168d9b4cd8d4560a6fbf6f0121c5fc34eb68/importlib_metadata-8.6.1-py3-none-any.whl", hash = "sha256:02a89390c1e15fdfdc0d7c6b25cb3e62650d0494005c97d6f148bf5b9787525e", size = 26971, upload-time = "2025-01-20T22:21:29.177Z" }, ] [[package]] @@ -2054,7 +2059,7 @@ ssm = [ { name = "tilelang" }, ] te = [ - { name = "megatron-core", extra = ["te"] }, + { name = "transformer-engine" }, ] tensor-inspect = [ { name = "nvdlfw-inspect" }, @@ -2122,7 +2127,6 @@ requires-dist = [ { name = "librosa", marker = "extra == 'audio'" }, { name = "mamba-ssm", marker = "extra == 'ssm'" }, { name = "megatron-core", extras = ["dev", "mlm"], editable = "3rdparty/Megatron-LM" }, - { name = "megatron-core", extras = ["te"], marker = "extra == 'te'", editable = "3rdparty/Megatron-LM" }, { name = "mistral-common", specifier = ">=1.10.0" }, { name = "mlflow", specifier = ">=3.15.1" }, { name = "nemo-run", marker = "extra == 'recipes'" }, @@ -2143,6 +2147,7 @@ requires-dist = [ { name = "timm" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", specifier = ">=4.67.1" }, + { 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" }, @@ -2205,103 +2210,107 @@ dependencies = [ [package.optional-dependencies] dev = [ { name = "av", marker = "sys_platform == 'never'" }, + { name = "causal-conv1d" }, { name = "datasets" }, { name = "einops" }, { name = "emerging-optimizers" }, - { name = "fast-hadamard-transform" }, { name = "fastapi" }, - { name = "flash-linear-attention" }, + { name = "flash-linear-attention", extra = ["tilelang"] }, { name = "flashinfer-python" }, { name = "hypercorn" }, + { name = "mamba-ssm" }, { name = "megatron-energon", extra = ["av-decode"] }, { name = "multi-storage-client" }, - { name = "nemo-lens" }, { name = "nvidia-cudnn-frontend" }, { name = "nvidia-modelopt" }, { name = "nvidia-resiliency-ext" }, + { name = "nvtx" }, { name = "onnxscript" }, { name = "openai", extra = ["aiohttp"] }, { name = "opentelemetry-api" }, - { name = "orjson" }, { name = "quart" }, { name = "tensorstore" }, { name = "tqdm" }, + { name = "transformer-engine" }, { name = "wget" }, - { name = "zstandard" }, ] mlm = [ { name = "accelerate" }, { name = "flask-restful" }, - { name = "omegaconf" }, { name = "sentencepiece" }, { name = "tiktoken" }, { name = "transformers" }, { name = "wandb" }, ] -te = [ - { name = "transformer-engine" }, -] [package.metadata] requires-dist = [ { name = "accelerate", marker = "extra == 'mlm'" }, { name = "accelerate", marker = "extra == 'training'" }, { name = "av", marker = "extra == 'dev'" }, - { name = "causal-conv1d", marker = "extra == 'ssm'", specifier = "~=1.5" }, + { name = "av", marker = "extra == 'lts'" }, + { name = "causal-conv1d", marker = "extra == 'dev'", specifier = "~=1.5" }, + { name = "causal-conv1d", marker = "extra == 'lts'", specifier = "~=1.5" }, { name = "datasets", marker = "extra == 'dev'" }, + { name = "datasets", marker = "extra == 'lts'" }, { name = "einops", marker = "extra == 'dev'", specifier = "~=0.8" }, - { name = "emerging-optimizers", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, - { name = "fast-hadamard-transform", marker = "extra == 'dev'", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, + { name = "einops", marker = "extra == 'lts'", specifier = "~=0.8" }, + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and extra == 'lts'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, - { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "==0.5.1" }, + { name = "fastapi", marker = "extra == 'lts'", specifier = "~=0.50" }, + { name = "flash-linear-attention", extras = ["tilelang"], marker = "extra == 'dev'", specifier = "==0.5.1" }, { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.7.0" }, + { name = "flashinfer-python", marker = "extra == 'lts'", specifier = ">=0.5.0,<0.7.0" }, { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, { name = "hypercorn", marker = "extra == 'dev'" }, - { name = "mamba-ssm", marker = "extra == 'ssm'", git = "https://github.com/state-spaces/mamba.git?rev=0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }, + { name = "mamba-ssm", marker = "extra == 'dev'", specifier = "~=2.2" }, + { name = "mamba-ssm", marker = "extra == 'lts'", specifier = "~=2.2" }, { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'dev'", specifier = "~=7.0" }, - { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.50" }, - { name = "nemo-lens", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, - { name = "nemo-lens", extras = ["sdk"], marker = "extra == 'otel'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, + { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'lts'", specifier = "~=6.0" }, + { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.27" }, + { name = "multi-storage-client", marker = "extra == 'lts'", specifier = "~=0.27" }, { name = "numpy" }, - { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", specifier = "==1.26.0" }, - { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'", specifier = ">=0.44" }, + { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/cudnn-frontend.git?rev=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" }, + { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'" }, { name = "nvidia-resiliency-ext", marker = "extra == 'dev'", specifier = "==0.6.0" }, - { name = "omegaconf", marker = "extra == 'mlm'" }, - { name = "omegaconf", marker = "extra == 'training'" }, + { name = "nvtx", marker = "extra == 'dev'", specifier = "~=0.2" }, + { name = "nvtx", marker = "extra == 'lts'", specifier = "~=0.2" }, { name = "onnxscript", marker = "extra == 'dev'" }, + { name = "onnxscript", marker = "extra == 'lts'" }, { name = "openai", extras = ["aiohttp"], marker = "extra == 'dev'" }, - { name = "opentelemetry-api", marker = "extra == 'dev'", specifier = ">=1.33.1,<2" }, - { name = "orjson", marker = "extra == 'dev'" }, + { name = "opentelemetry-api", marker = "extra == 'dev'", specifier = "~=1.33.1" }, + { name = "opentelemetry-api", marker = "extra == 'lts'", specifier = "~=1.33.1" }, { name = "packaging", specifier = ">=24.2" }, { name = "quart", marker = "extra == 'dev'" }, { name = "sentencepiece", marker = "extra == 'mlm'" }, { name = "sentencepiece", marker = "extra == 'training'" }, { name = "tensorstore", marker = "extra == 'dev'", specifier = "~=0.1,!=0.1.46,!=0.1.72" }, + { name = "tensorstore", marker = "extra == 'lts'", specifier = "~=0.1,!=0.1.46,!=0.1.72" }, { name = "tiktoken", marker = "extra == 'mlm'" }, { name = "tiktoken", marker = "extra == 'training'" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", marker = "extra == 'dev'" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689" }, + { name = "tqdm", marker = "extra == 'lts'" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=f031cf87bd054c7558b887df7bed93975456667f" }, { name = "transformers", marker = "extra == 'mlm'" }, { name = "transformers", marker = "extra == 'training'" }, { name = "wandb", marker = "extra == 'mlm'" }, { name = "wandb", marker = "extra == 'training'" }, { name = "wget", marker = "extra == 'dev'" }, - { name = "zstandard", marker = "extra == 'dev'" }, + { name = "wget", marker = "extra == 'lts'" }, ] -provides-extras = ["training", "mlm", "dev", "otel", "lts", "te", "ssm"] +provides-extras = ["training", "mlm", "dev", "lts"] [package.metadata.requires-dev] -batch-invariant = [{ name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }] build = [ { name = "cython", specifier = ">=3.0.0" }, { name = "hatchling" }, - { name = "nvidia-cudnn-frontend", specifier = ">=1.25.0" }, { name = "nvidia-mathdx" }, { name = "packaging", specifier = ">=24.2" }, { name = "pybind11" }, - { name = "setuptools", specifier = ">=80" }, + { name = "setuptools", specifier = ">=77.0.0,<80.0.0" }, { name = "torch" }, ] ci = [ @@ -2325,18 +2334,17 @@ linting = [ { name = "ruff", specifier = "~=0.9.0" }, ] no-pypi-wheels = [ - { name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }, - { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, + { name = "emerging-optimizers", marker = "python_full_version >= '3.12'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, + { name = "fast-hadamard-transform", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [ { name = "coverage" }, - { name = "mock" }, - { name = "nemo-run", git = "https://github.com/NVIDIA-NeMo/Run.git?rev=e3935393a290aed1822af52139b4b8ee270fed1f" }, + { name = "nemo-run", git = "https://github.com/NVIDIA-NeMo/Run.git?rev=17ae86b64d7f75653351664f5d8c9e466faede00" }, { name = "nltk" }, { name = "pydantic" }, { name = "pygithub" }, - { name = "pytest", specifier = "==9.1.1" }, + { name = "pytest", specifier = "==8.3.5" }, { name = "pytest-asyncio" }, { name = "pytest-cov" }, { name = "pytest-mock" }, @@ -2652,14 +2660,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/eb/dc/55481808fd70ef1567cf13540ffd4702af3f74b112e35427564b03f79c2d/narwhals-2.25.0-py3-none-any.whl", hash = "sha256:1f0f403e8c7e4463cde9bfe78b12fdd809e3ae3dda6d9b2f802934fb9c7a6a8f", size = 467373, upload-time = "2026-08-20T18:10:13.834Z" }, ] -[[package]] -name = "nemo-lens" -version = "0.2.0+b85578f" -source = { git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d#b85578fc2b736a1804705e537001b5f45e9c715d" } -dependencies = [ - { name = "opentelemetry-api" }, -] - [[package]] name = "nemo-run" version = "0.4.0" @@ -2934,6 +2934,18 @@ wheels = [ { url = "https://pypi.nvidia.com/nvidia-sphinx-theme/nvidia_sphinx_theme-0.0.9.post1-py3-none-any.whl", hash = "sha256:21ca60206dff2f380d7783d64bbaf71a5b9cacae53c7d0686f089c16b5a3d45a" }, ] +[[package]] +name = "nvtx" +version = "0.2.16" +source = { registry = "https://pypi.nvidia.com/" } +sdist = { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16.tar.gz", hash = "sha256:252fdb870308be4a2beecb3ed174b1d59099eeb63c883f269944a765ce5f7720" } +wheels = [ + { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a94c321fc3fe489e40f82e2d7fe7db7ef312882e85228aaaa97ee3dfadc067c" }, + { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b6230be4ee35edcfe00411c6f71d48de5d50147146c4efc06b710d2516086031" }, + { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-win_amd64.whl", hash = "sha256:916cc45efa8794d3ad4c1f91e32110eb419f2d7669af6a10f47143e9e0844855" }, + { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-win_arm64.whl", hash = "sha256:f70870ead7b6cf4227e4ae6ceba4e89a4bc3eecffbb62baca2c39a58d808df7e" }, +] + [[package]] name = "omegaconf" version = "2.3.1" @@ -3044,14 +3056,15 @@ aiohttp = [ [[package]] name = "opentelemetry-api" -version = "1.44.0" +version = "1.33.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "deprecated" }, + { name = "importlib-metadata" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ee/8b/aa9e2d8b8dfa7c946f7dec5d1f8f6ba8eca062f43509a06bdb5ce93d26c0/opentelemetry_api-1.44.0.tar.gz", hash = "sha256:67647e5e9566edcf421166fdf022b3537f818635daa852b289e34604dc6fb33a", size = 72406, upload-time = "2026-07-16T15:25:32.678Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/8d/1f5a45fbcb9a7d87809d460f09dc3399e3fbd31d7f3e14888345e9d29951/opentelemetry_api-1.33.1.tar.gz", hash = "sha256:1c6055fc0a2d3f23a50c7e17e16ef75ad489345fd3df1f8b8af7c0bbf8a109e8", size = 65002, upload-time = "2025-05-16T18:52:41.146Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ca/6f/a04e900f465ff3221ccc395522503e2d10e79fa21f2723c8e177aae1e0d1/opentelemetry_api-1.44.0-py3-none-any.whl", hash = "sha256:94b98c893a91b88657eaac1e3ba89618cdb85be6918196705354f34728b2cdef", size = 60018, upload-time = "2026-07-16T15:25:11.657Z" }, + { url = "https://files.pythonhosted.org/packages/05/44/4c45a34def3506122ae61ad684139f0bbc4e00c39555d4f7e20e0e001c8a/opentelemetry_api-1.33.1-py3-none-any.whl", hash = "sha256:4db83ebcf7ea93e64637ec6ee6fabee45c5cbe4abd9cf3da95c43828ddb50b83", size = 65771, upload-time = "2025-05-16T18:52:17.419Z" }, ] [[package]] @@ -3068,48 +3081,29 @@ wheels = [ [[package]] name = "opentelemetry-sdk" -version = "1.44.0" +version = "1.33.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-semantic-conventions" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5d/77/a6592cbc7c8d9bcc9d6757a9df45e04a7c585e3e6e7a13456da522b21109/opentelemetry_sdk-1.44.0.tar.gz", hash = "sha256:cebe7f65dc12f26ead75c6064de12fd2a9052e5060c0272d402cfa203aae123b", size = 208624, upload-time = "2026-07-16T15:25:46.078Z" } +sdist = { url = "https://files.pythonhosted.org/packages/67/12/909b98a7d9b110cce4b28d49b2e311797cffdce180371f35eba13a72dd00/opentelemetry_sdk-1.33.1.tar.gz", hash = "sha256:85b9fcf7c3d23506fbc9692fd210b8b025a1920535feec50bd54ce203d57a531", size = 161885, upload-time = "2025-05-16T18:52:52.832Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/23/ff077e61886ee020a17ce9c8b6fa11c601c8d8345b09ea24f605445df62a/opentelemetry_sdk-1.44.0-py3-none-any.whl", hash = "sha256:df081c4c6bcfdb1211e3e86140376792643128a25f8d72d1d27675936e7e96ad", size = 137221, upload-time = "2026-07-16T15:25:29.534Z" }, + { url = "https://files.pythonhosted.org/packages/df/8e/ae2d0742041e0bd7fe0d2dcc5e7cce51dcf7d3961a26072d5b43cc8fa2a7/opentelemetry_sdk-1.33.1-py3-none-any.whl", hash = "sha256:19ea73d9a01be29cacaa5d6c8ce0adc0b7f7b4d58cc52f923e4413609f670112", size = 118950, upload-time = "2025-05-16T18:52:37.297Z" }, ] [[package]] name = "opentelemetry-semantic-conventions" -version = "0.65b0" +version = "0.54b1" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "deprecated" }, { name = "opentelemetry-api" }, - { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/8f/73/0cbdebcb4cf545fdd328da14f5137e37d0770c3f26185e478b0d15d94f50/opentelemetry_semantic_conventions-0.65b0.tar.gz", hash = "sha256:f9b2b81e9d5b64f11bc952075e7e9c7fb0aab075c7fd1c46d597f1b919852d60", size = 148774, upload-time = "2026-07-16T15:25:46.902Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a6/0e/49df70d9b81fb5cbae4bbf2a49d865b09bcbcbc4eb53f5851b1027738d78/opentelemetry_semantic_conventions-0.65b0-py3-none-any.whl", hash = "sha256:1cacde7b0ad306f84c5ef08c3dbe1bbaf20165bba6f8bff43b670e555a086bcb", size = 204645, upload-time = "2026-07-16T15:25:30.688Z" }, -] - -[[package]] -name = "orjson" -version = "3.12.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/f3/742fb1f62b825f2c010697eaf4e828004bc2a81e7e806666989c132c7c42/orjson-3.12.0.tar.gz", hash = "sha256:d14203fb1aae2ad9b3d52f8a0e82aeb10197ef1c9bc61da7f358bd70b00123d5", size = 4142915, upload-time = "2026-08-14T16:13:30.607Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5b/2c/d7990fc1ffc82889d466e7cd680788ace44a26789809924813b164344393/opentelemetry_semantic_conventions-0.54b1.tar.gz", hash = "sha256:d1cecedae15d19bdaafca1e56b29a66aa286f50b5d08f036a145c7f3e9ef9cee", size = 118642, upload-time = "2025-05-16T18:52:53.962Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/be/4a/295da39c651c2faac8bd351a2a346f0fdedd9d50b847ee9dfc27d2207ef6/orjson-3.12.0-cp312-cp312-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:aa3e43a6846e91d7bde3d5a9c66090fcd8744f569a9b6cffc5e1ca38f6a461c0", size = 223427, upload-time = "2026-08-14T16:12:28.525Z" }, - { url = "https://files.pythonhosted.org/packages/29/98/758cf90fbeaaafb7f8141bfac75a432099959f3a2f5db93a412e876415d8/orjson-3.12.0-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:11edb4660a6680abee9788a3a9072208a2c96538cc1322bd79542065229d8e54", size = 123725, upload-time = "2026-08-14T16:12:30.013Z" }, - { url = "https://files.pythonhosted.org/packages/32/b5/5b934d251f8651f7e41df180ad0c57a6e1cabe15c7bd331638413a50ebc9/orjson-3.12.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.whl", hash = "sha256:2d3a9da945a4d96ae758fdaaca56742e6b73b6fd554c5d8876f252a6dad70b83", size = 113375, upload-time = "2026-08-14T16:12:31.209Z" }, - { url = "https://files.pythonhosted.org/packages/cd/d2/37efb5b12a176ce3ced29f4144f20da57d02757f78ce549637dc1b4e1fc8/orjson-3.12.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:92ffc09e07233a6ab6d4e067f7841edcbcc134cb4812155cf171ea5255a421d7", size = 129983, upload-time = "2026-08-14T16:12:32.721Z" }, - { url = "https://files.pythonhosted.org/packages/50/22/0644b87c73f13e0092df8f35a1fe280d991e5e90072087411e0dd7e44e0c/orjson-3.12.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bf44e374aadde77b1f6109f1030be51433eb61984379852766b6f4e187db7b1e", size = 130629, upload-time = "2026-08-14T16:12:34.084Z" }, - { url = "https://files.pythonhosted.org/packages/8c/57/80b986ebfecd9c6a177ddf1c2319717f0cd8feffb2b78946595a18a2fc88/orjson-3.12.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1192a7021b6d071aaf909864f6e924d6a2675ca360485b972b8401749311750b", size = 131245, upload-time = "2026-08-14T16:12:35.713Z" }, - { url = "https://files.pythonhosted.org/packages/80/3d/75c5ac5a69161f44492a68fbdde66f4cc4ce48cd5e1fb05918e46f0c8848/orjson-3.12.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:53c0c474a9d9aff9aebfc0c88de1f28f843d940e6e3a80729abdf6a20274356f", size = 135397, upload-time = "2026-08-14T16:12:37.128Z" }, - { url = "https://files.pythonhosted.org/packages/71/93/4d71f2df314a97ff0d27a4559bf5888fc8406e3c6dec90e92291e3511215/orjson-3.12.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:532ff8cd4bd59a327a953a7dcde922c7fc25b85e29721bb8633265430d3a3873", size = 127693, upload-time = "2026-08-14T16:12:38.627Z" }, - { url = "https://files.pythonhosted.org/packages/bc/1d/0dbc6be5adfd1730491072fb60beb6bcdf5d7b2596ee41b7fc2e298bfc09/orjson-3.12.0-cp312-cp312-win32.whl", hash = "sha256:a6cf4b18e7de173f209f2084ffbd736dd72389a396326ee80a7022168be232e5", size = 128000, upload-time = "2026-08-14T16:12:39.954Z" }, - { url = "https://files.pythonhosted.org/packages/2d/c9/97b1ce0112ebf5e949c775ed5b1755e562233179f3584579673cc24d6378/orjson-3.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:010811c1b69773450a01cef97727a67b223242f350b77d4ca000e59a9ef2155a", size = 122106, upload-time = "2026-08-14T16:12:41.324Z" }, - { url = "https://files.pythonhosted.org/packages/a8/6a/facd8b312e4a0d3a7fa978c7e15821f74a336adf1d65529faec33b48e18b/orjson-3.12.0-cp312-cp312-win_arm64.whl", hash = "sha256:ad29eece0c601737f2a60edc2752a84e7a0785df3efb62e3012834700a5afe0d", size = 126869, upload-time = "2026-08-14T16:12:42.651Z" }, + { url = "https://files.pythonhosted.org/packages/0a/80/08b1698c52ff76d96ba440bf15edc2f4bc0a279868778928e947c1004bdd/opentelemetry_semantic_conventions-0.54b1-py3-none-any.whl", hash = "sha256:29dab644a7e435b58d3a3918b58c333c92686236b30f7891d5e51f02933ca60d", size = 194938, upload-time = "2025-05-16T18:52:38.796Z" }, ] [[package]] @@ -4638,8 +4632,8 @@ wheels = [ [[package]] name = "transformer-engine" -version = "2.17.1+4329ff84" -source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689#4329ff84bfbdaa778a33cba02a15fb0807c64689" } +version = "2.14.0+f031cf87" +source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=f031cf87bd054c7558b887df7bed93975456667f#f031cf87bd054c7558b887df7bed93975456667f" } dependencies = [ { name = "einops" }, { name = "importlib-metadata" }, @@ -5070,28 +5064,3 @@ sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0 wheels = [ { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" }, ] - -[[package]] -name = "zstandard" -version = "0.25.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/82/fc/f26eb6ef91ae723a03e16eddb198abcfce2bc5a42e224d44cc8b6765e57e/zstandard-0.25.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7b3c3a3ab9daa3eed242d6ecceead93aebbb8f5f84318d82cee643e019c4b73b", size = 795738, upload-time = "2025-09-14T22:16:56.237Z" }, - { url = "https://files.pythonhosted.org/packages/aa/1c/d920d64b22f8dd028a8b90e2d756e431a5d86194caa78e3819c7bf53b4b3/zstandard-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:913cbd31a400febff93b564a23e17c3ed2d56c064006f54efec210d586171c00", size = 640436, upload-time = "2025-09-14T22:16:57.774Z" }, - { url = "https://files.pythonhosted.org/packages/53/6c/288c3f0bd9fcfe9ca41e2c2fbfd17b2097f6af57b62a81161941f09afa76/zstandard-0.25.0-cp312-cp312-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:011d388c76b11a0c165374ce660ce2c8efa8e5d87f34996aa80f9c0816698b64", size = 5343019, upload-time = "2025-09-14T22:16:59.302Z" }, - { url = "https://files.pythonhosted.org/packages/1e/15/efef5a2f204a64bdb5571e6161d49f7ef0fffdbca953a615efbec045f60f/zstandard-0.25.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dffecc361d079bb48d7caef5d673c88c8988d3d33fb74ab95b7ee6da42652ea", size = 5063012, upload-time = "2025-09-14T22:17:01.156Z" }, - { url = "https://files.pythonhosted.org/packages/b7/37/a6ce629ffdb43959e92e87ebdaeebb5ac81c944b6a75c9c47e300f85abdf/zstandard-0.25.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:7149623bba7fdf7e7f24312953bcf73cae103db8cae49f8154dd1eadc8a29ecb", size = 5394148, upload-time = "2025-09-14T22:17:03.091Z" }, - { url = "https://files.pythonhosted.org/packages/e3/79/2bf870b3abeb5c070fe2d670a5a8d1057a8270f125ef7676d29ea900f496/zstandard-0.25.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:6a573a35693e03cf1d67799fd01b50ff578515a8aeadd4595d2a7fa9f3ec002a", size = 5451652, upload-time = "2025-09-14T22:17:04.979Z" }, - { url = "https://files.pythonhosted.org/packages/53/60/7be26e610767316c028a2cbedb9a3beabdbe33e2182c373f71a1c0b88f36/zstandard-0.25.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a56ba0db2d244117ed744dfa8f6f5b366e14148e00de44723413b2f3938a902", size = 5546993, upload-time = "2025-09-14T22:17:06.781Z" }, - { url = "https://files.pythonhosted.org/packages/85/c7/3483ad9ff0662623f3648479b0380d2de5510abf00990468c286c6b04017/zstandard-0.25.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:10ef2a79ab8e2974e2075fb984e5b9806c64134810fac21576f0668e7ea19f8f", size = 5046806, upload-time = "2025-09-14T22:17:08.415Z" }, - { url = "https://files.pythonhosted.org/packages/08/b3/206883dd25b8d1591a1caa44b54c2aad84badccf2f1de9e2d60a446f9a25/zstandard-0.25.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aaf21ba8fb76d102b696781bddaa0954b782536446083ae3fdaa6f16b25a1c4b", size = 5576659, upload-time = "2025-09-14T22:17:10.164Z" }, - { url = "https://files.pythonhosted.org/packages/9d/31/76c0779101453e6c117b0ff22565865c54f48f8bd807df2b00c2c404b8e0/zstandard-0.25.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1869da9571d5e94a85a5e8d57e4e8807b175c9e4a6294e3b66fa4efb074d90f6", size = 4953933, upload-time = "2025-09-14T22:17:11.857Z" }, - { url = "https://files.pythonhosted.org/packages/18/e1/97680c664a1bf9a247a280a053d98e251424af51f1b196c6d52f117c9720/zstandard-0.25.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:809c5bcb2c67cd0ed81e9229d227d4ca28f82d0f778fc5fea624a9def3963f91", size = 5268008, upload-time = "2025-09-14T22:17:13.627Z" }, - { url = "https://files.pythonhosted.org/packages/1e/73/316e4010de585ac798e154e88fd81bb16afc5c5cb1a72eeb16dd37e8024a/zstandard-0.25.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:f27662e4f7dbf9f9c12391cb37b4c4c3cb90ffbd3b1fb9284dadbbb8935fa708", size = 5433517, upload-time = "2025-09-14T22:17:16.103Z" }, - { url = "https://files.pythonhosted.org/packages/5b/60/dd0f8cfa8129c5a0ce3ea6b7f70be5b33d2618013a161e1ff26c2b39787c/zstandard-0.25.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:99c0c846e6e61718715a3c9437ccc625de26593fea60189567f0118dc9db7512", size = 5814292, upload-time = "2025-09-14T22:17:17.827Z" }, - { url = "https://files.pythonhosted.org/packages/fc/5f/75aafd4b9d11b5407b641b8e41a57864097663699f23e9ad4dbb91dc6bfe/zstandard-0.25.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:474d2596a2dbc241a556e965fb76002c1ce655445e4e3bf38e5477d413165ffa", size = 5360237, upload-time = "2025-09-14T22:17:19.954Z" }, - { url = "https://files.pythonhosted.org/packages/ff/8d/0309daffea4fcac7981021dbf21cdb2e3427a9e76bafbcdbdf5392ff99a4/zstandard-0.25.0-cp312-cp312-win32.whl", hash = "sha256:23ebc8f17a03133b4426bcc04aabd68f8236eb78c3760f12783385171b0fd8bd", size = 436922, upload-time = "2025-09-14T22:17:24.398Z" }, - { url = "https://files.pythonhosted.org/packages/79/3b/fa54d9015f945330510cb5d0b0501e8253c127cca7ebe8ba46a965df18c5/zstandard-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:ffef5a74088f1e09947aecf91011136665152e0b4b359c42be3373897fb39b01", size = 506276, upload-time = "2025-09-14T22:17:21.429Z" }, - { url = "https://files.pythonhosted.org/packages/ea/6b/8b51697e5319b1f9ac71087b0af9a40d8a6288ff8025c36486e0c12abcc4/zstandard-0.25.0-cp312-cp312-win_arm64.whl", hash = "sha256:181eb40e0b6a29b3cd2849f825e0fa34397f649170673d385f3598ae17cca2e9", size = 462679, upload-time = "2025-09-14T22:17:23.147Z" }, -] From d74585b3e172a4fcca26d408ddd5335877ce4d7b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 27 Aug 2026 16:58:57 -0700 Subject: [PATCH 31/35] fix(deps): keep MCore dev pin separate Signed-off-by: Chen Cui --- .../scripts/test_docker_dependency_layers.sh | 8 +- 3rdparty/Megatron-LM | 2 +- pyproject.toml | 6 +- uv.lock | 187 ++++++++++-------- 4 files changed, 117 insertions(+), 86 deletions(-) diff --git a/.github/scripts/test_docker_dependency_layers.sh b/.github/scripts/test_docker_dependency_layers.sh index 10949032b6..49384a7836 100755 --- a/.github/scripts/test_docker_dependency_layers.sh +++ b/.github/scripts/test_docker_dependency_layers.sh @@ -242,14 +242,14 @@ if "$temporary_dir/revision-validator" "$revision_repo" not-a-full-sha; then exit 1 fi -if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq ' "transformer-engine[pytorch,core_cu13]",'; then - echo "Bridge's TE extra must install TransformerEngine's PyTorch and CUDA 13 packages" >&2 +if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq ' "megatron-core[te]",'; then + echo "Bridge's TE extra must enable the selected MCore ref's TE extra" >&2 exit 1 fi if ! grep -Fq \ - '{ name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == '\''te'\''" }' \ + '{ name = "megatron-core", extras = ["te"], marker = "extra == '\''te'\''", editable = "3rdparty/Megatron-LM" }' \ uv.lock; then - echo "Bridge's lock metadata must preserve the direct TransformerEngine extras" >&2 + echo "Bridge's lock metadata must preserve the MCore TE extra" >&2 exit 1 fi if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml; then diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 8a77f4a9b7..731b791469 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 8a77f4a9b78cd58dee862f234cc5cc84e47311be +Subproject commit 731b791469004f8fdcb896e65d610d6b1b0ebb32 diff --git a/pyproject.toml b/pyproject.toml index 486528fb0d..7526fb9e7c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -161,8 +161,8 @@ requires-dist = [] # Lane-specific metadata for VCS packages that otherwise compile CUDA extensions during lock. [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.14.0+f031cf87" -requires-dist = ["einops", "importlib-metadata", "nvdlfw-inspect", "onnx", "onnxscript", "packaging", "pydantic", "torch"] +version = "2.17.1+4329ff84" +requires-dist = ["pydantic<2.14", "importlib-metadata>=1.0", "packaging", "torch>=2.1", "einops", "onnxscript", "onnx", "nvdlfw-inspect"] [[tool.uv.dependency-metadata]] name = "fast-hadamard-transform" @@ -185,7 +185,7 @@ tensor-inspect = [ "nvdlfw-inspect==0.2.1", ] te = [ - "transformer-engine[pytorch,core_cu13]", + "megatron-core[te]", ] ssm = [ "mamba-ssm", diff --git a/uv.lock b/uv.lock index 479923e544..9a169b1307 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = "==3.12.*" [options] @@ -40,8 +40,8 @@ version = "1.0.0+b7643bd" [[manifest.dependency-metadata]] name = "transformer-engine" -version = "2.14.0+f031cf87" -requires-dist = ["einops", "importlib-metadata", "nvdlfw-inspect", "onnx", "onnxscript", "packaging", "pydantic", "torch"] +version = "2.17.1+4329ff84" +requires-dist = ["pydantic<2.14", "importlib-metadata>=1.0", "packaging", "torch>=2.1", "einops", "onnxscript", "onnx", "nvdlfw-inspect"] [[package]] name = "absl-py" @@ -964,8 +964,8 @@ wheels = [ [[package]] name = "emerging-optimizers" -version = "0.3.0" -source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0#b309e2f01cda75dc96a6dc1a2355a7b3b64b5e16" } +version = "0.2.0" +source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0#1effa026ff096b7fa1063ca2fba19d98be6e6cdf" } dependencies = [ { name = "absl-py" }, { name = "torch", marker = "sys_platform == 'never'" }, @@ -1098,11 +1098,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/68/3c/5819fb19dc071302ca818616a4e64d4454e1f1193929eb8365c8d38e9052/flash_linear_attention-0.5.1-py3-none-any.whl", hash = "sha256:9022862f0a238752372c81290694b8b2ed1cb2d13fc40d340ffeeb554d6bee5c", size = 403096, upload-time = "2026-06-18T18:17:14.025Z" }, ] -[package.optional-dependencies] -tilelang = [ - { name = "tilelang" }, -] - [[package]] name = "flash-mla" version = "1.0.0+b7643bd" @@ -1618,14 +1613,14 @@ wheels = [ [[package]] name = "importlib-metadata" -version = "8.6.1" +version = "9.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "zipp" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/33/08/c1395a292bb23fd03bdf572a1357c5a733d3eecbab877641ceacab23db6e/importlib_metadata-8.6.1.tar.gz", hash = "sha256:310b41d755445d74569f993ccfc22838295d9fe005425094fad953d7f15c8580", size = 55767, upload-time = "2025-01-20T22:21:30.429Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/79/9d/0fb148dc4d6fa4a7dd1d8378168d9b4cd8d4560a6fbf6f0121c5fc34eb68/importlib_metadata-8.6.1-py3-none-any.whl", hash = "sha256:02a89390c1e15fdfdc0d7c6b25cb3e62650d0494005c97d6f148bf5b9787525e", size = 26971, upload-time = "2025-01-20T22:21:29.177Z" }, + { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, ] [[package]] @@ -2059,7 +2054,7 @@ ssm = [ { name = "tilelang" }, ] te = [ - { name = "transformer-engine" }, + { name = "megatron-core", extra = ["te"] }, ] tensor-inspect = [ { name = "nvdlfw-inspect" }, @@ -2127,6 +2122,7 @@ requires-dist = [ { name = "librosa", marker = "extra == 'audio'" }, { name = "mamba-ssm", marker = "extra == 'ssm'" }, { name = "megatron-core", extras = ["dev", "mlm"], editable = "3rdparty/Megatron-LM" }, + { name = "megatron-core", extras = ["te"], marker = "extra == 'te'", editable = "3rdparty/Megatron-LM" }, { name = "mistral-common", specifier = ">=1.10.0" }, { name = "mlflow", specifier = ">=3.15.1" }, { name = "nemo-run", marker = "extra == 'recipes'" }, @@ -2147,7 +2143,6 @@ requires-dist = [ { name = "timm" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", specifier = ">=4.67.1" }, - { 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" }, @@ -2210,107 +2205,103 @@ dependencies = [ [package.optional-dependencies] dev = [ { name = "av", marker = "sys_platform == 'never'" }, - { name = "causal-conv1d" }, { name = "datasets" }, { name = "einops" }, { name = "emerging-optimizers" }, + { name = "fast-hadamard-transform" }, { name = "fastapi" }, - { name = "flash-linear-attention", extra = ["tilelang"] }, + { name = "flash-linear-attention" }, { name = "flashinfer-python" }, { name = "hypercorn" }, - { name = "mamba-ssm" }, { name = "megatron-energon", extra = ["av-decode"] }, { name = "multi-storage-client" }, + { name = "nemo-lens" }, { name = "nvidia-cudnn-frontend" }, { name = "nvidia-modelopt" }, { name = "nvidia-resiliency-ext" }, - { name = "nvtx" }, { name = "onnxscript" }, { name = "openai", extra = ["aiohttp"] }, { name = "opentelemetry-api" }, + { name = "orjson" }, { name = "quart" }, { name = "tensorstore" }, { name = "tqdm" }, - { name = "transformer-engine" }, { name = "wget" }, + { name = "zstandard" }, ] mlm = [ { name = "accelerate" }, { name = "flask-restful" }, + { name = "omegaconf" }, { name = "sentencepiece" }, { name = "tiktoken" }, { name = "transformers" }, { name = "wandb" }, ] +te = [ + { name = "transformer-engine" }, +] [package.metadata] requires-dist = [ { name = "accelerate", marker = "extra == 'mlm'" }, { name = "accelerate", marker = "extra == 'training'" }, { name = "av", marker = "extra == 'dev'" }, - { name = "av", marker = "extra == 'lts'" }, - { name = "causal-conv1d", marker = "extra == 'dev'", specifier = "~=1.5" }, - { name = "causal-conv1d", marker = "extra == 'lts'", specifier = "~=1.5" }, + { name = "causal-conv1d", marker = "extra == 'ssm'", specifier = "~=1.5" }, { name = "datasets", marker = "extra == 'dev'" }, - { name = "datasets", marker = "extra == 'lts'" }, { name = "einops", marker = "extra == 'dev'", specifier = "~=0.8" }, - { name = "einops", marker = "extra == 'lts'", specifier = "~=0.8" }, - { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, - { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and extra == 'lts'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, + { name = "emerging-optimizers", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, + { name = "fast-hadamard-transform", marker = "extra == 'dev'", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, - { name = "fastapi", marker = "extra == 'lts'", specifier = "~=0.50" }, - { name = "flash-linear-attention", extras = ["tilelang"], marker = "extra == 'dev'", specifier = "==0.5.1" }, + { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "==0.5.1" }, { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.7.0" }, - { name = "flashinfer-python", marker = "extra == 'lts'", specifier = ">=0.5.0,<0.7.0" }, { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, { name = "hypercorn", marker = "extra == 'dev'" }, - { name = "mamba-ssm", marker = "extra == 'dev'", specifier = "~=2.2" }, - { name = "mamba-ssm", marker = "extra == 'lts'", specifier = "~=2.2" }, + { name = "mamba-ssm", marker = "extra == 'ssm'", git = "https://github.com/state-spaces/mamba.git?rev=0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }, { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'dev'", specifier = "~=7.0" }, - { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'lts'", specifier = "~=6.0" }, - { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.27" }, - { name = "multi-storage-client", marker = "extra == 'lts'", specifier = "~=0.27" }, + { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.50" }, + { name = "nemo-lens", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, + { name = "nemo-lens", extras = ["sdk"], marker = "extra == 'otel'", git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, { name = "numpy" }, - { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/cudnn-frontend.git?rev=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" }, - { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'" }, + { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", specifier = "==1.26.0" }, + { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'", specifier = ">=0.44" }, { name = "nvidia-resiliency-ext", marker = "extra == 'dev'", specifier = "==0.6.0" }, - { name = "nvtx", marker = "extra == 'dev'", specifier = "~=0.2" }, - { name = "nvtx", marker = "extra == 'lts'", specifier = "~=0.2" }, + { name = "omegaconf", marker = "extra == 'mlm'" }, + { name = "omegaconf", marker = "extra == 'training'" }, { name = "onnxscript", marker = "extra == 'dev'" }, - { name = "onnxscript", marker = "extra == 'lts'" }, { name = "openai", extras = ["aiohttp"], marker = "extra == 'dev'" }, - { name = "opentelemetry-api", marker = "extra == 'dev'", specifier = "~=1.33.1" }, - { name = "opentelemetry-api", marker = "extra == 'lts'", specifier = "~=1.33.1" }, + { name = "opentelemetry-api", marker = "extra == 'dev'", specifier = ">=1.33.1,<2" }, + { name = "orjson", marker = "extra == 'dev'" }, { name = "packaging", specifier = ">=24.2" }, { name = "quart", marker = "extra == 'dev'" }, { name = "sentencepiece", marker = "extra == 'mlm'" }, { name = "sentencepiece", marker = "extra == 'training'" }, { name = "tensorstore", marker = "extra == 'dev'", specifier = "~=0.1,!=0.1.46,!=0.1.72" }, - { name = "tensorstore", marker = "extra == 'lts'", specifier = "~=0.1,!=0.1.46,!=0.1.72" }, { name = "tiktoken", marker = "extra == 'mlm'" }, { name = "tiktoken", marker = "extra == 'training'" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", marker = "extra == 'dev'" }, - { name = "tqdm", marker = "extra == 'lts'" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=f031cf87bd054c7558b887df7bed93975456667f" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689" }, { name = "transformers", marker = "extra == 'mlm'" }, { name = "transformers", marker = "extra == 'training'" }, { name = "wandb", marker = "extra == 'mlm'" }, { name = "wandb", marker = "extra == 'training'" }, { name = "wget", marker = "extra == 'dev'" }, - { name = "wget", marker = "extra == 'lts'" }, + { name = "zstandard", marker = "extra == 'dev'" }, ] -provides-extras = ["training", "mlm", "dev", "lts"] +provides-extras = ["training", "mlm", "dev", "otel", "lts", "te", "ssm"] [package.metadata.requires-dev] +batch-invariant = [{ name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }] build = [ { name = "cython", specifier = ">=3.0.0" }, { name = "hatchling" }, + { name = "nvidia-cudnn-frontend", specifier = ">=1.25.0" }, { name = "nvidia-mathdx" }, { name = "packaging", specifier = ">=24.2" }, { name = "pybind11" }, - { name = "setuptools", specifier = ">=77.0.0,<80.0.0" }, + { name = "setuptools", specifier = ">=80" }, { name = "torch" }, ] ci = [ @@ -2334,17 +2325,18 @@ linting = [ { name = "ruff", specifier = "~=0.9.0" }, ] no-pypi-wheels = [ - { name = "emerging-optimizers", marker = "python_full_version >= '3.12'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, - { name = "fast-hadamard-transform", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, + { name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }, + { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [ { name = "coverage" }, - { name = "nemo-run", git = "https://github.com/NVIDIA-NeMo/Run.git?rev=17ae86b64d7f75653351664f5d8c9e466faede00" }, + { name = "mock" }, + { name = "nemo-run", git = "https://github.com/NVIDIA-NeMo/Run.git?rev=e3935393a290aed1822af52139b4b8ee270fed1f" }, { name = "nltk" }, { name = "pydantic" }, { name = "pygithub" }, - { name = "pytest", specifier = "==8.3.5" }, + { name = "pytest", specifier = "==9.1.1" }, { name = "pytest-asyncio" }, { name = "pytest-cov" }, { name = "pytest-mock" }, @@ -2660,6 +2652,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/eb/dc/55481808fd70ef1567cf13540ffd4702af3f74b112e35427564b03f79c2d/narwhals-2.25.0-py3-none-any.whl", hash = "sha256:1f0f403e8c7e4463cde9bfe78b12fdd809e3ae3dda6d9b2f802934fb9c7a6a8f", size = 467373, upload-time = "2026-08-20T18:10:13.834Z" }, ] +[[package]] +name = "nemo-lens" +version = "0.2.0+b85578f" +source = { git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d#b85578fc2b736a1804705e537001b5f45e9c715d" } +dependencies = [ + { name = "opentelemetry-api" }, +] + [[package]] name = "nemo-run" version = "0.4.0" @@ -2934,18 +2934,6 @@ wheels = [ { url = "https://pypi.nvidia.com/nvidia-sphinx-theme/nvidia_sphinx_theme-0.0.9.post1-py3-none-any.whl", hash = "sha256:21ca60206dff2f380d7783d64bbaf71a5b9cacae53c7d0686f089c16b5a3d45a" }, ] -[[package]] -name = "nvtx" -version = "0.2.16" -source = { registry = "https://pypi.nvidia.com/" } -sdist = { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16.tar.gz", hash = "sha256:252fdb870308be4a2beecb3ed174b1d59099eeb63c883f269944a765ce5f7720" } -wheels = [ - { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a94c321fc3fe489e40f82e2d7fe7db7ef312882e85228aaaa97ee3dfadc067c" }, - { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b6230be4ee35edcfe00411c6f71d48de5d50147146c4efc06b710d2516086031" }, - { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-win_amd64.whl", hash = "sha256:916cc45efa8794d3ad4c1f91e32110eb419f2d7669af6a10f47143e9e0844855" }, - { url = "https://pypi.nvidia.com/nvtx/nvtx-0.2.16-cp312-cp312-win_arm64.whl", hash = "sha256:f70870ead7b6cf4227e4ae6ceba4e89a4bc3eecffbb62baca2c39a58d808df7e" }, -] - [[package]] name = "omegaconf" version = "2.3.1" @@ -3056,15 +3044,14 @@ aiohttp = [ [[package]] name = "opentelemetry-api" -version = "1.33.1" +version = "1.44.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "deprecated" }, - { name = "importlib-metadata" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9a/8d/1f5a45fbcb9a7d87809d460f09dc3399e3fbd31d7f3e14888345e9d29951/opentelemetry_api-1.33.1.tar.gz", hash = "sha256:1c6055fc0a2d3f23a50c7e17e16ef75ad489345fd3df1f8b8af7c0bbf8a109e8", size = 65002, upload-time = "2025-05-16T18:52:41.146Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/8b/aa9e2d8b8dfa7c946f7dec5d1f8f6ba8eca062f43509a06bdb5ce93d26c0/opentelemetry_api-1.44.0.tar.gz", hash = "sha256:67647e5e9566edcf421166fdf022b3537f818635daa852b289e34604dc6fb33a", size = 72406, upload-time = "2026-07-16T15:25:32.678Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/05/44/4c45a34def3506122ae61ad684139f0bbc4e00c39555d4f7e20e0e001c8a/opentelemetry_api-1.33.1-py3-none-any.whl", hash = "sha256:4db83ebcf7ea93e64637ec6ee6fabee45c5cbe4abd9cf3da95c43828ddb50b83", size = 65771, upload-time = "2025-05-16T18:52:17.419Z" }, + { url = "https://files.pythonhosted.org/packages/ca/6f/a04e900f465ff3221ccc395522503e2d10e79fa21f2723c8e177aae1e0d1/opentelemetry_api-1.44.0-py3-none-any.whl", hash = "sha256:94b98c893a91b88657eaac1e3ba89618cdb85be6918196705354f34728b2cdef", size = 60018, upload-time = "2026-07-16T15:25:11.657Z" }, ] [[package]] @@ -3081,29 +3068,48 @@ wheels = [ [[package]] name = "opentelemetry-sdk" -version = "1.33.1" +version = "1.44.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-semantic-conventions" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/67/12/909b98a7d9b110cce4b28d49b2e311797cffdce180371f35eba13a72dd00/opentelemetry_sdk-1.33.1.tar.gz", hash = "sha256:85b9fcf7c3d23506fbc9692fd210b8b025a1920535feec50bd54ce203d57a531", size = 161885, upload-time = "2025-05-16T18:52:52.832Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5d/77/a6592cbc7c8d9bcc9d6757a9df45e04a7c585e3e6e7a13456da522b21109/opentelemetry_sdk-1.44.0.tar.gz", hash = "sha256:cebe7f65dc12f26ead75c6064de12fd2a9052e5060c0272d402cfa203aae123b", size = 208624, upload-time = "2026-07-16T15:25:46.078Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/df/8e/ae2d0742041e0bd7fe0d2dcc5e7cce51dcf7d3961a26072d5b43cc8fa2a7/opentelemetry_sdk-1.33.1-py3-none-any.whl", hash = "sha256:19ea73d9a01be29cacaa5d6c8ce0adc0b7f7b4d58cc52f923e4413609f670112", size = 118950, upload-time = "2025-05-16T18:52:37.297Z" }, + { url = "https://files.pythonhosted.org/packages/e7/23/ff077e61886ee020a17ce9c8b6fa11c601c8d8345b09ea24f605445df62a/opentelemetry_sdk-1.44.0-py3-none-any.whl", hash = "sha256:df081c4c6bcfdb1211e3e86140376792643128a25f8d72d1d27675936e7e96ad", size = 137221, upload-time = "2026-07-16T15:25:29.534Z" }, ] [[package]] name = "opentelemetry-semantic-conventions" -version = "0.54b1" +version = "0.65b0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "deprecated" }, { name = "opentelemetry-api" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5b/2c/d7990fc1ffc82889d466e7cd680788ace44a26789809924813b164344393/opentelemetry_semantic_conventions-0.54b1.tar.gz", hash = "sha256:d1cecedae15d19bdaafca1e56b29a66aa286f50b5d08f036a145c7f3e9ef9cee", size = 118642, upload-time = "2025-05-16T18:52:53.962Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8f/73/0cbdebcb4cf545fdd328da14f5137e37d0770c3f26185e478b0d15d94f50/opentelemetry_semantic_conventions-0.65b0.tar.gz", hash = "sha256:f9b2b81e9d5b64f11bc952075e7e9c7fb0aab075c7fd1c46d597f1b919852d60", size = 148774, upload-time = "2026-07-16T15:25:46.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/0e/49df70d9b81fb5cbae4bbf2a49d865b09bcbcbc4eb53f5851b1027738d78/opentelemetry_semantic_conventions-0.65b0-py3-none-any.whl", hash = "sha256:1cacde7b0ad306f84c5ef08c3dbe1bbaf20165bba6f8bff43b670e555a086bcb", size = 204645, upload-time = "2026-07-16T15:25:30.688Z" }, +] + +[[package]] +name = "orjson" +version = "3.12.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/f3/742fb1f62b825f2c010697eaf4e828004bc2a81e7e806666989c132c7c42/orjson-3.12.0.tar.gz", hash = "sha256:d14203fb1aae2ad9b3d52f8a0e82aeb10197ef1c9bc61da7f358bd70b00123d5", size = 4142915, upload-time = "2026-08-14T16:13:30.607Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0a/80/08b1698c52ff76d96ba440bf15edc2f4bc0a279868778928e947c1004bdd/opentelemetry_semantic_conventions-0.54b1-py3-none-any.whl", hash = "sha256:29dab644a7e435b58d3a3918b58c333c92686236b30f7891d5e51f02933ca60d", size = 194938, upload-time = "2025-05-16T18:52:38.796Z" }, + { url = "https://files.pythonhosted.org/packages/be/4a/295da39c651c2faac8bd351a2a346f0fdedd9d50b847ee9dfc27d2207ef6/orjson-3.12.0-cp312-cp312-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:aa3e43a6846e91d7bde3d5a9c66090fcd8744f569a9b6cffc5e1ca38f6a461c0", size = 223427, upload-time = "2026-08-14T16:12:28.525Z" }, + { url = "https://files.pythonhosted.org/packages/29/98/758cf90fbeaaafb7f8141bfac75a432099959f3a2f5db93a412e876415d8/orjson-3.12.0-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:11edb4660a6680abee9788a3a9072208a2c96538cc1322bd79542065229d8e54", size = 123725, upload-time = "2026-08-14T16:12:30.013Z" }, + { url = "https://files.pythonhosted.org/packages/32/b5/5b934d251f8651f7e41df180ad0c57a6e1cabe15c7bd331638413a50ebc9/orjson-3.12.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.whl", hash = "sha256:2d3a9da945a4d96ae758fdaaca56742e6b73b6fd554c5d8876f252a6dad70b83", size = 113375, upload-time = "2026-08-14T16:12:31.209Z" }, + { url = "https://files.pythonhosted.org/packages/cd/d2/37efb5b12a176ce3ced29f4144f20da57d02757f78ce549637dc1b4e1fc8/orjson-3.12.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:92ffc09e07233a6ab6d4e067f7841edcbcc134cb4812155cf171ea5255a421d7", size = 129983, upload-time = "2026-08-14T16:12:32.721Z" }, + { url = "https://files.pythonhosted.org/packages/50/22/0644b87c73f13e0092df8f35a1fe280d991e5e90072087411e0dd7e44e0c/orjson-3.12.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bf44e374aadde77b1f6109f1030be51433eb61984379852766b6f4e187db7b1e", size = 130629, upload-time = "2026-08-14T16:12:34.084Z" }, + { url = "https://files.pythonhosted.org/packages/8c/57/80b986ebfecd9c6a177ddf1c2319717f0cd8feffb2b78946595a18a2fc88/orjson-3.12.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1192a7021b6d071aaf909864f6e924d6a2675ca360485b972b8401749311750b", size = 131245, upload-time = "2026-08-14T16:12:35.713Z" }, + { url = "https://files.pythonhosted.org/packages/80/3d/75c5ac5a69161f44492a68fbdde66f4cc4ce48cd5e1fb05918e46f0c8848/orjson-3.12.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:53c0c474a9d9aff9aebfc0c88de1f28f843d940e6e3a80729abdf6a20274356f", size = 135397, upload-time = "2026-08-14T16:12:37.128Z" }, + { url = "https://files.pythonhosted.org/packages/71/93/4d71f2df314a97ff0d27a4559bf5888fc8406e3c6dec90e92291e3511215/orjson-3.12.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:532ff8cd4bd59a327a953a7dcde922c7fc25b85e29721bb8633265430d3a3873", size = 127693, upload-time = "2026-08-14T16:12:38.627Z" }, + { url = "https://files.pythonhosted.org/packages/bc/1d/0dbc6be5adfd1730491072fb60beb6bcdf5d7b2596ee41b7fc2e298bfc09/orjson-3.12.0-cp312-cp312-win32.whl", hash = "sha256:a6cf4b18e7de173f209f2084ffbd736dd72389a396326ee80a7022168be232e5", size = 128000, upload-time = "2026-08-14T16:12:39.954Z" }, + { url = "https://files.pythonhosted.org/packages/2d/c9/97b1ce0112ebf5e949c775ed5b1755e562233179f3584579673cc24d6378/orjson-3.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:010811c1b69773450a01cef97727a67b223242f350b77d4ca000e59a9ef2155a", size = 122106, upload-time = "2026-08-14T16:12:41.324Z" }, + { url = "https://files.pythonhosted.org/packages/a8/6a/facd8b312e4a0d3a7fa978c7e15821f74a336adf1d65529faec33b48e18b/orjson-3.12.0-cp312-cp312-win_arm64.whl", hash = "sha256:ad29eece0c601737f2a60edc2752a84e7a0785df3efb62e3012834700a5afe0d", size = 126869, upload-time = "2026-08-14T16:12:42.651Z" }, ] [[package]] @@ -4632,8 +4638,8 @@ wheels = [ [[package]] name = "transformer-engine" -version = "2.14.0+f031cf87" -source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=f031cf87bd054c7558b887df7bed93975456667f#f031cf87bd054c7558b887df7bed93975456667f" } +version = "2.17.1+4329ff84" +source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689#4329ff84bfbdaa778a33cba02a15fb0807c64689" } dependencies = [ { name = "einops" }, { name = "importlib-metadata" }, @@ -5064,3 +5070,28 @@ sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0 wheels = [ { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" }, ] + +[[package]] +name = "zstandard" +version = "0.25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/fc/f26eb6ef91ae723a03e16eddb198abcfce2bc5a42e224d44cc8b6765e57e/zstandard-0.25.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7b3c3a3ab9daa3eed242d6ecceead93aebbb8f5f84318d82cee643e019c4b73b", size = 795738, upload-time = "2025-09-14T22:16:56.237Z" }, + { url = "https://files.pythonhosted.org/packages/aa/1c/d920d64b22f8dd028a8b90e2d756e431a5d86194caa78e3819c7bf53b4b3/zstandard-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:913cbd31a400febff93b564a23e17c3ed2d56c064006f54efec210d586171c00", size = 640436, upload-time = "2025-09-14T22:16:57.774Z" }, + { url = "https://files.pythonhosted.org/packages/53/6c/288c3f0bd9fcfe9ca41e2c2fbfd17b2097f6af57b62a81161941f09afa76/zstandard-0.25.0-cp312-cp312-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:011d388c76b11a0c165374ce660ce2c8efa8e5d87f34996aa80f9c0816698b64", size = 5343019, upload-time = "2025-09-14T22:16:59.302Z" }, + { url = "https://files.pythonhosted.org/packages/1e/15/efef5a2f204a64bdb5571e6161d49f7ef0fffdbca953a615efbec045f60f/zstandard-0.25.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dffecc361d079bb48d7caef5d673c88c8988d3d33fb74ab95b7ee6da42652ea", size = 5063012, upload-time = "2025-09-14T22:17:01.156Z" }, + { url = "https://files.pythonhosted.org/packages/b7/37/a6ce629ffdb43959e92e87ebdaeebb5ac81c944b6a75c9c47e300f85abdf/zstandard-0.25.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:7149623bba7fdf7e7f24312953bcf73cae103db8cae49f8154dd1eadc8a29ecb", size = 5394148, upload-time = "2025-09-14T22:17:03.091Z" }, + { url = "https://files.pythonhosted.org/packages/e3/79/2bf870b3abeb5c070fe2d670a5a8d1057a8270f125ef7676d29ea900f496/zstandard-0.25.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:6a573a35693e03cf1d67799fd01b50ff578515a8aeadd4595d2a7fa9f3ec002a", size = 5451652, upload-time = "2025-09-14T22:17:04.979Z" }, + { url = "https://files.pythonhosted.org/packages/53/60/7be26e610767316c028a2cbedb9a3beabdbe33e2182c373f71a1c0b88f36/zstandard-0.25.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a56ba0db2d244117ed744dfa8f6f5b366e14148e00de44723413b2f3938a902", size = 5546993, upload-time = "2025-09-14T22:17:06.781Z" }, + { url = "https://files.pythonhosted.org/packages/85/c7/3483ad9ff0662623f3648479b0380d2de5510abf00990468c286c6b04017/zstandard-0.25.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:10ef2a79ab8e2974e2075fb984e5b9806c64134810fac21576f0668e7ea19f8f", size = 5046806, upload-time = "2025-09-14T22:17:08.415Z" }, + { url = "https://files.pythonhosted.org/packages/08/b3/206883dd25b8d1591a1caa44b54c2aad84badccf2f1de9e2d60a446f9a25/zstandard-0.25.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aaf21ba8fb76d102b696781bddaa0954b782536446083ae3fdaa6f16b25a1c4b", size = 5576659, upload-time = "2025-09-14T22:17:10.164Z" }, + { url = "https://files.pythonhosted.org/packages/9d/31/76c0779101453e6c117b0ff22565865c54f48f8bd807df2b00c2c404b8e0/zstandard-0.25.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1869da9571d5e94a85a5e8d57e4e8807b175c9e4a6294e3b66fa4efb074d90f6", size = 4953933, upload-time = "2025-09-14T22:17:11.857Z" }, + { url = "https://files.pythonhosted.org/packages/18/e1/97680c664a1bf9a247a280a053d98e251424af51f1b196c6d52f117c9720/zstandard-0.25.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:809c5bcb2c67cd0ed81e9229d227d4ca28f82d0f778fc5fea624a9def3963f91", size = 5268008, upload-time = "2025-09-14T22:17:13.627Z" }, + { url = "https://files.pythonhosted.org/packages/1e/73/316e4010de585ac798e154e88fd81bb16afc5c5cb1a72eeb16dd37e8024a/zstandard-0.25.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:f27662e4f7dbf9f9c12391cb37b4c4c3cb90ffbd3b1fb9284dadbbb8935fa708", size = 5433517, upload-time = "2025-09-14T22:17:16.103Z" }, + { url = "https://files.pythonhosted.org/packages/5b/60/dd0f8cfa8129c5a0ce3ea6b7f70be5b33d2618013a161e1ff26c2b39787c/zstandard-0.25.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:99c0c846e6e61718715a3c9437ccc625de26593fea60189567f0118dc9db7512", size = 5814292, upload-time = "2025-09-14T22:17:17.827Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5f/75aafd4b9d11b5407b641b8e41a57864097663699f23e9ad4dbb91dc6bfe/zstandard-0.25.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:474d2596a2dbc241a556e965fb76002c1ce655445e4e3bf38e5477d413165ffa", size = 5360237, upload-time = "2025-09-14T22:17:19.954Z" }, + { url = "https://files.pythonhosted.org/packages/ff/8d/0309daffea4fcac7981021dbf21cdb2e3427a9e76bafbcdbdf5392ff99a4/zstandard-0.25.0-cp312-cp312-win32.whl", hash = "sha256:23ebc8f17a03133b4426bcc04aabd68f8236eb78c3760f12783385171b0fd8bd", size = 436922, upload-time = "2025-09-14T22:17:24.398Z" }, + { url = "https://files.pythonhosted.org/packages/79/3b/fa54d9015f945330510cb5d0b0501e8253c127cca7ebe8ba46a965df18c5/zstandard-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:ffef5a74088f1e09947aecf91011136665152e0b4b359c42be3373897fb39b01", size = 506276, upload-time = "2025-09-14T22:17:21.429Z" }, + { url = "https://files.pythonhosted.org/packages/ea/6b/8b51697e5319b1f9ac71087b0af9a40d8a6288ff8025c36486e0c12abcc4/zstandard-0.25.0-cp312-cp312-win_arm64.whl", hash = "sha256:181eb40e0b6a29b3cd2849f825e0fa34397f649170673d385f3598ae17cca2e9", size = 462679, upload-time = "2025-09-14T22:17:23.147Z" }, +] From 098656c4238fdb8ccc4516a2dd0732d2a02ba4c3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 27 Aug 2026 17:06:44 -0700 Subject: [PATCH 32/35] docs(model-card): verify DeepSeek V4 pretraining Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 19 ++++++++++--------- .../test_validate_card.py | 1 + 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 68410d996b..8d4f8a84f6 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -9,7 +9,7 @@ summary: > performance results. DeepSeek-V4-Flash verification covers GPU import at the item-pinned Bridge revision. - The latest 64-GPU GB200 MXFP8 library pretraining candidate uses + The verified 64-GPU GB200 MXFP8 library pretraining run uses PP4/VPP4/EP16 with imported model weights, natural routing, grouped GEMM, TE fused cross entropy, and selective recompute without activation offload. It completed 100 finite steps and saved a model-only grouped-MXFP8 checkpoint at step 100. @@ -50,8 +50,8 @@ verification_index: - megatron_to_hf_cpu training: GB200: - verified: [sft, sft_long_context, sft_export_inference, peft] - unverified: [pretrain, checkpoint_resume] + verified: [pretrain, sft, sft_long_context, sft_export_inference, peft] + unverified: [checkpoint_resume] performance: GB200: verified GB300: verified @@ -63,7 +63,7 @@ model: min_transformers_version: "5.12.0" verification_environment: - base_container: nvcr.io/nvidia/nemo:26.08.rc5 + base_container: nvcr.io/nvidia/nemo:26.08 bridge_commit: 07c5cc30d0510d437bfbbe7b12d9d98387295d89 # pragma: allowlist secret items: @@ -178,8 +178,9 @@ items: pretrain: GB200: - status: unverified + status: verified precision: fp8_mx + bridge_commit: dff6445009ebd054028a883ea6d84f9a16ad6b9c # pragma: allowlist secret enabled_features: moe_dispatcher: hybridep command: > @@ -206,7 +207,7 @@ items: final_loss: 3.237480 last_10_steps_step_time_ms_avg: 8482.070 last_10_steps_model_tflops_per_gpu_avg: 175.950 - last_10_steps_tokens_per_second_per_gpu_avg: null + last_10_steps_tokens_per_second_per_gpu_avg: 1931.603960 peak_allocated_memory_gib: 159.93 peak_reserved_memory_gib: 162.06 expected_result: > @@ -218,9 +219,9 @@ items: activation offload. Expert capacity, paged stash, CUDA graphs, and forced load balancing remain disabled. The run completed 100 finite LM/MTP-loss steps with no skipped or NaN iterations and wrote a complete model-only - grouped-MXFP8 checkpoint at step 100. Promote this item to verified only - after repeating the command from a clean Bridge commit that can be pinned - in this leaf. + grouped-MXFP8 checkpoint at step 100. The item-specific Bridge revision + pins the compatible Megatron-LM dev revision used by this workload without + changing the card's default environment for other verification items. sft: GB200: diff --git a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py index d3b3f7be43..fb448700ca 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py @@ -21,6 +21,7 @@ ("deepseek-v3", "pretrain_performance", "H100"): (4096, 16384, 1024), ("deepseek-v3", "pretrain_performance", "GB200"): (4096, 4096, 256), ("deepseek-v3", "pretrain_performance", "GB300"): (4096, 4096, 256), + ("deepseek-v4-flash", "pretrain", "GB200"): (4096, 256, 64), ("deepseek-v4-flash", "sft", "GB200"): (1024, 128, 32), ("deepseek-v4-flash", "peft", "GB200"): (1024, 128, 32), ("deepseek-v4-flash", "sft_long_context", "GB200"): (1024, 128, 64), From c0adc7079fce1a3b5e0888d3cb2d872e02a2f51e Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 27 Aug 2026 18:22:06 -0700 Subject: [PATCH 33/35] docs(model-card): verify DeepSeek V4 checkpoint resume Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 88 ++++++++++++------- .../test_validate_card.py | 1 + 2 files changed, 56 insertions(+), 33 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index 8d4f8a84f6..cc6a6e3aec 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -12,9 +12,9 @@ summary: > The verified 64-GPU GB200 MXFP8 library pretraining run uses PP4/VPP4/EP16 with imported model weights, natural routing, grouped GEMM, TE fused cross entropy, and selective recompute without activation offload. It completed - 100 finite steps and saved a model-only grouped-MXFP8 checkpoint at step 100. - Formal checkpoint-resume verification remains separate because it requires the - step-50 sentinel comparison below. A capacity-1.5 + 100 finite steps with full-state checkpoints at steps 50 and 100. A direct + step-50 continuation restored optimizer, scheduler, data-order, and RNG state, + completed steps 51-100, and passed both loss-sentinel comparisons. A capacity-1.5 performance screen averaged 206.56 TFLOP/s/GPU over steps 2-8, but dropped 66.16% of expert assignments and 64.23% of routing probability mass; that result is benchmark-only and is excluded from the natural-routing library recipe and verification metrics. CPU @@ -50,8 +50,7 @@ verification_index: - megatron_to_hf_cpu training: GB200: - verified: [pretrain, sft, sft_long_context, sft_export_inference, peft] - unverified: [checkpoint_resume] + verified: [pretrain, sft, sft_long_context, sft_export_inference, peft, checkpoint_resume] performance: GB200: verified GB300: verified @@ -195,21 +194,20 @@ items: validation.eval_interval=0 validation.eval_iters=0 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true - checkpoint.exit_on_missing_checkpoint=false checkpoint.load=null - checkpoint.load_optim=false checkpoint.load_rng=false checkpoint.save_optim=false + checkpoint.exit_on_missing_checkpoint=false checkpoint.finetune=false + checkpoint.load=null checkpoint.load_optim=false checkpoint.load_rng=false + checkpoint.save_optim=true checkpoint.save_rng=true checkpoint.async_save=false --save_dir work/model-verification/dsv4-flash/pretrain-ref - --save_interval 100 + --save_interval 50 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null dist.distributed_timeout_minutes=120 last_verified: 2026-08-27 metrics: - initial_loss: 7.250136 - final_loss: 3.237480 - last_10_steps_step_time_ms_avg: 8482.070 - last_10_steps_model_tflops_per_gpu_avg: 175.950 - last_10_steps_tokens_per_second_per_gpu_avg: 1931.603960 - peak_allocated_memory_gib: 159.93 - peak_reserved_memory_gib: 162.06 + initial_loss: 7.250204 + final_loss: 3.280722 + last_10_steps_step_time_ms_avg: 7856.370 + last_10_steps_model_tflops_per_gpu_avg: 189.690 + last_10_steps_tokens_per_second_per_gpu_avg: 2085.441495 expected_result: > On 64 GB200 (16 nodes x 4 GPUs), TP1/PP4/VPP4/EP16/CP1 uses dense DP16, expert DP1, and GBS/MBS 256/1. The command loads the imported model weights @@ -218,10 +216,11 @@ items: mla_up_proj, and MXFP8 parameter gather and gradient-buffer reuse without activation offload. Expert capacity, paged stash, CUDA graphs, and forced load balancing remain disabled. The run completed 100 finite LM/MTP-loss - steps with no skipped or NaN iterations and wrote a complete model-only - grouped-MXFP8 checkpoint at step 100. The item-specific Bridge revision - pins the compatible Megatron-LM dev revision used by this workload without - changing the card's default environment for other verification items. + steps with no skipped or NaN iterations and wrote complete grouped-MXFP8 + checkpoints containing model, optimizer, scheduler, data-order, and RNG + state at steps 50 and 100. The item-specific Bridge revision pins the + compatible Megatron-LM dev revision used by this workload without changing + the card's default environment for other verification items. sft: GB200: @@ -419,29 +418,52 @@ items: checkpoint_resume: GB200: - status: unverified + status: verified precision: fp8_mx + bridge_commit: dff6445009ebd054028a883ea6d84f9a16ad6b9c # pragma: allowlist secret depends_on: pretrain - command: null - last_verified: null + command: > + ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 + --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config + --mode pretrain --dataset megatron-indexed --max_steps 100 + 'dataset.blend=[["work/data/the-pile/my-gpt3_08_text_document"],null]' + dataset.path_to_cache=work/cache/the-pile dataset.num_workers=0 + dataset.random_seed=1234 rng.seed=1234 + scheduler.lr_warmup_iters=10 scheduler.lr_decay_iters=100 + validation.eval_interval=0 validation.eval_iters=0 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + checkpoint.exit_on_missing_checkpoint=false checkpoint.finetune=false + checkpoint.load_optim=true checkpoint.load_rng=true + checkpoint.save_optim=true checkpoint.save_rng=true checkpoint.async_save=false + --load_dir work/model-verification/dsv4-flash/pretrain-ref + --save_dir work/model-verification/dsv4-flash/pretrain-resumed + --save_interval 50 checkpoint.ckpt_step=50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + dist.distributed_timeout_minutes=120 + last_verified: 2026-08-27 metrics: - initial_loss: null - final_loss: null - last_10_steps_step_time_ms_avg: null - last_10_steps_model_tflops_per_gpu_avg: null - last_10_steps_tokens_per_second_per_gpu_avg: null + initial_loss: 4.028752 + final_loss: 3.261145 + last_10_steps_step_time_ms_avg: 7995.850 + last_10_steps_model_tflops_per_gpu_avg: 186.640 + last_10_steps_tokens_per_second_per_gpu_avg: 2049.062951 resume_comparison: reference_item: pretrain sentinel_steps: [51, 100] loss_relative_tolerance: 1.0e-2 loss_absolute_tolerance: 1.0e-6 - sentinels_match: false + sentinels_match: true expected_result: > - Verification requires a new uninterrupted reference run that saves a step-50 - checkpoint, followed by a fresh process that restores optimizer, scheduler, - data-order, and RNG state and runs steps 51-100 in a separate output root. - Step-51 and step-100 losses must match the uninterrupted reference within the - declared 1% relative and 1e-6 absolute tolerances, and all five metrics must be recorded. + The direct 64-GPU TP1/PP4/VPP4/EP16 continuation restores model, optimizer, + scheduler, data-order, and RNG state from the complete step-50 reference + checkpoint, starts at step 51, and finishes at step 100 in a distinct output + root. All 50 resumed steps are finite with no skipped or NaN iterations, all + five metrics are recorded, the post-setup configuration persists, and a + complete full-state step-100 checkpoint is saved. Step-51 resumed/reference + LM loss is 4.028752/4.028773, an absolute difference of 0.000021. Step-100 + resumed/reference LM loss is 3.261145/3.280722, an absolute difference of + 0.019577. Both sentinels pass the declared 1% relative and 1e-6 absolute gate. pretrain_performance: GB200: diff --git a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py index fb448700ca..85784a930f 100644 --- a/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py +++ b/tests/unit_tests/skills/create_model_verification_card/test_validate_card.py @@ -22,6 +22,7 @@ ("deepseek-v3", "pretrain_performance", "GB200"): (4096, 4096, 256), ("deepseek-v3", "pretrain_performance", "GB300"): (4096, 4096, 256), ("deepseek-v4-flash", "pretrain", "GB200"): (4096, 256, 64), + ("deepseek-v4-flash", "checkpoint_resume", "GB200"): (4096, 256, 64), ("deepseek-v4-flash", "sft", "GB200"): (1024, 128, 32), ("deepseek-v4-flash", "peft", "GB200"): (1024, 128, 32), ("deepseek-v4-flash", "sft_long_context", "GB200"): (1024, 128, 64), From 0c9e0dc0085ff44cd5479a307b586658f6788e02 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Fri, 28 Aug 2026 11:08:36 -0700 Subject: [PATCH 34/35] fix(recipes): address DeepSeek V4 review findings Signed-off-by: Chen Cui --- .../model_verification_cards/deepseek-v4-flash/card.yaml | 2 +- examples/models/deepseek_v4/README.md | 8 -------- tests/unit_tests/recipes/test_deepseek_recipes.py | 2 +- 3 files changed, 2 insertions(+), 10 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index cc6a6e3aec..f5c96e10c2 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -322,7 +322,7 @@ items: --recipe deepseek_v4_flash_sft_openmath_thinking_packed_config --step-func dsv4_step --pretrained_checkpoint work/models/deepseek-v4-flash-megatron - --save_dir work/model-verification/dsv4-flash/sft-ref + --save_dir work/model-verification/dsv4-flash/sft-long-context-ref --save_interval 100 --max_steps 100 model.pipeline_model_parallel_size=4 'model.pipeline_model_parallel_layout=Et*11|t*11|t*11|t*10mL' diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 3ce31bd319..6f78730045 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -28,14 +28,6 @@ supported training runtime. Switch the submodule before syncing the environment: uv sync ``` -Grouped-FP8 checkpoint saves also require the stateless grouped extra-state fix -merged in Megatron-LM -[#5997](https://github.com/NVIDIA/Megatron-LM/pull/5997). If the selected `dev` -revision does not contain that fix yet, apply the equivalent signed commit -`d05cb07ff618` from -[#6509](https://github.com/NVIDIA/Megatron-LM/pull/6509) to the compatible `dev` -checkout before enabling checkpoint saves. - `fast-hadamard-transform` is required by DSA and is installed from the pinned source dependency by `uv sync`. Run the examples in a CUDA-enabled Megatron Bridge container; see the diff --git a/tests/unit_tests/recipes/test_deepseek_recipes.py b/tests/unit_tests/recipes/test_deepseek_recipes.py index 99a0f770cb..0d571a2a34 100644 --- a/tests/unit_tests/recipes/test_deepseek_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_recipes.py @@ -492,7 +492,7 @@ def test_deepseek_v4_recipes_keep_hardware_qualified_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 + assert getattr(cfg.model, "moe_hybridep_num_sms", None) is None @pytest.mark.parametrize( From c6d73b7a308a0c1d8a50bfb751fe56e3877e3bee Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Sun, 30 Aug 2026 16:02:54 -0700 Subject: [PATCH 35/35] fix(recipes): align DeepSeek V4 GB200 config Signed-off-by: Chen Cui --- .../deepseek-v4-flash/card.yaml | 4 ++-- examples/models/deepseek_v4/README.md | 4 ++-- .../references/recipe-index.md | 2 +- src/megatron/bridge/recipes/deepseek/__init__.py | 4 ++-- .../bridge/recipes/deepseek/gb200/__init__.py | 4 ++-- .../bridge/recipes/deepseek/gb200/deepseek_v4.py | 12 +++++------- .../recipes/test_deepseek_v4_hardware_recipes.py | 6 +++--- 7 files changed, 17 insertions(+), 19 deletions(-) diff --git a/examples/model_verification_cards/deepseek-v4-flash/card.yaml b/examples/model_verification_cards/deepseek-v4-flash/card.yaml index f5c96e10c2..29e9b7c630 100644 --- a/examples/model_verification_cards/deepseek-v4-flash/card.yaml +++ b/examples/model_verification_cards/deepseek-v4-flash/card.yaml @@ -184,7 +184,7 @@ items: moe_dispatcher: hybridep command: > ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 - --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config + --recipe deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config --mode pretrain --dataset megatron-indexed --max_steps 100 --pretrained_checkpoint work/model-verification/dsv4-flash/import-gpu/iter_0000000 'dataset.blend=[["work/data/the-pile/my-gpt3_08_text_document"],null]' @@ -424,7 +424,7 @@ items: depends_on: pretrain command: > ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 - --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config + --recipe deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config --mode pretrain --dataset megatron-indexed --max_steps 100 'dataset.blend=[["work/data/the-pile/my-gpt3_08_text_document"],null]' dataset.path_to_cache=work/cache/the-pile dataset.num_workers=0 diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md index 6f78730045..6ef4e3a2cc 100644 --- a/examples/models/deepseek_v4/README.md +++ b/examples/models/deepseek_v4/README.md @@ -97,7 +97,7 @@ launcher. Hardware-qualified library recipes are defined under | GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config` | BF16 Adam | | GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config` | MXFP8 Adam | | GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config` | BF16 Muon | -| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | MXFP8 Adam | +| GB200, 64 GPUs | `deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config` | MXFP8 Adam | | GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config` | BF16 Adam | | GB300, 32 GPUs | `deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config` | MXFP8 Adam | @@ -105,7 +105,7 @@ For example, a short generated-data run of the GB200 library recipe is: ```bash ./scripts/training/train.sh --nodes 16 --gpus-per-node 4 \ - --recipe deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config \ + --recipe deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config \ --mode pretrain --dataset mock --max_steps 10 ``` diff --git a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md index 77add30b7d..dd2352f5c8 100644 --- a/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md +++ b/skills/nemo-mbridge-recipe-recommender/references/recipe-index.md @@ -74,7 +74,7 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | `deepseek_v2_pretrain_config` | Pretrain | 1 | 4 | 32 | 128 | | `deepseek_v3_pretrain_config` | Pretrain | 2 | 16 | 64 | 2048 | | `deepseek_v3_pretrain_config_32nodes` | Pretrain | 2 | 8 | 32 | 256 | -| `deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 4 | 16 | 64 (GB200; VPP4; natural unlimited capacity) | +| `deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config` | Pretrain | 1 | 4 | 16 | 64 (GB200; VPP4; natural unlimited capacity) | | `deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config` | SFT | 1 | 4 | 8 | 32 (GB200; offline packed) | | `deepseek_v4_flash_peft_openmath_thinking_packed_config` | PEFT | 1 | 4 | 8 | 32 (offline packed) | | `deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config` | PEFT | 1 | 4 | 8 | 32 (GB200; offline packed) | diff --git a/src/megatron/bridge/recipes/deepseek/__init__.py b/src/megatron/bridge/recipes/deepseek/__init__.py index 986a27991a..b22e03b3be 100644 --- a/src/megatron/bridge/recipes/deepseek/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/__init__.py @@ -52,7 +52,7 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, + deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config, deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config, ) from .gb300.deepseek_v4 import ( @@ -86,7 +86,7 @@ "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", - "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", + "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config", "deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config", "deepseek_v4_pro_pretrain_32gpu_gb300_bf16_config", "deepseek_v4_pro_pretrain_32gpu_gb300_fp8mx_config", diff --git a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py index b91ceffda7..8a46121edd 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/__init__.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/__init__.py @@ -18,14 +18,14 @@ deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config, deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config, deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config, + deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config, deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config, ) __all__ = [ "deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config", - "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config", + "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_config", "deepseek_v4_flash_pretrain_64gpu_gb200_bf16_muon_config", "deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config", diff --git a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py index 6334a77f1f..c6adda2d41 100644 --- a/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/recipes/deepseek/gb200/deepseek_v4.py @@ -191,14 +191,12 @@ def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: return cfg -def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigContainer: +def deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config() -> ConfigContainer: """Return the scalable real-training DeepSeek V4 Flash GB200 config. - The historical name is retained for compatibility. This variant uses PP4/VPP4, - EP16, and selective recompute without activation offload; its current - qualification is at 64 GPUs. Expert capacity, paged stash, CUDA graphs, and - forced load balancing remain disabled to preserve natural-routing training - semantics. + This 64-GPU variant uses PP4/VPP4, EP16, and selective recompute without + activation offload. Expert capacity, paged stash, CUDA graphs, and forced + load balancing remain disabled to preserve natural-routing training semantics. """ cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() @@ -222,7 +220,7 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config() -> ConfigCont cfg.model.use_transformer_engine_op_fuser = True cfg.model.recompute_modules = ["mhc", "mla_up_proj"] cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = [] + cfg.model.offload_modules = None cfg.model.fine_grained_offloading_max_inflight_offloads = None cfg.model.moe_pad_experts_for_cuda_graph_inference = False cfg.model.cuda_graph_impl = "none" diff --git a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py index c009b9b85c..0e2d287a17 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_hardware_recipes.py @@ -30,7 +30,7 @@ deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config as flash_fp8_config, ) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config as flash_library_config, + deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config as flash_library_config, ) from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config as flash_packed_sft_config, @@ -227,7 +227,7 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: ) assert cfg.model.recompute_modules == ["mhc", "mla_up_proj"] assert cfg.model.fine_grained_activation_offloading is False - assert cfg.model.offload_modules == [] + assert cfg.model.offload_modules is None assert cfg.model.fine_grained_offloading_max_inflight_offloads is None assert getattr(cfg.model, "moe_expert_rank_capacity_factor", None) is None assert getattr(cfg.model, "moe_paged_stash", False) is False @@ -240,5 +240,5 @@ def test_flash_high_scale_recipe_preserves_real_training_contract() -> None: def test_high_scale_deepseek_v4_recipes_are_exported() -> None: assert recipes.deepseek_v4_flash_peft_openmath_thinking_packed_gb200_config is flash_packed_peft_config - assert recipes.deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_library_config is flash_library_config + assert recipes.deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_library_config is flash_library_config assert recipes.deepseek_v4_flash_sft_openmath_thinking_packed_gb200_config is flash_packed_sft_config