Merge branch 'rghadia/nt3-ultra-gtp-cg' of https://github.com/ghadiaravi13/NT3-Ultra-Megatron-Bridge into rghadia/nt3-ultra-gtp-cg

Signed-off-by: rghadia <rghadia@nvidia.com>

Author: ghadiaravi13

src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py

@@ -249,8 +249,8 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256)
     cfg.model.gtp_remat_opt_in_modules = ["moe_latent_proj"]
 
     cfg.model.cuda_graph_impl = "local"
-    set_cuda_graph_modules(cfg.model, ["attn", "mamba", "moe_preprocess"])
-    # TODO: "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss
+    set_cuda_graph_modules(cfg.model, ["attn", "mamba"])
+    # TODO: "moe_preprocess", "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss
     # TE attention requires TE's RNG tracker while MCore local graphs are capturing.
     cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = True
 

src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py

@@ -260,8 +260,8 @@ def _nemotron_3_ultra_gb300_fp8mx_config(
     _apply_nemotron_3_ultra_gtp(cfg)
 
     cfg.model.cuda_graph_impl = "local"
-    set_cuda_graph_modules(cfg.model, ["attn", "mamba", "moe_preprocess"])
-    # TODO: "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss
+    set_cuda_graph_modules(cfg.model, ["attn", "mamba"])
+    # TODO: "moe_preprocess", "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss
     # TE attention requires TE's RNG tracker while MCore local graphs are capturing.
     cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = True