Skip to content

Commit bcd2d13

Browse files
cuichenxcuichenx
authored
[recipes] refactor: standardize Nemotron 3.5 Lightning naming (#5303)
Signed-off-by: Chen Cui <chcui@nvidia.com>
1 parent df68e2c commit bcd2d13

14 files changed

Lines changed: 293 additions & 283 deletions

File tree

examples/model_verification_cards/nemotron-3.5-nano/card.yaml renamed to examples/model_verification_cards/nemotron-3.5-lightning/card.yaml

Lines changed: 55 additions & 55 deletions
Large diffs are not rendered by default.

‎scripts/performance/dump_perf_configs.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -276,12 +276,12 @@
276276
("nemotronh", "nemotron_3_nano", "pretrain", 8, "b200", "nvfp4"),
277277
("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "bf16"),
278278
("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "fp8_cs"),
279-
# Nemotron 3.5 Nano
280-
("nemotronh", "nemotron_3_5_nano", "pretrain", 8, "gb200", "bf16"),
281-
("nemotronh", "nemotron_3_5_nano", "pretrain", 8, "gb200", "fp8_mx"),
282-
("nemotronh", "nemotron_3_5_nano", "pretrain", 8, "gb200", "nvfp4"),
283-
("nemotronh", "nemotron_3_5_nano", "pretrain", 16, "h100", "bf16"),
284-
("nemotronh", "nemotron_3_5_nano", "pretrain", 16, "h100", "fp8_cs"),
279+
# Nemotron 3.5 Lightning
280+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "bf16"),
281+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "fp8_mx"),
282+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "nvfp4"),
283+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 16, "h100", "bf16"),
284+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 16, "h100", "fp8_cs"),
285285
# Nemotron 3 Super
286286
("nemotronh", "nemotron_3_super", "pretrain", 64, "gb300", "bf16"),
287287
("nemotronh", "nemotron_3_super", "pretrain", 64, "gb300", "fp8_mx"),

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@
1919
nemotronh_56b_pretrain_64gpu_b300_fp8cs_config,
2020
)
2121
from megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh import (
22-
nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config,
23-
nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config,
24-
nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config,
25-
nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config,
22+
nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config,
23+
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config,
24+
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,
25+
nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config,
2626
nemotron_3_nano_pretrain_8gpu_gb200_bf16_config,
2727
nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config,
2828
nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config,
@@ -45,8 +45,8 @@
4545
nemotronh_56b_pretrain_256gpu_gb300_fp8cs_config,
4646
)
4747
from megatron.bridge.perf_recipes.nemotronh.h100.nemotronh import (
48-
nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config,
49-
nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config,
48+
nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config,
49+
nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config,
5050
nemotron_3_nano_pretrain_16gpu_h100_bf16_config,
5151
nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config,
5252
nemotronh_56b_pretrain_64gpu_h100_fp8cs_config,

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

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,9 @@
3333
from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules
3434

3535

36-
# Public Nemotron 3.5 Lightning checkpoint used by the legacy Nano recipe API.
37-
_NEMOTRON_3_5_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
38-
_NEMOTRON_3_5_NANO_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret
36+
# Public Nemotron 3.5 Lightning checkpoint used by the Lightning recipe API.
37+
_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
38+
_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret
3939

4040

4141
def nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config() -> ConfigContainer:
@@ -458,18 +458,18 @@ def nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer:
458458
return cfg
459459

460460

461-
def nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
462-
"""Nemotron 3.5 Nano pretrain: 8× GB200, BF16."""
461+
def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
462+
"""Nemotron 3.5 Lightning pretrain: 8× GB200, BF16."""
463463
cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config()
464464
cfg.model.mtp_num_layers = 2
465465
cfg.model.mtp_hybrid_override_pattern = "*E"
466466
cfg.model.mtp_use_repeated_layer = True
467467
cfg.model.keep_mtp_spec_in_bf16 = True
468468
cfg.model.mtp_loss_scaling_factor = 0.3
469-
cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID
470-
cfg.model.hf_model_revision = _NEMOTRON_3_5_NANO_MODEL_REVISION
471-
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID
472-
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_NANO_MODEL_REVISION}
469+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
470+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
471+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
472+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
473473
cfg.env_vars = {
474474
**COMMON_PERF_ENV_VARS,
475475
"CUDA_DEVICE_MAX_CONNECTIONS": 32,
@@ -489,18 +489,18 @@ def nemotron_3_5_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
489489
return cfg
490490

491491

492-
def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
493-
"""Nemotron 3.5 Nano pretrain: 8× GB200, MXFP8."""
492+
def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
493+
"""Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8."""
494494
cfg = nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config()
495495
cfg.model.mtp_num_layers = 2
496496
cfg.model.mtp_hybrid_override_pattern = "*E"
497497
cfg.model.mtp_use_repeated_layer = True
498498
cfg.model.keep_mtp_spec_in_bf16 = True
499499
cfg.model.mtp_loss_scaling_factor = 0.3
500-
cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID
501-
cfg.model.hf_model_revision = _NEMOTRON_3_5_NANO_MODEL_REVISION
502-
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID
503-
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_NANO_MODEL_REVISION}
500+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
501+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
502+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
503+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
504504
cfg.env_vars = {
505505
**COMMON_PERF_ENV_VARS,
506506
"CUDA_DEVICE_MAX_CONNECTIONS": 32,
@@ -520,9 +520,9 @@ def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
520520
return cfg
521521

522522

523-
def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config() -> ConfigContainer:
524-
"""Nemotron 3.5 Nano pretrain: 8× GB200, MXFP8, Megatron FSDP."""
525-
cfg = nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_config()
523+
def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config() -> ConfigContainer:
524+
"""Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8, Megatron FSDP."""
525+
cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config()
526526

527527
# FSDP reduces the model-state footprint enough to use the larger measured
528528
# microbatch. Megatron FSDP registers module hooks that Transformer Engine
@@ -565,18 +565,18 @@ def nemotron_3_5_nano_pretrain_8gpu_gb200_fp8mx_fsdp_config() -> ConfigContainer
565565
return cfg
566566

567567

568-
def nemotron_3_5_nano_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer:
569-
"""Nemotron 3.5 Nano pretrain: 8× GB200, NVFP4."""
568+
def nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer:
569+
"""Nemotron 3.5 Lightning pretrain: 8× GB200, NVFP4."""
570570
cfg = nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config()
571571
cfg.model.mtp_num_layers = 2
572572
cfg.model.mtp_hybrid_override_pattern = "*E"
573573
cfg.model.mtp_use_repeated_layer = True
574574
cfg.model.keep_mtp_spec_in_bf16 = True
575575
cfg.model.mtp_loss_scaling_factor = 0.3
576-
cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID
577-
cfg.model.hf_model_revision = _NEMOTRON_3_5_NANO_MODEL_REVISION
578-
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID
579-
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_NANO_MODEL_REVISION}
576+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
577+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
578+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
579+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
580580
cfg.env_vars = {
581581
**COMMON_PERF_ENV_VARS,
582582
"CUDA_DEVICE_MAX_CONNECTIONS": 32,

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@
2727
from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules
2828

2929

30-
# Public Nemotron 3.5 Lightning checkpoint used by the legacy Nano recipe API.
31-
_NEMOTRON_3_5_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
32-
_NEMOTRON_3_5_NANO_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret
30+
# Public Nemotron 3.5 Lightning checkpoint used by the Lightning recipe API.
31+
_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
32+
_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret
3333

3434

3535
def nemotronh_56b_pretrain_64gpu_h100_fp8cs_config() -> ConfigContainer:
@@ -175,8 +175,8 @@ def nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer:
175175
return cfg
176176

177177

178-
def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer:
179-
"""Nemotron 3.5 Nano benchmark pretrain: 16× H100, BF16."""
178+
def nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config() -> ConfigContainer:
179+
"""Nemotron 3.5 Lightning benchmark pretrain: 16× H100, BF16."""
180180
cfg = nemotron_3_nano_pretrain_16gpu_h100_bf16_config()
181181
# Keep the benchmark workload aligned with the GB200 BF16 recipe. The
182182
# hardware recipes may tune execution-only knobs such as microbatch size,
@@ -187,10 +187,10 @@ def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer:
187187
cfg.model.mtp_use_repeated_layer = True
188188
cfg.model.keep_mtp_spec_in_bf16 = True
189189
cfg.model.mtp_loss_scaling_factor = 0.3
190-
cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID
191-
cfg.model.hf_model_revision = _NEMOTRON_3_5_NANO_MODEL_REVISION
192-
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID
193-
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_NANO_MODEL_REVISION}
190+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
191+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
192+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
193+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
194194

