Star84
365
代码
代码
Issues57
Pull Requests142
项目讨论
Wiki
分析
Star84
365
已合并
提交PP自动并行算法的代码实现 #516
AtomGit-Bot创建于 2024年7月8日
提交PP自动并行算法的代码实现 #516
已合并
AtomGit-Bot创建于 2024年7月8日
从refs/pull/516/head合入到master
共 15 个文件变更+1695-26
Mdocs/features/automated-pipeline.md+23-13
@@ -1,34 +1,44 @@
1-# PP自动并行1+# PP自动并行算法
2 2 
3## 问题分析3## 问题分析
4 4 
5-靠近模型前面的流水线stage的内存占用多于模型后面的stage内存占用，并且内存占用差距有2~3倍，总体上模型规模受限于PP-Stage 0的显存。当前通过离线建模搜索手动配置PP层的分布可以缓解显存不平衡的问题，但存在多模型的泛化性问题，端到端训练效果无法得到保证。5+流水线并行是将模型网络层切分成多个stage，再把stage映射到不同的设备上，使得不同设备并行计算神经网络的不同部分。流水线并行大大缓解了单卡内存瓶颈问题，并通过多卡之间的流水训练提高了硬件的利用率。流水线并行成为了当前大模型训练最常用的并行方式之一。然而当前流水线并行在内存消耗和性能方面并非最优，主要存在两大问题：
6+ 
7+1）内存不均衡：当前流水线常用调度模式（1F1B）下，靠近模型前面层的流水线stage的内存占用远多于后面的stage内存占用，并且内存占用差距有2~3倍，总体上可训的模型规模受限于PP-Stage 0的显存消耗。
8+ 
9+2）流水线气泡：流水线1F1B调度策略在每个设备上交替进行小批次数据的前向后向计算，由于各流水设备之间计算负载不均衡或者网络通信的波动，导致设备与设备之间存在等待（流水线气泡），影响训练性能。
6 10 
7## 解决方案11## 解决方案
12+本系统基于在线profiling+PP建模搜索，通过使能内存优化模块、性能优化模块分别最大化流水线并行训练的内存和性能。内存优化模块旨在通过自动寻找流水线并行中stage的最优层分布和细粒度重计算模块，均匀分配每个卡上的显存，优化存在显存瓶颈的PP-stages，降低峰值内存；性能优化模块采用mbs序列和前向反向调度序列自动寻优和多流异步通信机制，压缩流水线气泡，提升训练性能。
8 13 
9-本算法通过自动寻找流水线并行中stage的最优网络层分布和细粒度重计算模块，均匀分配每个卡上的显存，优化存在显存瓶颈的PP-stages，降低峰值内存。14+### 内存优化模块
15+基于在线profiling+PP建模搜索，自动构建出最优的内存排布方案均衡化各个stage之间的内存开销，降低峰值内存的同时最小化端到端训练时间，具备较好的易用性和泛化性。具体而言，在层分布和细粒度重计算的联合搜索空间自动寻优内存排布方案：
16+① PP层分布切分：采用不均匀层切分策略，自动搜索最优层切分方式，均衡化每个卡消耗的显存，从而优化存在显存瓶颈的PP-stages，降低峰值内存。
17+② 细粒度重计算：利用流水线气泡时间来做重计算，保证性能不劣化，通过自动寻优细粒度的重计算策略，进一步降低峰值内存。
10 18 
11-### 解决思路19+### 性能优化模块
20+在满足训练峰值内存开销不超过设备最大内存容量的条件下，通过自动寻找流水线并行中最优的mbs序列及前向反向调度序列，最小化端到端训练时间。
21+① 动态mbs：在给定的gbs下，自动搜索最优mbs序列。通过小mbs加速流水线的启动与冷却，压缩气泡时间，稳态阶段自动寻找最高效的mbs进行计算，缩短稳态阶段计算时间，提升端到端训练性能。
22+② 前反向调度：通过调整流水线并行过程中前反向计算的顺序，结合多流异步通信机制，压缩流水线稳态气泡，提升训练性能。
12 23 
13-基于在线建模搜索自动构建出最优的内存排布方案使得各个stage之间的内存相对均衡，降低峰值内存的同时最小化端到端训练时间，具备更好的易用性和泛化性。通过数学理论建模+在线profiling的方式。从PP维度，在层分布和细粒度重计算的联合搜索空间自动寻优：24+PP自动并行系统如下图所示：
14-① PP层分布：搜索最优层切分方式，均匀分配每个卡上的显存，优化存在显存瓶颈的PP-stages，降低峰值内存。
15-② 细粒度重计算模块：通过自动寻优重计算策略，进一步降低峰值内存，同时保证性能不劣化。
16- 
17-PP自动并行1.0流程如下图所示：
18 25 
19<p align="center"> <img src="../../sources/images/auto_pipeline_parallel.png"></p>26<p align="center"> <img src="../../sources/images/auto_pipeline_parallel.png"></p>
20 27 
21 28 
22## 使用场景29## 使用场景
A
Aaijgnem12024年7月9日

这里是不是使用场景更好，因为之前模板都是使用场景

likedislike
23 30 
24-该特性主要用于训练过程中显存不足的场景，使用PP自动并行可有效降低显存的占用。31+该系统主要用于开启流水线并行的训练场景，使用PP自动并行系统可有效优化内存不足或流水线气泡占比过大的问题。
25-**注意：**使用条件：`--pipeline-model-parallel-size >= 2`32+**注意：**使用条件：
33+1. `--pipeline-model-parallel-size >= 2`；
34+2. 内存、性能优化模块不能同时使用。
26 35 
27 36 
28## 使用方法37## 使用方法
29 38 
30-启用PP自动并行，请首先在训练脚本中添加 `--automated-pipeline` 标志开启PP自动并行策略。39+（1）当内存不足时，可启用PP自动并行内存优化模块，请首先在训练脚本中添加 `--automated-pipeline` 标志启用功能。
40+（2）当流水线气泡过大导致训练性能不优时，可启用PP自动并行性能优化模块，请首先在训练脚本中添加 `--automated-pipeline-perf` 标志启用功能。
31 41 
32## 使用效果42## 使用效果
33 43 
34-LLaMA2-7B，LLaMA-13B，LLaMA2-70B等使用流水线并行PP配置训练的模型，叠加本算法后平均峰值内存减少5%，平均性能劣化小于1%。44+PP自动并行内存优化模块收益：LLaMA2-7B，LLaMA-13B，LLaMA2-70B等使用流水线并行PP配置训练的模型，叠加本算法后平均峰值内存减少11.5%，平均性能劣化小于1%。性能优化模块收益：LLaMA2-7B，LLaMA-13B，LLaMA3-8B等使用流水线并行PP配置训练的模型，叠加本算法后平均性能提升7.6%。
Mmindspeed/arguments.py+27-1
@@ -228,8 +228,12 @@ def _add_automated_pipeline_args(parser):
228 group = parser.add_argument_group(title='automated_pipeline_allocation')228 group = parser.add_argument_group(title='automated_pipeline_allocation')
229 group.add_argument('--automated-pipeline',229 group.add_argument('--automated-pipeline',
230 action='store_true',230 action='store_true',
231- help='To enable automated pipeline process'231+ help='To enable automated pipeline memory saving process'
232 )232 )
233+ group.add_argument('--automated-pipeline-perf',
234+ action='store_true',
235+ help='To enable automated pipeline performance acceleration process'
236+ )
233 group.add_argument('--save-memory-ratio',237 group.add_argument('--save-memory-ratio',
234 type=float, default=0.20,238 type=float, default=0.20,
235 help='To set memory saving rate in automated pipeline'239 help='To set memory saving rate in automated pipeline'
@@ -245,6 +249,22 @@ def _add_automated_pipeline_args(parser):
245 help='To store the recompute type of automated pipeline, 0 for mlp block '249 help='To store the recompute type of automated pipeline, 0 for mlp block '
246 '1 for attention block and 2 for transformer layer'250 '1 for attention block and 2 for transformer layer'
247 )251 )
252+ group.add_argument('--optimized-mbs-list',
253+ type=str,
254+ help='To store the optimized mbs policy of automated pipeline performance'
255+ )
256+ group.add_argument('--mbs-idx',
257+ type=int,
258+ help='To store the index of mbs list'
259+ )
260+ group.add_argument('--pp-schedule-list',
261+ type=str,
262+ help='To store the pipeline schedule policy of automated pipeline performance'
263+ )
264+ group.add_argument('--optimized-mbs-mode',
265+ action='store_false',
266+ help='To store the status of optimized mbs in automated pipeline performance'
267+ )
248 group.add_argument('--memory-fragmentation',268 group.add_argument('--memory-fragmentation',
249 action='store_true', default=False,269 action='store_true', default=False,
250 help='Enable the memory fragmentation feature.')270 help='Enable the memory fragmentation feature.')
@@ -444,6 +464,12 @@ def validate_args_wrapper(validate_args):
444 if args.optimize_recomp_communication_level > 0:464 if args.optimize_recomp_communication_level > 0:
445 print("[WARNING] disable optimize recomp communication level when enabling automated pipeline")465 print("[WARNING] disable optimize recomp communication level when enabling automated pipeline")
446 args.optimize_recomp_communication_level = 0466 args.optimize_recomp_communication_level = 0
467+ if args.automated_pipeline_perf:
468+ if args.automated_pipeline:
469+ print("[WARNING] disable automated pipeline when enabling automated pipeline performance version")
470+ args.automated_pipeline = False
471+ if args.num_layers_per_virtual_pipeline_stage is not None:
M
Mmingzhenwang2024年7月8日

automated_pipeline_perf和automated_pipeline同时开启打warning并自动选择perf没问题 但是这里VPP打开了打warning置为none不合理，建议抛出AssertionError

