Skip to content

Commit bf07e56

Browse files
authored
vr200 cfgs match w/ gb300 (#4574)
Signed-off-by: Malay Nagda <malayn@nvidia.com>
1 parent 59ae409 commit bf07e56

14 files changed

Lines changed: 141 additions & 274 deletions

File tree

scripts/performance/configs/deepseek/deepseek_workload_base_configs.py

Lines changed: 9 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -108,25 +108,11 @@
108108
DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_NVFP4_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_V1
109109

110110

111-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_V1 = replace(
112-
BASE_DEEPSEEK_V3_CONFIG,
113-
num_gpus=128,
114-
global_batch_size=4096,
115-
micro_batch_size=1,
116-
pipeline_model_parallel_size=2,
117-
virtual_pipeline_model_parallel_size=8,
118-
pp_layout="Et*4|(t*4|)*14tmL",
119-
expert_model_parallel_size=64,
120-
moe_flex_dispatcher_backend="hybridep",
121-
moe_a2a_overlap=False,
122-
cuda_graph_impl="transformer_engine",
123-
cuda_graph_scope=["attn", "moe_router", "moe_preprocess"],
124-
recompute_modules=["mla_up_proj"],
125-
)
126-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_BF16_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_V1
127-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_CS_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_V1
128-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_V1
129-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_NVFP4_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_V1
111+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_V1
112+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_BF16_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_BF16_V1
113+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_CS_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_FP8_CS_V1
114+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_FP8_MX_V1
115+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_NVFP4_V1 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_NVFP4_V1
130116

131117

132118
DEEPSEEK_V3_PRETRAIN_CONFIG_B300_V1 = replace(
@@ -230,10 +216,10 @@
230216
DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_NVFP4_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_V2
231217

232218

233-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_BF16_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_BF16_V2
234-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_CS_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_FP8_CS_V2
235-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_FP8_MX_V2
236-
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_NVFP4_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB200_NVFP4_V2
219+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_BF16_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_BF16_V2
220+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_CS_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_FP8_CS_V2
221+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_FP8_MX_V2
222+
DEEPSEEK_V3_PRETRAIN_CONFIG_VR200_NVFP4_V2 = DEEPSEEK_V3_PRETRAIN_CONFIG_GB300_NVFP4_V2
237223

238224

239225
DEEPSEEK_V3_PRETRAIN_CONFIG_B300_V2 = replace(

scripts/performance/configs/gpt_oss/gpt_oss_workload_base_configs.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -106,9 +106,9 @@
106106
GPT_OSS_20B_PRETRAIN_CONFIG_GB300_NVFP4_V2 = GPT_OSS_20B_PRETRAIN_CONFIG_GB200_NVFP4_V2
107107
GPT_OSS_20B_PRETRAIN_CONFIG_GB300_FP8_MX_V2 = GPT_OSS_20B_PRETRAIN_CONFIG_GB200_NVFP4_V2
108108
GPT_OSS_20B_PRETRAIN_CONFIG_GB300_FP8_MX_V3 = GPT_OSS_20B_PRETRAIN_CONFIG_GB200_FP8_MX_V3
109-
GPT_OSS_20B_PRETRAIN_CONFIG_VR200_NVFP4_V1 = GPT_OSS_20B_PRETRAIN_CONFIG_B300_FP8_MX_V1
110-
GPT_OSS_20B_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = GPT_OSS_20B_PRETRAIN_CONFIG_VR200_NVFP4_V1
111-
GPT_OSS_20B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = GPT_OSS_20B_PRETRAIN_CONFIG_B300_NVFP4_V2
109+
GPT_OSS_20B_PRETRAIN_CONFIG_VR200_NVFP4_V1 = GPT_OSS_20B_PRETRAIN_CONFIG_GB300_NVFP4_V1
110+
GPT_OSS_20B_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = GPT_OSS_20B_PRETRAIN_CONFIG_GB300_FP8_MX_V1
111+
GPT_OSS_20B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = GPT_OSS_20B_PRETRAIN_CONFIG_GB300_NVFP4_V2
112112

113113

114114
# GPT-OSS 120B Pretrain - V1 (GBS=512)
@@ -223,8 +223,8 @@
223223
GPT_OSS_120B_PRETRAIN_CONFIG_B200_FP8_MX_V2 = GPT_OSS_120B_PRETRAIN_CONFIG_B200_BF16_V2
224224
GPT_OSS_120B_PRETRAIN_CONFIG_H100_FP8_CS_V2 = GPT_OSS_120B_PRETRAIN_CONFIG_H100_BF16_V2
225225

226-
GPT_OSS_120B_PRETRAIN_CONFIG_VR200_BF16_V2 = GPT_OSS_120B_PRETRAIN_CONFIG_GB200_BF16_V2
227-
GPT_OSS_120B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = GPT_OSS_120B_PRETRAIN_CONFIG_GB200_FP8_MX_V2
226+
GPT_OSS_120B_PRETRAIN_CONFIG_VR200_BF16_V2 = GPT_OSS_120B_PRETRAIN_CONFIG_GB300_BF16_V2
227+
GPT_OSS_120B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = GPT_OSS_120B_PRETRAIN_CONFIG_GB300_FP8_MX_V2
228228

229229

230230
__all__ = [

scripts/performance/configs/kimi/kimi_workload_base_configs.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,8 @@
6060
KIMI_K2_PRETRAIN_CONFIG_GB200_FP8_CS = KIMI_K2_PRETRAIN_CONFIG_GB200
6161
KIMI_K2_PRETRAIN_CONFIG_GB200_FP8_MX = KIMI_K2_PRETRAIN_CONFIG_GB200
6262

63-
KIMI_K2_PRETRAIN_CONFIG_VR200_BF16 = KIMI_K2_PRETRAIN_CONFIG_GB200_BF16
64-
KIMI_K2_PRETRAIN_CONFIG_VR200_FP8_MX = KIMI_K2_PRETRAIN_CONFIG_GB200_FP8_MX
63+
KIMI_K2_PRETRAIN_CONFIG_VR200_BF16 = KIMI_K2_PRETRAIN_CONFIG_GB300_BF16
64+
KIMI_K2_PRETRAIN_CONFIG_VR200_FP8_MX = KIMI_K2_PRETRAIN_CONFIG_GB300_FP8_MX
6565

6666

6767
KIMI_K2_PRETRAIN_CONFIG_B300 = replace(

scripts/performance/configs/llama/llama31_workload_base_configs.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -289,9 +289,9 @@
289289
)
290290

291291

292-
LLAMA31_405B_PRETRAIN_CONFIG_VR200_BF16_V2 = LLAMA31_405B_PRETRAIN_CONFIG_GB200_BF16_V2
293-
LLAMA31_405B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = LLAMA31_405B_PRETRAIN_CONFIG_GB200_FP8_MX_V2
294-
LLAMA31_405B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = LLAMA31_405B_PRETRAIN_CONFIG_GB200_NVFP4_V2
292+
LLAMA31_405B_PRETRAIN_CONFIG_VR200_BF16_V2 = LLAMA31_405B_PRETRAIN_CONFIG_GB300_BF16_V2
293+
LLAMA31_405B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = LLAMA31_405B_PRETRAIN_CONFIG_GB300_FP8_MX_V2
294+
LLAMA31_405B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = LLAMA31_405B_PRETRAIN_CONFIG_GB300_NVFP4_V2
295295

296296

297297
LLAMA31_405B_PRETRAIN_CONFIG_B300_BF16_V2 = LLAMA31_405B_PRETRAIN_CONFIG_GB300_BF16_V2

scripts/performance/configs/llama/llama3_workload_base_configs.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -264,9 +264,9 @@
264264
)
265265

266266

267-
LLAMA3_70B_PRETRAIN_CONFIG_VR200_BF16_V2 = LLAMA3_70B_PRETRAIN_CONFIG_GB200_BF16_V2
268-
LLAMA3_70B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = LLAMA3_70B_PRETRAIN_CONFIG_GB200_FP8_MX_V2
269-
LLAMA3_70B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = LLAMA3_70B_PRETRAIN_CONFIG_GB200_NVFP4_V2
267+
LLAMA3_70B_PRETRAIN_CONFIG_VR200_BF16_V2 = LLAMA3_70B_PRETRAIN_CONFIG_GB300_BF16_V2
268+
LLAMA3_70B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = LLAMA3_70B_PRETRAIN_CONFIG_GB300_FP8_MX_V2
269+
LLAMA3_70B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = LLAMA3_70B_PRETRAIN_CONFIG_GB300_NVFP4_V2
270270

271271

272272
LLAMA3_70B_PRETRAIN_CONFIG_B300_BF16_V2 = replace(
@@ -411,10 +411,10 @@
411411
cuda_graph_scope="full_iteration",
412412
)
413413

414-
LLAMA3_8B_PRETRAIN_CONFIG_VR200_BF16_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB200_BF16_V1
415-
LLAMA3_8B_PRETRAIN_CONFIG_VR200_FP8_CS_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB200_FP8_CS_V1
416-
LLAMA3_8B_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB200_FP8_MX_V1
417-
LLAMA3_8B_PRETRAIN_CONFIG_VR200_NVFP4_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB200_NVFP4_V1
414+
LLAMA3_8B_PRETRAIN_CONFIG_VR200_BF16_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB300_BF16_V1
415+
LLAMA3_8B_PRETRAIN_CONFIG_VR200_FP8_CS_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB300_FP8_CS_V1
416+
LLAMA3_8B_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB300_FP8_MX_V1
417+
LLAMA3_8B_PRETRAIN_CONFIG_VR200_NVFP4_V1 = LLAMA3_8B_PRETRAIN_CONFIG_GB300_NVFP4_V1
418418

419419
LLAMA3_8B_PRETRAIN_CONFIG_B300_BF16_V1 = replace(
420420
BASE_LLAMA3_8B_CONFIG,

scripts/performance/configs/nemotronh/nemotron_3_workload_base_configs.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -53,9 +53,9 @@
5353
NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB200_FP8_MX_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB200_BF16_V1
5454
NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB200_NVFP4_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB200_BF16_V1
5555

56-
NEMOTRON_3_NANO_PRETRAIN_CONFIG_VR200_BF16_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB200_BF16_V1
57-
NEMOTRON_3_NANO_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB200_FP8_MX_V1
58-
NEMOTRON_3_NANO_PRETRAIN_CONFIG_VR200_NVFP4_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB200_NVFP4_V1
56+
NEMOTRON_3_NANO_PRETRAIN_CONFIG_VR200_BF16_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB300_BF16_V1
57+
NEMOTRON_3_NANO_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB300_FP8_MX_V1
58+
NEMOTRON_3_NANO_PRETRAIN_CONFIG_VR200_NVFP4_V1 = NEMOTRON_3_NANO_PRETRAIN_CONFIG_GB300_NVFP4_V1
5959

6060
NEMOTRON_3_NANO_PRETRAIN_CONFIG_B300_BF16_V1 = replace(
6161
BASE_NEMOTRON_3_NANO_CONFIG,
@@ -121,9 +121,9 @@
121121
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB200_FP8_MX_V1 = BASE_NEMOTRON_3_SUPER_CONFIG_GB200
122122
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB200_NVFP4_V1 = BASE_NEMOTRON_3_SUPER_CONFIG_GB200
123123

124-
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_VR200_BF16_V1 = NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB200_BF16_V1
125-
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB200_FP8_MX_V1
126-
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_VR200_NVFP4_V1 = NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB200_NVFP4_V1
124+
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_VR200_BF16_V1 = NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB300_BF16_V1
125+
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB300_FP8_MX_V1
126+
NEMOTRON_3_SUPER_PRETRAIN_CONFIG_VR200_NVFP4_V1 = NEMOTRON_3_SUPER_PRETRAIN_CONFIG_GB300_NVFP4_V1
127127

128128

129129
BASE_NEMOTRON_3_SUPER_CONFIG_B300 = replace(

scripts/performance/configs/qwen/qwen3_workload_base_configs.py

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -238,12 +238,9 @@
238238
QWEN3_235B_A22B_PRETRAIN_CONFIG_GB200_NVFP4_V2 = QWEN3_235B_A22B_PRETRAIN_CONFIG_GB200_FP8_CS_V2
239239

240240

241-
QWEN3_235B_A22B_PRETRAIN_CONFIG_VR200_BF16_V2 = QWEN3_235B_A22B_PRETRAIN_CONFIG_GB200_BF16_V2
242-
QWEN3_235B_A22B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = replace(
243-
QWEN3_235B_A22B_PRETRAIN_CONFIG_GB200_FP8_CS_V2,
244-
virtual_pipeline_model_parallel_size=3,
245-
)
246-
QWEN3_235B_A22B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = QWEN3_235B_A22B_PRETRAIN_CONFIG_GB200_NVFP4_V2
241+
QWEN3_235B_A22B_PRETRAIN_CONFIG_VR200_BF16_V2 = QWEN3_235B_A22B_PRETRAIN_CONFIG_GB300_BF16_V2
242+
QWEN3_235B_A22B_PRETRAIN_CONFIG_VR200_FP8_MX_V2 = QWEN3_235B_A22B_PRETRAIN_CONFIG_GB300_FP8_MX_V2
243+
QWEN3_235B_A22B_PRETRAIN_CONFIG_VR200_NVFP4_V2 = QWEN3_235B_A22B_PRETRAIN_CONFIG_GB300_NVFP4_V2
247244

248245

249246
QWEN3_235B_A22B_PRETRAIN_CONFIG_B300_BF16_V2 = replace(
@@ -398,15 +395,8 @@
398395
)
399396

400397

401-
QWEN3_30B_A3B_PRETRAIN_CONFIG_VR200_BF16_V1 = QWEN3_30B_A3B_PRETRAIN_CONFIG_GB200_BF16_V1
402-
QWEN3_30B_A3B_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = replace(
403-
BASE_QWEN3_30B_A3B_CONFIG,
404-
num_gpus=8,
405-
micro_batch_size=4,
406-
moe_flex_dispatcher_backend="hybridep",
407-
cuda_graph_impl="transformer_engine",
408-
cuda_graph_scope=["moe_router", "moe_preprocess"],
409-
)
398+
QWEN3_30B_A3B_PRETRAIN_CONFIG_VR200_BF16_V1 = QWEN3_30B_A3B_PRETRAIN_CONFIG_GB300_BF16_V1
399+
QWEN3_30B_A3B_PRETRAIN_CONFIG_VR200_FP8_MX_V1 = QWEN3_30B_A3B_PRETRAIN_CONFIG_GB300_FP8_MX_V1
410400

411401

412402
QWEN3_30B_A3B_PRETRAIN_CONFIG_B300_BF16_V1 = replace(

src/megatron/bridge/perf_recipes/deepseek/vr200/deepseek_v3.py

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@
2323
deepseek_v3_pretrain_config,
2424
set_deepseek_v3_pipeline_model_parallel_layout,
2525
)
26-
from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v3 import (
27-
deepseek_v3_pretrain_256gpu_gb200_bf16_config,
28-
deepseek_v3_pretrain_256gpu_gb200_fp8cs_config,
29-
deepseek_v3_pretrain_256gpu_gb200_fp8mx_config,
30-
deepseek_v3_pretrain_256gpu_gb200_nvfp4_config,
26+
from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v3 import (
27+
deepseek_v3_pretrain_256gpu_gb300_bf16_config,
28+
deepseek_v3_pretrain_256gpu_gb300_fp8cs_config,
29+
deepseek_v3_pretrain_256gpu_gb300_fp8mx_config,
30+
deepseek_v3_pretrain_256gpu_gb300_nvfp4_config,
3131
)
3232

3333

@@ -148,22 +148,20 @@ def deepseek_v3_pretrain_128gpu_vr200_nvfp4_config() -> ConfigContainer:
148148

149149

150150
def deepseek_v3_pretrain_256gpu_vr200_bf16_config() -> ConfigContainer:
151-
"""DeepSeek V3 pretrain: 256× VR200, BF16 (alias of GB200)."""
152-
return deepseek_v3_pretrain_256gpu_gb200_bf16_config()
151+
"""DeepSeek V3 pretrain: 256× VR200, BF16 (alias of GB300)."""
152+
return deepseek_v3_pretrain_256gpu_gb300_bf16_config()
153153

154154

155155
def deepseek_v3_pretrain_256gpu_vr200_fp8cs_config() -> ConfigContainer:
156-
"""DeepSeek V3 pretrain: 256× VR200, FP8-CS (alias of GB200)."""
157-
return deepseek_v3_pretrain_256gpu_gb200_fp8cs_config()
156+
"""DeepSeek V3 pretrain: 256× VR200, FP8-CS (alias of GB300)."""
157+
return deepseek_v3_pretrain_256gpu_gb300_fp8cs_config()
158158

159159

160160
def deepseek_v3_pretrain_256gpu_vr200_fp8mx_config() -> ConfigContainer:
161-
"""DeepSeek V3 pretrain: 256× VR200, FP8-MX (alias of GB200)."""
162-
cfg = deepseek_v3_pretrain_256gpu_gb200_fp8mx_config()
163-
cfg.model.fp8_output_proj = False
164-
return cfg
161+
"""DeepSeek V3 pretrain: 256× VR200, FP8-MX (alias of GB300)."""
162+
return deepseek_v3_pretrain_256gpu_gb300_fp8mx_config()
165163

166164

167165
def deepseek_v3_pretrain_256gpu_vr200_nvfp4_config() -> ConfigContainer:
168-
"""DeepSeek V3 pretrain: 256× VR200, NVFP4 (alias of GB200)."""
169-
return deepseek_v3_pretrain_256gpu_gb200_nvfp4_config()
166+
"""DeepSeek V3 pretrain: 256× VR200, NVFP4 (alias of GB300)."""
167+
return deepseek_v3_pretrain_256gpu_gb300_nvfp4_config()

src/megatron/bridge/perf_recipes/gpt_oss/vr200/gpt_oss.py

Lines changed: 11 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -14,43 +14,24 @@
1414
"""VR200 performance recipes for GPT-OSS."""
1515

1616
from megatron.bridge.perf_recipes.gpt_oss.common import (
17-
CommOverlapConfig,
1817
ConfigContainer,
1918
_apply_gpt_oss_20b_common_configs,
2019
_apply_gpt_oss_20b_transformer_engine_graph_configs,
2120
_benchmark_common,
2221
_gpt_oss_20b_fp8mx_precision,
2322
_gpt_oss_20b_nvfp4_precision,
24-
_perf_precision,
2523
gpt_oss_20b_pretrain_config,
26-
gpt_oss_120b_pretrain_config,
24+
)
25+
from megatron.bridge.perf_recipes.gpt_oss.gb300.gpt_oss import (
26+
gpt_oss_20b_pretrain_8gpu_gb300_nvfp4_config,
27+
gpt_oss_120b_pretrain_64gpu_gb300_bf16_config,
28+
gpt_oss_120b_pretrain_64gpu_gb300_fp8mx_config,
2729
)
2830

2931

3032
def gpt_oss_20b_pretrain_8gpu_vr200_nvfp4_config() -> ConfigContainer:
31-
"""GPT-OSS 20B pretrain: 8× VR200, NVFP4."""
32-
cfg = gpt_oss_20b_pretrain_config()
33-
cfg.mixed_precision = _gpt_oss_20b_nvfp4_precision()
34-
35-
cfg.model.tensor_model_parallel_size = 1
36-
cfg.model.context_parallel_size = 2
37-
cfg.model.expert_model_parallel_size = 4
38-
cfg.model.expert_tensor_parallel_size = 1
39-
cfg.model.sequence_parallel = False
40-
cfg.train.global_batch_size = 4
41-
cfg.train.micro_batch_size = 1
42-
43-
_benchmark_common(cfg)
44-
_apply_gpt_oss_20b_common_configs(cfg)
45-
_apply_gpt_oss_20b_transformer_engine_graph_configs(cfg)
46-
47-
cfg.model.cuda_graph_warmup_steps = 1
48-
cfg.optimizer.lr = 0.0006
49-
cfg.optimizer.min_lr = 0.0006
50-
cfg.validation.eval_interval = 384
51-
cfg.validation.eval_iters = 43
52-
cfg.scheduler.lr_warmup_iters = 64
53-
return cfg
33+
"""GPT-OSS 20B pretrain: 8× VR200, NVFP4 (alias of GB300)."""
34+
return gpt_oss_20b_pretrain_8gpu_gb300_nvfp4_config()
5435

5536

5637
def gpt_oss_20b_pretrain_8gpu_vr200_fp8mx_config() -> ConfigContainer:
@@ -106,65 +87,10 @@ def gpt_oss_20b_pretrain_64gpu_vr200_nvfp4_config() -> ConfigContainer:
10687

10788

10889
def gpt_oss_120b_pretrain_64gpu_vr200_bf16_config() -> ConfigContainer:
109-
"""GPT-OSS 120B pretrain: 64× VR200, BF16."""
110-
cfg = gpt_oss_120b_pretrain_config()
111-
cfg.mixed_precision = _perf_precision("bf16")
112-
cfg.model.moe_router_fusion = True
113-
cfg.model.moe_router_force_load_balancing = True
114-
115-
cfg.model.tensor_model_parallel_size = 1
116-
cfg.model.pipeline_model_parallel_size = 1
117-
cfg.model.expert_model_parallel_size = 64
118-
cfg.model.sequence_parallel = False
119-
cfg.train.global_batch_size = 1280
120-
cfg.train.micro_batch_size = 4
121-
122-
cfg.model.recompute_modules = ["layernorm", "moe_act"]
123-
cfg.model.recompute_granularity = "selective"
124-
125-
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
126-
cfg.comm_overlap.delay_wgrad_compute = True
127-
cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
128-
129-
_benchmark_common(cfg)
130-
return cfg
90+
"""GPT-OSS 120B pretrain: 64× VR200, BF16 (alias of GB300)."""
91+
return gpt_oss_120b_pretrain_64gpu_gb300_bf16_config()
13192

13293

13394
def gpt_oss_120b_pretrain_64gpu_vr200_fp8mx_config() -> ConfigContainer:
134-
"""GPT-OSS 120B pretrain: 64× VR200, FP8-MX."""
135-
cfg = gpt_oss_120b_pretrain_config()
136-
cfg.mixed_precision = _perf_precision("fp8_mx")
137-
cfg.model.moe_router_fusion = True
138-
cfg.model.moe_router_force_load_balancing = True
139-
140-
cfg.model.tensor_model_parallel_size = 1
141-
cfg.model.pipeline_model_parallel_size = 1
142-
cfg.model.expert_model_parallel_size = 64
143-
cfg.model.sequence_parallel = False
144-
cfg.train.global_batch_size = 1280
145-
cfg.train.micro_batch_size = 4
146-
147-
cfg.model.recompute_modules = ["layernorm", "moe_act"]
148-
cfg.model.recompute_granularity = "selective"
149-
150-
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
151-
cfg.comm_overlap.delay_wgrad_compute = True
152-
cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
153-
154-
_benchmark_common(cfg)
155-
156-
cfg.model.cuda_graph_impl = "full_iteration"
157-
cfg.model.cuda_graph_scope = []
158-
cfg.model.moe_flex_dispatcher_backend = "hybridep"
159-
cfg.model.moe_token_dispatcher_type = "flex"
160-
cfg.model.high_priority_a2a_comm_stream = True
161-
cfg.model.moe_hybridep_num_sms = 32
162-
cfg.model.moe_hybridep_num_sms_preprocessing = 32
163-
cfg.model.moe_mlp_glu_interleave_size = 32
164-
cfg.model.use_te_rng_tracker = True
165-
cfg.model.use_transformer_engine_op_fuser = True
166-
cfg.mixed_precision.fp8_dot_product_attention = True
167-
cfg.rng.te_rng_tracker = True
168-
cfg.model.recompute_granularity = "selective"
169-
cfg.model.recompute_modules = []
170-
return cfg
95+
"""GPT-OSS 120B pretrain: 64× VR200, FP8-MX (alias of GB300)."""
96+
return gpt_oss_120b_pretrain_64gpu_gb300_fp8mx_config()

src/megatron/bridge/perf_recipes/kimi/vr200/kimi_k2.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,17 @@
1616
from megatron.bridge.perf_recipes.kimi.common import (
1717
ConfigContainer,
1818
)
19-
from megatron.bridge.perf_recipes.kimi.gb200.kimi_k2 import (
20-
kimi_k2_pretrain_256gpu_gb200_bf16_config,
21-
kimi_k2_pretrain_256gpu_gb200_fp8mx_config,
19+
from megatron.bridge.perf_recipes.kimi.gb300.kimi_k2 import (
20+
kimi_k2_pretrain_256gpu_gb300_bf16_config,
21+
kimi_k2_pretrain_256gpu_gb300_fp8mx_config,
2222
)
2323

2424

2525
def kimi_k2_pretrain_256gpu_vr200_bf16_config() -> ConfigContainer:
26-
"""Kimi K2 pretrain: 256× VR200, BF16 (alias of GB200)."""
27-
return kimi_k2_pretrain_256gpu_gb200_bf16_config()
26+
"""Kimi K2 pretrain: 256× VR200, BF16 (alias of GB300)."""
27+
return kimi_k2_pretrain_256gpu_gb300_bf16_config()
2828

2929

3030
def kimi_k2_pretrain_256gpu_vr200_fp8mx_config() -> ConfigContainer:
31-
"""Kimi K2 pretrain: 256× VR200, MXFP8 (alias of GB200)."""
32-
return kimi_k2_pretrain_256gpu_gb200_fp8mx_config()
31+
"""Kimi K2 pretrain: 256× VR200, MXFP8 (alias of GB300)."""
32+
return kimi_k2_pretrain_256gpu_gb300_fp8mx_config()

0 commit comments

Comments
 (0)