m-bridge for nemo container 26.02

Author: srivatsankrishnan

conf/experimental/megatron_bridge/test/b200/megatron_bridge_qwen_30b.toml

@@ -28,8 +28,8 @@ mount_as = "/opt/Megatron-Bridge"
 [cmd_args]
 gpu_type = "b200"
 container_image = "nvcr.io#nvidia/nemo:25.11.01"
-model_name = "qwen3"
-model_size = "30b_a3b"
+model_family_name = "qwen3"
+model_recipe_name = "30b_a3b"
 gpus_per_node = 4
 num_gpus = 8
 domain = "llm"

conf/experimental/megatron_bridge/test/gb200/megatron_bridge_qwen_30b.toml

@@ -28,8 +28,8 @@ mount_as = "/opt/Megatron-Bridge"
 [cmd_args]
 gpu_type = "gb200"
 container_image = "nvcr.io#nvidia/nemo:25.11.01"
-model_name = "qwen3"
-model_size = "30b_a3b"
+model_family_name = "qwen3"
+model_recipe_name = "30b_a3b"
 gpus_per_node = 4
 num_gpus = 8
 domain = "llm"

conf/experimental/megatron_bridge/test/gb300/megatron_bridge_qwen_30b.toml

@@ -28,8 +28,8 @@ mount_as = "/opt/Megatron-Bridge"
 [cmd_args]
 gpu_type = "gb300"
 container_image = "nvcr.io#nvidia/nemo:25.11.01"
-model_name = "qwen3"
-model_size = "30b_a3b"
+model_family_name = "qwen3"
+model_recipe_name = "30b_a3b"
 gpus_per_node = 4
 num_gpus = 8
 domain = "llm"

conf/experimental/megatron_bridge/test/h100/megatron_bridge_qwen_30b.toml

@@ -28,8 +28,8 @@ mount_as = "/opt/Megatron-Bridge"
 [cmd_args]
 gpu_type = "h100"
 container_image = "nvcr.io#nvidia/nemo:25.11.01"
-model_name = "qwen3"
-model_size = "30b_a3b"
+model_family_name = "qwen3"
+model_recipe_name = "30b_a3b"
 gpus_per_node = 8
 num_gpus = 16
 domain = "llm"

doc/workloads/megatron_bridge.rst

@@ -19,8 +19,8 @@ Test TOML example:
    # Container can be an NGC/enroot URL (nvcr.io#...) or a local .sqsh path.
    container_image = "nvcr.io#nvidia/nemo:25.11.01"
    
-   model_name = "qwen3"
-   model_size = "30b_a3b"
+   model_family_name = "qwen3"
+   model_recipe_name = "30b_a3b"
    task = "pretrain"
    domain = "llm"
    compute_dtype = "fp8_mx"
@@ -55,8 +55,8 @@ Test-in-Scenario example:
    
      [Tests.cmd_args]
      container_image = "nvcr.io#nvidia/nemo:25.11.01"
-     model_name = "qwen3"
-     model_size = "30b_a3b"
+     model_family_name = "qwen3"
+     model_recipe_name = "30b_a3b"
      task = "pretrain"
      domain = "llm"
      compute_dtype = "fp8_mx"

src/cloudai/workloads/megatron_bridge/megatron_bridge.py

@@ -40,17 +40,23 @@ class MegatronBridgeCmdArgs(CmdArgs):
     detach: Optional[bool] = Field(default=None)
 
     # Model/task
-    model_name: str = Field(default="")
-    model_size: str = Field(default="")
+    model_family_name: str = Field(default="")
+    model_recipe_name: str = Field(default="")
+    use_recipes: Optional[bool] = Field(default=None)
     domain: str = Field(default="llm")
     task: str = Field(default="pretrain")
     compute_dtype: str = Field(default="bf16")
     fp8_recipe: Optional[str] = Field(default=None)
     hf_token: Optional[str] = Field(default=None)
     nemo_home: Optional[str] = Field(default=None)
     wandb_key: Optional[str] = Field(default=None)
-    wandb_prj_name: Optional[str] = Field(default=None)
-    wandb_exp_name: Optional[str] = Field(default=None)
+    wandb_project_name: Optional[str] = Field(default=None)
+    wandb_entity_name: Optional[str] = Field(default=None)
+    wandb_experiment_name: Optional[str] = Field(default=None)
+    wandb_save_dir: Optional[str] = Field(default=None)
+
+    # Retries
+    max_retries: Optional[int] = Field(default=None)
 
     # Feature flags (allow sweeps)
     use_tokendrop: Optional[Union[bool, List[bool]]] = Field(default=None)
@@ -69,6 +75,43 @@ class MegatronBridgeCmdArgs(CmdArgs):
     # Batch sizes
     mb: Optional[Union[int, List[int]]] = Field(default=None)
     gb: Optional[Union[int, List[int]]] = Field(default=None)
