提交Flex Parallel：多维并行配置自动寻优算法的代码实现-MindSpeed-AtomGit

Adocs/features/Automatic_Parallelism.md+157-0

@@ -0,0 +1,157 @@
	1	+## Automatic Parallelism
	2	+
	3	+## 问题分析
	4	+
	5	+当前主流的大模型并行训练方法有PP、TP、DP、SP、CP、Ulyssess Parallel（UP）、VPP、EP等，在内存、计算、通信方面都有不同的优化，直接叠加。大模型端到端训练性能由模型结构、集群规模、并行配置、batch_size等因素共同决定，在调优时需要综合考虑。当前并行配置人工调优需要大量的专家经验、人工分析和实验调优，预计数天~数周，实验成本高。相似模型的最优并行配置也并不相同，仍需花费时间进行优化。随着搜索空间变大，依赖手工调优变的不可行。例如，llama65B模型在4*8的集群规模下，仅考虑PP、TP、DP、SP、VP、mbs六个维度，配置组合有812种，手工调优时间成本太高。因此，需要构建自动并行系统根据模型结构和集群规模给用户自动推荐一个性能较优的并行配置策略。
	6	+
	7	+## 解决方案
	8	+
	9	+针对该问题场景提出多维并行配置自动寻优算法，在给定模型结构、集群配置的条件下，用户仅需要在启动脚本中配置相关参数即可启动多维并行配置自动寻优，在规定时间内找到较优的并行配置推荐给用户。算法原理图如下：
	10	+
	11	+* 内存自适应感知的搜索空间构建：考虑模型结构和集群信息约束，采用内存灰盒模型排除OOM并行配置，缩小搜索空间；
	12	+* 基于算子不确定性估计的高保序性Cost Model建模方法：引入低保真数据（单算子调用）作为先验信息，结合算子整网性能数据构建算子执行耗时的不确定性模型，结合通信耗时根据并行策略合成得到端到端性能的概率分布模型。
	13	+* 基于概率匹配的高效搜索算法：基于Thompson Sampling方法探索并行策略，以高概率探索高价值并行配置，提高探索效率，灵活支持探索早停，提高易用性。
	14	+
	15	+
	16	+![1](../../sources/images/auto_parallel_1.png)
	17	+
	18	+并行配置的支持情况：
	19	+
	20	+已支持搜索的并行配置维度：
	21	+
	22	+- [x] PP
	23	+- [x] TP
	24	+- [x] DP
	25	+- [x] CP
	26	+- [x] DeepSpeed-Ulyssess
	Z Z zhao-yifan272024年7月27日这里是不是跟仓上的长序列并行描述统一一下，CP Ring Attention和CP Ulyssess
	27	+- [x] Megatron-SP
	28	+- [x] mbs
	29	+
	30	+正在支持的并行配置维度：
	31	+
	32	+- [ ] MOE
	33	+- [ ] VP
	34	+- [ ] 自适应重计算
	35	+
	36	+## 使用方法
	37	+
	38	+在使用多维自动并行特性时，需使用python作为脚本启动器，在所有的节点上拉起脚本，并配置多维自动并行相关的参数。相关参数及其函数如下表所示：
	39	+
	40	+\| 参数名 \| 参数含义 \|
	41	+\| ---------------- \| ------------------------------------------------- \|
	42	+\| --auto-parallel \| 多维自动并行特性总开关 \|
	43	+\| --nodes \| 集群中节点的个数 \|
	44	+\| --nproc-per-node \| 每个节点中NPU的个数 \|
	45	+\| --master-addr \| 集群中主节点的IP地址 \|
	46	+\| --master-port \| 用于通信的端口号，各节点需要配置相同的端口号 \|
	47	+\| --node-rank \| 集群中节点的rank，主节点为0，其他节点为1,2,······ \|
	48	+
	49	+注：算法最长搜索时间为8小时，支持灵活提前退出，无需人工干预。
	50	+
	51	+下面是基于llama7B模型的配置示例：
	52	+
	53	+```shell
	54	+#!/bin/bash
	55	+export CUDA_DEVICE_MAX_CONNECTIONS=1
	56	+export NPU_ASD_ENABLE=0
	57	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	58	+
	59	+MASTER_ADDR=localhost
	60	+MASTER_PORT=6001
	61	+GPUS_PER_NODE=8
	62	+NNODES=1
	63	+NODE_RANK=0
	64	+
	65	+LOAD_CHECKPOINT_PATH=./ckpt
	66	+SAVE_CHECKPOINT_PATH=./ckpt
	67	+DATA_PATH={your dataset path}
	68	+TOKENIZER_MODEL={your tokenizer model path}
	69	+TP=1
	70	+PP=8
	71	+
	72	+DISTRIBUTED_ARGS="
	73	+ --nproc_per_node $GPUS_PER_NODE \
	74	+ --nnodes $NNODES \
	75	+ --node_rank $NODE_RANK \
	76	+ --master_addr $MASTER_ADDR \
	77	+ --master_port $MASTER_PORT
	78	+"
	79	+
	80	+GPT_ARGS="
	81	+ --tensor-model-parallel-size ${TP} \
	82	+ --pipeline-model-parallel-size ${PP} \
	83	+ --sequence-parallel \
	84	+ --num-layers 32 \
	85	+ --hidden-size 4096 \
	86	+ --ffn-hidden-size 11008 \
	87	+ --num-attention-heads 32 \
	88	+ --tokenizer-type Llama2Tokenizer \
	89	+ --tokenizer-model ${TOKENIZER_MODEL} \
	90	+ --seq-length 2048 \
	91	+ --max-position-embeddings 2048 \
	92	+ --micro-batch-size 4 \
	93	+ --global-batch-size 256 \
	94	+ --make-vocab-size-divisible-by 1 \
	95	+ --lr 1.0e-6 \
	96	+ --train-iters 5000 \
	97	+ --lr-decay-style cosine \
	98	+ --untie-embeddings-and-output-weights \
	99	+ --disable-bias-linear \
	100	+ --attention-dropout 0.0 \
	101	+ --init-method-std 0.01 \
	102	+ --hidden-dropout 0.0 \
	103	+ --position-embedding-type rope \
	104	+ --normalization RMSNorm \
	105	+ --use-fused-rmsnorm \
	106	+ --swiglu \
	107	+ --use-flash-attn \
	108	+ --no-masked-softmax-fusion \
	109	+ --attention-softmax-in-fp32 \
	110	+ --min-lr 1.0e-7 \
	111	+ --weight-decay 1e-1 \
	112	+ --lr-warmup-fraction 0.01 \
	113	+ --clip-grad 1.0 \
	114	+ --adam-beta1 0.9 \
	115	+ --initial-loss-scale 65536 \
	116	+ --adam-beta2 0.95 \
	117	+ --no-gradient-accumulation-fusion \
	118	+ --load ${LOAD_CHECKPOINT_PATH} \
	119	+ --no-load-optim \
	120	+ --no-load-rng \
	121	+ --fp16
	122	+"
	123	+
	124	+DATA_ARGS="
	125	+ --data-path $DATA_PATH \
	126	+ --split 100,0,0
	127	+"
	128	+
	129	+OUTPUT_ARGS="
	130	+ --log-interval 1 \
	131	+ --save-interval 10000 \
	132	+ --eval-interval 1000 \
	133	+ --eval-iters 0 \
	134	+"
	135	+
	136	+SEARCH_ARGS="
	137	+ --auto-parallel \
	138	+ --nnodes $NNODES \
	139	+ --nproc-per-node $GPUS_PER_NODE \
	140	+ --master-addr $MASTER_ADDR \
	141	+ --master-port $MASTER_PORT \
	142	+ --node-rank $NODE_RANK \
	143	+"
	144	+
	145	+python pretrain_gpt.py \
	146	+ $GPT_ARGS \
	147	+ $DATA_ARGS \
	148	+ $OUTPUT_ARGS \
	149	+ $SEARCH_ARGS \
	150	+ --distributed-backend nccl \
	151	+ \| tee logs/search_llama_7b.txt
	152	+```
	153	+
	154	+## 使用效果
	155	+
	156	+![2](../../sources/images/auto_parallel_2.png)
	157	+

Mmindspeed/arguments.py+17-0

	@@ -41,6 +41,7 @@ def process_args(parser):
41	parser = _add_ndmm_args(parser)	41	parser = _add_ndmm_args(parser)
42	parser = _add_coc_args(parser)	42	parser = _add_coc_args(parser)
43	parser = _add_profile_args(parser)	43	parser = _add_profile_args(parser)
		44	+ parser = _add_auto_parallel_args(parser)
44		45
45	return parser	46	return parser
46		47
	@@ -591,3 +592,19 @@ def _add_ndmm_args(parser):
591	group.add_argument('--nd2-dim1-size', type=int, default=1,	592	group.add_argument('--nd2-dim1-size', type=int, default=1,
592	help='Dim1 of the second nd matmul when use-3d-matmul is True')	593	help='Dim1 of the second nd matmul when use-3d-matmul is True')
593	return parser	594	return parser
		595	+
		596	+
		597	+def _add_auto_parallel_args(parser):
		598	+ group = parser.add_argument_group(title='auto_parallel')
		599	+ group.add_argument('--auto-parallel', action='store_true',
		600	+ help='enable automatic parallelism with auto-parallel')
		601	+ group.add_argument('--nnodes', type=int, default=1, help='the number of node in the cluster')
		602	+ group.add_argument('--nproc-per-node', type=int, default=8, help='the number of NPU on each node')
		603	+ group.add_argument('--master-addr', type=str, default=None, help='the ip-address of master node')
		604	+ group.add_argument('--master-port', type=str, default=None, help='the ip-port of master node')
		605	+ group.add_argument('--node-rank', type=int, default=0,
		606	+ help='the rank of nodes in the cluster, starting from 0 and increment by 1')
		607	+ group.add_argument('--profile-operator', action='store_true', help='')
		608	+ group.add_argument('--profile-memory', action='store_true', help='')
		609	+ group.add_argument('--prof-file', type=str, default=None, help='')
		610	+ return parser

Amindspeed/core/auto_parallel/__init__.py+243-0

