2d张量并行-MindSpeed-AtomGit

MREADME.md+1-0

	@@ -174,6 +174,7 @@ MindSpeed特性由六大模块组成，分别为：megetron特性支持、并行
174	\| 特性 \| 介绍 \|	174	\| 特性 \| 介绍 \|
175	\|------------------------------\|-----------------------------------------------------------\|	175	\|------------------------------\|-----------------------------------------------------------\|
176	\| Ascend nano-pipe流水线并行 \| [link](docs/features/nanopipe-pipeline-parallel.md) \|	176	\| Ascend nano-pipe流水线并行 \| [link](docs/features/nanopipe-pipeline-parallel.md) \|
		177	+\| Ascend 高维张量并行 \| [link](docs/features/tensor-parallel-2d.md) \|
177		178
178	## 关键场景特性	179	## 关键场景特性
179	\| 特性 \| 介绍 \|	180	\| 特性 \| 介绍 \|

Adocs/features/tensor-parallel-2d.md+131-0

@@ -0,0 +1,131 @@
	1	+# 高维张量并行
	2	+
	3	+## 问题分析
	4	+
	5	+大模型训练时，张量并行（TP）将模型参数切分到多个设备上以减少其内存的占用，在训练过程中为了更新参数梯度信息等，需要引入allreduce通信。当集群规模较大时，如果设置TP域很大时，其通信开销会变得很大，使得训练效率降低。
	6	+
	7	+## 解决方案
	8	+
	9	+为了提高大规模TP域通信效率，采用高维张量并行，其将激活值和参数同时切分到多个计算设备上，相对1D-TP降低了通信域、减少通信次数，从而减少通信时间，提升模型训练的性能。
	10	+
	11	+### 解决思路
	12	+
	13	+#### 2D张量并行策略
	14	+
	15	+给定TP域大小，通过建立多通信域，在原Megatron（ColumnParallelLinear、RowParallelLinear）增加了一维的切分维度。将原tp通信域进行分解为两个子通信域tp_x和tp_y，需要满足`tp = tp_x * tp_y`。以MLP层为例，其实现过程如下：
	16	+
	17	+![img](../../sources/images/tensor-parallel-2d.png)
	B B bingobb2024年10月19日图画的有问题，第二个allgather应该是对y做
	18	+
	19	+#### 分布式normalization
	20	+
	21	+在transformer网络中，normalization会将每一层神经元的输入都转成均值方差都一样的，加快其收敛。在MLP和attention层分别进行2D张量并行时，其输入和输出都分别在first-dim和last-dim做了tp_x和tp_y的切分，如果继续使用原LayerNorm或者RMSNorm需要先将input进行沿first-dim进行all-gather(x)和沿last-dim进行all-gather(y)操作，才能保证input数据的完整性。为了提升这部分的性能，采用了分布式normalization。其处理流程如下：
	22	+
	23	+##### 步骤1：计算输入的总和
	24	+
	25	+首先，计算输入张量$\mathbf{x}$ 在最后一个维度上的总和：
	26	+
	27	+$$
	28	+e_x = \sum_{i=1}^{H} x_i
	29	+\
	30	+$$
	31	+
	32	+##### 步骤2：分布式归约操作（All-Reduce）
	33	+
	34	+将步骤1中的总和 $e_x$ 在所有tp_y通信域进程中进行归约（求和），确保每个进程都拥有其通信域全局总和：
	35	+$$
	36	+\
	37	+e_x^{\text{global}} = \text{AllReduce}\left( e_x \right) = \sum_{p=1}^{P} \sum_{i=1}^{H} x_i^{(p)}
	38	+\
	39	+$$
	40	+
	41	+其中：
	42	+- $P$ 是分布式进程的数量。
	43	+- $x_i^{(p)}$ 表示第 $p$ 个进程中第 $i$ 个元素的值。
	44	+
	45	+##### 步骤3：计算输入元素的平方和
	46	+
	47	+接下来，计算输入张量每个元素的平方和：
	48	+
	49	+$$
	50	+s_x = \sum_{i=1}^{H} x_i^2
	51	+$$
	52	+
	53	+##### 步骤4：分布式归约操作（All-Reduce）
	54	+
	55	+将步骤3中的平方和 $s_x$ 在所有tp_y通信域进程中进行归约（求和），确保每个进程都拥有其通信域全局平方和：
	56	+
	57	+$$
	58	+s_x^{\text{global}} = \text{AllReduce}\left( s_x \right) = \sum_{p=1}^{P} \sum_{i=1}^{H} \left( x_i^{(p)} \right)^2
	59	+$$
	60	+
	61	+##### 步骤5：中心化输入数据
	62	+
	63	+将输入数据 $\mathbf{x}$ 中心化，即减去平均值。平均值 $\mu$ 计算如下：
	64	+
	65	+$$
	66	+\mu = \frac{e_x^{\text{global}}}{H}
	67	+$$
	68	+
	69	+然后，中心化输入：
	70	+
	71	+$$
	72	+x'_i = x_i - \mu \quad \forall i \in \{1, 2, \dots, H\}
	73	+$$
	74	+
	75	+##### 步骤6：计算总和的平方
	76	+
	77	+计算全局总和的平方：
	78	+
	79	+$$
	80	+e_x'^2 = \left( e_x^{\text{global}} \right)^2
	81	+$$
	82	+
	83	+##### 步骤7：计算归一化因子
	84	+
	85	+计算归一化因子 $\gamma$，用于标准化输入数据。公式如下：
	86	+
	87	+$$
	88	+\gamma = \frac{1}{\sqrt{ \left( \frac{s_x^{\text{global}}}{H} \right) - e_x'^2 + \epsilon }}
	89	+$$
	90	+
	91	+这里：
	92	+- $\frac{s_x^{\text{global}}}{H}$ 是全局平方和的平均值。
	93	+- $e_x'^2$ 是全局总和的平方。
	94	+- $\epsilon$ 是一个小常数，防止分母为零，增加数值稳定性。
	95	+
	96	+##### 步骤8：标准化输入数据
	97	+
	98	+将中心化后的输入数据 $\mathbf{x}'$ 与归一化因子 $\gamma$ 相乘，得到标准化后的数据 $\mathbf{\hat{x}}$：
	99	+
	100	+$$
	101	+\hat{x}_i = x'_i \cdot \gamma \quad \forall i \in \{1, 2, \dots, H\}
	102	+$$
	103	+
	104	+##### 步骤9：应用权重和偏置
	105	+
	106	+最后，将标准化后的数据与权重向量 $\mathbf{W}$ 相乘，并根据是否存在偏置向量 $\mathbf{b}$ 来决定最终输出。
	107	+
	108	+- 如果存在偏置：
	109	+
	110	+$$
	111	+\text{output}_i = b_i + W_i \cdot \hat{x}_i \quad \forall i \in \{1, 2, \dots, H\}
	112	+$$
	113	+
	114	+- 如果不存在偏置：
	115	+
	116	+$$
	117	+\text{output}_i = W_i \cdot \hat{x}_i \quad \forall i \in \{1, 2, \dots, H\}
	118	+$$
	119	+
	120	+
	121	+## 使用场景
	122	+
	123	+当TP通信域需要设置较大时，通信效率较低，需要通过分解通信域来提升其通信效率。
	124	+
	125	+## 使用方法
	126	+
	127	+在训练脚本的参数列表中加入 `--tp-2d`，开启2D张量并行，`--tp-x N1`和`--tp-y N2`分别设置其x轴、y轴的切分大小，其中需满足`tp = N1 * N2`。
	128	+
	129	+## 使用效果
	130	+
	131	+在llama3-405B模型训练时，tp=16情况下，开启2D张量并行，tp_x=8，tp_y=2，相比原Megatron 1D张量并行性能提升5%+。

Mmindspeed/arguments.py+33-0

	@@ -38,6 +38,7 @@ def process_args(parser):
38	parser = _add_automated_pipeline_args(parser)	38	parser = _add_automated_pipeline_args(parser)
39	parser = _add_alibi_args(parser)	39	parser = _add_alibi_args(parser)
40	parser = _add_ndmm_args(parser)	40	parser = _add_ndmm_args(parser)
		41	+ parser = _add_2d_tp_args(parser)
41	parser = _add_coc_args(parser)	42	parser = _add_coc_args(parser)
42	parser = _add_profile_args(parser)	43	parser = _add_profile_args(parser)
43	parser = _add_auto_parallel_args(parser)	44	parser = _add_auto_parallel_args(parser)
	@@ -694,6 +695,23 @@ def validate_args_wrapper(validate_args):
694	args.use_pipe_experts = False	695	args.use_pipe_experts = False
695	args.pipe_experts_multi_stream = False	696	args.pipe_experts_multi_stream = False
696	args.pipe_experts_multi_data = 1	697	args.pipe_experts_multi_data = 1
		698	+ if args.tp_2d:
		699	+ if args.sequence_parallel:
		700	+ raise AssertionError('2d tp does not support sequence parallel')
		701	+ if args.use_fused_rmsnorm:
		702	+ raise AssertionError('2d tp does not support fused rmsnorm')
		703	+ if args.use_nanopipe:
		704	+ raise AssertionError('tp-2d does not support nano-pipe')
		705	+ if args.ampipe_degree > 1:
		706	+ raise AssertionError('tp-2d does not support ampipe')
		707	+ if args.context_parallel_algo not in ['megatron_cp_algo', 'ulysses_cp_algo']:
		708	+ raise AssertionError('tp-2d now only support megatron_cp_algo or ulysses_cp_algo')
		709	+ if args.use_ascend_coc:
		710	+ raise AssertionError('tp-2d does not support ascend coc')
		711	+ if args.tensor_model_parallel_size // args.tp_x != args.tp_y:
		712	+ raise AssertionError('need satisfy tp = tp_x * tp_y')
		713	+ if args.expert_model_parallel_size > 1:
		714	+ raise AssertionError('2d tp does not support moe')
697		715
698	from megatron.training.arguments import _print_args	716	from megatron.training.arguments import _print_args
699	_print_args('arguments', args, True)	717	_print_args('arguments', args, True)
	@@ -761,3 +779,18 @@ def _add_auto_parallel_args(parser):
761	group.add_argument('--profile-memory', action='store_true', help='')	779	group.add_argument('--profile-memory', action='store_true', help='')
762	group.add_argument('--prof-file', type=str, default=None, help='')	780	group.add_argument('--prof-file', type=str, default=None, help='')
763	return parser	781	return parser
		782	+
		783	+
		784	+def _add_2d_tp_args(parser):
		785	+ group = parser.add_argument_group(title='2d-tp')
		786	+ group.add_argument('--tp-2d', action='store_true', default=False,
		787	+ help='use use-2d-tp to replace megatron-style tensor parallel')
		788	+ group.add_argument('--enable-overlap-ag-with-matmul', action='store_true', default=False,
		789	+ help='use enable-overlap-ag-with-matmul to overlap allgather with matmul')
		790	+ group.add_argument('--enable-overlap-matmul-with-rs', action='store_true', default=False,
		791	+ help='use enable-overlap-matmul-with-rs to overlap matmul with reduce-scatter')
		792	+ group.add_argument('--tp-x', type=int, default=1,
		793	+ help='the fist dim tensor parallel size for Linear')
		794	+ group.add_argument('--tp-y', type=int, default=1,
		795	+ help='the second dim tensor parallel size for Linear')
		796	+ return parser

Amindspeed/core/mindspeed_parallel_group.py+75-0

@@ -0,0 +1,75 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+import abc
	16	+
	17	+import torch
	18	+import torch.distributed
	19	+from mindspeed.core.simple_parallel_cfg import SimpleParallelCfg
	20	+
	21	+
	22	+class MindspeedParallelGroup:
	23	+ def __init__(
	24	+ self,
	25	+ parallel_cfg: SimpleParallelCfg = None,
	26	+ pg_name: str = None,
	27	+ overlap_gp_name: str = None,
	28	+ nccl_comm_cfgs=None,
	29	+ ):
	30	+ """Parallel group interface, any type of parallelism class can implement this class.
	31	+
	32	+ :param parallel_cfg: Parallel configuration.
	33	+ :param pg_name: parallel process group name.
	34	+ :param overlap_gp_name: overlap process name, for the send/rcv parallel.
	35	+ :param nccl_comm_cfgs:
	36	+ """
	37	+ self._pg_name = pg_name
	38	+ self._overlap_pg_name = overlap_gp_name
	39	+ self._group, self._global_ranks, self._overlap_group = self.init_group(
	40	+ parallel_cfg, pg_name, overlap_gp_name, nccl_comm_cfgs
	41	+ )
	42	+
	43	+ @staticmethod
	44	+ @abc.abstractmethod
	45	+ def init_group(
	46	+ parallel_cfg: SimpleParallelCfg,
	47	+ pg_name: str,
	48	+ overlap_gp_name: str = None,
	49	+ nccl_comm_cfgs=None,
	50	+ ):
	51	+ raise NotImplementedError
	52	+
	53	+ @property
	54	+ def group(self):
	55	+ return self._group
	56	+
	57	+ @property
	58	+ def overlap_group(self):
	59	+ return self._overlap_group
	60	+
	61	+ @property
	62	+ def global_ranks(self):
	63	+ return self._global_ranks
	64	+
	65	+ def get_parallel_rank(self):
	66	+ if torch.distributed.is_available() and torch.distributed.is_initialized():
	67	+ return torch.distributed.get_rank(group=self.group)
	68	+ else:
	69	+ raise AssertionError("The distribution is not available or not initialized.")
	70	+
	71	+ def get_parallel_group_world_size(self):
	72	+ if torch.distributed.is_available() and torch.distributed.is_initialized():
	73	+ return torch.distributed.get_world_size(group=self.group)
	74	+ else:
	75	+ return 0

Amindspeed/core/models/common/embeddings/language_model_embedding.py+34-0

@@ -0,0 +1,34 @@
	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 functools import wraps
	17	+from megatron.training import get_args
	18	+from mindspeed.core.tensor_parallel.comm_autograd_function import (
	19	+ auto_grad_scatter_along_first_dim_then_last_dim,
	20	+)
	21	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm
	22	+from mindspeed.core.tensor_parallel.comm_group_api import TPYCollectiveComm
	23	+
	24	+
	25	+def language_model_embedding_forward_wrapper(forward):
	26	+ @wraps(forward)
	27	+ def wrapper(self, args, *kwargs):
	28	+ encoder_input = forward(self, args, *kwargs)
	29	+ if get_args().tp_2d:
	30	+ encoder_input = auto_grad_scatter_along_first_dim_then_last_dim(
	31	+ encoder_input, TPXCollectiveComm, TPYCollectiveComm
	32	+ )
	33	+ return encoder_input
	34	+ return wrapper

Mmindspeed/core/models/common/embeddings/rotary_pos_embedding.py+107-1

	@@ -16,6 +16,7 @@ from mindspeed.core.parallel_state import (get_context_parallel_for_hybrid_ulyss
16	get_context_parallel_for_hybrid_ring_world_size,	16	get_context_parallel_for_hybrid_ring_world_size,
17	get_context_parallel_for_hybrid_ring_rank)	17	get_context_parallel_for_hybrid_ring_rank)
18	from mindspeed.core.context_parallel.utils import get_remapped_seq_order	18	from mindspeed.core.context_parallel.utils import get_remapped_seq_order
		19	+from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP
19		20
20		21
21	def yarn_find_correction_dim(	22	def yarn_find_correction_dim(
	@@ -156,13 +157,19 @@ def rotary_forward(self, max_seq_len: int, offset: int = 0) -> Tensor:
156	def get_pos_emb_on_this_cp_rank(pos_emb, seq_dim):	157	def get_pos_emb_on_this_cp_rank(pos_emb, seq_dim):
157	args = get_args()	158	args = get_args()
158		159
		160	+ cp_expanded_by_2d_tp = args.tp_y > 1
159	if args.context_parallel_algo == 'megatron_cp_algo':	161	if args.context_parallel_algo == 'megatron_cp_algo':
160	if args.cp_attention_mask_type == 'general':	162	if args.cp_attention_mask_type == 'general':
161	pos_emb = _get_pos_emb_on_this_cp_rank_in_ulysses_cp(pos_emb, seq_dim)	163	pos_emb = _get_pos_emb_on_this_cp_rank_in_ulysses_cp(pos_emb, seq_dim)
		164	+ elif cp_expanded_by_2d_tp:
		165	+ pos_emb = _get_pos_emb_on_this_tp_y_cp_rank_in_megatron_cp(pos_emb, seq_dim)
162	else:	166	else:
163	pos_emb = _get_pos_emb_on_this_cp_rank_in_megatron_cp(pos_emb, seq_dim)	167	pos_emb = _get_pos_emb_on_this_cp_rank_in_megatron_cp(pos_emb, seq_dim)
164	elif args.context_parallel_algo == 'ulysses_cp_algo':	168	elif args.context_parallel_algo == 'ulysses_cp_algo':
165	- pos_emb = _get_pos_emb_on_this_cp_rank_in_ulysses_cp(pos_emb, seq_dim)	169	+ if cp_expanded_by_2d_tp:
		170	+ pos_emb = _get_pos_emb_on_this_tp_y_cp_rank_in_ulysses_cp(pos_emb, seq_dim)
		171	+ else:
		172	+ pos_emb = _get_pos_emb_on_this_cp_rank_in_ulysses_cp(pos_emb, seq_dim)
166	elif args.context_parallel_algo == 'hybrid_cp_algo':	173	elif args.context_parallel_algo == 'hybrid_cp_algo':
167	if args.cp_attention_mask_type == 'general':	174	if args.cp_attention_mask_type == 'general':
168	pos_emb = _get_pos_emb_on_this_cp_rank_in_hybrid_cp_general(pos_emb, seq_dim)	175	pos_emb = _get_pos_emb_on_this_cp_rank_in_hybrid_cp_general(pos_emb, seq_dim)
	@@ -189,6 +196,37 @@ def _get_pos_emb_on_this_cp_rank_in_megatron_cp(pos_emb, seq_dim):
189	return pos_emb	196	return pos_emb
190		197
191		198
		199	+def _get_pos_emb_on_this_tp_y_cp_rank_in_megatron_cp(pos_emb, seq_dim):
		200	+ origin_pos_emb_shape = pos_emb.shape
		201	+ tp_y_cp_group = TensorParallelYUnionCP()
		202	+ tp_y_cp_size = tp_y_cp_group.get_parallel_group_world_size()
		203	+ # [s, 1, 1, head_dim] ---> [2tp_y_cp_size, s/(2tp_y_cp_size), 1, 1, head_dim]
		204	+ pos_emb = pos_emb.view(
		205	+ pos_emb.shape[:seq_dim], 2 tp_y_cp_size, -1, *pos_emb.shape[(seq_dim + 1) :]
		206	+ )
		207	+ rearrange_index = []
		208	+ for i in range(tp_y_cp_size):
		209	+ rearrange_index.extend([i, 2 * tp_y_cp_size - 1 - i])
		210	+
		211	+ rearrange_idx_tensor = torch.tensor(rearrange_index, device=pos_emb.device)
		212	+
		213	+ # Reorder pos embedding according dataset handling.
		214	+ # selected res shape: [2 * tp_y_cp_size, s / (2 * tp_y_cp_size), 1, 1, head_dim]
		215	+ pos_emb = pos_emb.index_select(seq_dim, index=rearrange_idx_tensor)
		216	+ pos_emb = pos_emb.view(*origin_pos_emb_shape)
		217	+ # viewed res shape: [tp_y_cp_sz, s/tp_y_cp_sz, 1, head_dim]
		218	+ pos_emb = pos_emb.view(
		219	+ *pos_emb.shape[0:seq_dim],
		220	+ tp_y_cp_size,
		221	+ pos_emb.shape[seq_dim] // tp_y_cp_size,
		222	+ *pos_emb.shape[(seq_dim + 1):],
		223	+ )
		224	+ # cur_rank_pos_emb shape: [s/cp, 1, 1, head_dim]
		225	+ tp_y_cp_rank = tp_y_cp_group.get_parallel_rank()
		226	+ cur_rank_pos_emb = pos_emb[tp_y_cp_rank].squeeze(axis=0)
		227	+ return cur_rank_pos_emb
		228	+
		229	+
192	def _get_pos_emb_on_this_cp_rank_in_ulysses_cp(pos_emb, seq_dim):	230	def _get_pos_emb_on_this_cp_rank_in_ulysses_cp(pos_emb, seq_dim):
193	cp_size = parallel_state.get_context_parallel_world_size()	231	cp_size = parallel_state.get_context_parallel_world_size()
194	cp_rank = parallel_state.get_context_parallel_rank()	232	cp_rank = parallel_state.get_context_parallel_rank()
	@@ -259,3 +297,71 @@ def _get_pos_emb_on_this_cp_rank_in_hybrid_adaptive_cp(pos_emd, seq_dim):
259		297
260	return pos_emd	298	return pos_emd
261		299
		300	+
		301	+def rotary_embedding_forward(self, max_seq_len: int, offset: int = 0) -> Tensor:
		302	+ """Forward pass of RoPE embedding.
		303	+
		304	+ Args:
		305	+ max_seq_len (int): Maximum size of sequence
		306	+ offset (int, optional): _description_. Defaults to 0.
		307	+
		308	+ Returns:
		309	+ Tensor: Embeddings after applying RoPE.
		310	+ """
		311	+ seq = (
		312	+ torch.arange(max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
		313	+ + offset
		314	+ )
		315	+
		316	+ if self.seq_len_interpolation_factor is not None:
		317	+ seq *= 1 / self.seq_len_interpolation_factor
		318	+
		319	+ freqs = torch.outer(seq, self.inv_freq)
		320	+ # first part even vector components, second part odd vector components,
		321	+ # 2 * dim in dimension size
		322	+ if not self.rotary_interleaved:
		323	+ emb = torch.cat((freqs, freqs), dim=-1)
		324	+ else:
		325	+ emb = torch.stack((freqs.view(-1, 1), freqs.view(-1, 1)), dim=-1).view(
		326	+ freqs.shape[0], -1
		327	+ )
		328	+ # emb [seq_length, .., dim]
		329	+ emb = emb[:, None, None, :]
		330	+ global_args = get_args()
		331	+ cp = global_args.context_parallel_size
		332	+ if global_args.tp_2d:
		333	+ tp_y_cp_sz = cp * global_args.tp_y
		334	+ else:
		335	+ tp_y_cp_sz = cp
		336	+ if tp_y_cp_sz > 1:
		337	+ # slice rotary_pos_emb along sequence dimension and select the parition of the current CP rank
		338	+ emb = get_pos_emb_on_this_cp_rank(emb, 0)
		339	+ return emb
		340	+
		341	+
		342	+def rotary_embedding_forward_wrapper(fn):
		343	+ @wraps(fn)
		344	+ def wrapper(self, max_seq_len: int, offset: int = 0):
		345	+ return rotary_embedding_forward(self, max_seq_len, offset)
		346	+
		347	+ return wrapper
		348	+
		349	+
		350	+def _get_pos_emb_on_this_tp_y_cp_rank_in_ulysses_cp(pos_emb, seq_dim):
		351	+ tp_y_cp_group = TensorParallelYUnionCP()
		352	+ tp_y_cp_size = tp_y_cp_group.get_parallel_group_world_size()
		353	+
		354	+ cp_rank = tp_y_cp_group.get_parallel_rank()
		355	+ pos_emb = pos_emb.chunk(tp_y_cp_size, dim=seq_dim)[cp_rank]
		356	+ return pos_emb
		357	+
		358	+
		359	+def rotary_embedding_get_rotary_seq_len_wrapper(fn):
		360	+ @wraps(fn)
		361	+ def wrapper(self, inference_params, transformer, transformer_input, transformer_config,):
		362	+ rotary_seq_len = fn(self, inference_params, transformer, transformer_input, transformer_config,)
		363	+ global_args = get_args()
		364	+ if global_args.tp_2d:
		365	+ rotary_seq_len *= global_args.tp_x
		366	+ return rotary_seq_len
		367	+ return wrapper

Mmindspeed/core/parallel_state.py+168-9

	@@ -22,6 +22,9 @@ from datetime import timedelta
22	import torch	22	import torch
23	import megatron	23	import megatron
24		24
		25	+from mindspeed.core.simple_parallel_cfg import SimpleParallelCfg
		26	+from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP
		27	+
25	_CONTEXT_PARALLEL_GROUP_FOR_SEND_RECV_OVERLAP = None	28	_CONTEXT_PARALLEL_GROUP_FOR_SEND_RECV_OVERLAP = None
26	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_ULYSSES = None	29	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_ULYSSES = None
27	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_RING = None	30	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_RING = None
	@@ -36,7 +39,8 @@ _CONTEXT_PARALLEL_RANKS_FOR_RING_INTER_WINDOW_DKV = None
36	_CONTEXT_PARALLEL_GROUP_FOR_RING_INTRA_WINDOW = None	39	_CONTEXT_PARALLEL_GROUP_FOR_RING_INTRA_WINDOW = None
37	_CONTEXT_PARALLEL_GROUP_FOR_RING_INTRA_WINDOW_SEND_RECV_OVERLAP = None	40	_CONTEXT_PARALLEL_GROUP_FOR_RING_INTRA_WINDOW_SEND_RECV_OVERLAP = None
38		41
39	-	42	+_TP_X_PARALLEL_RING_RANKS = None
		43	+_TP_Y_PARALLEL_RING_RANKS = None
40		44
41	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1 = None	45	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1 = None
42	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2 = None	46	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2 = None
	@@ -46,6 +50,12 @@ _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND1_DIM1 = None
46	_TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND1_DIM2 = None	50	_TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND1_DIM2 = None
47	_TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM1 = None	51	_TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM1 = None
48	_TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM2 = None	52	_TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM2 = None
		53	+_TP_X_SD_RCV_OVERLAP_GROUP = None
		54	+_TP_Y_SD_RCV_OVERLAP_GROUP = None
		55	+_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_RANK = None
		56	+_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_RANK = None
		57	+_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_WORLD_SIZE = None
		58	+_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_WORLD_SIZE = None
49		59
50	_TENSOR_AND_CONTEXT_PARALLEL_GROUP = None	60	_TENSOR_AND_CONTEXT_PARALLEL_GROUP = None
51	_TENSOR_AND_CONTEXT_PARALLEL_GLOBAL_RANKS = None	61	_TENSOR_AND_CONTEXT_PARALLEL_GLOBAL_RANKS = None
	@@ -227,14 +237,32 @@ def initialize_model_parallel_wrapper(initialize_model_parallel):
227	if rank in ranks:	237	if rank in ranks:
228	_PIPELINE_MODEL_PARALLEL_GROUP_FOR_NEW_STREAM = group	238	_PIPELINE_MODEL_PARALLEL_GROUP_FOR_NEW_STREAM = group
229		239
230	-	240	+ from megatron.training import get_args
		241	+ args = get_args()
		242	+ nd1_dim1_sz = args.nd1_dim1_size if args.use_nd_matmul else args.tp_x
		243	+ nd2_dim1_sz = args.nd2_dim1_size if args.use_nd_matmul else args.tp_y
231	initialize_ndmm_parallel_group(	244	initialize_ndmm_parallel_group(
232	nccl_comm_cfgs,	245	nccl_comm_cfgs,
233	tensor_model_parallel_size=tensor_model_parallel_size,	246	tensor_model_parallel_size=tensor_model_parallel_size,
234	- nd1_dim1_size=args.nd1_dim1_size,	247	+ nd1_dim1_size=nd1_dim1_sz,
235	- nd2_dim1_size=args.nd2_dim1_size,	248	+ nd2_dim1_size=nd2_dim1_sz,
236	)	249	)
237	-	250	+ if args.tp_2d:
		251	+ tp_y_cp_group = TensorParallelYUnionCP(
		252	+ parallel_cfg=SimpleParallelCfg(
		253	+ dp=data_parallel_size,
		254	+ pp=pipeline_model_parallel_size,
		255	+ tp=tensor_model_parallel_size,
		256	+ cp=context_parallel_size,
		257	+ ep=expert_model_parallel_size,
		258	+ tp_x=get_args().tp_x,
		259	+ tp_y=get_args().tp_y,
		260	+ ),
		261	+ pg_name="tp-y-cp",
		262	+ overlap_gp_name="tp-y-cp-overlap",
		263	+ nccl_comm_cfgs=nccl_comm_cfgs
		264	+ )
		265	+ print(f'tp_y_cp_group.global_ranks={tp_y_cp_group.global_ranks} for rank {rank}')
238	return wrapper	266	return wrapper
239		267
240		268
	@@ -280,7 +308,8 @@ def initialize_context_parallel_group_for_send_recv_overlap(
280	from megatron.training import get_args	308	from megatron.training import get_args
281	if not get_args().use_cp_send_recv_overlap:	309	if not get_args().use_cp_send_recv_overlap:
282	return	310	return
283	-	311	+ if get_args().tp_2d:
		312	+ return
284	rank = torch.distributed.get_rank()	313	rank = torch.distributed.get_rank()
285	world_size: int = torch.distributed.get_world_size()	314	world_size: int = torch.distributed.get_world_size()
286	num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size	315	num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size
	@@ -378,6 +407,8 @@ def initialize_context_parallel_group_for_double_ring(
378	):	407	):
379	from megatron.training import get_args	408	from megatron.training import get_args
380	args = get_args()	409	args = get_args()
		410	+ if args.tp_2d:
		411	+ return
381	if context_parallel_size == 1 or args.context_parallel_algo not in ['megatron_cp_algo', 'hybrid_cp_algo']:	412	if context_parallel_size == 1 or args.context_parallel_algo not in ['megatron_cp_algo', 'hybrid_cp_algo']:
382	return	413	return
383		414
	@@ -917,6 +948,18 @@ def get_context_parallel_for_hybrid_ring_global_ranks():
917	return _CONTEXT_PARALLEL_RANKS_FOR_HYBRID_RING	948	return _CONTEXT_PARALLEL_RANKS_FOR_HYBRID_RING
918		949
919		950
		951	+def get_tp_x_ring_global_ranks():
		952	+ global _TP_X_PARALLEL_RING_RANKS
		953	+ assert (_TP_X_PARALLEL_RING_RANKS is not None), 'TP-X parallel group for ring is not initialized'
		954	+ return _TP_X_PARALLEL_RING_RANKS
		955	+
		956	+
		957	+def get_tp_y_ring_global_ranks():
		958	+ global _TP_Y_PARALLEL_RING_RANKS
		959	+ assert (_TP_Y_PARALLEL_RING_RANKS is not None), 'TP-Y parallel group for ring is not initialized'
		960	+ return _TP_Y_PARALLEL_RING_RANKS
		961	+
		962	+
920	def destroy_model_parallel_wrapper(destroy_model_parallel):	963	def destroy_model_parallel_wrapper(destroy_model_parallel):
921	@wraps(destroy_model_parallel)	964	@wraps(destroy_model_parallel)
922	def wrapper():	965	def wrapper():
	@@ -928,16 +971,48 @@ def destroy_model_parallel_wrapper(destroy_model_parallel):
928	global _CONTEXT_PARALLEL_GROUP_FOR_HYBRID_ULYSSES	971	global _CONTEXT_PARALLEL_GROUP_FOR_HYBRID_ULYSSES
929	global _CONTEXT_PARALLEL_RANKS_FOR_HYBRID_RING	972	global _CONTEXT_PARALLEL_RANKS_FOR_HYBRID_RING
930	global _CONTEXT_PARALLEL_RANKS_FOR_HYBRID_ULYSSES	973	global _CONTEXT_PARALLEL_RANKS_FOR_HYBRID_ULYSSES
		974	+ global _TP_X_PARALLEL_RING_RANKS
		975	+ global _TP_Y_PARALLEL_RING_RANKS
		976	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1
		977	+ global _TP_X_SD_RCV_OVERLAP_GROUP
		978	+ global _TP_Y_SD_RCV_OVERLAP_GROUP
		979	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2
		980	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_RANK
		981	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_RANK
		982	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_WORLD_SIZE
		983	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_WORLD_SIZE
		984	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM1
		985	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM2
		986	+ global _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND1_DIM1
		987	+ global _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND1_DIM2
		988	+ global _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM1
		989	+ global _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM2
		990	+ global _TENSOR_AND_CONTEXT_PARALLEL_GROUP
		991	+ global _TENSOR_AND_CONTEXT_PARALLEL_GLOBAL_RANKS
931	_CONTEXT_PARALLEL_GROUP_FOR_SEND_RECV_OVERLAP = None	992	_CONTEXT_PARALLEL_GROUP_FOR_SEND_RECV_OVERLAP = None
932	_PIPELINE_MODEL_PARALLEL_GROUP_FOR_NEW_STREAM = None	993	_PIPELINE_MODEL_PARALLEL_GROUP_FOR_NEW_STREAM = None
933	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_RING = None	994	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_RING = None
934	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_ULYSSES = None	995	_CONTEXT_PARALLEL_GROUP_FOR_HYBRID_ULYSSES = None
935	_CONTEXT_PARALLEL_RANKS_FOR_HYBRID_RING = None	996	_CONTEXT_PARALLEL_RANKS_FOR_HYBRID_RING = None
936	_CONTEXT_PARALLEL_RANKS_FOR_HYBRID_ULYSSES = None	997	_CONTEXT_PARALLEL_RANKS_FOR_HYBRID_ULYSSES = None
937	- global _TENSOR_AND_CONTEXT_PARALLEL_GROUP
938	- global _TENSOR_AND_CONTEXT_PARALLEL_GLOBAL_RANKS
939	_TENSOR_AND_CONTEXT_PARALLEL_GROUP = None	998	_TENSOR_AND_CONTEXT_PARALLEL_GROUP = None
940	_TENSOR_AND_CONTEXT_PARALLEL_GLOBAL_RANKS = None	999	_TENSOR_AND_CONTEXT_PARALLEL_GLOBAL_RANKS = None
		1000	+ _TP_X_PARALLEL_RING_RANKS = None
		1001	+ _TP_Y_PARALLEL_RING_RANKS = None
		1002	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1 = None
		1003	+ _TP_X_SD_RCV_OVERLAP_GROUP = None
		1004	+ _TP_Y_SD_RCV_OVERLAP_GROUP = None
		1005	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2 = None
		1006	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_RANK = None
		1007	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_RANK = None
		1008	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_WORLD_SIZE = None
		1009	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_WORLD_SIZE = None
		1010	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM1 = None
		1011	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM2 = None
		1012	+ _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND1_DIM1 = None
		1013	+ _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND1_DIM2 = None
		1014	+ _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM1 = None
		1015	+ _TENSOR_MODEL_PARALLEL_WORLD_SIZE_FOR_ND2_DIM2 = None
941		1016
942	global _GROBAL_PROCESS_GROUP_GLOO	1017	global _GROBAL_PROCESS_GROUP_GLOO
943	_GROBAL_PROCESS_GROUP_GLOO = None	1018	_GROBAL_PROCESS_GROUP_GLOO = None
	@@ -951,6 +1026,18 @@ def get_tensor_model_parallel_group_for_nd1_dim1(check_initialized=True):
951	return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1	1026	return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1
952		1027
953		1028
		1029	+def get_tp_x_sd_rcv_overlap_group(check_initialized=True):
		1030	+ if check_initialized and _TP_X_SD_RCV_OVERLAP_GROUP is None:
		1031	+ raise AssertionError('tp-x send recv overlap group is not initialized')
		1032	+ return _TP_X_SD_RCV_OVERLAP_GROUP
		1033	+
		1034	+
		1035	+def get_tp_y_sd_rcv_overlap_group(check_initialized=True):
		1036	+ if check_initialized and _TP_Y_SD_RCV_OVERLAP_GROUP is None:
		1037	+ raise AssertionError('tp-y send recv overlap group is not initialized')
		1038	+ return _TP_Y_SD_RCV_OVERLAP_GROUP
		1039	+
		1040	+
954	def get_tensor_model_parallel_group_for_nd1_dim2(check_initialized=True):	1041	def get_tensor_model_parallel_group_for_nd1_dim2(check_initialized=True):
955	if check_initialized and _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2 is None:	1042	if check_initialized and _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2 is None:
956	raise AssertionError('tensor model parallel group for nd1 dim2 is not initialized')	1043	raise AssertionError('tensor model parallel group for nd1 dim2 is not initialized')
	@@ -963,6 +1050,42 @@ def get_tensor_model_parallel_group_for_nd2_dim1(check_initialized=True):
963	return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM1	1050	return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM1
964		1051
965		1052
		1053	+def get_tensor_model_parallel_group_for_nd1_dim1_rank():
		1054	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_RANK
		1055	+ if _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_RANK is None:
		1056	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_RANK = torch.distributed.get_rank(
		1057	+ group=get_tensor_model_parallel_group_for_nd1_dim1())
		1058	+
		1059	+ return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_RANK
		1060	+
		1061	+
		1062	+def get_tensor_model_parallel_group_for_nd1_dim2_rank():
		1063	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_RANK
		1064	+ if _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_RANK is None:
		1065	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_RANK = torch.distributed.get_rank(
		1066	+ group=get_tensor_model_parallel_group_for_nd1_dim2())
		1067	+
		1068	+ return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_RANK
		1069	+
		1070	+
		1071	+def get_tensor_model_parallel_group_for_nd1_dim1_world_size():
		1072	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_WORLD_SIZE
		1073	+ if _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_WORLD_SIZE is None:
		1074	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_WORLD_SIZE = torch.distributed.get_world_size(
		1075	+ group=get_tensor_model_parallel_group_for_nd1_dim1())
		1076	+
		1077	+ return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1_WORLD_SIZE
		1078	+
		1079	+
		1080	+def get_tensor_model_parallel_group_for_nd1_dim2_world_size():
		1081	+ global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_WORLD_SIZE
		1082	+ if _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_WORLD_SIZE is None:
		1083	+ _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_WORLD_SIZE = torch.distributed.get_world_size(
		1084	+ group=get_tensor_model_parallel_group_for_nd1_dim2())
		1085	+
		1086	+ return _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2_WORLD_SIZE
		1087	+
		1088	+
966	def get_tensor_model_parallel_group_for_nd2_dim2(check_initialized=True):	1089	def get_tensor_model_parallel_group_for_nd2_dim2(check_initialized=True):
967	if check_initialized and _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM2 is None:	1090	if check_initialized and _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM2 is None:
968	raise AssertionError('tensor model parallel group for nd2 dim2 is not initialized')	1091	raise AssertionError('tensor model parallel group for nd2 dim2 is not initialized')
	@@ -1016,7 +1139,7 @@ def initialize_ndmm_parallel_group(
1016	from megatron.training.global_vars import _ensure_var_is_not_initialized	1139	from megatron.training.global_vars import _ensure_var_is_not_initialized
1017		1140
1018	args = get_args()	1141	args = get_args()
1019	- if not args.use_nd_matmul:	1142	+ if not (args.use_nd_matmul or args.tp_2d):
1020	return	1143	return
1021		1144
1022	global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1	1145	global _TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1
	@@ -1039,6 +1162,18 @@ def initialize_ndmm_parallel_group(
1039	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM2, 'nd2_dim2'	1162	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND2_DIM2, 'nd2_dim2'
1040	)	1163	)
1041		1164
		1165	+ global _TP_X_PARALLEL_RING_RANKS
		1166	+ _ensure_var_is_not_initialized(_TP_X_PARALLEL_RING_RANKS, 'tp_x_ring_ranks')
		1167	+
		1168	+ global _TP_Y_PARALLEL_RING_RANKS
		1169	+ _ensure_var_is_not_initialized(_TP_Y_PARALLEL_RING_RANKS, 'tp_y_ring_ranks')
		1170	+
		1171	+ global _TP_X_SD_RCV_OVERLAP_GROUP
		1172	+ _ensure_var_is_not_initialized(_TP_X_SD_RCV_OVERLAP_GROUP, 'tp_x_overlap_ranks')
		1173	+
		1174	+ global _TP_Y_SD_RCV_OVERLAP_GROUP
		1175	+ _ensure_var_is_not_initialized(_TP_Y_SD_RCV_OVERLAP_GROUP, 'tp_y_overlap_ranks')
		1176	+
1042	if tensor_model_parallel_size % nd1_dim1_size != 0:	1177	if tensor_model_parallel_size % nd1_dim1_size != 0:
1043	raise RuntimeError(	1178	raise RuntimeError(
1044	f"tensor_model_parallel_size can't divisible by nd1_dim1_size"	1179	f"tensor_model_parallel_size can't divisible by nd1_dim1_size"
	@@ -1053,17 +1188,30 @@ def initialize_ndmm_parallel_group(
1053	world_size: int = torch.distributed.get_world_size()	1188	world_size: int = torch.distributed.get_world_size()
1054	num_tensor_model_parallel_group: int = world_size // tensor_model_parallel_size	1189	num_tensor_model_parallel_group: int = world_size // tensor_model_parallel_size
1055		1190
		1191	+ tp_nd1_dim1_groups = []
		1192	+ tp_nd1_dim2_groups = []
		1193	+ tp_nd2_dim1_groups = []
		1194	+ tp_nd2_dim2_groups = []
1056	for i in range(num_tensor_model_parallel_group):	1195	for i in range(num_tensor_model_parallel_group):
1057	for j in range(tensor_model_parallel_size // nd1_dim1_size):	1196	for j in range(tensor_model_parallel_size // nd1_dim1_size):
1058	ranks = range(	1197	ranks = range(
1059	i * tensor_model_parallel_size + j * nd1_dim1_size,	1198	i * tensor_model_parallel_size + j * nd1_dim1_size,
1060	i * tensor_model_parallel_size + (j + 1) * nd1_dim1_size	1199	i * tensor_model_parallel_size + (j + 1) * nd1_dim1_size
1061	)	1200	)
		1201	+ tp_nd1_dim1_groups.append(list(ranks))
1062	group = torch.distributed.new_group(	1202	group = torch.distributed.new_group(
1063	ranks, pg_options=ps.get_nccl_options('nd1_dim1', nccl_comm_cfgs)	1203	ranks, pg_options=ps.get_nccl_options('nd1_dim1', nccl_comm_cfgs)
1064	)	1204	)
		1205	+ if args.enable_overlap_ag_with_matmul:
		1206	+ tp_x_ag_overlap_group = torch.distributed.new_group(
		1207	+ ranks, pg_options=ps.get_nccl_options('ag_x_sd_rcv_overlap', nccl_comm_cfgs)
		1208	+ )
		1209	+ else:
		1210	+ tp_x_ag_overlap_group = None
1065	if rank in ranks:	1211	if rank in ranks:
1066	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1 = group	1212	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM1 = group
		1213	+ _TP_X_SD_RCV_OVERLAP_GROUP = tp_x_ag_overlap_group
		1214	+ _TP_X_PARALLEL_RING_RANKS = ranks
1067		1215
1068	nd1_dim2_size = tensor_model_parallel_size // nd1_dim1_size	1216	nd1_dim2_size = tensor_model_parallel_size // nd1_dim1_size
1069	for j in range(tensor_model_parallel_size // nd1_dim2_size):	1217	for j in range(tensor_model_parallel_size // nd1_dim2_size):
	@@ -1072,17 +1220,27 @@ def initialize_ndmm_parallel_group(
1072	(i + 1) * tensor_model_parallel_size,	1220	(i + 1) * tensor_model_parallel_size,
1073	nd1_dim1_size	1221	nd1_dim1_size
1074	)	1222	)
		1223	+ tp_nd1_dim2_groups.append(list(ranks))
1075	group = torch.distributed.new_group(	1224	group = torch.distributed.new_group(
1076	ranks, pg_options=ps.get_nccl_options('nd1_dim2', nccl_comm_cfgs)	1225	ranks, pg_options=ps.get_nccl_options('nd1_dim2', nccl_comm_cfgs)
1077	)	1226	)
		1227	+ if args.enable_overlap_ag_with_matmul:
		1228	+ tp_y_ag_overlap_group = torch.distributed.new_group(
		1229	+ ranks, pg_options=ps.get_nccl_options('ag_y_sd_rcv_overlap', nccl_comm_cfgs)
		1230	+ )
		1231	+ else:
		1232	+ tp_y_ag_overlap_group = None
1078	if rank in ranks:	1233	if rank in ranks:
1079	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2 = group	1234	_TENSOR_MODEL_PARALLEL_GROUP_FOR_ND1_DIM2 = group
		1235	+ _TP_Y_SD_RCV_OVERLAP_GROUP = tp_y_ag_overlap_group
		1236	+ _TP_Y_PARALLEL_RING_RANKS = ranks
1080		1237
1081	for j in range(tensor_model_parallel_size // nd2_dim1_size):	1238	for j in range(tensor_model_parallel_size // nd2_dim1_size):
1082	ranks = range(	1239	ranks = range(
1083	i * tensor_model_parallel_size + j * nd2_dim1_size,	1240	i * tensor_model_parallel_size + j * nd2_dim1_size,
1084	i * tensor_model_parallel_size + (j + 1) * nd2_dim1_size	1241	i * tensor_model_parallel_size + (j + 1) * nd2_dim1_size
1085	)	1242	)
		1243	+ tp_nd2_dim1_groups.append(list(ranks))
1086	group = torch.distributed.new_group(	1244	group = torch.distributed.new_group(
1087	ranks, pg_options=ps.get_nccl_options('nd2_dim1', nccl_comm_cfgs)	1245	ranks, pg_options=ps.get_nccl_options('nd2_dim1', nccl_comm_cfgs)
1088	)	1246	)
	@@ -1096,6 +1254,7 @@ def initialize_ndmm_parallel_group(
1096	(i + 1) * tensor_model_parallel_size,	1254	(i + 1) * tensor_model_parallel_size,
1097	nd2_dim1_size	1255	nd2_dim1_size
1098	)	1256	)
		1257	+ tp_nd2_dim2_groups.append(list(ranks))
1099	group = torch.distributed.new_group(	1258	group = torch.distributed.new_group(
1100	ranks, pg_options=ps.get_nccl_options('nd2_dim2', nccl_comm_cfgs)	1259	ranks, pg_options=ps.get_nccl_options('nd2_dim2', nccl_comm_cfgs)
1101	)	1260	)

Mmindspeed/core/pipeline_parallel/flexible_schedules.py+612-2

	@@ -636,7 +636,8 @@ def forward_backward_pipelining_with_interleaving_nano_pipe(
636	tensor_shape[0] = tensor_shape[0] // parallel_state.get_context_parallel_world_size()	636	tensor_shape[0] = tensor_shape[0] // parallel_state.get_context_parallel_world_size()
637	if config.sequence_parallel:	637	if config.sequence_parallel:
638	tensor_shape[0] = tensor_shape[0] // parallel_state.get_tensor_model_parallel_world_size()	638	tensor_shape[0] = tensor_shape[0] // parallel_state.get_tensor_model_parallel_world_size()
639	-	639	+ tensor_shape[0] = tensor_shape[0] // args.tp_x
		640	+ tensor_shape[-1] = tensor_shape[-1] // args.tp_y
640	# Compute number of warmup and remaining microbatches.	641	# Compute number of warmup and remaining microbatches.
641	num_model_chunks = len(model)	642	num_model_chunks = len(model)
642	total_num_microbatches = num_microbatches * num_model_chunks	643	total_num_microbatches = num_microbatches * num_model_chunks
	@@ -1148,4 +1149,613 @@ def forward_backward_pipelining_with_interleaving_nano_pipe(
1148	# embedding all-reduce for pipeline parallelism).	1149	# embedding all-reduce for pipeline parallelism).
1149	config.finalize_model_grads_func(model)	1150	config.finalize_model_grads_func(model)
1150		1151
1151	- return forward_data_store	1152	+ return forward_data_store
		1153	+
		1154	+
		1155	+def forward_backward_pipelining_with_interleaving_patch(
		1156	+ *,
		1157	+ forward_step_func,
		1158	+ data_iterator: Union[Iterator, List[Iterator]],
		1159	+ model: Union[torch.nn.Module, List[torch.nn.Module]],
		1160	+ num_microbatches: int,
		1161	+ seq_length: int,
		1162	+ micro_batch_size: int,
		1163	+ decoder_seq_length: int = None,
		1164	+ forward_only: bool = False,
		1165	+ collect_non_loss_data: bool = False,
		1166	+ first_val_step: bool = None,
		1167	+):
		1168	+ """Run interleaved 1F1B schedule (model split into model chunks), with
		1169	+ communication between pipeline stages as needed.
		1170	+
		1171	+ Returns dictionary with losses if the last stage, empty dict otherwise."""
		1172	+ if not isinstance(model, list):
		1173	+ raise AssertionError("interleaved pipeline parallelism expected model chunking")
		1174	+ if not all(isinstance(chunk, torch.nn.Module) for chunk in model):
		1175	+ raise AssertionError("invalid model chunking")
		1176	+ if not isinstance(data_iterator, list):
		1177	+ raise AssertionError("interleaved pipeline parallelism expected each model chunk to have a data iterator")
		1178	+ config = get_model_config(model[0])
		1179	+ if config.overlap_p2p_comm and config.batch_p2p_comm:
		1180	+ raise ValueError("Can not use both overlap_p2p_comm and batch_p2p_comm")
		1181	+
		1182	+ if config.timers is not None:
		1183	+ config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)
		1184	+
		1185	+ # Disable async grad reductions
		1186	+ no_sync_func = config.no_sync_func
		1187	+ if isinstance(no_sync_func, list):
		1188	+
		1189	+ def multi_no_sync():
		1190	+ stack = contextlib.ExitStack()
		1191	+ for model_chunk_no_sync_func in config.no_sync_func:
		1192	+ stack.enter_context(model_chunk_no_sync_func())
		1193	+ return stack
		1194	+
		1195	+ no_sync_func = multi_no_sync
		1196	+ if no_sync_func is None:
		1197	+ no_sync_func = contextlib.nullcontext
		1198	+ no_sync_context = None
		1199	+
		1200	+ if config.grad_sync_func is not None and not isinstance(config.grad_sync_func, list):
		1201	+ config.grad_sync_func = [config.grad_sync_func for _ in model]
		1202	+
		1203	+ if config.param_sync_func is not None and not isinstance(config.param_sync_func, list):
		1204	+ config.param_sync_func = [config.param_sync_func for _ in model]
		1205	+
		1206	+ def disable_grad_sync():
		1207	+ """Disable asynchronous grad reductions"""
		1208	+ nonlocal no_sync_context
		1209	+ if no_sync_context is None:
		1210	+ no_sync_context = no_sync_func()
		1211	+ no_sync_context.__enter__()
		1212	+
		1213	+ def enable_grad_sync():
		1214	+ """Enable asynchronous grad reductions"""
		1215	+ nonlocal no_sync_context
		1216	+ if no_sync_context is not None:
		1217	+ no_sync_context.__exit__(None, None, None)
		1218	+ no_sync_context = None
		1219	+
		1220	+ disable_grad_sync()
		1221	+
		1222	+ # Model chunk IDs with synchronized grads
		1223	+ synchronized_model_chunks = set()
		1224	+
		1225	+ input_tensors = [[] for _ in range(len(model))]
		1226	+ output_tensors = [[] for _ in range(len(model))]
		1227	+ total_num_tokens = torch.tensor(0, dtype=torch.int).cuda()
		1228	+
		1229	+ forward_data_store = []
		1230	+ if not forward_only:
		1231	+ output_tensor_grads = [[] for _ in range(len(model))]
		1232	+
		1233	+ pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
		1234	+ pipeline_parallel_rank = parallel_state.get_pipeline_model_parallel_rank()
		1235	+
		1236	+ if num_microbatches % pipeline_parallel_size != 0:
		1237	+ msg = f'number of microbatches ({num_microbatches}) is not divisible by '
		1238	+ msg += f'pipeline-model-parallel-size ({pipeline_parallel_size}) '
		1239	+ msg += 'when using interleaved schedule'
		1240	+ raise RuntimeError(msg)
		1241	+
		1242	+ model_type = get_model_type(model[0])
		1243	+ if model_type == ModelType.encoder_and_decoder:
		1244	+ raise RuntimeError("Interleaving is not supported with an encoder and decoder model.")
		1245	+
		1246	+ if decoder_seq_length is not None and decoder_seq_length != seq_length:
		1247	+ raise RuntimeError(
		1248	+ "Interleaving is not supported with a different decoder sequence length."
		1249	+ )
		1250	+
		1251	+ tensor_shape = [seq_length, micro_batch_size, config.hidden_size]
		1252	+ tensor_shape[0] = tensor_shape[0] // parallel_state.get_context_parallel_world_size()
		1253	+ if config.sequence_parallel:
		1254	+ tensor_shape[0] = tensor_shape[0] // parallel_state.get_tensor_model_parallel_world_size()
		1255	+ tensor_shape[0] = tensor_shape[0] // get_args().tp_x
		1256	+ tensor_shape[-1] = tensor_shape[-1] // get_args().tp_y
		1257	+ # Compute number of warmup and remaining microbatches.
		1258	+ num_model_chunks = len(model)
		1259	+ total_num_microbatches = num_microbatches * num_model_chunks
		1260	+ all_warmup_microbatches = False
		1261	+ if forward_only:
		1262	+ num_warmup_microbatches = total_num_microbatches
		1263	+ else:
		1264	+ # Run all forward passes and then all backward passes if number of
		1265	+ # microbatches is just the number of pipeline stages.
		1266	+ # Otherwise, perform (num_model_chunks-1)*pipeline_parallel_size on
		1267	+ # all workers, followed by more microbatches after depending on
		1268	+ # stage ID (more forward passes for earlier stages, later stages can
		1269	+ # immediately start with 1F1B).
		1270	+ if num_microbatches == pipeline_parallel_size:
		1271	+ num_warmup_microbatches = total_num_microbatches
		1272	+ all_warmup_microbatches = True
		1273	+ else:
		1274	+ num_warmup_microbatches = (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
		1275	+ num_warmup_microbatches += (num_model_chunks - 1) * pipeline_parallel_size
		1276	+ num_warmup_microbatches = min(num_warmup_microbatches, total_num_microbatches)
		1277	+ num_microbatches_remaining = total_num_microbatches - num_warmup_microbatches
		1278	+
		1279	+ # Checkpoint the activations of partial Transformer layers in a number of micro-batches
		1280	+ # within the maximum outstanding micro-batch backpropagations.
		1281	+ # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
		1282	+ # checkpoint partial Transformer layers (or skip checkpointing) and
		1283	+ # the rest of micro-batches within a window of micro-batches checkpoint
		1284	+ # all Transformer layers. The window of micro-batches is set by the maximum
		1285	+ # outstanding backpropagations and becomes smaller at later pipeline stages.
		1286	+ # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
		1287	+ max_outstanding_backprops = None
		1288	+ if config.num_microbatches_with_partial_activation_checkpoints is not None:
		1289	+ max_outstanding_backprops = num_warmup_microbatches + 1
		1290	+
		1291	+ # Synchronize params for first two model chunks
		1292	+ if config.param_sync_func is not None:
		1293	+ config.param_sync_func[0](model[0].parameters())
		1294	+ config.param_sync_func[1](model[1].parameters())
		1295	+
		1296	+ def get_model_chunk_id(microbatch_id, forward):
		1297	+ """Helper method to get the model chunk ID given the iteration number."""
		1298	+ microbatch_id_in_group = microbatch_id % (pipeline_parallel_size * num_model_chunks)
		1299	+ model_chunk_id = microbatch_id_in_group // pipeline_parallel_size
		1300	+ if not forward:
		1301	+ model_chunk_id = num_model_chunks - model_chunk_id - 1
		1302	+ return model_chunk_id
		1303	+
		1304	+ def get_microbatch_id_in_model_chunk(iteration_id, forward):
		1305	+ """Helper method to get the microbatch_id within model chunk given the iteration number."""
		1306	+ assert forward
		1307	+ iteration_group_id = iteration_id // (pipeline_parallel_size * num_model_chunks)
		1308	+ microbatch_id_in_model_chunk = (iteration_group_id * pipeline_parallel_size) + (
		1309	+ iteration_id % pipeline_parallel_size
		1310	+ )
		1311	+ return microbatch_id_in_model_chunk
		1312	+
		1313	+ def is_first_microbatch_for_model_chunk(microbatch_id: int) -> bool:
		1314	+ """Check if an iteration is the first for a model chunk."""
		1315	+ microbatch_group_size = pipeline_parallel_size * num_model_chunks
		1316	+ num_microbatch_groups = total_num_microbatches // microbatch_group_size
		1317	+ microbatch_group_id = microbatch_id // microbatch_group_size
		1318	+ microbatch_id_in_group = microbatch_id % microbatch_group_size
		1319	+ if microbatch_group_id == 0:
		1320	+ return microbatch_id_in_group % pipeline_parallel_size == 0
		1321	+ else:
		1322	+ return False
		1323	+
		1324	+ def is_last_microbatch_for_model_chunk(microbatch_id: int) -> bool:
		1325	+ """Check if an iteration is the last for a model chunk."""
		1326	+ microbatch_group_size = pipeline_parallel_size * num_model_chunks
		1327	+ num_microbatch_groups = total_num_microbatches // microbatch_group_size
		1328	+ microbatch_group_id = microbatch_id // microbatch_group_size
		1329	+ microbatch_id_in_group = microbatch_id % microbatch_group_size
		1330	+ if microbatch_group_id == num_microbatch_groups - 1:
		1331	+ return microbatch_id_in_group % pipeline_parallel_size == pipeline_parallel_size - 1
		1332	+ else:
		1333	+ return False
		1334	+
		1335	+ def forward_step_helper(microbatch_id, current_microbatch, checkpoint_activations_microbatch):
		1336	+ """Helper method to run forward step with model split into chunks
		1337	+ (run set_virtual_pipeline_model_parallel_rank() before calling
		1338	+ forward_step())."""
		1339	+ model_chunk_id = get_model_chunk_id(microbatch_id, forward=True)
		1340	+ parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id)
		1341	+
		1342	+ # launch param synchronization for next model chunk
		1343	+ # Note: Asynchronous communication tends to slow down compute.
		1344	+ # To reduce idling from mismatched microbatch times, we launch
		1345	+ # asynchronous communication at the same time across the
		1346	+ # pipeline-parallel group.
		1347	+ if config.param_sync_func is not None:
		1348	+ param_sync_microbatch_id = microbatch_id + pipeline_parallel_rank
		1349	+ if (
		1350	+ param_sync_microbatch_id < total_num_microbatches
		1351	+ and is_first_microbatch_for_model_chunk(param_sync_microbatch_id)
		1352	+ ):
		1353	+ param_sync_chunk_id = get_model_chunk_id(param_sync_microbatch_id, forward=True) + 1
		1354	+ if 1 < param_sync_chunk_id < num_model_chunks:
		1355	+ config.param_sync_func[param_sync_chunk_id](
		1356	+ model[param_sync_chunk_id].parameters()
		1357	+ )
		1358	+
		1359	+ # forward step
		1360	+ if parallel_state.is_pipeline_first_stage():
		1361	+ if len(input_tensors[model_chunk_id]) == len(output_tensors[model_chunk_id]):
		1362	+ input_tensors[model_chunk_id].append(None)
		1363	+ input_tensor = input_tensors[model_chunk_id][-1]
		1364	+
		1365	+ output_tensor, num_tokens = forward_step(
		1366	+ forward_step_func,
		1367	+ data_iterator[model_chunk_id],
		1368	+ model[model_chunk_id],
		1369	+ num_microbatches,
		1370	+ input_tensor,
		1371	+ forward_data_store,
		1372	+ config,
		1373	+ collect_non_loss_data,
		1374	+ checkpoint_activations_microbatch,
		1375	+ check_first_val_step(
		1376	+ first_val_step, forward_only, is_first_microbatch_for_model_chunk(microbatch_id),
		1377	+ ),
		1378	+ current_microbatch=current_microbatch,
		1379	+ )
		1380	+ output_tensors[model_chunk_id].append(output_tensor)
		1381	+
		1382	+ nonlocal total_num_tokens
		1383	+ total_num_tokens += num_tokens.item()
		1384	+
		1385	+ # if forward-only, no need to save tensors for a backward pass
		1386	+ if forward_only:
		1387	+ input_tensors[model_chunk_id].pop()
		1388	+ output_tensors[model_chunk_id].pop()
		1389	+
		1390	+ return output_tensor
		1391	+
		1392	+ def backward_step_helper(microbatch_id):
		1393	+ """Helper method to run backward step with model split into chunks
		1394	+ (run set_virtual_pipeline_model_parallel_rank() before calling
		1395	+ backward_step())."""
		1396	+ model_chunk_id = get_model_chunk_id(microbatch_id, forward=False)
		1397	+ parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id)
		1398	+
		1399	+ # launch grad synchronization (default)
		1400	+ if config.grad_sync_func is None and is_last_microbatch_for_model_chunk(microbatch_id):
		1401	+ enable_grad_sync()
		1402	+ synchronized_model_chunks.add(model_chunk_id)
		1403	+
		1404	+ if parallel_state.is_pipeline_last_stage():
		1405	+ if len(output_tensor_grads[model_chunk_id]) == 0:
		1406	+ output_tensor_grads[model_chunk_id].append(None)
		1407	+ input_tensor = input_tensors[model_chunk_id].pop(0)
		1408	+ output_tensor = output_tensors[model_chunk_id].pop(0)
		1409	+ output_tensor_grad = output_tensor_grads[model_chunk_id].pop(0)
		1410	+ input_tensor_grad = backward_step(
		1411	+ input_tensor, output_tensor, output_tensor_grad, model_type, config
		1412	+ )
		1413	+
		1414	+ # launch grad synchronization (custom grad sync)
		1415	+ # Note: Asynchronous communication tends to slow down compute.
		1416	+ # To reduce idling from mismatched microbatch times, we launch
		1417	+ # asynchronous communication at the same time across the
		1418	+ # pipeline-parallel group.
		1419	+ if config.grad_sync_func is not None:
		1420	+ grad_sync_microbatch_id = microbatch_id - pipeline_parallel_rank
		1421	+ if grad_sync_microbatch_id >= 0 and is_last_microbatch_for_model_chunk(
		1422	+ grad_sync_microbatch_id
		1423	+ ):
		1424	+ grad_sync_chunk_id = get_model_chunk_id(grad_sync_microbatch_id, forward=False)
		1425	+ enable_grad_sync()
		1426	+ config.grad_sync_func[grad_sync_chunk_id](model[grad_sync_chunk_id].parameters())
		1427	+ synchronized_model_chunks.add(grad_sync_chunk_id)
		1428	+ disable_grad_sync()
		1429	+
		1430	+ return input_tensor_grad
		1431	+
		1432	+ # Run warmup forward passes.
		1433	+ parallel_state.set_virtual_pipeline_model_parallel_rank(0)
		1434	+ input_tensors[0].append(p2p_communication.recv_forward(tensor_shape, config))
		1435	+
		1436	+ fwd_wait_handles = None
		1437	+ bwd_wait_handles = None
		1438	+
		1439	+ for k in range(num_warmup_microbatches):
		1440	+
		1441	+ if fwd_wait_handles is not None:
		1442	+ for req in fwd_wait_handles:
		1443	+ req.wait()
		1444	+
		1445	+ cur_model_chunk_id = get_model_chunk_id(k, forward=True)
		1446	+ # Decide to checkpoint all layers' activations of the current micro-batch
		1447	+ if max_outstanding_backprops is not None:
		1448	+ checkpoint_activations_microbatch = (
		1449	+ k % max_outstanding_backprops
		1450	+ >= config.num_microbatches_with_partial_activation_checkpoints
		1451	+ )
		1452	+ else:
		1453	+ checkpoint_activations_microbatch = None
		1454	+
		1455	+ current_microbatch = get_microbatch_id_in_model_chunk(k, forward=True)
		1456	+ output_tensor = forward_step_helper(
		1457	+ k, current_microbatch, checkpoint_activations_microbatch
		1458	+ )
		1459	+
		1460	+ # Determine if tensor should be received from previous stage.
		1461	+ next_forward_model_chunk_id = get_model_chunk_id(k + 1, forward=True)
		1462	+ recv_prev = True
		1463	+ if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
		1464	+ if next_forward_model_chunk_id == 0:
		1465	+ recv_prev = False
		1466	+ if k == (total_num_microbatches - 1):
		1467	+ recv_prev = False
		1468	+
		1469	+ # Don't send tensor downstream if on last stage.
		1470	+ if parallel_state.is_pipeline_last_stage():
		1471	+ output_tensor = None
		1472	+
		1473	+ # Send and receive tensors as appropriate (send tensors computed
		1474	+ # in this iteration; receive tensors for next iteration).
		1475	+ if not config.overlap_p2p_comm:
		1476	+ if (
		1477	+ k == (num_warmup_microbatches - 1)
		1478	+ and not forward_only
		1479	+ and not all_warmup_microbatches
		1480	+ ):
		1481	+ input_tensor_grad = None
		1482	+ recv_next = True
		1483	+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
		1484	+ recv_next = False
		1485	+ (
		1486	+ input_tensor,
		1487	+ output_tensor_grad,
		1488	+ ) = p2p_communication.send_forward_backward_recv_forward_backward(
		1489	+ output_tensor,
		1490	+ input_tensor_grad,
		1491	+ recv_prev=recv_prev,
		1492	+ recv_next=recv_next,
		1493	+ tensor_shape=tensor_shape,
		1494	+ config=config,
		1495	+ )
		1496	+ output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad)
		1497	+ else:
		1498	+ input_tensor = p2p_communication.send_forward_recv_forward(
		1499	+ output_tensor, recv_prev=recv_prev, tensor_shape=tensor_shape, config=config
		1500	+ )
		1501	+ input_tensors[next_forward_model_chunk_id].append(input_tensor)
		1502	+ else:
		1503	+ input_tensor, fwd_wait_handles = p2p_communication.send_forward_recv_forward(
		1504	+ output_tensor,
		1505	+ recv_prev=recv_prev,
		1506	+ tensor_shape=tensor_shape,
		1507	+ config=config,
		1508	+ overlap_p2p_comm=True,
		1509	+ )
		1510	+
		1511	+ if (
		1512	+ k == (num_warmup_microbatches - 1)
		1513	+ and not forward_only
		1514	+ and not all_warmup_microbatches
		1515	+ ):
		1516	+ input_tensor_grad = None
		1517	+ recv_next = True
		1518	+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
		1519	+ recv_next = False
		1520	+
		1521	+ (
		1522	+ output_tensor_grad,
		1523	+ bwd_wait_handles,
		1524	+ ) = p2p_communication.send_backward_recv_backward(
		1525	+ input_tensor_grad,
		1526	+ recv_next=recv_next,
		1527	+ tensor_shape=tensor_shape,
		1528	+ config=config,
		1529	+ overlap_p2p_comm=True,
		1530	+ )
		1531	+
		1532	+ output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad)
		1533	+ input_tensors[next_forward_model_chunk_id].append(input_tensor)
		1534	+
		1535	+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
		1536	+
		1537	+ # Run 1F1B in steady state.
		1538	+ for k in range(num_microbatches_remaining):
		1539	+ # Forward pass.
		1540	+ forward_k = k + num_warmup_microbatches
		1541	+
		1542	+ # Decide to checkpoint all layers' activations of the current micro-batch
		1543	+ if max_outstanding_backprops is not None:
		1544	+ checkpoint_activations_microbatch = (
		1545	+ forward_k % max_outstanding_backprops
		1546	+ >= config.num_microbatches_with_partial_activation_checkpoints
		1547	+ )
		1548	+ else:
		1549	+ checkpoint_activations_microbatch = None
		1550	+
		1551	+ cur_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
		1552	+ current_microbatch = get_microbatch_id_in_model_chunk(forward_k, forward=True)
		1553	+ if config.overlap_p2p_comm:
		1554	+ if fwd_wait_handles is not None:
		1555	+ for req in fwd_wait_handles:
		1556	+ req.wait()
		1557	+
		1558	+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
		1559	+
		1560	+ output_tensor = forward_step_helper(
		1561	+ forward_k, current_microbatch, checkpoint_activations_microbatch
		1562	+ )
		1563	+
		1564	+ # Determine if current stage has anything to send in either direction,
		1565	+ # otherwise set tensor to None.
		1566	+ forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
		1567	+ parallel_state.set_virtual_pipeline_model_parallel_rank(forward_model_chunk_id)
		1568	+
		1569	+ # Last virtual stage no activation tensor to send
		1570	+ if parallel_state.is_pipeline_last_stage():
		1571	+ output_tensor = None
		1572	+
		1573	+ # Determine if peers are sending, and where in data structure to put
		1574	+ # received tensors.
		1575	+ recv_prev = True
		1576	+ if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
		1577	+ # First stage is ahead of last stage by (pipeline_parallel_size - 1).
		1578	+ next_forward_model_chunk_id = get_model_chunk_id(
		1579	+ forward_k - (pipeline_parallel_size - 1), forward=True
		1580	+ )
		1581	+ if next_forward_model_chunk_id == (num_model_chunks - 1):
		1582	+ recv_prev = False
		1583	+ next_forward_model_chunk_id += 1
		1584	+ else:
		1585	+ next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1, forward=True)
		1586	+
		1587	+ # If last iteration, don't receive; we already received one extra
		1588	+ # before the start of the for loop.
		1589	+ if k == (num_microbatches_remaining - 1):
		1590	+ recv_prev = False
		1591	+
		1592	+ # Send activation tensor to the next stage and receive activation tensor from the
		1593	+ # previous stage
		1594	+ input_tensor, fwd_wait_handles = p2p_communication.send_forward_recv_forward(
		1595	+ output_tensor,
		1596	+ recv_prev=recv_prev,
		1597	+ tensor_shape=tensor_shape,
		1598	+ config=config,
		1599	+ overlap_p2p_comm=True,
		1600	+ )
		1601	+ # assert fwd_wait_handles is not None
		1602	+
		1603	+ if bwd_wait_handles is not None:
		1604	+ for req in bwd_wait_handles:
		1605	+ req.wait()
		1606	+
		1607	+ # Backward pass.
		1608	+ backward_k = k
		1609	+ input_tensor_grad = backward_step_helper(backward_k)
		1610	+
		1611	+ backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
		1612	+ parallel_state.set_virtual_pipeline_model_parallel_rank(backward_model_chunk_id)
		1613	+
		1614	+ # First virtual stage no activation gradient tensor to send
		1615	+ if parallel_state.is_pipeline_first_stage():
		1616	+ input_tensor_grad = None
		1617	+
		1618	+ # Determine if the current virtual stage has an activation gradient tensor to receive
		1619	+ recv_next = True
		1620	+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
		1621	+ # Last stage is ahead of first stage by (pipeline_parallel_size - 1).
		1622	+ next_backward_model_chunk_id = get_model_chunk_id(
		1623	+ backward_k - (pipeline_parallel_size - 1), forward=False
		1624	+ )
		1625	+ if next_backward_model_chunk_id == 0:
		1626	+ recv_next = False
		1627	+ next_backward_model_chunk_id -= 1
		1628	+ else:
		1629	+ next_backward_model_chunk_id = get_model_chunk_id(backward_k + 1, forward=False)
		1630	+
		1631	+ output_tensor_grad, bwd_wait_handles = p2p_communication.send_backward_recv_backward(
		1632	+ input_tensor_grad,
		1633	+ recv_next=recv_next,
		1634	+ tensor_shape=tensor_shape,
		1635	+ config=config,
		1636	+ overlap_p2p_comm=True,
		1637	+ )
		1638	+
		1639	+ else: # no p2p overlap
		1640	+ output_tensor = forward_step_helper(
		1641	+ forward_k, current_microbatch, checkpoint_activations_microbatch
		1642	+ )
		1643	+
		1644	+ # Backward pass.
		1645	+ backward_k = k
		1646	+ input_tensor_grad = backward_step_helper(backward_k)
		1647	+
		1648	+ # Send output_tensor and input_tensor_grad, receive input_tensor
		1649	+ # and output_tensor_grad.
		1650	+
		1651	+ # Determine if current stage has anything to send in either direction,
		1652	+ # otherwise set tensor to None.
		1653	+ forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
		1654	+ parallel_state.set_virtual_pipeline_model_parallel_rank(forward_model_chunk_id)
		1655	+ if parallel_state.is_pipeline_last_stage():
		1656	+ output_tensor = None
		1657	+
		1658	+ backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
		1659	+ parallel_state.set_virtual_pipeline_model_parallel_rank(backward_model_chunk_id)
		1660	+ if parallel_state.is_pipeline_first_stage():
		1661	+ input_tensor_grad = None
		1662	+
		1663	+ # Determine if peers are sending, and where in data structure to put
		1664	+ # received tensors.
		1665	+ recv_prev = True
		1666	+ if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
		1667	+ # First stage is ahead of last stage by (pipeline_parallel_size - 1).
		1668	+ next_forward_model_chunk_id = get_model_chunk_id(
		1669	+ forward_k - (pipeline_parallel_size - 1), forward=True
		1670	+ )
		1671	+ if next_forward_model_chunk_id == (num_model_chunks - 1):
		1672	+ recv_prev = False
		1673	+ next_forward_model_chunk_id += 1
		1674	+ else:
		1675	+ next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1, forward=True)
		1676	+
		1677	+ recv_next = True
		1678	+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
		1679	+ # Last stage is ahead of first stage by (pipeline_parallel_size - 1).
		1680	+ next_backward_model_chunk_id = get_model_chunk_id(
		1681	+ backward_k - (pipeline_parallel_size - 1), forward=False
		1682	+ )
		1683	+ if next_backward_model_chunk_id == 0:
		1684	+ recv_next = False
		1685	+ next_backward_model_chunk_id -= 1
		1686	+ else:
		1687	+ next_backward_model_chunk_id = get_model_chunk_id(backward_k + 1, forward=False)
		1688	+
		1689	+ # If last iteration, don't receive; we already received one extra
		1690	+ # before the start of the for loop.
		1691	+ if k == (num_microbatches_remaining - 1):
		1692	+ recv_prev = False
		1693	+
		1694	+ # Communicate tensors.
		1695	+ (
		1696	+ input_tensor,
		1697	+ output_tensor_grad,
		1698	+ ) = p2p_communication.send_forward_backward_recv_forward_backward(
		1699	+ output_tensor,
		1700	+ input_tensor_grad,
		1701	+ recv_prev=recv_prev,
		1702	+ recv_next=recv_next,
		1703	+ tensor_shape=tensor_shape,
		1704	+ config=config,
		1705	+ )
		1706	+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
		1707	+
		1708	+ # Put input_tensor and output_tensor_grad in data structures in the
		1709	+ # right location.
		1710	+ if recv_prev:
		1711	+ input_tensors[next_forward_model_chunk_id].append(input_tensor)
		1712	+ if recv_next:
		1713	+ output_tensor_grads[next_backward_model_chunk_id].append(output_tensor_grad)
		1714	+
		1715	+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
		1716	+
		1717	+ # Run cooldown backward passes (flush out pipeline).
		1718	+ if not forward_only:
		1719	+ if config.overlap_p2p_comm and bwd_wait_handles is not None:
		1720	+ for wait_handle in bwd_wait_handles:
		1721	+ wait_handle.wait()
		1722	+
		1723	+ if all_warmup_microbatches:
		1724	+ output_tensor_grads[num_model_chunks - 1].append(
		1725	+ p2p_communication.recv_backward(tensor_shape, config=config)
		1726	+ )
		1727	+ for k in range(num_microbatches_remaining, total_num_microbatches):
		1728	+ input_tensor_grad = backward_step_helper(k)
		1729	+ next_backward_model_chunk_id = get_model_chunk_id(k + 1, forward=False)
		1730	+ recv_next = True
		1731	+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
		1732	+ if next_backward_model_chunk_id == (num_model_chunks - 1):
		1733	+ recv_next = False
		1734	+ if k == (total_num_microbatches - 1):
		1735	+ recv_next = False
		1736	+ output_tensor_grads[next_backward_model_chunk_id].append(
		1737	+ p2p_communication.send_backward_recv_backward(
		1738	+ input_tensor_grad, recv_next=recv_next, tensor_shape=tensor_shape, config=config
		1739	+ )
		1740	+ )
		1741	+
		1742	+ # Launch any remaining grad reductions.
		1743	+ enable_grad_sync()
		1744	+ if config.grad_sync_func is not None:
		1745	+ for model_chunk_id in range(num_model_chunks):
		1746	+ if model_chunk_id not in synchronized_model_chunks:
		1747	+ config.grad_sync_func[model_chunk_id](model[model_chunk_id].parameters())
		1748	+ synchronized_model_chunks.add(model_chunk_id)
		1749	+
		1750	+ if config.finalize_model_grads_func is not None and not forward_only:
		1751	+ # Finalize model grads (perform full grad all-reduce / reduce-scatter for
		1752	+ # data parallelism, layernorm all-reduce for sequence parallelism, and
		1753	+ # embedding all-reduce for pipeline parallelism).
		1754	+ config.finalize_model_grads_func(
		1755	+ model, total_num_tokens if config.calculate_per_token_loss else None
		1756	+ )
		1757	+
		1758	+ if config.timers is not None:
		1759	+ config.timers('forward-backward').stop()
		1760	+
		1761	+ return forward_data_store

Mmindspeed/core/pipeline_parallel/ripipe_schedules.py+2-1

	@@ -139,7 +139,8 @@ def forward_backward_ripipe_pipelining(
139	tensor_shape[0] = tensor_shape[0] // parallel_state.get_context_parallel_world_size()	139	tensor_shape[0] = tensor_shape[0] // parallel_state.get_context_parallel_world_size()
140	if config.sequence_parallel:	140	if config.sequence_parallel:
141	tensor_shape[0] = tensor_shape[0] // parallel_state.get_tensor_model_parallel_world_size()	141	tensor_shape[0] = tensor_shape[0] // parallel_state.get_tensor_model_parallel_world_size()
142	-	142	+ tensor_shape[0] = tensor_shape[0] // args.tp_x
		143	+ tensor_shape[-1] = tensor_shape[-1] // args.tp_y
143	# Compute number of warmup and remaining microbatches.	144	# Compute number of warmup and remaining microbatches.
144	num_model_chunks = len(model)	145	num_model_chunks = len(model)
145	total_num_microbatches = num_microbatches * num_model_chunks	146	total_num_microbatches = num_microbatches * num_model_chunks

Mmindspeed/core/pipeline_parallel/schedules.py+14-0

	@@ -153,3 +153,17 @@ def forward_step(
153	if unwrap_output_tensor:	153	if unwrap_output_tensor:
154	return output_tensor, num_tokens	154	return output_tensor, num_tokens
155	return [output_tensor], num_tokens	155	return [output_tensor], num_tokens
		156	+
		157	+
		158	+def get_tensor_shapes_wrapper(get_tensor_shapes):
		159	+ @wraps(get_tensor_shapes)
		160	+ def wrapper(args, *kwargs):
		161	+ # [s, b, h]
		162	+ tensor_shapes = get_tensor_shapes(args, *kwargs)
		163	+ arguments = get_args()
		164	+ if arguments.tp_2d:
		165	+ tensor_shapes = [[tensor_shape[0] // arguments.tp_x, tensor_shape[1], tensor_shape[2] // arguments.tp_y]
		166	+ for tensor_shape in tensor_shapes]
		167	+
		168	+ return tensor_shapes
		169	+ return wrapper

Amindspeed/core/simple_parallel_cfg.py+19-0

@@ -0,0 +1,19 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+from collections import namedtuple
	16	+
	17	+SimpleParallelCfg = namedtuple(
	18	+ "SimpleParallelCfg", field_names=["dp", "pp", "tp", "cp", "ep", "tp_x", "tp_y"]
	19	+)

Amindspeed/core/singleton_meta.py+27-0

@@ -0,0 +1,27 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+import threading
	16	+
	17	+
	18	+class SingletonMeta(type):
	19	+ _instances = {}
	20	+ _lock = threading.Lock()
	21	+
	22	+ def __call__(cls, args, *kwargs):
	23	+ with cls._lock:
	24	+ if cls not in cls._instances:
	25	+ instance = super().__call__(args, *kwargs)
	26	+ cls._instances[cls] = instance
	27	+ return cls._instances[cls]

Amindspeed/core/tensor_parallel/comm_autograd_function.py+204-0

@@ -0,0 +1,204 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+import torch
	16	+import torch.distributed
	17	+
	18	+from mindspeed.core.tensor_parallel.comm_group_api import CollectiveCommIntf
	19	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm
	20	+from mindspeed.core.tensor_parallel.comm_group_api import TPYCollectiveComm
	21	+from mindspeed.core.tensor_parallel.comm_utils import _gather_along_last_dim
	22	+from mindspeed.core.tensor_parallel.comm_utils import _split_along_first_dim
	23	+from mindspeed.core.tensor_parallel.comm_utils import _split_along_last_dim
	24	+from mindspeed.core.tensor_parallel.comm_utils import sync_gather_along_first_dim
	25	+from mindspeed.core.tensor_parallel.comm_utils import sync_gather_along_last_dim
	26	+
	27	+
	28	+class _SyncGatherAlongFirstDim(torch.autograd.Function):
	29	+ """Gather the input from model parallel X region and concatinate."""
	30	+
	31	+ @staticmethod
	32	+ def symbolic(graph, input_):
	33	+ return sync_gather_along_first_dim(input_, TPXCollectiveComm)
	34	+
	35	+ @staticmethod
	36	+ def forward(ctx, input_, comm_intf: CollectiveCommIntf):
	37	+ ctx.comm_intf = comm_intf
	38	+ return sync_gather_along_first_dim(input_, comm_intf, "_SyncGatherAlongFirstDim")
	39	+
	40	+ @staticmethod
	41	+ def backward(ctx, grad_output):
	42	+ return _split_along_first_dim(grad_output, ctx.comm_intf), None
	43	+
	44	+
	45	+class _SyncGatherAlongLastDim(torch.autograd.Function):
	46	+ """Gather the input from model parallel Y region and concatinate."""
	47	+
	48	+ @staticmethod
	49	+ def symbolic(graph, input_):
	50	+ return sync_gather_along_last_dim(input_, TPYCollectiveComm)
	51	+
	52	+ @staticmethod
	53	+ def forward(ctx, input_, comm_intf: CollectiveCommIntf):
	54	+ ctx.comm_intf = comm_intf
	55	+ return sync_gather_along_last_dim(input_, comm_intf)
	56	+
	57	+ @staticmethod
	58	+ def backward(ctx, grad_output):
	59	+ return _split_along_last_dim(grad_output, ctx.comm_intf), None
	60	+
	61	+
	62	+def _reduce(input_, tp_intf: CollectiveCommIntf = TPXCollectiveComm):
	63	+ """All-reduce the input tensor across model parallel group."""
	64	+
	65	+ # Bypass the function if we are using only 1 GPU.
	66	+ if tp_intf.get_comm_group_world_size() == 1:
	67	+ return input_
	68	+
	69	+ # All-reduce.
	70	+ torch.distributed.all_reduce(input_, group=tp_intf.get_comm_group())
	71	+ return input_
	72	+
	73	+
	74	+class _ReduceFromModelParallelRegion(torch.autograd.Function):
	75	+ """All-reduce the input from the model parallel region."""
	76	+
	77	+ @staticmethod
	78	+ def symbolic(graph, input_, tp_intf: CollectiveCommIntf = TPXCollectiveComm):
	79	+ return _reduce(input_, tp_intf), None
	80	+
	81	+ @staticmethod
	82	+ def forward(ctx, input_, tp_intf: CollectiveCommIntf = TPXCollectiveComm):
	83	+ return _reduce(input_, tp_intf)
	84	+
	85	+ @staticmethod
	86	+ def backward(ctx, grad_output):
	87	+ return grad_output, None
	88	+
	89	+
	90	+class _GatherFromParallelRegion(torch.autograd.Function):
	91	+ """Gather the input from model parallel region and concatinate."""
	92	+
	93	+ @staticmethod
	94	+ def symbolic(graph, input_):
	95	+ return _gather_along_last_dim(input_)
	96	+
	97	+ @staticmethod
	98	+ def forward(ctx, input_, comm_intf: CollectiveCommIntf):
	99	+ ctx.comm_intf = comm_intf
	100	+ return _gather_along_last_dim(input_, comm_intf)
	101	+
	102	+ @staticmethod
	103	+ def backward(ctx, grad_output):
	104	+ return _split_along_last_dim(grad_output, ctx.comm_intf), None
	105	+
	106	+
	107	+class _ScatterAlongLastDim(torch.autograd.Function):
	108	+ """Split the input and keep only the corresponding chuck to the rank."""
	109	+
	110	+ @staticmethod
	111	+ def symbolic(graph, input_, comm_intf: CollectiveCommIntf):
	112	+ return _split_along_last_dim(input_, comm_intf)
	113	+
	114	+ @staticmethod
	115	+ def forward(ctx, input_, comm_intf: CollectiveCommIntf):
	116	+ ctx.comm_intf = comm_intf
	117	+ return _split_along_last_dim(input_, comm_intf)
	118	+
	119	+ @staticmethod
	120	+ def backward(ctx, grad_output):
	121	+ return _gather_along_last_dim(grad_output, ctx.comm_intf), None
	122	+
	123	+
	124	+class _ScatterAlongFirstDim(torch.autograd.Function):
	125	+ """Split the input and keep only the corresponding chuck to the rank."""
	126	+
	127	+ @staticmethod
	128	+ def symbolic(graph, input_, comm_intf: CollectiveCommIntf):
	129	+ return _split_along_first_dim(input_, comm_intf)
	130	+
	131	+ @staticmethod
	132	+ def forward(ctx, input_, comm_intf: CollectiveCommIntf):
	133	+ ctx.comm_intf = comm_intf
	134	+ return _split_along_first_dim(input_, comm_intf)
	135	+
	136	+ @staticmethod
	137	+ def backward(ctx, grad_output):
	138	+ return sync_gather_along_first_dim(grad_output, ctx.comm_intf, "_ScatterAlongFirstDim"), None
	139	+
	140	+
	141	+class _ScatterAlongFirstDimThenLastDim(torch.autograd.Function):
	142	+ """Split the input and keep only the corresponding chuck to the rank."""
	143	+
	144	+ @staticmethod
	145	+ def symbolic(graph, local_rank_input, first_dim_comm_intf, last_dim_comm_intf):
	146	+ graph.first_dim_comm_intf = first_dim_comm_intf
	147	+ graph.last_dim_comm_intf = last_dim_comm_intf
	148	+
	149	+ first_dim_split_output = _split_along_first_dim(local_rank_input, first_dim_comm_intf)
	150	+ return _split_along_last_dim(first_dim_split_output, last_dim_comm_intf)
	151	+
	152	+ @staticmethod
	153	+ def forward(ctx, local_rank_input, first_dim_comm_intf, last_dim_comm_intf):
	154	+ ctx.first_dim_comm_intf = first_dim_comm_intf
	155	+ ctx.last_dim_comm_intf = last_dim_comm_intf
	156	+
	157	+ first_dim_split_output = _split_along_first_dim(local_rank_input, first_dim_comm_intf)
	158	+ return _split_along_last_dim(first_dim_split_output, last_dim_comm_intf)
	159	+
	160	+ @staticmethod
	161	+ def backward(ctx, grad_output):
	162	+ last_dim_gather_output = _gather_along_last_dim(grad_output, ctx.last_dim_comm_intf)
	163	+ first_dim_gather_output = sync_gather_along_first_dim(
	164	+ last_dim_gather_output, ctx.first_dim_comm_intf, "_ScatterAlongFirstDimThenLastDim"
	165	+ )
	166	+ return first_dim_gather_output, None, None
	167	+
	168	+
	169	+def auto_grad_sync_gather_along_first_dim(input_, comm_intf: CollectiveCommIntf):
	170	+ return _SyncGatherAlongFirstDim.apply(input_, comm_intf)
	171	+
	172	+
	173	+def auto_grad_sync_gather_along_last_dim(input_, comm_intf: CollectiveCommIntf):
	174	+ return _SyncGatherAlongLastDim.apply(input_, comm_intf)
	175	+
	176	+
	177	+def scatter_to_tensor_parallel_y_region(input_):
	178	+ return _ScatterAlongLastDim.apply(input_)
	179	+
	180	+
	181	+def auto_grad_scatter_along_last_dim(input_, comm_intf: CollectiveCommIntf):
	182	+ return _ScatterAlongLastDim.apply(input_, comm_intf)
	183	+
	184	+
	185	+def auto_grad_scatter_along_first_dim(input_, comm_intf: CollectiveCommIntf):
	186	+ return _ScatterAlongFirstDim.apply(input_, comm_intf)
	187	+
	188	+
	189	+def auto_grad_scatter_along_first_dim_then_last_dim(
	190	+ local_rank_input: torch.Tensor,
	191	+ first_dim_comm_intf: CollectiveCommIntf,
	192	+ last_dim_comm_intf: CollectiveCommIntf,
	193	+):
	194	+ return _ScatterAlongFirstDimThenLastDim.apply(
	195	+ local_rank_input, first_dim_comm_intf, last_dim_comm_intf
	196	+ )
	197	+
	198	+
	199	+def reduce_from_parallel_region(input_, tp_intf: CollectiveCommIntf = TPXCollectiveComm):
	200	+ return _ReduceFromModelParallelRegion.apply(input_, tp_intf)
	201	+
	202	+
	203	+def gather_from_parallel_region(input_, comm_intf: CollectiveCommIntf):
	204	+ return _GatherFromParallelRegion.apply(input_, comm_intf)

Amindspeed/core/tensor_parallel/comm_group_api.py+154-0

@@ -0,0 +1,154 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+import abc
	16	+
	17	+from megatron.core.parallel_state import get_context_parallel_group
	18	+from megatron.core.parallel_state import get_context_parallel_rank
	19	+from megatron.core.parallel_state import get_context_parallel_world_size
	20	+from megatron.core.parallel_state import get_tensor_model_parallel_group
	21	+from megatron.core.parallel_state import get_tensor_model_parallel_rank
	22	+from megatron.core.parallel_state import get_tensor_model_parallel_world_size
	23	+
	24	+from mindspeed.core.parallel_state import get_tensor_model_parallel_group_for_nd1_dim1
	25	+from mindspeed.core.parallel_state import get_tensor_model_parallel_group_for_nd1_dim1_rank
	26	+from mindspeed.core.parallel_state import get_tensor_model_parallel_group_for_nd1_dim1_world_size
	27	+from mindspeed.core.parallel_state import get_tensor_model_parallel_group_for_nd1_dim2
	28	+from mindspeed.core.parallel_state import get_tensor_model_parallel_group_for_nd1_dim2_rank
	29	+from mindspeed.core.parallel_state import get_tensor_model_parallel_group_for_nd1_dim2_world_size
	30	+from mindspeed.core.parallel_state import get_tp_x_ring_global_ranks
	31	+from mindspeed.core.parallel_state import get_tp_x_sd_rcv_overlap_group
	32	+from mindspeed.core.parallel_state import get_tp_y_ring_global_ranks
	33	+from mindspeed.core.parallel_state import get_tp_y_sd_rcv_overlap_group
	34	+
	35	+
	36	+class CollectiveCommIntf:
	37	+ def __init__(self, comm_group_name):
	38	+ self.comm_group_name = comm_group_name
	39	+
	40	+ @classmethod
	41	+ @abc.abstractmethod
	42	+ def get_comm_group_world_size(cls):
	43	+ raise NotImplementedError
	44	+
	45	+ @classmethod
	46	+ @abc.abstractmethod
	47	+ def get_comm_group(cls):
	48	+ raise NotImplementedError
	49	+
	50	+ @classmethod
	51	+ @abc.abstractmethod
	52	+ def get_comm_rank(cls):
	53	+ raise NotImplementedError
	54	+
	55	+ def get_comm_group_name(self):
	56	+ return self.comm_group_name
	57	+
	58	+
	59	+class OverlapCollectiveIntf(CollectiveCommIntf):
	60	+ @classmethod
	61	+ @abc.abstractmethod
	62	+ def get_ring_global_ranks(cls):
	63	+ raise NotImplementedError
	64	+
	65	+
	66	+class CPCollectiveComm(CollectiveCommIntf):
	67	+ @classmethod
	68	+ def get_comm_group_world_size(cls):
	69	+ return get_context_parallel_world_size()
	70	+
	71	+ @classmethod
	72	+ def get_comm_group(cls):
	73	+ return get_context_parallel_group()
	74	+
	75	+ @classmethod
	76	+ def get_comm_rank(cls):
	77	+ return get_context_parallel_rank()
	78	+
	79	+
	80	+class TPXCollectiveComm(CollectiveCommIntf):
	81	+ def __init__(self, name="tp-x"):
	82	+ super().__init__(name)
	83	+
	84	+ @classmethod
	85	+ def get_comm_rank(cls):
	86	+ return get_tensor_model_parallel_group_for_nd1_dim1_rank()
	87	+
	88	+ @classmethod
	89	+ def get_comm_group_world_size(cls):
	90	+ return get_tensor_model_parallel_group_for_nd1_dim1_world_size()
	91	+
	92	+ @classmethod
	93	+ def get_comm_group(cls):
	94	+ return get_tensor_model_parallel_group_for_nd1_dim1()
	95	+
	96	+
	97	+class TPXOverlapCollectiveComm(TPXCollectiveComm, OverlapCollectiveIntf):
	98	+ def __init__(self):
	99	+ super().__init__("tp-x-overlap")
	100	+
	101	+ @classmethod
	102	+ def get_comm_group(cls):
	103	+ return get_tp_x_sd_rcv_overlap_group()
	104	+
	105	+ @classmethod
	106	+ def get_ring_global_ranks(cls):
	107	+ return get_tp_x_ring_global_ranks()
	108	+
	109	+
	110	+class TPYCollectiveComm(CollectiveCommIntf):
	111	+ def __init__(self, name="tp-y"):
	112	+ super().__init__(name)
	113	+
	114	+ @classmethod
	115	+ def get_comm_rank(cls):
	116	+ return get_tensor_model_parallel_group_for_nd1_dim2_rank()
	117	+
	118	+ @classmethod
	119	+ def get_comm_group_world_size(cls):
	120	+ return get_tensor_model_parallel_group_for_nd1_dim2_world_size()
	121	+
	122	+ @classmethod
	123	+ def get_comm_group(cls):
	124	+ return get_tensor_model_parallel_group_for_nd1_dim2()
	125	+
	126	+
	127	+class TPYOverlapCollectiveComm(TPYCollectiveComm, OverlapCollectiveIntf):
	128	+ def __init__(self):
	129	+ super().__init__("y-overlap")
	130	+
	131	+ @classmethod
	132	+ def get_comm_group(cls):
	133	+ return get_tp_y_sd_rcv_overlap_group()
	134	+
	135	+ @classmethod
	136	+ def get_ring_global_ranks(cls):
	137	+ return get_tp_y_ring_global_ranks()
	138	+
	139	+
	140	+class TPXYCollectiveComm(CollectiveCommIntf):
	141	+ def __init__(self):
	142	+ super().__init__("tp-xy")
	143	+
	144	+ @classmethod
	145	+ def get_comm_rank(cls):
	146	+ return get_tensor_model_parallel_rank()
	147	+
	148	+ @classmethod
	149	+ def get_comm_group_world_size(cls):
	150	+ return get_tensor_model_parallel_world_size()
	151	+
	152	+ @classmethod
	153	+ def get_comm_group(cls):
	154	+ return get_tensor_model_parallel_group()

Amindspeed/core/tensor_parallel/comm_utils.py+270-0

@@ -0,0 +1,270 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+
	16	+import torch
	17	+from torch import Tensor
	18	+from torch import distributed
	19	+import torch.distributed as dist
	20	+
	21	+from megatron.core.parallel_state import get_global_memory_buffer
	22	+from mindspeed.core.tensor_parallel.comm_group_api import CollectiveCommIntf
	23	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm
	24	+
	25	+
	26	+def _split_along_last_dim(
	27	+ local_rank_input: Tensor, comm_intf: CollectiveCommIntf = TPXCollectiveComm
	28	+):
	29	+ """Split the tensor along its last dimension and keep the
	30	+ corresponding slice."""
	31	+
	32	+ world_size = comm_intf.get_comm_group_world_size()
	33	+ # Bypass the function if we are using only 1 GPU.
	34	+ if world_size == 1:
	35	+ return local_rank_input
	36	+
	37	+ # Split along last dimension.
	38	+ last_dim = local_rank_input.dim() - 1
	39	+ last_dim_size = local_rank_input.size()[last_dim] // world_size
	40	+ # Split.
	41	+ tensor_list = torch.split(local_rank_input, last_dim_size, dim=last_dim)
	42	+
	43	+ # Note: torch.split does not create contiguous tensors by default.
	44	+ rank = comm_intf.get_comm_rank()
	45	+ output = tensor_list[rank].contiguous()
	46	+
	47	+ return output
	48	+
	49	+
	50	+def _split_along_first_dim(local_rank_input, comm_intf: CollectiveCommIntf = TPXCollectiveComm):
	51	+ """Split the tensor along its first dimension and keep the
	52	+ corresponding slice."""
	53	+
	54	+ world_size = comm_intf.get_comm_group_world_size()
	55	+ # Bypass the function if we are using only 1 GPU.
	56	+ if world_size == 1:
	57	+ return local_rank_input
	58	+
	59	+ # Split along first dimension.
	60	+ dim_size = local_rank_input.size()[0]
	61	+ if dim_size % world_size:
	62	+ raise AssertionError("First dimension of the tensor should be divisible by parallel size")
	63	+ local_dim_size = dim_size // world_size
	64	+ rank = comm_intf.get_comm_rank()
	65	+ dim_offset = rank * local_dim_size
	66	+
	67	+ output = local_rank_input[dim_offset : dim_offset + local_dim_size].contiguous()
	68	+
	69	+ return output
	70	+
	71	+
	72	+def _gather_along_last_dim(
	73	+ local_rank_input: Tensor, ag_comm_intf: CollectiveCommIntf = TPXCollectiveComm
	74	+):
	75	+ """Gather tensors and concatinate along the last dimension."""
	76	+
	77	+ world_size = ag_comm_intf.get_comm_group_world_size()
	78	+ # Bypass the function if we are using only 1 GPU.
	79	+ if world_size == 1:
	80	+ return local_rank_input
	81	+
	82	+ tensor_list = [torch.empty_like(local_rank_input) for _ in range(world_size)]
	83	+ torch.distributed.all_gather(
	84	+ tensor_list, local_rank_input, group=ag_comm_intf.get_comm_group(), async_op=False
	85	+ )
	86	+
	87	+ # Note: torch.cat already creates a contiguous tensor.
	88	+ last_dim = local_rank_input.dim() - 1
	89	+ output = torch.cat(tensor_list, dim=last_dim).contiguous()
	90	+ return output
	91	+
	92	+
	93	+def sync_gather_along_last_dim(
	94	+ local_rank_tensor: Tensor, ag_comm_intf: CollectiveCommIntf = TPXCollectiveComm
	95	+):
	96	+ """Gather tensors and concatinate along the last dimension synchronously.
	97	+
	98	+ :param local_rank_tensor: input of current rank.
	99	+ :param ag_comm_intf: the communication process group interface.
	100	+ :return: the AllGather-ed result.
	101	+ """
	102	+
	103	+ world_size = ag_comm_intf.get_comm_group_world_size()
	104	+ # Bypass the function if we are using only 1 GPU/NPU.
	105	+ if world_size == 1:
	106	+ return local_rank_tensor
	107	+
	108	+ gathered_tensors = [torch.empty_like(local_rank_tensor) for _ in range(world_size)]
	109	+ torch.distributed.all_gather(
	110	+ gathered_tensors,
	111	+ local_rank_tensor.contiguous(),
	112	+ group=ag_comm_intf.get_comm_group(),
	113	+ async_op=False,
	114	+ )
	115	+
	116	+ return torch.cat(gathered_tensors, dim=local_rank_tensor.dim() - 1).contiguous()
	117	+
	118	+
	119	+def async_gather_tensors(
	120	+ local_rank_input: Tensor,
	121	+ ag_comm_intf: CollectiveCommIntf = TPXCollectiveComm,
	122	+ buffer_name="mpu-async-tp-2d",
	123	+):
	124	+ """Gather tensors and concatinate along the last dimension asynchronously.
	125	+
	126	+ :param local_rank_input: input of current rank.
	127	+ :param ag_comm_intf: the AllGather communication process group interface.
	128	+ :param buffer_name: buffer name of str type.
	129	+ :return: the AllGather op handle and tensor list storing the op result tensors.
	130	+
	131	+ Note: the result tensors may be handled as following according to your need:
	132	+ output = torch.cat(gathered_tensors, dim=xx_dim).contiguous()
	133	+ """
	134	+
	135	+ world_size = ag_comm_intf.get_comm_group_world_size()
	136	+ # Bypass the function if we are using only 1 NPU/GPU.
	137	+ if world_size == 1:
	138	+ return None, local_rank_input
	139	+
	140	+ dim_size = list(local_rank_input.size())
	141	+ dim_size[0] *= world_size
	142	+
	143	+ all_gather_buffer = get_global_memory_buffer().get_tensor(
	144	+ dim_size, local_rank_input.dtype, buffer_name
	145	+ )
	146	+ handle = torch.distributed._all_gather_base(
	147	+ all_gather_buffer, local_rank_input, group=ag_comm_intf.get_comm_group(), async_op=True
	148	+ )
	149	+
	150	+ return handle, all_gather_buffer
	151	+
	152	+
	153	+def sync_gather_along_first_dim(
	154	+ local_rank_input: Tensor,
	155	+ comm_intf: CollectiveCommIntf = TPXCollectiveComm,
	156	+ buffer_name="mpu-sync-tp-2d",
	157	+):
	158	+ """Gather tensors and concatinate along the first dimension."""
	159	+
	160	+ world_size = comm_intf.get_comm_group_world_size()
	161	+ # Bypass the function if we are using only 1 GPU.
	162	+ if world_size == 1:
	163	+ return local_rank_input
	164	+
	165	+ dim_size = list(local_rank_input.size())
	166	+ dim_size[0] *= world_size
	167	+
	168	+ output = get_global_memory_buffer().get_tensor(dim_size, local_rank_input.dtype, buffer_name)
	169	+ torch.distributed._all_gather_base(
	170	+ output, local_rank_input.contiguous(), group=comm_intf.get_comm_group()
	171	+ )
	172	+
	173	+ return output
	174	+
	175	+
	176	+def sync_reduce_scatter_along_first_dim(
	177	+ local_rank_input, comm_intf: CollectiveCommIntf = TPXCollectiveComm
	178	+):
	179	+ """Reduce-scatter the input tensor across specified parallel group."""
	180	+ world_size = comm_intf.get_comm_group_world_size()
	181	+ # Bypass the function if we are using only 1 GPU.
	182	+ if world_size == 1:
	183	+ return local_rank_input
	184	+
	185	+ dim_size = list(local_rank_input.size())
	186	+ if dim_size[0] % world_size:
	187	+ raise AssertionError("First dimension of the tensor should be divisible by tensor parallel size")
	188	+
	189	+ dim_size[0] = dim_size[0] // world_size
	190	+
	191	+ output = torch.empty(dim_size, dtype=local_rank_input.dtype, device=torch.cuda.current_device())
	192	+ dist.reduce_scatter_tensor(
	193	+ output, local_rank_input.contiguous(), group=comm_intf.get_comm_group(), async_op=False
	194	+ )
	195	+
	196	+ return output
	197	+
	198	+
	199	+def async_reduce_scatter_along_first_dim(
	200	+ local_rank_input, comm_intf: CollectiveCommIntf = TPXCollectiveComm
	201	+):
	202	+ """Reduce-scatter the input tensor across parallel group specified by comm_intf."""
	203	+ world_size = comm_intf.get_comm_group_world_size()
	204	+ # Bypass the function if we are using only 1 GPU.
	205	+ if world_size == 1:
	206	+ return None, local_rank_input
	207	+
	208	+ dim_size = list(local_rank_input.size())
	209	+ if dim_size[0] % world_size:
	210	+ raise AssertionError("First dimension of the tensor should be divisible by parallel size")
	211	+
	212	+ dim_size[0] = dim_size[0] // world_size
	213	+
	214	+ rs_output = torch.empty(
	215	+ dim_size, dtype=local_rank_input.dtype, device=torch.cuda.current_device()
	216	+ )
	217	+ handle = dist.reduce_scatter_tensor(
	218	+ rs_output, local_rank_input.contiguous(), group=comm_intf.get_comm_group(), async_op=True
	219	+ )
	220	+ return handle, rs_output
	221	+
	222	+
	223	+def async_gather_along_last_dim(input_, comm_intf: CollectiveCommIntf = TPXCollectiveComm):
	224	+ world_size = comm_intf.get_comm_group_world_size()
	225	+ # Bypass the function if we are using only 1 GPU/NPU.
	226	+ if world_size == 1:
	227	+ return None, input_
	228	+
	229	+ gathered_tensors = [torch.empty_like(input_) for _ in range(world_size)]
	230	+ handle = torch.distributed.all_gather(
	231	+ gathered_tensors, input_.contiguous(), group=comm_intf.get_comm_group(), async_op=True,
	232	+ )
	233	+
	234	+ return handle, gathered_tensors
	235	+
	236	+
	237	+def sync_reduce_scatter_along_last_dim(
	238	+ local_rank_input, rs_comm_intf: CollectiveCommIntf = TPXCollectiveComm
	239	+):
	240	+ """Reduce-scatter the input tensor across specified parallel group."""
	241	+ world_size = rs_comm_intf.get_comm_group_world_size()
	242	+ # Bypass the function if we are using only 1 GPU.
	243	+ if world_size == 1:
	244	+ return local_rank_input
	245	+
	246	+ local_rank_input = local_rank_input.transpose(0, 2)
	247	+ output = sync_reduce_scatter_along_first_dim(local_rank_input, rs_comm_intf)
	248	+ return output.transpose(0, 2).contiguous()
	249	+
	250	+
	251	+def async_reduce_scatter_along_last_dim(
	252	+ local_rank_input, rs_comm_intf: CollectiveCommIntf = TPXCollectiveComm
	253	+):
	254	+ """Reduce-scatter the input tensor across model parallel group.
	255	+
	256	+ :param local_rank_input: input of local rank
	257	+ :param rs_comm_intf: Reduce scatter comm intf.
	258	+ :return:
	259	+
	260	+ Note: the result tensors should be handled as following:
	261	+ rs_output = rs_output.transpose(0, 2).contiguous()
	262	+
	263	+ """
	264	+ world_size = rs_comm_intf.get_comm_group_world_size()
	265	+ # Bypass the function if we are using only 1 GPU.
	266	+ if world_size == 1:
	267	+ return None, local_rank_input
	268	+
	269	+ local_rank_input = local_rank_input.transpose(0, 2)
	270	+ return async_reduce_scatter_along_first_dim(local_rank_input, rs_comm_intf)

Mmindspeed/core/tensor_parallel/layers.py+58-0

	@@ -1592,6 +1592,64 @@ class Mc2RowSeqParallelLinear(torch.autograd.Function):
1592	return grad_input, grad_weight, grad_bias, None	1592	return grad_input, grad_weight, grad_bias, None
1593		1593
1594		1594
		1595	+def _initialize_affine_weight_cpu_2d(
		1596	+ weight,
		1597	+ output_size,
		1598	+ input_size,
		1599	+ input_size_per_partition,
		1600	+ output_size_per_partition,
		1601	+ partition_dim,
		1602	+ init_method,
		1603	+ stride=1,
		1604	+ return_master_weight=False,
		1605	+ *,
		1606	+ params_dtype=torch.float32
		1607	+):
		1608	+ """Initialize affine weight for model parallel when use tp-2d"""
		1609	+ set_tensor_model_parallel_attributes(
		1610	+ tensor=weight, is_parallel=True, dim=partition_dim, stride=stride
		1611	+ )
		1612	+
		1613	+ # Initialize master weight
		1614	+ master_weight = torch.empty(output_size, input_size, dtype=torch.float, requires_grad=False)
		1615	+ init_method(master_weight)
		1616	+
		1617	+ master_weight = master_weight.to(dtype=params_dtype)
		1618	+ # Split and copy
		1619	+ rank = ps.get_tensor_model_parallel_rank()
		1620	+ world_size = ps.get_tensor_model_parallel_world_size()
		1621	+
		1622	+ def compute_target_rank(rank, row_num, col_num):
		1623	+ return rank % row_num * col_num + rank // row_num
		1624	+
		1625	+ # The weight positions of nd and megatron are different. So weight needs to be rearranged.
		1626	+ # This rearrangement is only to make the calculations of nd and megatron consistent.
		1627	+ # Even if this rearrangement is removed, it will not affect the correctness of nd calculation.
		1628	+ if partition_dim == 0:
		1629	+ row_num = input_size // input_size_per_partition
		1630	+ col_num = output_size // output_size_per_partition
		1631	+ else:
		1632	+ col_num = input_size // input_size_per_partition
		1633	+ row_num = output_size // output_size_per_partition
		1634	+ weight_list = torch.split(master_weight, master_weight.size()[partition_dim] // world_size, dim=partition_dim)
		1635	+ tensor_list = [weight_list[compute_target_rank(i, row_num, col_num)] for i in range(world_size)]
		1636	+ master_weight = torch.cat(tensor_list, dim=partition_dim)
		1637	+ weight_list_1 = torch.split(master_weight, input_size_per_partition, dim=1)
		1638	+ if partition_dim == 0:
		1639	+ weight_1 = weight_list_1[rank // col_num]
		1640	+ else:
		1641	+ weight_1 = weight_list_1[rank % col_num]
		1642	+ weight_list_2 = torch.split(weight_1, output_size_per_partition, dim=0)
		1643	+ if partition_dim == 0:
		1644	+ my_weight_list = weight_list_2[rank % col_num:: world_size]
		1645	+ else:
		1646	+ my_weight_list = weight_list_2[rank // col_num:: world_size]
		1647	+ with torch.no_grad():
		1648	+ torch.cat(my_weight_list, dim=partition_dim, out=weight)
		1649	+ if return_master_weight:
		1650	+ return master_weight
		1651	+
		1652	+
1595	def _initialize_affine_weight_cpu_nd(	1653	def _initialize_affine_weight_cpu_nd(
1596	weight,	1654	weight,
1597	output_size,	1655	output_size,

Amindspeed/core/tensor_parallel/tp_2d/layernorm_2d.py+179-0

@@ -0,0 +1,179 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+from typing import Any
	16	+from typing import Tuple
	17	+
	18	+import torch
	19	+import torch.distributed as dist
	20	+from torch import Tensor
	21	+from torch.cuda.amp import custom_bwd
	22	+from torch.cuda.amp import custom_fwd
	23	+from torch.nn import Parameter
	24	+
	25	+from megatron.core.utils import divide
	26	+from mindspeed.core.tensor_parallel.comm_group_api import CollectiveCommIntf
	27	+from mindspeed.core.tensor_parallel.comm_group_api import TPYCollectiveComm
	28	+
	29	+
	30	+class LayerNorm2D(torch.nn.Module):
	31	+ """LayerNorm2D layer with row and column parallelism.
	32	+
	33	+ Arguments:
	34	+ hidden_size (int): input normalized size from an expected input of size
	35	+ eps: a value added to the denominator for numerical stability. Default: 1e-5
	36	+ bias: (bool, optional): Whether to add a bias, defaults to ``True``.
	37	+ dtype: (:class:`torch.dtype`, optional): The dtype of parameters, defaults to None.
	38	+ last_dim_split_comm_intf: Reduce scatter comm intf.
	39	+ """
	40	+
	41	+ def __init__(
	42	+ self,
	43	+ hidden_size: int,
	44	+ eps: float = 1e-5,
	45	+ bias: bool = True,
	46	+ dtype=None,
	47	+ last_dim_split_comm_intf: CollectiveCommIntf = TPYCollectiveComm(),
	48	+ ) -> None:
	49	+ super(LayerNorm2D, self).__init__()
	50	+ # layer norm config
	51	+ self.hidden_size = hidden_size
	52	+ self.epsilon = eps
	53	+
	54	+ # parallel setting
	55	+ self.last_dim_split_comm_intf = last_dim_split_comm_intf
	56	+ self.rs_comm_world_sz = self.last_dim_split_comm_intf.get_comm_group_world_size()
	57	+ # partitioning dimension
	58	+ self.partitioned_dim = divide(hidden_size, self.rs_comm_world_sz)
	59	+ # create parameters
	60	+ factory_kwargs = {"device": torch.cuda.current_device(), "dtype": dtype}
	61	+
	62	+ # [H/(xy)]
	63	+ self.weight = Parameter(torch.ones(self.partitioned_dim, **factory_kwargs))
	64	+ if bias:
	65	+ # [H/(xy)]
	66	+ self.bias = Parameter(torch.zeros(self.partitioned_dim, **factory_kwargs))
	67	+ else:
	68	+ self.bias = None
	69	+
	70	+ # set sequence parallelism flag on weight and bias parameters
	71	+ setattr(self.weight, "2d_tp", True)
	72	+ setattr(self.bias, "2d_tp", True)
	73	+
	74	+ def forward(self, x: Tensor) -> Tensor:
	75	+ return _ParallelLayerNorm2D.apply(
	76	+ x,
	77	+ self.weight,
	78	+ self.bias,
	79	+ self.epsilon,
	80	+ self.hidden_size,
	81	+ self.last_dim_split_comm_intf,
	82	+ )
	83	+
	84	+
	85	+class _ParallelLayerNorm2D(torch.autograd.Function):
	86	+ @staticmethod
	87	+ @custom_fwd
	88	+ def forward(
	89	+ ctx: Any,
	90	+ input_: Tensor,
	91	+ weight,
	92	+ bias,
	93	+ epsilon,
	94	+ hidden_size: int,
	95	+ last_dim_split_comm_intf: CollectiveCommIntf
	96	+ ) -> Tensor:
	97	+ """
	98	+
	99	+ :param ctx:
	100	+ :param input_: [s/(cp*x), b, H/y]
	101	+ :param weight: [H/(xy)]
	102	+ :param bias: [H/(xy)]
	103	+ :param epsilon:
	104	+ :param hidden_size: H
	105	+ :param last_dim_split_comm_intf:
	106	+ :return:
	107	+ """
	108	+ # [s/(cpx), b, H/y]---> [s/(cpx), b, 1]
	109	+ e_x = torch.sum(input_, dim=-1, keepdim=True)
	110	+ # [s/(cp*x), b, 1]
	111	+ handle_ex = torch.distributed.all_reduce(
	112	+ e_x, group=last_dim_split_comm_intf.get_comm_group(), async_op=True
	113	+ )
	114	+
	115	+ # [s/(cpx), b, H/y]---> [s/(cpx), b, 1]
	116	+ var_x = torch.sum(input_.float().pow(2), dim=-1, keepdim=True)
	117	+ if handle_ex:
	118	+ handle_ex.wait()
	119	+
	120	+ handle_var = torch.distributed.all_reduce(
	121	+ var_x, group=last_dim_split_comm_intf.get_comm_group(), async_op=True
	122	+ )
	123	+
	124	+ input_.sub_(e_x.div_(hidden_size))
	125	+ e_x.mul_(e_x)
	126	+ if handle_var:
	127	+ handle_var.wait()
	128	+
	129	+ var_x = torch.rsqrt(var_x.div_(hidden_size).sub_(e_x).add_(epsilon))
	130	+
	131	+ ctx.hidden_size = hidden_size
	132	+ ctx.last_dim_split_comm_intf = last_dim_split_comm_intf
	133	+ # [s/(cpx), b, H/y] [s/(cpx), b, 1] --> [s/(cpx), b, H/y]
	134	+ norm_x = torch.mul(input_, var_x)
	135	+
	136	+ if bias is not None:
	137	+ # bias + weight * norm, [H/y] + [H/y] * [s/(cp*x), b, H/y]
	138	+ output = torch.addcmul(bias, weight, norm_x)
	139	+ else:
	140	+ output = torch.mul(weight, norm_x)
	141	+
	142	+ ctx.save_for_backward(norm_x, var_x, bias, weight)
	143	+ return output
	144	+
	145	+ @staticmethod
	146	+ @custom_bwd
	147	+ def backward(ctx: Any, output_grad: Tensor) -> Tuple[Tensor, ...]:
	148	+ x, var_x, bias, weight = ctx.saved_tensors
	149	+ # calculate grad_bias
	150	+ if bias is None:
	151	+ grad_bias = None
	152	+ else:
	153	+ grad_bias = output_grad.sum(dim=(0, 1))
	154	+
	155	+ # calculate grad_input
	156	+ grad_norm_x = torch.mul(output_grad, weight)
	157	+ output_grad_sum = torch.sum(grad_norm_x, dim=-1, keepdim=True)
	158	+ handle_grad_sum = torch.distributed.all_reduce(
	159	+ output_grad_sum, group=ctx.last_dim_split_comm_intf.get_comm_group(), async_op=True
	160	+ )
	161	+ output_grad_mul_x_sum = torch.sum(grad_norm_x * x, dim=-1, keepdim=True)
	162	+
	163	+ # calculate grad_weight
	164	+ grad_weight = torch.mul(output_grad, x)
	165	+ grad_weight = grad_weight.sum(dim=(0, 1))
	166	+
	167	+ if handle_grad_sum:
	168	+ handle_grad_sum.wait()
	169	+
	170	+ handle_grad_mul_x = torch.distributed.all_reduce(
	171	+ output_grad_mul_x_sum, group=ctx.last_dim_split_comm_intf.get_comm_group(), async_op=True
	172	+ )
	173	+ output_grad_sum.div_(ctx.hidden_size)
	174	+ grad_input = grad_norm_x.sub(output_grad_sum)
	175	+ if handle_grad_mul_x:
	176	+ handle_grad_mul_x.wait()
	177	+
	178	+ grad_input = (grad_input - x * (output_grad_mul_x_sum / ctx.hidden_size)) * var_x
	179	+ return grad_input, grad_weight, grad_bias, None, None, None

Amindspeed/core/tensor_parallel/tp_2d/linear_2d_split_along_first_dim.py+319-0

@@ -0,0 +1,319 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+
	16	+import torch
	17	+import torch_npu
	18	+from torch import distributed as torch_dist
	19	+from torch.cuda.amp import custom_bwd
	20	+from torch.cuda.amp import custom_fwd
	21	+from megatron.training import get_args
	22	+
	23	+from mindspeed.core.tensor_parallel.comm_autograd_function import CollectiveCommIntf
	24	+from mindspeed.core.tensor_parallel.comm_group_api import OverlapCollectiveIntf
	25	+from mindspeed.core.tensor_parallel.comm_utils import async_gather_tensors
	26	+from mindspeed.core.tensor_parallel.comm_utils import async_reduce_scatter_along_first_dim
	27	+from mindspeed.core.tensor_parallel.comm_utils import sync_gather_along_first_dim
	28	+from mindspeed.core.tensor_parallel.comm_utils import sync_reduce_scatter_along_first_dim
	29	+
	30	+
	31	+class Linear2DSplitAlongFirstDim(torch.autograd.Function):
	32	+ """2D Linear out axe communication implementation."""
	33	+
	34	+ @staticmethod
	35	+ @custom_fwd
	36	+ def forward(
	37	+ ctx,
	38	+ activation_input,
	39	+ weight,
	40	+ bias,
	41	+ ag_comm_intf: CollectiveCommIntf,
	42	+ ag_overlap_comm_intf: OverlapCollectiveIntf,
	43	+ rs_comm_intf: CollectiveCommIntf,
	44	+ rs_overlap_comm_intf: OverlapCollectiveIntf,
	45	+ enable_overlap_ag_with_matmul=False,
	46	+ enable_overlap_matmul_with_rs=False,
	47	+ gradient_accumulation_fusion=False,
	48	+ ):
	49	+ """
	50	+ :param ctx: context to save some tensors or vars for backward use.
	51	+ :param activation_input: with shape: [s/(x*cp), b, h/y]
	52	+ :param weight: with shape: [h/y, E/x], E means the output size.
	53	+ :param bias: bias parameter tensor.
	54	+ :param ag_comm_intf: AllGather communication process group interface.
	55	+ :param ag_overlap_comm_intf: AllGather communication overlap send and recv comm group
	56	+ :param rs_comm_intf: ReduceScatter communication process group interface.
	57	+ :param rs_overlap_comm_intf: ReduceScatter communication overlap send and recv comm group
	58	+ :param enable_overlap_ag_with_matmul: enable overlap all-gather with matmul
	59	+ :param enable_overlap_matmul_with_rs: enable overlap matmul with reduce-scatter
	60	+ :param gradient_accumulation_fusion: enable gradient accumulation fusion
	61	+ :return: forward result tensor.
	62	+ """
	63	+ ctx.save_for_backward(activation_input)
	64	+ ctx.weight = weight
	65	+ ctx.use_bias = bias is not None
	66	+ ctx.rs_comm_intf = rs_comm_intf
	67	+ ctx.ag_comm_intf = ag_comm_intf
	68	+ ctx.ag_overlap_comm_intf = ag_overlap_comm_intf
	69	+ ctx.rs_overlap_comm_intf = rs_overlap_comm_intf
	70	+ ctx.gradient_accumulation_fusion = gradient_accumulation_fusion
	71	+ if enable_overlap_matmul_with_rs:
	72	+ activation_input = activation_input.contiguous()
	73	+ return Linear2DSplitAlongFirstDim._do_mm_overlap_reducescatter(
	74	+ activation_input, weight.t(), bias, ag_comm_intf, rs_comm_intf
	75	+ )
	76	+
	77	+ # first_linear forward: [s/cp, b, H/y] @ [H/y, e/x] -> [s/cp, b, e/x]
	78	+ if enable_overlap_ag_with_matmul:
	79	+ matmul_res, _ = Linear2DSplitAlongFirstDim._do_allgather_left_tensor_and_matmul_overlap(
	80	+ ag_comm_intf,
	81	+ ag_overlap_comm_intf,
	82	+ part_left_tensor=activation_input,
	83	+ full_right_tensor=weight.t(),
	84	+ )
	85	+
	86	+ if bias is not None:
	87	+ matmul_res += bias
	88	+ elif get_args().use_ascend_mc2: # MC2 适配
	89	+ ag_group = ag_comm_intf.get_comm_group()
	90	+ rank = ag_comm_intf.get_comm_rank()
	91	+ if torch.__version__ > "2.0":
	92	+ global_rank = torch.distributed.get_global_rank(ag_group, rank)
	93	+ hcomm_info = ag_group._get_backend(torch.device("npu")).get_hccl_comm_name(
	94	+ global_rank
	95	+ )
	96	+ else:
	97	+ hcomm_info = ag_group.get_hccl_comm_name(rank)
	98	+
	99	+ x = activation_input.reshape(-1, activation_input.size(2))
	100	+
	101	+ # [s/(x*cp), b, H/y] --AG(x)--> [s/cp, b, H/y]
	102	+ # [s/cp, b, H/y] @ [H/y, e/x] -> [sb/cp, e/x]
	103	+ output, _ = torch_npu.npu_all_gather_base_mm(
	104	+ x,
	105	+ weight.t(),
	106	+ hcomm_info,
	107	+ ag_comm_intf.get_comm_group_world_size(),
	108	+ bias=bias,
	109	+ gather_index=0,
	110	+ gather_output=False,
	111	+ )
	112	+ # [sb/cp, e/x]---> [s/cp, b, e/x]
	113	+ matmul_res = output.view(
	114	+ output.shape[0] // activation_input.shape[1], activation_input.shape[1], output.shape[1]
	115	+ )
	116	+ else:
	117	+ # [s/(x*cp), b, H/y] -> [s/cp, b, H/y]
	118	+ activation_input = activation_input.contiguous()
	119	+ total_input = sync_gather_along_first_dim(activation_input, ag_comm_intf)
	120	+ # [s/cp, b, H/y] @ [H/y, e/x] -> [s/cp, b, e/x]
	121	+ matmul_res = torch.matmul(total_input, weight.t())
	122	+ # [s/cp, b, E/x] -> [s/(y*cp), b, E/x]
	123	+ matmul_res = matmul_res.contiguous()
	124	+ matmul_res = sync_reduce_scatter_along_first_dim(matmul_res, rs_comm_intf)
	125	+ return matmul_res
	126	+
	127	+ @staticmethod
	128	+ @custom_bwd
	129	+ def backward(ctx, grad_output):
	130	+ """Backward implementation of Linear2DSplitAlongFirstDim, the computation and communication
	131	+ overlap:
	132	+
	133	+ ----------------------------------------------------------------------------->time
	134	+ \| AG(grad_o, Y\|X)
	135	+ \| AG(activation_input, X\|Y)
	136	+ \| part_grad_act = MM(tot_grad_o, weight)
	137	+ \| RS(part_grad_act, X\|Y)
	138	+ \| MM(tot_grad_o^T, tot_act_input)
	139	+
	140	+
	141	+ :param ctx: context
	142	+ :param grad_output: with shape: [s/cp, b, E/(xy)]
	143	+ :return:grads of all the input para of forward function as a tuple
	144	+ """
	145	+ # activation_input shape: [s/(x*cp), b, h/y]
	146	+ # weight shape: [h/y, E/x]
	147	+ activation_input, = ctx.saved_tensors
	148	+ weight = ctx.weight
	149	+ use_bias = ctx.use_bias
	150	+ # first we prepare the total inputs needed to compute grad_input, grad_weight.
	151	+ # [s/(y*cp), b, E/x]---AG(y)---> [s/cp, b, E/x]
	152	+ # Use sync AG to avoid communication competition, for the bandwidth is shared for 910C.
	153	+ grad_output = grad_output.contiguous()
	154	+ total_grad_output = sync_gather_along_first_dim(grad_output, ctx.rs_comm_intf)
	155	+
	156	+ # prepare total activation_input for computing grad weight.
	157	+ # [s/(x*cp), b, h/y]---AG(X)--->[s/cp, b, h/y]
	158	+ activation_input = activation_input.contiguous()
	159	+ gather_input_handle, gathered_tensors = async_gather_tensors(
	160	+ local_rank_input=activation_input, ag_comm_intf=ctx.ag_comm_intf
	161	+ )
	162	+
	163	+ # [s/cp, b, E/x] @ [E/x, H/y]--> [s/cp, b, H/y] (partial sum)
	164	+ partial_grad_input = total_grad_output.matmul(weight).contiguous()
	165	+ # [s/cp, b, H/y] (partial sum)---RS(X)--->[s/cp, b, H/(xy)] (full sum)
	166	+ rs_grad_input_handle, grad_input = async_reduce_scatter_along_first_dim(
	167	+ partial_grad_input, comm_intf=ctx.ag_comm_intf
	168	+ )
	169	+
	170	+ # Convert the tensor shapes to 2D for execution compatibility
	171	+ sb = total_grad_output.shape[0] * total_grad_output.shape[1]
	172	+ # [s/cp, b, E/x]--view--> [sb/cp, E/x]
	173	+ total_grad_output = total_grad_output.view(sb, total_grad_output.shape[2])
	174	+
	175	+ if gather_input_handle:
	176	+ gather_input_handle.wait()
	177	+
	178	+ # [s/(x*cp), b, h/y]---AG(X)--->[s/cp, b, h/y]
	179	+ total_activation_input = gathered_tensors
	180	+ # [s/cp, b, h/y]--view--> [sb/cp, h/y]
	181	+ total_activation_input = total_activation_input.view(sb, total_activation_input.shape[2])
	182	+ if ctx.gradient_accumulation_fusion:
	183	+ import fused_weight_gradient_mlp_cuda
	184	+ total_grad_output = total_grad_output.contiguous()
	185	+ if weight.main_grad.dtype == torch.float32:
	186	+ fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
	187	+ total_activation_input, total_grad_output, weight.main_grad
	188	+ )
	189	+ elif weight.main_grad.dtype in (torch.float16, torch.bfloat16):
	190	+ fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(
	191	+ total_activation_input, total_grad_output, weight.main_grad
	192	+ )
	193	+ else:
	194	+ raise RuntimeError("Unsupported gradient type for gradient accumulation fusion")
	195	+
	196	+ if hasattr(weight, 'grad_added_to_main_grad'):
	197	+ # When overlap_grad_reduce is True, need to ensure that backward hooks
	198	+ # are all run on the main backprop thread to prevent deadlocks. Setup
	199	+ # dummy grad_weight tensor to prevent backward hooks from being run
	200	+ # in a background thread.
	201	+ if getattr(weight, 'zero_out_wgrad', False):
	202	+ grad_weight = torch.zeros(
	203	+ weight.main_grad.shape,
	204	+ dtype=activation_input.dtype,
	205	+ device=torch.cuda.current_device(),
	206	+ requires_grad=False,
	207	+ )
	208	+ else:
	209	+ grad_weight = torch.empty(
	210	+ weight.main_grad.shape,
	211	+ dtype=activation_input.dtype,
	212	+ device=torch.cuda.current_device(),
	213	+ requires_grad=False,
	214	+ )
	215	+ weight.grad_added_to_main_grad = True
	216	+ else:
	217	+ grad_weight = None
	218	+ else:
	219	+ # [E/x, sb/cp] @ [sb/cp, h/y] ---> [E/x, h/y]
	220	+ grad_weight = total_grad_output.t().matmul(total_activation_input)
	221	+ grad_bias = total_grad_output.sum(dim=0) if use_bias else None
	222	+
	223	+ if rs_grad_input_handle:
	224	+ rs_grad_input_handle.wait()
	225	+ return grad_input, grad_weight, grad_bias, None, None, None, None, None, None, None
	226	+
	227	+ @staticmethod
	228	+ def _do_allgather_left_tensor_and_matmul_overlap(
	229	+ ag_comm_intf, ag_overlap_comm_intf, part_left_tensor, full_right_tensor, return_ag_res=False
	230	+ ):
	231	+ cur_ag_rank = ag_comm_intf.get_comm_rank()
	232	+ ag_world_sz = ag_comm_intf.get_comm_group_world_size()
	233	+
	234	+ # do tp-x times matmul and reduce the partial res.
	235	+ matmul_res = [None] * ag_world_sz
	236	+ cur_step_rcv_handle = None
	237	+ ring_ag_ranks = ag_overlap_comm_intf.get_ring_global_ranks()
	238	+ next_rank = ring_ag_ranks[(cur_ag_rank + ag_world_sz - 1) % ag_world_sz]
	239	+ prev_rank = ring_ag_ranks[(cur_ag_rank + 1) % ag_world_sz]
	240	+ ag_comm_group = ag_comm_intf.get_comm_group()
	241	+ ag_overlap_comm_group = ag_overlap_comm_intf.get_comm_group()
	242	+ cur_step_tensor_to_send = part_left_tensor
	243	+
	244	+ # 下一次要计算的数据（本次要从上一个 rank 接收的 tensor。）
	245	+ cur_step_rcv_input = torch.empty_like(part_left_tensor)
	246	+ all_ag_res = None
	247	+ if return_ag_res:
	248	+ all_ag_res = [None] * ag_world_sz
	249	+ all_ag_res[cur_ag_rank] = part_left_tensor
	250	+
	251	+ # first_linear forward: [H/y, e/x] -> [H/(xy), e/x]
	252	+ for step in range(ag_world_sz):
	253	+ if step < ag_world_sz - 1 and cur_ag_rank % 2 == 0: # 偶数 rank 先发再收
	254	+ torch_dist.isend(cur_step_tensor_to_send, next_rank, ag_comm_group)
	255	+ cur_step_rcv_handle = torch_dist.irecv(
	256	+ cur_step_rcv_input, prev_rank, ag_overlap_comm_group
	257	+ )
	258	+ elif step < ag_world_sz - 1 and cur_ag_rank % 2 == 1: # 奇数 rank 先收再发
	259	+ cur_step_rcv_handle = torch_dist.irecv(cur_step_rcv_input, prev_rank, ag_comm_group)
	260	+ torch_dist.isend(cur_step_tensor_to_send, next_rank, ag_overlap_comm_group)
	261	+
	262	+ # compute: part_left_tensor @ split_right(split by inner dim)
	263	+ # [e/x, h/(xy)]
	264	+ cur_tensor_idx = (step + cur_ag_rank) % ag_world_sz
	265	+ if return_ag_res and step > 0:
	266	+ all_ag_res[cur_tensor_idx] = cur_step_tensor_to_send.clone()
	267	+
	268	+ # first linear forward: [s/(xcp), b, H/y] @ [H/y, e/x] -> [s/(xcp), b, e/x]
	269	+ cur_step_matmul_res = torch.matmul(cur_step_tensor_to_send, full_right_tensor)
	270	+ matmul_res[cur_tensor_idx] = cur_step_matmul_res
	271	+
	272	+ if step < ag_world_sz - 1:
	273	+ cur_step_rcv_handle.wait()
	274	+ cur_step_tensor_to_send = cur_step_rcv_input.clone()
	275	+
	276	+ final_matmul_res = torch.cat(matmul_res)
	277	+
	278	+ return final_matmul_res, all_ag_res
	279	+
	280	+ @staticmethod
	281	+ def _do_mm_overlap_reducescatter(activation_input, weight, bias, ag_comm_intf, rs_comm_intf):
	282	+ # [s/(x*cp), b, H/y] -> [s/cp, b, H/y]
	283	+ activation_input = activation_input.contiguous()
	284	+ total_input = sync_gather_along_first_dim(activation_input, ag_comm_intf)
	285	+ # [s/cp, b, H/y] @ [H/y, e/x] -> [s/cp, b, e/x]
	286	+ chunk_num = rs_comm_intf.get_comm_group_world_size()
	287	+ rs_chunks = []
	288	+ rs_handle_and_tmp_tensors = []
	289	+ # convert tuple to list to free used tensors ahead.
	290	+ seq_len, b, h = total_input.size()
	291	+ chunk_size = seq_len // chunk_num
	292	+ input_chunks = torch.reshape(total_input.view(chunk_size, -1, h).transpose(0, 1), (chunk_num, -1, h))
	293	+ rs_res = torch.empty((chunk_size, b, weight.size(1)), dtype=weight.dtype, device=weight.device)
	294	+ for idx in range(chunk_num):
	295	+ input_chunk = input_chunks[idx].reshape(chunk_size, -1, h)
	296	+ # [s/(cpy), b, H/y] @ [H/y, e/x] -> [s/(cpy), b, e/x]
	297	+ chunk_matmul_res = torch.matmul(input_chunk, weight).contiguous()
	298	+ if bias is not None:
	299	+ chunk_matmul_res += bias
	300	+
	301	+ # [s/(cpy), b, e/x]--rs--> [s/(cpy*y), b, e/x]
	302	+ rs_handle, rs_chunk = async_reduce_scatter_along_first_dim(
	303	+ chunk_matmul_res, rs_comm_intf
	304	+ )
	305	+ rs_chunks.append(rs_chunk)
	306	+ rs_handle_and_tmp_tensors.append((idx, rs_handle, chunk_matmul_res))
	307	+
	308	+ offset = 0
	309	+ sub_chunk_size = chunk_size // chunk_num
	310	+ for idx, rs_handle, chunk_matmul_res_tensor in rs_handle_and_tmp_tensors:
	311	+ if rs_handle:
	312	+ rs_handle.wait()
	313	+ chunk_matmul_res_tensor.untyped_storage().resize_(0)
	314	+ rs_res[offset:offset + sub_chunk_size] = rs_chunks[idx]
	315	+ offset += sub_chunk_size
	316	+
	317	+ # [s / (cp * y * y), b, e / x] -> [s/(cp*y), b, e/x]
	318	+ final_res = torch.reshape(rs_res.view(chunk_num, -1, weight.size(1)).transpose(0, 1), (chunk_size, -1, weight.size(1)))
	319	+ return final_res

Amindspeed/core/tensor_parallel/tp_2d/norm_factory.py+73-0

@@ -0,0 +1,73 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+from functools import wraps
	16	+from typing import List
	17	+
	18	+import torch
	19	+import torch.distributed as dist
	20	+from torch._utils import _flatten_dense_tensors
	21	+from torch._utils import _unflatten_dense_tensors
	22	+
	23	+from megatron.core.transformer import TransformerConfig
	24	+from megatron.core.utils import get_attr_wrapped_model
	25	+from megatron.training import get_args
	26	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm
	27	+from mindspeed.core.tensor_parallel.comm_group_api import TPYCollectiveComm
	28	+from mindspeed.core.tensor_parallel.tp_2d.layernorm_2d import LayerNorm2D
	29	+from mindspeed.core.tensor_parallel.tp_2d.rms_norm_2d import RMSNorm2D
	30	+
	31	+
	32	+def _allreduce_layernorm_grads_wrapper(function):
	33	+ @wraps(function)
	34	+ def wrapper(model: List[torch.nn.Module], config: TransformerConfig):
	35	+ function(model, config)
	36	+ layer_norm_2d_grads = []
	37	+ for model_chunk in model:
	38	+ for name, param in get_attr_wrapped_model(model_chunk, "named_parameters")():
	39	+ if param.requires_grad and getattr(param, "2d_tp", False):
	40	+ layer_norm_2d_grad = param.main_grad
	41	+ layer_norm_2d_grads.append(layer_norm_2d_grad.data)
	42	+
	43	+ if layer_norm_2d_grads:
	44	+ coalesced = _flatten_dense_tensors(layer_norm_2d_grads)
	45	+ torch.distributed.all_reduce(coalesced, group=TPXCollectiveComm.get_comm_group())
	46	+ for buf, synced in zip(
	47	+ layer_norm_2d_grads, _unflatten_dense_tensors(coalesced, layer_norm_2d_grads)
	48	+ ):
	49	+ buf.copy_(synced)
	50	+
	51	+ return wrapper
	52	+
	53	+
	54	+def get_norm_tp_2d(config):
	55	+ args = get_args()
	56	+ if args.normalization == "LayerNorm":
	57	+ return LayerNorm2D(
	58	+ config.hidden_size,
	59	+ eps=config.layernorm_epsilon,
	60	+ last_dim_split_comm_intf=TPYCollectiveComm(),
	61	+ )
	62	+ elif args.normalization == "RMSNorm":
	63	+ if args.apply_layernorm_1p:
	64	+ raise NotImplementedError(
	65	+ "RMSNorm does not currently support the layernorm_1p formulation."
	66	+ )
	67	+ return RMSNorm2D(
	68	+ config.hidden_size,
	69	+ eps=config.layernorm_epsilon,
	70	+ last_dim_split_comm_intf=TPYCollectiveComm(),
	71	+ )
	72	+ else:
	73	+ raise Exception(f"unsupported norm type '{args.normalization}'.")

Amindspeed/core/tensor_parallel/tp_2d/parallel_linear_2d.py+196-0

@@ -0,0 +1,196 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+from typing import Callable
	16	+
	17	+import torch
	18	+
	19	+from megatron.core import ModelParallelConfig
	20	+from megatron.core.tensor_parallel.layers import _initialize_affine_weight_gpu
	21	+from megatron.core.utils import divide
	22	+from mindspeed.core.tensor_parallel.comm_group_api import CollectiveCommIntf
	23	+from mindspeed.core.tensor_parallel.comm_group_api import OverlapCollectiveIntf
	24	+from mindspeed.core.tensor_parallel.layers import _initialize_affine_weight_cpu_2d
	25	+from mindspeed.core.tensor_parallel.tp_2d.linear_2d_split_along_first_dim import Linear2DSplitAlongFirstDim
	26	+
	27	+
	28	+class ParallelLinear2D(torch.nn.Module):
	29	+ """Linear2D layer with row and column parallelism.
	30	+
	31	+ The linear layer is defined as Y = XA + b. A is parallelized along
	32	+ its second dimension as A = [A_1, ..., A_p].
	33	+
	34	+ Arguments:
	35	+ input_size: first dimension of matrix A.
	36	+ output_size: second dimension of matrix A.
	37	+
	38	+ Keyword Arguments
	39	+ bias: If true, add bias
	40	+ gather_output: If true, call all-gather on output and make Y available
	41	+ to all GPUs, otherwise, every GPU will have its output
	42	+ which is Y_i = XA_i
	43	+ init_method: method to initialize weights. Note that bias is always set
	44	+ to zero.
	45	+ stride: For the strided linear layers.
	46	+ keep_master_weight_for_test: This was added for testing and should be
	47	+ set to False. It returns the master weights
	48	+ used for initialization.
	49	+ skip_bias_add: If True, do not add the bias term, instead
	50	+ return it to be added by the caller. This
	51	+ enables performance optimations where bias can
	52	+ be fused with other elementwise operations.
	53	+ skip_weight_param_allocation: If True, weight parameter is not allocated and must be passed
	54	+ as a keyword argument `weight` during the forward pass. Note
	55	+ that this does not affect bias, which will be allocated if
	56	+ bias is True. Defaults to False.
	57	+ is_expert: If True, the layer is treated as an MoE expert layer.
	58	+ config: ModelParallelConfig object
	59	+ tp_comm_buffer_name: Communication buffer name is not used in
	60	+ non-Transformer-Engine modules.
	61	+ partition_dim: divide with dim, column parallel set 0, row parallel set 1
	62	+
	63	+ """
	64	+
	65	+ def __init__(
	66	+ self,
	67	+ input_size,
	68	+ output_size,
	69	+ *,
	70	+ config: ModelParallelConfig,
	71	+ init_method: Callable,
	72	+ add_bias=True,
	73	+ gather_output=False,
	74	+ stride=1,
	75	+ keep_master_weight_for_test=False,
	76	+ skip_bias_add=True,
	77	+ skip_weight_param_allocation: bool = False,
	78	+ is_expert: bool = False,
	79	+ ag_comm_intf: CollectiveCommIntf = None,
	80	+ ag_sd_rcv_overlap_comm_intf: OverlapCollectiveIntf = None,
	81	+ rs_comm_intf: CollectiveCommIntf = None,
	82	+ rs_sd_rcv_overlap_comm_intf: OverlapCollectiveIntf = None,
	83	+ enable_overlap_ag_with_matmul=False,
	84	+ enable_overlap_matmul_with_rs=False,
	85	+ partition_dim: int = 0,
	86	+ ):
	87	+ super().__init__()
	88	+ self.mp_config: ModelParallelConfig = config
	89	+ self.para_init_method = init_method
	90	+ self.stride = stride
	91	+ self.keep_master_weight_for_test = keep_master_weight_for_test
	92	+ self.add_bias = add_bias
	93	+ self.input_size = input_size
	94	+ self.output_size = output_size
	95	+ self.ag_comm_intf = ag_comm_intf
	96	+ self.rs_comm_intf = rs_comm_intf
	97	+ self.ag_comm_world_sz = ag_comm_intf.get_comm_group_world_size()
	98	+ self.rs_comm_world_sz = rs_comm_intf.get_comm_group_world_size()
	99	+ # when AG comm group is small, do overlap AG with matmul.
	100	+ self.enable_overlap_ag_with_matmul = enable_overlap_ag_with_matmul
	101	+ self.enable_overlap_matmul_with_rs = enable_overlap_matmul_with_rs
	102	+ self.ag_overlap_comm_intf = ag_sd_rcv_overlap_comm_intf
	103	+ self.rs_sd_rcv_overlap_comm_intf = rs_sd_rcv_overlap_comm_intf
	104	+
	105	+ if input_size % self.rs_comm_world_sz:
	106	+ raise AssertionError("input size should be divisible by tp-y")
	107	+ if output_size % self.ag_comm_world_sz:
	108	+ raise AssertionError("output size should be divisible by tp-x")
	109	+
	110	+ self.input_size_per_partition = divide(input_size, self.rs_comm_world_sz)
	111	+ self.output_size_per_partition = divide(output_size, self.ag_comm_world_sz)
	112	+ self.skip_bias_add = skip_bias_add
	113	+ self.is_expert = is_expert
	114	+ self.expert_parallel = config.expert_model_parallel_size > 1
	115	+ self.gradient_accumulation_fusion = config.gradient_accumulation_fusion
	116	+ if config.sequence_parallel:
	117	+ raise RuntimeError(
	118	+ "Nd_matmul cannot be used with sequence_parallel."
	119	+ "If you want to train long sequences, "
	120	+ "you can use ulysess or context_parallel that is compatible with nd_matmul."
	121	+ )
	122	+ self.partition_dim = partition_dim
	123	+ self.init_linear_weights()
	124	+
	125	+ def init_linear_weights(self):
	126	+ init_with_cpu = self.mp_config.use_cpu_initialization
	127	+ device = None if init_with_cpu else torch.cuda.current_device()
	128	+
	129	+ self.weight = torch.nn.Parameter(
	130	+ torch.empty(
	131	+ self.output_size_per_partition,
	132	+ self.input_size_per_partition,
	133	+ device=device,
	134	+ dtype=self.mp_config.params_dtype,
	135	+ )
	136	+ )
	137	+ if self.add_bias:
	138	+ self.bias = torch.nn.Parameter(
	139	+ torch.empty(self.output_size_per_partition, dtype=self.mp_config.params_dtype, device=device)
	140	+ )
	141	+ else:
	142	+ self.register_parameter("bias", None)
	143	+
	144	+ if init_with_cpu and self.mp_config.perform_initialization:
	145	+ _initialize_affine_weight_cpu_2d(
	146	+ self.weight,
	147	+ self.output_size,
	148	+ self.input_size,
	149	+ self.input_size_per_partition,
	150	+ self.output_size_per_partition,
	151	+ self.partition_dim,
	152	+ self.para_init_method,
	153	+ stride=self.stride,
	154	+ return_master_weight=self.keep_master_weight_for_test,
	155	+ params_dtype=self.mp_config.params_dtype,
	156	+ )
	157	+ elif self.mp_config.perform_initialization:
	158	+ _initialize_affine_weight_gpu(
	159	+ self.weight,
	160	+ self.para_init_method,
	161	+ partition_dim=self.partition_dim,
	162	+ stride=self.stride,
	163	+ expert_parallel=False,
	164	+ )
	165	+
	166	+ setattr(self.weight, "allreduce", True)
	167	+
	168	+ if self.add_bias and self.mp_config.perform_initialization:
	169	+ with torch.no_grad():
	170	+ self.bias.zero_()
	171	+
	172	+ setattr(self.bias, "allreduce", True)
	173	+ setattr(self.bias, "sequence_parallel", False)
	174	+
	175	+ def forward(self, activation_input):
	176	+ matmul_output = Linear2DSplitAlongFirstDim.apply(
	177	+ activation_input,
	178	+ self.weight,
	179	+ self.bias,
	180	+ self.ag_comm_intf,
	181	+ self.ag_overlap_comm_intf,
	182	+ self.rs_comm_intf,
	183	+ self.rs_sd_rcv_overlap_comm_intf,
	184	+ self.enable_overlap_ag_with_matmul,
	185	+ self.enable_overlap_matmul_with_rs,
	186	+ self.gradient_accumulation_fusion,
	187	+ )
	188	+
	189	+ if not self.skip_bias_add:
	190	+ output = (matmul_output + self.bias) if self.bias is not None else matmul_output
	191	+ output_bias = None
	192	+ else:
	193	+ output = matmul_output
	194	+ output_bias = self.bias
	195	+
	196	+ return output, output_bias

Amindspeed/core/tensor_parallel/tp_2d/rms_norm_2d.py+98-0

@@ -0,0 +1,98 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+from typing import Any
	16	+from typing import Tuple
	17	+
	18	+import torch
	19	+import torch.distributed as dist
	20	+from torch import Tensor
	21	+from torch import nn
	22	+from torch.cuda.amp import custom_bwd
	23	+from torch.cuda.amp import custom_fwd
	24	+from megatron.core.utils import divide
	25	+from mindspeed.core.tensor_parallel.comm_group_api import CollectiveCommIntf
	26	+from mindspeed.core.tensor_parallel.comm_group_api import TPYCollectiveComm
	27	+
	28	+
	29	+class RMSNorm2D(torch.nn.Module):
	30	+
	31	+ def __init__(self,
	32	+ hidden_size: int,
	33	+ eps: float = 1e-6,
	34	+ last_dim_split_comm_intf: CollectiveCommIntf = TPYCollectiveComm()):
	35	+ """RMS Normaliation 2d module
	36	+
	37	+ Args:
	38	+ hidden_size (int): The width of input, i.e. hidden size
	39	+ eps (float): epsilon to use for the norm, default to 1e-6
	40	+ last_dim_split_comm_intf: All-reduce at last dim comm intf.
	41	+ """
	42	+ super().__init__()
	43	+ self.eps = eps
	44	+ self.hidden_size = hidden_size
	45	+ self.last_dim_split_comm_intf = last_dim_split_comm_intf
	46	+ self.last_dim_split_comm_world_sz = self.last_dim_split_comm_intf.get_comm_group_world_size()
	47	+ # partitioning dimension
	48	+ self.partitioned_dim = divide(hidden_size, self.last_dim_split_comm_world_sz)
	49	+ self.weight = nn.Parameter(torch.ones(self.partitioned_dim))
	50	+
	51	+ setattr(self.weight, "2d_tp", True)
	52	+
	53	+ def forward(self, x):
	54	+ return _ParallelRMSNorm2D.apply(
	55	+ x,
	56	+ self.weight,
	57	+ self.eps,
	58	+ self.hidden_size,
	59	+ self.last_dim_split_comm_intf,
	60	+ )
	61	+
	62	+
	63	+class _ParallelRMSNorm2D(torch.autograd.Function):
	64	+ @staticmethod
	65	+ @custom_fwd
	66	+ def forward(
	67	+ ctx: Any,
	68	+ input_: Tensor,
	69	+ weight,
	70	+ epsilon,
	71	+ hidden_size: int,
	72	+ last_dim_split_comm_intf: CollectiveCommIntf,
	73	+ ) -> Tensor:
	74	+ # input_ inner: [s/cp, b, h/xy]
	75	+ # input_ outer: [s/(cp*x), b, h/y]
	76	+ ctx.last_dim_split_comm_intf = last_dim_split_comm_intf
	77	+ ctx.hidden_size = hidden_size
	78	+ pow_mean = input_.float().pow(2).mean(-1, keepdim=True)
	79	+ torch.distributed.all_reduce(pow_mean, group=last_dim_split_comm_intf.get_comm_group())
	80	+ var_x = torch.rsqrt(pow_mean.div_(last_dim_split_comm_intf.get_comm_group_world_size()) + epsilon)
	81	+ norm_x = torch.mul(input_, var_x).type_as(input_)
	82	+ ctx.save_for_backward(norm_x, var_x, weight)
	83	+ return norm_x * weight
	84	+
	85	+ @staticmethod
	86	+ @custom_bwd
	87	+ def backward(ctx: Any, output_grad: Tensor) -> Tuple[Tensor, ...]:
	88	+ x, var_x, weight = ctx.saved_tensors
	89	+ grad_norm_x = torch.mul(output_grad, weight)
	90	+ output_grad_mul_x_sum = torch.sum(grad_norm_x * x, dim=-1, keepdim=True)
	91	+ handle_grad = torch.distributed.all_reduce(
	92	+ output_grad_mul_x_sum, group=ctx.last_dim_split_comm_intf.get_comm_group(), async_op=True)
	93	+ # calculate grad_weight
	94	+ grad_weight = torch.mul(output_grad, x)
	95	+ if handle_grad:
	96	+ handle_grad.wait()
	97	+ grad_input = (grad_norm_x - x * (output_grad_mul_x_sum / ctx.hidden_size)) * var_x
	98	+ return grad_input, grad_weight, None, None, None

Amindspeed/core/tensor_parallel_y_union_cp.py+99-0

@@ -0,0 +1,99 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+
	16	+import torch
	17	+import torch.distributed
	18	+
	19	+from megatron.core.parallel_state import get_nccl_options
	20	+
	21	+from mindspeed.core.mindspeed_parallel_group import MindspeedParallelGroup
	22	+from mindspeed.core.simple_parallel_cfg import SimpleParallelCfg
	23	+from mindspeed.core.singleton_meta import SingletonMeta
	24	+
	25	+
	26	+class TensorParallelYUnionCP(MindspeedParallelGroup, metaclass=SingletonMeta):
	27	+ def __init__(
	28	+ self,
	29	+ parallel_cfg: SimpleParallelCfg = None,
	30	+ pg_name: str = None,
	31	+ overlap_gp_name: str = None,
	32	+ nccl_comm_cfgs=None,
	33	+ ):
	34	+ super().__init__(parallel_cfg, pg_name, overlap_gp_name, nccl_comm_cfgs)
	35	+
	36	+ @staticmethod
	37	+ def init_group(
	38	+ parallel_cfg: SimpleParallelCfg,
	39	+ pg_name: str,
	40	+ overlap_gp_name: str = None,
	41	+ nccl_comm_cfgs=None,
	42	+ ):
	43	+ pp = parallel_cfg.pp
	44	+ tp = parallel_cfg.tp
	45	+ cp = parallel_cfg.cp
	46	+ tp_x = parallel_cfg.tp_x
	47	+
	48	+ rank = torch.distributed.get_rank()
	49	+ world_size: int = torch.distributed.get_world_size()
	50	+ num_pp_groups: int = world_size // pp
	51	+ dp = world_size // (tp * pp * cp)
	52	+
	53	+ all_cp_grps = []
	54	+ for i in range(pp):
	55	+ for j in range(dp):
	56	+ start_rank = i * num_pp_groups + j * tp * cp
	57	+ end_rank = i * num_pp_groups + (j + 1) * tp * cp
	58	+ for k in range(tp):
	59	+ ranks = range(start_rank + k, end_rank, tp)
	60	+ all_cp_grps.append(ranks)
	61	+
	62	+ all_tp_x_grps = []
	63	+ all_tp_y_grps = []
	64	+ num_tp_grps: int = world_size // tp
	65	+ for i in range(num_tp_grps):
	66	+ for j in range(tp // tp_x):
	67	+ ranks = range(i * tp + j * tp_x, i * tp + (j + 1) * tp_x)
	68	+ all_tp_x_grps.append(list(ranks))
	69	+
	70	+ num_tp_x_group = tp // tp_x
	71	+ for j in range(tp // num_tp_x_group):
	72	+ ranks = range(i * tp + j, (i + 1) * tp, tp_x)
	73	+ all_tp_y_grps.append(list(ranks))
	74	+
	75	+ # Build the tensor model-parallel-y-cp groups.
	76	+ res_group, res_overlap_group, res_global_ranks = None, None, None
	77	+ all_tp_y_cp_grps = []
	78	+ for cp_grp in all_cp_grps:
	79	+ tp_y_cp_grp_ranks = []
	80	+ for cp_rank in cp_grp:
	81	+ for tp_y_grp in all_tp_y_grps:
	82	+ if cp_rank in tp_y_grp:
	83	+ tp_y_cp_grp_ranks += tp_y_grp
	84	+
	85	+ if tp_y_cp_grp_ranks not in all_tp_y_cp_grps:
	86	+ all_tp_y_cp_grps.append(tp_y_cp_grp_ranks)
	87	+ res_group = torch.distributed.new_group(
	88	+ tp_y_cp_grp_ranks, pg_options=get_nccl_options(pg_name, nccl_comm_cfgs)
	89	+ )
	90	+ # for send-recv in parallel
	91	+ if overlap_gp_name:
	92	+ res_overlap_group = torch.distributed.new_group(
	93	+ tp_y_cp_grp_ranks,
	94	+ pg_options=get_nccl_options(overlap_gp_name, nccl_comm_cfgs),
	95	+ )
	96	+ if rank in tp_y_cp_grp_ranks:
	97	+ res_global_ranks = tp_y_cp_grp_ranks
	98	+ return res_group, res_global_ranks, res_overlap_group
	99	+ return res_group, res_global_ranks, res_overlap_group

Mmindspeed/core/transformer/attention.py+48-1

	@@ -10,10 +10,18 @@ from megatron.core.transformer import TransformerConfig, ModuleSpec, build_modul
10	from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules	10	from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
11	from megatron.core.transformer.enums import AttnMaskType	11	from megatron.core.transformer.enums import AttnMaskType
12	from megatron.core import mpu	12	from megatron.core import mpu
		13	+from megatron.core.utils import divide
13	from megatron.training import get_args	14	from megatron.training import get_args
14		15
15	from mindspeed.core.context_parallel.ulysses_context_parallel import UlyssesContextAttention	16	from mindspeed.core.context_parallel.ulysses_context_parallel import UlyssesContextAttention
16	from mindspeed.core.parallel_state import get_context_parallel_group_for_hybrid_ulysses	17	from mindspeed.core.parallel_state import get_context_parallel_group_for_hybrid_ulysses
		18	+from mindspeed.core.context_parallel.ulysses_context_parallel import UlyssesContextAttention
		19	+from mindspeed.core.parallel_state import get_context_parallel_group_for_hybrid_ulysses, \
		20	+ get_tensor_model_parallel_world_size_for_nd1_dim1
		21	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPXOverlapCollectiveComm, \
		22	+ TPYCollectiveComm, TPYOverlapCollectiveComm
		23	+from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP
		24	+from mindspeed.core.tensor_parallel.tp_2d.parallel_linear_2d import ParallelLinear2D
17		25
18		26
19	@dataclass	27	@dataclass
	@@ -34,7 +42,11 @@ def attention_init_wrapper(fn):
34		42
35	args = get_args()	43	args = get_args()
36	if args.context_parallel_size > 1 and args.context_parallel_algo in ['ulysses_cp_algo', 'hybrid_cp_algo', 'hybrid_adaptive_cp_algo']:	44	if args.context_parallel_size > 1 and args.context_parallel_algo in ['ulysses_cp_algo', 'hybrid_cp_algo', 'hybrid_adaptive_cp_algo']:
37	- ulysses_group = mpu.get_context_parallel_group()	45	+ if args.tp_2d:
		46	+ tp_y_cp = TensorParallelYUnionCP()
		47	+ ulysses_group = tp_y_cp.group
		48	+ else:
		49	+ ulysses_group = mpu.get_context_parallel_group()
38	if args.context_parallel_algo == 'hybrid_cp_algo' or args.context_parallel_algo == 'hybrid_adaptive_cp_algo':	50	if args.context_parallel_algo == 'hybrid_cp_algo' or args.context_parallel_algo == 'hybrid_adaptive_cp_algo':
39	ulysses_group = get_context_parallel_group_for_hybrid_ulysses()	51	ulysses_group = get_context_parallel_group_for_hybrid_ulysses()
40	self.core_attention = UlyssesContextAttention(self.core_attention, ulysses_group)	52	self.core_attention = UlyssesContextAttention(self.core_attention, ulysses_group)
	@@ -143,6 +155,41 @@ def self_attention_init_wrapper(fn):
143	tp_comm_buffer_name='proj',	155	tp_comm_buffer_name='proj',
144	)	156	)
145		157
		158	+ if args.tp_2d:
		159	+ attn_heads_split_num = get_tensor_model_parallel_world_size_for_nd1_dim1()
		160	+ self.num_attention_heads_per_partition = divide(self.config.num_attention_heads, attn_heads_split_num)
		161	+ self.num_query_groups_per_partition = divide(self.config.num_query_groups, attn_heads_split_num)
		162	+ self.linear_qkv = ParallelLinear2D(
		163	+ self.config.hidden_size,
		164	+ self.query_projection_size + 2 * self.kv_projection_size,
		165	+ config=self.config,
		166	+ init_method=self.config.init_method,
		167	+ add_bias=self.config.add_bias_linear,
		168	+ skip_bias_add=True,
		169	+ ag_comm_intf=TPXCollectiveComm,
		170	+ ag_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
		171	+ rs_comm_intf=TPYCollectiveComm,
		172	+ rs_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
		173	+ enable_overlap_ag_with_matmul=False,
		174	+ enable_overlap_matmul_with_rs=False,
		175	+ partition_dim=0
		176	+ )
		177	+ self.linear_proj = ParallelLinear2D(
		178	+ self.query_projection_size,
		179	+ self.config.hidden_size,
		180	+ config=self.config,
		181	+ init_method=self.config.output_layer_init_method,
		182	+ add_bias=self.config.add_bias_linear,
		183	+ skip_bias_add=True,
		184	+ ag_comm_intf=TPYCollectiveComm,
		185	+ ag_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
		186	+ rs_comm_intf=TPXCollectiveComm,
		187	+ rs_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
		188	+ enable_overlap_ag_with_matmul=args.enable_overlap_ag_with_matmul,
		189	+ enable_overlap_matmul_with_rs=False,
		190	+ partition_dim=1
		191	+ )
		192	+
146	return wrapper	193	return wrapper
147		194
148		195

Mmindspeed/core/transformer/custom_layers/transformer_engine.py+31-10

	@@ -4,6 +4,11 @@ import torch.nn as nn
4		4
5	from megatron.core.transformer.transformer_config import TransformerConfig	5	from megatron.core.transformer.transformer_config import TransformerConfig
6	from megatron.legacy.model.rms_norm import RMSNorm	6	from megatron.legacy.model.rms_norm import RMSNorm
		7	+from megatron.training import get_args
		8	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm
		9	+from mindspeed.core.tensor_parallel.comm_group_api import TPYCollectiveComm
		10	+from mindspeed.core.tensor_parallel.tp_2d.layernorm_2d import LayerNorm2D
		11	+from mindspeed.core.tensor_parallel.tp_2d.rms_norm_2d import RMSNorm2D
7		12
8		13
9	class PTNorm:	14	class PTNorm:
	@@ -14,18 +19,34 @@ class PTNorm:
14	def __new__(	19	def __new__(
15	cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5,	20	cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5,
16	):	21	):
		22	+ args = get_args()
17	if config.normalization == "LayerNorm":	23	if config.normalization == "LayerNorm":
18	- instance = nn.LayerNorm(	24	+ if args.tp_2d:
19	- normalized_shape=hidden_size,	25	+ instance = LayerNorm2D(
20	- eps=eps,	26	+ hidden_size,
21	- )	27	+ eps=eps,
		28	+ last_dim_split_comm_intf=TPYCollectiveComm(),
		29	+ )
		30	+ else:
		31	+ instance = nn.LayerNorm(
		32	+ normalized_shape=hidden_size,
		33	+ eps=eps,
		34	+ )
22	elif config.normalization == "RMSNorm":	35	elif config.normalization == "RMSNorm":
23	- instance = RMSNorm(	36	+ if args.tp_2d:
24	- dim=hidden_size,	37	+ instance = RMSNorm2D(
25	- eps=eps,	38	+ hidden_size,
26	- sequence_parallel=config.sequence_parallel,	39	+ eps=eps,
27	- )	40	+ last_dim_split_comm_intf=TPYCollectiveComm(),
28	- instance.use_fused_rmsnorm = True	41	+ )
		42	+ instance.use_fused_rmsnorm = False
		43	+ else:
		44	+ instance = RMSNorm(
		45	+ dim=hidden_size,
		46	+ eps=eps,
		47	+ sequence_parallel=config.sequence_parallel,
		48	+ )
		49	+ instance.use_fused_rmsnorm = True
29	else:	50	else:
30	raise Exception('Only LayerNorm and RMSNorm are curently supported')	51	raise Exception('Only LayerNorm and RMSNorm are curently supported')
31		52

Mmindspeed/core/transformer/dot_product_attention.py+35-14

	@@ -23,6 +23,7 @@ from mindspeed.core.parallel_state import (get_context_parallel_group_for_hybrid
23	get_ring_ranks_for_inter_window_dkv,	23	get_ring_ranks_for_inter_window_dkv,
24	get_ring_group_for_intra_window,	24	get_ring_group_for_intra_window,
25	get_ring_group_for_intra_window_send_recv_overlap)	25	get_ring_group_for_intra_window_send_recv_overlap)
		26	+from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP
26	from mindspeed.model.transformer import get_attention_mask	27	from mindspeed.model.transformer import get_attention_mask
27	from mindspeed.utils import get_actual_seq_len	28	from mindspeed.utils import get_actual_seq_len
28	from mindspeed.core.context_parallel.adaptive_context_parallel import adaptive_attn_context_parallel	29	from mindspeed.core.context_parallel.adaptive_context_parallel import adaptive_attn_context_parallel
	@@ -108,23 +109,35 @@ def dot_product_attention_forward(
108		109
109	scale = 1.0 / math.sqrt(self.hidden_size_per_attention_head) if self.scale_mask_softmax.scale is None else self.softmax_scale	110	scale = 1.0 / math.sqrt(self.hidden_size_per_attention_head) if self.scale_mask_softmax.scale is None else self.softmax_scale
110		111
111	- if args.context_parallel_size > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo',	112	+ cp_expanded_by_2d_tp = args.tp_2d and args.tp_y > 1
		113	+ if cp_expanded_by_2d_tp:
		114	+ tp_y_cp_sz = args.context_parallel_size * args.tp_y
		115	+ else:
		116	+ tp_y_cp_sz = args.context_parallel_size
		117	+ if tp_y_cp_sz > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo',
112	'adaptive_cp_algo', 'hybrid_adaptive_cp_algo']:	118	'adaptive_cp_algo', 'hybrid_adaptive_cp_algo']:
113	in_hybrid_mode = False	119	in_hybrid_mode = False
114	if get_context_parallel_group_for_hybrid_ring(check_initialized=False) is not None:	120	if get_context_parallel_group_for_hybrid_ring(check_initialized=False) is not None:
115	in_hybrid_mode = True	121	in_hybrid_mode = True
116	-	122	+
117	if not in_hybrid_mode:	123	if not in_hybrid_mode:
118	- cp_group = mpu.get_context_parallel_group()	124	+ if cp_expanded_by_2d_tp:
119	- cp_size = mpu.get_context_parallel_world_size()	125	+ tp_y_cp = TensorParallelYUnionCP()
120	- rank = mpu.get_context_parallel_rank()	126	+ cp_group = tp_y_cp.group
121	- cp_global_ranks = mpu.get_context_parallel_global_ranks()	127	+ cp_size = tp_y_cp.get_parallel_group_world_size()
		128	+ rank = tp_y_cp.get_parallel_rank()
		129	+ cp_global_ranks = tp_y_cp.global_ranks
		130	+ else:
		131	+ cp_group = mpu.get_context_parallel_group()
		132	+ cp_size = mpu.get_context_parallel_world_size()
		133	+ rank = mpu.get_context_parallel_rank()
		134	+ cp_global_ranks = mpu.get_context_parallel_global_ranks()
122	else:	135	else:
123	cp_group = get_context_parallel_group_for_hybrid_ring()	136	cp_group = get_context_parallel_group_for_hybrid_ring()
124	cp_size = get_context_parallel_for_hybrid_ring_world_size()	137	cp_size = get_context_parallel_for_hybrid_ring_world_size()
125	rank = get_context_parallel_for_hybrid_ring_rank()	138	rank = get_context_parallel_for_hybrid_ring_rank()
126	cp_global_ranks = get_context_parallel_for_hybrid_ring_global_ranks()	139	cp_global_ranks = get_context_parallel_for_hybrid_ring_global_ranks()
127	-	140	+
128	cp_para = dict()	141	cp_para = dict()
129	cp_para['causal'] = args.cp_attention_mask_type == 'causal'	142	cp_para['causal'] = args.cp_attention_mask_type == 'causal'
130	cp_para['cp_group'] = cp_group	143	cp_para['cp_group'] = cp_group
	@@ -134,15 +147,23 @@ def dot_product_attention_forward(
134	query, key, value = [rearrange(x, 's b h d -> s b (h d)') for x in [query, key, value]]	147	query, key, value = [rearrange(x, 's b h d -> s b (h d)') for x in [query, key, value]]
135	if args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo']:	148	if args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo']:
136	cp_para['cp_global_ranks'] = cp_global_ranks	149	cp_para['cp_global_ranks'] = cp_global_ranks
137	- cp_para['cp_group_for_send_recv_overlap'] = mpu.get_context_parallel_group_for_send_recv_overlap() \	150	+ if args.use_cp_send_recv_overlap:
138	- if args.use_cp_send_recv_overlap else None	151	+ if cp_expanded_by_2d_tp:
		152	+ cp_para['cp_group_for_send_recv_overlap'] = tp_y_cp.overlap_group
		153	+ else:
		154	+ cp_para['cp_group_for_send_recv_overlap'] = mpu.get_context_parallel_group_for_send_recv_overlap()
		155	+ else:
		156	+ cp_para['cp_group_for_send_recv_overlap'] = None
139	cp_para['pse'] = self.pse	157	cp_para['pse'] = self.pse
140	cp_para['pse_type'] = self.pse_type	158	cp_para['pse_type'] = self.pse_type
141	- cp_para['cp_inner_ranks'] = get_ring_ranks_for_intra_window()	159	+
142	- cp_para['cp_outer_ranks'] = get_ring_ranks_for_inter_window_kv()	160	+ if args.context_parallel_size > 1 and not args.tp_2d:
143	- cp_para['cp_dkv_outer_ranks'] = get_ring_ranks_for_inter_window_dkv()	161	+ cp_para['cp_inner_ranks'] = get_ring_ranks_for_intra_window()
144	- cp_para['cp_group_for_intra_window'] = get_ring_group_for_intra_window()	162	+ cp_para['cp_outer_ranks'] = get_ring_ranks_for_inter_window_kv()
145	- cp_para['cp_group_for_intra_window_send_recv_overlap'] = get_ring_group_for_intra_window_send_recv_overlap()	163	+ cp_para['cp_dkv_outer_ranks'] = get_ring_ranks_for_inter_window_dkv()
		164	+ cp_para['cp_group_for_intra_window'] = get_ring_group_for_intra_window()
		165	+ cp_para['cp_group_for_intra_window_send_recv_overlap'] = get_ring_group_for_intra_window_send_recv_overlap()
		166	+
146	output = ringattn_context_parallel(query, key, value, n_head, cp_para, scale, attention_mask, self.attention_dropout.p,	167	output = ringattn_context_parallel(query, key, value, n_head, cp_para, scale, attention_mask, self.attention_dropout.p,
147	actual_seq_len, actual_seq_len)	168	actual_seq_len, actual_seq_len)
148	else:	169	else:

Amindspeed/core/transformer/mlp.py+63-0

@@ -0,0 +1,63 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+from functools import wraps
	16	+from megatron.training import get_args
	17	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPXOverlapCollectiveComm, \
	18	+ TPYCollectiveComm, TPYOverlapCollectiveComm
	19	+from mindspeed.core.tensor_parallel.tp_2d.parallel_linear_2d import ParallelLinear2D
	20	+
	21	+
	22	+def mlp_init_wrapper(fn):
	23	+ @wraps(fn)
	24	+ def wrapper(self, arg, *kwargs):
	25	+ fn(self, arg, *kwargs)
	26	+ args = get_args()
	27	+ if args.tp_2d:
	28	+ ffn_hidden_size = self.config.ffn_hidden_size
	29	+ if self.config.gated_linear_unit:
	30	+ ffn_hidden_size *= 2
	31	+ self.linear_fc1 = ParallelLinear2D(
	32	+ self.config.hidden_size,
	33	+ ffn_hidden_size,
	34	+ config=self.config,
	35	+ init_method=self.config.init_method,
	36	+ add_bias=self.config.add_bias_linear,
	37	+ skip_bias_add=True,
	38	+ is_expert=False,
	39	+ ag_comm_intf=TPXCollectiveComm,
	40	+ ag_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
	41	+ rs_comm_intf=TPYCollectiveComm,
	42	+ rs_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
	43	+ enable_overlap_ag_with_matmul=False,
	44	+ enable_overlap_matmul_with_rs=args.enable_overlap_matmul_with_rs,
	45	+ partition_dim=0,
	46	+ )
	47	+ self.linear_fc2 = ParallelLinear2D(
	48	+ self.config.ffn_hidden_size,
	49	+ self.config.hidden_size,
	50	+ config=self.config,
	51	+ init_method=self.config.output_layer_init_method,
	52	+ add_bias=self.config.add_bias_linear,
	53	+ skip_bias_add=True,
	54	+ is_expert=False,
	55	+ ag_comm_intf=TPYCollectiveComm,
	56	+ ag_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
	57	+ rs_comm_intf=TPXCollectiveComm,
	58	+ rs_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
	59	+ enable_overlap_ag_with_matmul=args.enable_overlap_ag_with_matmul,
	60	+ enable_overlap_matmul_with_rs=False,
	61	+ partition_dim=1
	62	+ )
	63	+ return wrapper

Mmindspeed/core/transformer/transformer_block.py+14-0

	@@ -8,6 +8,9 @@ from megatron.core.transformer.module import MegatronModule
8	from megatron.core.transformer.spec_utils import build_module	8	from megatron.core.transformer.spec_utils import build_module
9	from megatron.training import get_args	9	from megatron.training import get_args
10	from megatron.core.transformer.custom_layers.transformer_engine import TENorm	10	from megatron.core.transformer.custom_layers.transformer_engine import TENorm
		11	+from mindspeed.core.tensor_parallel.comm_autograd_function import auto_grad_sync_gather_along_last_dim, \
		12	+ auto_grad_sync_gather_along_first_dim
		13	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPYCollectiveComm
11		14
12		15
13	def transformer_block_checkpointed_forward_wrapper(forward_func):	16	def transformer_block_checkpointed_forward_wrapper(forward_func):
	@@ -210,3 +213,14 @@ def _build_layers(self):
210	hidden_size=self.config.hidden_size,	213	hidden_size=self.config.hidden_size,
211	eps=self.config.layernorm_epsilon,	214	eps=self.config.layernorm_epsilon,
212	)	215	)
		216	+
		217	+
		218	+def transformer_block_forward_wrapper(fn):
		219	+ @wraps(fn)
		220	+ def wrapper(args, *kwargs):
		221	+ hidden_states = fn(args, *kwargs)
		222	+ if get_args().tp_2d and parallel_state.is_pipeline_last_stage():
		223	+ hidden_states = auto_grad_sync_gather_along_first_dim(hidden_states, TPXCollectiveComm)
		224	+ hidden_states = auto_grad_sync_gather_along_last_dim(hidden_states, TPYCollectiveComm)
		225	+ return hidden_states
		226	+ return wrapper

Mmindspeed/megatron_adaptor.py+31-1

	@@ -177,6 +177,12 @@ def mcore_models_adaptation(aspm, mindspeed_args):
177	get_gpt_layer_local_spec)	177	get_gpt_layer_local_spec)
178	aspm.register_patch('megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding.__init__',	178	aspm.register_patch('megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding.__init__',
179	rotary_embedding_init_wrapper)	179	rotary_embedding_init_wrapper)
		180	+ from .core.models.common.embeddings.language_model_embedding import language_model_embedding_forward_wrapper
		181	+ aspm.register_patch('megatron.core.models.common.embeddings.language_model_embedding.LanguageModelEmbedding.forward',
		182	+ language_model_embedding_forward_wrapper)
		183	+ from .core.models.common.embeddings.rotary_pos_embedding import rotary_embedding_get_rotary_seq_len_wrapper
		184	+ aspm.register_patch('megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding.get_rotary_seq_len',
		185	+ rotary_embedding_get_rotary_seq_len_wrapper)
180		186
181	if not mindspeed_args.automated_pipeline and mindspeed_args.noop_layers:	187	if not mindspeed_args.automated_pipeline and mindspeed_args.noop_layers:
182	from .core.transformer.transformer_block import _build_layers	188	from .core.transformer.transformer_block import _build_layers
	@@ -187,6 +193,7 @@ def mcore_models_adaptation(aspm, mindspeed_args):
187	aspm.register_patch('megatron.training.training.num_floating_point_operations', num_floating_point_wrapper)	193	aspm.register_patch('megatron.training.training.num_floating_point_operations', num_floating_point_wrapper)
188	aspm.register_patch('megatron.core.transformer.moe.moe_utils.track_moe_metrics', track_moe_metrics)	194	aspm.register_patch('megatron.core.transformer.moe.moe_utils.track_moe_metrics', track_moe_metrics)
189		195
		196	+
190	if mindspeed_args.recompute_norm:	197	if mindspeed_args.recompute_norm:
191	from .core.models.gpt.gpt_layer_specs import build_norm_recompute_layer_wrapper	198	from .core.models.gpt.gpt_layer_specs import build_norm_recompute_layer_wrapper
192	aspm.register_patch('megatron.core.transformer.transformer_block.TransformerBlock._build_layers', build_norm_recompute_layer_wrapper)	199	aspm.register_patch('megatron.core.transformer.transformer_block.TransformerBlock._build_layers', build_norm_recompute_layer_wrapper)
	@@ -220,6 +227,8 @@ def mcore_transformer_adaptation(aspm):
220	from .core.transformer.transformer_block import transformer_block_checkpointed_forward_wrapper	227	from .core.transformer.transformer_block import transformer_block_checkpointed_forward_wrapper
221	from .core.transformer.transformer import parallel_transformer_layer_init_wrapper	228	from .core.transformer.transformer import parallel_transformer_layer_init_wrapper
222	from .core.transformer.transformer import core_mlp_forward_wrapper	229	from .core.transformer.transformer import core_mlp_forward_wrapper
		230	+ from .core.transformer.mlp import mlp_init_wrapper
		231	+ from .core.transformer.transformer_block import transformer_block_forward_wrapper
223	aspm.register_patch('megatron.core.transformer.attention.SelfAttentionSubmodules', SelfAttentionSubmodules)	232	aspm.register_patch('megatron.core.transformer.attention.SelfAttentionSubmodules', SelfAttentionSubmodules)
224	aspm.register_patch('megatron.core.transformer.attention.SelfAttention.__init__', self_attention_init_wrapper)	233	aspm.register_patch('megatron.core.transformer.attention.SelfAttention.__init__', self_attention_init_wrapper)
225	aspm.register_patch("megatron.core.transformer.attention.Attention.forward", attention_forward_wrapper)	234	aspm.register_patch("megatron.core.transformer.attention.Attention.forward", attention_forward_wrapper)
	@@ -232,6 +241,9 @@ def mcore_transformer_adaptation(aspm):
232	parallel_transformer_layer_init_wrapper)	241	parallel_transformer_layer_init_wrapper)
233	aspm.register_patch('megatron.core.transformer.mlp.MLP.forward',	242	aspm.register_patch('megatron.core.transformer.mlp.MLP.forward',
234	core_mlp_forward_wrapper)	243	core_mlp_forward_wrapper)
		244	+ aspm.register_patch('megatron.core.transformer.mlp.MLP.__init__', mlp_init_wrapper)
		245	+ aspm.register_patch('megatron.core.transformer.transformer_block.TransformerBlock.forward',
		246	+ transformer_block_forward_wrapper)
235		247
236		248
237	def mcore_parallel_state_adaptation(aspm):	249	def mcore_parallel_state_adaptation(aspm):
	@@ -302,7 +314,7 @@ def mcore_optimizer_adapation(aspm):
302		314
303	def mcore_pipeline_parallel_adaptation(aspm):	315	def mcore_pipeline_parallel_adaptation(aspm):
304	from .core.pipeline_parallel.p2p_communication import _communicate_shapes	316	from .core.pipeline_parallel.p2p_communication import _communicate_shapes
305	- from .core.pipeline_parallel.schedules import get_forward_backward_func_wrapper	317	+ from .core.pipeline_parallel.schedules import get_tensor_shapes_wrapper, get_forward_backward_func_wrapper
306	from .core.performance.auto_pipeline_perf.schedules import get_forward_backward_func_decorator, \	318	from .core.performance.auto_pipeline_perf.schedules import get_forward_backward_func_decorator, \
307	backward_step_decorator, forward_step_decorator	319	backward_step_decorator, forward_step_decorator
308		320
	@@ -316,6 +328,8 @@ def mcore_pipeline_parallel_adaptation(aspm):
316	forward_step_decorator)	328	forward_step_decorator)
317	aspm.register_patch('megatron.core.pipeline_parallel.p2p_communication._communicate_shapes',	329	aspm.register_patch('megatron.core.pipeline_parallel.p2p_communication._communicate_shapes',
318	_communicate_shapes)	330	_communicate_shapes)
		331	+ aspm.register_patch('megatron.core.pipeline_parallel.schedules.get_tensor_shapes',
		332	+ get_tensor_shapes_wrapper)
319		333
320		334
321	def mcore_multiparam_pipeline_parallel_adaptation(aspm, mindspeed_args):	335	def mcore_multiparam_pipeline_parallel_adaptation(aspm, mindspeed_args):
	@@ -676,6 +690,21 @@ def zero3_adaptation(aspm, args):
676	distributed_data_parallel_zero_grad_wrapper)	690	distributed_data_parallel_zero_grad_wrapper)
677		691
678		692
		693	+def tensor_2d_adaptation(aspm, args):
		694	+ if args.tp_2d:
		695	+ from mindspeed.core.tensor_parallel.tp_2d.norm_factory import get_norm_tp_2d
		696	+ from mindspeed.core.tensor_parallel.tp_2d.norm_factory import _allreduce_layernorm_grads_wrapper
		697	+ from mindspeed.core.models.common.embeddings.rotary_pos_embedding import rotary_embedding_forward_wrapper
		698	+ from mindspeed.core.pipeline_parallel.flexible_schedules import forward_backward_pipelining_with_interleaving_patch
		699	+ aspm.register_patch('megatron.legacy.model.utils.get_norm', get_norm_tp_2d)
		700	+ aspm.register_patch('megatron.core.distributed.finalize_model_grads._allreduce_layernorm_grads',
		701	+ _allreduce_layernorm_grads_wrapper)
		702	+ aspm.register_patch('megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding.forward',
		703	+ rotary_embedding_forward_wrapper)
		704	+ aspm.register_patch('megatron.core.pipeline_parallel.schedules.forward_backward_pipelining_with_interleaving',
		705	+ forward_backward_pipelining_with_interleaving_patch)
		706	+
		707	+
679	def adaptation_l0(aspm):	708	def adaptation_l0(aspm):
680	"""	709	"""
681	The minimum patch set for megatron to adapt to NPU	710	The minimum patch set for megatron to adapt to NPU
	@@ -731,6 +760,7 @@ def adaptation_l2(aspm, mindspeed_args):
731	mcore_moe_adaptation(aspm, mindspeed_args)	760	mcore_moe_adaptation(aspm, mindspeed_args)
732	deepspeed_moe_adaptation(aspm, mindspeed_args)	761	deepspeed_moe_adaptation(aspm, mindspeed_args)
733	zero3_adaptation(aspm, mindspeed_args)	762	zero3_adaptation(aspm, mindspeed_args)
		763	+ tensor_2d_adaptation(aspm, mindspeed_args)
734		764
735		765
736	def exe_adaptation():	766	def exe_adaptation():

Mmindspeed/model/language_model.py+12-1

	@@ -17,6 +17,10 @@ from functools import wraps
17	import torch	17	import torch
18	from megatron.training import get_args	18	from megatron.training import get_args
19	from megatron.core import mpu, tensor_parallel	19	from megatron.core import mpu, tensor_parallel
		20	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPYCollectiveComm
		21	+from mindspeed.core.tensor_parallel.comm_autograd_function import (auto_grad_sync_gather_along_first_dim,
		22	+ auto_grad_sync_gather_along_last_dim,
		23	+ auto_grad_scatter_along_first_dim_then_last_dim)
