Skip to content

Commit 2eca468

Browse files
ghadiaravi13rghadia
andauthored
feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes (#6228)
Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> Signed-off-by: rghadia <rghadia@nsc-lhr-slurm-1-vscode-01.nvidia.com> Signed-off-by: rghadia <rghadia@nvidia.com> Co-authored-by: rghadia <rghadia@nsc-lhr-slurm-1-vscode-01.nvidia.com>
1 parent b1bb0a9 commit 2eca468

11 files changed

Lines changed: 160 additions & 73 deletions

File tree

‎src/megatron/bridge/perf_recipes/nemotronh/common.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,27 @@ def _apply_nemotron_3_ultra_perf_defaults(cfg: ConfigContainer) -> None:
206206
cfg.ddp.num_buckets = 48
207207

208208

209+
def _apply_nemotron_3_ultra_gtp(cfg: ConfigContainer) -> None:
210+
"""Apply rack-local dense GTP settings for Nemotron 3 Ultra on GB300.
211+
212+
Dense weights are sharded across one 64-GPU NVLink domain. Routed-expert
213+
weights remain unsharded by GTP because EP64 already spans that domain;
214+
this avoids expert-weight all-gathers across racks.
215+
"""
216+
cfg.dist.use_megatron_fsdp = False
217+
cfg.ddp.use_megatron_fsdp = False
218+
219+
# GTP is controlled through the user-facing total weight-shard counts.
220+
# TP1 x GTP64 gives one dense-weight shard per GPU within a GB300 rack.
221+
cfg.model.tensor_parallel_num_weight_shards = _GB300_NVLINK_DOMAIN_GPUS
222+
cfg.model.expert_tensor_parallel_num_weight_shards = cfg.model.expert_tensor_parallel_size
223+
224+
# GTP currently supports only one distributed-optimizer instance.
225+
cfg.ddp.num_distributed_optimizer_instances = 1
226+
cfg.ddp.data_parallel_sharding_strategy = "no_shard"
227+
cfg.ddp.outer_dp_sharding_strategy = "no_shard"
228+
229+
209230
def _apply_nemotron_3_ultra_fsdp_hsdp(cfg: ConfigContainer, num_gpus: int) -> None:
210231
"""Apply Megatron-FSDP (HSDP) settings for Nemotron 3 Ultra on GB300.
211232

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

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,15 @@
1313
# limitations under the License.
1414
"""GB200 performance recipes for Nemotron 3 and 3.5."""
1515

16+
from megatron.bridge.perf_recipes._common import _enable_ncclep
1617
from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS
1718
from megatron.bridge.perf_recipes.nemotronh.common import (
1819
_TE_QUANT_CFG_PATH,
1920
ConfigContainer,
2021
_apply_nemotron_3_nano_perf_defaults,
2122
_apply_nemotron_3_super_perf_defaults,
2223
_apply_nemotron_3_ultra_fsdp_hsdp,
24+
_apply_nemotron_3_ultra_gtp,
2325
_apply_nemotron_3_ultra_perf_defaults,
2426
_benchmark_common,
2527
_enable_nemotron_3_super_full_iteration,
@@ -185,10 +187,10 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer:
185187

186188

187189
def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256) -> ConfigContainer:
188-
"""Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP).
190+
"""Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, GTP.
189191
190192
TP2 + SP (due to smaller GB200 HBM) / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed
191-
precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute +
193+
precision, dense GTP32, NCCLEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute +
192194
fine-grained activation offload of the expert MLP, MTP=2.
193195
194196
Args:
@@ -198,15 +200,17 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256)
198200

199201
expert_model_parallel_size = 64
200202
global_batch_size = 256
201-
hybrid_ep_ranks_per_nvlink_domain = 64
202203

203204
cfg = nemotron_3_ultra_pretrain_config()
204205
cfg.mixed_precision = _perf_precision("fp8_mx")
205206

206207
_apply_nemotron_3_ultra_perf_defaults(cfg)
207208

208-
# Apply HSDP / FSDP dtype overrides last so they win over the generic defaults.
209-
_apply_nemotron_3_ultra_fsdp_hsdp(cfg, num_gpus=num_gpus)
209+
# Shard dense weights with GTP within each 64-GPU rack. GTP currently
210+
# requires a single distributed-optimizer instance.
211+
_apply_nemotron_3_ultra_gtp(cfg)
212+
# Also, enable GTP sharding for the moe_latent_proj module due to memory pressure on GB200
213+
cfg.model.gtp_remat_opt_in_modules = ["moe_latent_proj"]
210214

211215
# Parallelism
212216
cfg.model.tensor_model_parallel_size = 2
@@ -241,6 +245,20 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256)
241245
cfg.model.recompute_granularity = "selective"
242246
cfg.model.recompute_modules = ["moe_act"]
243247

248+
# Enable GTP sharding for the moe_latent_proj module
249+
cfg.model.gtp_remat_opt_in_modules = ["moe_latent_proj"]
250+
251+
cfg.model.cuda_graph_impl = "local"
252+
set_cuda_graph_modules(cfg.model, ["attn", "mamba"])
253+
# TODO: "moe_preprocess", "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss
254+
# TE attention requires TE's RNG tracker while MCore local graphs are capturing.
255+
cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = True
256+
257+
_enable_ncclep(cfg)
258+
# Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
259+
# host every layer, which serializes the CPU behind the GPU when dispatch is fast.
260+
cfg.model.moe_use_grouped_tensor = True
261+
244262
# Keep process settings next to the recipe so users can see the exact benchmark environment.
245263
cfg.env_vars = {
246264
**COMMON_PERF_ENV_VARS,
@@ -255,14 +273,14 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256)
255273
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
256274
# NCCL user-buffer and launch settings.
257275
"NCCL_NVLS_ENABLE": 0,
258-
# HybridEP topology for the target system.
259-
"NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": hybrid_ep_ranks_per_nvlink_domain,
260-
"NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
261-
"NVLINK_DOMAIN_SIZE": 72,
262-
"USE_MNNVL": 1,
276+
# NCCL EP dispatcher mode and one GPU per rank.
277+
"NCCL_EP_HT_EM_PULL_PUSH": 1,
263278
# Transformer Engine overlap settings for this model.
264279
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
265280
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
281+
# Use cuDNN normalization kernels for the VR200 performance candidate.
282+
"NVTE_NORM_BWD_USE_CUDNN": 1,
283+
"NVTE_NORM_FWD_USE_CUDNN": 1,
266284
# Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid
267285
# offloading weights;
268286
"NVTE_CPU_OFFLOAD_V1": 1,

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

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
_apply_nemotron_3_nano_perf_defaults,
2222
_apply_nemotron_3_super_perf_defaults,
2323
_apply_nemotron_3_ultra_fsdp_hsdp,
24+
_apply_nemotron_3_ultra_gtp,
2425
_apply_nemotron_3_ultra_perf_defaults,
2526
_benchmark_common,
2627
_enable_nemotron_3_5_lightning_full_iteration,
@@ -209,13 +210,15 @@ def nemotron_3_super_pretrain_64gpu_gb300_nvfp4_config() -> ConfigContainer:
209210
def _nemotron_3_ultra_gb300_fp8mx_config(
210211
*, num_gpus: int, expert_model_parallel_size: int, global_batch_size: int
211212
) -> ConfigContainer:
212-
"""Shared builder for Nemotron 3 Ultra GB300 MXFP8 Megatron-FSDP perf recipes."""
213+
"""Shared builder for Nemotron 3 Ultra GB300 MXFP8 GTP perf recipes."""
213214
cfg = nemotron_3_ultra_pretrain_config()
214215
cfg.mixed_precision = _perf_precision("fp8_mx")
215216

217+
_apply_nemotron_3_ultra_perf_defaults(cfg)
218+
216219
"""
217220
Uses TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1,
218-
seq 8192, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective
221+
seq 8192, NCCLEP flex dispatcher, CuteDSL fused grouped MLP, selective
219222
recompute + activation offload of the expert MLP, MTP=2. The MoE architecture
220223
(512 experts, latent MoE, MTP, squared-relu, hybrid Mamba/attention pattern,
221224
...) is inherited from the base recipe via ``AutoBridge``.
@@ -252,20 +255,30 @@ def _nemotron_3_ultra_gb300_fp8mx_config(
252255
cfg.model.recompute_granularity = "selective"
253256
cfg.model.recompute_modules = ["moe_act"]
254257

255-
_apply_nemotron_3_ultra_perf_defaults(cfg)
258+
# Shard dense weights with GTP within each 64-GPU rack. GTP currently
259+
# requires a single distributed-optimizer instance.
260+
_apply_nemotron_3_ultra_gtp(cfg)
256261

257-
# Apply HSDP / FSDP dtype overrides last so they win over the generic defaults.
258-
_apply_nemotron_3_ultra_fsdp_hsdp(cfg, num_gpus=num_gpus)
262+
cfg.model.cuda_graph_impl = "local"
263+
set_cuda_graph_modules(cfg.model, ["attn", "mamba"])
264+
# TODO: "moe_preprocess", "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss
265+
# TE attention requires TE's RNG tracker while MCore local graphs are capturing.
266+
cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = True
267+
268+
_enable_ncclep(cfg)
269+
# Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
270+
# host every layer, which serializes the CPU behind the GPU when dispatch is fast.
271+
cfg.model.moe_use_grouped_tensor = True
259272

260273
return cfg
261274

262275

263276
def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256) -> ConfigContainer:
264-
"""Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, MXFP8, Megatron-FSDP (HSDP).
277+
"""Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, MXFP8, GTP.
265278
266279
TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed
267-
precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective
268-
recompute + fine-grained activation offload of the expert MLP, MTP=2.
280+
precision, dense GTP64, NCCLEP flex dispatcher, CuteDSL fused grouped MLP,
281+
selective recompute + fine-grained activation offload of the expert MLP, MTP=2.
269282
270283
Args:
271284
num_gpus: Launch GPU count used to size HSDP optimizer instances.
@@ -276,10 +289,6 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256)
276289
expert_model_parallel_size=64,
277290
global_batch_size=256,
278291
)
279-
_enable_ncclep(cfg)
280-
# Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
281-
# host every layer, which serializes the CPU behind the GPU when dispatch is fast.
282-
cfg.model.moe_use_grouped_tensor = True
283292
# Keep process settings next to the recipe so users can see the exact benchmark environment.
284293
cfg.env_vars = {
285294
**COMMON_PERF_ENV_VARS,
@@ -296,6 +305,9 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256)
296305
# Transformer Engine overlap settings for this model.
297306
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
298307
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
308+
# Use cuDNN normalization kernels for the VR200 performance candidate.
309+
"NVTE_NORM_BWD_USE_CUDNN": 1,
310+
"NVTE_NORM_FWD_USE_CUDNN": 1,
299311
# Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid
300312
# offloading weights;
301313
"NVTE_CPU_OFFLOAD_V1": 1,

