cleanup

Author: gdengk

nemo/collections/llm/recipes/deepseek_v3.py

@@ -122,8 +122,8 @@ def pretrain_recipe(
     recipe.log.ckpt.train_time_interval = run.Config(timedelta, minutes=60)
 
     # recompute
-    recipe.model.config.recompute_granularity = None #"selective"
-    recipe.model.config.recompute_modules = None #["mla_up_proj", "layernorm"]
+    recipe.model.config.recompute_granularity = "selective"
+    recipe.model.config.recompute_modules = ["mla_up_proj", "layernorm"]
 
     # DeepEP
     deepep_callback = run.Config(DeepEPCallback)
@@ -135,11 +135,10 @@ def pretrain_recipe(
     )
     comm_overlap_callback = run.Config(
         MegatronCommOverlapCallback,
-        tp_comm_overlap=True,
-        tp_comm_bootstrap_backend='nccl',
+        tp_comm_overlap=False,
     )
 
-    # recipe.trainer.callbacks.append(deepep_callback)
+    recipe.trainer.callbacks.append(deepep_callback)
     recipe.trainer.callbacks.append(garbage_collection_callback)
     recipe.trainer.callbacks.append(comm_overlap_callback)
 

nemo/lightning/pytorch/strategies/megatron_strategy.py

@@ -574,6 +574,7 @@ def setup_distributed(self) -> None:
         """Setups dist env"""
         setup_parallel_ranks(self)
 
+        # Capture Cudagraph on a side stream
         if self.model.config.external_cuda_graph:
             torch.cuda.set_stream(torch.cuda.Stream())
 
@@ -725,7 +726,7 @@ def training_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OUTP
         assert self.lightning_module is not None
         assert isinstance(self.model, MegatronParallel)
 
-        # (TODO) Capture the cuda graph for the first step for now
+        # (TODO:) Capture the cuda graph for the first step
         if self.trainer.global_step == 0 and self.model.config.external_cuda_graph:
             # disable prehook
             if self.ddp_config.use_distributed_optimizer and self.ddp_config.overlap_param_gather:

scripts/performance/argument_parser.py

@@ -104,16 +104,6 @@ def parse_cli_args():
         help="Enable Nsys profiling. Diabled by default",
         action="store_true",
     )