+    seq_length: Optional[Union[int, List[int]]] = Field(default=None)
+
+    # Optimizer
+    lr: Optional[Union[float, List[float]]] = Field(default=None)
+    min_lr: Optional[Union[float, List[float]]] = Field(default=None)
+    warmup_iters: Optional[Union[int, List[int]]] = Field(default=None)
+
+    # Checkpointing
+    pretrained_checkpoint: Optional[str] = Field(default=None)
+    save_dir: Optional[str] = Field(default=None)
+    load_dir: Optional[str] = Field(default=None)
+    save_interval: Optional[int] = Field(default=None)
+    most_recent_k: Optional[int] = Field(default=None)
+    save_config_filepath: Optional[str] = Field(default=None)
+
+    # Data / Tokenizer
+    data: Optional[str] = Field(default=None)
+    dataset_paths: Optional[Union[str, List[str]]] = Field(default=None)
+    dataset_root: Optional[str] = Field(default=None)
+    index_mapping_dir: Optional[str] = Field(default=None)
+    dataset_name: Optional[str] = Field(default=None)
+    packed_sequence: Optional[bool] = Field(default=None)
+    head_only: Optional[bool] = Field(default=None)
+    tokenizer_type: Optional[str] = Field(default=None)
+    tokenizer_model: Optional[str] = Field(default=None)
+    vocab_size: Optional[int] = Field(default=None)
+
+    # Profiling (performance group in argument_parser.py)
+    pytorch_profiler: Optional[bool] = Field(default=None)
+    profiling_start_step: Optional[int] = Field(default=None)
+    profiling_stop_step: Optional[int] = Field(default=None)
+    record_memory_history: Optional[bool] = Field(default=None)
+    profiling_gpu_metrics: Optional[bool] = Field(default=None)
+    profiling_ranks: Optional[Union[int, List[int]]] = Field(default=None)
+
+    # Performance
+    nccl_ub: Optional[Union[bool, List[bool]]] = Field(default=None)
 
     # Perf/tuning
     moe_a2a_overlap: Optional[Union[bool, List[bool]]] = Field(default=None)

src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py