20	from mindspeed.moe.utils import get_slice_indices_from_disorder_to_order, get_slice_indices_from_order_to_disorder	24	from mindspeed.moe.utils import get_slice_indices_from_disorder_to_order, get_slice_indices_from_order_to_disorder
21		25
22		26
	@@ -41,6 +45,10 @@ def parallel_lm_logits(
41	if args.use_nd_matmul:	45	if args.use_nd_matmul:
42	input_parallel = tensor_parallel.gather_from_tensor_model_parallel_region(input_parallel)	46	input_parallel = tensor_parallel.gather_from_tensor_model_parallel_region(input_parallel)
43		47
		48	+ if args.tp_2d:
		49	+ input_parallel = auto_grad_sync_gather_along_first_dim(input_parallel, TPXCollectiveComm)
		50	+ input_parallel = auto_grad_sync_gather_along_last_dim(input_parallel, TPYCollectiveComm)
		51	+
44	# Matrix multiply.	52	# Matrix multiply.
45	logits_parallel = tensor_parallel.linear_with_grad_accumulation_and_async_allreduce(	53	logits_parallel = tensor_parallel.linear_with_grad_accumulation_and_async_allreduce(
46	input=input_parallel,	54	input=input_parallel,
	@@ -50,7 +58,6 @@ def parallel_lm_logits(
50	async_grad_allreduce=async_grad_allreduce,	58	async_grad_allreduce=async_grad_allreduce,
51	sequence_parallel=args.sequence_parallel)	59	sequence_parallel=args.sequence_parallel)
52	# Gather if needed.	60	# Gather if needed.
53	-
54	if parallel_output:	61	if parallel_output:
55	return logits_parallel	62	return logits_parallel
56		63
	@@ -63,6 +70,10 @@ def embedding_forward_wrapper(forward):
63	encoder_input = forward(self, args, *kwargs)	70	encoder_input = forward(self, args, *kwargs)
64	if get_args().use_nd_matmul:	71	if get_args().use_nd_matmul:
65	encoder_input = tensor_parallel.scatter_to_tensor_model_parallel_region(encoder_input)	72	encoder_input = tensor_parallel.scatter_to_tensor_model_parallel_region(encoder_input)
		73	+ if get_args().tp_2d:
		74	+ encoder_input = auto_grad_scatter_along_first_dim_then_last_dim(
		75	+ encoder_input, TPXCollectiveComm, TPYCollectiveComm
		76	+ )
66	return encoder_input	77	return encoder_input
67	return wrapper	78	return wrapper
68		79

Mmindspeed/model/transformer.py+122-15

	@@ -55,7 +55,14 @@ from mindspeed.core.parallel_state import (get_context_parallel_group_for_hybrid
55	get_ring_group_for_intra_window,	55	get_ring_group_for_intra_window,
56	get_ring_group_for_intra_window_send_recv_overlap)	56	get_ring_group_for_intra_window_send_recv_overlap)
57	from mindspeed.core.fusions.fused_bias_swiglu import fused_swiglu	57	from mindspeed.core.fusions.fused_bias_swiglu import fused_swiglu
		58	+from mindspeed.core.parallel_state import get_tensor_model_parallel_world_size_for_nd1_dim1
		59	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm
		60	+from mindspeed.core.tensor_parallel.comm_group_api import TPXOverlapCollectiveComm
		61	+from mindspeed.core.tensor_parallel.comm_group_api import TPYCollectiveComm
		62	+from mindspeed.core.tensor_parallel.comm_group_api import TPYOverlapCollectiveComm
		63	+from mindspeed.core.tensor_parallel.tp_2d.parallel_linear_2d import ParallelLinear2D
58	from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput	64	from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput
		65	+from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP
59	from mindspeed.moe.ampipe.ampipe import AttMoEPipe	66	from mindspeed.moe.ampipe.ampipe import AttMoEPipe
60	from mindspeed.ops.fusion_attention_v2 import npu_fusion_attention	67	from mindspeed.ops.fusion_attention_v2 import npu_fusion_attention
61	from mindspeed.core.tensor_parallel.layers import Nd_ParallelLinear	68	from mindspeed.core.tensor_parallel.layers import Nd_ParallelLinear
	@@ -68,7 +75,6 @@ from mindspeed.moe.utils import (get_slice_indices_from_order_to_disorder,
68	from mindspeed.core.context_parallel.adaptive_context_parallel import adaptive_attn_context_parallel	75	from mindspeed.core.context_parallel.adaptive_context_parallel import adaptive_attn_context_parallel
69	from mindspeed.core.context_parallel.utils import get_scheduling_info	76	from mindspeed.core.context_parallel.utils import get_scheduling_info
70		77
71	-
72	try:	78	try:
73	from einops import rearrange	79	from einops import rearrange
74	except ImportError:	80	except ImportError:
	@@ -767,6 +773,36 @@ def parallel_mlp_init_wrapper(fn):
767	is_expert=is_expert,	773	is_expert=is_expert,
768	matmul_id=2	774	matmul_id=2
769	)	775	)
		776	+ elif _args.tp_2d:
		777	+ self.dense_h_to_4h = ParallelLinear2D(
		778	+ config.hidden_size,
		779	+ ffn_hidden_size,
		780	+ config=config,
		781	+ init_method=config.init_method,
		782	+ add_bias=self.add_bias,
		783	+ skip_bias_add=True,
		784	+ is_expert=is_expert,
		785	+ ag_comm_intf=TPXCollectiveComm,
		786	+ ag_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
		787	+ rs_comm_intf=TPYCollectiveComm,
		788	+ rs_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
		789	+ enable_overlap_ag_with_matmul=False,
		790	+ enable_overlap_matmul_with_rs=_args.enable_overlap_matmul_with_rs,
		791	+ partition_dim=0)
		792	+ self.dense_4h_to_h = ParallelLinear2D(
		793	+ config.ffn_hidden_size,
		794	+ config.hidden_size,
		795	+ config=config,
		796	+ init_method=config.output_layer_init_method,
		797	+ add_bias=self.add_bias,
		798	+ skip_bias_add=True,
		799	+ ag_comm_intf=TPYCollectiveComm,
		800	+ ag_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
		801	+ rs_comm_intf=TPXCollectiveComm,
		802	+ rs_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
		803	+ enable_overlap_ag_with_matmul=_args.enable_overlap_ag_with_matmul,
		804	+ enable_overlap_matmul_with_rs=False,
		805	+ partition_dim=1)