@@ -0,0 +1,243 @@
	1	+# coding=utf-8
	2	+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+import os
	A A aijgnem12024年7月29日 license
	17	+import json
	18	+import operator
	19	+from functools import reduce
	20	+import datetime
	21	+import threading
	A A aijgnem12024年7月29日引用顺序按照先引用标准库再引用第三方库，再引用自定义模块的顺序，且三种引用间使用空行分割。
	22	+
	23	+import torch
	24	+import numpy as np
	25	+import pandas as pd
	26	+
	27	+KV_STORE = None
	28	+ITERATION_LOOP_TIME = 5
	29	+BAND_WIDTH_UNIDIRECTIONAL = 25 # GB/s
	30	+SEARCH_CACHE_PATH = None
	31	+MODULE_PATTERN = 'PP{}_TP{}_DP{}_CP{}_UP{}_MBS{}_MODULE.json'
	32	+OPERATOR_PATTERN = 'PP{}_TP{}_DP{}_CP{}_UP{}_MBS{}_OPERATOR'
	33	+
	34	+
	35	+# Operator dims after merging
	36	+ARD_NUM_DIMS = {
	37	+ 'MatMul': 3,
	38	+ 'BatchMatMul': 4,
	39	+ 'Softmax': 4,
	40	+ 'SoftmaxGrad': 4,
	41	+ 'RmsNorm': 3,
	42	+ 'RmsNormGrad': 3,
	43	+ 'FlashAttentionScore': 3,
	44	+ 'FlashAttentionScoreGrad': 3
	45	+}
	46	+
	47	+
	48	+# profiling data filed
	49	+class KeyField:
	50	+ OpType = 'Type'
	51	+ InputShapes = 'Input Shapes'
	52	+ OutputShapes = 'Output Shapes'
	53	+ Duration = 'Duration(us)'
	54	+ FwdTime = 'fwd_time'
	55	+ BwdTime = 'bwd_time'
	56	+
	57	+
	58	+class GlobalMemoryBuffer:
	59	+ buffers_length = [0, 0, 0]
	60	+ buffers = [None, None, None]
	61	+
	62	+ @staticmethod
	63	+ def get_tensor(shape: list, index):
	64	+ if index not in (0, 1, 2):
	65	+ raise AssertionError('index must be 0, 1, 2')
	66	+ data_type = torch.float16
	67	+ required_len = reduce(operator.mul, shape, 1)
	68	+ if GlobalMemoryBuffer.buffers_length[index] < required_len:
	69	+ GlobalMemoryBuffer.buffers[index] = torch.empty(
	70	+ required_len, dtype=data_type, requires_grad=False, device=torch.cuda.current_device()
	71	+ )
	72	+ GlobalMemoryBuffer.buffers_length[index] = required_len
	73	+ return GlobalMemoryBuffer.buffers[index][0:required_len].view(*shape).uniform_()
	74	+
	75	+
	76	+class SingletonType(type):
	77	+ single_lock = threading.RLock()
	78	+
	79	+ def __call__(cls, args, *kwargs):
	80	+ with SingletonType.single_lock:
	81	+ if not hasattr(cls, "_instance"):
	82	+ cls._instance = super(SingletonType, cls).__call__(args, *kwargs)
	83	+ return cls._instance
	84	+
	85	+
	86	+class SampleCache:
	87	+ def __init__(self):
	88	+ self.MatMul = {}
	89	+ self.RmsNorm = {}
	90	+ self.RmsNormGrad = {}
	91	+ self.BatchMatMul = {}
	92	+ self.Add = {}
	93	+ self.LayerNorm = {}
	94	+ self.LayerNormGrad = {}
	95	+ self.ScaledMaskedSoftmax = {}
	96	+ self.ScaledMaskedSoftmaxGrad = {}
	97	+ self.FastGeluGrad = {}
	98	+ self.FastGelu = {}
	99	+ self.Mul = {}
	100	+ self.Softmax = {}
	101	+ self.SoftmaxGrad = {}
	102	+ self.FlashAttentionScore = {}
	103	+ self.FlashAttentionScoreGrad = {}
	104	+
	105	+ def clear_cache(self):
	106	+ for attr in self.__dict__:
	107	+ setattr(self, attr, {})
	108	+
	109	+
	110	+class ModelManager:
	111	+ def __init__(self, npu_type='910B'):
	112	+ self.models = {}
	113	+ self.npu_type = npu_type
	114	+
	115	+ def cache_model(self, model, op):
	116	+ self.models[op] = model
	117	+
	118	+ def get_cached_model(self, model_name: str):
	119	+ return self.models.get(model_name, None)
	120	+
	121	+ def load_model(self, model, op, model_dir):
	122	+ if not os.path.exists(model_dir):
	123	+ raise FileNotFoundError(f"Can't find '{model_dir}'.")
	A A aijgnem12024年7月30日应该是 `raise FileNotFoundError`
	124	+ path = os.path.join(model_dir, f"{op}_{self.npu_type}.pth")
	125	+ weight = torch.load(path)
	126	+ model.set_model_info(weight.popitem()[1])
	127	+ model.load_state_dict(weight)
	128	+ # if use model to predict,need to set training=False,otherwise require inputs dims==model_train_inputs dims
	129	+ # during fit,after clear model cache(self.train()),training's value will be reset True
	130	+ model.training = False
	131	+ self.models[op] = model
	132	+
	133	+ def save_model(self, model, op, model_dir):
	134	+ if not os.path.exists(model_dir):
	135	+ os.makedirs(model_dir, exist_ok=False)
	A A aijgnem12024年7月30日已有`exist_ok=True`，不需要`if not os.path.exists(model_dir):`
	136	+ weight = model.state_dict()
	137	+ weight['model_info'] = model.get_model_info()
	138	+ torch.save(weight, f'{model_dir}/{op}_{self.npu_type}.pth')
	139	+
	140	+ def save_models(self, model_dir):
	141	+ for op, op_model in self.models.items():
	142	+ self.save_model(op_model, op, model_dir)
	143	+
	144	+
	145	+class OperateProfileCache(metaclass=SingletonType):
	146	+ def __init__(self):
	147	+ self.data_frame = pd.DataFrame(
	148	+ columns=[KeyField.OpType, KeyField.InputShapes, KeyField.OutputShapes, KeyField.FwdTime, KeyField.BwdTime]
	149	+ )
	150	+
	151	+ def record(self, op_type: str, input_shapes: list, output_shapes: list, fwd_time: float, bwd_time: float):
	152	+ _, _, exist = self.find(op_type, input_shapes)
	153	+ if not exist:
	154	+ input_shapes_str = OperateProfileCache.shapes_to_str(input_shapes)
	155	+ output_shape_str = OperateProfileCache.shapes_to_str(output_shapes)
	156	+ self.data_frame.loc[len(self.data_frame.index)] = [
	157	+ op_type, input_shapes_str, output_shape_str, fwd_time, bwd_time
	158	+ ]
	159	+
	160	+ def find(self, op_type: str, input_shapes: list):
	161	+ input_shapes_str = OperateProfileCache.shapes_to_str(input_shapes)
	162	+ data = self.data_frame[
	163	+ (self.data_frame[KeyField.OpType] == op_type) &
	164	+ (self.data_frame[KeyField.InputShapes] == input_shapes_str)
	165	+ ]
	166	+ fwd_time = data[KeyField.FwdTime].mean()
	167	+ bwd_time = data[KeyField.BwdTime].mean()
	168	+ from_cache = False if np.isnan(fwd_time) and np.isnan(bwd_time) else True
	169	+ return fwd_time, bwd_time, from_cache
	170	+
	171	+ @staticmethod
	172	+ def shapes_to_str(shapes):
	173	+ result = ''
	174	+ index = 0
	175	+ for shape in shapes:
	176	+ result += ','.join(map(lambda x: str(x), shape)) if isinstance(shape, list) else str(shape)
	177	+ if index < len(shapes) - 1:
	178	+ result += ';' if isinstance(shape, list) else ','
	179	+ index += 1
	180	+ result = '"' + result
	181	+ result = result + '"'
	182	+ return result
	183	+
	184	+
	185	+def get_cache_path():
	186	+ global SEARCH_CACHE_PATH
	187	+ if SEARCH_CACHE_PATH is None:
	188	+ SEARCH_CACHE_PATH = os.getcwd() + os.sep + 'autoparallel_temp_cache' + os.sep
	189	+ try:
	190	+ os.makedirs(SEARCH_CACHE_PATH, exist_ok=True)
	191	+ print(f"Create cache: {SEARCH_CACHE_PATH}")
	192	+ except Exception:
	193	+ print(f'Create cache directory failed')
	194	+ SEARCH_CACHE_PATH = os.getcwd()
	195	+ return SEARCH_CACHE_PATH
	196	+
	197	+
	198	+def analyse_module_profile(profile_file, key):
	199	+ if key not in ('step_time', 'transformer_act_mem'):
	200	+ raise AssertionError('key[{}] error'.format(key))
	201	+
	202	+ if not os.path.exists(path=profile_file):
	203	+ return float('inf')
	204	+
	205	+ with open(profile_file, 'r') as file:
	206	+ try:
	207	+ content = file.read()
	208	+ content = json.loads(content)
	209	+ return float(content.get(key))
	210	+ except Exception:
	211	+ return float('inf')
	A A aijgnem12024年7月30日函数`analyse_memory`和`analyse_module_profile`可以考虑合并
	212	+
	213	+
	214	+def set_kv_store(args):
	215	+ global KV_STORE
	216	+ if args.node_rank == 0:
	217	+ KV_STORE = torch.distributed.TCPStore(
	218	+ host_name=args.master_addr,
	219	+ port=int(args.master_port) + 2,
	220	+ world_size=args.nnodes,
	221	+ is_master=True,
	222	+ timeout=datetime.timedelta(seconds=30)
	223	+ )
	224	+ else:
	225	+ KV_STORE = torch.distributed.TCPStore(
	226	+ host_name=args.master_addr,
	227	+ port=int(args.master_port) + 2,
	228	+ world_size=args.nnodes,
	229	+ is_master=False
	230	+ )
	231	+
	232	+
	233	+def get_kv_store():
	234	+ global KV_STORE
	235	+ if KV_STORE is None:
	236	+ raise AssertionError('KV_STORE must be initialized')
	237	+ return KV_STORE
	238	+
	239	+
	240	+# init SingletonType class
	241	+model_manager = ModelManager()
	242	+sample_cache = SampleCache()
	243	+operator_cache = OperateProfileCache()

Amindspeed/core/auto_parallel/auto_parallel_apply.py+156-0