likedislike
472+ raise AssertionError('automated pipeline performance is temporarily incompatible with virtual pipeline')
447 if args.use_ascend_mc2 and args.use_ascend_coc:473 if args.use_ascend_mc2 and args.use_ascend_coc:
448 raise AssertionError('--mc2 and coc can not be used together')474 raise AssertionError('--mc2 and coc can not be used together')
449 if args.use_nd_matmul:475 if args.use_nd_matmul:
Mmindspeed/core/memory/auto_pipeline/autopipeline_solver.py+3-1
@@ -432,7 +432,9 @@ def destroy_model_parallel_profiling_wrapper(destroy_model_parallel):
432 @wraps(destroy_model_parallel)432 @wraps(destroy_model_parallel)
433 def wrapper(*args, **kwargs):433 def wrapper(*args, **kwargs):
434 argument = get_args()434 argument = get_args()
435- if argument.automated_pipeline and not argument.num_layer_list:435+ enable_profiling_destroy = (argument.automated_pipeline and not argument.num_layer_list) \
436+ or (argument.automated_pipeline_perf and not argument.optimized_mbs_list)
437+ if enable_profiling_destroy:
436 destroy_global_parallel_group()438 destroy_global_parallel_group()
437 else:439 else:
438 destroy_model_parallel(*args, **kwargs)440 destroy_model_parallel(*args, **kwargs)
Amindspeed/core/performance/auto_pipeline_perf/autopipeline_perf.py+400-0
@@ -0,0 +1,400 @@
1+import time
2+from functools import partial
3+_TRAIN_START_TIME = time.time()
4+import json
5+import os.path
6+import gc
7+import copy
8+import torch
9+import torch.nn
10+import torch_npu
11+from megatron.training import print_rank_0
12+from megatron.training.arguments import parse_args
13+from megatron.core.parallel_state import get_embedding_group
14+from megatron.training import get_args
15+from megatron.training import get_timers
16+from megatron.training import training
17+from megatron.training.training import print_datetime
18+from megatron.core.pipeline_parallel import p2p_communication
19+from megatron.core import mpu, tensor_parallel
20+from megatron.training.initialize import initialize_megatron
21+from megatron.training.initialize import set_jit_fusion_options
22+ 
23+ 
24+profile_context = {"fwd_time":[], "bwd_time":[]}
25+ 
26+ 
27+class AutoPipeline_Perf:
28+ autopipeline_perf = None
29+ 
30+ def __init__(self, args):
31+ self.args = copy.deepcopy(args)
32+ self.context = {
33+ 'module': []
34+ }
35+ self.modules_hooks = []
36+ self.profiling_step = 0
37+ self.stop_profiling_step = 3
38+ self.unit_mb = 1024 * 1024
39+ 
40+ @staticmethod
41+ def get_memory_status():
42+ used_memory = torch.npu.memory_allocated()
43+ reserved_memory = torch.npu.memory_reserved()
44+ return used_memory, reserved_memory
45+ 
46+ def _cal_tensor_size(self, tensor):
47+ try:
48+ return tensor.numel() * tensor.element_size() / self.unit_mb
49+ except ZeroDivisionError:
50+ return 0
51+ 
52+ def pre_hook_func(self, state, sync: bool, *args, **kargs):
53+ used_memory, _ = self.get_memory_status()
54+ torch.npu.reset_max_memory_allocated()
55+ state['memory'] = used_memory
56+ size = 0
57+ for arg in args:
58+ if isinstance(arg, torch.Tensor):
59+ size += self._cal_tensor_size(arg)
60+ elif isinstance(arg, tuple) or isinstance(arg, list):
61+ for t in arg:
62+ if isinstance(t, torch.Tensor):
63+ size += self._cal_tensor_size(t)
64+ state['input'] = size
65+ 
66+ def post_hook_func(self, state, sync: bool, *args, **kargs):
67+ used_memory, _ = self.get_memory_status()
68+ max_mem = torch.npu.max_memory_allocated()
69+ state['peak_memory'] = max_mem - state['memory']
70+ state['memory'] = (used_memory - state['memory']) // self.unit_mb
71+ 
72+ def forward_pre_hook(self, name, parent_ctx, ctx):
73+ if self.profiling_step < self.stop_profiling_step:
74+ ctx['name'] = name
75+ if 'layers' in parent_ctx:
76+ parent_ctx['layers'].append(ctx)
77+ 
78+ def hook(module, *args, **kargs):
79+ if self.profiling_step < self.stop_profiling_step:
80+ if 'module' in self.context:
81+ self.context['module'].append(ctx)
82+ self.pre_hook_func(ctx, True, *args, **kargs)
83+ 
84+ return hook
85+ 
86+ def forward_post_hook(self, ctx):
87+ def hook(module, *args, **kargs):
88+ if self.profiling_step < self.stop_profiling_step:
89+ self.post_hook_func(ctx, True, *args)
90+ if 'module' in self.context:
91+ self.context['module'].pop()
92+ 
93+ return hook
94+ 
95+ def register_recursive_hook(self, prefix_name, model, ctx):
96+ for name, module in model.named_children():
97+ if 'layers' not in ctx:
98+ ctx['layers'] = []
99+ current_ctx = {}
100+ 
101+ next_name = prefix_name + "." + name if prefix_name != "" else name
102+ if next_name == "module.module":
103+ pre_hook = module.register_forward_pre_hook(self.forward_pre_hook(name, ctx, current_ctx))
104+ post_hook = module.register_forward_hook(self.forward_post_hook(current_ctx))
105+ self.modules_hooks.append(pre_hook)
106+ self.modules_hooks.append(post_hook)
107+ self.register_recursive_hook(next_name, module, current_ctx)
108+ 
109+ def step_hook(self, model):
110+ self.profiling_step += 1
111+ 
112+ def hook_step_func(self, step_func, models):
113+ def custom_step_func(*args, **kargs):
114+ result = step_func(*args, **kargs)
115+ if self.profiling_step < self.stop_profiling_step:
116+ used_memory, reserved_memory = self.get_memory_status()
117+ self.context['used_mem'] = used_memory // self.unit_mb
118+ if isinstance(models, list):
119+ for model in models:
120+ self.step_hook(model)
121+ else:
122+ self.step_hook(models)
123+ return result
124+ 
125+ return custom_step_func
126+ 
127+ def remove_outliers(self, data, m=2):
128+ data = sorted(data)
129+ median = data[len(data) // 2]
130+ deviation = [x for x in data if median - m * median < x < median + m * median]
131+ return deviation
132+ 
133+ def get_forward_context(self):
134+ global profile_context
135+ if "fwd_time" in profile_context:
136+ fwd_time_list = self.remove_outliers(profile_context["fwd_time"])
137+ try:
138+ self.context["fwd_time"] = sum(fwd_time_list) / len(fwd_time_list)
139+ except ZeroDivisionError:
140+ print("[Error] Divided by zero.")
141+ else:
142+ self.context["fwd_time"] = 0
143+ 
144+ def get_backward_context(self):
145+ global profile_context
146+ if "bwd_time" in profile_context:
147+ bwd_time_list = self.remove_outliers(profile_context["bwd_time"])
148+ try:
149+ self.context["bwd_time"] = sum(bwd_time_list) / len(bwd_time_list)
150+ except ZeroDivisionError:
151+ print("[Error] Divided by zero.")
152+ else:
153+ self.context["bwd_time"] = 0
154+ 
155+ def clear_global_context(self):
156+ global profile_context
157+ profile_context["fwd_time"] = []
158+ profile_context["bwd_time"] = []
159+ 
160+ def get_comm_time(self, config, sync: bool):
161+ if torch.distributed.get_rank() == 0:
162+ if sync:
163+ torch.cuda.synchronize()
164+ input_tensor = torch.ones(self.args.seq_length, self.args.micro_batch_size, self.args.hidden_size)
165+ start_time = time.time()
166+ p2p_communication.send_backward(input_tensor, config)
167+ comm_time = (time.time() - start_time) * 1000
168+ self.context['comm_time'] = comm_time
169+ else:
170+ self.context['comm_time'] = 0.028
171+ 
172+ def get_peak_memory(self, sync: bool):
173+ if sync:
174+ torch.cuda.synchronize()
175+ max_mem = torch.npu.max_memory_allocated() / (1 << 20)
176+ self.context['peak_memory'] = max_mem
177+ 
178+ def get_smi_peak_memory(self, sync: bool):
179+ if sync:
180+ torch.cuda.synchronize()
181+ mem_infos = torch.npu.mem_get_info()
182+ smi_peak_memory = (mem_infos[1] - mem_infos[0]) / (1 << 20)
183+ self.context['smi_peak_memory'] = smi_peak_memory
184+ 
185+ def get_smi_left_memory(self, sync: bool):
186+ if sync:
187+ torch.cuda.synchronize()
188+ mem_infos = torch.npu.mem_get_info()
189+ smi_left_memory = mem_infos[0] / (1 << 20)
190+ self.context['smi_left_memory'] = smi_left_memory
191+ 
192+ def get_data_parallel_size(self, data_parallel_size):
193+ if data_parallel_size:
194+ self.context['data_parallel_size'] = data_parallel_size
195+ else:
196+ self.context['data_parallel_size'] = 1
197+ 
198+ def broadcast_param_in_ranks(self, src_rank, param, init_memory):
199+ if torch.distributed.get_rank() == src_rank:
200+ try:
201+ param = torch.npu.max_memory_allocated() / self.unit_mb - init_memory
202+ except ZeroDivisionError:
203+ print("[Error] Divided by zero.")
204+ tmp_param = torch.cuda.IntTensor([param])
205+ torch.distributed.broadcast(tmp_param, src=src_rank)
206+ param = tmp_param.item()
207+ return param
208+ 
209+ def update_args_for_profiling(self, micro_batch_size=None):
210+ args = get_args()
211+ args.train_iters = self.stop_profiling_step
212+ if micro_batch_size:
213+ args.micro_batch_size = micro_batch_size
214+ args.global_batch_size = args.micro_batch_size * 16
215+ args.save = False
216+ args.log_interval = 10
217+ 
218+ def restore_args_for_training(self):
219+ args = get_args()
220+ if args.num_layers_per_virtual_pipeline_stage is None:
221+ args.num_layers = self.args.num_layers
222+ args.encoder_num_layers = self.args.num_layers
223+ args.train_iters = self.args.train_iters
224+ args.micro_batch_size = self.args.micro_batch_size
225+ args.global_batch_size = self.args.global_batch_size
226+ args.save = self.args.save
227+ args.log_interval = self.args.log_interval
228+ 
229+ 
230+def check_equal_model_configs(args, parsed_contents):
231+ model_index = 0
232+ for model_instance in parsed_contents:
233+ if args.hidden_size == model_instance.get("model_configs", {}).get("hidden_size") \
A
Aaijgnem12024年7月9日

建议使用 get获取model_instance.get("model_configs", {}).get("hidden_size")

likedislike
234+ and args.ffn_hidden_size == model_instance.get("model_configs", {}).get("ffn_hidden_size") \
235+ and args.seq_length == model_instance.get("model_configs", {}).get("seq_length") \
236+ and args.num_attention_heads == model_instance.get("model_configs", {}).get("num_attention_heads"):
237+ return model_index
238+ else:
239+ model_index += 1
240+ return -1
241+ 
242+ 
243+def check_equal_parallel_configs(args, parsed_content):
244+ for parallel_instance in parsed_content.get("optimpipeline_policy"):
245+ if args.num_layers == parallel_instance.get("num_layers") \
246+ and args.pipeline_model_parallel_size == parallel_instance.get("pipeline_model_parallel_size") \
247+ and args.tensor_model_parallel_size == parallel_instance.get("tensor_model_parallel_size") \
248+ and args.micro_batch_size == parallel_instance.get("micro_batch_size") \
249+ and args.global_batch_size == parallel_instance.get("global_batch_size"):
250+ return parallel_instance.get("enable_scheduler"), parallel_instance.get("optimized_mbs_list"), parallel_instance.get(
251+ "pp_schedule_list"), parallel_instance.get("optimal_layers")
252+ return None, None, None, None
253+ 
254+ 
255+def check_skip_profiling(args, config_file):
256+ if os.path.exists(config_file):
257+ with open(config_file) as config_json:
258+ config_contents = config_json.read()
259+ parsed_contents = json.loads(config_contents)
260+ index = check_equal_model_configs(args, parsed_contents)
261+ if index != -1:
262+ optimized_type, optimized_mbs_list, pp_schedule_list, optimal_layers = check_equal_parallel_configs(args, parsed_contents[index])
263+ if optimized_mbs_list or pp_schedule_list:
264+ return True, (optimized_type, optimized_mbs_list, pp_schedule_list, optimal_layers)
265+ return False, (None, None, None, None)
266+ 
267+ 
268+def check_out_of_memory(args, context, mbs_tries):
269+ total_memory = torch_npu.npu.get_device_properties(0).total_memory / (1 << 20)
270+ per_activation_memory_allocated = context["layers"][0]["memory"] // mbs_tries
271+ predict_next_max_memory_allocated = context["smi_peak_memory"] + per_activation_memory_allocated * args.pipeline_model_parallel_size + 1000
272+ if predict_next_max_memory_allocated > total_memory:
273+ return True
274+ else:
275+ return False
276+ 
277+ 
278+def broadcast_skip_in_ranks(src_rank, policy):
279+ is_skip = [False]
280+ if torch.distributed.get_rank() == src_rank:
281+ is_skip = [policy]
282+ tmp_is_skip = torch.cuda.BoolTensor(is_skip)
283+ torch.distributed.broadcast(tmp_is_skip, src=src_rank)
284+ return tmp_is_skip.item()
285+ 
286+ 
287+def calculate_num_of_activations(context):
288+ total_memory = torch_npu.npu.get_device_properties(0).total_memory / (1 << 20)
289+ activation_memory_allocated = context["layers"][0]["memory"]
290+ num_of_activations_left = (total_memory - context["smi_peak_memory"]) // activation_memory_allocated
291+ return int(num_of_activations_left)
292+ 
293+ 
294+def get_autopipeline_perf(args):
295+ AutoPipeline_Perf.autopipeline_perf = AutoPipeline_Perf(args)
296+ return AutoPipeline_Perf.autopipeline_perf
297+ 
298+ 
299+def autopipelineperf_profiling(mbs_tries, model_provider, model_type, forward_step_func, train_valid_test_dataset_provider,
300+ process_non_loss_data_func):
301+ initialize_megatron(extra_args_provider=None,
302+ args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})
303+ set_jit_fusion_options()
304+ global _TRAIN_START_TIME
305+ start_time_tensor = torch.cuda.DoubleTensor([_TRAIN_START_TIME])
306+ torch.distributed.all_reduce(start_time_tensor,
307+ op=torch.distributed.ReduceOp.MIN)
308+ _TRAIN_START_TIME = start_time_tensor.item()
309+ print_rank_0('time to initialize megatron (seconds): {:.3f}'.format(
310+ time.time() - _TRAIN_START_TIME))
311+ print_datetime('after megatron is initialized')
312+ args = get_args()
313+ pipelining = get_autopipeline_perf(args)
314+ pipelining.update_args_for_profiling(mbs_tries)
315+ models, optimizer, lr_scheduler = training.setup_model_and_optimizer(model_provider, model_type)
316+ optimizer.step = pipelining.hook_step_func(optimizer.step, models)
317+ config = training.get_model_config(models[0])
318+ 
319+ if args.virtual_pipeline_model_parallel_size is not None:
320+ train_data_iterator = []
321+ valid_data_iterator = []
322+ for i in range(len(models)):
323+ mpu.set_virtual_pipeline_model_parallel_rank(i)
324+ iterators = training.build_train_valid_test_data_iterators(
325+ train_valid_test_dataset_provider)
326+ train_data_iterator.append(iterators[0])
327+ valid_data_iterator.append(iterators[1])
328+ else:
329+ train_data_iterator, valid_data_iterator, _ = training.build_train_valid_test_data_iterators(
330+ train_valid_test_dataset_provider)
331+ if isinstance(models, list):
332+ for model in models:
333+ pipelining.register_recursive_hook("module", model, pipelining.context)
334+ else:
335+ pipelining.register_recursive_hook("module", models, pipelining.context)
336+ training.train(forward_step_func, models, optimizer, lr_scheduler, train_data_iterator, valid_data_iterator,
337+ process_non_loss_data_func, config)
338+ pipelining.get_smi_peak_memory(sync=True)
339+ pipelining.get_smi_left_memory(sync=True)
340+ pipelining.get_comm_time(config, sync=True)
341+ pipelining.get_peak_memory(sync=True)
342+ pipelining.get_data_parallel_size(args.data_parallel_size)
343+ pipelining.get_forward_context()
344+ pipelining.get_backward_context()
345+ pipelining.clear_global_context()
346+ 
347+ timers = get_timers()
348+ if timers('interval-time'):
349+ timers('interval-time').stop(barrier=True)
350+ 
351+ for hook_handle in pipelining.modules_hooks:
352+ hook_handle.remove()
353+ pipelining.modules_hooks.clear()
354+ pipelining.restore_args_for_training()
355+
356+ if hasattr(optimizer, 'chained_optimizers'):
357+ for op in optimizer.chained_optimizers:
358+ for key, value in op.optimizer.state.items():
359+ key.detach()
360+ key.grad = None
361+ key.storage().resize_(0)
362+ if "momentum_buffer" in value:
363+ value["momentum_buffer"].detach()
364+ value["momentum_buffer"].grad = None
365+ value["momentum_buffer"].storage().resize_(0)
366+ for ofg in op.param_groups:
367+ if "params" in ofg:
368+ for og in ofg["params"]:
369+ og.detach()
370+ og.grad = None
371+ og.storage().resize_(0)
372+ else:
373+ for key, value in optimizer.optimizer.state.items():
374+ key.detach()
375+ key.grad = None
376+ key.storage().resize_(0)
377+ if "momentum_buffer" in value:
378+ value["momentum_buffer"].detach()
379+ value["momentum_buffer"].grad = None
380+ value["momentum_buffer"].storage().resize_(0)
381+ for ofg in optimizer.param_groups:
382+ if "params" in ofg:
383+ for og in ofg["params"]:
384+ og.detach()
385+ og.grad = None
386+ og.storage().resize_(0)
387+ for md in models:
388+ for param in md.parameters():
389+ param.detach()
390+ param.grad = None
391+ param.storage().resize_(0)
392+ for param_tensor in md.state_dict():
393+ if md.state_dict()[param_tensor] is not None:
394+ md.state_dict()[param_tensor].detach()
395+ md.state_dict()[param_tensor].grad = None
396+ md.state_dict()[param_tensor].storage().resize_(0)
397+ 
398+ gc.collect()
399+ torch_npu.npu.empty_cache()
400+ return pipelining.context
Amindspeed/core/performance/auto_pipeline_perf/data_samplers.py+71-0
@@ -0,0 +1,71 @@
1+import random
2+from functools import wraps
3+import numpy as np
4+import torch
5+from torch.utils.data import Dataset
6+from megatron.training import get_args
7+from megatron.core import mpu
8+ 
9+ 
10+def build_pretraining_data_loader_decorator(build_pretraining_data_loader):
11+ @wraps(build_pretraining_data_loader)
12+ def wrapper(*args, **kwargs):
13+ if args[0] is None:
14+ return None
15+ argument = get_args()
16+ if argument.dataloader_type == 'single' and argument.automated_pipeline_perf and argument.optimized_mbs_list:
17+ batch_sampler = DynamicMicroBatchPretrainingSampler(
18+ total_samples=len(args[0]),
19+ consumed_samples=args[1],
20+ micro_batch_size=argument.micro_batch_size,
21+ data_parallel_rank=mpu.get_data_parallel_rank(),
22+ data_parallel_size=mpu.get_data_parallel_world_size())
23+ return torch.utils.data.DataLoader(args[0],
24+ batch_sampler=batch_sampler,
25+ num_workers=argument.num_workers,
26+ pin_memory=True)
27+ else:
28+ dataloader = build_pretraining_data_loader(*args, **kwargs)
29+ return dataloader
30+ return wrapper
31+ 
32+ 
33+class DynamicMicroBatchPretrainingSampler:
34+ 
35+ def __init__(self, total_samples, consumed_samples, micro_batch_size,
36+ data_parallel_rank, data_parallel_size, drop_last=True):
37+ 
38+ args = get_args()
39+ self.total_samples = total_samples
40+ self.consumed_samples = consumed_samples
41+ self.micro_batch_size = micro_batch_size
42+ self.data_parallel_rank = data_parallel_rank
43+ self.drop_last = drop_last
44+ self.dynamic_micro_batch_size = args.optimized_mbs_list
45+ self.micro_batch_times_data_parallel_size = [
46+ self.dynamic_micro_batch_size[i] * data_parallel_size \
47+ for i in range(len(self.dynamic_micro_batch_size))
48+ ]
49+ 
50+ def __len__(self):
51+ return self.total_samples
52+ 
53+ def get_start_end_idx(self, n_mbs):
54+ start_idx = self.data_parallel_rank * self.dynamic_micro_batch_size[n_mbs]
55+ end_idx = start_idx + self.dynamic_micro_batch_size[n_mbs]
56+ return start_idx, end_idx
57+ 
58+ def __iter__(self):
59+ batch = []
60+ n_mbs = 0
61+ for idx in range(self.consumed_samples, self.total_samples):
62+ batch.append(idx)
63+ if len(batch) == self.micro_batch_times_data_parallel_size[n_mbs]:
64+ start_idx, end_idx = self.get_start_end_idx(n_mbs)
65+ yield batch[start_idx:end_idx]
66+ batch = []
67+ n_mbs = (n_mbs + 1) % len(self.micro_batch_times_data_parallel_size)
68+ 
69+ if len(batch) > 0 and not self.drop_last:
70+ start_idx, end_idx = self.get_start_end_idx()
71+ yield batch[start_idx:end_idx]
Amindspeed/core/performance/auto_pipeline_perf/global_vars.py+16-0
@@ -0,0 +1,16 @@
1+from functools import wraps
2+from megatron.training import get_args
3+ 
4+ 
5+def get_num_microbatches_wrapper(get_num_microbatches):
6+ @wraps(get_num_microbatches)
7+ def wrapper(*args, **kwargs):
8+ argument = get_args()
9+ automated_pipeline_profile = argument.automated_pipeline_perf and not argument.optimized_mbs_list
10+ if argument.automated_pipeline_perf and argument.optimized_mbs_list and argument.optimized_mbs_mode:
11+ return len(argument.optimized_mbs_list)
12+ elif automated_pipeline_profile:
13+ return argument.global_batch_size // argument.data_parallel_size // argument.micro_batch_size
14+ else:
15+ return get_num_microbatches(*args, **kwargs)
16+ return wrapper
Amindspeed/core/performance/auto_pipeline_perf/optimpipeline_solver.py+304-0
@@ -0,0 +1,304 @@
1+import os
2+import json
3+import math
4+import time
5+from datetime import datetime
6+from itertools import product
7+import numpy as np
8+import torch
9+from megatron.training import get_args
10+from megatron.training.arguments import parse_args
11+from mindspeed.arguments import parse_args_wrapper
12+from .autopipeline_perf import check_equal_model_configs
13+ 
14+ 
15+class Parallel_Paras:
16+ def __init__(self,
17+ num_stages,
18+ fwd_durations,
19+ bwd_durations,
20+ num_microbatch,
21+ comm_matrix):
22+ self.num_stages = num_stages
23+ self.num_microbatch = num_microbatch
24+ self.fwd_durations = fwd_durations
25+ self.bwd_durations = bwd_durations
26+ self.comm_matrix = comm_matrix
27+ 
28+ 
29+def dynamic_mbs_1f1b(paras):
30+ num_stages = paras.num_stages
31+ num_microbatch = paras.num_microbatch
32+ computation_placement = list(range(num_stages)) + list(range(num_stages - 1, -1, -1))
33+ fwd_durations = paras.fwd_durations
34+ bwd_durations = paras.bwd_durations
35+ comm_matrix = paras.comm_matrix
36+ 
37+ fwd_bwd_order = ([f'F_{i}' for i in range(num_stages)] +
38+ [f'B_{i}' for i in range(num_stages - 1, -1, -1)])
39+ fwd_bwd_chunk_stage = dict(zip(fwd_bwd_order, computation_placement))
40+ 
41+ def get_stage_list(fwd_seq, bwd_seq, num_advanced):
42+ stage_order = []
43+ n = len(fwd_seq)
44+ for idx in range(n):
45+ if idx < num_advanced:
46+ stage_order.append(fwd_seq[idx])
47+ else:
48+ stage_order.append(fwd_seq[idx])
49+ stage_order.append(bwd_seq[idx - num_advanced])
50+ if idx == n - 1:
51+ for i in range(num_advanced):
52+ stage_order.append(bwd_seq[i - num_advanced])
53+ 
54+ return stage_order
55+ 
56+ def get_stage_schedule(all_jobs_array, comp_placement):
57+ stage_list = []
58+ for s in range(num_stages):
59+ stage_chunk_id = [index for index, element in enumerate(comp_placement) if element == s]
60+ warmup = num_stages - s
61+ stage_s_list = get_stage_list(all_jobs_array[stage_chunk_id[0]],
62+ all_jobs_array[stage_chunk_id[1]],
63+ warmup - 1)
64+ stage_list.append(stage_s_list)
65+ 
66+ return stage_list
67+ 
68+ all_jobs = np.array([[s + f'-{i}' for i in range(num_microbatch)] for s in fwd_bwd_order])
69+ stage_list = get_stage_schedule(all_jobs, computation_placement)
70+ 
71+ fwd_bwd_list = ([f"F_{j}-{i}" for i in range(num_microbatch) for j in range(num_stages)]
72+ + [f"B_{j}-{i}" for i in range(num_microbatch) for j in range(num_stages)])
73+ values = [0 for _ in range(num_stages * num_microbatch * 2)]
74+ start_time = dict(zip(fwd_bwd_list, values))
75+ fwd_bwd_durations = dict()
76+ for j in range(num_stages):
77+ for i in range(num_microbatch):
78+ fwd_bwd_durations[f"F_{j}-{i}"] = fwd_durations[j, i]
79+ fwd_bwd_durations[f"B_{j}-{i}"] = bwd_durations[j, i]
80+ 
81+ for n in range(num_stages - 1):
82+ for s in range(n + 1):
83+ start_time[f"F_{s}-{n - s + 1}"] = max(start_time[f"F_{s}-{n - s + 1}"],
84+ start_time[f"F_{s}-{n - s}"] + fwd_durations[s, n - s] + comm_matrix[s][s + 1])
85+ start_time[f"F_{s + 1}-{n - s}"] = max(start_time[f"F_{s + 1}-{n - s}"],
86+ start_time[f"F_{s}-{n - s}"] + fwd_durations[s, n - s] + comm_matrix[s][s + 1])
87+ 
88+ def get_prev_job_time(comp_start_time, pp_list, pp_id, mb_idx,
89+ comp_chunk_stage, comp_order, model_chunk_times,
90+ comm_time_matrix):
91+ current_job = pp_list[pp_id][mb_idx]
92+ prev_job_stage = pp_list[pp_id][mb_idx - 1]
93+ chunk_prev_job_stage, _ = prev_job_stage.split('-')
94+ stage_id_prev_job = comp_chunk_stage[chunk_prev_job_stage]
95+ chunk_position = comp_order.index(chunk_prev_job_stage)
96+ if chunk_position < len(comp_order) - 1:
97+ stage_id_next = comp_chunk_stage[comp_order[chunk_position + 1]]
98+ comm_time = comm_time_matrix[stage_id_prev_job][stage_id_next]
99+ else:
100+ comm_time = 0
101+ end_time_prev_job_stage = (comp_start_time[prev_job_stage] + model_chunk_times[prev_job_stage]
102+ + comm_time)
103+ 
104+ cur_model_chunk, cur_mb = current_job.split('-')
105+ chunk_position = comp_order.index(cur_model_chunk)
106+ if chunk_position > 0:
107+ prev_model_chunk = comp_order[chunk_position - 1]
108+ prev_job_batch = prev_model_chunk + '-' + cur_mb
109+ comm_time = comm_time_matrix[comp_chunk_stage[prev_model_chunk]][comp_chunk_stage[cur_model_chunk]]
110+ end_time_prev_job_batch = comp_start_time[prev_job_batch] + model_chunk_times[prev_job_batch] + comm_time
111+ completed_flag = comp_start_time[prev_job_stage] > 0 and comp_start_time[prev_job_batch] > 0
112+ else:
113+ end_time_prev_job_batch = 0
114+ completed_flag = comp_start_time[prev_job_stage] > 0
115+ 
116+ return end_time_prev_job_stage, end_time_prev_job_batch, completed_flag
117+ 
118+ begin_up = [num_stages - s for s in range(num_stages)]
119+ remaining = [num_microbatch * 2 - begin_up[p] for p in range(num_stages)]
120+ remaining_flag = True
121+ while remaining_flag:
122+ ids_old = []
123+ ids_new = []
124+ for s in range(num_stages):
125+ ids_old.append(remaining[s])
126+ if remaining[s]:
127+ idx = len(stage_list[0]) - remaining[s]
128+ end_time_prev_stage, end_time_prev_batch, job_flag = get_prev_job_time(start_time, stage_list, s, idx,
129+ fwd_bwd_chunk_stage,
130+ fwd_bwd_order,
131+ fwd_bwd_durations,
132+ comm_matrix)
133+ 
134+ if job_flag:
135+ start_time[stage_list[s][idx]] = max(end_time_prev_stage, end_time_prev_batch)
136+ remaining[s] = remaining[s] - 1
137+ 
138+ ids_new.append(remaining[s])
139+ if all(item == 0 for item in remaining):
140+ remaining_flag = False
141+ if ids_old == ids_new:
142+ break
143+ 
144+ e2e_time = start_time[f'B_0-{num_microbatch-1}'] + bwd_durations[0, -1]
145+ stage_start_time = [[start_time[job_name] for job_name in stage_list[s]] for s in range(num_stages)]
146+ return e2e_time, stage_start_time, stage_list, start_time
147+ 
148+ 
149+def find_integer_solutions(coefficients, global_batch_size):
150+ n = len(coefficients)
151+ mbs_max_value = (n + 1) // 2
152+ solutions = []
153+ all_comb = []
154+ for i in range(n):
155+ if i == mbs_max_value - 1:
156+ batch_using = sum(coefficients[0:mbs_max_value - 1] * 4)
157+ all_comb.append(list(range((global_batch_size - batch_using) // mbs_max_value,
158+ global_batch_size // mbs_max_value + 1)))
159+ else:
160+ all_comb.append(list(range(4)))
161+ 
162+ for x in product(*all_comb):
163+ if sum(coefficients[i] * x[i] for i in range(n)) == global_batch_size:
164+ solutions.append(x)
165+ 
166+ return solutions
167+ 
168+ 
169+def dynamic_mbs_search(num_stages, global_batch_size, fwd_mbs, bwd_mbs, comm_matrix):
170+ comp_mbs_ratio = [value / (index + 1) for index, value in enumerate(fwd_mbs)]
171+ fwd_mbs_selected = fwd_mbs[0:comp_mbs_ratio.index(min(comp_mbs_ratio)) + 1]
172+ bwd_mbs_selected = bwd_mbs[0:comp_mbs_ratio.index(min(comp_mbs_ratio)) + 1]
173+ mbs_max_value = len(fwd_mbs_selected)
174+ bwd_mbs_stages = [fwd_mbs_selected] * num_stages
175+ fwd_mbs_stages = [bwd_mbs_selected] * num_stages
176+ 
177+ coefficients = list(range(1, mbs_max_value + 1)) + list(range(mbs_max_value - 1, 0, -1))
178+ solutions = find_integer_solutions(coefficients, global_batch_size)
179+ 
180+ mbs_list = sum([solutions[0][i] * [coefficients[i]] for i in range(len(solutions[0]))], [])
181+ num_microbatch = len(mbs_list)
182+ fwd_durations = np.zeros([num_stages, num_microbatch])
183+ bwd_durations = np.zeros([num_stages, num_microbatch])
184+ for j in range(num_microbatch):
185+ for i in range(num_stages):
186+ fwd_durations[i, j] = fwd_mbs_stages[i][mbs_list[j] - 1]
187+ bwd_durations[i, j] = bwd_mbs_stages[i][mbs_list[j] - 1]
188+ 
189+ paras = Parallel_Paras(num_stages, fwd_durations, bwd_durations, num_microbatch, comm_matrix)
190+ e2e_time = []
191+ for sol in solutions:
192+ mbs_list = sum([sol[i] * [coefficients[i]] for i in range(len(sol))], [])
193+ num_microbatch = len(mbs_list)
194+ fwd_durations = np.zeros([num_stages, num_microbatch])
195+ bwd_durations = np.zeros([num_stages, num_microbatch])
196+ for j in range(num_microbatch):
197+ for i in range(num_stages):
198+ fwd_durations[i, j] = fwd_mbs_stages[i][mbs_list[j] - 1]
199+ bwd_durations[i, j] = bwd_mbs_stages[i][mbs_list[j] - 1]
200+ 
201+ paras.fwd_durations = fwd_durations
202+ paras.bwd_durations = bwd_durations
203+ paras.num_microbatch = num_microbatch
204+ 
205+ e2e_time0, stage_start_time0, stage_list0, start_time0 = dynamic_mbs_1f1b(paras)
206+ e2e_time.append(e2e_time0)
207+ 
208+ e2e_time_array = np.array(e2e_time)
209+ optimal_solution = solutions[e2e_time_array.argmin()]
210+ return optimal_solution, e2e_time_array.min()
211+ 
212+ 
213+def broadcast_oom_in_ranks(src_rank, policy):
214+ is_oom = [True]
215+ if torch.distributed.get_rank() == src_rank:
216+ is_oom = [policy]
217+ tmp_is_oom = torch.cuda.BoolTensor(is_oom)
218+ torch.distributed.broadcast(tmp_is_oom, src=src_rank)
219+ return tmp_is_oom.item()
220+ 
221+ 
222+def broadcast_mbs_in_ranks(src_rank, optimal_solution):
223+ args = get_args()
224+ solution_length = [0]
225+ if torch.distributed.get_rank() == src_rank:
226+ solution_length = [len(optimal_solution)]
227+ tmp_solution_length = torch.cuda.IntTensor(solution_length)
228+ torch.distributed.broadcast(tmp_solution_length, src=src_rank)
229+ solution_length = tmp_solution_length.item()
230+ 
231+ tmp_optimal_solution = [0] * solution_length
232+ if torch.distributed.get_rank() == src_rank:
233+ tmp_optimal_solution = optimal_solution
234+ tmp_optimal_solution = torch.cuda.IntTensor(tmp_optimal_solution)
235+ torch.distributed.broadcast(tmp_optimal_solution, src=src_rank)
236+ tmp_optimal_solution = tmp_optimal_solution.tolist()
237+ mbs_max_value = math.ceil(len(tmp_optimal_solution) / 2)
238+ coefficients = list(range(1, mbs_max_value + 1)) + list(range(mbs_max_value - 1, 0, -1))
239+ optimal_mbs_list = sum([tmp_optimal_solution[i] * [coefficients[i]] for i in range(len(tmp_optimal_solution))], [])
240+ args.optimized_mbs_list = optimal_mbs_list
241+ return optimal_mbs_list
242+ 
243+ 
244+def get_profiling_data(policy, args):
245+ instance = {"model_configs": {
246+ "hidden_size": args.hidden_size,
247+ "ffn_hidden_size": args.ffn_hidden_size,
248+ "seq_length": args.seq_length,
249+ "num_attention_heads": args.num_attention_heads
250+ }, "optimpipeline_policy": [{
251+ "num_layers": args.num_layers,
252+ "pipeline_model_parallel_size": args.pipeline_model_parallel_size,
253+ "tensor_model_parallel_size": args.tensor_model_parallel_size,
254+ "micro_batch_size": args.micro_batch_size,
255+ "global_batch_size": args.global_batch_size,
256+ "enable_scheduler": policy[0],
257+ "optimized_mbs_list": policy[1],
258+ "pp_schedule_list": policy[2],
259+ "optimal_layers": policy[3]
260+ }]}
261+ return instance
262+ 
263+ 
264+def save_profiling_data(policy, config_file):
265+ if torch.distributed.get_rank() % int(os.getenv('GPUS_PER_NODE', '8')) == 0:
266+ new_parse_args = parse_args_wrapper(parse_args)
267+ args = new_parse_args(None, False)
268+ instance = get_profiling_data(policy, args)
269+ if os.path.exists(config_file):
270+ with open(config_file, "r") as config_json:
271+ config_contents = config_json.read()
272+ parsed_contents = json.loads(config_contents)
273+ index = check_equal_model_configs(args, parsed_contents)
274+ if index != -1:
275+ if "optimpipeline_policy" in parsed_contents[index]:
276+ parsed_contents[index]["optimpipeline_policy"].append(instance["optimpipeline_policy"][0])
277+ else:
278+ parsed_contents.append(instance)
279+ with open(config_file, "w") as f:
280+ json.dump(parsed_contents, f, ensure_ascii=False)
281+ os.chmod(config_file, 0o644)
282+ else:
283+ with open(config_file, "w") as f:
284+ json.dump([instance], f, ensure_ascii=False)
285+ os.chmod(config_file, 0o644)
286+ 
287+ 
288+def solve_optimpipeline(args, data_parallel_size, global_context):
289+ mbs_max_value = len(global_context)
290+ coefficients = list(range(1, mbs_max_value + 1)) + list(range(mbs_max_value - 1, 0, -1))
291+ optimal_solution = [0] * len(coefficients)
292+ optimal_time = 0
293+ if torch.distributed.get_rank() == 0:
294+ num_stages = args.pipeline_model_parallel_size
295+ global_batch_size = args.global_batch_size // data_parallel_size
296+ fwd_mbs = [item[0] for item in global_context]
297+ bwd_mbs = [item[1] for item in global_context]
298+ comm_matrix = [[0.05] * num_stages for _ in range(num_stages)]
299+ for i in range(num_stages):
300+ comm_matrix[i][i] = 0
301+ 
302+ optimal_solution, optimal_time = dynamic_mbs_search(num_stages, global_batch_size, fwd_mbs, bwd_mbs, comm_matrix)
303+ torch.distributed.barrier()
304+ return optimal_solution, optimal_time
Amindspeed/core/performance/auto_pipeline_perf/schedulepipeline_solver.py+445-0
@@ -0,0 +1,445 @@
1+import time
2+import json
3+import numpy as np
4+import torch
5+import torch_npu
6+from megatron.training import get_args
7+from megatron.training import print_rank_0
8+ 
9+ 
10+class PipelineParallelParas:
11+ def __init__(self,
12+ num_stages,
13+ fwd_durations,
14+ bwd_durations,
15+ num_microbatches,
16+ comm_matrix,
17+ num_layers):
18+ self.num_stages = num_stages
19+ self.num_microbatches = num_microbatches
20+ self.fwd_durations = fwd_durations
21+ self.bwd_durations = bwd_durations
22+ self.comm_matrix = comm_matrix
23+ self.num_layers = num_layers
24+ 
25+ 
26+ 
27+def time_model_1f1b(paras):
28+ # obtain the E2E time for 1F1B scheme
29+ num_stages = paras.num_stages
30+ num_micro_batches = paras.num_microbatches
31+ fwd_durations = paras.fwd_durations
32+ bwd_durations = paras.bwd_durations
33+ p2p_matrix = paras.comm_matrix
34+ fwd_start = np.zeros([num_stages, num_micro_batches])
35+ bwd_start = np.zeros([num_stages, num_micro_batches])
36+ 
37+ warmup = [num_stages - s for s in range(num_stages)]
38+ remaining = [num_micro_batches - warmup[p] for p in range(num_stages)]
39+ # warm_up stage-0
40+ for m in range(num_stages):
41+ fwd_start[0, m] = m * fwd_durations[0]
42+ # warm_up stage
43+ for s in range(1, num_stages, 1):
44+ fwd_start[s, 0] = fwd_start[s - 1, 0] + fwd_durations[s - 1] + p2p_matrix[s - 1][s]
45+ for m in range(1, num_stages - s, 1):
46+ fwd_start[s, m] = max(fwd_start[s - 1, m] + fwd_durations[s - 1] + p2p_matrix[s - 1][s],
47+ fwd_start[s, m - 1] + fwd_durations[s])
48+ 
49+ # 0 micro batch at last stage bwd start
50+ bwd_start[num_stages - 1, 0] = fwd_start[num_stages - 1, 0] + fwd_durations[num_stages - 1]
51+ for s in range(num_stages - 2, -1, -1):
52+ bwd_start[s, 0] = bwd_start[s + 1, 0] + bwd_durations[s + 1] + p2p_matrix[s + 1][s]
53+ 
54+ # steady state
55+ for m in range(1, num_micro_batches, 1):
56+ # forward time
57+ for s in range(num_stages):
58+ if m > remaining[s]:
59+ continue
60+ if s == 0:
61+ fwd_start[s, m + num_stages - 1] = bwd_start[s, m - 1] + bwd_durations[s]
62+ else:
63+ fwd_start[s, m + num_stages - s - 1] = max(
64+ fwd_start[s - 1, m + num_stages - s - 1] + fwd_durations[s - 1] + p2p_matrix[s - 1][s],
65+ bwd_start[s, m - 1] + bwd_durations[s])
66+ 
67+ # backward time
68+ for s in range(num_stages - 1, -1, -1):
69+ # cool down stage
70+ if m + num_stages - s > num_micro_batches:
71+ bwd_start[s, m] = bwd_start[s + 1, m] + bwd_durations[s + 1] + p2p_matrix[s + 1][s]
72+ continue
73+ 
74+ if s == num_stages - 1:
75+ bwd_start[s, m] = fwd_start[s, m] + fwd_durations[s]
76+ else:
77+ bwd_start[s, m] = max(bwd_start[s + 1, m] + bwd_durations[s + 1] + p2p_matrix[s + 1][s],
78+ fwd_start[s, m + num_stages - s - 1] + fwd_durations[s])
79+ 
80+ e2e_time = bwd_start[0, -1] + bwd_durations[0]
81+ return e2e_time, fwd_start, bwd_start
82+ 
83+ 
84+def time_model_nfmb(paras, stage_schedule):
85+ # 给定一个调度序列，计算端到端时间
86+ num_stages = paras.num_stages
87+ num_mb = paras.num_microbatches
88+ comm_matrix = paras.comm_matrix
89+ chunk_placement = list(range(num_stages)) + list(range(num_stages - 1, -1, -1))
90+ # Fwd Bwd执行顺序
91+ fwd_bwd_comp_order = ([f'F_{i}' for i in range(num_stages)] +
92+ [f'B_{i}' for i in range(num_stages - 1, -1, -1)])
93+ chunk_stage_map = dict(zip(fwd_bwd_comp_order, chunk_placement))
94+ 
95+ if isinstance(stage_schedule, dict):
96+ stage_list = []
97+ for s in range(num_stages):
98+ fb_list = stage_schedule[f"stage{s}"]
99+ stage_list.append([element[0]+f"_{s}-"+element[1:] for element in fb_list])
100+ else:
101+ stage_list = stage_schedule
102+ 
103+ # 初始化
104+ fwd_bwd_list = ([f"F_{j}-{i}" for i in range(num_mb) for j in range(num_stages)]
105+ + [f"B_{j}-{i}" for i in range(num_mb) for j in range(num_stages)])
106+ values = [0 for _ in range(num_stages * num_mb * 2)]
107+ start_time = dict(zip(fwd_bwd_list, values))
108+ fwd_bwd_durations = dict()
109+ fwd_durations = np.array(paras.fwd_durations * num_mb).reshape(num_mb, num_stages).transpose()
110+ bwd_durations = np.array(paras.bwd_durations * num_mb).reshape(num_mb, num_stages).transpose()
111+ for j in range(num_stages):
112+ for i in range(num_mb):
113+ fwd_bwd_durations[f"F_{j}-{i}"] = fwd_durations[j, i]
114+ fwd_bwd_durations[f"B_{j}-{i}"] = bwd_durations[j, i]
115+ 
116+ start_time[f"F_{0}-{0}"] = 0.1
117+ for s in range(num_stages - 1):
118+ start_time[f"F_{s + 1}-{0}"] = start_time[f"F_{s}-{0}"] + fwd_durations[s, 0] + comm_matrix[s][s + 1]
119+ 
120+ # 获取当前任务的上一个任务以及依赖任务的结束时间
121+ def get_prev_task_time(task_start_time, task_list, pp_stage_id, mb_idx,
122+ chunk_stage_map, comp_order, model_chunk_times,
123+ comm_time_matrix):
124+ current_task = task_list[pp_stage_id][mb_idx]
125+ prev_task_same_stage = task_list[pp_stage_id][mb_idx - 1]
126+ chunk_id_prev_task_same_stage, _ = prev_task_same_stage.split('-')
127+ stage_id_prev_task = chunk_stage_map[chunk_id_prev_task_same_stage]
128+ chunk_position = comp_order.index(chunk_id_prev_task_same_stage)
129+ # 前一个任务计算完成后的通信时间
130+ if chunk_position < len(comp_order) - 1:
131+ stage_id_next = chunk_stage_map[comp_order[chunk_position + 1]]
132+ comm_time = comm_time_matrix[stage_id_prev_task][stage_id_next]
133+ else:
134+ comm_time = 0.01
135+ # 同一个stage上，前一个任务完成时间
136+ end_time_prev_task_stage = (task_start_time[prev_task_same_stage]
137+ + model_chunk_times[prev_task_same_stage]
138+ + comm_time)
139+ 
140+ # 相同micro batch id，上一个model chunk上的计算时间
141+ cur_model_chunk, cur_mb = current_task.split('-')
142+ chunk_position = comp_order.index(cur_model_chunk)
143+ if chunk_position > 0:
144+ prev_model_chunk = comp_order[chunk_position - 1]
145+ prev_task_batch = prev_model_chunk + '-' + cur_mb
146+ comm_time = comm_time_matrix[chunk_stage_map[prev_model_chunk]][chunk_stage_map[cur_model_chunk]]
147+ end_time_dependent_task_batch = (task_start_time[prev_task_batch]
148+ + model_chunk_times[prev_task_batch]
149+ + comm_time)
150+ completed_flag = task_start_time[prev_task_same_stage] > 0 and task_start_time[prev_task_batch] > 0
151+ else:
152+ end_time_dependent_task_batch = 0.1
153+ completed_flag = task_start_time[prev_task_same_stage] > 0
154+ 
155+ return end_time_prev_task_stage, end_time_dependent_task_batch, completed_flag
156+ 
157+ # 更新计算时间
158+ begin_up = [1] * num_stages
A
Aaijgnem12024年7月9日

规则6.7 禁止通过注释的方式删除无用的功能代码。

likedislike
159+ remaining = [num_mb * 2 - begin_up[p] for p in range(num_stages)]
160+ remaining_flag = True
161+ count = 0
162+ while remaining_flag:
163+ ids_old = []
164+ ids_new = []
165+ for s in range(num_stages):
166+ ids_old.append(remaining[s])
167+ if remaining[s]:
168+ microbatch_idx = len(stage_list[0]) - remaining[s]
169+ (end_time_prev_task_same_stage,
170+ end_time_dependent_task_same_microbatch,
171+ job_flag) = get_prev_task_time(start_time, stage_list, s, microbatch_idx, chunk_stage_map,
172+ fwd_bwd_comp_order, fwd_bwd_durations, comm_matrix)
173+ 
174+ if job_flag:
175+ start_time[stage_list[s][microbatch_idx]] = max(end_time_prev_task_same_stage,
176+ end_time_dependent_task_same_microbatch)
177+ remaining[s] = remaining[s] - 1
178+ 
179+ ids_new.append(remaining[s])
180+ 
181+ if all(item == 0 for item in remaining):
182+ remaining_flag = False
183+ 
184+ if ids_old == ids_new:
185+ count += 1
186+ if count == 3:
187+ start_time[f'B_0-{num_mb - 1}'] = 1e7
188+ break
189+ 
190+ e2e_time = start_time[f'B_0-{num_mb - 1}'] + bwd_durations[0, -1]
191+ stage_start_time = [[start_time[job_name] for job_name in stage_list[s]] for s in range(num_stages)]
192+ 
193+ return e2e_time, stage_start_time
194+ 
195+ 
196+def get_schedule_1f1b(paras):
197+ # generate 1f1b schedule list
198+ num_stages = paras.num_stages
199+ num_microbatches = paras.num_microbatches
200+ computation_placement = list(range(num_stages)) + list(range(num_stages - 1, -1, -1))
201+ 
202+ # Fwd Bwd执行顺序
203+ fwd_bwd_order = ([f'F_{i}' for i in range(num_stages)] +
204+ [f'B_{i}' for i in range(num_stages - 1, -1, -1)])
205+ 
206+ # 根据1F1B策略生成每个stage上的调度顺序
207+ def get_stage_list(fwd_seq, bwd_seq, num_advanced):
208+ stage_order = []
209+ n = len(fwd_seq)
210+ for idx in range(n):
211+ if idx < num_advanced:
212+ stage_order.append(fwd_seq[idx])
213+ else:
214+ stage_order.append(fwd_seq[idx])
215+ stage_order.append(bwd_seq[idx - num_advanced])
216+ if idx == n - 1:
217+ for i in range(num_advanced):
218+ stage_order.append(bwd_seq[i - num_advanced])
219+ 
220+ return stage_order
221+ 
222+ def get_stage_schedule(all_jobs_array, comp_placement, num_stages):
223+ stage_list = []
224+ for s in range(num_stages):
225+ stage_chunk_id = [index for index, element in enumerate(comp_placement) if element == s]
226+ warmup = num_stages - s
227+ stage_s_list = get_stage_list(all_jobs_array[stage_chunk_id[0]],
228+ all_jobs_array[stage_chunk_id[1]],
229+ warmup - 1)
230+ stage_list.append(stage_s_list)
231+ return stage_list
232+ 
233+ all_jobs = np.array([[s + f'-{i}' for i in range(num_microbatches)] for s in fwd_bwd_order])
234+ stage_list = get_stage_schedule(all_jobs, computation_placement, num_stages)
235+ stage_schedule_dict = dict()
236+ for s in range(paras.num_stages):
237+ stage_s_list = []
238+ for element in stage_list[s]:
239+ item1, item2 = element.split("-")
240+ stage_s_list.append(item1[0] + item2)
241+ stage_schedule_dict[f"stage{s}"] = stage_s_list
242+ return stage_schedule_dict
243+ 
244+ 
245+def get_schedule_eager1f1b(paras, num_forwards, layers_placement):
246+ # generate 1f1b schedule list
247+ num_stages = paras.num_stages
248+ num_microbatches = paras.num_microbatches
249+ # 将原始模型切分为多个model chunk，chunk在PP stage上的放置顺序
250+ chunk_placement = list(range(num_stages)) + list(range(num_stages - 1, -1, -1))
251+ 
252+ # Fwd Bwd执行顺序
253+ fwd_bwd_comp_order = ([f'F_{i}' for i in range(num_stages)] +
254+ [f'B_{i}' for i in range(num_stages - 1, -1, -1)])
255+ 
256+ # 根据1F1B策略生成每个stage上的调度顺序
257+ def get_stage_list(fwd_seq, bwd_seq, num_advanced):
258+ stage_order = []
259+ n = len(fwd_seq)
260+ for idx in range(n):
261+ if idx < num_advanced:
262+ stage_order.append(fwd_seq[idx])
263+ else:
264+ stage_order.append(fwd_seq[idx])
265+ stage_order.append(bwd_seq[idx - num_advanced])
266+ if idx == n - 1:
267+ for i in range(num_advanced):
268+ stage_order.append(bwd_seq[i - num_advanced])
269+ 
270+ return stage_order
271+ 
272+ def get_stage_schedule(all_jobs_array, comp_placement, num_advanced, paras, layers_placement):
273+ stage_list = []
274+ activations_num = int(paras.num_layers // paras.num_stages) * (num_advanced + paras.num_stages)
275+ nums_under_memory = [int(activations_num // layers_placement[i]) for i in range(paras.num_stages)]
276+ warmups = [min(nums_under_memory[s] - s - 1,
277+ 2 * paras.num_stages - 2 * s - 2) for s in range(paras.num_stages)]
278+ for i in range(paras.num_stages - 1):
279+ warmups[i + 1] = min(warmups[i] - 1, warmups[i + 1])
280+ warmups[i + 1] = max(warmups[i + 1], 0)
281+ 
282+ for s in range(paras.num_stages):
283+ stage_chunk_id = [index for index, element in enumerate(comp_placement) if element == s]
284+ num = sum(np.array(paras.bwd_durations[s + 1:])
285+ + np.array(paras.fwd_durations[s + 1:])) // np.array(paras.fwd_durations[s])
286+ stage_s_list = get_stage_list(all_jobs_array[stage_chunk_id[0]],
287+ all_jobs_array[stage_chunk_id[1]],
288+ warmups[s])
289+ stage_list.append(stage_s_list)
290+ return stage_list
291+ 
292+ all_jobs = np.array([[s + f'-{i}' for i in range(num_microbatches)] for s in fwd_bwd_comp_order])
293+ stage_list = get_stage_schedule(all_jobs, chunk_placement, num_forwards, paras, layers_placement)
294+ 
295+ # 转换为dictionary
296+ stage_schedule_dict = dict()
297+ for s in range(paras.num_stages):
298+ stage_s_list = []
299+ for element in stage_list[s]:
300+ item1, item2 = element.split("-")
301+ stage_s_list.append(item1[0] + item2)
302+ stage_schedule_dict[f"stage{s}"] = stage_s_list
303+ 
304+ return stage_schedule_dict
305+ 
306+ 
307+def schedule_layers(paras, num_mb_for_remaining_memory):
308+ # 调整层分布，对比层分布改变后，1F1B建模时间
309+ stage_layers = int(paras.num_layers // paras.num_stages)
310+ if paras.num_stages > 2:
311+ fwd_time_per_layer = sum(paras.fwd_durations[1:-1]) / (paras.num_stages - 2) / stage_layers
312+ bwd_time_per_layer = sum(paras.bwd_durations[1:-1]) / (paras.num_stages - 2) / stage_layers
313+ else:
314+ fwd_time_per_layer = paras.fwd_durations[0] / stage_layers
315+ bwd_time_per_layer = paras.bwd_durations[0] / stage_layers
316+ 
317+ # 1f1b as baseline
318+ e2e_time = np.ones([2, paras.num_stages]) * 1e9
319+ paras_all = []
320+ layers_placement = []
321+ schedule_1f1b = get_schedule_1f1b(paras)
322+ e2e_time[0, 0], stage_start_time1 = time_model_nfmb(paras, schedule_1f1b)
323+ paras_all.append(paras)
324+ layers_p1 = [stage_layers] * paras.num_stages
325+ layers_placement.append(layers_p1)
326+ # 调度序列
327+ schedule_eager_1f1b = get_schedule_eager1f1b(paras, num_mb_for_remaining_memory, layers_p1)
328+ e2e_time[1, 0], stage_start_time2 = time_model_nfmb(paras, schedule_eager_1f1b)
329+ 
330+ if stage_layers >= 2:
331+ for i in range(paras.num_stages - 1):
332+ fwd_new = np.array(paras.fwd_durations)
333+ fwd_new[i] += fwd_time_per_layer
334+ fwd_new[-1] -= fwd_time_per_layer
335+ bwd_new = np.array(paras.bwd_durations)
336+ bwd_new[i] += bwd_time_per_layer
337+ bwd_new[-1] -= bwd_time_per_layer
338+ paras1 = PipelineParallelParas(paras.num_stages,
339+ fwd_new.tolist(),
340+ bwd_new.tolist(),
341+ paras.num_microbatches,
342+ paras.comm_matrix,
343+ paras.num_layers)
344+ e2e_time[0, i + 1], stage_start_time1 = time_model_nfmb(paras1, schedule_1f1b)
345+ paras_all.append(paras1)
346+ layers_p1 = [stage_layers] * paras.num_stages
347+ layers_p1[i] += 1
348+ layers_p1[-1] -= 1
349+ layers_placement.append(layers_p1)
350+ schedule_eager_1f1b = get_schedule_eager1f1b(paras1, num_mb_for_remaining_memory, layers_p1)
351+ e2e_time[1, i + 1], stage_start_time2 = time_model_nfmb(paras1, schedule_eager_1f1b)
352+ 
353+ optimal_paras = paras_all[e2e_time[1, :].argmin()]
354+ optimal_layer = layers_placement[e2e_time[1, :].argmin()]
355+ schedule_scheme = get_schedule_eager1f1b(optimal_paras, num_mb_for_remaining_memory, optimal_layer)
356+ 
357+ return schedule_scheme, optimal_layer, e2e_time[1, :].min()
358+ 
359+ 
360+def broadcast_enable_schedule_in_ranks(src_rank, policy):
361+ enable_schedule = [False]
362+ if torch.distributed.get_rank() == src_rank:
363+ enable_schedule = [policy]
364+ tmp_enable_schedule = torch.cuda.BoolTensor(enable_schedule)
365+ torch.distributed.broadcast(tmp_enable_schedule, src=src_rank)
366+ return tmp_enable_schedule.item()
367+ 
368+ 
369+def broadcast_scheduler_in_ranks(src_rank, policy):
370+ args = get_args()
371+ policy_str = json.dumps(policy)
372+ byte_tensor = torch.cuda.ByteTensor(list(policy_str.encode()))
373+ torch.distributed.broadcast(byte_tensor, src_rank)
374+ if torch.distributed.get_rank() != 0:
375+ received_byte_tensor = torch.cuda.ByteTensor([0] * len(byte_tensor))
376+ else:
377+ received_byte_tensor = byte_tensor.clone()
378+ torch.distributed.broadcast(received_byte_tensor, src_rank)
379+ received_policy_str = ''.join([chr(byte) for byte in received_byte_tensor.tolist()])
380+ received_policy_data = json.loads(received_policy_str)
381+ args.pp_schedule_list = received_policy_data
382+ return received_policy_data
383+ 
384+ 
385+def broadcast_layer_in_ranks(src_rank, policy):
386+ args = get_args()
387+ num_layer_list = args.pipeline_model_parallel_size * [0]
388+ if torch.distributed.get_rank() == 0:
389+ num_layer_list = policy
390+ tmp_layer_list = torch.cuda.IntTensor(num_layer_list)
391+ torch.distributed.broadcast(tmp_layer_list, src=src_rank)
392+ args.num_layer_list = tmp_layer_list.tolist()
393+ return tmp_layer_list.tolist()
394+ 
395+ 
396+def all_gather_time(args, gather_time):
397+ recv_gather_time_list = []
398+ world_size = torch.distributed.get_world_size()
399+ gather_time = torch.cuda.FloatTensor([gather_time])
400+ gathered_tensors = [torch.zeros_like(gather_time) for _ in range(world_size)]
401+ torch.distributed.all_gather(gathered_tensors, gather_time)
402+ for rank, tensor in enumerate(gathered_tensors):
403+ pipeline_stage_rank = get_pipeline_stage_rank(world_size, args.pipeline_model_parallel_size, rank)
404+ recv_gather_time_list.append((pipeline_stage_rank, tensor.item()))
405+ return recv_gather_time_list
406+ 
407+ 
408+def average_time_by_rank(time_list):
409+ time_dict = {}
410+ for item in time_list:
411+ if item[0] not in time_dict:
412+ time_dict[item[0]] = item[1]
413+ else:
414+ time_dict[item[0]] += item[1]
415+ time_dict[item[0]] /= 2
416+ return time_dict
417+ 
418+ 
419+def get_pipeline_stage_rank(world_size, num_stages, global_rank):
420+ assert world_size % num_stages == 0, "World size must be divisible by the number of stages"
421+ assert global_rank < world_size, "Global rank must be less than world size"
422+ 
423+ stage_size = world_size // num_stages
424+ return global_rank // stage_size
425+ 
426+ 
427+def solve_pipelineschedule(args, data_parallel_size, num_forwards_first_stage, forward_time_dict, backward_time_dict):
428+ pipeline_stages = args.pipeline_model_parallel_size
429+ forward_time_each_stage = [forward_time_dict[rank] for rank in forward_time_dict]
430+ backward_time_each_stage = [backward_time_dict[rank] for rank in backward_time_dict]
431+ comm_matrix = [[0.05] * pipeline_stages for _ in range(pipeline_stages)]
432+ num_micro_batches = args.global_batch_size // data_parallel_size // args.micro_batch_size
433+ num_layers = args.num_layers
434+ 
435+ pp_paras = PipelineParallelParas(pipeline_stages,
436+ forward_time_each_stage,
437+ backward_time_each_stage,
438+ num_micro_batches,
439+ comm_matrix,
440+ num_layers)
441+ # FB schedule
442+ start_time = time.time()
443+ schedule_list, optimal_layers, schedule_time = schedule_layers(pp_paras, num_forwards_first_stage)
444+ end_time = time.time()
445+ return schedule_list, schedule_time, optimal_layers
Amindspeed/core/performance/auto_pipeline_perf/schedules.py+276-0
@@ -0,0 +1,276 @@
1+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
2+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+#
4+# Licensed under the Apache License, Version 2.0 (the "License");
5+# you may not use this file except in compliance with the License.
6+# You may obtain a copy of the License at
7+#
8+# http://www.apache.org/licenses/LICENSE-2.0
9+#
10+# Unless required by applicable law or agreed to in writing, software
11+# distributed under the License is distributed on an "AS IS" BASIS,
12+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+# See the License for the specific language governing permissions and
14+# limitations under the License.
15+import time
A
Aaijgnem12024年7月9日

license

likedislike
16+from functools import wraps
17+import contextlib
18+from typing import Callable, Iterator, List, Optional, Union
19+import torch
20+from torch.autograd.variable import Variable
21+from megatron.core import parallel_state
22+from megatron.core.enums import ModelType
23+from megatron.core.pipeline_parallel import p2p_communication
24+from megatron.core.utils import get_model_config, get_model_type
25+from megatron.training import get_args
26+from megatron.core.pipeline_parallel.schedules import forward_step, backward_step, deallocate_output_tensor, check_first_val_step
27+from mindspeed.core.performance.auto_pipeline_perf.autopipeline_perf import profile_context
28+import mindspeed.core.training as training
29+ 
30+ 
31+def get_forward_backward_func_decorator(get_forward_backward_func):
32+ @wraps(get_forward_backward_func)
33+ def wrapper(*args, **kwargs):
34+ argument = get_args()
35+ pipeline_model_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
36+ if pipeline_model_parallel_size > 1 and argument.automated_pipeline_perf and argument.optimized_mbs_list:
37+ forward_backward_func = optimized_forward_backward_pipelining
38+ else:
39+ forward_backward_func = get_forward_backward_func(*args, **kwargs)
40+ return forward_backward_func
41+ return wrapper
42+ 
43+ 
44+def forward_step_decorator(fn):
45+ @wraps(fn)
46+ def wrapper(*args, **kwargs):
47+ argument = get_args()
48+ if argument.automated_pipeline_perf and not (argument.optimized_mbs_list or argument.pp_schedule_list):
49+ torch.cuda.synchronize()
50+ start_time = time.time()
51+ output_tensor = fn(*args, **kwargs)
52+ torch.cuda.synchronize()
53+ profile_context["fwd_time"].append((time.time() - start_time) * 1000)
54+ else:
55+ output_tensor = fn(*args, **kwargs)
56+ return output_tensor
57+ 
58+ return wrapper
59+ 
60+ 
61+def backward_step_decorator(fn):
62+ @wraps(fn)
63+ def wrapper(*args, **kwargs):
64+ argument = get_args()
65+ if argument.automated_pipeline_perf and not (argument.optimized_mbs_list or argument.pp_schedule_list):
66+ torch.cuda.synchronize()
67+ start_time = time.time()
68+ input_tensor_grad = fn(*args, **kwargs)
69+ torch.cuda.synchronize()
70+ profile_context["bwd_time"].append((time.time() - start_time) * 1000)
71+ else:
72+ input_tensor_grad = fn(*args, **kwargs)
73+ return input_tensor_grad
74+ return wrapper
75+ 
76+ 
77+def get_tensor_shapes():
78+ args = get_args()
79+ tensor_shapes = []
80+ mbs = args.optimized_mbs_list
81+ for m in mbs:
82+ tensor_shapes.append((args.seq_length // parallel_state.get_context_parallel_world_size() // parallel_state.get_tensor_model_parallel_world_size(), m, args.hidden_size))
83+ return tensor_shapes
84+ 
85+ 
86+def optimized_forward_backward_pipelining(
87+ *,
88+ forward_step_func,
89+ data_iterator: Union[Iterator, List[Iterator]],
90+ model: Union[torch.nn.Module, List[torch.nn.Module]],
91+ num_microbatches: int,
92+ seq_length: int,
93+ micro_batch_size: int,
94+ decoder_seq_length: int = None,
95+ forward_only: bool = False,
96+ collect_non_loss_data: bool = False,
97+ first_val_step: bool = None,
98+):
99+ """Run non-interleaved 1F1B schedule, with reduced pipeline bubble.
100+ Returns dictionary with losses if the last stage, empty dict otherwise.
101+ """
102+ if isinstance(model, list):
103+ model = model[0]
104+ if isinstance(data_iterator, list):
105+ data_iterator = data_iterator[0]
106+ argument = get_args()
107+ config = get_model_config(model)
108+ model_type = get_model_type(model)
109+ tensor_shapes = get_tensor_shapes()
110+ cnt_fwd, cnt_bwd = 0, 0
111+ argument.mbs_idx = cnt_fwd
112+ argument.optimized_mbs_mode = True
113+ num_microbatches = len(argument.optimized_mbs_list)
114+ if config.overlap_p2p_comm:
115+ raise ValueError(
116+ "Optimized pipeline parallelism does not support overlapping p2p communication"
117+ )
118+ 
119+ # Disable async grad reductions
120+ no_sync_func = config.no_sync_func
121+ if no_sync_func is None:
122+ no_sync_func = contextlib.nullcontext
123+ no_sync_context = None
124+ 
125+ def disable_grad_sync():
126+ """Disable asynchronous grad reductions"""
127+ nonlocal no_sync_context
128+ if no_sync_context is None:
129+ no_sync_context = no_sync_func()
130+ no_sync_context.__enter__()
131+ 
132+ def enable_grad_sync():
133+ """Enable asynchronous grad reductions"""
134+ nonlocal no_sync_context
135+ if no_sync_context is not None:
136+ no_sync_context.__exit__(None, None, None)
137+ no_sync_context = None
138+ 
139+ disable_grad_sync()
140+ 
141+ # Compute number of warmup microbatches.
142+ num_warmup_microbatches = \
143+ (parallel_state.get_pipeline_model_parallel_world_size() -
144+ parallel_state.get_pipeline_model_parallel_rank() - 1)
145+ num_warmup_microbatches = min(
146+ num_warmup_microbatches,
147+ num_microbatches)
148+ num_microbatches_remaining = \
149+ num_microbatches - num_warmup_microbatches
150+ 
151+ input_tensors = []
152+ output_tensors = []
153+ forward_data_store = []
154+ rank = parallel_state.get_pipeline_model_parallel_rank()
155+ 
156+ # Run warmup forward passes.
157+ for i in range(num_warmup_microbatches):
158+ input_tensor = p2p_communication.recv_forward(config=config,
159+ tensor_shape=tensor_shapes[cnt_fwd])
160+ argument.micro_batch_size = argument.optimized_mbs_list[cnt_fwd]
161+ output_tensor = forward_step(
162+ forward_step_func,
163+ data_iterator,
164+ model,
165+ num_microbatches,
166+ input_tensor,
167+ forward_data_store,
168+ config,
169+ collect_non_loss_data,
170+ None,
171+ check_first_val_step(first_val_step, forward_only, i == 0),
172+ )
173+ p2p_communication.send_forward(output_tensor, config=config)
174+ cnt_fwd += 1
175+ input_tensors.append(input_tensor)
176+ output_tensors.append(output_tensor)
177+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
178+ 
179+ # Before running 1F1B, need to receive first forward tensor.
180+ # If all microbatches are run in warmup / cooldown phase, then no need to
181+ # receive this tensor here.
182+ if num_microbatches_remaining > 0:
183+ input_tensor = p2p_communication.recv_forward(config=config,
184+ tensor_shape=tensor_shapes[cnt_fwd])
185+ 
186+ # Run 1F1B in steady state.
187+ for i in range(num_microbatches_remaining):
188+ last_iteration = (i == (num_microbatches_remaining - 1))
189+ argument.micro_batch_size = argument.optimized_mbs_list[cnt_fwd]
190+ output_tensor = forward_step(
191+ forward_step_func,
192+ data_iterator,
193+ model,
194+ num_microbatches,
195+ input_tensor,
196+ forward_data_store,
197+ config,
198+ collect_non_loss_data,
199+ None,
200+ check_first_val_step(
201+ first_val_step, forward_only, (i == 0) and (num_warmup_microbatches == 0)
202+ ),
203+ )
204+ if forward_only:
205+ p2p_communication.send_forward(output_tensor, config=config)
206+ if not last_iteration:
207+ input_tensor = p2p_communication.recv_forward(tensor_shapes=tensor_shapes[cnt_fwd], config=config)
208+ else:
209+ output_tensor_grad = \
210+ p2p_communication.send_forward_recv_backward(output_tensor,
211+ tensor_shape=tensor_shapes[cnt_bwd], config=config)
212+ 
213+ cnt_fwd += 1
214+ #if argument.mbs_idx < len(argument.optimized_mbs_list):
215+ # argument.micro_batch_size = argument.optimized_mbs_list[argument.mbs_idx]
216+ # Add input_tensor and output_tensor to end of list, then pop from the
217+ # start of the list for backward pass.
218+ input_tensors.append(input_tensor)
219+ output_tensors.append(output_tensor)
220+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
221+ 
222+ if forward_only:
223+ if not last_iteration:
224+ input_tensor = p2p_communication.recv_forward(config=config,
225+ tensor_shape=tensor_shapes[cnt_fwd])
226+ else:
227+ input_tensor, output_tensor = input_tensors.pop(0), output_tensors.pop(0)
228+ if num_warmup_microbatches == 0 and last_iteration:
229+ if config.grad_sync_func is None or rank == 0:
230+ enable_grad_sync()
231+ 
232+ input_tensor_grad = \
233+ backward_step(input_tensor, output_tensor,
234+ output_tensor_grad, model_type, config)
235+ 
236+ if last_iteration:
237+ input_tensor = None
238+ p2p_communication.send_backward(input_tensor_grad, config=config)
239+ else:
240+ input_tensor = \
241+ p2p_communication.send_backward_recv_forward(
242+ input_tensor_grad, tensor_shape=tensor_shapes[cnt_fwd], config=config)
243+ cnt_bwd += 1
244+ 
245+ # Run cooldown backward passes.
246+ if not forward_only:
247+ for _ in range(num_warmup_microbatches):
248+ if i == num_warmup_microbatches - 1:
249+ if config.grad_sync_func is None or rank == 0:
250+ enable_grad_sync()
251+ 
252+ input_tensor = input_tensors.pop(0)
253+ output_tensor = output_tensors.pop(0)
254+ output_tensor_grad = p2p_communication.recv_backward(
255+ tensor_shape=tensor_shapes[cnt_bwd], config=config)
256+ 
257+ input_tensor_grad = \
258+ backward_step(input_tensor, output_tensor,
259+ output_tensor_grad, model_type, config)
260+ p2p_communication.send_backward(input_tensor_grad, config)
261+ cnt_bwd += 1
262+ 
263+ # Launch any remaining grad reductions.
264+ if no_sync_context is not None:
265+ enable_grad_sync()
266+ if config.grad_sync_func is not None:
267+ config.grad_sync_func(model.parameters())
268+ 
269+ if config.finalize_model_grads_func is not None and not forward_only:
270+ # Finalize model grads (perform full grad all-reduce / reduce-scatter for
271+ # data parallelism, layernorm all-reduce for sequence parallelism, and
272+ # embedding all-reduce for pipeline parallelism).
273+ config.finalize_model_grads_func([model])
274+ argument.optimized_mbs_mode = False
275+ argument.micro_batch_size = training.ORIGIN_MBS
276+ return forward_data_store
Amindspeed/core/performance/auto_pipeline_perf/transformer.py+17-0
@@ -0,0 +1,17 @@
1+from functools import wraps
2+from megatron.training import get_args
3+from mindspeed.model.transformer import generate_attention_mask
4+import mindspeed.model.transformer
5+ 
6+ 
7+def get_attention_mask_wrapper(get_attention_mask):
8+ @wraps(get_attention_mask)
9+ def wrapper(*args, **kwargs):
10+ argument = get_args()
11+ automated_pipeline_perf = argument.automated_pipeline_perf and argument.optimized_mbs_list
12+ if automated_pipeline_perf:
13+ generate_attention_mask()
14+ else:
15+ get_attention_mask(*args, **kwargs)
16+ return mindspeed.model.transformer._GLOBAL_ATTN_MASK
17+ return wrapper
Mmindspeed/core/pipeline_parallel/flexible_schedules.py+6-2
@@ -41,6 +41,7 @@ from megatron.core.pipeline_parallel.p2p_communication import (
41from megatron.core.parallel_state import get_pipeline_model_parallel_group41from megatron.core.parallel_state import get_pipeline_model_parallel_group
42from mindspeed.core.parallel_state import get_pipeline_parallel_group_for_new_stream42from mindspeed.core.parallel_state import get_pipeline_parallel_group_for_new_stream
43from mindspeed.core.weight_grad_store import WeightGradStore43from mindspeed.core.weight_grad_store import WeightGradStore
44+from megatron.training import get_args
44 45 
45 46 
46forward_comm_stream = None47forward_comm_stream = None
@@ -436,10 +437,13 @@ def forward_backward_pipelining_without_interleaving(
436 default_stream = torch.cuda.default_stream()437 default_stream = torch.cuda.default_stream()
437 438 
438 global scheduler_plan439 global scheduler_plan
439- if scheduler_plan is None:440+ arguments = get_args()
441+ key = 'stage{}'.format(parallel_state.get_pipeline_model_parallel_rank())
442+ if scheduler_plan is None and arguments.pp_schedule_list:
443+ scheduler_plan = arguments.pp_schedule_list.get(key)
444+ elif scheduler_plan is None and arguments.pp_schedule_list is None:
440 scheduler_plan = generate_1f1b_scheduler_plan(parallel_state.get_pipeline_model_parallel_world_size(),445 scheduler_plan = generate_1f1b_scheduler_plan(parallel_state.get_pipeline_model_parallel_world_size(),
441 num_microbatches)446 num_microbatches)
442- key = 'stage{}'.format(parallel_state.get_pipeline_model_parallel_rank())
443 scheduler_plan = scheduler_plan.get(key)447 scheduler_plan = scheduler_plan.get(key)
444 448 
445 config.batch_p2p_comm = False449 config.batch_p2p_comm = False
Mmindspeed/core/pipeline_parallel/schedules.py+3-0
@@ -27,6 +27,9 @@ def get_forward_backward_func_wrapper(get_forward_backward_func):
27 if arguments.optimize_send_recv_comm and arguments.num_layers_per_virtual_pipeline_stage is None:27 if arguments.optimize_send_recv_comm and arguments.num_layers_per_virtual_pipeline_stage is None:
28 return flexible_schedules.forward_backward_pipelining_without_interleaving28 return flexible_schedules.forward_backward_pipelining_without_interleaving
29 29 
30+ if arguments.automated_pipeline_perf and arguments.pp_schedule_list:
31+ return flexible_schedules.forward_backward_pipelining_without_interleaving
32+ 
30 if arguments.recompute_in_bubble or arguments.recompute_in_advance:33 if arguments.recompute_in_bubble or arguments.recompute_in_advance:
31 return forward_backward_ripipe_pipelining34 return forward_backward_ripipe_pipelining
32 35 
Mmindspeed/core/training.py+91-9
@@ -12,30 +12,103 @@ from megatron.core import parallel_state
12from megatron.training import print_rank_012from megatron.training import print_rank_0
13from megatron.training.arguments import parse_args13from megatron.training.arguments import parse_args
14from mindspeed.core.memory.auto_pipeline.autopipeline import autopipeline_profiling14from mindspeed.core.memory.auto_pipeline.autopipeline import autopipeline_profiling
15+from mindspeed.core.performance.auto_pipeline_perf.autopipeline_perf import (autopipelineperf_profiling, check_out_of_memory,
16+ calculate_num_of_activations, check_skip_profiling,
17+ broadcast_skip_in_ranks)
18+from mindspeed.core.performance.auto_pipeline_perf.optimpipeline_solver import solve_optimpipeline, broadcast_oom_in_ranks, broadcast_mbs_in_ranks, save_profiling_data
19+from mindspeed.core.performance.auto_pipeline_perf.schedulepipeline_solver import (solve_pipelineschedule, broadcast_enable_schedule_in_ranks,
20+ broadcast_scheduler_in_ranks, broadcast_layer_in_ranks,
21+ all_gather_time, average_time_by_rank)
15from mindspeed.core.memory.auto_pipeline.autopipeline_apply import apply_autopipeline22from mindspeed.core.memory.auto_pipeline.autopipeline_apply import apply_autopipeline
16from mindspeed.core.memory.auto_pipeline.autopipeline_solver import solve_autopipeline, broadcast_policy_in_ranks, destroy_global_vars23from mindspeed.core.memory.auto_pipeline.autopipeline_solver import solve_autopipeline, broadcast_policy_in_ranks, destroy_global_vars
17from mindspeed.arguments import parse_args_wrapper24from mindspeed.arguments import parse_args_wrapper
18 25 
19 26 
20-policy = None27+POLICY = None
28+OPTIMIZED_MBS_LIST = None
29+PP_SCHEDULE_LIST = None
30+OPTIMAL_LAYERS = None
31+ORIGIN_MBS = None
32+DATA_PARALLEL_SIZE = 1
33+ENABLE_SCHEDULER = False
21 34 
22 35 
23def pretrain_decorator(pretrain):36def pretrain_decorator(pretrain):
24 @wraps(pretrain)37 @wraps(pretrain)
25 def wrapper(*args, **kwargs):38 def wrapper(*args, **kwargs):
39+ global POLICY
40+ global OPTIMIZED_MBS_LIST
41+ global PP_SCHEDULE_LIST
42+ global OPTIMAL_LAYERS
43+ global ORIGIN_MBS
44+ global DATA_PARALLEL_SIZE
45+ global ENABLE_SCHEDULER
26 new_parse_args = parse_args_wrapper(parse_args)46 new_parse_args = parse_args_wrapper(parse_args)
27 argument = new_parse_args(None, False)47 argument = new_parse_args(None, False)
28 if argument.automated_pipeline and not argument.num_layer_list:48 if argument.automated_pipeline and not argument.num_layer_list:
29- global policy49+ context, POLICY = autopipeline_profiling(args[1], args[2], args[3],
30- context, policy = autopipeline_profiling(args[1], args[2], args[3],50+ args[0], None, argument)
31- args[0], None, argument)
32 if context:51 if context:
33- policy = solve_autopipeline(context)52+ POLICY = solve_autopipeline(context)
34 parallel_state.destroy_global_memory_buffer()53 parallel_state.destroy_global_memory_buffer()
35 parallel_state.destroy_model_parallel()54 parallel_state.destroy_model_parallel()
36 destroy_global_vars()55 destroy_global_vars()
37 gc.collect()56 gc.collect()
38 torch.cuda.empty_cache()57 torch.cuda.empty_cache()
58+ 
59+ if argument.automated_pipeline_perf:
60+ ORIGIN_MBS = argument.micro_batch_size
61+ is_skip, exist_policy = check_skip_profiling(argument, config_file="autopipeline_perf_config.json")
62+ if not is_skip:
63+ global_context = []
64+ mbs_time, pp_schedule_time = 0, 0
65+ mbs_tries = 1
66+ num_forwards_first_stage = 0
67+ is_oom = False
68+ forward_time_dict = {}
69+ backward_time_dict = {}
70+ 
71+ while mbs_tries < ORIGIN_MBS + 2:
72+ context = autopipelineperf_profiling(mbs_tries, args[1], args[2], args[3],
73+ args[0], None)
74+ if mbs_tries == ORIGIN_MBS:
75+ schedule_context = context
76+ forward_time_list = all_gather_time(argument, schedule_context['fwd_time'])
77+ forward_time_dict = average_time_by_rank(forward_time_list)
78+ backward_time_list = all_gather_time(argument, schedule_context['bwd_time'])
79+ backward_time_dict = average_time_by_rank(backward_time_list)
80+ num_forwards_first_stage = calculate_num_of_activations(schedule_context)
81+ 
82+ parallel_state.destroy_global_memory_buffer()
83+ parallel_state.destroy_model_parallel()
84+ destroy_global_vars()
85+ gc.collect()
86+ torch.cuda.empty_cache()
87+ global_context.append((context['fwd_time'], context['bwd_time'], context['comm_time']))
88+ DATA_PARALLEL_SIZE = context['data_parallel_size']
89+ if not is_oom:
90+ is_oom = check_out_of_memory(argument, context, mbs_tries)
91+ is_oom = broadcast_oom_in_ranks(0, is_oom)
92+ mbs_tries += 1
93+ if mbs_tries <= ORIGIN_MBS and is_oom:
94+ raise AssertionError(
95+ 'A risk of Out of Memory could occur, please '
96+ 'reset to a smaller micro batch size.')
97+ if mbs_tries > ORIGIN_MBS and is_oom:
98+ break
99+ if len(global_context) > 0:
100+ OPTIMIZED_MBS_LIST, mbs_time = solve_optimpipeline(argument, DATA_PARALLEL_SIZE, global_context)
101+ PP_SCHEDULE_LIST, pp_schedule_time, OPTIMAL_LAYERS = solve_pipelineschedule(argument, DATA_PARALLEL_SIZE, num_forwards_first_stage, forward_time_dict, backward_time_dict)
102+ if torch.distributed.get_rank() == 0 and mbs_time > pp_schedule_time and num_forwards_first_stage > 2:
103+ ENABLE_SCHEDULER = True
104+ ENABLE_SCHEDULER = broadcast_enable_schedule_in_ranks(0, ENABLE_SCHEDULER)
105+ optimized_policy = (ENABLE_SCHEDULER, OPTIMIZED_MBS_LIST, PP_SCHEDULE_LIST, OPTIMAL_LAYERS)
106+ save_profiling_data(optimized_policy, config_file="autopipeline_perf_config.json")
107+ else:
108+ ENABLE_SCHEDULER = exist_policy[0]
109+ OPTIMIZED_MBS_LIST = exist_policy[1]
110+ PP_SCHEDULE_LIST = exist_policy[2]
111+ OPTIMAL_LAYERS = exist_policy[3]
39 pretrain(*args, **kwargs)112 pretrain(*args, **kwargs)
40 return wrapper113 return wrapper
41 114 
@@ -43,15 +116,24 @@ def pretrain_decorator(pretrain):
43def setup_model_and_optimizer_decorator(setup_model_and_optimizer):116def setup_model_and_optimizer_decorator(setup_model_and_optimizer):
44 @wraps(setup_model_and_optimizer)117 @wraps(setup_model_and_optimizer)
45 def wrapper(*args, **kwargs):118 def wrapper(*args, **kwargs):
46- global policy119+ global POLICY
47- if policy:120+ global OPTIMIZED_MBS_LIST
121+ global PP_SCHEDULE_LIST
122+ global OPTIMAL_LAYERS
123+ global ENABLE_SCHEDULER
124+ argument = get_args()
125+ if argument.automated_pipeline and POLICY:
48 if torch.distributed.get_rank() == 0:126 if torch.distributed.get_rank() == 0:
49 broadcast_policy_in_ranks(0, policy)127 broadcast_policy_in_ranks(0, policy)
50 else:128 else:
51 broadcast_policy_in_ranks(0)129 broadcast_policy_in_ranks(0)
130+ if argument.automated_pipeline_perf and ENABLE_SCHEDULER:
131+ broadcast_scheduler_in_ranks(0, PP_SCHEDULE_LIST)
132+ broadcast_layer_in_ranks(0, OPTIMAL_LAYERS)
133+ elif argument.automated_pipeline_perf and OPTIMIZED_MBS_LIST:
134+ broadcast_mbs_in_ranks(0, OPTIMIZED_MBS_LIST)
52 model, optimizer, opt_param_scheduler = setup_model_and_optimizer(*args, **kwargs)135 model, optimizer, opt_param_scheduler = setup_model_and_optimizer(*args, **kwargs)
53- args = get_args()136+ if argument.recompute_module_list:
54- if args.recompute_module_list:
55 apply_autopipeline(model)137 apply_autopipeline(model)
56 return model, optimizer, opt_param_scheduler138 return model, optimizer, opt_param_scheduler
57 return wrapper139 return wrapper
Mmindspeed/megatron_adaptor.py+13-0
@@ -146,6 +146,7 @@ def megatron_core_adaptation(aspm):
146 from .optimizer.optimizer import (mixed_precision_optimizer_step, \146 from .optimizer.optimizer import (mixed_precision_optimizer_step, \
147 reuse_fp32_param_init_wrapper, optimizer_config_init_wrapper)147 reuse_fp32_param_init_wrapper, optimizer_config_init_wrapper)
148 from .core.pipeline_parallel.schedules import get_forward_backward_func_wrapper148 from .core.pipeline_parallel.schedules import get_forward_backward_func_wrapper
149+ from .core.performance.auto_pipeline_perf.schedules import get_forward_backward_func_decorator, backward_step_decorator, forward_step_decorator
149 from .core.pipeline_parallel.p2p_communication import _communicate_shapes150 from .core.pipeline_parallel.p2p_communication import _communicate_shapes
150 from .optimizer.distrib_optimizer import reuse_fp32_param_distrib_optimizer_init_wrapper151 from .optimizer.distrib_optimizer import reuse_fp32_param_distrib_optimizer_init_wrapper
151 from .core.models.common.embeddings.rotary_pos_embedding import get_pos_emb_on_this_cp_rank152 from .core.models.common.embeddings.rotary_pos_embedding import get_pos_emb_on_this_cp_rank
@@ -195,6 +196,12 @@ def megatron_core_adaptation(aspm):
195 dot_product_attention_forward_wrapper)196 dot_product_attention_forward_wrapper)
196 aspm.register_patch('megatron.core.pipeline_parallel.schedules.get_forward_backward_func',197 aspm.register_patch('megatron.core.pipeline_parallel.schedules.get_forward_backward_func',
197 get_forward_backward_func_wrapper)198 get_forward_backward_func_wrapper)
199+ aspm.register_patch('megatron.core.pipeline_parallel.schedules.get_forward_backward_func',
200+ get_forward_backward_func_decorator)
201+ aspm.register_patch('megatron.core.pipeline_parallel.schedules.backward_step',
202+ backward_step_decorator)
203+ aspm.register_patch('megatron.core.pipeline_parallel.schedules.forward_step',
204+ forward_step_decorator)
198 aspm.register_patch('megatron.core.pipeline_parallel.p2p_communication._communicate_shapes',205 aspm.register_patch('megatron.core.pipeline_parallel.p2p_communication._communicate_shapes',
199 _communicate_shapes)206 _communicate_shapes)
200 207 
@@ -240,6 +247,10 @@ def megatron_legacy_adaptation(aspm):
240 from .model.transformer import switch_mlp_init_wrapper, switch_mlp_forward_wrapper, \247 from .model.transformer import switch_mlp_init_wrapper, switch_mlp_forward_wrapper, \
241 parallel_transformer_layer_init_wrapper248 parallel_transformer_layer_init_wrapper
242 from .model.language_model import parallel_lm_logits, embedding_forward_wrapper249 from .model.language_model import parallel_lm_logits, embedding_forward_wrapper
250+ from .core.performance.auto_pipeline_perf.data_samplers import build_pretraining_data_loader_decorator
251+ from .core.performance.auto_pipeline_perf.transformer import get_attention_mask_wrapper
252+ aspm.register_patch('mindspeed.model.transformer.get_attention_mask', get_attention_mask_wrapper)
253+ aspm.register_patch('megatron.legacy.data.data_samplers.build_pretraining_data_loader', build_pretraining_data_loader_decorator)
243 aspm.register_patch('megatron.legacy.model.language_model.parallel_lm_logits', parallel_lm_logits)254 aspm.register_patch('megatron.legacy.model.language_model.parallel_lm_logits', parallel_lm_logits)
244 aspm.register_patch('megatron.legacy.model.language_model.Embedding.forward', embedding_forward_wrapper)255 aspm.register_patch('megatron.legacy.model.language_model.Embedding.forward', embedding_forward_wrapper)
245 from .model.gpt_model import post_language_model_processing_wrapper256 from .model.gpt_model import post_language_model_processing_wrapper
@@ -284,6 +295,7 @@ def megatron_legacy_adaptation(aspm):
284 295 
285 296 
286def megatron_training_adaptation(aspm):297def megatron_training_adaptation(aspm):
298+ from .core.performance.auto_pipeline_perf.global_vars import get_num_microbatches_wrapper
287 from .initialize import _compile_dependencies, set_jit_fusion_options_wrapper299 from .initialize import _compile_dependencies, set_jit_fusion_options_wrapper
288 from .utils import get_batch_on_this_cp_rank300 from .utils import get_batch_on_this_cp_rank
289 from .training import pretrain301 from .training import pretrain
@@ -291,6 +303,7 @@ def megatron_training_adaptation(aspm):
291 from .tokenizer import build_tokenizer_wrapper303 from .tokenizer import build_tokenizer_wrapper
292 from .yaml_arguments import core_transformer_config_from_yaml_wrapper, print_args_wrapper304 from .yaml_arguments import core_transformer_config_from_yaml_wrapper, print_args_wrapper
293 from .core.training import pretrain_decorator, setup_model_and_optimizer_decorator, save_checkpoint_and_time_decorator305 from .core.training import pretrain_decorator, setup_model_and_optimizer_decorator, save_checkpoint_and_time_decorator
306+ aspm.register_patch('megatron.training.global_vars.get_num_microbatches', get_num_microbatches_wrapper)
294 aspm.register_patch('megatron.training.training.pretrain', pretrain_decorator)307 aspm.register_patch('megatron.training.training.pretrain', pretrain_decorator)
295 aspm.register_patch('megatron.training.training.setup_model_and_optimizer', setup_model_and_optimizer_decorator)308 aspm.register_patch('megatron.training.training.setup_model_and_optimizer', setup_model_and_optimizer_decorator)
296 aspm.register_patch('megatron.training.training.save_checkpoint_and_time', save_checkpoint_and_time_decorator)309 aspm.register_patch('megatron.training.training.save_checkpoint_and_time', save_checkpoint_and_time_decorator)