770	else:	806	else:
771	self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(	807	self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(
772	config.hidden_size,	808	config.hidden_size,
	@@ -973,18 +1009,29 @@ def flash_self_attention_forward(self, q, k, v, attention_mask):
973	scale = 1.0 / math.sqrt(head_dim) if self.softmax_scale is None else self.softmax_scale	1009	scale = 1.0 / math.sqrt(head_dim) if self.softmax_scale is None else self.softmax_scale
974	except Exception as e:	1010	except Exception as e:
975	raise ValueError('Invalid head_dim: {}'.format(head_dim)) from e	1011	raise ValueError('Invalid head_dim: {}'.format(head_dim)) from e
976	-	1012	+ cp_expanded_by_2d_tp = args.tp_2d and args.tp_y > 1
977	- if args.context_parallel_size > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo',	1013	+ if cp_expanded_by_2d_tp:
		1014	+ tp_y_cp_sz = args.context_parallel_size * args.tp_y
		1015	+ else:
		1016	+ tp_y_cp_sz = args.context_parallel_size
		1017	+ if tp_y_cp_sz > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo',
978	'adaptive_cp_algo', 'hybrid_adaptive_cp_algo']:	1018	'adaptive_cp_algo', 'hybrid_adaptive_cp_algo']:
979	in_hybrid_mode = False	1019	in_hybrid_mode = False
980	if get_context_parallel_group_for_hybrid_ring(check_initialized=False) is not None:	1020	if get_context_parallel_group_for_hybrid_ring(check_initialized=False) is not None:
981	in_hybrid_mode = True	1021	in_hybrid_mode = True
982		1022
983	if not in_hybrid_mode:	1023	if not in_hybrid_mode:
984	- cp_group = mpu.get_context_parallel_group()	1024	+ if cp_expanded_by_2d_tp:
985	- cp_size = mpu.get_context_parallel_world_size()	1025	+ tp_y_cp = TensorParallelYUnionCP()
986	- rank = mpu.get_context_parallel_rank()	1026	+ cp_group = tp_y_cp.group
987	- cp_global_ranks = mpu.get_context_parallel_global_ranks()	1027	+ cp_size = tp_y_cp.get_parallel_group_world_size()
		1028	+ rank = tp_y_cp.get_parallel_rank()
		1029	+ cp_global_ranks = tp_y_cp.global_ranks
		1030	+ else:
		1031	+ cp_group = mpu.get_context_parallel_group()
		1032	+ cp_size = mpu.get_context_parallel_world_size()
		1033	+ rank = mpu.get_context_parallel_rank()
		1034	+ cp_global_ranks = mpu.get_context_parallel_global_ranks()