195195
# This mock-data, force-balanced benchmark uses lower-precision Adam moments
196196
# to make the measured selective-recompute and activation-offload stack fit.
@@ -229,18 +229,18 @@ def nemotron_3_5_nano_pretrain_16gpu_h100_bf16_config() -> ConfigContainer:
229229
return cfg
230230

231231

232-
def nemotron_3_5_nano_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer:
233-
"""Nemotron 3.5 Nano pretrain: 16× H100, FP8 current-scaling."""
232+
def nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config() -> ConfigContainer:
233+
"""Nemotron 3.5 Lightning pretrain: 16× H100, FP8 current-scaling."""
234234
cfg = nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config()
235235
cfg.model.mtp_num_layers = 2
236236
cfg.model.mtp_hybrid_override_pattern = "*E"
237237
cfg.model.mtp_use_repeated_layer = True
238238
cfg.model.keep_mtp_spec_in_bf16 = True
239239
cfg.model.mtp_loss_scaling_factor = 0.3
240-
cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID
241-
cfg.model.hf_model_revision = _NEMOTRON_3_5_NANO_MODEL_REVISION
242-
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID
243-
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_NANO_MODEL_REVISION}
240+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
241+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
242+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
243+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
244244
cfg.env_vars = {
245245
**COMMON_PERF_ENV_VARS,
246246
"CUDA_DEVICE_MAX_CONNECTIONS": 32,

‎src/megatron/bridge/recipes/nemotronh/__init__.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,19 @@
1313
# limitations under the License.
1414

1515
# Nemotron Nano v2 models
16-
# Nemotron 3 and 3.5 Nano models
16+
# Nemotron 3 Nano and Nemotron 3.5 Lightning models
1717
from megatron.bridge.recipes.nemotronh.gb200 import (
18-
nemotron_3_5_nano_pretrain_8k_config,
19-
nemotron_3_5_nano_pretrain_8k_fsdp_config,
20-
nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config,
18+
nemotron_3_5_lightning_pretrain_8k_config,
19+
nemotron_3_5_lightning_pretrain_8k_fsdp_config,
20+
nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config,
2121
nemotron_3_nano_gb200_pretrain_config,
2222
nemotron_3_nano_pretrain_8gpu_gb200_bf16_config,
2323
)
2424
from megatron.bridge.recipes.nemotronh.nemotron_3_nano import (
25-
nemotron_3_5_nano_peft_config,
26-
nemotron_3_5_nano_pretrain_config,
27-
nemotron_3_5_nano_sft_config,
28-
nemotron_3_5_nano_sft_openmathinstruct2_packed_config,
25+
nemotron_3_5_lightning_peft_config,
26+
nemotron_3_5_lightning_pretrain_config,
27+
nemotron_3_5_lightning_sft_config,
28+
nemotron_3_5_lightning_sft_openmathinstruct2_packed_config,
2929
nemotron_3_nano_peft_config,
3030
nemotron_3_nano_pretrain_config,
3131
nemotron_3_nano_sft_config,
@@ -95,14 +95,14 @@
9595
"nemotron_nano_12b_v2_sft_config",
9696
"nemotron_nano_9b_v2_peft_config",
9797
"nemotron_nano_12b_v2_peft_config",
98-
# Nemotron 3 and 3.5 Nano models
99-
"nemotron_3_5_nano_peft_config",
100-
"nemotron_3_5_nano_pretrain_8k_config",
101-
"nemotron_3_5_nano_pretrain_8k_fsdp_config",
102-
"nemotron_3_5_nano_pretrain_config",
103-
"nemotron_3_5_nano_sft_config",
104-
"nemotron_3_5_nano_sft_openmathinstruct2_packed_config",
105-
"nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config",
98+
# Nemotron 3 Nano and Nemotron 3.5 Lightning models
99+
"nemotron_3_5_lightning_peft_config",
100+
"nemotron_3_5_lightning_pretrain_8k_config",
101+
"nemotron_3_5_lightning_pretrain_8k_fsdp_config",
102+
"nemotron_3_5_lightning_pretrain_config",
103+
"nemotron_3_5_lightning_sft_config",
104+
"nemotron_3_5_lightning_sft_openmathinstruct2_packed_config",
105+
"nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config",
106106
"nemotron_3_nano_pretrain_config",
107107
"nemotron_3_nano_sft_config",
108108
"nemotron_3_nano_peft_config",

‎src/megatron/bridge/recipes/nemotronh/gb200/__init__.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,18 +13,18 @@
1313
# limitations under the License.
1414

1515
from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano import (
16-
nemotron_3_5_nano_pretrain_8k_config,
17-
nemotron_3_5_nano_pretrain_8k_fsdp_config,
18-
nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config,
16+
nemotron_3_5_lightning_pretrain_8k_config,
17+
nemotron_3_5_lightning_pretrain_8k_fsdp_config,
18+
nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config,
1919
nemotron_3_nano_gb200_pretrain_config,
2020
nemotron_3_nano_pretrain_8gpu_gb200_bf16_config,
2121
)
2222

2323

2424
__all__ = [
25-
"nemotron_3_5_nano_pretrain_8k_config",
26-
"nemotron_3_5_nano_pretrain_8k_fsdp_config",
27-
"nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config",
25+
"nemotron_3_5_lightning_pretrain_8k_config",
26+
"nemotron_3_5_lightning_pretrain_8k_fsdp_config",
27+
"nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config",
2828
"nemotron_3_nano_gb200_pretrain_config",
2929
"nemotron_3_nano_pretrain_8gpu_gb200_bf16_config",
3030
]

‎src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
from megatron.bridge import AutoBridge
2020
from megatron.bridge.recipes.common import _pretrain_common
2121
from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import (
22-
nemotron_3_5_nano_pretrain_config,
23-
nemotron_3_5_nano_sft_openmathinstruct2_packed_config,
22+
nemotron_3_5_lightning_pretrain_config,
23+
nemotron_3_5_lightning_sft_openmathinstruct2_packed_config,
2424
)
2525
from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS
2626
from megatron.bridge.training.comm_overlap import CommOverlapConfig
@@ -157,9 +157,9 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
157157
return cfg
158158

159159

160-
def nemotron_3_5_nano_pretrain_8k_config() -> ConfigContainer:
161-
"""Return the Nemotron 3.5 Nano BF16 8K pretraining config."""
162-
cfg = nemotron_3_5_nano_pretrain_config()
160+
def nemotron_3_5_lightning_pretrain_8k_config() -> ConfigContainer:
161+
"""Return the Nemotron 3.5 Lightning BF16 8K pretraining config."""
162+
cfg = nemotron_3_5_lightning_pretrain_config()
163163

164164
# Keep the H100 convergence contract intact and change only GB200
165165
# execution/performance settings.
@@ -180,9 +180,9 @@ def nemotron_3_5_nano_pretrain_8k_config() -> ConfigContainer:
180180
return cfg
181181

182182

183-
def nemotron_3_5_nano_pretrain_8k_fsdp_config() -> ConfigContainer:
184-
"""Return the Nemotron 3.5 Nano BF16 8K Megatron FSDP config."""
185-
cfg = nemotron_3_5_nano_pretrain_8k_config()
183+
def nemotron_3_5_lightning_pretrain_8k_fsdp_config() -> ConfigContainer:
184+
"""Return the Nemotron 3.5 Lightning BF16 8K Megatron FSDP config."""
185+
cfg = nemotron_3_5_lightning_pretrain_8k_config()
186186

187187
# Megatron FSDP module hooks are incompatible with Transformer Engine CUDA
188188
# graph capture. Keep all convergence settings inherited from the BF16
@@ -205,9 +205,9 @@ def nemotron_3_5_nano_pretrain_8k_fsdp_config() -> ConfigContainer:
205205
return cfg
206206

207207

208-
def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContainer:
208+
def nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config() -> ConfigContainer:
209209
"""Return the optimized TP1 4K packed OpenMathInstruct-2 SFT config."""
210-
cfg = nemotron_3_5_nano_sft_openmathinstruct2_packed_config()
210+
cfg = nemotron_3_5_lightning_sft_openmathinstruct2_packed_config()
211211

212212
cfg.model.tensor_model_parallel_size = 1
213213
cfg.model.sequence_parallel = False
@@ -245,9 +245,9 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContain
245245

246246

247247
__all__ = [
248-
"nemotron_3_5_nano_pretrain_8k_config",
249-
"nemotron_3_5_nano_pretrain_8k_fsdp_config",
250-
"nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config",
248+
"nemotron_3_5_lightning_pretrain_8k_config",
249+
"nemotron_3_5_lightning_pretrain_8k_fsdp_config",
250+
"nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config",
251251
"nemotron_3_nano_gb200_pretrain_config",
252252
"nemotron_3_nano_pretrain_8gpu_gb200_bf16_config",
253253
]

‎src/megatron/bridge/recipes/nemotronh/h100/__init__.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,10 @@
2525
"nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config",
2626
"nemotron_3_nano_4b_sft_8gpu_h100_bf16_32k_config",
2727
"nemotron_3_nano_4b_sft_8gpu_h100_bf16_config",
28-
"nemotron_3_5_nano_peft_config",
29-
"nemotron_3_5_nano_pretrain_config",
30-
"nemotron_3_5_nano_sft_config",
31-
"nemotron_3_5_nano_sft_openmathinstruct2_packed_config",
28+
"nemotron_3_5_lightning_peft_config",
29+
"nemotron_3_5_lightning_pretrain_config",
30+
"nemotron_3_5_lightning_sft_config",
31+
"nemotron_3_5_lightning_sft_openmathinstruct2_packed_config",
3232
"nemotron_3_nano_peft_8gpu_h100_bf16_config",
3333
"nemotron_3_nano_pretrain_8gpu_h100_bf16_config",
3434
"nemotron_3_nano_sft_8gpu_h100_bf16_config",

0 commit comments

Comments
 (0)