‎src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -240,11 +240,8 @@ def nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config(*, num_gpus: int = 256
240240
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
241241
# NCCL user-buffer and launch settings.
242242
"NCCL_NVLS_ENABLE": 0,
243-
# Use 64 HybridEP ranks within the 72-GPU NVLink domain.
244-
"NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64,
245-
"NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
246-
"NVLINK_DOMAIN_SIZE": 72,
247-
"USE_MNNVL": 1,
243+
# NCCL EP dispatcher mode and one GPU per rank.
244+
"NCCL_EP_HT_EM_PULL_PUSH": 1,
248245
# Transformer Engine overlap settings for this model.
249246
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
250247
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,

‎src/megatron/bridge/training/train.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@
5757
VisionTECudaGraphHelper,
5858
get_vision_cuda_graph_seq_length,
5959
)
60+
from megatron.core.transformer.cuda_graphs import (
61+
delete_cuda_graphs as delete_local_cuda_graphs,
62+
)
6063
from megatron.core.utils import (
6164
check_param_hashes_across_dp_replicas,
6265
get_attr_wrapped_model,
@@ -1786,6 +1789,9 @@ def _delete_cuda_graphs(cuda_graph_helper: TECudaGraphHelper | None):
17861789
if cuda_graph_helper is not None and cuda_graph_helper.graphs_created():
17871790
cuda_graph_helper.delete_cuda_graphs()
17881791

1792+
# Cleanup per-layer CUDA graphs owned by MCore CudaGraphManager (cuda_graph_impl="local").
1793+
delete_local_cuda_graphs()
1794+
17891795
# Run GC to collect the freshed object
17901796
gc.collect()
17911797

‎src/megatron/bridge/utils/cuda_graph.py‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -140,16 +140,21 @@ def validate_cuda_graph_configuration(config: Any, *, config_name: str = "model"
140140

141141
graph_modules = cuda_graph_module_names(config)
142142
inference_scopes = _member_names(getattr(config, "inference_cuda_graph_scope", None))
143+
is_local_per_layer_graph = cuda_graph_impl == "local" and bool(graph_modules)
143144
is_local_inference_graph = (
144145
cuda_graph_impl == "local" and not graph_modules and bool(inference_scopes & {"layer", "block"})
145146
)
146-
if cuda_graph_impl == "local" and not is_full_iteration_cuda_graph(config) and not is_local_inference_graph:
147-
modules_msg = f" with scopes {graph_modules}" if graph_modules else ""
147+
if (
148+
cuda_graph_impl == "local"
149+
and not is_full_iteration_cuda_graph(config)
150+
and not is_local_per_layer_graph
151+
and not is_local_inference_graph
152+
):
148153
raise ValueError(
149154
f'Megatron Bridge supports {config_name}.cuda_graph_impl="local" only for '
150-
f"full-iteration CUDA graphs{modules_msg}. Use "
151-
f'{config_name}.cuda_graph_impl="transformer_engine" for scoped CUDA graph '
152-
'capture such as "attn" or "mlp".'
155+
"full-iteration CUDA graphs or per-layer CUDA graphs with explicit "
156+
f'{config_name}.cuda_graph_modules such as ["attn", "mamba"]. Whole-layer local '
157+
"capture with empty cuda_graph_modules is not supported."
153158
)
154159

155160

‎tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -177,17 +177,17 @@ def test_vrnvl72_perf_recipe_uses_nvl72_ultra_topology() -> None:
177177
assert cfg.model.expert_model_parallel_size == 64
178178
assert cfg.model.expert_tensor_parallel_size == 1
179179
assert cfg.model.moe_token_dispatcher_type == "flex"
180-
assert cfg.model.moe_flex_dispatcher_backend == "hybridep"
181-
assert cfg.model.cuda_graph_impl == "none"
180+
assert cfg.model.moe_flex_dispatcher_backend == "ncclep"
181+
assert cfg.model.cuda_graph_impl == "local"
182182
assert cfg.train.global_batch_size == 256
183183
assert cfg.train.micro_batch_size == 1
184-
assert cfg.ddp.use_megatron_fsdp is True
185-
assert cfg.ddp.num_distributed_optimizer_instances == 4
186-
assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 64
187-
assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72
184+
assert cfg.ddp.use_megatron_fsdp is False
185+
assert cfg.model.tensor_parallel_num_weight_shards == 64
186+
assert cfg.model.expert_tensor_parallel_num_weight_shards == 1
187+
assert cfg.ddp.num_distributed_optimizer_instances == 1
188+
assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1
188189
assert cfg.env_vars["NVTE_NORM_BWD_USE_CUDNN"] == 1
189190
assert cfg.env_vars["NVTE_NORM_FWD_USE_CUDNN"] == 1
190-
assert cfg.env_vars["USE_MNNVL"] == 1
191191

192192

193193
@pytest.mark.unit

‎tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
"""Every GB300 MoE perf recipe switched to NCCL EP carries the full NCCL EP stack, and the VRNVL72
15-
aliases built on the shared HybridEP bases stay on HybridEP."""
15+
aliases built on the shared HybridEP bases stay on HybridEP. VRNVL72 aliases built on an NCCL EP base
16+
carry the same NCCL EP stack."""
1617

1718
import importlib
1819

@@ -53,6 +54,9 @@
5354
("qwen_vl.gb300.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_gb300_bf16_config"),
5455
)
5556

57+
# VRNVL72 aliases that share the GB300 NCCL EP builder instead of a HybridEP base.
58+
_VRNVL72_NCCLEP_RECIPES = (("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config"),)
59+
5660
# VRNVL72 aliases of the same recipes: they reuse the shared HybridEP builders and must not pick up NCCL EP.
5761
_VRNVL72_ALIASES = (
5862
("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config"),
@@ -63,7 +67,6 @@
6367
("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config"),
6468
("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_bf16_config"),
6569
("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_nvfp4_config"),
66-
("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config"),
6770
("nemotronh.vrnvl72.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config"),
6871
("nemotronh.vrnvl72.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_fp8mx_config"),
6972
("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config"),
@@ -94,7 +97,9 @@ def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None:
9497
patch_recipe_construction_dependencies(monkeypatch)
9598

9699

97-
@pytest.mark.parametrize(("module_name", "factory_name"), _GB300_NCCLEP_RECIPES, ids=lambda value: value)
100+
@pytest.mark.parametrize(
101+
("module_name", "factory_name"), (*_GB300_NCCLEP_RECIPES, *_VRNVL72_NCCLEP_RECIPES), ids=lambda value: value
102+
)
98103
def test_gb300_moe_perf_recipes_default_to_the_nccl_ep_stack(module_name: str, factory_name: str) -> None:
99104
cfg = _build(module_name, factory_name)
100105

‎tests/unit_tests/recipes/test_override_precedence.py‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -435,19 +435,23 @@ def test_local_full_iteration_is_normalized(self):
435435
assert recipe.rng.te_rng_tracker is True
436436
assert recipe.rerun_state_machine.check_for_nan_in_loss is False
437437

438-
def test_local_scoped_graph_is_rejected(self):
439-
"""Local graphs cannot capture individual layer modules in Bridge."""
438+
def test_local_scoped_graph_is_normalized(self):
439+
"""Local per-layer graphs populate current MCore module values."""
440440
from utils.overrides import _set_cuda_graph_overrides
441441

442-
with pytest.raises(
443-
ValueError,
444-
match='cuda_graph_impl="local".*cuda_graph_impl="transformer_engine"',
445-
):
446-
_set_cuda_graph_overrides(
447-
_minimal_cuda_graph_recipe(),
448-
cuda_graph_impl="local",
449-
cuda_graph_scope="mlp",
450-
)
442+
from megatron.bridge.utils.cuda_graph import cuda_graph_module_names, is_full_iteration_cuda_graph
443+
444+
recipe = _set_cuda_graph_overrides(
445+
_minimal_cuda_graph_recipe(),
446+
cuda_graph_impl="local",
447+
cuda_graph_scope="mlp",
448+
)
449+
450+
assert recipe.model.cuda_graph_impl == "local"
451+
assert cuda_graph_module_names(recipe.model) == ["mlp"]
452+
assert not is_full_iteration_cuda_graph(recipe.model)
453+
assert recipe.model.use_te_rng_tracker is True
454+
assert recipe.rng.te_rng_tracker is True
451455

452456
def test_transformer_engine_scoped_graph_is_normalized(self):
453457
"""TE scoped graphs populate current MCore module values."""

0 commit comments

Comments
 (0)