@@ -0,0 +1,156 @@
	1	+# coding=utf-8
	2	+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+import json
	17	+import time
	18	+import math
	19	+
	20	+import torch
	21	+from megatron.training.global_vars import get_args
	22	+
	23	+from mindspeed.core.auto_parallel import set_kv_store
	24	+from mindspeed.core.auto_parallel.auto_parallel_optimizer import SearchByGreyBox
	25	+from mindspeed.core.auto_parallel.auto_parallel_memory import MemoryCostModel
	26	+from mindspeed.core.auto_parallel.auto_parallel_profiling import (
	27	+ DistributedMemoryProfiler,
	28	+ DistributedOperateProfiler,
	29	+ DistributedPerformanceProfiler
	30	+)
	31	+
	32	+
	33	+def filter_unvalid_configs(search_spaces):
	34	+ memory_model = MemoryCostModel()
	35	+ fitting_configs = memory_model.get_fitting_configurations(search_spaces)
	36	+ for config in fitting_configs:
	37	+ mem = DistributedMemoryProfiler().launch(config)
	38	+ if not math.isinf(mem):
	39	+ memory_model.profiled_configs.append(config)
	40	+ memory_model.profiled_configs_memory.append(mem)
	41	+
	42	+ print(f"profiled_configs: {memory_model.profiled_configs}")
	43	+ print(f"profiled_configs_mem: {memory_model.profiled_configs_memory}")
	44	+
	45	+ memory_model.fit_model()
	46	+ valid_configs, valid_configs_memory = [], []
	47	+ for config in search_spaces:
	48	+ cost_memory = memory_model.get_peak_memory(config)
	B B bingobb2024年7月27日这里-1我看是mbs，为啥大于16就跳过了呢
	49	+ if not memory_model.is_oom(cost_memory):
	50	+ valid_configs.append(config)
	51	+ valid_configs_memory.append(cost_memory)
	52	+ return valid_configs
	53	+
	54	+
	55	+def build_initial_spaces(args):
	56	+ world_size = args.nproc_per_node * args.nnodes
	57	+ device_count = args.nproc_per_node
	58	+
	59	+ solutions = []
	60	+ for pp in range(1, world_size + 1):
	61	+ if world_size % pp != 0 or args.num_layers % pp != 0:
	62	+ continue
	63	+
	64	+ for i in range(device_count):
	65	+ tp = 2 ** i
	66	+ if tp > device_count or tp > (world_size // pp):
	67	+ break
	68	+ if (args.num_query_groups > 1 and args.num_query_groups % tp != 0) \
	69	+ or (args.num_attention_heads % tp != 0):
	70	+ break
	71	+
	72	+ max_cp_size = world_size // (pp * tp)
	73	+ for cp_size in range(1, max_cp_size + 1):
	74	+ if world_size % (pp * tp * cp_size) != 0 or \
	75	+ args.global_batch_size % (world_size // (pp * tp * cp_size)) != 0:
	B B bingobb2024年7月27日 context_parallel_size 这个注释是不是可以去掉，还有上面的tp和pp。这里应该不难理解
	76	+ continue
	77	+
	78	+ for up in range(1, cp_size + 1):
	79	+ if cp_size % up != 0:
	80	+ continue
	81	+ cp = cp_size // up
	82	+ head, remainder = divmod(args.num_attention_heads, up * tp)
	83	+ if (head < 1 or remainder != 0) or (args.seq_length % (2 * cp) != 0):
	84	+ continue
	85	+
	86	+ dp = world_size // (pp * tp * cp_size)
	87	+ dp_group_batch_size = args.global_batch_size // dp
	88	+ for num_mb in range(1, dp_group_batch_size + 1):
	89	+ if dp_group_batch_size % num_mb != 0:
	90	+ continue
	91	+ mbs = dp_group_batch_size // num_mb
	92	+ solutions.append([pp, tp, dp, cp, up, mbs])
	93	+ return solutions
	94	+
	95	+
	96	+def monitor_train_task():
	97	+ while True:
	98	+ message = torch.tensor([0 for _ in range(7)], dtype=torch.int)
	99	+ torch.distributed.broadcast(message, 0)
	100	+ task_type = message[-1].item()
	101	+ config = [m.item() for m in message[:-1]]
	102	+ if task_type == -1:
	103	+ break
	104	+ elif task_type == 0:
	105	+ DistributedMemoryProfiler().launch(config)
	106	+ elif task_type == 1:
	107	+ DistributedOperateProfiler().launch(config)
	108	+ elif task_type == 2:
	109	+ DistributedPerformanceProfiler().launch(config)
	110	+
	111	+
	112	+def export_results(config):
	113	+ results = {}
	114	+ results['optimal_parallel_strategy'] = {}
	115	+ results['optimal_parallel_strategy']['pipeline-model-parallel-size'] = config[0]
	116	+ results['optimal_parallel_strategy']['tensor-model-parallel-size'] = config[1]
	117	+ results['optimal_parallel_strategy']['data-parallel-size'] = config[2]
	118	+ results['optimal_parallel_strategy']['micro-batch-size'] = config[-1]
	119	+ if config[3] > 1 and config[4] > 1:
	120	+ results['optimal_parallel_strategy']['context-parallel-algo'] = 'hybrid_cp_algo'
	121	+ results['optimal_parallel_strategy']['context-parallel-size'] = config[3] * config[4]
	122	+ results['optimal_parallel_strategy']['ulysses-degree-in-cp'] = config[4]
	123	+ elif config[3] > 1 and config[4] == 1:
	124	+ results['optimal_parallel_strategy']['context-parallel-algo'] = 'megatron_cp_algo'
	125	+ results['optimal_parallel_strategy']['context-parallel-size'] = config[3]
	126	+ elif config[3] == 1 and config[4] > 1:
	127	+ results['optimal_parallel_strategy']['context-parallel-algo'] = 'ulysses_cp_algo'
	128	+ results['optimal_parallel_strategy']['context-parallel-size'] = config[4]
	129	+ return json.dumps(results)
	130	+
	131	+
	132	+def search_optimal_configuration(args):
	133	+ set_kv_store(args)
	134	+
	135	+ init_method = 'tcp://{}:{}'.format(args.master_addr, int(args.master_port) + 1)
	136	+ torch.distributed.init_process_group(
	137	+ backend=torch.distributed.Backend.GLOO,
	138	+ init_method=init_method,
	139	+ rank=args.node_rank,
	140	+ world_size=args.nnodes
	141	+ )
	142	+
	143	+ if args.node_rank == 0:
	144	+ search_space = build_initial_spaces(args)
	145	+ search_space = filter_unvalid_configs(search_space)
	146	+ print(f"filter search_space: {len(search_space)}")
	147	+ print("\n".join(str(item) for item in search_space), flush=True)
	148	+
	149	+ start_time = time.time()
	150	+ config, _ = SearchByGreyBox().search(get_args(), search_space)
	151	+ torch.distributed.broadcast(torch.tensor([-1 for _ in range(7)], dtype=torch.int), 0)
	152	+
	153	+ results = export_results(config)
	154	+ print(f"find optimal configuration: {results}, cost_time: {time.time() - start_time}")
	155	+ else:
	156	+ monitor_train_task()

Amindspeed/core/auto_parallel/auto_parallel_memory.py+168-0

@@ -0,0 +1,168 @@
	1	+# coding=utf-8
	2	+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+from itertools import product
	17	+
	18	+import numpy as np
	19	+import torch
	20	+from megatron.training.global_vars import get_args
	21	+
	22	+from mindspeed.core.auto_parallel import SingletonType
	23	+
	24	+
	25	+class MemoryCostModel(metaclass=SingletonType):
	26	+ def __init__(self):
	27	+ args = get_args()
	28	+ self.num_layers = args.num_layers
	29	+ self.num_attn_heads = args.num_attention_heads
	30	+ self.hidden_size = args.hidden_size
	31	+ self.seq_length = args.seq_length
	32	+ self.ffn_hidden_size = args.ffn_hidden_size
	33	+ if not self.ffn_hidden_size:
	34	+ self.ffn_hidden_size = 4 * self.hidden_size
	35	+
	36	+ self.model = None
	37	+ self.profiled_configs = []
	38	+ self.profiled_configs_memory = []
	39	+ self.max_available_memory = None
	40	+
	41	+ @staticmethod
	42	+ def cal_coeff(config):
	43	+ _, tp, _, cp, up, b = config
	44	+ coeff = [
	45	+ 1,
	46	+ b * (1 / tp) * (1 / cp) * (1 / up),
	47	+ b * (1 / tp) * (1 / cp) * (1 / cp) * (1 / up),
	48	+ b * (1 / cp) * (1 / up)
	49	+ ]
	50	+ return np.array(coeff)
	51	+
	52	+ @staticmethod
	53	+ def cal_coeff_matrix(configs):
	54	+ coeff_matrix = []
	55	+ for config in configs:
	56	+ _, tp, _, cp, up, b = config
	57	+ coeff_matrix.append([
	58	+ 1,
	59	+ b * (1 / tp) * (1 / cp) * (1 / up),
	60	+ b * (1 / tp) * (1 / cp) * (1 / cp) * (1 / up),
	61	+ b * (1 / cp) * (1 / up)
	62	+ ])
	63	+ return np.array(coeff_matrix)
	64	+
	65	+ def is_oom(self, cost_memory):
	66	+ if self.max_available_memory is None:
	67	+ properties = torch.npu.get_device_properties(0)
	68	+ self.max_available_memory = properties.total_memory / (1024 ** 3)
	69	+ # 以单卡最大可用内存的1.2倍作为OOM阈值
	70	+ return cost_memory > (self.max_available_memory * 1.2)
	A A aijgnem12024年7月29日 `1.2`魔鬼数字。 “魔鬼数字”通常指的是代码中出现的未命名的常量或魔法数字（magic numbers），这些数字的意义不明，使得代码难以理解和维护。优化这些魔鬼数字的方法主要是通过命名这些常量，使它们具有明确的意义。下面是一些具体的步骤和建议： 1.定义常量：将这些数字定义为常量，并给予有意义的名字。使用大写字母表示常量是一种常见的约定。 2.封装到配置文件：如果这些数字代表的是配置参数，可以将它们放在配置文件中，如JSON、YAML或属性文件等。通过读取配置文件来获取这些值。 3.文档说明：在代码中添加注释来解释这些数字的意义。如果数字的含义非常重要但又不适合定义为常量，至少要在注释中说明。
	71	+
	72	+ def get_fitting_configurations(self, search_spaces):
	73	+ search_spaces_matrix = np.array(search_spaces)
	74	+ temp_search_spaces = [config for config in search_spaces if config[-1] < 8]
	B B bingobb2024年7月27日这里小于8的设置是为啥呢
	75	+
	76	+ tp_group = []
	77	+ max_tp = search_spaces_matrix[:, 1].max()
	78	+ for config in temp_search_spaces:
	79	+ _, tp, _, cp, up, _ = config
	80	+ if cp == 1 and up == 1 and tp == max_tp:
	81	+ tp_group.append(config)
	82	+
	83	+ cp_group = []
	84	+ min_cp = search_spaces_matrix[:, 3].min()
	85	+ for config in temp_search_spaces:
	86	+ pp, tp, _, cp, up, _ = config
	87	+ if tp > 1 or up > 1:
	88	+ continue
	89	+ if pp > 1 and cp > min_cp:
	90	+ cp_group.append(config)
	91	+
	92	+ up_group = []
	93	+ min_up = search_spaces_matrix[:, 4].min()
	94	+ for config in temp_search_spaces:
	95	+ pp, tp, _, cp, up, _ = config
	96	+ if tp > 1 or cp > 1:
	97	+ continue
	98	+ if pp > 1 and up > min_up:
	99	+ up_group.append(config)
	100	+
	101	+ cp_up_group = []
	102	+ for config in temp_search_spaces:
	103	+ _, tp, _, cp, up, _ = config
	104	+ if tp == 1 and cp > 1 and up > 1:
	105	+ cp_up_group.append(config)
	106	+
	107	+ tp_cp_up_group = []
	108	+ for config in temp_search_spaces:
	109	+ _, tp, _, cp, up, _ = config
	110	+ if tp > 1 and cp > 1 and up > 1:
	111	+ tp_cp_up_group.append(config)
	112	+
	113	+ product_iter = product(*[tp_group, cp_group, up_group, cp_up_group, tp_cp_up_group])
	114	+ fitting_group, cur_condition_number = None, float('inf')
	115	+
	116	+ for group in product_iter:
	117	+ # 条件数小于100的矩阵的逆矩阵数值更稳定，拟合效果更好
	118	+ if cur_condition_number < 100:
	119	+ break
	120	+
	121	+ empty_set = set([row[-1] for row in group])
	122	+ if len(empty_set) < 2:
	123	+ continue
	124	+
	125	+ coeff_matrix = MemoryCostModel.cal_coeff_matrix(group)
	126	+ coeff_matrix = coeff_matrix.transpose() @ coeff_matrix
	127	+ if np.linalg.matrix_rank(coeff_matrix) == coeff_matrix.shape[0]:
	128	+ con_num = np.linalg.cond(coeff_matrix)
	129	+ if con_num < cur_condition_number:
	130	+ fitting_group = group
	131	+ cur_condition_number = con_num
	132	+
	133	+ print(f"fitting_group: {fitting_group} condition_number: {cur_condition_number}", flush=True)
	134	+ return fitting_group
	135	+
	136	+
	137	+ def fit_model(self):
	138	+ coeff_matrix = MemoryCostModel.cal_coeff_matrix(self.profiled_configs)
	139	+ profiled_configs_memory = np.array(self.profiled_configs_memory)
	140	+ self.model = np.linalg.inv(coeff_matrix.transpose() @ coeff_matrix) \
	141	+ @ coeff_matrix.transpose() \
	142	+ @ profiled_configs_memory
	143	+
	144	+ def predict(self, config):
	145	+ config_matrix = MemoryCostModel.cal_coeff(config)
	146	+ pred_memory = config_matrix @ self.model
	147	+ return pred_memory
	148	+
	149	+ def get_peak_memory(self, config):
	150	+ args = get_args()
	151	+ pp, tp, _ = config[0], config[1], config[-1]
	152	+ hidden_size = self.hidden_size
	153	+ ffn_hidden_size = self.ffn_hidden_size
	154	+ if args.swiglu:
	155	+ ffn_hidden_size *= 2
	156	+ transformer_params_count = (4 * hidden_size * hidden_size + 2 * hidden_size * ffn_hidden_size) / tp
	157	+ total_params_count = transformer_params_count * (self.num_layers // pp)
	158	+
	159	+ mem_para = 2 * total_params_count
	160	+ mem_grad = 2 * total_params_count
	161	+ mem_optimizer = 12 * total_params_count if args.reuse_fp32_param else 16 * total_params_count
	162	+ mem_activation_layer = abs(self.predict(config)) * (1024 ** 3)
	163	+ mem_activation_batch = mem_activation_layer * (self.num_layers // pp)
	164	+ mem_activation = mem_activation_batch * pp
	165	+ m1 = mem_para + mem_optimizer + mem_activation
	166	+ m2 = mem_para + mem_optimizer + mem_activation + mem_grad - mem_activation_batch
	167	+ peak_memory = max(m1, m2)
	168	+ return peak_memory / (1024 ** 3) + 4

Amindspeed/core/auto_parallel/auto_parallel_model.py+462-0

@@ -0,0 +1,462 @@
	1	+# coding=utf-8
	2	+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+import time
	17	+import math
	18	+from functools import reduce
	19	+
	20	+import numpy as np
	21	+import torch
	22	+import torch_npu
	23	+from megatron.training.global_vars import get_args
	24	+
	25	+from mindspeed.core.auto_parallel import (
	26	+ ITERATION_LOOP_TIME,
	27	+ BAND_WIDTH_UNIDIRECTIONAL,
	28	+ operator_cache,
	29	+ GlobalMemoryBuffer
	30	+)
	31	+from mindspeed.core.auto_parallel.auto_parallel_rectify import Sampler
	32	+from mindspeed.core.auto_parallel.auto_parallel_profiling import CommProfiling
	33	+from mindspeed.model.transformer import (
	34	+ get_attention_mask,
	35	+ generate_attention_mask
	36	+)
	37	+
	38	+
	39	+class Linear(torch.nn.Module):
	40	+ def __init__(self):
	41	+ super(Linear, self).__init__()
	42	+
	43	+ def forward(self, inputs):
	44	+ x, y = inputs
	45	+ return torch.matmul(x, y.t())
	46	+
	47	+
	48	+class LayerNormV3(torch.nn.Module):
	49	+ def __init__(self, hidden_size, eps=1e-5):
	50	+ super(LayerNormV3, self).__init__()
	51	+ self.layer_norm = torch.nn.LayerNorm(normalized_shape=hidden_size, eps=eps)
	52	+
	53	+ def forward(self, x):
	54	+ return self.layer_norm(*x)
	55	+
	56	+
	57	+class FusedRmsNorm(torch.nn.Module):
	58	+ def __init__(self, hidden_size, eps=1e-6) -> None:
	59	+ super().__init__()
	60	+ self.weight = torch.nn.Parameter(torch.ones(hidden_size, dtype=torch.float16)).npu()
	61	+ self.eps = eps
	62	+
	63	+ def forward(self, x):
	64	+ return torch_npu.npu_rms_norm(x[0], self.weight, epsilon=self.eps)[0]
	65	+
	66	+
	67	+class BatchMatMul(torch.nn.Module):
	68	+ def __init__(self):
	69	+ super(BatchMatMul, self).__init__()
	70	+
	71	+ def forward(self, inputs):
	72	+ x, y = inputs
	73	+ return torch.bmm(x, y)
	74	+
	75	+
	76	+class FlashAttention(torch.nn.Module):
	77	+ def __init__(self, head_dim):
	78	+ super().__init__()
	79	+ self.head_dim = head_dim
	80	+ self.scale = 1.0 / math.sqrt(self.head_dim)
	81	+ self.pre_tockens = 65536
	82	+ self.next_tockens = 0
	83	+
	84	+ generate_attention_mask()
	85	+ self.attention_mask = get_attention_mask()
	86	+
	87	+ def forward(self, x):
	88	+ q, k, v = x
	89	+ seq_length, _, hd = q.shape[0], q.shape[1], q.shape[2]
	90	+ head_num = hd // self.head_dim
	91	+ output = torch_npu.npu_fusion_attention(
	92	+ q, k, v, head_num, 'SBH',
	93	+ pse=None,
	94	+ padding_mask=None,
	95	+ atten_mask=self.attention_mask,
	96	+ scale=self.scale,
	97	+ pre_tockens=self.pre_tockens,
	98	+ next_tockens=self.next_tockens,
	99	+ keep_prob=1.0,
	100	+ inner_precise=0,
	101	+ sparse_mode=get_args().sparse_mode
	102	+ )[0]
	103	+ return output
	104	+
	105	+
	106	+class TransformerBlock:
	107	+ def __init__(self):
	108	+ self.number_sample = 100
	109	+ self.noise_model = OperatorNoiseSampler(self.number_sample)
	110	+
	111	+ def norm(self):
	112	+ args = get_args()
	113	+ tp = args.tensor_model_parallel_size
	114	+ cp = args.context_parallel_size // args.ulysses_degree_in_cp
	115	+ up = args.ulysses_degree_in_cp
	116	+ input_shape = [args.seq_length // cp // tp // up, args.micro_batch_size, args.hidden_size]
	117	+ if args.normalization == 'RMSNorm':
	118	+ ftime, btime = self.noise_model.fused_rms_norm(input_shape, input_shape, args.hidden_size)
	119	+ else:
	120	+ ftime, btime = self.noise_model.layernorm(input_shape, input_shape, args.hidden_size)
	121	+ return ftime, btime
	122	+
	123	+ def self_attention_with_fa(self):
	124	+ args = get_args()
	125	+ tp = args.tensor_model_parallel_size
	126	+ cp = args.context_parallel_size // args.ulysses_degree_in_cp
	127	+ up = args.ulysses_degree_in_cp
	128	+ ftime, btime = self.noise_model.flash_attention(
	129	+ [args.seq_length // cp, args.micro_batch_size, args.hidden_size // tp // up],
	130	+ [args.seq_length // cp, args.micro_batch_size, args.hidden_size // tp // up],
	131	+ [args.seq_length // cp, args.micro_batch_size, args.hidden_size // tp // up],
	132	+ [args.seq_length // cp, args.micro_batch_size, args.hidden_size // tp // up],
	133	+ args.hidden_size // args.num_attention_heads,
	134	+ )
	135	+ return ftime, btime
	136	+
	137	+ def get_block_time(self):
	BB B bingobb2024年7月27日是不是可以根据模型结构，做一些分类，，增加一个分类的入参 B bingobb2024年7月27日是不是可以根据模型结构，做一些分类，，增加一个分类的入参类似PaLM这种
	138	+ args = get_args()
	139	+ s = args.seq_length
	140	+ a = args.num_attention_heads
	141	+ h = args.hidden_size
	142	+ ffn = args.ffn_hidden_size if args.ffn_hidden_size is not None else 4 * args.hidden_size
	143	+ d = args.hidden_size // args.num_attention_heads
	144	+ b = args.micro_batch_size
	145	+ tp = args.tensor_model_parallel_size
	146	+ cp = args.context_parallel_size // args.ulysses_degree_in_cp
	147	+ up = args.ulysses_degree_in_cp
	148	+
	149	+ fwd_time = np.array([0 for _ in range(self.number_sample)]).astype(np.float64)
	150	+ bwd_time = np.array([0 for _ in range(self.number_sample)]).astype(np.float64)
	151	+
	152	+ ftime, btime = self.norm()
	153	+ fwd_time += ftime
	154	+ bwd_time += btime
	155	+
	156	+ all_gather_time = CommProfiling.get_comm_time([s // cp // up // tp, b, h], tp, 'all_gather')
	157	+ reduce_scatter_time = CommProfiling.get_comm_time([s // cp // up, b, h], tp, 'reduce_scatter')
	158	+ fwd_time += all_gather_time
	159	+ bwd_time += reduce_scatter_time
	160	+
	161	+ ftime, btime = self.noise_model.matmul(
	162	+ [s // cp // up * b, h],
	163	+ [3 * h // tp, h],
	164	+ [s // cp // up * b, 3 * h // tp]
	165	+ )
	166	+ fwd_time += ftime
	167	+ bwd_time += btime
	168	+
	169	+ if not args.use_flash_attn:
	170	+ raise AssertionError('the auto-parallel only support FA')
	171	+ else:
	172	+ alltoall_time = CommProfiling.get_comm_time([s // cp // up, b, a // tp, d], up, 'alltoall')
	173	+ fwd_time += (3 * alltoall_time)
	174	+ bwd_time += (3 * alltoall_time)
	175	+
	176	+ send_recv_time = CommProfiling.get_send_recv_time([2, 2, s // cp // 2, b, a // tp // up * d])
	177	+ ftime, btime = self.self_attention_with_fa()
	178	+ for _ in range(cp - 1):
	179	+ fwd_time += max([ftime.max(), send_recv_time])
	180	+ bwd_time += max([btime.max(), send_recv_time])
	181	+ fwd_time += ftime
	182	+ bwd_time += btime
	183	+
	184	+ alltoall_time = CommProfiling.get_comm_time([s // cp, b, a // tp // up, d], up, 'alltoall')
	185	+ fwd_time += alltoall_time
	186	+ bwd_time += alltoall_time
	187	+
	188	+ ftime, btime = self.noise_model.matmul([s // cp // up * b, h // tp], [h, h // tp], [s // cp // up * b, h])
	189	+ fwd_time += ftime
	190	+ bwd_time += btime
	191	+
	192	+ reduce_scatter_time = CommProfiling.get_comm_time([s // cp // up, b, h], tp, 'reduce_scatter')
	193	+ all_gather_time = CommProfiling.get_comm_time([s // cp // up // tp, b, h], tp, 'all_gather')
	194	+ fwd_time += reduce_scatter_time
	195	+ bwd_time += all_gather_time
	196	+
	197	+ ftime, btime = self.norm()
	198	+ fwd_time += ftime
	199	+ bwd_time += btime
	200	+
	201	+ all_gather_time = CommProfiling.get_comm_time([s // cp // up // tp, b, h], tp, 'all_gather')
	B B bingobb2024年7月27日如果是PaLM那种attention和mlp并行的模型结构，是不是就不准了，因为通信变少了
	202	+ reduce_scatter_time = CommProfiling.get_comm_time([s // cp // up, b, h], tp, 'reduce_scatter')
	203	+ fwd_time += all_gather_time
	204	+ bwd_time += reduce_scatter_time
	205	+
	206	+ ftime, btime = self.noise_model.matmul([s // cp // up * b, h], [ffn // tp, h], [s // cp // up * b, ffn // tp])
	207	+ fwd_time += ftime
	208	+ bwd_time += btime
	209	+
	210	+ # 4h->h
	211	+ ftime, btime = self.noise_model.matmul([s // cp // up * b, ffn // tp], [h, ffn // tp], [s // cp // up * b, h])
	212	+ fwd_time += ftime
	213	+ bwd_time += btime
	214	+
	215	+ reduce_scatter_time = CommProfiling.get_comm_time([s // cp // up, b, h], tp, 'reduce_scatter')
	216	+ all_gather_time = CommProfiling.get_comm_time([s // cp // up // tp, b, h], tp, 'all_gather')
	217	+ fwd_time += reduce_scatter_time
	218	+ bwd_time += all_gather_time
	219	+
	220	+ return fwd_time, bwd_time
	221	+
	222	+
	223	+class OperatorNoiseSampler:
	224	+ def __init__(self, num_sample=100):
	225	+ self.sampling = Sampler(num_sample=num_sample)
	226	+
	227	+ @staticmethod
	228	+ def measure_matmul_time(left_shape, left_transpose, right_shape, right_transpose):
	229	+ left_matrix = GlobalMemoryBuffer.get_tensor(left_shape, 0)
	230	+ left_matrix = left_matrix if not left_transpose else left_matrix.t()
	231	+ right_matrix = GlobalMemoryBuffer.get_tensor(right_shape, 1)
	232	+ right_matrix = right_matrix if not right_transpose else right_matrix.t()
	233	+
	234	+ for _ in range(ITERATION_LOOP_TIME):
	235	+ torch.matmul(left_matrix, right_matrix)
	236	+
	237	+ torch.npu.synchronize()
	238	+ start_time = time.time()
	239	+ for _ in range(ITERATION_LOOP_TIME):
	240	+ torch.matmul(left_matrix, right_matrix)
	241	+ torch.npu.synchronize()
	242	+ return (time.time() - start_time) * 1e6 / ITERATION_LOOP_TIME
	243	+
	244	+ @staticmethod
	245	+ def measure_batchmatmul_time(left_shape, left_transpose, right_shape, right_transpose):
	246	+ left_matrix = GlobalMemoryBuffer.get_tensor(left_shape, 0)
	247	+ left_matrix = left_matrix if not left_transpose else left_matrix.permute(0, 2, 1)
	248	+ right_matrix = GlobalMemoryBuffer.get_tensor(right_shape, 0)
	249	+ right_matrix = right_matrix if not right_transpose else right_matrix.permute(0, 2, 1)
	250	+
	251	+ for _ in range(ITERATION_LOOP_TIME):
	252	+ torch.bmm(left_matrix, right_matrix)
	253	+
	254	+ torch.npu.synchronize()
	255	+ start_time = time.time()
	256	+ for _ in range(ITERATION_LOOP_TIME):
	257	+ torch.bmm(left_matrix, right_matrix)
	258	+ torch.npu.synchronize()
	259	+ return (time.time() - start_time) * 1e6 / ITERATION_LOOP_TIME
	260	+
	261	+ def matmul(self, input_shape1, input_shape2, output_shape):
	262	+ ftime, _, from_cache = operator_cache.find('MatMul', [input_shape1, input_shape2])
	263	+ if not from_cache:
	264	+ ftime = self.measure_matmul_time(input_shape1, False, input_shape2, True)
	265	+ ftime_uncertainty = self.sampling.run('MatMul', ftime, output_shape, input_shape1, input_shape2)
	266	+ operator_cache.record('MatMul', [input_shape1, input_shape2], output_shape, ftime, 0)
	267	+
	268	+ btime1, _, from_cache = operator_cache.find('MatMul', [output_shape, input_shape2])
	269	+ if not from_cache:
	270	+ btime1 = self.measure_matmul_time(output_shape, False, input_shape2, False)
	271	+ btime1_uncertainty = self.sampling.run('MatMul', btime1, input_shape1, output_shape, input_shape2)
	272	+ operator_cache.record('MatMul', [output_shape, input_shape2], input_shape1, btime1, 0)
	273	+
	274	+ btime2, _, from_cache = operator_cache.find('MatMul', [output_shape, input_shape1])
	275	+ if not from_cache:
	276	+ btime2 = self.measure_matmul_time(output_shape, True, input_shape1, False)
	277	+ btime2_uncertainty = self.sampling.run('MatMul', btime2, input_shape2, output_shape, input_shape1)
	278	+ operator_cache.record('MatMul', [output_shape, input_shape1], input_shape2, btime2, 0)
	279	+ return ftime_uncertainty, btime1_uncertainty + btime2_uncertainty
	280	+
	281	+ def batch_matmul(self, input_shape1, input_shape2, output_shape):
	282	+ ftime, _, from_cache = operator_cache.find('BatchMatMul', [input_shape1, input_shape2])
	283	+ if not from_cache:
	284	+ ftime = self.measure_batchmatmul_time(input_shape1, False, input_shape2, False)
	285	+ ftime_uncertainty = self.sampling.run('BatchMatMul', ftime, output_shape, input_shape1, input_shape2)
	286	+ operator_cache.record('BatchMatMul', [input_shape1, input_shape2], output_shape, ftime, 0)
	287	+
	288	+ btime1, _, from_cache = operator_cache.find('BatchMatMul', [input_shape1, output_shape])
	289	+ if not from_cache:
	290	+ btime1 = self.measure_batchmatmul_time(input_shape1, True, output_shape, False)
	291	+ btime1_uncertainty = self.sampling.run('BatchMatMul', btime1, input_shape2, input_shape1, output_shape)
	292	+ operator_cache.record('BatchMatMul', [input_shape1, output_shape], input_shape2, btime1, 0)
	293	+
	294	+ btime2, _, from_cache = operator_cache.find('BatchMatMul', [output_shape, input_shape2])
	295	+ if not from_cache:
	296	+ btime2 = self.measure_batchmatmul_time(output_shape, False, input_shape2, True)
	297	+ btime2_uncertainty = self.sampling.run('BatchMatMul', btime2, input_shape1, output_shape, input_shape2)
	298	+ operator_cache.record('BatchMatMul', [output_shape, input_shape2], input_shape1, btime2, 0)
	299	+ return ftime_uncertainty, btime1_uncertainty + btime2_uncertainty
	300	+
	301	+ def layernorm(self, input_shape, output_shape, hidden_size, eps=1e-5):
	302	+ layernorm = LayerNormV3(hidden_size, eps)
	303	+ ftime, btime, from_cache = operator_cache.find('LayerNormV3', input_shape)
	304	+ if not from_cache:
	305	+ ftime, btime = TimeCostModel.profile(layernorm, [input_shape])
	306	+ ftime_uncertainty = self.sampling.run('LayerNormV3', ftime, output_shape, input_shape)
	307	+ btime_uncertainty = self.sampling.run('LayerNormGrad', btime, input_shape, output_shape)
	308	+ operator_cache.record('LayerNormV3', input_shape, output_shape, ftime, btime)
	309	+ return ftime_uncertainty, btime_uncertainty
	310	+
	311	+ def fused_rms_norm(self, input_shape, output_shape, hidden_size, eps=1e-6):
	312	+ fused_rms_norm = FusedRmsNorm(hidden_size, eps)
	313	+ ftime, btime, from_cache = operator_cache.find('RmsNorm', input_shape)
	314	+ if not from_cache:
	315	+ ftime, btime = TimeCostModel.profile(fused_rms_norm, [input_shape])
	316	+ ftime_uncertainty = self.sampling.run('RmsNorm', ftime, output_shape, input_shape)
	317	+ btime_uncertainty = self.sampling.run('RmsNormGrad', btime, output_shape, input_shape)
	318	+ operator_cache.record('RmsNorm', input_shape, output_shape, ftime, btime)
	319	+ return ftime_uncertainty, btime_uncertainty
	320	+
	321	+ def flash_attention(self, q, k, v, output_shape, head_dim):
	322	+ flash_attn = FlashAttention(head_dim)
	323	+ ftime, btime, from_cache = operator_cache.find('FlashAttentionScore', [q, k, v])
	324	+ if not from_cache:
	325	+ ftime, btime = TimeCostModel.profile(flash_attn, [q, k, v])
	326	+ ftime_uncertainty = self.sampling.run('FlashAttentionScore', ftime, output_shape, q, k, v)
	327	+ btime_uncertainty = self.sampling.run('FlashAttentionScoreGrad', btime, output_shape, q, k, v)
	328	+ operator_cache.record('FlashAttentionScore', [q, k, v], q, ftime, btime)
	329	+ return ftime_uncertainty, btime_uncertainty
	330	+
	331	+
	332	+class TimeCostModel(object):
	333	+ def __init__(self):
	334	+ args = get_args()
	335	+ self.seq_length = args.seq_length
	336	+ self.hidden_size = args.hidden_size
	337	+ self.pp_size = args.pipeline_model_parallel_size
	338	+ self.dp_size = args.data_parallel_size
	339	+ self.micro_batch_size = args.micro_batch_size
	340	+ self.num_layers_per_stage = args.num_layers // args.pipeline_model_parallel_size
	341	+ self.num_micro_batch = args.global_batch_size // args.micro_batch_size // args.data_parallel_size
	342	+
	343	+ def get_iteration_time(self):
	344	+ transformer_block = TransformerBlock()
	345	+ fwd_time, bwd_time = transformer_block.get_block_time()
	346	+ fwd_time *= self.num_layers_per_stage
	347	+ bwd_time *= self.num_layers_per_stage
	348	+ iteration_times = np.array([0 for _ in range(fwd_time.shape[0])]).astype(np.float64)
	349	+ for i in range(fwd_time.shape[0]):
	350	+ iteration_times[i] = self.pipeline_costmodel(fwd_time[i], bwd_time[i])
	351	+ return iteration_times
	352	+
	353	+ def pipeline_costmodel(self, fwd_time, bwd_time):
	354	+ if self.pp_size == 1:
	355	+ return (fwd_time + bwd_time) * self.num_micro_batch
	356	+
	357	+ send_recv_time = CommProfiling.get_send_recv_time(
	358	+ [self.seq_length, self.micro_batch_size, self.hidden_size]
	359	+ )
	360	+ # p and m start with 1
	361	+ SF = np.zeros((self.pp_size + 1, self.num_micro_batch + 1), np.float64)
	362	+ SB = np.zeros((self.pp_size + 1, self.num_micro_batch + 1), np.float64)
	363	+ EF = np.zeros((self.pp_size + 1, self.num_micro_batch + 1), np.float64)
	364	+ EB = np.zeros((self.pp_size + 1, self.num_micro_batch + 1), np.float64)
	365	+
	366	+ warmup = [self.pp_size - p - 1 for p in range(self.pp_size)]
	367	+ remaining = [self.num_micro_batch - warmup[p] for p in range(self.pp_size)]
	368	+
	369	+ # warmup
	370	+ for p in range(1, self.pp_size + 1):
	371	+ for m in range(1, warmup[p - 1] + 1):
	372	+ if p == 1:
	373	+ SF[p][m] = (m - 1) * fwd_time
	374	+ EF[p][m] = m * fwd_time
	375	+ else:
	376	+ SF[p][m] = max(EF[p][m - 1], EF[p - 1][m] + send_recv_time)
	377	+ EF[p][m] = SF[p][m] + fwd_time
	378	+
	379	+ # 1f1b
	380	+ for num_1f1b in range(1, self.num_micro_batch + 1):
	381	+ # forward of 1f1b
	382	+ for p in range(1, self.pp_size + 1):
	383	+ if num_1f1b > remaining[p - 1]:
	384	+ # cool down phase
	385	+ continue
	386	+ m = warmup[p - 1] + num_1f1b
	387	+ if p == 1:
	388	+ SF[p][m] = EB[p][m + p - self.pp_size - 1]
	389	+ EF[p][m] = SF[p][m] + fwd_time
	390	+ else:
	391	+ SF[p][m] = max(EB[p][m + p - self.pp_size - 1], EF[p - 1][m] + send_recv_time)
	392	+ EF[p][m] = SF[p][m] + fwd_time
	393	+
	394	+ # backward of 1f1b
	395	+ for p in range(self.pp_size, 0, -1):
	396	+ m = num_1f1b
	397	+ if num_1f1b > remaining[p - 1]:
	398	+ # cool down phase
	399	+ continue
	400	+ if p == self.pp_size:
	401	+ SB[p][m] = EF[p][m]
	402	+ else:
	403	+ SB[p][m] = max(EF[p][m + self.pp_size - p], EB[p + 1][m] + send_recv_time)
	404	+ EB[p][m] = SB[p][m] + bwd_time
	405	+
	406	+ # cool down phase
	407	+ for p in range(self.pp_size, 0, -1):
	408	+ m = num_1f1b
	409	+ if num_1f1b <= remaining[p - 1]:
	410	+ continue
	411	+ SB[p][m] = max(EB[p][m - 1], EB[p + 1][m] + send_recv_time)
	412	+ EB[p][m] = SB[p][m] + bwd_time
	413	+
	414	+ e2e_time = max([max(EB[p]) for p in range(self.pp_size)])
	415	+ # allreduce_gradients
	416	+ e2e_time += 0.0
	417	+ return e2e_time
	418	+
	419	+ @staticmethod
	420	+ def profile(model, shapes):
	421	+ model.to(torch.cuda.current_device())
	422	+
	423	+ input_tensors = []
	424	+ index = 0
	425	+ for shape in shapes:
	426	+ tensor = GlobalMemoryBuffer.get_tensor(shape, index).requires_grad_()
	427	+ input_tensors.append(tensor)
	428	+ index += 1
	429	+
	430	+ sum_z = None
	431	+ for _ in range(3):
	432	+ sum_z = model(input_tensors)
	433	+
	434	+ # forward_time
	435	+ torch.npu.synchronize()
	436	+ start_time = time.time()
	437	+ for _ in range(ITERATION_LOOP_TIME):
	438	+ model(input_tensors)
	439	+ torch.npu.synchronize()
	440	+ fwd_time = (time.time() - start_time) * 1e6 / ITERATION_LOOP_TIME
	441	+
	442	+ for _ in range(3):
	443	+ z = model(input_tensors)
	444	+ loss = torch.sum(z)
	445	+ loss.backward()
	446	+
	447	+ torch.npu.synchronize()
	448	+ start_time = time.time()
	449	+ for _ in range(ITERATION_LOOP_TIME):
	450	+ torch.sum(sum_z)
	451	+ torch.npu.synchronize()
	452	+ loss_time = (time.time() - start_time) * 1e6 / ITERATION_LOOP_TIME
	453	+
	454	+ torch.npu.synchronize()
	455	+ start_time = time.time()
	456	+ for i in range(ITERATION_LOOP_TIME):
	457	+ z = model(input_tensors)
	458	+ loss = torch.sum(z)
	459	+ loss.backward()
	460	+ torch.npu.synchronize()
	461	+ bwd_time = (time.time() - start_time) * 1e6 / ITERATION_LOOP_TIME - fwd_time - loss_time
	462	+ return fwd_time, bwd_time

Amindspeed/core/auto_parallel/auto_parallel_optimizer.py+127-0

@@ -0,0 +1,127 @@
	1	+# coding=utf-8
	2	+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+import os
	17	+import time
	18	+import math
	19	+import random
	20	+from pathlib import Path
	21	+
	22	+import pandas as pd
	23	+import gpytorch
	24	+
	25	+from mindspeed.core.auto_parallel import (
	26	+ model_manager,
	27	+ sample_cache,
	28	+ operator_cache,
	29	+)
	30	+from mindspeed.core.auto_parallel.auto_parallel_rectify import ExactGPModel
	31	+from mindspeed.core.auto_parallel.auto_parallel_model import TimeCostModel
	32	+from mindspeed.core.auto_parallel.auto_parallel_profiling import (
	33	+ BaseLaunch,
	34	+ DistributedOperateProfiler,
	35	+ DistributedPerformanceProfiler
	36	+)
	37	+
	38	+
	39	+class SearchByGreyBox:
	40	+ def __init__(self, stop_threshold=0.05):
	41	+ self.operators = [
	42	+ 'MatMul',
	43	+ 'RmsNorm',
	44	+ 'RmsNormGrad',
	45	+ 'FlashAttentionScore',
	46	+ 'FlashAttentionScoreGrad'
	47	+ ]
	48	+ self.stop_threshold = stop_threshold
	49	+ self.config_performances = {}
	50	+ self.exist_config = []
	51	+ self.e2e_log = pd.DataFrame()
	52	+
	53	+ @staticmethod
	54	+ def find_csv(operator_profile, key='kernel_details'):
	55	+ csv_files = []
	56	+ for cf in list(Path(operator_profile).rglob('*.csv')):
	57	+ if key in str(cf):
	58	+ csv_files.append(os.path.abspath(str(cf)))
	59	+ if len(csv_files) <= 0:
	60	+ print(f"not find kernel_details.csv")
	61	+ return None
	62	+ return sorted(csv_files)[0]
	63	+
	64	+ @staticmethod
	65	+ def theory_modeling(config):
	66	+ base_launch = BaseLaunch()
	67	+ base_launch.update_args(config)
	68	+ cost_time = TimeCostModel().get_iteration_time()
	69	+ base_launch.recover_args()
	70	+ return cost_time
	71	+
	72	+ def save(self, config, cost_time):
	73	+ self.e2e_log[str(config)] = cost_time
	74	+
	75	+ def generate_config(self):
	76	+ best_config = self.e2e_log.apply(lambda col: col.idxmin(), axis=1).values
	77	+ rest_config = [i for i in best_config if str(i) not in self.exist_config]
	78	+ prop = len(rest_config) / len(best_config)
	79	+ if prop > self.stop_threshold:
	80	+ sample = random.choice(rest_config)
	81	+ self.exist_config.append(sample)
	82	+ return eval(sample)
	83	+ print(f'Unexplored proportion: {prop} < stop_thd :{self.stop_threshold}, early stop triggered.')
	84	+ return None
	85	+
	86	+ def train(self, train_profiling_file, train_operator_data):
	87	+ for operator in self.operators:
	88	+ model = model_manager.get_cached_model(operator)
	89	+ model.fit(train_profiling_file, train_operator_data)
	90	+
	91	+ def load_base_model(self, model_dir):
	92	+ for operator in self.operators:
	93	+ likelihood = gpytorch.likelihoods.GaussianLikelihood(gpytorch.priors.NormalPrior(1e-3, 0.02))
	94	+ model = ExactGPModel(operator=operator, likelihood=likelihood)
	95	+ model_manager.load_model(model, operator, model_dir)
	96	+
	97	+ def search(self, args, search_spaces):
	98	+ start_time = time.time()
	99	+ self.load_base_model(os.path.dirname(os.path.abspath(__file__)) + os.sep + 'noise_predict_ckpt')
	100	+ while ((time.time() - start_time) / 3600) < 8 \
	101	+ and len(self.config_performances) < len(search_spaces):
	102	+ for config in search_spaces:
	103	+ cost_time = SearchByGreyBox.theory_modeling(config)
	104	+ self.save(config, cost_time)
	105	+ print(f"complete model config: {config}")
	106	+
	107	+ next_config = self.generate_config()
	108	+ if next_config is None:
	109	+ break
	110	+ print(f"next_config={next_config}")
	111	+
	112	+ operator_profile_path = DistributedOperateProfiler().launch(next_config)
	113	+ duration_time = DistributedPerformanceProfiler().launch(next_config)
	114	+ self.config_performances[duration_time] = str(next_config)
	115	+ if math.isinf(duration_time):
	116	+ search_spaces.remove(next_config)
	117	+
	118	+ operator_data = operator_cache.data_frame
	119	+ operator_profile = SearchByGreyBox.find_csv(operator_profile_path)
	120	+ if operator_profile is not None:
	121	+ print(f"operator_data: {operator_data}\noperator_profile: {operator_profile}")
	122	+ self.train(operator_profile, operator_data)
	123	+ sample_cache.clear_cache()
	124	+
	125	+ model_manager.save_models('final_model')
	126	+ min_key = min(self.config_performances.keys())
	127	+ return eval(self.config_performances.get(min_key)), min_key

Amindspeed/core/auto_parallel/auto_parallel_profiling.py+400-0

@@ -0,0 +1,400 @@
	1	+# coding=utf-8
	2	+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+import os
	17	+import stat
	18	+import sys
	19	+import time
	20	+import json
	21	+import copy
	22	+import re
	23	+import operator
	24	+import functools
	25	+import subprocess
	26	+import signal
	27	+import threading
	28	+
	29	+import pandas as pd
	30	+import torch
	31	+import torch_npu
	32	+from torch_npu.profiler.profiler import analyse
	33	+from megatron.training.global_vars import set_args, get_args
	34	+
	35	+from mindspeed.core.auto_parallel import (
	36	+ SingletonType,
	37	+ get_cache_path,
	38	+ get_kv_store,
	39	+ analyse_module_profile,
	40	+ MODULE_PATTERN,
	41	+ OPERATOR_PATTERN,
	42	+ BAND_WIDTH_UNIDIRECTIONAL
	43	+)
	44	+
	45	+
	46	+class BaseLaunch:
	47	+ def __init__(self):
	48	+ self.old_args = None
	49	+
	50	+ def launch(self, config):
	51	+ def update_or_append_param(argv: list, key, value=None):
	52	+ if not value:
	53	+ argv.append(key)
	54	+ return
	55	+
	56	+ if key in argv:
	57	+ argv[argv.index(key) + 1] = value
	58	+ else:
	59	+ argv.extend([key, value])
	60	+
	61	+ def remove_param(argv: list, key, has_value=False):
	62	+ if key in argv:
	63	+ pos = argv.index(key)
	64	+ argv.pop(pos)
	65	+ if has_value:
	66	+ argv.pop(pos)
	67	+
	68	+ def monitor_exit(process):
	69	+ while True:
	70	+ exit_flag = get_kv_store().get("exit_flag")
	71	+ if int(exit_flag) == 1:
	72	+ try:
	73	+ process_group_id = os.getpgid(process.pid)
	74	+ os.killpg(process_group_id, signal.SIGKILL)
	75	+ break
	76	+ except ProcessLookupError:
	77	+ break
	78	+ time.sleep(60)
	79	+
	80	+ args = get_args()
	81	+ argv: list = sys.argv[1:]
	82	+ update_or_append_param(argv, '--eval-iters', '0')
	83	+ update_or_append_param(argv, '--train-iters', '5')
	84	+ update_or_append_param(argv, '--global-batch-size', str(args.global_batch_size))
	85	+ update_or_append_param(argv, '--num-layers', str(args.num_layers))
	86	+ update_or_append_param(argv, '--pipeline-model-parallel-size', str(args.pipeline_model_parallel_size))
	87	+ update_or_append_param(argv, '--tensor-model-parallel-size', str(args.tensor_model_parallel_size))
	88	+ update_or_append_param(argv, '--micro-batch-size', str(args.micro_batch_size))
	89	+ update_or_append_param(argv, '--sequence-parallel')
	90	+ if args.profile:
	91	+ update_or_append_param(argv, '--profile')
	92	+ if args.profile_memory:
	93	+ update_or_append_param(argv, '--profile-memory')
	94	+ if args.module_profile_path:
	95	+ update_or_append_param(argv, '--prof-file', str(args.module_profile_path))
	96	+ if args.context_parallel_algo == 'hybrid_cp_algo':
	97	+ update_or_append_param(argv, '--context-parallel-algo', 'hybrid_cp_algo')
	98	+ update_or_append_param(argv, '--context-parallel-size', str(args.context_parallel_size))
	99	+ update_or_append_param(argv, '--ulysses-degree-in-cp', str(args.ulysses_degree_in_cp))
	100	+ if args.context_parallel_algo == 'megatron_cp_algo':
	101	+ update_or_append_param(argv, '--context-parallel-algo', 'megatron_cp_algo')
	102	+ update_or_append_param(argv, '--context-parallel-size', str(args.context_parallel_size))
	103	+ if args.context_parallel_algo == 'ulysses_cp_algo':
	104	+ update_or_append_param(argv, '--context-parallel-algo', 'ulysses_cp_algo')
	105	+ update_or_append_param(argv, '--context-parallel-size', str(args.context_parallel_size))
	106	+ remove_param(argv, '--auto-parallel')
	107	+
	108	+ command = [
	109	+ 'torchrun',
	110	+ '--nproc_per_node', str(args.nproc_per_node),
	111	+ '--nnodes', str(args.nnodes),
	112	+ '--node-rank', str(args.node_rank),
	113	+ '--master_addr', str(args.master_addr),
	114	+ '--master_port', str(args.master_port),
	115	+ str(sys.argv[0])
	116	+ ] + argv
	117	+
	118	+ get_kv_store().set("exit_flag", "0")
	119	+ process = subprocess.Popen(command, shell=False, preexec_fn=lambda: os.setpgrp())
	120	+ monitor_thread = threading.Thread(target=monitor_exit, args=(process,))
	121	+ monitor_thread.start()
	122	+ process.wait()
	123	+ get_kv_store().set("exit_flag", "1")
	124	+ torch.distributed.barrier()
	125	+
	126	+ def update_args(self, config):
	127	+ args = get_args()
	128	+ self.old_args = copy.deepcopy(args)
	129	+
	130	+ args.pipeline_model_parallel_size = config[0]
	131	+ args.tensor_model_parallel_size = config[1]
	132	+ args.data_parallel_size = config[2]
	133	+ args.context_parallel_size = config[3] * config[4]
	134	+ args.ulysses_degree_in_cp = config[4]
	135	+ args.micro_batch_size = config[5]
	136	+ if config[3] > 1 and config[4] > 1:
	137	+ args.context_parallel_algo = 'hybrid_cp_algo'
	138	+ args.use_cp_send_recv_overlap = True
	139	+ elif config[3] > 1 and config[4] == 1:
	140	+ args.context_parallel_algo = 'megatron_cp_algo'
	141	+ args.use_cp_send_recv_overlap = True
	142	+ elif config[3] == 1 and config[4] > 1:
	143	+ args.context_parallel_algo = 'ulysses_cp_algo'
	144	+
	145	+ def recover_args(self):
	146	+ set_args(self.old_args)
	147	+
	148	+
	149	+class DistributedMemoryProfiler(BaseLaunch):
	150	+ def update_args(self, config):
	151	+ super().update_args(config)
	152	+ args = get_args()
	153	+ args.module_profile_path = (get_cache_path() + MODULE_PATTERN).format(*config)
	154	+ args.global_batch_size = args.pipeline_model_parallel_size * args.data_parallel_size * args.micro_batch_size
	155	+ args.num_layers = args.pipeline_model_parallel_size
	156	+ args.profile_memory = True
	157	+
	158	+ def launch(self, config):
	159	+ args = get_args()
	160	+ if args.node_rank != 0:
	161	+ self.update_args(config)
	162	+ super().launch(config)
	163	+ super().recover_args()
	164	+ return None
	165	+
	166	+ self.update_args(config)
	167	+ module_profile_path = get_args().module_profile_path
	168	+ if os.path.exists(module_profile_path):
	169	+ super().recover_args()
	170	+ return analyse_module_profile(module_profile_path, key='transformer_act_mem')
	171	+
	172	+ buffer = config + [0]
	173	+ torch.distributed.broadcast(torch.tensor(buffer, dtype=torch.int), 0)
	174	+
	175	+ super().launch(config)
	176	+ super().recover_args()
	177	+ return analyse_module_profile(module_profile_path, key='transformer_act_mem')
	178	+
	179	+
	180	+class DistributedOperateProfiler(BaseLaunch):
	181	+ def update_args(self, config):
	182	+ super().update_args(config)
	183	+ args = get_args()
	184	+ args.module_profile_path = None
	185	+ args.operator_profile_path = (get_cache_path() + OPERATOR_PATTERN).format(*config)
	186	+ args.global_batch_size = 4 * args.pipeline_model_parallel_size * args.data_parallel_size * args.micro_batch_size
	187	+ args.num_layers = 2 * args.pipeline_model_parallel_size
	188	+ args.profile_operator = True
	189	+
	190	+ def launch(self, config):
	191	+ self.update_args(config)
	192	+ args = get_args()
	193	+ if args.node_rank != 0:
	194	+ super().launch(config)
	195	+ super().recover_args()
	196	+ return None
	197	+
	198	+ operator_profile_path = args.operator_profile_path
	199	+ if os.path.exists(operator_profile_path):
	200	+ super().recover_args()
	201	+ return operator_profile_path
	202	+
	203	+ buffer = config + [1]
	204	+ torch.distributed.broadcast(torch.tensor(buffer, dtype=torch.int), 0)
	205	+
	206	+ os.environ['ASCEND_WORK_PATH'] = operator_profile_path
	207	+ os.makedirs(operator_profile_path)
	208	+ super().launch(config)
	209	+ super().recover_args()
	210	+
	211	+ analyse_thread = threading.Thread(
	212	+ target=analyse, args=(operator_profile_path + os.sep + 'profiling_data', 32)
	213	+ )
	214	+ analyse_thread.daemon = True
	215	+ analyse_thread.start()
	216	+ return operator_profile_path
	217	+
	218	+
	219	+class DistributedPerformanceProfiler(BaseLaunch):
	220	+ def update_args(self, config):
	221	+ super().update_args(config)
	222	+ args = get_args()
	223	+ args.module_profile_path = (get_cache_path() + MODULE_PATTERN).format(*config)
	224	+
	225	+ def launch(self, config):
	226	+ self.update_args(config)
	227	+ args = get_args()
	228	+ if args.node_rank != 0:
	229	+ super().launch(config)
	230	+ super().recover_args()
	231	+ return None
	232	+
	233	+ module_profile_path = get_args().module_profile_path
	234	+ if os.path.exists(module_profile_path):
	235	+ super().recover_args()
	236	+ return analyse_module_profile(module_profile_path, key='step_time')
	237	+
	238	+ buffer = config + [2]
	239	+ torch.distributed.broadcast(torch.tensor(buffer, dtype=torch.int), 0)
	240	+ super().launch(config)
	241	+ super().recover_args()
	242	+ return analyse_module_profile(module_profile_path, key='step_time')
	243	+
	244	+
	245	+class OperateProfile(metaclass=SingletonType):
	246	+ def __init__(self, args):
	247	+ experimental_config = torch_npu.profiler._ExperimentalConfig(
	248	+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
	249	+ data_simplification=False
	250	+ )
	251	+ activities = [torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU]
	252	+ self.op_profiler = torch_npu.profiler.profile(
	253	+ activities=activities,
	254	+ record_shapes=True,
	255	+ schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=2),
	256	+ experimental_config=experimental_config,
	257	+ )
	258	+ self.op_profiler.start()
	259	+ self.node_rank = torch.distributed.get_rank() // args.nproc_per_node
	260	+
	261	+ def step(self):
	262	+ if self.node_rank in (0,):
	263	+ self.op_profiler.step()
	264	+
	265	+ def stop(self):
	266	+ if self.node_rank in (0,):
	267	+ self.op_profiler.stop()
	268	+
	269	+
	270	+class Profiling(metaclass=SingletonType):
	271	+ MEMORY_UNIT = 1024 ** 3
	272	+
	273	+ def __init__(self, args, warmup_step=3, stop_step=5):
	274	+ self.args = args
	275	+ self.warmup_step = warmup_step
	276	+ self.stop_step = stop_step
	277	+ self.curr_step = 0
	278	+ self.pattern = r'^module.module.language_model.encoder.layers.\d+$'
	279	+ self.context = {
	280	+ 'step_time': 0,
	281	+ 'tranformer_act_mem': 0
	282	+ }
	283	+
	284	+ def should_profiling(self):
	285	+ rank = torch.distributed.get_rank()
	286	+ if rank in self.args.profile_ranks and \
	287	+ self.warmup_step <= self.curr_step < self.stop_step:
	288	+ return True
	289	+ return False
	290	+
	291	+ def forward_pre_hook(self):
	292	+ def hook(module, args, *kwargs):
	293	+ if torch.distributed.get_rank() in self.args.profile_ranks:
	294	+ torch.npu.synchronize()
	295	+ self.start_memory = torch.npu.memory_allocated()
	296	+ torch.npu.reset_max_memory_allocated()
	297	+ return hook
	298	+
	299	+ def forward_post_hook(self):
	300	+ def hook(module, args, *kwargs):
	301	+ if torch.distributed.get_rank() in self.args.profile_ranks:
	302	+ torch.npu.synchronize()
	303	+ self.end_memory = torch.npu.max_memory_allocated()
	304	+ tranformer_act_mem = (self.end_memory - self.start_memory) / Profiling.MEMORY_UNIT
	305	+ self.context['tranformer_act_mem'] = tranformer_act_mem
	306	+ return hook
	307	+
	308	+ def register_recursive_hook(self, prefix_name, model):
	309	+ model = model[0] if isinstance(model, list) else model
	310	+ for name, module in model.named_children():
	311	+ next_name = prefix_name + "." + name if prefix_name != "" else name
	312	+ if re.fullmatch(self.pattern, next_name):
	313	+ module.register_forward_pre_hook(self.forward_pre_hook())
	314	+ module.register_forward_hook(self.forward_post_hook())
	315	+ break
	316	+ self.register_recursive_hook(next_name, module)
	317	+
	318	+ def hook_train_step(self, train_step):
	319	+ def custom_train_step(args, *kwargs):
	320	+ start_time = time.time()
	321	+ result = train_step(args, *kwargs)
	322	+ torch.cuda.synchronize()
	323	+ step_time = time.time() - start_time
	324	+ if self.should_profiling():
	325	+ cur_step_time = self.context.get('step_time')
	326	+ cur_step_time += (step_time - cur_step_time) / (self.curr_step - self.warmup_step + 1)
	327	+ self.context['step_time'] = cur_step_time
	328	+ self.export_to_file()
	329	+ self.curr_step += 1
	330	+ return result
	331	+ return custom_train_step
	332	+
	333	+ def export_to_file(self):
	334	+ if torch.distributed.get_rank() in self.args.profile_ranks:
	335	+ flags = os.O_WRONLY \| os.O_CREAT \| os.O_TRUNC
	336	+ modes = stat.S_IWUSR \| stat.S_IRUSR
	337	+ with os.fdopen(os.open(self.args.prof_file, flags, modes), 'w') as fout:
	338	+ fout.write(json.dumps(self.context))
	339	+
	340	+
	341	+class CommProfiling:
	342	+ @staticmethod
	343	+ def get_comm_time(shape, domains, op):
	344	+ if domains == 1:
	345	+ return 0
	346	+
	347	+ if op == 'all_reduce':
	348	+ return CommProfiling.cal_all_reduce(shape, domains)
	349	+ if op == 'all_gather':
	350	+ return CommProfiling.cal_all_gather(shape, domains)
	351	+ if op == 'alltoall':
	352	+ return CommProfiling.cal_alltoall(shape, domains)
	353	+ if op == 'reduce_scatter':
	354	+ return CommProfiling.cal_reduce_scatter(shape, domains)
	355	+ raise AssertionError('communicate operator type error')
	356	+
	357	+ @staticmethod
	358	+ def cal_all_reduce(shape, domains):
	359	+ data_size = CommProfiling.get_data_size(shape)
	360	+ data_size = data_size / domains * (domains - 1) * domains * 2
	361	+ band_width = domains * (domains - 1) / 2 * BAND_WIDTH_UNIDIRECTIONAL
	362	+ return CommProfiling.div(data_size, band_width)
	363	+
	364	+ @staticmethod
	365	+ def cal_all_gather(shape, domains):
	366	+ data_size = CommProfiling.get_data_size(shape)
	367	+ data_size = data_size / domains * (domains - 1) * domains
	368	+ band_width = domains * (domains - 1) / 2 * BAND_WIDTH_UNIDIRECTIONAL
	369	+ return CommProfiling.div(data_size, band_width)
	370	+
	371	+ @staticmethod
	372	+ def cal_alltoall(shape, domains):
	373	+ data_size = CommProfiling.get_data_size(shape)
	374	+ data_size = data_size / domains * (domains - 1) * domains
	375	+ band_width = domains * (domains - 1) / 2 * BAND_WIDTH_UNIDIRECTIONAL
	376	+ return CommProfiling.div(data_size, band_width)
	377	+
	378	+ @staticmethod
	379	+ def cal_reduce_scatter(shape, domains):
	380	+ data_size = CommProfiling.get_data_size(shape)
	381	+ data_size = data_size / domains * (domains - 1) * domains
	382	+ band_width = domains * (domains - 1) / 2 * BAND_WIDTH_UNIDIRECTIONAL
	383	+ return CommProfiling.div(data_size, band_width)
	384	+
	385	+ @staticmethod
	386	+ def get_send_recv_time(shape):
	387	+ data_size = CommProfiling.get_data_size(shape)
	388	+ return (data_size / BAND_WIDTH_UNIDIRECTIONAL) * 1e6
	389	+
	390	+ @staticmethod
	391	+ def get_data_size(shape):
	392	+ return functools.reduce(operator.mul, shape) * 2 // 1024**3
	393	+
	394	+ @staticmethod
	395	+ def div(data_size, band_width):
	396	+ try:
	397	+ return data_size / band_width * 1e6
	398	+ except ZeroDivisionError:
	399	+ print(f"band_width is zero")
	400	+ return 0

Amindspeed/core/auto_parallel/auto_parallel_rectify.py+420-0

@@ -0,0 +1,420 @@
	1	+# coding=utf-8
	2	+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved.
	3	+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	4	+#
	5	+# Licensed under the Apache License, Version 2.0 (the "License");
	6	+# you may not use this file except in compliance with the License.
	7	+# You may obtain a copy of the License at
	8	+#
	9	+# http://www.apache.org/licenses/LICENSE-2.0
	10	+#
	11	+# Unless required by applicable law or agreed to in writing, software
	12	+# distributed under the License is distributed on an "AS IS" BASIS,
	13	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	14	+# See the License for the specific language governing permissions and
	15	+# limitations under the License.
	16	+import os
	17	+import glob
	18	+import copy
	19	+import warnings
	20	+from typing import Optional
	21	+
	22	+import pandas as pd
	23	+import numpy as np
	24	+import gpytorch
	25	+import torch
	26	+
	27	+from mindspeed.core.auto_parallel import (
	28	+ ARD_NUM_DIMS,
	29	+ KeyField,
	30	+ sample_cache,
	31	+ model_manager
	32	+)
	33	+
	34	+
	35	+class ExactGPModel(gpytorch.models.ExactGP):
	36	+ def __init__(self, operator, train_inputs=None,
	37	+ train_targets=None, raw_lengthscale=None,
	38	+ likelihood=None, dtype=torch.float64):
	39	+ super(ExactGPModel, self).__init__(train_inputs, train_targets, likelihood=likelihood)
	40	+ self.operator = operator
	41	+ self.dtype = dtype
	42	+
	43	+ self.mean_module = gpytorch.means.ConstantMean()
	44	+ self.covar_module = gpytorch.kernels.ScaleKernel(
	45	+ gpytorch.kernels.MaternKernel(nu=0.5, ard_num_dims=ARD_NUM_DIMS[operator],
	46	+ lengthscale_constraint=gpytorch.constraints.GreaterThan(3e-2)))
	47	+ if raw_lengthscale is not None:
	48	+ self.covar_module.base_kernel.raw_lengthscale.data \
	49	+ = self.raw_lengthscale * torch.ones_like(self.covar_module.base_kernel.raw_lengthscale.data)
	50	+
	51	+ self.train_round = 0
	52	+ self.train_data = pd.DataFrame()
	B B bingobb2024年7月27日注释的风格建议统一，都单独另起一行
	53	+
	54	+ self.y_train_mean: Optional[torch.Tensor] = None
	55	+ self.y_train_std: Optional[torch.Tensor] = None
	56	+ self.x_train_std: Optional[torch.Tensor] = None
	57	+
	58	+ def get_model_info(self):
	59	+ return self.train_data, self.train_round
	60	+
	61	+ def set_model_info(self, values):
	62	+ self.train_data, self.train_round = values
	63	+ # set model info by train_data
	64	+ self.data_standardize()
	65	+
	66	+ def forward(self, x):
	67	+ mean = self.mean_module(x)
	68	+ covar = self.covar_module(x)
	69	+ return gpytorch.distributions.MultivariateNormal(mean, covar)
	70	+
	71	+ def fit(self, profiling_file, multi_operator_data, num_iter=3000, lr=0.03):
	72	+ hd = DataHandler(profiling_file, multi_operator_data)
	73	+ data = hd.generate_data(self.operator)
	74	+ # merge self.train_data with new train_data
	75	+ self.update_data(data)
	76	+ # set model train_inputs and target_inputs
	77	+ self.data_standardize()
	78	+ # clear cache
	79	+ self.train()
	80	+ self.likelihood.train()
	81	+ optimizer = torch.optim.Adam(self.parameters(), lr=lr)
	82	+ mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.likelihood, self)
	83	+ for i in range(num_iter):
	84	+ optimizer.zero_grad()
	85	+ output = self(self.train_inputs[0])
	86	+ loss = -mll(output, self.train_targets)
	87	+ loss.backward()
	88	+ if i % 100 == 0:
	89	+ logs = 'Iter %d/%d - Loss: %.5f outputscale: %.5f noise: %.5f' % (
	90	+ i + 1, num_iter, loss.item(),
	91	+ self.covar_module.outputscale.item(),
	92	+ self.likelihood.noise.item()
	93	+ ) + ' lengthscale: ' + str(
	94	+ np.round(self.covar_module.base_kernel.lengthscale.detach().cpu().numpy()[0], 5))
	95	+ print(logs)
	96	+ optimizer.step()
	97	+ self.eval()
	98	+ self.likelihood.eval()
	99	+ self.train_round += 1
	100	+
	101	+ def update_data(self, data: pd.DataFrame):
	102	+ """
	103	+ :param data columns = [shape error count]
	104	+ """
	105	+ if not self.train_data.empty:
	106	+ exits_shapes = self.train_data.loc[:, KeyField.InputShapes].values.tolist()
	107	+ for index, rows in data.iterrows():
	108	+ shape = getattr(rows, KeyField.InputShapes)
	109	+ # update existent input_shape
	110	+ if shape in exits_shapes:
	111	+ error, number = data[data[KeyField.InputShapes] == shape].iloc[:, 1:3].values.flatten()
	112	+ current_train_data = self.train_data[self.train_data[KeyField.InputShapes] == shape]
	113	+ train_error, train_number = current_train_data.iloc[:, 1:3].values.flatten()
	114	+ count = int(number + train_number)
	115	+ new_error = (error * number + train_error * train_number) / count
	116	+ self.train_data[self.train_data[KeyField.InputShapes] == shape] = [shape, new_error, count]
	117	+ else:
	118	+ # save new input_shape
	119	+ self.train_data = pd.concat([self.train_data, rows.to_frame().T], ignore_index=True)
	120	+ else:
	121	+ self.train_data = data
	122	+
	123	+ def data_standardize(self):
	124	+ y_train = torch.tensor(self.train_data['error'], dtype=self.dtype)
	125	+ x_train = self.train_data[KeyField.InputShapes].str.split(',', expand=True).values.astype(int)
	126	+ x_train = torch.tensor(x_train, dtype=self.dtype).log()
	127	+ if x_train.shape[0] == 1:
	128	+ self.x_train_std = torch.tensor(np.ones(x_train.shape), dtype=self.dtype)
	129	+ self.y_train_std = torch.tensor(1, dtype=self.dtype)
	130	+ else:
	131	+ self.x_train_std, self.y_train_std = torch.std(x_train, dim=0), torch.std(y_train, dim=0)
	132	+ self.x_train_std[self.x_train_std == 0] = 1.
	133	+ self.y_train_std[self.y_train_std == 0] = 1.
	134	+ x_train /= self.x_train_std
	135	+ self.y_train_mean = torch.mean(y_train, dim=0)
	136	+ y_train = (y_train - self.y_train_mean) / self.y_train_std
	137	+ self.set_train_data(x_train, y_train, strict=False)
	138	+
	139	+
	140	+class Sampler:
	141	+ def __init__(self, num_sample=10, pre_thd=0):
	142	+ self.pre_thd = pre_thd
	143	+ self.num_sample = torch.Size([num_sample])
	144	+
	145	+ def run(self, operator, direct_time, output_shape: list, *input_shape):
	146	+ input_shape = copy.deepcopy(input_shape)
	147	+ output_shape = copy.deepcopy(output_shape)
	148	+ # modify input_shape
	149	+ input_shape = Sampler.reduce_dim(operator, output_shape, input_shape)
	150	+ # check cache
	151	+ cached_samples = getattr(sample_cache, operator)
	152	+ sample = cached_samples.get(input_shape, None)
	153	+ if sample is not None:
	154	+ return sample
	155	+ # load model
	156	+ model = model_manager.get_cached_model(operator)
	157	+ if model is None:
	158	+ raise AssertionError("Can't find {operator} model")
	159	+
	160	+ # predict
	161	+ input_shape_np = np.array(input_shape).reshape(1, -1)
	162	+ fixed_shape = np.concatenate([input_shape_np, input_shape_np], axis=0)
	163	+ x = torch.tensor(fixed_shape, dtype=torch.float64).log()
	164	+ with torch.no_grad(), gpytorch.settings.fast_pred_var():
	165	+ pred = model(x / model.x_train_std)
	166	+ pred = pred * model.y_train_std.item() + model.y_train_mean.item()
	167	+ relative_error = pred.sample(self.num_sample).cpu().numpy()[:, 0]
	168	+ sample = direct_time * (relative_error + 1.).flatten()
	169	+ negative_indices = np.where(sample <= self.pre_thd)[0]
	170	+ if negative_indices.size > 0:
	171	+ sample[negative_indices] = 0
	172	+ warnings.warn(f'Uncertainty of {operator} is too large, input shape: {input_shape}', Warning)
	173	+ # save prediction data
	174	+ cached_samples[input_shape] = sample
	175	+ return sample
	176	+
	177	+ @staticmethod
	178	+ def reduce_dim(operator, output_shape, input_shapes):
	179	+ input_shapes = copy.deepcopy(input_shapes)
	180	+ output_shape = copy.deepcopy(output_shape)
	181	+ if operator in ['LayerNorm', 'LayerNormGrad', 'LayerNormV3']:
	182	+ input_shape = input_shapes[0]
	183	+ elif operator in ['FastGelu', 'FastGeluGrad']:
	184	+ input_shape = output_shape
	185	+ elif operator in ['Softmax', 'SoftmaxGrad']:
	186	+ input_shape = output_shape
	187	+ elif operator == 'Add' or operator == 'Mul':
	188	+ if len(input_shapes[0]) >= len(input_shapes[1]):
	189	+ max_dims, min_dims = input_shapes
	190	+ else:
	191	+ min_dims, max_dims = input_shapes
	192	+ if len(max_dims) == 2:
	193	+ max_dims.insert(0, 1)
	194	+ if len(max_dims) == 1:
	195	+ max_dims = [1, 1, max_dims[0]]
	196	+ if len(min_dims) == 3:
	197	+ min_dims = [1, 1, 1]
	198	+ elif len(min_dims) == 2:
	199	+ min_dims = [2, 1, 1]
	200	+ else:
	201	+ min_dims = [2, 2, 1]
	202	+ max_dims.extend(min_dims)
	203	+ input_shape = max_dims
	204	+ elif operator == 'BatchMatMul':
	205	+ if len(input_shapes) != 2:
	206	+ raise AssertionError(f"Dim of BatchMatMul is {len(input_shapes)}")
	207	+ b, k, m = output_shape[0], output_shape[2], output_shape[1]
	208	+ n = input_shapes[0][1:] + input_shapes[1][1:]
	209	+ for shape in output_shape[1:]:
	210	+ n.remove(shape)
	211	+ input_shape = [b, m, n[0], k]
	212	+ elif operator == 'MatMul':
	213	+ if len(input_shapes) != 2:
	214	+ raise AssertionError(f"Dim of MatMul is {len(input_shapes)}")
	215	+ input_shape = input_shapes[0]
	216	+ input_shape.extend(input_shapes[1])
	217	+ for shape in output_shape:
	218	+ input_shape.remove(shape)
	219	+ output_shape.insert(1, input_shape[0])
	220	+ input_shape = output_shape
	221	+ elif operator == 'RmsNorm' or operator == 'RmsNormGrad':
	222	+ input_shape = input_shapes[0]
	223	+ elif operator == 'FlashAttentionScore' or operator == 'FlashAttentionScoreGrad':
	224	+ input_shape = input_shapes[0]
	225	+ else:
	226	+ raise ValueError(f"{operator} not supported.")
	227	+
	228	+ return tuple(input_shape)
	229	+
	230	+
	231	+class DataHandler:
	232	+ def __init__(self, profiling_file, multi_operator_data: pd.DataFrame):
	233	+ self.sample_data = multi_operator_data
	234	+ self.profiling = self.extract_target_data(profiling_file)
	235	+ self.current_profiling_operator = None
	236	+ self.current_sample_operator = None
	237	+ self.backward_flag = False
	238	+
	239	+ @staticmethod
	240	+ def extract_target_data(file):
	241	+ if os.path.isdir(file):
	242	+ file = glob.glob(os.path.join(file, "*.csv"))
	243	+ data = pd.concat((pd.read_csv(f) for f in file), ignore_index=True).loc[:,
	244	+ [KeyField.OpType, KeyField.InputShapes, KeyField.OutputShapes, KeyField.Duration]]
	245	+ else:
	246	+ data = pd.read_csv(file).loc[:,
	247	+ [KeyField.OpType, KeyField.InputShapes, KeyField.OutputShapes, KeyField.Duration]]
	248	+ data.loc[data['Type'].str.startswith('MatMul'), 'Type'] = 'MatMul'
	249	+ data.loc[data['Type'].str.startswith('BatchMatMul'), 'Type'] = 'BatchMatMul'
	250	+ data.loc[
	251	+ (data['Type'].str.startswith('Softmax') & ~(data['Type'].str.endswith('Grad'))), 'Type'] = 'Softmax'
	252	+ # filter
	253	+ data = data[(data[KeyField.Duration] > 5) & (data[KeyField.InputShapes].str.len() > 4)].reset_index(drop=True)
	254	+ return data
	255	+
	256	+ @staticmethod
	257	+ def convert_dim(data):
	258	+ new_input_shape = []
	259	+ for index, tmp_data in data[[KeyField.OpType, KeyField.InputShapes, KeyField.OutputShapes]].iterrows():
	260	+ op, input_shape, output_shape = tmp_data.tolist()
	261	+ input_shape, output_shape = eval(input_shape), eval(output_shape)
	262	+ if op == 'LayerNormV3' or op == 'LayerNormGrad':
	263	+ input_shape = input_shape.split(';')[0]
	264	+ elif op == 'Add' or op == 'Mul':
	265	+ dims = input_shape.split(';')
	266	+ d0_l, d1_l = dims[0].split(','), dims[1].split(',')
	267	+ if len(d0_l) >= len(d1_l):
	268	+ max_length_dim = d0_l
	269	+ min_length_dim = d1_l
	270	+ else:
	271	+ max_length_dim = d1_l
	272	+ min_length_dim = d0_l
	273	+ if len(max_length_dim) == 2:
	274	+ max_length_dim = ['1', '1', max_length_dim[0], max_length_dim[1]]
	B B bingobb2024年7月27日无效代码的注释要删掉
	275	+ elif len(max_length_dim) == 1:
	276	+ max_length_dim = ['1', '1', '1', max_length_dim[0]]
	277	+ elif len(max_length_dim) == 3:
	278	+ max_length_dim.insert(0, '1')
	279	+ if len(min_length_dim) == 3:
	280	+ min_length_dim = ['2', '1', '1', '1']
	281	+ elif len(min_length_dim) == 2:
	282	+ min_length_dim = ['2', '2', '1', '1']
	283	+ elif len(min_length_dim) == 1:
	284	+ min_length_dim = ['2', '2', '2', '1']
	285	+ elif len(min_length_dim) == 4:
	286	+ min_length_dim = ['1', '1', '1', '1']
	287	+ max_length_dim.extend(min_length_dim)
	288	+ input_shape = ','.join(max_length_dim)
	289	+ elif op == 'BatchMatMul':
	290	+ output_shape = output_shape.split(',')
	291	+ b, k, m = output_shape[0], output_shape[2], output_shape[1]
	292	+ input_shapes = input_shape.split(';')
	293	+ n = input_shapes[0].split(',')[1:] + input_shapes[1].split(',')[1:]
	294	+ for shape in output_shape[1:]:
	295	+ n.remove(shape)
	296	+ input_shape = ','.join([b, m, n[0], k])
	297	+ elif op == 'MatMul':
	298	+ input_shape = input_shape.replace(';', ',').split(',')
	299	+ output_shape = output_shape.split(',')
	300	+ for shape in output_shape:
	301	+ input_shape.remove(shape)
	302	+ output_shape.insert(1, input_shape[0])
	303	+ input_shape = ','.join(output_shape)
	304	+ elif op == 'Softmax' or op.startswith('SoftmaxGrad'):
	305	+ input_shape = input_shape.split(';')[0]
	306	+ elif op == 'RmsNorm' or op == 'RmsNormGrad':
	307	+ input_shape = input_shape.split(';')[0]
	308	+ elif op == 'FlashAttentionScore' or op == 'FlashAttentionScoreGrad':
	309	+ input_shape = input_shape.split(';')[0]
	310	+ else:
	311	+ raise TypeError(f"{op} don't support")
	312	+ new_input_shape.append(input_shape)
	313	+ return new_input_shape
	314	+
	315	+ def handle_transpose(self):
	316	+ input_shapes = []
	317	+ for index, sample in self.current_profiling_operator.iterrows():
	318	+ input_shape = sample[KeyField.InputShapes]
	319	+ input_shape = eval(input_shape).split(';')
	320	+ input_shape = [list(map(lambda x: int(x), s.split(','))) for s in input_shape]
	321	+ output_shape = eval(sample[KeyField.OutputShapes]).split(',')
	322	+ output_shape = [int(s) for s in output_shape]
	323	+ if sample[KeyField.OpType] == 'BatchMatMul':
	324	+ if output_shape[1] != input_shape[0][1]:
	325	+ input_shape[0][1], input_shape[0][2] = input_shape[0][2], input_shape[0][1]
	326	+ if output_shape[-1] != input_shape[1][-1]:
	327	+ input_shape[1][1], input_shape[1][2] = input_shape[1][2], input_shape[1][1]
	328	+ elif sample[KeyField.OpType] == 'MatMul':
	329	+ if output_shape[0] != input_shape[0][0]:
	330	+ input_shape[0][0], input_shape[0][1] = input_shape[0][1], input_shape[0][0]
	331	+ if output_shape[-1] != input_shape[1][-1]:
	332	+ input_shape[1][0], input_shape[1][1] = input_shape[1][1], input_shape[1][0]
	333	+ input_shape1 = ','.join([str(i) for i in input_shape[0]])
	334	+ input_shape2 = ','.join([str(i) for i in input_shape[1]])
	335	+ input_shape_sum = input_shape1 + ';' + input_shape2
	336	+ input_shapes.append(f'"{input_shape_sum}"')
	337	+ self.current_profiling_operator.loc[:, KeyField.InputShapes] = input_shapes
	338	+
	339	+ def handle_layer_norm_backward(self, operator):
	340	+ v2 = self.profiling[KeyField.OpType] == operator[1]
	341	+ v3 = self.profiling[KeyField.OpType] == operator[0]
	342	+ profiling = self.profiling[v2 \| v3].reset_index(drop=True)
	343	+ # 一个前向layer_norm对应layer_norm_v2和layer_norm_v3，时间做合并
	344	+ back_grad_data = pd.DataFrame()
	345	+ for index in range(0, profiling.shape[0], 2):
	346	+ sum_duration = profiling.loc[index, KeyField.Duration] + profiling.loc[
	347	+ index + 1, KeyField.Duration]
	348	+ input_shape = profiling.loc[index, KeyField.InputShapes].split(';')[0] + '"'
	349	+ back_grad_data.loc[index, KeyField.OpType] = 'LayerNormGrad'
	350	+ back_grad_data.loc[index, KeyField.InputShapes] = input_shape
	351	+ back_grad_data.loc[index, KeyField.OutputShapes] = input_shape
	352	+ back_grad_data.loc[index, KeyField.Duration] = sum_duration
	353	+ return back_grad_data.reset_index(drop=True)
	354	+
	355	+ def handle_fv(self):
	356	+ condition = self.current_profiling_operator[KeyField.InputShapes].str.replace('"', '').str.split(';').map(
	357	+ lambda x: x[:3]).map(lambda x: x[0] == x[1] == x[2])
	358	+ self.current_profiling_operator = self.current_profiling_operator[condition]
	359	+ # 对FV_grad的input_shape可能出现的异常情况容错处理
	360	+ target_shape = self.current_sample_operator[KeyField.InputShapes].values[0]
	361	+ current_shape = self.current_profiling_operator[KeyField.InputShapes].values[0]
	362	+ if target_shape.split(';')[1] != current_shape.split(';')[1]:
	363	+ self.current_profiling_operator[KeyField.InputShapes] = target_shape
	364	+
	365	+ def generate_data(self, operator):
	366	+ # 串行处理各个算子
	367	+ if len(operator) == 2:
	368	+ # layer_norm反向特殊处理
	369	+ self.current_profiling_operator = self.handle_layer_norm_backward(operator)
	370	+ operator = self.current_profiling_operator.loc[0][KeyField.OpType]
	371	+ else:
	372	+ self.current_profiling_operator = self.profiling[self.profiling[KeyField.OpType] == operator]
	B B bingobb2024年7月27日这里也是，无效代码要删掉
	373	+ self.backward_flag = False
	374	+ if operator.endswith('Grad'):
	375	+ self.backward_flag = True
	376	+ operator = operator.split('Grad')[0]
	377	+ # matmul和batch_matmul需要考虑转置情况
	378	+ if operator in ['MatMul', 'BatchMatMul']:
	379	+ self.handle_transpose()
	380	+ # convert sample input_shape
	381	+ self.current_sample_operator = self.sample_data[
	382	+ self.sample_data[KeyField.OpType].str.startswith(operator)].reset_index(
	383	+ drop=True)
	384	+ # 删除负载均衡产生的shape和对FVGrad可能出现的异常Input_shape容错处理.
	385	+ if operator.startswith('FlashAttention'):
	386	+ self.handle_fv()
	387	+ # convert profiling input_shape
	388	+ self.current_profiling_operator.loc[:, KeyField.InputShapes] = self.convert_dim(
	389	+ self.current_profiling_operator
	390	+ )
	391	+ self.current_sample_operator[KeyField.InputShapes] = self.convert_dim(self.current_sample_operator)
	392	+ # 获取当前算子的所有input_shape
	393	+ set_operator = self.current_sample_operator[KeyField.InputShapes].drop_duplicates().tolist()
	394	+ errors_df = pd.DataFrame()
	395	+ # 计算每个input_shape的相对误差
	396	+ for shape in set_operator:
	397	+ # 获取profiling数据当前input_shape的所有样本
	398	+ tmp_data = self.current_profiling_operator[
	399	+ self.current_profiling_operator[KeyField.InputShapes] == shape].copy()
	400	+ if self.backward_flag:
	401	+ direct_mean = self.current_sample_operator[
	402	+ self.current_sample_operator[KeyField.InputShapes] == shape
	403	+ ]['bwd_time'].values[0]
	404	+ else:
	405	+ direct_mean = self.current_sample_operator[
	406	+ self.current_sample_operator[KeyField.InputShapes] == shape
	407	+ ]['fwd_time'].values[0]
	408	+ # 计算相对误差
	409	+ tmp_data['error'] = (tmp_data[KeyField.Duration] - direct_mean) / direct_mean
	410	+ tmp_data['direct_mean'] = direct_mean
	411	+ errors_df = pd.concat([errors_df, tmp_data], axis=0)
	412	+ if errors_df.empty:
	413	+ raise AssertionError('profiling_shape mismatch operator_shape')
	414	+
	415	+ # 分组平均和计数
	416	+ train_data = errors_df.groupby(KeyField.InputShapes).agg(
	417	+ {'error': 'mean', KeyField.InputShapes: 'count'})
	418	+ train_data.rename(columns={KeyField.InputShapes: 'sample_number'}, inplace=True)
	419	+ train_data.reset_index(inplace=True)
	420	+ return train_data

Amindspeed/core/auto_parallel/noise_predict_ckpt/FlashAttentionScoreGrad_910B.pth+0-0