fix(nemotron): disable moe_router in CUDA graph configurations

- Removed "moe_router" from the set_cuda_graph_modules function calls in both GB200 and GB300 configurations.
- Added a TODO comment to indicate the need for this change due to MCore PR 7046, which affects CPU-to-GPU copy for per token loss.

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

Author: ghadiaravi13

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

@@ -249,7 +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_router", "moe_preprocess"])
+    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
     # 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,7 +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_router", "moe_preprocess"])
+    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
     # 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