feat(recipes): modernize DeepSeek V4 library configs (#5554)

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

.dev.commit

@@ -1 +1 @@
-2f2f8ebaef90f4d62a18a43808398223b1fc5809
+8a77f4a9b78cd58dee862f234cc5cc84e47311be

examples/model_verification_cards/deepseek-v4-flash/card.yaml

@@ -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,

examples/models/deepseek_v4/README.md

@@ -132,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"),
@@ -420,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
@@ -496,17 +505,19 @@ def main(argv: list[str] | None = None) -> None:
             world_size=benchmark_world_size,
         )
     configuration_mode = _train_mode(args.mode)
-
-    if benchmark_metadata is not None:
+    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"
         step_mode = benchmark_metadata.task
     else:
-        recipe = apply_runtime_environment(recipe)
         execution_mode = configuration_mode
         step_mode = configuration_mode
 

scripts/inference/setup_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,

scripts/training/run_recipe.py

@@ -74,6 +74,10 @@ 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_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) |
 
 ### GLM-4.5
 

skills/create-model-verification-card/scripts/verify_hf_inference.py

@@ -944,11 +944,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
@@ -962,6 +961,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,

skills/nemo-mbridge-recipe-recommender/references/recipe-index.md

@@ -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,9 +48,12 @@
     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,
+    deepseek_v4_flash_pretrain_64gpu_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,
@@ -66,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",
@@ -80,6 +86,8 @@
     "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_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",
     "set_deepseek_v4_pipeline_model_parallel_layout",

src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ruff: noqa: F401
-"""Compatibility aliases for legacy recipe names."""
+"""Hardware-agnostic DeepSeek V4 recipes and compatibility aliases."""
 
 from __future__ import annotations
 
 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,34 @@
 ]
 
 
+_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
+    # 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:
     """DSv4 Flash SFT on OpenMathInstruct-2 with thinking channel and offline-packed sequences.
 
@@ -93,6 +123,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;
@@ -104,3 +137,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