988	else:	1035	else:
989	cp_group = get_context_parallel_group_for_hybrid_ring()	1036	cp_group = get_context_parallel_group_for_hybrid_ring()
990	cp_size = get_context_parallel_for_hybrid_ring_world_size()	1037	cp_size = get_context_parallel_for_hybrid_ring_world_size()
	@@ -999,15 +1046,21 @@ def flash_self_attention_forward(self, q, k, v, attention_mask):
999		1046
1000	if args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo']:	1047	if args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo']:
1001	cp_para['cp_global_ranks'] = cp_global_ranks	1048	cp_para['cp_global_ranks'] = cp_global_ranks
1002	- cp_para['cp_group_for_send_recv_overlap'] = mpu.get_context_parallel_group_for_send_recv_overlap() \	1049	+ if args.use_cp_send_recv_overlap:
1003	- if args.use_cp_send_recv_overlap else None	1050	+ if cp_expanded_by_2d_tp:
		1051	+ cp_para['cp_group_for_send_recv_overlap'] = tp_y_cp.overlap_group
		1052	+ else:
		1053	+ cp_para['cp_group_for_send_recv_overlap'] = mpu.get_context_parallel_group_for_send_recv_overlap()
		1054	+ else:
		1055	+ cp_para['cp_group_for_send_recv_overlap'] = None
