Skip to content

Commit b18ee07

Browse files
yaoyu-33yaoyu-33
committed
fix(perf): restore llama flat recipe parity
1 parent 930f650 commit b18ee07

13 files changed

Lines changed: 242 additions & 194 deletions

File tree

‎src/megatron/bridge/perf_recipes/llama/b200/llama3.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
from megatron.bridge.perf_recipes.llama.common import (
1717
CommOverlapConfig,
1818
ConfigContainer,
19-
_benchmark_common,
2019
_enable_overlap_param_gather_with_optimizer_step,
20+
_llama_benchmark_common,
2121
_perf_precision,
2222
_with_global_batch_size,
2323
llama3_8b_pretrain_config,
@@ -51,7 +51,7 @@ def llama3_8b_pretrain_8gpu_b200_bf16_config() -> ConfigContainer:
5151

5252
cfg.model.moe_token_dispatcher_type = "alltoall"
5353

54-
_benchmark_common(cfg)
54+
_llama_benchmark_common(cfg)
5555
return cfg
5656

5757

@@ -78,7 +78,7 @@ def llama3_8b_pretrain_8gpu_b200_fp8cs_config() -> ConfigContainer:
7878

7979
cfg.model.moe_token_dispatcher_type = "alltoall"
8080

81-
_benchmark_common(cfg)
81+
_llama_benchmark_common(cfg)
8282
return cfg
8383

8484

@@ -103,7 +103,7 @@ def llama3_8b_pretrain_8gpu_b200_fp8mx_config() -> ConfigContainer:
103103

104104
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
105105

106-
_benchmark_common(cfg)
106+
_llama_benchmark_common(cfg)
107107
return cfg
108108

109109

@@ -124,7 +124,7 @@ def llama3_8b_pretrain_8gpu_b200_nvfp4_config() -> ConfigContainer:
124124

125125
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
126126

127-
_benchmark_common(cfg)
127+
_llama_benchmark_common(cfg)
128128
return cfg
129129

130130

@@ -151,7 +151,7 @@ def llama3_70b_pretrain_64gpu_b200_bf16_config() -> ConfigContainer:
151151

152152
cfg.model.moe_token_dispatcher_type = "alltoall"
153153

154-
_benchmark_common(cfg)
154+
_llama_benchmark_common(cfg)
155155
_enable_overlap_param_gather_with_optimizer_step(cfg)
156156
return cfg
157157

@@ -189,7 +189,7 @@ def llama3_70b_pretrain_64gpu_b200_fp8cs_config() -> ConfigContainer:
189189

190190
cfg.model.moe_token_dispatcher_type = "alltoall"
191191

192-
_benchmark_common(cfg)
192+
_llama_benchmark_common(cfg)
193193
return cfg
194194

195195

@@ -210,7 +210,7 @@ def llama3_70b_pretrain_64gpu_b200_fp8mx_config() -> ConfigContainer:
210210

211211
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192
212212

213-
_benchmark_common(cfg)
213+
_llama_benchmark_common(cfg)
214214
return cfg
215215

216216

@@ -232,7 +232,7 @@ def llama3_70b_pretrain_64gpu_b200_nvfp4_config() -> ConfigContainer:
232232
cfg.comm_overlap.tp_comm_overlap = False
233233
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192
234234

235-
_benchmark_common(cfg)
235+
_llama_benchmark_common(cfg)
236236
return cfg
237237

238238

@@ -265,7 +265,7 @@ def llama3_70b_lora_8gpu_b200_bf16_config() -> ConfigContainer:
265265
cfg.dataset.packed_sequence_specs.pad_cu_seqlens = True
266266
cfg.dataset.dataset_kwargs = {"pad_to_max_length": True}
267267

268-
_benchmark_common(cfg)
268+
_llama_benchmark_common(cfg)
269269
return cfg
270270

271271

@@ -287,6 +287,7 @@ def llama3_70b_lora_8gpu_b200_fp8cs_config() -> ConfigContainer:
287287
cfg.model.tensor_model_parallel_size = 1
288288
cfg.model.pipeline_model_parallel_size = 2
289289
cfg.model.context_parallel_size = 1
290+
cfg.model.virtual_pipeline_model_parallel_size = 20
290291
cfg.train.global_batch_size = 32
291292
cfg.train.micro_batch_size = 1
292293

@@ -298,7 +299,8 @@ def llama3_70b_lora_8gpu_b200_fp8cs_config() -> ConfigContainer:
298299
cfg.dataset.packed_sequence_specs.pad_cu_seqlens = True
299300
cfg.dataset.dataset_kwargs = {"pad_to_max_length": True}
300301

301-
_benchmark_common(cfg)
302+
_llama_benchmark_common(cfg)
303+
_enable_overlap_param_gather_with_optimizer_step(cfg)
302304
return cfg
303305

304306

@@ -331,7 +333,7 @@ def llama3_70b_lora_8gpu_b200_fp8mx_config() -> ConfigContainer:
331333
cfg.dataset.packed_sequence_specs.pad_cu_seqlens = True
332334
cfg.dataset.dataset_kwargs = {"pad_to_max_length": True}
333335

334-
_benchmark_common(cfg)
336+
_llama_benchmark_common(cfg)
335337
return cfg
336338

337339

‎src/megatron/bridge/perf_recipes/llama/b200/llama31.py‎

Lines changed: 27 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@
1515

1616
from megatron.bridge.perf_recipes.llama.common import (
1717
ConfigContainer,
18-
_benchmark_common,
18+
_enable_overlap_param_gather_with_optimizer_step,
19+
_llama_benchmark_common,
1920
_perf_precision,
2021
llama31_405b_pretrain_config,
2122
userbuffers_bf16_b200_h16384_tp4_cp2_mbs1_seqlen8192,
@@ -39,18 +40,19 @@ def llama31_405b_pretrain_128gpu_b200_bf16_config() -> ConfigContainer:
3940
cfg.dataset.seq_length = 8192
4041

4142
cfg.model.tensor_model_parallel_size = 4
42-
cfg.model.pipeline_model_parallel_size = 8
43-
cfg.model.context_parallel_size = 2
43+
cfg.model.pipeline_model_parallel_size = 16
44+
cfg.model.context_parallel_size = 1
4445
cfg.model.virtual_pipeline_model_parallel_size = 8
4546
cfg.model.sequence_parallel = True
46-
cfg.train.global_batch_size = 64
47+
cfg.train.global_batch_size = 768
4748
cfg.train.micro_batch_size = 1
4849

4950
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_bf16_b200_h16384_tp4_cp2_mbs1_seqlen8192
5051

5152
cfg.model.moe_token_dispatcher_type = "alltoall"
5253

53-
_benchmark_common(cfg)
54+
_llama_benchmark_common(cfg)
55+
_enable_overlap_param_gather_with_optimizer_step(cfg)
5456
return cfg
5557

5658

@@ -64,18 +66,19 @@ def llama31_405b_pretrain_128gpu_b200_fp8cs_config() -> ConfigContainer:
6466
cfg.dataset.seq_length = 8192
6567

6668
cfg.model.tensor_model_parallel_size = 4
67-
cfg.model.pipeline_model_parallel_size = 8
68-
cfg.model.context_parallel_size = 2
69-
cfg.model.virtual_pipeline_model_parallel_size = 8
69+
cfg.model.pipeline_model_parallel_size = 16
70+
cfg.model.context_parallel_size = 1
71+
cfg.model.virtual_pipeline_model_parallel_size = 4
7072
cfg.model.sequence_parallel = True
71-
cfg.train.global_batch_size = 64
73+
cfg.train.global_batch_size = 768
7274
cfg.train.micro_batch_size = 1
7375

7476
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h16384_tp4_cp2_mbs1_seqlen8192
7577

7678
cfg.model.moe_token_dispatcher_type = "alltoall"
7779

78-
_benchmark_common(cfg)
80+
_llama_benchmark_common(cfg)
81+
_enable_overlap_param_gather_with_optimizer_step(cfg)
7982
return cfg
8083

8184

@@ -89,16 +92,16 @@ def llama31_405b_pretrain_128gpu_b200_fp8mx_config() -> ConfigContainer:
8992
cfg.dataset.seq_length = 8192
9093

9194
cfg.model.tensor_model_parallel_size = 4
92-
cfg.model.pipeline_model_parallel_size = 8
93-
cfg.model.context_parallel_size = 2
95+
cfg.model.pipeline_model_parallel_size = 16
96+
cfg.model.context_parallel_size = 1
9497
cfg.model.virtual_pipeline_model_parallel_size = 8
9598
cfg.model.sequence_parallel = True
96-
cfg.train.global_batch_size = 64
99+
cfg.train.global_batch_size = 768
97100
cfg.train.micro_batch_size = 1
98101

99102
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h16384_tp4_cp2_mbs1_seqlen8192
100103

101-
_benchmark_common(cfg)
104+
_llama_benchmark_common(cfg)
102105
return cfg
103106

104107

@@ -116,19 +119,15 @@ def llama31_405b_pretrain_128gpu_b200_nvfp4_config() -> ConfigContainer:
116119
cfg.model.context_parallel_size = 1
117120
cfg.model.virtual_pipeline_model_parallel_size = 8
118121
cfg.model.sequence_parallel = True
119-
cfg.train.global_batch_size = 64
122+
cfg.train.global_batch_size = 768
120123
cfg.train.micro_batch_size = 1
121124

122125
cfg.model.cuda_graph_impl = "none"
123126

124-
cfg.model.recompute_granularity = "full"
125-
cfg.model.recompute_method = "block"
126-
cfg.model.recompute_num_layers = 1
127-
128127
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h16384_tp4_cp2_mbs1_seqlen8192
129128
cfg.comm_overlap.tp_comm_overlap = False
130129

131-
_benchmark_common(cfg)
130+
_llama_benchmark_common(cfg)
132131
return cfg
133132

134133

@@ -138,4 +137,11 @@ def llama31_405b_pretrain_128gpu_b200_nvfp4_config() -> ConfigContainer:
138137

139138
llama31_405b_pretrain_256gpu_b200_fp8mx_config = llama31_405b_pretrain_256gpu_gb200_fp8mx_config
140139

141-
llama31_405b_pretrain_256gpu_b200_nvfp4_config = llama31_405b_pretrain_256gpu_gb200_nvfp4_config
140+
def llama31_405b_pretrain_256gpu_b200_nvfp4_config() -> ConfigContainer:
141+
"""Llama3.1 405B pretrain: 256x B200, NVFP4."""
142+
cfg = llama31_405b_pretrain_256gpu_gb200_nvfp4_config()
143+
cfg.ddp.overlap_param_gather = True
144+
cfg.optimizer.overlap_param_gather = False
145+
cfg.comm_overlap.overlap_param_gather = None
146+
cfg.comm_overlap.align_param_gather = None
147+
return cfg

‎src/megatron/bridge/perf_recipes/llama/b300/llama3.py‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from megatron.bridge.perf_recipes.llama.common import (
1717
CommOverlapConfig,
1818
ConfigContainer,
19-
_benchmark_common,
19+
_llama_benchmark_common,
2020
_perf_precision,
2121
llama3_8b_pretrain_config,
2222
llama3_70b_peft_config,
@@ -47,7 +47,7 @@ def llama3_8b_pretrain_8gpu_b300_bf16_config() -> ConfigContainer:
4747

4848
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
4949

50-
_benchmark_common(cfg)
50+
_llama_benchmark_common(cfg)
5151
return cfg
5252

5353

@@ -72,7 +72,7 @@ def llama3_8b_pretrain_8gpu_b300_fp8cs_config() -> ConfigContainer:
7272

7373
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
7474

75-
_benchmark_common(cfg)
75+
_llama_benchmark_common(cfg)
7676
return cfg
7777

7878

@@ -97,7 +97,7 @@ def llama3_8b_pretrain_8gpu_b300_fp8mx_config() -> ConfigContainer:
9797

9898
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
9999

100-
_benchmark_common(cfg)
100+
_llama_benchmark_common(cfg)
101101
return cfg
102102

103103

@@ -118,7 +118,7 @@ def llama3_8b_pretrain_8gpu_b300_nvfp4_config() -> ConfigContainer:
118118

119119
cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
120120

121-
_benchmark_common(cfg)
121+
_llama_benchmark_common(cfg)
122122
return cfg
123123

124124

@@ -149,7 +149,7 @@ def llama3_70b_pretrain_64gpu_b300_bf16_config() -> ConfigContainer:
149149

150150
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_bf16_b200_h8192_tp2_mbs1_seqlen8192
151151

152-
_benchmark_common(cfg)
152+
_llama_benchmark_common(cfg)
153153
return cfg
154154

155155

@@ -180,7 +180,7 @@ def llama3_70b_pretrain_64gpu_b300_fp8cs_config() -> ConfigContainer:
180180

181181
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192
182182

183-
_benchmark_common(cfg)
183+
_llama_benchmark_common(cfg)
184184
return cfg
185185

186186

@@ -201,7 +201,7 @@ def llama3_70b_pretrain_64gpu_b300_fp8mx_config() -> ConfigContainer:
201201

202202
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192
203203

204-
_benchmark_common(cfg)
204+
_llama_benchmark_common(cfg)
205205
return cfg
206206

207207

@@ -227,7 +227,7 @@ def llama3_70b_pretrain_64gpu_b300_nvfp4_config() -> ConfigContainer:
227227
cfg.comm_overlap.tp_comm_overlap = False
228228
cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192
229229

230-
_benchmark_common(cfg)
230+
_llama_benchmark_common(cfg)
231231
return cfg
232232

233233

@@ -260,7 +260,7 @@ def llama3_70b_lora_8gpu_b300_bf16_config() -> ConfigContainer:
260260
cfg.dataset.packed_sequence_specs.pad_cu_seqlens = True
261261
cfg.dataset.dataset_kwargs = {"pad_to_max_length": True}
262262

263-
_benchmark_common(cfg)
263+
_llama_benchmark_common(cfg)
264264
return cfg
265265

266266

@@ -293,7 +293,7 @@ def llama3_70b_lora_8gpu_b300_fp8cs_config() -> ConfigContainer:
293293
cfg.dataset.packed_sequence_specs.pad_cu_seqlens = True
294294
cfg.dataset.dataset_kwargs = {"pad_to_max_length": True}
295295

296-
_benchmark_common(cfg)
296+
_llama_benchmark_common(cfg)
297297
return cfg
298298

299299

@@ -326,5 +326,5 @@ def llama3_70b_lora_8gpu_b300_fp8mx_config() -> ConfigContainer:
326326
cfg.dataset.packed_sequence_specs.pad_cu_seqlens = True
327327
cfg.dataset.dataset_kwargs = {"pad_to_max_length": True}
328328

329-
_benchmark_common(cfg)
329+
_llama_benchmark_common(cfg)
330330
return cfg

0 commit comments

Comments
 (0)