Skip to content

Commit 84fba28

Browse files
author
mollys
committed
add llama2 70b peft recipe
1 parent 6124fbe commit 84fba28

8 files changed

Lines changed: 689 additions & 2 deletions

File tree

‎scripts/performance/argument_parser.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -695,7 +695,7 @@ def parse_cli_args():
695695
"-c",
696696
"--compute_dtype",
697697
type=str,
698-
choices=["bf16", "fp8_cs", "fp8_mx", "fp8_sc", "nvfp4"],
698+
choices=["bf16", "fp8_ds","fp8_cs", "fp8_mx", "fp8_sc", "nvfp4"],
699699
help="Compute precision. Options- bf16 or fp8. Defaults to bf16",
700700
required=False,
701701
default="bf16",

‎scripts/performance/utils/precision.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from megatron.bridge.training.comm_overlap import *
1717
from megatron.bridge.training.mixed_precision import (
1818
bf16_mixed,
19+
bf16_with_fp8_delayed_scaling_mixed,
1920
bf16_with_fp8_current_scaling_mixed,
2021
bf16_with_fp8_subchannel_scaling_mixed,
2122
bf16_with_mxfp8_mixed,
@@ -34,6 +35,8 @@ def get_precision_config(compute_dtype: str):
3435
return bf16_with_mxfp8_mixed()
3536
elif compute_dtype == "fp8_sc":
3637
return bf16_with_fp8_subchannel_scaling_mixed()
38+
elif compute_dtype == "fp8_ds":
39+
return bf16_with_fp8_delayed_scaling_mixed()
3740
elif compute_dtype == "bf16":
3841
return bf16_mixed()
3942
elif compute_dtype == "nvfp4":

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

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,3 +56,126 @@ def _llama_benchmark_common(cfg: ConfigContainer) -> None:
5656
cfg.model.moe_token_dispatcher_type = "alltoall"
5757
if cuda_graph_impl == "none" and te_rng_disabled:
5858
cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = False
59+
60+
61+
def _llama2_70b_precision_config(compute_dtype: str):
62+
"""Get the precision configs for the given compute dtype and FP8 recipe."""
63+
precision_config = _perf_precision(compute_dtype)
64+
precision_config.fp4_param = False
65+
precision_config.fp4_param_gather = False
66+
precision_config.fp8_param = False
67+
precision_config.fp8_param_gather = False
68+
precision_config.reuse_grad_buf_for_mxfp8_param_ag = False
69+
precision_config.num_layers_at_start_in_bf16 = 0
70+
precision_config.num_layers_at_end_in_bf16 = 0
71+
precision_config.first_last_layers_bf16 = False
72+
if compute_dtype == "fp8_ds":
73+
precision_config.fp8_param = True
74+
precision_config.fp8_param_gather = True
75+
return precision_config
76+
77+
78+
def _llama2_benchmark_common(cfg: ConfigContainer) -> None:
79+
"""Apply legacy Llama 2 benchmark defaults shared by the flat recipes."""
80+
cfg.optimizer.bf16 = True
81+
cfg.optimizer.clip_grad = 0.3
82+
cfg.optimizer.fp8_recipe = "delayed"
83+
cfg.optimizer.lr = 0.0005
84+
cfg.optimizer.min_lr = 0
85+
cfg.optimizer.overlap_params = True
86+
cfg.optimizer.params_dtype = torch.bfloat16
87+
cfg.optimizer.use_distributed_optimizer = True
88+
cfg.optimizer.weight_decay = 0.0001
89+
cfg.ddp.bucket_size = 45000000
90+
cfg.ddp.overlap_grad_reduce = True
91+
cfg.ddp.overlap_param_gather = True
92+
cfg.ddp.use_distributed_optimizer = True
93+
cfg.ddp.grad_reduce_in_fp32 = False
94+
cfg.model.hidden_size = 8192
95+
cfg.model.ffn_hidden_size = 28672
96+
cfg.model.num_layers = 80
97+
cfg.model.num_attention_heads = 64
98+
cfg.model.num_query_groups = 8
99+
cfg.model.kv_channels = 128
100+
cfg.model.normalization = "RMSNorm"
101+
cfg.model.gated_linear_unit = True
102+
cfg.model.add_bias_linear = False
103+
cfg.model.bf16 = True
104+
cfg.model.autocast_dtype = torch.bfloat16
105+
cfg.model.params_dtype = torch.bfloat16
106+
cfg.model.fp8_amax_compute_algo = "max"
107+
cfg.model.fp8_amax_history_len = 4
108+
cfg.model.fp8_dot_product_attention = True
109+
cfg.model.num_layers_at_start_in_bf16 = 0
110+
cfg.model.num_layers_at_end_in_bf16 = 0
111+
cfg.model.apply_rope_fusion = True
112+
cfg.model.fused_single_qkv_rope = 1
113+
cfg.model.bias_activation_fusion = True
114+
cfg.model.bias_dropout_fusion = True
115+
cfg.model.gradient_accumulation_fusion = True
116+
cfg.model.masked_softmax_fusion = True
117+
cfg.model.cross_entropy_loss_fusion = True
118+
cfg.model.cross_entropy_fusion_impl = "te"
119+
cfg.model.persist_layer_norm = True
120+
cfg.model.recompute_modules = ['core_attn']
121+
cfg.model.use_transformer_engine_op_fuser = 1
122+
cfg.model.use_te_rng_tracker = True
123+
cfg.rng.te_rng_tracker = True
124+
cfg.model.cp_comm_type = "a2a"
125+
cfg.model.cpu_offloading = False
126+
cfg.model.cuda_graph_modules = []
127+
cfg.model.cuda_graph_warmup_steps = 5
128+
cfg.model.deallocate_pipeline_outputs = True
129+
cfg.model.disable_parameter_transpose_cache = True
130+
cfg.model.attention_dropout = 0.0
131+
cfg.model.hidden_dropout = 0.0
132+
cfg.model.attention_softmax_in_fp32 = False
133+
cfg.model.embedding_init_method_std = 0.02
134+
cfg.model.expert_tensor_parallel_size = 1
135+
cfg.model.microbatch_group_size_per_vp_stage = 1
136+
cfg.model.pipeline_dtype = None
137+
cfg.mixed_precision.fp8_amax_compute_algo = "max"
138+
cfg.mixed_precision.fp8_amax_history_len = 4
139+
cfg.mixed_precision.fp8_dot_product_attention = True
140+
cfg.mixed_precision.grad_reduce_in_fp32 = False
141+
cfg.mixed_precision.pipeline_dtype = None
142+
cfg.train.manual_gc = True
143+
cfg.train.manual_gc_eval = False
144+
cfg.train.manual_gc_interval = 10000
145+
cfg.train.skip_sync_grad_norm_across_mp = True
146+
cfg.train.check_optimizer_step_success = False
147+
cfg.train.decrease_batch_size_if_needed = False
148+
cfg.train.empty_unused_memory_level = 0
149+
cfg.validation.eval_micro_batch_size = 1
150+
cfg.validation.full_validation = False
151+
cfg.validation.multiple_validation_sets = False
152+
cfg.train.skip_train = False
153+
cfg.train.test_mode = False
154+
cfg.scheduler.lr_decay_style = "cosine"
155+
cfg.scheduler.lr_warmup_fraction = 0.0
156+
cfg.scheduler.lr_warmup_iters = 0
157+
cfg.scheduler.lr_warmup_steps = 0.0
158+
cfg.scheduler.start_weight_decay = 0.0001
159+
cfg.scheduler.end_weight_decay = 0.0001
160+
cfg.scheduler.weight_decay_incr_style = "constant"
161+
cfg.scheduler.lr_wsd_decay_style = "exponential"
162+
cfg.scheduler.override_opt_param_scheduler = False
163+
cfg.scheduler.use_checkpoint_opt_param_scheduler = False
164+
cfg.dataset.seq_length = 8192
165+
cfg.dataset.enable_offline_packing = True
166+
cfg.dataset.create_attention_mask = False
167+
cfg.dataset.dataloader_type = "batch"
168+
cfg.dataset.data_sharding = True
169+
cfg.dataset.drop_last = True
170+
cfg.dataset.persistent_workers = True
171+
cfg.dataset.pin_memory = True
172+
cfg.dataset.memmap_workers = 1
173+
cfg.dataset.do_validation = True
174+
cfg.dataset.do_test = False
175+
cfg.checkpoint.finetune = True
176+
cfg.checkpoint.ckpt_format = "torch_dist"
177+
cfg.checkpoint.fully_parallel_save = True
178+
cfg.checkpoint.load_optim = False
179+
cfg.checkpoint.load_rng = False
180+
cfg.tokenizer.tokenizer_type = "HuggingFaceTokenizer"
181+
cfg.tokenizer.tensor_model_parallel_size = 1

‎src/megatron/bridge/perf_recipes/llama/gb200/llama2.py‎

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""GB200 performance recipes for Llama 2."""
15+
16+
from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS
17+
from megatron.bridge.perf_recipes.llama.common import (
18+
CommOverlapConfig,
19+
ConfigContainer,
20+
_llama2_benchmark_common,
21+
_llama2_70b_precision_config,
22+
llama2_70b_peft_config,
23+
)
24+
25+
def llama2_70b_peft_4gpu_gb200_fp8ds_config() -> ConfigContainer:
26+
"""Llama 2 70B LoRA: 4× GB200, FP8 DS, GBS=8, seq_length=2048."""
27+
cfg = llama2_70b_peft_config(peft_scheme="lora")
28+
cfg.mixed_precision = _llama2_70b_precision_config("fp8_ds")
29+
_llama2_benchmark_common(cfg)
30+
31+
cfg.model.tensor_model_parallel_size = 1
32+
cfg.model.pipeline_model_parallel_size = 1
33+
cfg.model.context_parallel_size = 1
34+
cfg.train.global_batch_size = 8
35+
cfg.train.micro_batch_size = 1
36+
cfg.model.cuda_graph_impl = None
37+
cfg.model.cuda_graph_scope = None
38+
cfg.model.cpu_offloading_num_layers = 20
39+
cfg.model.cpu_offloading = True
40+
cfg.validation.eval_global_batch_size = 4
41+
cfg.validation.eval_interval = 48
42+
cfg.validation.eval_iters = 44
43+
cfg.validation.start_at_eval_iter = 192
44+
cfg.scheduler.lr_decay_iters = 800
45+
cfg.scheduler.lr_decay_steps = 6400
46+
cfg.scheduler.wd_incr_steps = 6400
47+
cfg.dataset.max_train_samples = 6432
48+
cfg.dataset.num_workers = 4
49+
cfg.dataset.seed = 30339
50+
cfg.rng.seed = 30339
51+
# Keep process settings next to the recipe so users can see the exact benchmark environment.
52+
cfg.env_vars = {
53+
**COMMON_PERF_ENV_VARS,
54+
# CUDA stream scheduling for this model and parallel layout.
55+
"CUDA_DEVICE_MAX_CONNECTIONS": 1,
56+
# CUDA graph and allocator behavior for this recipe.
57+
"NCCL_GRAPH_REGISTER": 0,
58+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
59+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
60+
# NCCL user-buffer and launch settings.
61+
"NCCL_NVLS_ENABLE": 1,
62+
# Transformer Engine overlap settings for this model.
63+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
64+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
65+
}
66+
return cfg
67+
68+
69+
def llama2_70b_peft_8gpu_gb200_fp8ds_config() -> ConfigContainer:
70+
"""Llama 2 70B LoRA: 8× GB200, FP8 DS, GBS=8, seq_length=2048."""
71+
cfg = llama2_70b_peft_config(peft_scheme="lora")
72+
cfg.mixed_precision = _llama2_70b_precision_config("fp8_ds")
73+
_llama2_benchmark_common(cfg)
74+
75+
cfg.model.tensor_model_parallel_size = 1
76+
cfg.model.pipeline_model_parallel_size = 1
77+
cfg.model.context_parallel_size = 1
78+
cfg.train.global_batch_size = 8
79+
cfg.train.micro_batch_size = 1
80+
cfg.model.cuda_graph_impl = "local"
81+
cfg.model.cuda_graph_scope = "full_iteration"
82+
cfg.model.cpu_offloading_num_layers = 11
83+
cfg.model.cpu_offloading = True
84+
cfg.validation.eval_global_batch_size = 8
85+
cfg.validation.eval_interval = 48
86+
cfg.validation.eval_iters = 22
87+
cfg.validation.start_at_eval_iter = 192
88+
cfg.scheduler.lr_decay_iters = 800
89+
cfg.scheduler.lr_decay_steps = 6400
90+
cfg.scheduler.wd_incr_steps = 6400
91+
cfg.dataset.max_train_samples = 6432
92+
cfg.dataset.num_workers = 4
93+
cfg.dataset.seed = 27208
94+
cfg.rng.seed = 27208
95+
# Keep process settings next to the recipe so users can see the exact benchmark environment.
96+
cfg.env_vars = {
97+
**COMMON_PERF_ENV_VARS,
98+
# CUDA stream scheduling for this model and parallel layout.
99+
"CUDA_DEVICE_MAX_CONNECTIONS": 1,
100+
# CUDA graph and allocator behavior for this recipe.
101+
"NCCL_GRAPH_REGISTER": 0,
102+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
103+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
104+
# NCCL user-buffer and launch settings.
105+
"NCCL_NVLS_ENABLE": 1,
106+
# Transformer Engine overlap settings for this model.
107+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
108+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
109+
}
110+
return cfg
111+
112+
113+
def llama2_70b_peft_8gpu_gb200_nvfp4_config() -> ConfigContainer:
114+
"""Llama 2 70B LoRA: 8× GB200, NVFP4, GBS=8, seq_length=2048."""
115+
cfg = llama2_70b_peft_config(peft_scheme="lora")
116+
cfg.mixed_precision = _llama2_70b_precision_config("nvfp4")
117+
_llama2_benchmark_common(cfg)
118+
119+
cfg.model.tensor_model_parallel_size = 1
120+
cfg.model.pipeline_model_parallel_size = 1
121+
cfg.model.context_parallel_size = 2
122+
cfg.train.global_batch_size = 8
123+
cfg.train.micro_batch_size = 1
124+
cfg.model.cuda_graph_impl = "local"
125+
cfg.model.cuda_graph_scope = "full_iteration"
126+
cfg.model.cpu_offloading_num_layers = 20
127+
cfg.optimizer.lr = 0.0006
128+
cfg.validation.eval_global_batch_size = 4
129+
cfg.validation.eval_interval = 48
130+
cfg.validation.eval_iters = 44
131+
cfg.validation.start_at_eval_iter = 192
132+
cfg.scheduler.lr_decay_iters = 700
133+
cfg.scheduler.lr_decay_steps = 5600
134+
cfg.scheduler.wd_incr_steps = 5600
135+
cfg.dataset.max_train_samples = 5628
136+
cfg.dataset.num_workers = 4
137+
cfg.dataset.seed = 23829
138+
cfg.rng.seed = 23829
139+
# Keep process settings next to the recipe so users can see the exact benchmark environment.
140+
cfg.env_vars = {
141+
**COMMON_PERF_ENV_VARS,
142+
# CUDA stream scheduling for this model and parallel layout.
143+
"CUDA_DEVICE_MAX_CONNECTIONS": 1,
144+
# CUDA graph and allocator behavior for this recipe.
145+
"NCCL_GRAPH_REGISTER": 0,
146+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
147+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
148+
# NCCL user-buffer and launch settings.
149+
"NCCL_NVLS_ENABLE": 1,
150+
# Transformer Engine overlap settings for this model.
151+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
152+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
153+
}
154+
return cfg
155+
156+
157+
def llama2_70b_peft_72gpu_gb200_fp8ds_config() -> ConfigContainer:
158+
"""Llama 2 70B LoRA: 72× GB200, FP8 DS, GBS=8, seq_length=2048."""
159+
cfg = llama2_70b_peft_config(peft_scheme="lora")
160+
cfg.mixed_precision = _llama2_70b_precision_config("fp8_ds")
161+
_llama2_benchmark_common(cfg)
162+
163+
cfg.model.tensor_model_parallel_size = 1
164+
cfg.model.pipeline_model_parallel_size = 1
165+
cfg.model.context_parallel_size = 8
166+
cfg.train.global_batch_size = 9
167+
cfg.train.micro_batch_size = 1
168+
cfg.model.cuda_graph_impl = "local"
169+
cfg.model.cuda_graph_scope = "full_iteration"
170+
cfg.model.cpu_offloading_num_layers = 20
171+
cfg.model.cpu_offloading = True
172+
cfg.validation.eval_global_batch_size = 8
173+
cfg.validation.eval_global_batch_size = 36
174+
cfg.validation.eval_interval = 43
175+
cfg.validation.eval_iters = 5
176+
cfg.validation.start_at_eval_iter = 172
177+
cfg.scheduler.lr_decay_iters = 800
178+
cfg.scheduler.lr_decay_steps = 7200
179+
cfg.scheduler.wd_incr_steps = 7200
180+
cfg.dataset.max_train_samples = 7236
181+
cfg.dataset.num_workers = 2
182+
cfg.dataset.seed = 16584
183+
cfg.rng.seed = 16584
184+
# Keep process settings next to the recipe so users can see the exact benchmark environment.
185+
cfg.env_vars = {
186+
**COMMON_PERF_ENV_VARS,
187+
# CUDA stream scheduling for this model and parallel layout.
188+
"CUDA_DEVICE_MAX_CONNECTIONS": 1,
189+
# CUDA graph and allocator behavior for this recipe.
190+
"NCCL_GRAPH_REGISTER": 0,
191+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
192+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
193+
# NCCL user-buffer and launch settings.
194+
"NCCL_NVLS_ENABLE": 1,
195+
# Transformer Engine overlap settings for this model.
196+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
197+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
198+
}
199+
return cfg
200+
201+
202+
def llama2_70b_peft_512gpu_gb200_fp8ds_config() -> ConfigContainer:
203+
"""Llama 2 70B LoRA: 512× GB200, FP8 DS, GBS=8, seq_length=2048."""
204+
cfg = llama2_70b_peft_config(peft_scheme="lora")
205+
cfg.mixed_precision = _llama2_70b_precision_config("fp8_ds")
206+
_llama2_benchmark_common(cfg)
207+
208+
cfg.model.tensor_model_parallel_size = 1
209+
cfg.model.pipeline_model_parallel_size = 1
210+
cfg.model.context_parallel_size = 8
211+
cfg.train.global_batch_size = 64
212+
cfg.train.micro_batch_size = 1
213+
cfg.model.cuda_graph_impl = "local"
214+
cfg.model.cuda_graph_scope = "full_iteration"
215+
cfg.model.cpu_offloading_num_layers = 20
216+
cfg.model.cpu_offloading = True
217+
cfg.optimizer.lr = 0.0006
218+
cfg.validation.eval_global_batch_size = 64
219+
cfg.validation.eval_interval = 6
220+
cfg.validation.eval_iters = 3
221+
cfg.validation.start_at_eval_iter = 66
222+
cfg.scheduler.lr_decay_iters = 600
223+
cfg.scheduler.lr_decay_steps = 38400
224+
cfg.scheduler.wd_incr_steps = 38400
225+
cfg.dataset.max_train_samples = 38592
226+
cfg.dataset.num_workers = 2
227+
cfg.dataset.seed = 22205
228+
cfg.rng.seed = 22205
229+
# Keep process settings next to the recipe so users can see the exact benchmark environment.
230+
cfg.env_vars = {
231+
**COMMON_PERF_ENV_VARS,
232+
# CUDA stream scheduling for this model and parallel layout.
233+
"CUDA_DEVICE_MAX_CONNECTIONS": 1,
234+
# CUDA graph and allocator behavior for this recipe.
235+
"NCCL_GRAPH_REGISTER": 0,
236+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
237+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
238+
# NCCL user-buffer and launch settings.
239+
"NCCL_NVLS_ENABLE": 1,
240+
# Transformer Engine overlap settings for this model.
241+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
242+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
243+
}
244+
return cfg

0 commit comments

Comments
 (0)