1004	cp_para['pse'] = self.pse	1056	cp_para['pse'] = self.pse
1005	cp_para['pse_type'] = self.pse_type	1057	cp_para['pse_type'] = self.pse_type
1006	- cp_para['cp_inner_ranks'] = get_ring_ranks_for_intra_window()	1058	+ if args.context_parallel_size > 1 and not args.tp_2d:
1007	- cp_para['cp_outer_ranks'] = get_ring_ranks_for_inter_window_kv()	1059	+ cp_para['cp_inner_ranks'] = get_ring_ranks_for_intra_window()
1008	- cp_para['cp_dkv_outer_ranks'] = get_ring_ranks_for_inter_window_dkv()	1060	+ cp_para['cp_outer_ranks'] = get_ring_ranks_for_inter_window_kv()
1009	- cp_para['cp_group_for_intra_window'] = get_ring_group_for_intra_window()	1061	+ cp_para['cp_dkv_outer_ranks'] = get_ring_ranks_for_inter_window_dkv()
1010	- cp_para['cp_group_for_intra_window_send_recv_overlap'] = get_ring_group_for_intra_window_send_recv_overlap()	1062	+ cp_para['cp_group_for_intra_window'] = get_ring_group_for_intra_window()
		1063	+ cp_para['cp_group_for_intra_window_send_recv_overlap'] = get_ring_group_for_intra_window_send_recv_overlap()
1011	output = ringattn_context_parallel(q, k, v, head_num, cp_para, scale, attention_mask, self.dropout_p)	1064	output = ringattn_context_parallel(q, k, v, head_num, cp_para, scale, attention_mask, self.dropout_p)
1012	else:	1065	else:
1013	cp_para['scheduling_info'] = get_scheduling_info()	1066	cp_para['scheduling_info'] = get_scheduling_info()
	@@ -1047,7 +1100,27 @@ def parallel_attention_init_wrapper(fn):
1047	@wraps(fn)	1100	@wraps(fn)
1048	def wrapper(self, args, *kwargs):	1101	def wrapper(self, args, *kwargs):
1049	fn(self, args, *kwargs)	1102	fn(self, args, *kwargs)
		1103	+ # patch for 2d-tp
1050	config = args[0]	1104	config = args[0]
		1105	+ training_args = get_args()
		1106	+ attn_heads_split_num = (
		1107	+ get_tensor_model_parallel_world_size_for_nd1_dim1()
		1108	+ if training_args.tp_2d
		1109	+ else mpu.get_tensor_model_parallel_world_size()
		1110	+ )
		1111	+
		1112	+ # Per attention head and per partition values.
		1113	+ self.num_attention_heads_per_partition = config.num_attention_heads // attn_heads_split_num
		1114	+
		1115	+ if self.group_query_attention:
		1116	+ if training_args.num_query_groups % attn_heads_split_num != 0:
		1117	+ raise NotImplementedError(
		1118	+ "Currently the num_query_groups should be a multiple of the tensor parallel size"
		1119	+ )
		1120	+ self.num_query_groups_per_partition = training_args.num_query_groups // attn_heads_split_num
		1121	+ else:
		1122	+ self.num_query_groups_per_partition = self.num_attention_heads_per_partition
		1123	+
1051	query_projection_size = config.kv_channels * config.num_attention_heads	1124	query_projection_size = config.kv_channels * config.num_attention_heads
1052	_args = get_args()	1125	_args = get_args()
1053	if _args.group_query_attention:	1126	if _args.group_query_attention:
	@@ -1057,7 +1130,11 @@ def parallel_attention_init_wrapper(fn):
1057	# qkv bias	1130	# qkv bias
1058	bias = _args.add_qkv_bias or _args.add_bias_linear	1131	bias = _args.add_qkv_bias or _args.add_bias_linear
1059	if args[0].context_parallel_size > 1 and args[0].context_parallel_algo in ['ulysses_cp_algo', 'hybrid_cp_algo', 'hybrid_adaptive_cp_algo']:	1132	if args[0].context_parallel_size > 1 and args[0].context_parallel_algo in ['ulysses_cp_algo', 'hybrid_cp_algo', 'hybrid_adaptive_cp_algo']:
1060	- ulysses_group = mpu.get_context_parallel_group()	1133	+ if training_args.tp_2d:
		1134	+ tp_y_cp = TensorParallelYUnionCP()
		1135	+ ulysses_group = tp_y_cp.group
		1136	+ else:
		1137	+ ulysses_group = mpu.get_context_parallel_group()
1061	if args[0].context_parallel_algo == 'hybrid_cp_algo' or args[0].context_parallel_algo == 'hybrid_adaptive_cp_algo':	1138	if args[0].context_parallel_algo == 'hybrid_cp_algo' or args[0].context_parallel_algo == 'hybrid_adaptive_cp_algo':
1062	ulysses_group = get_context_parallel_group_for_hybrid_ulysses()	1139	ulysses_group = get_context_parallel_group_for_hybrid_ulysses()
1063	if self.use_flash_attn:	1140	if self.use_flash_attn:
	@@ -1076,6 +1153,21 @@ def parallel_attention_init_wrapper(fn):
1076	input_is_parallel=True,	1153	input_is_parallel=True,
1077	matmul_id=1	1154	matmul_id=1
1078	)	1155	)
		1156	+ elif _args.tp_2d:
		1157	+ self.query_key_value = ParallelLinear2D(
		1158	+ config.hidden_size,
		1159	+ query_projection_size + 2 * kv_projection_size,
		1160	+ config=config,
		1161	+ init_method=config.init_method,
		1162	+ add_bias=bias,
		1163	+ skip_bias_add=True,
		1164	+ ag_comm_intf=TPXCollectiveComm,
		1165	+ ag_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
		1166	+ rs_comm_intf=TPYCollectiveComm,
		1167	+ rs_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
		1168	+ enable_overlap_ag_with_matmul=False,
		1169	+ enable_overlap_matmul_with_rs=False,
		1170	+ partition_dim=0)
1079	else:	1171	else:
1080	self.query_key_value = tensor_parallel.ColumnParallelLinear(	1172	self.query_key_value = tensor_parallel.ColumnParallelLinear(
1081	config.hidden_size,	1173	config.hidden_size,
	@@ -1100,6 +1192,21 @@ def parallel_attention_init_wrapper(fn):
1100	input_is_parallel=True,	1192	input_is_parallel=True,
1101	matmul_id=2	1193	matmul_id=2
1102	)	1194	)
		1195	+ elif _args.tp_2d:
		1196	+ self.dense = ParallelLinear2D(
		1197	+ query_projection_size,
		1198	+ config.hidden_size,
		1199	+ config=config,
		1200	+ init_method=config.output_layer_init_method,
		1201	+ add_bias=bias,
		1202	+ skip_bias_add=True,
		1203	+ ag_comm_intf=TPYCollectiveComm,
		1204	+ ag_sd_rcv_overlap_comm_intf=TPYOverlapCollectiveComm,
		1205	+ rs_comm_intf=TPXCollectiveComm,
		1206	+ rs_sd_rcv_overlap_comm_intf=TPXOverlapCollectiveComm,
		1207	+ enable_overlap_ag_with_matmul=_args.enable_overlap_ag_with_matmul,
		1208	+ enable_overlap_matmul_with_rs=False,
		1209	+ partition_dim=1)
1103	else:	1210	else:
1104	self.dense = tensor_parallel.RowParallelLinear(	1211	self.dense = tensor_parallel.RowParallelLinear(
1105	query_projection_size,	1212	query_projection_size,

Mmindspeed/utils.py+49-2

	@@ -78,13 +78,16 @@ def get_batch_on_this_cp_rank(batch):
78	position_ids = position_ids.transpose(0, 1).contiguous()	78	position_ids = position_ids.transpose(0, 1).contiguous()
79	set_position_ids(position_ids)	79	set_position_ids(position_ids)
80		80
81	- cp_size = args.context_parallel_size	81	+ tp_y_cp_size = args.context_parallel_size * args.tp_y if args.tp_2d else args.context_parallel_size
82	- if not cp_size > 1:	82	+ if not tp_y_cp_size > 1:
83	return batch	83	return batch
84		84
		85	+ cp_expanded_by_2d_tp = args.tp_y > 1
85	if args.context_parallel_algo == 'megatron_cp_algo':	86	if args.context_parallel_algo == 'megatron_cp_algo':
86	if args.cp_attention_mask_type == 'general':	87	if args.cp_attention_mask_type == 'general':
87	batch = _get_batch_on_this_cp_rank_in_megatron_cp_general(batch)	88	batch = _get_batch_on_this_cp_rank_in_megatron_cp_general(batch)
		89	+ elif cp_expanded_by_2d_tp:
		90	+ batch = _get_batch_on_this_tp_y_cp_rank_in_megatron_cp(batch)
88	else:	91	else:
89	batch = _get_batch_on_this_cp_rank_in_megatron_cp(batch)	92	batch = _get_batch_on_this_cp_rank_in_megatron_cp(batch)
90	elif args.context_parallel_algo == 'ulysses_cp_algo':	93	elif args.context_parallel_algo == 'ulysses_cp_algo':
	@@ -415,4 +418,48 @@ def _get_batch_on_this_cp_rank_in_hybrid_adaptive_cp(batch):
415	index = torch.tensor(remapped_seq_order[which_per * per:(which_per + 1) * per], device=val.device)	418	index = torch.tensor(remapped_seq_order[which_per * per:(which_per + 1) * per], device=val.device)
416	val = val.index_select(seq_dim, index)	419	val = val.index_select(seq_dim, index)
417	batch[key] = val	420	batch[key] = val
		421	+
		422	+ return batch
		423	+
		424	+
		425	+def _get_batch_on_this_tp_y_cp_rank_in_megatron_cp(batch):
		426	+ cp_rank = mpu.get_context_parallel_rank()
		427	+ cp_size = mpu.get_context_parallel_world_size()
		428	+
		429	+ args = get_args()
		430	+ tp_y_cp_size = args.context_parallel_size * args.tp_y
		431	+
		432	+ for key, val in batch.items():
		433	+ if key == 'attention_mask' or val is None:
		434	+ continue
		435	+
		436	+ seq_dim = 1
		437	+ b = val.shape[0]
		438	+
		439	+ # [b, s] -> [b, 2tp_y_cp_sz, s/(2tp_y_cp_sz)]
		440	+ val = val.view(
		441	+ *val.shape[0:seq_dim],
		442	+ 2 * tp_y_cp_size,
		443	+ val.shape[seq_dim] // (2 * tp_y_cp_size),
		444	+ *val.shape[(seq_dim + 1):],
		445	+ )
		446	+
		447	+ rearrange_index = []
		448	+ for i in range(tp_y_cp_size):
		449	+ rearrange_index.extend([i, 2 * tp_y_cp_size - 1 - i])
		450	+ rearrange_idx_tensor = torch.tensor(rearrange_index, device=val.device)
		451	+
		452	+ val = val.index_select(seq_dim, index=rearrange_idx_tensor)
		453	+
		454	+ # [b, 2 * tp_y_cp_sz, s / (2 * tp_y_cp_sz)] -> [b, cp, s/cp]
		455	+ val = val.view(
		456	+ *val.shape[0:seq_dim],
		457	+ cp_size,
		458	+ val.shape[seq_dim] // cp_size,
		459	+ *val.shape[(seq_dim + 1):],
		460	+ )
		461	+ # [b, 1, s/cp] -> [b, s/cp]
		462	+ val = val[:, cp_rank].view(b, -1)
		463	+ batch[key] = val
		464	+
418	return batch	465	return batch

Asources/images/tensor-parallel-2d.png+0-0

Atests_extend/system_tests/llama2/pretrain_llama2_70B_8k_tp8_x4_y2_2d.sh+105-0

@@ -0,0 +1,105 @@
	1	+#!/bin/bash
	2	+export ASCEND_LAUNCH_BLOCKING=1
	3	+export CUDA_DEVICE_MAX_CONNECTIONS=1
	4	+
	5	+GPUS_PER_NODE=8
	6	+MASTER_ADDR=<master_ip_address>
	7	+MASTER_PORT=6000
	8	+NNODES=2
	9	+NODE_RANK=<local_rank>
	10	+WORLD_SIZE=$(($GPUS_PER_NODE * $NNODES))
	11	+
	12	+DISTRIBUTED_ARGS="
	13	+ --nproc_per_node $GPUS_PER_NODE \
	14	+ --nnodes $NNODES \
	15	+ --node_rank $NODE_RANK \
	16	+ --master_addr $MASTER_ADDR \
	17	+ --master_port $MASTER_PORT
	18	+"
	19	+
	20	+echo "NODE_RANK ${NODE_RANK}"
	21	+
	22	+DATA_PATH="/home/dataset/llama2/alpaca_text_document"
	23	+TOKENIZER_MODEL="/home/dataset/model/llama-2-7b-hf/tokenizer.model"
	24	+
	25	+TP=8
	26	+PP=1
	27	+CP=1
	28	+SEQ_LEN=$((8*1024))
	29	+NUM_LAYERS=32
	30	+
	31	+DISTRIBUTED_ARGS="
	32	+ --nproc_per_node $GPUS_PER_NODE \
	33	+ --nnodes $NNODES \
	34	+ --node_rank $NODE_RANK \
	35	+ --master_addr $MASTER_ADDR \
	36	+ --master_port $MASTER_PORT
	37	+"
	38	+
	39	+GPT_ARGS="
	40	+ --tensor-model-parallel-size ${TP} \
	41	+ --pipeline-model-parallel-size ${PP} \
	42	+ --tp-2d \
	43	+ --tp-x 4 \
	44	+ --tp-y 2 \
	45	+ --context-parallel-size ${CP} \
	46	+ --context-parallel-algo megatron_cp_algo \
	47	+ --use-cpu-initialization \
	48	+ --num-layers ${NUM_LAYERS} \
	49	+ --hidden-size 4096 \
	50	+ --ffn-hidden-size 11008 \
	51	+ --num-attention-heads 32 \
	52	+ --tokenizer-type Llama2Tokenizer \
	53	+ --tokenizer-model ${TOKENIZER_MODEL} \
	54	+ --seq-length $SEQ_LEN \
	55	+ --max-position-embeddings $SEQ_LEN \
	56	+ --micro-batch-size 2 \
	57	+ --global-batch-size 16 \
	58	+ --make-vocab-size-divisible-by 1 \
	59	+ --lr 1.25e-6 \
	60	+ --train-iters 2100 \
	61	+ --lr-decay-style cosine \
	62	+ --untie-embeddings-and-output-weights \
	63	+ --disable-bias-linear \
	64	+ --attention-dropout 0.0 \
	65	+ --init-method-std 0.01 \
	66	+ --hidden-dropout 0.0 \
	67	+ --position-embedding-type rope \
	68	+ --normalization RMSNorm \
	69	+ --swiglu \
	70	+ --use-flash-attn \
	71	+ --no-masked-softmax-fusion \
	72	+ --attention-softmax-in-fp32 \
	73	+ --min-lr 1.25e-7 \
	74	+ --weight-decay 1e-1 \
	75	+ --lr-warmup-fraction 0.01 \
	76	+ --clip-grad 1.0 \
	77	+ --adam-beta1 0.9 \
	78	+ --initial-loss-scale 65536 \
	79	+ --adam-beta2 0.95 \
	80	+ --no-gradient-accumulation-fusion \
	81	+ --no-load-optim \
	82	+ --no-load-rng \
	83	+ --use-distributed-optimizer \
	84	+ --use-cp-send-recv-overlap \
	85	+ --bf16
	86	+"
	87	+
	88	+DATA_ARGS="
	89	+ --data-path $DATA_PATH \
	90	+ --split 100,0,0
	91	+"
	92	+
	93	+OUTPUT_ARGS="
	94	+ --log-interval 1 \
	95	+ --save-interval 10000 \
	96	+ --eval-interval 1000 \
	97	+ --eval-iters 0 \
	98	+"
	99	+
	100	+torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
	101	+ $GPT_ARGS \
	102	+ $DATA_ARGS \
	103	+ $OUTPUT_ARGS \
	104	+ --distributed-backend nccl \
	105	+ \| tee logs/train_llama_2d.log

Atests_extend/unit_tests/mindspeed/core/tensor_parallel/test_layernorm2d.py+146-0

@@ -0,0 +1,146 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+import pytest
	16	+import torch
	17	+import torch.nn as nn
	18	+import torch.optim as optim
	19	+from unit_tests.common import DistributedTest
	20	+from commons import set_random_seed
	21	+
	22	+from mindspeed import megatron_adaptor
	23	+
	24	+from megatron.core.parallel_state import destroy_model_parallel, initialize_model_parallel
	25	+from megatron.core import tensor_parallel
	26	+from megatron.core.tensor_parallel import mappings
	27	+from megatron.training.global_vars import set_args
	28	+from megatron.training.arguments import parse_args
	29	+from mindspeed.core.tensor_parallel.comm_autograd_function import auto_grad_scatter_along_first_dim, \
	30	+ auto_grad_scatter_along_last_dim, auto_grad_sync_gather_along_first_dim, auto_grad_sync_gather_along_last_dim
	31	+from mindspeed.core.tensor_parallel.comm_utils import sync_reduce_scatter_along_first_dim
	32	+
	33	+from mindspeed.core.tensor_parallel.tp_2d.layernorm_2d import LayerNorm2D
	34	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPYCollectiveComm
	35	+
	36	+
	37	+# 定义一个包含LayerNorm的简单模型
	38	+class SimpleModel(nn.Module):
	39	+ def __init__(self, norm_size):
	40	+ super(SimpleModel, self).__init__()
	41	+ self.layer_norm = nn.LayerNorm(normalized_shape=norm_size)
	42	+
	43	+ def forward(self, x):
	44	+ output = self.layer_norm(x)
	45	+ return output
	46	+
	47	+
	48	+class _OnlyAllGatherFromTensorParallelRegion(torch.autograd.Function):
	49	+ """Gather the input from model parallel region and concatenate."""
	50	+
	51	+ @staticmethod
	52	+ def symbolic(graph, input_):
	53	+ return mappings._gather_along_last_dim(input_)
	54	+
	55	+ @staticmethod
	56	+ def forward(ctx, input_):
	57	+ return mappings._gather_along_last_dim(input_,)
	58	+
	59	+ @staticmethod
	60	+ def backward(ctx, grad_output):
	61	+ return mappings._split_along_last_dim(grad_output)
	62	+
	63	+
	64	+def only_all_gather_last_dim_from_tensor_parallel_region(input_):
	65	+ return _OnlyAllGatherFromTensorParallelRegion.apply(input_)
	66	+
	67	+
	68	+class SimpleDistModel(nn.Module):
	69	+ def __init__(self, norm_size):
	70	+ super(SimpleDistModel, self).__init__()
	71	+ last_dim_split_comm_intf = TPYCollectiveComm()
	72	+ self.layer_norm = LayerNorm2D(norm_size, last_dim_split_comm_intf=last_dim_split_comm_intf)
	73	+
	74	+ def forward(self, x):
	75	+ output = self.layer_norm(x)
	76	+ return output
	77	+
	78	+
	79	+class TestLayernorm2dRsFirstDim(DistributedTest):
	80	+ world_size = 4
	81	+ diff_value = 1e-5
	82	+
	83	+ @staticmethod
	84	+ def get_layernorm_grad(dist_schedule, h, input_x, targets):
	85	+ # 生成一些随机输入数据和目标输出
	86	+ input_x_data = input_x.clone().detach()
	87	+ targets_data = targets.clone().detach()
	88	+ input_x_data.requires_grad_()
	89	+ # 创建模型实例
	90	+ if dist_schedule:
	91	+ model = SimpleDistModel(h).npu()
	92	+ else:
	93	+ model = SimpleModel(h).npu()
	94	+
	95	+ # 创建损失函数和优化器
	96	+ criterion = nn.MSELoss()
	97	+ optimizer = optim.SGD(model.parameters(), lr=0.01)
	98	+ # 前向传播：计算预测值
	99	+ if dist_schedule:
	100	+ # s,b,E -> s/x,b,h/y
	101	+ input_x_data = auto_grad_scatter_along_first_dim(input_x_data, TPXCollectiveComm)
	102	+ input_x_data = auto_grad_scatter_along_last_dim(input_x_data, TPYCollectiveComm)
	103	+ outputs = model(input_x_data)
	104	+ outputs = auto_grad_sync_gather_along_first_dim(outputs, TPXCollectiveComm)
	105	+ outputs = auto_grad_sync_gather_along_last_dim(outputs, TPYCollectiveComm)
	106	+ else:
	107	+ outputs = model(input_x_data)
	108	+
	109	+ # 计算损失
	110	+ loss = criterion(outputs, targets_data)
	111	+ # 反向传播：计算梯度
	112	+ loss.backward()
	113	+ return outputs, model.layer_norm.weight.grad, model.layer_norm.bias.grad
	114	+
	115	+ def test_layer_norm_should_be_allclose_give_same_input_when_norm_1d_2d(self):
	116	+ set_random_seed(1)
	117	+ args = parse_args(None, True)
	118	+ set_args(args)
	119	+ tp, tp_x, tp_y = 4, 2, 2
	120	+ pp = self.world_size // tp
	121	+ h = 8
	122	+ seq = 16
	123	+ b = 1
	124	+ args.tp_2d = True
	125	+ args.tp_x = tp_x
	126	+ args.tp_y = tp_y
	127	+ # 生成一些随机输入数据和目标输出
	128	+ input_x = torch.randn(seq, b, h, requires_grad=True).npu()
	129	+ targets = torch.randn(seq, b, h).npu()
	130	+ destroy_model_parallel()
	131	+ initialize_model_parallel(
	132	+ tensor_model_parallel_size=tp,
	133	+ pipeline_model_parallel_size=pp,
	134	+ virtual_pipeline_model_parallel_size=None,
	135	+ pipeline_model_parallel_split_rank=None)
	136	+ # 2d layer_norm
	137	+ output_2d, weight_grad_2d, bias_grad_2d = self.get_layernorm_grad(dist_schedule=1, h=h, input_x=input_x, targets=targets)
	138	+ # 1d layer_norm
	139	+ output_1d, weight_grad_1d, bias_grad_1d = self.get_layernorm_grad(dist_schedule=0, h=h, input_x=input_x, targets=targets)
	140	+ assert torch.allclose(output_2d, output_1d, rtol=self.diff_value, atol=self.diff_value)
	141	+ weight_grad_1d = tensor_parallel.scatter_to_tensor_model_parallel_region(weight_grad_1d)
	142	+ bias_grad_1d = tensor_parallel.scatter_to_tensor_model_parallel_region(bias_grad_1d)
	143	+ weight_grad_2d = sync_reduce_scatter_along_first_dim(weight_grad_2d, TPXCollectiveComm)
	144	+ bias_grad_2d = sync_reduce_scatter_along_first_dim(bias_grad_2d, TPXCollectiveComm)
	145	+ assert torch.allclose(weight_grad_2d, weight_grad_1d, rtol=self.diff_value, atol=self.diff_value)
	146	+ assert torch.allclose(bias_grad_2d, bias_grad_1d, rtol=self.diff_value, atol=self.diff_value)

Atests_extend/unit_tests/mindspeed/core/tensor_parallel/test_rms_norm2d.py+143-0

@@ -0,0 +1,143 @@
	1	+# Copyright 2024 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+import pytest
	16	+import torch
	17	+import torch.nn as nn
	18	+import torch.optim as optim
	19	+from unit_tests.common import DistributedTest
	20	+from commons import set_random_seed
	21	+
	22	+from mindspeed import megatron_adaptor
	23	+
	24	+from megatron.core.parallel_state import destroy_model_parallel, initialize_model_parallel
	25	+from megatron.core import tensor_parallel
	26	+from megatron.core.tensor_parallel import mappings
	27	+from megatron.training.global_vars import set_args
	28	+from megatron.training.arguments import parse_args
	29	+from megatron.legacy.model import RMSNorm
	30	+from mindspeed.core.tensor_parallel.comm_autograd_function import auto_grad_scatter_along_first_dim, \
	31	+ auto_grad_scatter_along_last_dim, auto_grad_sync_gather_along_first_dim, auto_grad_sync_gather_along_last_dim
	32	+from mindspeed.core.tensor_parallel.comm_utils import sync_reduce_scatter_along_first_dim
	33	+
	34	+from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPYCollectiveComm
	35	+from mindspeed.core.tensor_parallel.tp_2d.rms_norm_2d import RMSNorm2D
	36	+
	37	+
	38	+class SimpleRMSNormModel(nn.Module):
	39	+ def __init__(self, norm_size):
	40	+ super(SimpleRMSNormModel, self).__init__()
	41	+ self.rms_norm = RMSNorm(dim=norm_size)
	42	+
	43	+ def forward(self, x):
	44	+ output = self.rms_norm(x)
	45	+ return output
	46	+
	47	+
	48	+class _OnlyAllGatherFromTensorParallelRegion(torch.autograd.Function):
	49	+ """Gather the input from model parallel region and concatenate."""
	50	+
	51	+ @staticmethod
	52	+ def symbolic(graph, input_):
	53	+ return mappings._gather_along_last_dim(input_)
	54	+
	55	+ @staticmethod
	56	+ def forward(ctx, input_):
	57	+ return mappings._gather_along_last_dim(input_,)
	58	+
	59	+ @staticmethod
	60	+ def backward(ctx, grad_output):
	61	+ return mappings._split_along_last_dim(grad_output)
	62	+
	63	+
	64	+def only_all_gather_last_dim_from_tensor_parallel_region(input_):
	65	+ return _OnlyAllGatherFromTensorParallelRegion.apply(input_)
	66	+
	67	+
	68	+class SimpleDistRMSNormModel(nn.Module):
	69	+ def __init__(self, norm_size):
	70	+ super(SimpleDistRMSNormModel, self).__init__()
	71	+ last_dim_split_comm_intf = TPYCollectiveComm()
	72	+ self.rms_norm = RMSNorm2D(norm_size, last_dim_split_comm_intf=last_dim_split_comm_intf)
	73	+
	74	+ def forward(self, x):
	75	+ output = self.rms_norm(x)
	76	+ return output
	77	+
	78	+
	79	+class TestRMSNorm2dRsFirstDim(DistributedTest):
	80	+ world_size = 4
	81	+ diff_value = 1e-5
	82	+
	83	+ @staticmethod
	84	+ def get_rms_norm_grad(dist_schedule, h, input_x, targets):
	85	+ # 生成一些随机输入数据和目标输出
	86	+ input_x_data = input_x.clone().detach()
	87	+ targets_data = targets.clone().detach()
	88	+ input_x_data.requires_grad_()
	89	+ # 创建模型实例
	90	+ if dist_schedule:
	91	+ model = SimpleDistRMSNormModel(h).npu()
	92	+ else:
	93	+ model = SimpleRMSNormModel(h).npu()
	94	+
	95	+ # 创建损失函数和优化器
	96	+ criterion = nn.MSELoss()
	97	+ optimizer = optim.SGD(model.parameters(), lr=0.01)
	98	+ # 前向传播：计算预测值
	99	+ if dist_schedule:
	100	+ # s,b,E -> s/x,b,h/y
	101	+ input_x_data = auto_grad_scatter_along_first_dim(input_x_data, TPXCollectiveComm)
	102	+ input_x_data = auto_grad_scatter_along_last_dim(input_x_data, TPYCollectiveComm)
	103	+ outputs = model(input_x_data)
	104	+ outputs = auto_grad_sync_gather_along_first_dim(outputs, TPXCollectiveComm)
	105	+ outputs = auto_grad_sync_gather_along_last_dim(outputs, TPYCollectiveComm)
	106	+ else:
	107	+ outputs = model(input_x_data)
	108	+
	109	+ # 计算损失
	110	+ loss = criterion(outputs, targets_data)
	111	+ # 反向传播：计算梯度
	112	+ loss.backward()
	113	+ return outputs, model.rms_norm.weight.grad
	114	+
	115	+ def test_rms_norm_should_be_allclose_give_same_input_when_norm_1d_2d(self):
	116	+ set_random_seed(1)
	117	+ args = parse_args(None, True)
	118	+ set_args(args)
	119	+ tp, tp_x, tp_y = 4, 2, 2
	120	+ pp = self.world_size // tp
	121	+ h = 16
	122	+ seq = 8
	123	+ b = 1
	124	+ args.tp_2d = True
	125	+ args.tp_x = tp_x
	126	+ args.tp_y = tp_y
	127	+ # 生成一些随机输入数据和目标输出
	128	+ input_x = torch.randn(seq, b, h, requires_grad=True).npu()
	129	+ targets = torch.randn(seq, b, h).npu()
	130	+ destroy_model_parallel()
	131	+ initialize_model_parallel(
	132	+ tensor_model_parallel_size=tp,
	133	+ pipeline_model_parallel_size=pp,
	134	+ virtual_pipeline_model_parallel_size=None,
	135	+ pipeline_model_parallel_split_rank=None)
	136	+ # 2d rms_norm
	137	+ output_2d, weight_grad_2d = self.get_rms_norm_grad(dist_schedule=1, h=h, input_x=input_x, targets=targets)
	138	+ # 1d rms_norm
	139	+ output_1d, weight_grad_1d = self.get_rms_norm_grad(dist_schedule=0, h=h, input_x=input_x, targets=targets)
	140	+ assert torch.allclose(output_2d, output_1d, rtol=self.diff_value, atol=self.diff_value)
	141	+ weight_grad_1d = tensor_parallel.scatter_to_tensor_model_parallel_region(weight_grad_1d)
	142	+ weight_grad_2d = sync_reduce_scatter_along_first_dim(weight_grad_2d, TPXCollectiveComm)
	143	+ assert torch.allclose(weight_grad_2d, weight_grad_1d, rtol=self.diff_value, atol=self.diff_value)