-    parser.add_argument(
-        "--run_local",
-        help="Run local. Diabled by default",
-        action="store_true",
-    )
-    parser.add_argument(
-        "--partial_cg",
-        help="Run local. Diabled by default",
-        action="store_true",
-    )
     parser.add_argument(
         "-em",
         "--enable_memory_profile",

scripts/performance/llm/pretrain_deepseek_v3.py

@@ -54,7 +54,7 @@ def override_recipe_configs(
     """
     DeepSeek V3 pre-train recipe aimed at achieving best possible performance.
     """
-    recipe = pretrain_recipe(performance_mode=True, use_mtp=False)
+    recipe = pretrain_recipe(performance_mode=True)
 
     # reset recompute args in the default recipe
     if args.recompute_modules is None:
@@ -69,16 +69,16 @@ def override_recipe_configs(
     # Token dispatcher configs. For H100 we use deepEP and for Blackwell,
     # because deepEP is not supported yet, we use all-to-all dispatcher with
     # token drop. After deepEP is supported, we can use deepEP dispatcher.
-    # if args.gpu.lower() in ['h100']:
-    #     recipe.model.config.moe_token_dispatcher_type = "flex"
-    #     recipe.model.config.moe_enable_deepep = True
-    #     recipe.model.config.moe_shared_expert_overlap = False  # not supported for deepEP
-    # else:
-    recipe.model.config.moe_token_dispatcher_type = "alltoall"
-    recipe.model.config.moe_enable_deepep = False
-    recipe.model.config.moe_shared_expert_overlap = True
-    if USE_TOKEN_DROP:
-        recipe.trainer.callbacks.append(run.Config(MegatronTokenDropCallback))
+    if args.gpu.lower() in ['h100']:
+        recipe.model.config.moe_token_dispatcher_type = "flex"
+        recipe.model.config.moe_enable_deepep = True
+        recipe.model.config.moe_shared_expert_overlap = False  # not supported for deepEP
+    else:
+        recipe.model.config.moe_token_dispatcher_type = "alltoall"
+        recipe.model.config.moe_enable_deepep = False
+        recipe.model.config.moe_shared_expert_overlap = True
+        if USE_TOKEN_DROP:
+            recipe.trainer.callbacks.append(run.Config(MegatronTokenDropCallback))
 
     # Performance optimization knobs
     recipe.model.config.moe_permute_fusion = True
@@ -159,52 +159,16 @@ def override_recipe_configs(
         )
     recipe.model.tokenizer = recipe.data.tokenizer
 
-
-    if args.run_local:
-        recipe.model.config.num_moe_experts = 16
-    recipe.model.config.num_layers=3
-    recipe.model.config.moe_layer_freq=[0,1,1]
-
-
-    # add the partial cg support
-    USE_PARTIAL_CG = args.partial_cg
-    if USE_PARTIAL_CG:
-        recipe.model.config.external_cuda_graph = True
-        recipe.model.config.cuda_graph_scope = "attn"
-        recipe.trainer.strategy.use_te_rng_tracker = True
-        recipe.model.config.enable_cuda_graph = False
-        
-    
     return recipe
 
 
 if __name__ == "__main__":
     args = parse_cli_args().parse_args()
     args_sanity_check(args)
 
-    # kwargs = get_user_configs(args.gpu.lower(), "pre_train", "deepseek", "v3", args)
-    # (
-    #     num_nodes,
-    #     mbs,
-    #     gbs,
-    #     tp_size,
-    #     pp_size,
-    #     cp_size,
-    #     vp_size,
-    #     ep_size,
-    #     etp_size,
-    #     enable_cuda_graphs,
-    #     use_mcore_fsdp,
-    #     recompute_layers,
-    #     activation_offload_layers,
-    #     recompute_modules,
-    #     _,  # keep_fsdp_fp8_transpose_cache
-    #     use_user_buffer_registration,
-    #     use_sharp,
-    # ) = kwargs[:17]
-
-    if args.run_local:
-        ( num_nodes,
+    kwargs = get_user_configs(args.gpu.lower(), "pre_train", "deepseek", "v3", args)
+    (
+        num_nodes,
         mbs,
         gbs,
         tp_size,
@@ -220,25 +184,9 @@ def override_recipe_configs(
         recompute_modules,
         _,  # keep_fsdp_fp8_transpose_cache
         use_user_buffer_registration,
-        use_sharp,) = 1, 1, 32, 1, 1, 1, 1, 8, 1, (not args.partial_cg) , False, 0, 0, None, False, False, False
-    else:
-        ( num_nodes,
-        mbs,
-        gbs,
-        tp_size,
-        pp_size,
-        cp_size,
-        vp_size,
-        ep_size,
-        etp_size,
-        enable_cuda_graphs,
-        use_mcore_fsdp,
-        recompute_layers,
-        activation_offload_layers,
-        recompute_modules,
-        use_user_buffer_registration,
-        use_sharp,) = 1, 1, 32, 1, 1, 1, 1, 8, 1, False, False, 0, 0, None, False, False, False
-    
+        use_sharp,
+    ) = kwargs[:17]
+
     recipe = override_recipe_configs(
         args,
         num_nodes,
@@ -262,25 +210,22 @@ def override_recipe_configs(
     exp_config = f"{num_nodes}nodes_tp{tp_size}_pp{pp_size}_cp{cp_size}_vp{vp_size}_ep{ep_size}_{mbs}mbs_{gbs}gbs"
     exp_name = f"{splitext(basename(__file__))[0]}_{args.compute_dtype}_{exp_config}"
 
-    if not args.run_local:
-        executor = slurm_executor(
-            args.gpu.lower(),
-            args.account,
-            args.partition,
-            args.log_dir,
-            num_nodes,
+    executor = slurm_executor(
+        args.gpu.lower(),
+        args.account,
+        args.partition,
+        args.log_dir,
+        num_nodes,
         args.gpus_per_node,
         args.time_limit,
         args.container_image,
         custom_mounts=args.custom_mounts,
         custom_env_vars={},
         hf_token=args.hf_token,
         nemo_home=args.nemo_home,
-            wandb_key=args.wandb_key,
-            network='sharp' if use_sharp else None,
-        )
-    else:
-        executor = run.LocalExecutor(ntasks_per_node=8, launcher="torchrun", env_vars={})
+        wandb_key=args.wandb_key,
+        network='sharp' if use_sharp else None,
+    )
 
     plugins = [
         PerfEnvPlugin(
@@ -295,11 +240,6 @@ def override_recipe_configs(
         assert args.memory_profile_out_path is not None
         plugins.append(MemoryProfilePlugin(dir=args.memory_profile_out_path))
 
-
-    if args.run_local:
-        run.run(recipe, executor=executor, name=exp_name, plugins=plugins)
-        exit()
-    
     with run.Experiment(exp_name) as exp:
         exp.add(
             recipe,