@@ -154,8 +154,8 @@ def _build_launcher_parts(  # noqa: C901
     ) -> list[str]:
         fields_set = args.model_fields_set
         force_fields = {
-            "model_name",
-            "model_size",
+            "model_family_name",
+            "model_recipe_name",
             "num_gpus",
             "gpus_per_node",
             "hf_token",
@@ -198,6 +198,15 @@ def add(flag: str, value: Any) -> None:
                 return
             if isinstance(value, bool):
                 parts.extend([flag, "true" if value else "false"])
+            elif isinstance(value, (list, tuple)):
+                if not value:
+                    return
+                if flag == "--dataset_paths":
+                    parts.extend([flag, *[str(x) for x in value]])
+                elif flag == "--profiling_ranks":
+                    parts.extend([flag, ",".join(str(x) for x in value)])
+                else:
+                    parts.extend([flag, str(value[0]) if len(value) == 1 else ",".join(str(x) for x in value)])
             else:
                 sv = str(value)
                 if sv != "":
@@ -222,31 +231,38 @@ def add_field(field: str, flag: str, value: Any) -> None:
         add_field("hf_token", "-hf", args.hf_token)
         add_field("nemo_home", "-nh", args.nemo_home)
         add_field("wandb_key", "-wdk", args.wandb_key)
-        add_field("wandb_prj_name", "-wdp", args.wandb_prj_name)
-        add_field("wandb_exp_name", "-wdj", args.wandb_exp_name)
+        add_field("wandb_project_name", "-wdp", args.wandb_project_name)
+        add_field("wandb_entity_name", "-wde", args.wandb_entity_name)
+        add_field("wandb_experiment_name", "-wdj", args.wandb_experiment_name)
+        add_field("wandb_save_dir", "-wds", args.wandb_save_dir)
+        add_field("max_retries", "--max_retries", args.max_retries)
         if args.dryrun and "dryrun" in fields_set:
             parts.append("-d")
         add_field("num_gpus", "-ng", args.num_gpus)
         add_field("gpus_per_node", "-gn", args.gpus_per_node)
         if mounts:
             add("-cm", ",".join(mounts))
 
-        # Model flags (Megatron-Bridge r0.2.0 API)
+        # Model flags (Megatron-Bridge main-branch API)
+        if args.use_recipes and "use_recipes" in fields_set:
+            parts.append("--use_recipes")
         if "enable_vboost" in fields_set:
             add_field("enable_vboost", "-vb", bool(args.enable_vboost))
-        if not args.model_name:
-            raise RuntimeError("Missing required cmd_args.model_name (maps to -m/--model_name).")
-        if not args.model_size:
-            raise RuntimeError("Missing required cmd_args.model_size (maps to -s/--model_size).")
-        add_field("model_name", "-m", args.model_name)
-        add_field("model_size", "-s", args.model_size)
+        if not args.model_family_name:
+            raise RuntimeError("Missing required cmd_args.model_family_name (maps to -m/--model_family_name).")
+        if not args.model_recipe_name:
+            raise RuntimeError("Missing required cmd_args.model_recipe_name (maps to -mr/--model_recipe_name).")
+        add_field("model_family_name", "-m", args.model_family_name)
+        add_field("model_recipe_name", "-mr", args.model_recipe_name)
         if args.enable_nsys and "enable_nsys" in fields_set:
             parts.append("-en")
         add_field("domain", "--domain", args.domain)
         if "use_tokendrop" in fields_set and args.use_tokendrop is not None:
             add_field("use_tokendrop", "--use_tokendrop", bool(args.use_tokendrop))
         if "use_megatron_fsdp" in fields_set and args.use_megatron_fsdp is not None:
             add_field("use_megatron_fsdp", "--use_megatron_fsdp", bool(args.use_megatron_fsdp))
+        if "nccl_ub" in fields_set and args.nccl_ub is not None:
+            add_field("nccl_ub", "--nccl_ub", bool(args.nccl_ub))
         add_field("cuda_graph_impl", "--cuda_graph_impl", args.cuda_graph_impl)
         if args.cuda_graph_scope and "cuda_graph_scope" in fields_set:
             add_field(
@@ -264,6 +280,7 @@ def add_field(field: str, flag: str, value: Any) -> None:
         # Batch
         add_field("mb", "-mb", args.mb)
         add_field("gb", "-gb", args.gb)
+        add_field("seq_length", "-sl", args.seq_length)
 
         # Misc
         if "moe_a2a_overlap" in fields_set:
@@ -273,11 +290,44 @@ def add_field(field: str, flag: str, value: Any) -> None:
         add_field("activation_offload_layers", "-ol", args.activation_offload_layers)
         if args.recompute_modules and "recompute_modules" in fields_set:
             parts.extend(["--recompute_modules", self._normalize_recompute_modules(args.recompute_modules)])
-        # r0.2.0 supports `--detach` / `--no-detach` flags (no boolean value)
-        if args.detach is True and "detach" in fields_set:
-            parts.append("--detach")
-        elif args.detach is False and "detach" in fields_set:
-            parts.append("--no-detach")
+        if "detach" in fields_set and args.detach is not None:
+            parts.extend(["--detach", "true" if args.detach else "false"])
+
+        # Optimizer
+        add_field("lr", "--lr", args.lr)
+        add_field("min_lr", "--min_lr", args.min_lr)
+        add_field("warmup_iters", "--warmup_iters", args.warmup_iters)
+
+        # Checkpointing
+        add_field("pretrained_checkpoint", "--pretrained_checkpoint", args.pretrained_checkpoint)
+        add_field("save_dir", "--save_dir", args.save_dir)
+        add_field("load_dir", "--load_dir", args.load_dir)
+        add_field("save_interval", "--save_interval", args.save_interval)
+        add_field("most_recent_k", "--most_recent_k", args.most_recent_k)
+        add_field("save_config_filepath", "--save_config_filepath", args.save_config_filepath)
+
+        # Data / Tokenizer
+        add_field("data", "--data", args.data)
+        add_field("dataset_paths", "--dataset_paths", args.dataset_paths)
+        add_field("dataset_root", "--dataset_root", args.dataset_root)
+        add_field("index_mapping_dir", "--index_mapping_dir", args.index_mapping_dir)
+        add_field("dataset_name", "--dataset_name", args.dataset_name)
+        if args.packed_sequence and "packed_sequence" in fields_set:
+            parts.append("--packed_sequence")
+        if args.head_only and "head_only" in fields_set:
+            parts.append("--head_only")
+        add_field("tokenizer_type", "--tokenizer_type", args.tokenizer_type)
+        add_field("tokenizer_model", "--tokenizer_model", args.tokenizer_model)
+        add_field("vocab_size", "--vocab_size", args.vocab_size)
+
+        # Profiling (performance group)
+        add_field("pytorch_profiler", "-pyp", args.pytorch_profiler)
+        add_field("profiling_start_step", "--profiling_start_step", args.profiling_start_step)
+        add_field("profiling_stop_step", "--profiling_stop_step", args.profiling_stop_step)
+        add_field("record_memory_history", "-mh", args.record_memory_history)
+        if args.profiling_gpu_metrics and "profiling_gpu_metrics" in fields_set:
+            parts.append("--profiling_gpu_metrics")
+        add_field("profiling_ranks", "--profiling_ranks", args.profiling_ranks)
 
         # Extra user args (dict -> string)
         if tdef.extra_cmd_args: