已开启
[bug] PP+FSDP PTA后端 和standalone 进行精度比较时，loss偏低 #193
MengXY107创建于 6月8日
MengXY107
6月8日 评论：
6月8日 评论：
import os
import sys
os.environ["HYPER_PARALLEL_PLATFORM"] = "torch"
sys.path.insert(0, os.path.abspath(os.path.join(__file__, "../../..")))
import time
import numpy as np
import logging
import torch
import torch.nn as nn
import torch.distributed as dist
from hyper_parallel import init_device_mesh
from hyper_parallel import SkipDTensorDispatch
from hyper_parallel.core.pipeline_parallel import (PipelineStage,
ScheduleGPipe,
Schedule1F1B,
ScheduleInterleaved1F1B,
BatchDimSpec,
)
from hyper_parallel.core.pipeline_parallel.stage import SharedParameterInfo
from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat
from moe_transformer_parallel_demo import (PanGuMoETransformer,
ModelConfig,
get_single_node_8gpu_config,
TransformerBlock,
RMSNorm,
get_rank,
get_world_size)
from hyper_parallel import init_device_mesh, ContextParallel, parallelize_module, ColwiseParallel, RowwiseParallel
from hyper_parallel.core.fully_shard.api import fully_shard
from hyper_parallel.core.expert_parallel.expert_parallel import ExpertParallel
enable_mindspore_backward_compat()
def compare_nparray(self, data_expected, data_me, rtol, atol, equal_nan=True):
if np.any(np.isnan(data_expected)):
assert np.allclose(data_expected, data_me, rtol, atol, equal_nan=equal_nan)
elif not np.allclose(data_expected, data_me, rtol, atol, equal_nan=equal_nan):
assert data_expected.shape == data_me.shape
total_count = len(data_expected.flatten())
error = np.abs(data_expected - data_me)
greater = np.greater(error, atol + np.abs(data_me) * rtol)
nan_diff = np.not_equal(np.isnan(data_expected), np.isnan(data_me))
inf_diff = np.not_equal(np.isinf(data_expected), np.isinf(data_me))
# ICKTGQ
if data_expected.dtype in ('complex64', 'complex128'):
greater = greater + nan_diff + inf_diff
else:
neginf_diff = np.not_equal(np.isneginf(data_expected), np.isneginf(data_me))
greater = greater + nan_diff + inf_diff + neginf_diff
loss_count = np.count_nonzero(greater)
assert (loss_count / total_count) < rtol, \
"\ndata_expected_std:{0}\ndata_me_error:{1}\nloss:{2}". \
format(data_expected[greater], data_me[greater], error[greater])
else:
assert np.array(data_expected).shape == np.array(data_me).shape
class PanGuMoETransformerSingle(PanGuMoETransformer):
"""
Pipeline 感知版 PanGuMoETransformer
每个 stage 只持有部分层，通过 pipeline_stage 和 num_stages 参数控制
"""
def __init__(self, config: ModelConfig):
super().__init__(config=config)
def forward(self, input_ids, labels=None):
x = self.embedding(input_ids).transpose(0, 1).contiguous()
seq_len = input_ids.shape[1]
mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device))
for layer in self.layers:
x = layer(x, mask)
x = self.final_norm(x).transpose(0, 1)
x = self.lm_head(x)
x = nn.functional.cross_entropy(x.reshape(-1, x.shape[-1]), labels.reshape(-1), reduction='mean')
return x
# ==================== PanGuMoETransformerPP 模型定义 ====================
class PanGuMoETransformerPP(nn.Module):
"""
Pipeline 感知版 PanGuMoETransformer
每个 stage 只持有部分层，通过 pipeline_stage 和 num_stages 参数控制
"""
def __init__(self, config, pipeline_stage=0, num_stages=1):
super().__init__()
self.config = config
self.pipeline_stage = pipeline_stage
self.num_stages = num_stages
tp_cfg = config.parallel.tensor_parallel if config.parallel else None
ep_cfg = config.parallel.expert_parallel if config.parallel else None
tp_rank = 0
tp_world_size = 1
ep_rank = 0
ep_world_size = 1
if tp_cfg and tp_cfg.enabled:
tp_rank = get_rank() % tp_cfg.world_size
tp_world_size = tp_cfg.world_size
if ep_cfg and ep_cfg.enabled:
ep_rank = get_rank() // tp_cfg.world_size
ep_world_size = ep_cfg.world_size
self.tp_rank = tp_rank
self.tp_world_size = tp_world_size
self.ep_rank = ep_rank
self.ep_world_size = ep_world_size
# 计算当前 stage 负责的层范围
num_layers = config.num_layers
layers_per_stage = (num_layers + num_stages - 1) // num_stages
start_idx = pipeline_stage * layers_per_stage
end_idx = min(start_idx + layers_per_stage, num_layers)
# Stage 0: 创建 embedding
if pipeline_stage == 0:
self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
# 当前 stage 对应的 layers
self.layers = nn.ModuleList([
TransformerBlock(config, i, tp_rank, tp_world_size, ep_rank, ep_world_size)
for i in range(start_idx, end_idx)
])
# 最后一个 stage: 创建 final_norm 和 lm_head
if pipeline_stage == num_stages - 1:
self.final_norm = RMSNorm(config.hidden_size)
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
self.transpose = ops_Transpose()
self.seq_len = config.seq_len
self.register_buffer("mask", torch.tril(torch.ones((self.seq_len, self.seq_len))))
def forward(self, input_ids, labels=None):
# Stage 0: 接收 input_ids
if self.pipeline_stage == 0:
x = self.embedding(input_ids)
x = self.transpose(x, (1, 0, 2)).contiguous()
else:
# 其他 stage: 接收上一个 stage 的 hidden states
x = input_ids
# 执行当前 stage 的 layers
for layer in self.layers:
x = layer(x, self.mask)
# 最后一个 stage: 执行 final_norm + lm_head
if self.pipeline_stage == self.num_stages - 1:
x = self.final_norm(x)
x = self.transpose(x, (1, 0, 2))
x = self.lm_head(x)
if labels is not None:
x = nn.functional.cross_entropy(x.reshape(-1, x.shape[-1]), labels.reshape(-1), reduction='mean')
return x
return x
def ops_Transpose():
"""包装 torch.transpose 操作"""
return TransposeOp()
class TransposeOp:
"""类似于 MindSpore ops.Transpose 的转置操作"""
def __call__(self, x, perms):
return x.permute(perms).contiguous()
# ==================== 配置和工具函数 ====================
def get_model_config(num_layers=4, vocab_size=1000, hidden_size=256, seq_len=32, num_experts=16, top_k=4,
moe_enabled=False):
"""获取测试模型配置"""
parallel_config = get_single_node_8gpu_config()
parallel_config.tensor_parallel.enabled = False
parallel_config.expert_parallel.enabled = False
parallel_config.sequence_parallel.enabled = False
parallel_config.micro_batch_size = 8
model_config = ModelConfig()
model_config.num_layers = num_layers
model_config.vocab_size = vocab_size
model_config.hidden_size = hidden_size
model_config.seq_len = seq_len
model_config.moe.num_experts = num_experts
model_config.moe.top_k = top_k
if not moe_enabled:
model_config.moe.layers_replaced_as_dense = list(range(num_layers))
model_config.parallel = parallel_config
logging.warning(f"ModelConfig====> 模型配置:")
logging.warning(f"ModelConfig====> 层数->num_layers：{model_config.num_layers}")
logging.warning(f"ModelConfig====> 隐藏层->hidden_size：{model_config.hidden_size}")
if moe_enabled:
logging.warning(f"ModelConfig====> 专家数->num_experts：{model_config.moe.num_experts}")
logging.warning(f"ModelConfig====> Top-K->top_k: {model_config.moe.top_k}")
return model_config
def setup_distributed_logging(rank, world_size):
# 主进程打印到 stdout
if rank == 0:
handler = logging.StreamHandler()
else:
handler = logging.FileHandler(f"rank{rank}.log")
formatter = logging.Formatter(
f"[R{rank}/{world_size}] %(asctime)s %(levelname)s %(message)s"
)
handler.setFormatter(formatter)
root = logging.getLogger()
root.setLevel(logging.WARNING)
root.addHandler(handler)
def setup_function():
"""测试前准备"""
torch.manual_seed(42)
np.random.seed(42)
os.environ['PYTHONHASHSEED'] = '42'
torch_npu.npu.manual_seed(42)
torch_npu.npu.manual_seed_all(42) # 多 NPU 卡
# CANN 确定性模式（核心）
os.environ['HYPER_THREAD'] = '0'
os.environ['ENABLE_MPSANDBOX'] = '0'
# 启用确定性算子
torch_npu.npu.set_compile_mode(jit_compile=False)
torch.use_deterministic_algorithms(True)
"""初始化分布式环境"""
if not dist.is_initialized():
dist.init_process_group(backend='hccl')
rank = dist.get_rank()
world = dist.get_world_size()
torch.npu.set_device(rank)
setup_distributed_logging(rank, world)
def teardown_function():
"""测试后清理"""
if dist.is_initialized():
dist.destroy_process_group()
class GeneratorDataset:
"""简单的 PyTorch 数据集模拟"""
def __init__(self, data, column_names, shuffle=False):
self.data = data
self.column_names = column_names
self.shuffle = shuffle
def __iter__(self):
indices = list(range(len(self.data)))
if self.shuffle:
np.random.shuffle(indices)
for idx in indices:
input_ids = self.data[idx][0]
labels = self.data[idx][1]
yield input_ids, labels
def generator_dataset(size, batch_size=1000, seq_len=32, dynamic=False, dtype=np.int32):
"""生成训练数据"""
np.random.seed(0)
for _ in range(size):
if dynamic:
batch_size = 8 * np.random.randint(1, 32)
input_ids = np.random.randint(0, batch_size, size=(batch_size, seq_len)).astype(dtype)
input_mask = np.random.randint(0, batch_size, size=(batch_size, seq_len)).astype(dtype)
yield input_ids, input_mask
def run_train_with_moe(model, train_dataset, epochs=10, max_steps=10, schedule=None, pp_rank=0, pp_size=4,
stages_param=None, clip_norm=False, amp=False, micro_batch=8):
"""PyTorch 版本的训练循环"""
# 获取可训练参数
if stages_param is not None:
trainable_params = stages_param
else:
trainable_params = list(filter(lambda p: p.requires_grad, model.parameters()))
optimizer = torch.optim.Adam(trainable_params, lr=1e-4)
loss_list = []
step = 0
for _ in range(epochs):
for _, (input_ids, labels) in enumerate(train_dataset):
input_ids = torch.tensor(input_ids, dtype=torch.long).npu()
labels = torch.tensor(labels, dtype=torch.long).npu()
# 前向传播
optimizer.zero_grad()
with SkipDTensorDispatch():
if schedule is not None:
# Pipeline parallel 训练
if amp:
with torch.cuda.amp.autocast():
loss = schedule.run(input_ids) if pp_rank == 0 else schedule.run(labels=labels)
else:
loss = schedule.run(input_ids) if pp_rank == 0 else schedule.run(labels=labels)
if pp_rank == pp_size - 1:
# 获取 loss
if isinstance(loss, (list, tuple)):
loss_values = [ls.item() if hasattr(ls, 'item') else ls for ls in loss]
loss = np.mean(loss_values)
print(f"hyper step: {step} loss: {loss}", flush=True)
# 从 stage 参数获取梯度
if clip_norm:
torch.nn.utils.clip_grad_norm_(trainable_params, max_norm=1.0)
# ClipGradNorm(max_norm=1.0)(tuple(p.grad for p in stages_param))
with torch.no_grad():
for p in stages_param:
if p.grad is not None:
p.grad.div_(micro_batch)
else:
# 单卡训练
if amp:
with torch.cuda.amp.autocast():
loss = model(input_ids, labels)
else:
loss = model(input_ids, labels)
loss.backward()
if get_rank() == 0:
print(f"standalone step: {step} loss: {loss}", flush=True)
if clip_norm:
torch.nn.utils.clip_grad_norm_(trainable_params, max_norm=1.0)
optimizer.step()
if isinstance(loss, (list, tuple)):
loss = float(np.mean([l.item() if hasattr(l, 'item') else l for l in loss]))
loss_list.append(loss if isinstance(loss, float) else loss.item())
step += 1
if step >= max_steps:
break
if step >= max_steps:
break
return loss_list, []
def run_train_with_moe_multi_model(models, train_dataset, epochs=10, max_steps=10, schedule=None, pp_rank=0, pp_size=4, micro_batch=8):
"""多模型版本的训练循环 (用于 ScheduleInterleaved1F1B)"""
model_params = [tuple(p for p in model.parameters() if p.requires_grad) for model in models]
optimizers = [torch.optim.Adam(_params, lr=1e-4) for _params in model_params]
loss_list = []
step = 0
time_list = []
for _ in range(epochs):
for _, (input_ids, labels) in enumerate(train_dataset):
input_ids = torch.tensor(input_ids, dtype=torch.long).npu()
labels = torch.tensor(labels, dtype=torch.long).npu()
start = time.time()
# 清零梯度
for m_optimizer in optimizers:
m_optimizer.zero_grad()
loss = schedule.run(input_ids) if pp_rank == 0 else schedule.run(labels=labels)
if pp_rank == pp_size - 1:
# 获取 loss
if isinstance(loss, (list, tuple)):
loss_values = [ls.item() if hasattr(ls, 'item') else ls for ls in loss]
loss = np.mean(loss_values)
# 更新每个模型的参数
for m_optimizer, m_params in zip(optimizers, model_params):
for p in m_params:
if p.grad is not None:
p.grad.div_(micro_batch)
m_optimizer.step()
end = time.time()
if step != 0:
time_list.append(end - start)
loss_list.append(loss)
step += 1
if step >= max_steps:
break
if step >= max_steps:
break
return loss_list, time_list
# ==================== 测试用例 ====================
def check_loss_converge(loss_list, slope_threshold=-0.01, compare_ratio=0.95):
"""
判断loss是否呈收敛趋势 方法1: 线性回归斜率 < threshold 方法2: 后50%平均loss / 前50%平均loss < ratio 两种方法同时满足才判定为收敛 """
loss = np.array(loss_list)
if len(loss) < 4:
return False, "loss数据点太少"
n = len(loss)
x = np.arange(n)
# 方法1: 线性回归斜率
coeffs = np.polyfit(x, loss, 1)
slope = coeffs[0]
if slope >= slope_threshold:
return False, f"线性斜率 {slope:.6f} >= {slope_threshold}，不收敛"
# 方法2: 前后半段比较
mid = n // 2
first_half_avg = np.mean(loss[:mid])
second_half_avg = np.mean(loss[mid:])
ratio_val = second_half_avg / first_half_avg if first_half_avg > 0 else 1.0
if ratio_val >= compare_ratio:
return False, f"前后半段比值 {ratio_val:.4f} >= {compare_ratio}，不收敛"
logging.warning(f"MindSporeTest====> 收敛: 斜率={slope:.6f}, 前后比值={ratio_val:.4f}")
return True
# ==================== 测试用例 ====================
"""
TEST_SUMMARY: 并行特性交互：FSDP+pp
Test Steps:
1. 构建测试网络
2. 初始化mesh
mesh = init_device_mesh(device_type="npu", mesh_shape=(4, 2), mesh_dim_names=("pp", "dp"))
pp_mesh = mesh["pp"]
3.设置 fsdp+pp 并行：
ScheduleGPipe + batch_size=256, micro_batch=8, stage_num=4
4.执行 10 step 训练
Expected Result:
1.验证不 hang、loss 正确
"""
def test_pipeline_parallel_feature_interaction_torch_check_004():
# 获取 device 信息
rank = get_rank()
world_size = get_world_size()
rank = get_rank()
# 训练参数
epochs = 100
max_steps = 500
batch_size = 256
seq_len = 32
micro_batch_num = 8
pp_size = 4
dp_size = world_size // pp_size
# 生成训练数据
all_data = list(generator_dataset(size=1000, batch_size=batch_size, seq_len=seq_len))
train_dataset_1 = GeneratorDataset(all_data[:], ["data", "label"], shuffle=False)
train_dataset_2 = GeneratorDataset(all_data[:], ["data", "label"], shuffle=False)
# 创建配置
model_config = get_model_config()
# 创建单卡训练模型
single_model = PanGuMoETransformerSingle(model_config)
for param in single_model.parameters():
param.data = param.data.npu()
single_model = single_model.to('npu')
# 创建pp训练模型
mesh = init_device_mesh(device_type="npu", mesh_shape=(pp_size, dp_size), mesh_dim_names=("pp", "dp"))
dp_mesh = mesh["dp"]
pp_mesh = mesh["pp"]
pp_rank = pp_mesh.get_local_rank()
chunk_model = PanGuMoETransformerPP(model_config, pipeline_stage=pp_rank, num_stages=pp_size)
# 权重同步
# 先同步单卡模型到NPU，确保state_dict的权重在NPU上
single_state = single_model.state_dict()
chunk_state = chunk_model.state_dict()
for key in chunk_state:
s_key = key.replace("layers.0.", f"layers.{rank // dp_size}.")
if s_key in single_state:
chunk_state[key] = single_state[s_key]
chunk_model.load_state_dict(chunk_state)
# 同步chunk_model参数和buffer到NPU
for param in chunk_model.parameters():
param.data = param.data.npu()
chunk_model = chunk_model.to('npu')
for name, buf in chunk_model.named_buffers():
chunk_model.register_buffer(name, single_state[name] if name in single_state else buf.to("npu"))
# FSDP: 对模型应用数据并行分片
fully_shard(chunk_model.layers[0], mesh=mesh["dp"])
fully_shard(chunk_model, mesh=mesh["dp"])
chunk_param = [param for param in chunk_model.parameters() if param.requires_grad]
# 执行端到端训练
logging.warning("MindSporeTest====> 开始单卡训练")
run_train_with_moe(single_model, train_dataset_1, epochs=epochs, max_steps=max_steps)
stages = [PipelineStage(chunk_model, stage_index=rank // dp_size, stage_num=pp_size, group=pp_mesh.get_group(),
mesh=pp_mesh, device='npu')]
schedule = ScheduleGPipe(stages, micro_batch_num=micro_batch_num)
logging.warning("MindSporeTest====> 开始 pp 训练")
run_train_with_moe(chunk_model, train_dataset_2, epochs=epochs, max_steps=max_steps,
schedule=schedule, pp_rank=pp_rank, stages_param=chunk_param, micro_batch=micro_batch_num)
# 对比同名参数的梯度
for name, p in single_model.named_parameters():
if p.grad is not None:
for name2, p2 in chunk_model.named_parameters():
s_name = name2.replace("layers.0.", f"layers.{rank // dp_size}.")
if name == s_name and p2.grad is not None:
# single_model 没有分片，直接取全量梯度
local_grad = p.grad.cpu().numpy()
# FSDP 梯度是 DTensor，all_gather 拼成全量梯度
g_local = p2.grad.to_local() if hasattr(p2.grad, "to_local") else p2.grad
full_grad_list = [torch.zeros_like(g_local) for _ in range(dp_size)]
torch.distributed.all_gather(full_grad_list, g_local, group=dp_mesh.get_group())
full_grad = torch.cat(full_grad_list, dim=0).cpu().numpy()
# 比较单卡全量梯度和 FSDP 收集后的全量梯度
compare_nparray(full_grad, local_grad, 0.002, 0.002)
if __name__ == "__main__":
import torch_npu
torch_npu.npu.is_available = lambda: True
logging.warning("MindSporeTest::set torch_npu.npu.is_available=True")
setup_function()
test_pipeline_parallel_feature_interaction_torch_check_004()
teardown_function()
MengXY107
6月8日 评论：
6月8日 评论：
Standalone 方式构建梯度累加：
per = batch // mb # micro-batch size, matches the PP/GPipe micro split
is_owner = (pp_rank == pp_size - 1) and (dp_mesh.get_local_rank() == 0)
f_std = f_pp = None
if is_owner:
f_std = open("standalone.txt", "w")
f_pp = open("pp_fsdp.txt", "w")
single_params = [p for p in single.parameters() if p.requires_grad]
for step in range(STEPS):
x = all_data[step][0]
y = all_data[step][1]
xt = torch.tensor(x, dtype=torch.long).npu()
yt = torch.tensor(y, dtype=torch.long).npu()
# ---- standalone in the SAME micro-batch accumulation口径 as PP+GPipe ----
# 8 micro-batches of `per`, mean-CE each, accumulate grads, then div_(mb)
# (mirrors the PP path's p.grad.div_(micro_batch)); loss = mean of micro means.
sopt.zero_grad()
s_micro = []
for i in range(mb):
xs = xt[i * per:(i + 1) * per]
ys = yt[i * per:(i + 1) * per]
l = single(xs, ys)
l.backward()
s_micro.append(l.item())
with torch.no_grad():
for p in single_params:
if p.grad is not None:
p.grad.div_(mb)
sval = float(np.mean(s_micro))
MengXY107
6月10日 评论：
6月10日 评论：
确认是问题，master 当前实现，torch侧的fully_shard的reduce_params只处理TP引入的mirrors AllReduce。因此导致sharded_param.grad并没有被挂载上。导致了精度出现问题。
MengXY107
6月12日 评论：
6月12日 评论：
moe_transformer_demo_mindspore.py
"""
简化版 MoE Transformer 示例 - MindSpore 版本
核心特性:
- MoE (Mixture of Experts): 192 专家选 10
- MLA (Multi-Head Latent Attention): 低秩注意力
- SwiGLU 激活
- RMSNorm 归一化
"""
import math
import numpy as np
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor, dtype as mstype
from typing import List, Tuple, Optional
# ==================== 配置类 ====================
class MoEConfig:
"""MoE 配置 (对应 yaml 的 moe 部分)"""
num_experts: int = 192 # 专家总数
top_k: int = 10 # 激活专家数
expert_hidden_size: int = 512 # 专家隐藏层维度
shared_hidden_size: int = 1024 # 共享专家维度
load_balancing: str = 'noaux_tc'
aux_loss_coeff: float = 0.0002
sigmoid_gating: bool = True
routed_scaling: float = 2.5
layers_replaced_as_dense: List[int] = [0] # 保持稠密的层 (默认第 0 层)
class ModelConfig:
"""模型配置 (从 yaml 提取关键参数)"""
# 基础
vocab_size: int = 153600
hidden_size: int = 2560
num_layers: int = 45
seq_len: int = 4096
# Attention
num_heads: int = 32
q_lora_rank: int = 768
kv_lora_rank: int = 512
qk_rope_head_dim: int = 64
# FFN
ffn_hidden: int = 6144
activation: str = 'swiglu'
# MoE
moe: MoEConfig = MoEConfig()
# ==================== 核心组件 ====================
class RMSNorm(nn.Cell):
"""RMS 归一化 (替代 LayerNorm)"""
def __init__(self, hidden_size: int, eps: float = 1e-5):
super().__init__()
self.weight = ms.Parameter(Tensor(np.ones(hidden_size), mstype.float32))
self.eps = eps
def construct(self, x: Tensor) -> Tensor:
return ops.rms_norm(x, self.weight, self.eps)[0]
class MoEGating(nn.Cell):
"""MoE 路由器 (Router) - 决定 token 分配给哪些专家"""
def __init__(self, config: MoEConfig, hidden_size: int):
super().__init__()
self.config = config
self.gate = nn.Dense(hidden_size, config.num_experts, has_bias=False)
if config.sigmoid_gating:
self.sigmoid = ops.Sigmoid()
self.top_k = config.top_k
self.routed_scaling = config.routed_scaling
def construct(self, x: Tensor) -> Tuple[Tensor, Tensor]:
"""
Args:
x: [seq_len * batch, hidden_size]
Returns:
topk_indices: [seq_len * batch, top_k] 专家索引
topk_weights: [seq_len * batch, top_k] 路由权重
"""
# 计算路由分数
gate_logits = self.gate(x) # [N, num_experts]
if self.config.sigmoid_gating:
gate_logits = self.sigmoid(gate_logits)
# Top-K 选择
topk_weights, topk_indices = ops.TopK()(gate_logits, self.top_k)
# 归一化权重 (softmax over top-k)
topk_weights = ops.Softmax(axis=-1)(topk_weights) * self.routed_scaling
return topk_indices, topk_weights
class Expert(nn.Cell):
"""单个专家网络 (简单的 MLP)"""
def __init__(self, config: MoEConfig, hidden_size: int):
super().__init__()
self.gate_proj = nn.Dense(hidden_size, config.expert_hidden_size)
self.up_proj = nn.Dense(hidden_size, config.expert_hidden_size)
self.down_proj = nn.Dense(config.expert_hidden_size, hidden_size)
self.silu = ops.Sigmoid()
def construct(self, x: Tensor) -> Tensor:
return self.down_proj(self.silu(self.gate_proj(x)) * self.up_proj(x))
class SharedExpert(nn.Cell):
"""共享专家 (所有 token 都经过)"""
def __init__(self, config: MoEConfig, hidden_size: int):
super().__init__()
self.gate_proj = nn.Dense(hidden_size, config.shared_hidden_size)
self.up_proj = nn.Dense(hidden_size, config.shared_hidden_size)
self.down_proj = nn.Dense(config.shared_hidden_size, hidden_size)
self.silu = ops.Sigmoid()
def construct(self, x: Tensor) -> Tensor:
return self.down_proj(self.silu(self.gate_proj(x)) * self.up_proj(x))
class DenseFFN(nn.Cell):
"""稠密 FFN 层 (用于第 0 层)"""
def __init__(self, config: ModelConfig):
super().__init__()
self.gate_proj = nn.Dense(config.hidden_size, config.ffn_hidden)
self.up_proj = nn.Dense(config.hidden_size, config.ffn_hidden)
self.down_proj = nn.Dense(config.ffn_hidden, config.hidden_size)
self.silu = ops.Sigmoid()
def construct(self, x: Tensor) -> Tensor:
return self.down_proj(self.silu(self.gate_proj(x)) * self.up_proj(x))
class MoELayer(nn.Cell):
"""
MoE 层 - 核心组件
工作流程:
1. Router 为每个 token 选择 top-k 专家
2. 每个专家独立处理分配给它的 token
3. 加权组合专家输出 + 共享专家输出
"""
def __init__(self, config: MoEConfig, hidden_size: int, ffn_hidden: int = None):
super().__init__()
self.config = config
self.hidden_size = hidden_size
self.ffn_hidden = ffn_hidden or config.expert_hidden_size
# 路由器
self.gate = MoEGating(config, hidden_size)
# 192 个专家 (用 CellList 管理)
self.experts = nn.CellList([
Expert(config, hidden_size) for _ in range(config.num_experts)
])
# 共享专家 (始终激活)
self.shared_expert = SharedExpert(config, hidden_size)
def construct(self, x: Tensor) -> Tensor:
"""
Args:
x: [seq_len, batch, hidden_size]
Returns:
output: [seq_len, batch, hidden_size]
"""
seq_len, batch, hidden = x.shape
# 展平为 [seq*batch, hidden] 方便路由
x_flat = x.reshape(-1, hidden)
# 1. 路由决策
topk_indices, topk_weights = self.gate(x_flat)
# 2. 初始化输出
output = ops.zeros_like(x_flat)
# 3. 对每个专家，处理分配给它的 token
for expert_idx, expert in enumerate(self.experts):
# 找到选择此专家的 token
token_mask = (topk_indices == expert_idx) # [N, top_k]
if ops.any(token_mask):
# 获取对应的 token 和权重
selected_mask = ops.any(token_mask, axis=1)
selected_tokens = x_flat[selected_mask]
# 专家处理
expert_out = expert(selected_tokens)
# 加权累加 - 简化处理
output[selected_mask] += expert_out * ops.mean(topk_weights[token_mask], -1, True)
# 4. 加上共享专家输出 (所有 token 都经过)
shared_out = self.shared_expert(x_flat)
output = output + shared_out
# 恢复形状
return output.reshape(seq_len, batch, hidden)
class MLAAttention(nn.Cell):
"""
简化版 MLA (Multi-Head Latent Attention)
核心思想：用低秩压缩减少 KV cache 大小
"""
def __init__(self, config: ModelConfig):
super().__init__()
self.hidden_size = config.hidden_size
self.num_heads = config.num_heads
self.head_dim = config.qk_rope_head_dim
# LoRA 低秩投影
self.q_lora_down = nn.Dense(config.hidden_size, config.q_lora_rank)
self.q_lora_up = nn.Dense(config.q_lora_rank, config.num_heads * self.head_dim)
self.kv_lora_down = nn.Dense(config.hidden_size, config.kv_lora_rank)
self.kv_lora_up = nn.Dense(config.kv_lora_rank,
config.num_heads * (self.head_dim + self.head_dim))
# 输出投影
self.out_proj = nn.Dense(config.num_heads * self.head_dim, config.hidden_size)
self.softmax = ops.Softmax(axis=-1)
self.transpose = ops.Transpose()
self.split = ops.Split(axis=-1, output_num=2)
def construct(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
"""
Args:
x: [seq_len, batch, hidden_size]
"""
seq_len, batch, _ = x.shape
# 低秩投影
q_latent = self.q_lora_down(x)
q = self.q_lora_up(q_latent)
kv_latent = self.kv_lora_down(x)
kv = self.kv_lora_up(kv_latent)
# 拆分 K 和 V
k, v = self.split(kv)
# 重塑为多头
q = self.transpose(q.view(seq_len, batch, self.num_heads, self.head_dim), (1, 2, 0, 3))
k = self.transpose(k.view(seq_len, batch, self.num_heads, self.head_dim), (1, 2, 0, 3))
v = self.transpose(v.view(seq_len, batch, self.num_heads, self.head_dim), (1, 2, 0, 3))
# 注意力计算
scores = ops.matmul(q, self.transpose(k, (0, 1, 3, 2))) / math.sqrt(self.head_dim)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn = self.softmax(scores)
out = ops.matmul(attn, v)
# 合并头
out = self.transpose(out, (0, 2, 1, 3)).reshape(seq_len, batch, -1)
return self.out_proj(out)
class TransformerBlock(nn.Cell):
"""Transformer 块 (包含 MLA + MoE)"""
def __init__(self, config: ModelConfig, layer_idx: int):
super().__init__()
self.layer_idx = layer_idx
# 归一化
self.norm1 = RMSNorm(config.hidden_size)
self.norm2 = RMSNorm(config.hidden_size)
# 注意力
self.attention = MLAAttention(config)
# MoE 层 (第 0 层是稠密层)
if layer_idx in config.moe.layers_replaced_as_dense:
self.moe = DenseFFN(config)
else:
self.moe = MoELayer(config.moe, config.hidden_size, config.ffn_hidden)
def construct(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
# Pre-Norm 架构
h = x + self.attention(self.norm1(x), mask)
out = h + self.moe(self.norm2(h))
return out
class PanGuMoETransformer(nn.Cell):
"""
MoE Transformer 简化版
架构:
Embedding -> [TransformerBlock × 45] -> RMSNorm -> LM Head
"""
def __init__(self, config: ModelConfig):
super().__init__()
self.config = config
# 词嵌入
self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
# Transformer 层
self.layers = nn.CellList([
TransformerBlock(config, i) for i in range(config.num_layers)
])
# 最终归一化
self.final_norm = RMSNorm(config.hidden_size)
# LM Head
self.lm_head = nn.Dense(config.hidden_size, config.vocab_size, has_bias=False)
self.transpose = ops.Transpose()
def construct(self, input_ids: Tensor) -> Tensor:
"""
Args:
input_ids: [batch, seq_len]
Returns:
logits: [batch, seq_len, vocab_size]
"""
# 嵌入
x = self.embedding(input_ids)
x = self.transpose(x, (1, 0, 2)) # [seq, batch, hidden]
# 生成因果 mask
seq_len = input_ids.shape[1]
mask = ops.tril(ops.ones((seq_len, seq_len), mstype.float32))
# 通过所有层
for layer in self.layers:
x = layer(x, mask)
# 最终归一化
x = self.final_norm(x)
# LM Head
x = self.transpose(x, (1, 0, 2))
logits = self.lm_head(x)
return logits
# ==================== 使用示例 ====================
if __name__ == '__main__':
print("=" * 60)
print(" MoE Transformer 简化示例 - MindSpore 版本")
print("=" * 60)
# 设置 MindSpore 环境
ms.set_context(mode=ms.PYNATIVE_MODE, device_target="Ascend")
# 1. 创建配置
config = ModelConfig()
print(f"\n模型配置:")
print(f" - 隐藏层维度：{config.hidden_size}")
print(f" - 层数：{config.num_layers}")
print(f" - 专家数：{config.moe.num_experts}")
print(f" - 激活专家：{config.moe.top_k}")
print(f" - 词表大小：{config.vocab_size}")
# 2. 创建模型 (用小配置演示)
print("\n创建模型...")
demo_config = ModelConfig()
demo_config.num_layers = 2
demo_config.vocab_size = 1000
demo_config.hidden_size = 256
demo_config.moe.num_experts = 8
demo_config.moe.top_k = 2
model = PanGuMoETransformer(demo_config)
# 打印参数量
total_params = sum(p.size for p in model.get_parameters())
print(f"模型参数量：{total_params / 1e6:.2f}M")
# 3. 前向传播示例
print("\n前向传播测试:")
batch_size = 2
seq_len = 32
input_ids = Tensor(np.random.randint(0, demo_config.vocab_size, (batch_size, seq_len)), mstype.int32)
logits = model(input_ids)
print(f" 输入形状：{input_ids.shape}")
print(f" 输出形状：{logits.shape}")
print(f" [OK] 前向传播成功!")
print("\n" + "=" * 60)
print("示例完成！")
print("=" * 60)
MengXY107
6月12日 评论：
6月12日 评论：
"""
Test suite for fully_shard ZeRO-2/ZeRO-3 on MindSpore.
包含:
- 共享基础函数 (模型配置、数据集、训练函数、权重同步、结果对比)
- 12 个主测试用例 (test_fully_shard_ms_check_001 ~ 012)
注意:
- 使用 network/mindspore/pangu/moe_transformer_parallel_demo_mindspore.py 中的 PanGuMoETransformer
- 大部分用例在测试函数中应用 fully_shard
- 只有特殊用例 (如 check_003 单层测试) 在模型 __init__ 中对 layer 调用 fully_shard
"""
import os
import sys
import time
import pytest
import numpy as np
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops.operations as P
from mindspore import mint
from mindspore import ops
from mindspore import log as logger
from mindspore.common import _no_grad
from mindspore.communication import get_group_size, get_rank, init
from mindspore.dataset import GeneratorDataset
from mindspore import value_and_grad
from mindspore._c_expression import NoFallbackGuard
from ...share.context_base import contextbase
from hyper_parallel import init_device_mesh, fully_shard, SkipDTensorDispatch
from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat
from ..network.mindspore.pangu.moe_transformer_parallel_demo_mindspore import (
PanGuMoETransformer, ModelConfig, get_single_node_8gpu_config,
TransformerBlock, RMSNorm, get_rank_fun, MoEConfig
)
from ..network.mindspore.clip_grad import ClipGradNorm
from ...share.marker import (
Author, Level0, Level1, Manual, Env_Cards,
SKIP_ENV_CPU, SKIP_ENV_910A, SKIP_MODE_GRAPH
)
from hyper_parallel.core.fully_shard.utils import CPUOffloadPolicy,MixedPrecisionPolicy
from hyper_parallel.core.dtensor.placement_types import Shard
enable_mindspore_backward_compat()
ms.runtime.launch_blocking()
def setup_function():
"""每个测试函数执行前的 setup"""
contextbase.case_prepare()
ms.set_deterministic(True)
ms.set_seed(0)
np.random.seed(0)
def teardown_function():
"""每个测试函数执行后的 teardown"""
contextbase.case_cleanup()
try:
from mindspore.communication import destroy_group
destroy_group()
except:
pass
def get_model_config(num_layers=4, vocab_size=1000, hidden_size=256, seq_len=32,
num_experts=16, top_k=4, moe_enabled=False, tp=False, ep=False, sp=False):
"""
创建模型配置
Args:
num_layers: Transformer 层数
vocab_size: 词表大小
hidden_size: 隐藏层维度
seq_len: 序列长度
num_experts: MoE 专家数量
top_k: MoE top-k 选择数
moe_enabled: 是否启用 MoE
Returns:
ModelConfig 对象
"""
parallel_config = get_single_node_8gpu_config()
parallel_config.tensor_parallel.enabled = tp
parallel_config.expert_parallel.enabled = ep
parallel_config.sequence_parallel.enabled = sp
parallel_config.micro_batch_size = 8
model_config = ModelConfig()
model_config.num_layers = num_layers
model_config.vocab_size = vocab_size
model_config.hidden_size = hidden_size
model_config.seq_len = seq_len
model_config.moe.num_experts = num_experts
model_config.moe.top_k = top_k
if not moe_enabled:
model_config.moe.layers_replaced_as_dense = list(range(num_layers))
model_config.parallel = parallel_config
logger.warning(f"MindSporeTest====>模型配置:")
logger.warning(f"MindSporeTest====> 层数->num_layers: {model_config.num_layers}")
logger.warning(f"MindSporeTest====> 隐藏层->hidden_size: {model_config.hidden_size}")
if moe_enabled:
logger.warning(f"MindSporeTest====> 专家数->num_experts: {model_config.moe.num_experts}")
logger.warning(f"MindSporeTest====> Top-K->top_k: {model_config.moe.top_k}")
return model_config
def generator_dataset(size, batch_size=2, seq_len=32, vocab_size=1000, dynamic=False, dtype=np.int32):
"""
生成训练数据 (参考 moe_transformer_parallel_demo_mindspore.py main 方法)
Args:
size: 数据样本数
batch_size: 批次大小
seq_len: 序列长度
vocab_size: 词表大小 (用于生成 input_ids 和 labels)
dynamic: 是否动态 batch_size
dtype: 数据类型
Yields:
input_ids, labels
"""
np.random.seed(0)
for _ in range(size):
if dynamic:
batch_size = 8 * np.random.randint(1, 32)
input_ids = np.random.randint(0, vocab_size, size=(batch_size, seq_len)).astype(dtype)
labels = np.random.randint(0, vocab_size, size=(batch_size, seq_len)).astype(dtype)
yield input_ids, labels
class PanGuMoETransformerSingle(PanGuMoETransformer):
"""
单卡参考模型 - 用于对比验证
重写 construct 以支持 labels 参数计算 loss
"""
def __init__(self, config: ModelConfig):
super().__init__(config=config)
ms.set_seed(0)
np.random.seed(0)
def construct(self, input_ids, labels=None):
x = self.embedding(input_ids)
x = self.transpose(x, (1, 0, 2))
seq_len = input_ids.shape[1]
mask = ops.tril(ops.ones((seq_len, seq_len), ms.float32))
for layer in self.layers:
x = layer(x, mask)
x = self.final_norm(x)
x = self.transpose(x, (1, 0, 2))
logits = self.lm_head(x)
if labels is not None:
loss = mint.nn.functional.cross_entropy(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1), reduction='mean')
return loss
return logits
class PanGuMoETransformerFSInit(PanGuMoETransformer):
"""
fully_shard 版本模型 - 在 __init__ 中对每个 layer 调用 fully_shard
特殊用例使用 (如 check_003 单层测试、check_004 多层嵌套测试)
重写 construct 以支持 labels 参数计算 loss
"""
def __init__(self, config: ModelConfig, mesh=None, reshard_after_forward=True,
apply_to_layers=False, **fully_shard_kwargs):
super().__init__(config)
ms.set_seed(0)
np.random.seed(0)
self.fs_mesh = mesh
self.fs_reshard = reshard_after_forward
# 特殊用例：在 __init__ 中对每个 layer 应用 fully_shard
if apply_to_layers and mesh is not None:
for idx, layer in enumerate(self.layers):
logger.warning(f'\nlayer: {layer}')
fully_shard(layer, mesh=mesh, reshard_after_forward=reshard_after_forward, **fully_shard_kwargs)
layer.set_requires_gradient_sync(True)
def construct(self, input_ids, labels=None):
x = self.embedding(input_ids)
x = self.transpose(x, (1, 0, 2))
seq_len = input_ids.shape[1]
mask = ops.tril(ops.ones((seq_len, seq_len), ms.float32))
for layer in self.layers:
x = layer(x, mask)
x = self.final_norm(x)
x = self.transpose(x, (1, 0, 2))
logits = self.lm_head(x)
if labels is not None:
loss = mint.nn.functional.cross_entropy(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1), reduction='mean')
return loss
return logits
# ============================================================================
# 共享训练函数
# ============================================================================
def run_train_single(model, train_dataset, epochs=1, max_steps=10,
rank=0, monitor_memory=False):
"""
执行单卡训练 (用于对比)
Args:
model: 待训练模型
train_dataset: 训练数据集
epochs: 训练轮数
max_steps: 最大步数
rank: 进程 rank (用于日志)
monitor_memory: 是否监控显存
Returns:
loss_list, grad_dict, memory_list (if monitor_memory=True)
"""
# 创建优化器
trainable_params = list(filter(lambda p: p.requires_grad, model.get_parameters()))
optimizer = nn.Adam(trainable_params, learning_rate=1e-4)
#optimizer = nn.Adam(model.trainable_params(), learning_rate=1e-4)
# 定义训练 step - 单卡训练使用 value_and_grad
def train_step(inputs, labels):
def _zero_grad(params) -> None:
for param in params:
param.grad = None
_zero_grad([param for param in model.trainable_params()])
loss = model(inputs, labels)
with SkipDTensorDispatch():
loss.backward()
grads = tuple(p.grad for p in model.trainable_params())
with _no_grad():
optimizer(grads)
return loss
loss_list = []
grad_dict = {}
memory_list = []
step = 0
if rank == 0 and monitor_memory:
logger.warning(f'\nStarting single-card training for {max_steps} steps...')
for _ in range(epochs):
for _, (input_ids, labels) in enumerate(train_dataset):
# Record memory usage before step
if monitor_memory:
mem_before = get_memory_usage()
if rank == 0:
print_memory_usage(f'Step {step+1} - Before', mem_before)
loss = train_step(input_ids, labels)
# Record memory usage after step
if monitor_memory:
mem_after = get_memory_usage()
if rank == 0:
print_memory_usage(f'Step {step+1} - After', mem_after)
memory_list.append((mem_before, mem_after))
loss_list.append(loss.item())
# 打印 loss (仅 rank 0)
if rank == 0:
logger.warning(f' Step {step+1}: Loss = {loss.item():.6f}')
# 保存梯度用于对比
if step == 0:
for i, param in enumerate(optimizer.parameters):
grad_dict[f'layer_{i}'] = param.grad.asnumpy() if param.grad is not None else None
step += 1
if step >= max_steps:
break
if step >= max_steps:
break
del optimizer
if monitor_memory:
return loss_list, grad_dict, memory_list
return loss_list, grad_dict
def policy_fn(ctx, op, *args, **kwargs):
return CheckpointPolicy.MUST_RECOMPUTE
def apply_fully_shard(model, mesh=None, reshard_after_forward=True,
apply_root_fs=True, apply_to_layers=True,
comm_fusion_zero_copy=False, mp_policy=None,
replicate_params=None, ignored_params=None,
offload_policy=None, outer_comm_fusion=False,
recompute_policy=None):
"""
对模型应用 fully_shard 配置
Args:
model: 待应用 fully_shard 的模型
mesh: 设备网格 (None 表示 ZeRO-1 模式)
reshard_after_forward: 是否 ZeRO-3 模式 (False 表示 ZeRO-1/ZeRO-2)
apply_root_fs: 是否对整网应用 fully_shard
apply_to_layers: 是否对每个 layer 单独应用 (特殊用例)
comm_fusion_zero_copy: 是否启用 comm_fusion_zero_copy (默认 False)
mp_policy: MixedPrecisionPolicy 配置 (默认 None)
replicate_params: 复制参数集合 (lambda m: {m.embedding} 或 None 表示全分片)
ignored_params: 忽略参数集合 (lambda m: {m.w1} 或 None 表示全部分片)
offload_policy: CPU offload 策略 (默认 None)
outer_comm_fusion: 最外层 comm_fusion 配置 (默认 False)
recompute_policy: 重计算策略 ("MUST_RECOMPUTE", "MUST_SWAP", etc., 默认 None)
"""
# ZeRO-1: mesh=None，只对模型整体应用
if mesh is None:
# mp_policy 为 None 时不传递该参数
if mp_policy is None:
fully_shard(model, reshard_after_forward=reshard_after_forward,
comm_fusion_zero_copy=comm_fusion_zero_copy,
replicate_params=replicate_params, ignored_params=ignored_params,
offload_policy=offload_policy)
else:
fully_shard(model, reshard_after_forward=reshard_after_forward,
comm_fusion_zero_copy=comm_fusion_zero_copy, mp_policy=mp_policy,
replicate_params=replicate_params, ignored_params=ignored_params,
offload_policy=offload_policy)
else:
# ZeRO-2/ZeRO-3: mesh 不为 None
if apply_to_layers:
# 对每个 layer 应用 fully_shard
for idx, layer in enumerate(model.layers):
if mp_policy is None:
fully_shard(layer, mesh=mesh, reshard_after_forward=reshard_after_forward,
comm_fusion_zero_copy=comm_fusion_zero_copy,
replicate_params=replicate_params, ignored_params=ignored_params,
offload_policy=offload_policy)
else:
fully_shard(layer, mesh=mesh, reshard_after_forward=reshard_after_forward,
comm_fusion_zero_copy=comm_fusion_zero_copy, mp_policy=mp_policy,
replicate_params=replicate_params, ignored_params=ignored_params,
offload_policy=offload_policy)
# 对整个模型应用 fully_shard，最外层可配置 comm_fusion
if apply_root_fs:
if mp_policy is None:
fully_shard(model, mesh=mesh, reshard_after_forward=reshard_after_forward,
comm_fusion_zero_copy=comm_fusion_zero_copy,
replicate_params=replicate_params, ignored_params=ignored_params,
offload_policy=offload_policy)
else:
fully_shard(model, mesh=mesh, reshard_after_forward=reshard_after_forward,
comm_fusion_zero_copy=comm_fusion_zero_copy,
mp_policy=mp_policy,
replicate_params=replicate_params, ignored_params=ignored_params,
offload_policy=offload_policy)
def apply_fs_to_model_with_layer_override(model, mesh, reshard_after_forward=True,
mp_policy=None, target_layer_idx=0,
layer_comm_fusion_zero_copy=True,
layer_shard_placement_fn=None):
"""
对模型应用 fully_shard，并对指定 layer 应用特殊配置
用于 test_fully_shard_ms_check_002 等需要覆盖某 layer 配置的用例
Args:
model: 待应用 fully_shard 的模型
mesh: 设备网格
reshard_after_forward: 是否 ZeRO-3 模式
mp_policy: MixedPrecisionPolicy 配置
target_layer_idx: 目标 layer 索引 (对该 layer 应用特殊配置)
layer_comm_fusion_zero_copy: 目标 layer 的 comm_fusion_zero_copy 配置
layer_shard_placement_fn: 目标 layer 的 shard_placement_fn 配置
"""
# 对每个 layer 应用 fully_shard
for idx, layer in enumerate(model.layers):
if idx == target_layer_idx and layer_shard_placement_fn is not None:
# 对目标 layer 应用特殊配置
fully_shard(
layer,
mesh=mesh,
reshard_after_forward=reshard_after_forward,
mp_policy=mp_policy,
comm_fusion_zero_copy=layer_comm_fusion_zero_copy,
shard_placement_fn=layer_shard_placement_fn
)
else:
# 其他 layer 使用默认配置
fully_shard(
layer,
mesh=mesh,
reshard_after_forward=reshard_after_forward,
mp_policy=mp_policy
)
# 对整个模型应用 fully_shard
fully_shard(
model,
mesh=mesh,
reshard_after_forward=reshard_after_forward,
mp_policy=mp_policy
)
model.set_requires_gradient_sync(True)
def run_train_with_fully_shard(model, train_dataset, epochs=1, max_steps=10,
mesh=None, reshard_after_forward=True,
clip_norm=False, apply_fs=True, apply_to_layers=True,
comm_fusion_zero_copy=False, mp_policy=None,
replicate_params=None, ignored_params=None,
offload_policy=None, outer_comm_fusion=False,
recompute_policy=None, # "MUST_RECOMPUTE", "MUST_SWAP", etc.
set_gradient_sync_at_step=None,
set_gradient_sync_value=True, # True=启用梯度同步，False=关闭
set_reshard_after_forward_at_step=None,
set_is_last_backward_at_step=None,
set_reshard_after_backward=None,
rank=0, monitor_memory=True,
loss_scale_manager=None,
set_gradient_sync_at_first=False):
"""
执行 fully_shard 训练
Args:
model: 待训练模型
train_dataset: 训练数据集
epochs: 训练轮数
max_steps: 最大步数
mesh: 设备网格 (None 表示 ZeRO-1 模式)
reshard_after_forward: 是否 ZeRO-3 模式 (False 表示 ZeRO-1/ZeRO-2)
clip_norm: 是否梯度裁剪
apply_fs: 是否应用 fully_shard
apply_to_layers: 是否对每个 layer 单独应用 (特殊用例)
comm_fusion_zero_copy: 是否启用 comm_fusion_zero_copy (默认 False)
mp_policy: MixedPrecisionPolicy 配置 (默认 None)
replicate_params: 复制参数集合 (lambda m: {m.embedding} 或 None 表示全分片)
ignored_params: 忽略参数集合 (lambda m: {m.w1} 或 None 表示全部分片)
offload_policy: CPU offload 策略 (默认 None)
outer_comm_fusion: 最外层 comm_fusion 配置 (默认 False)
recompute_policy: 重计算策略 ("MUST_RECOMPUTE", "MUST_SWAP", etc., 默认 None)
set_gradient_sync_at_step: 在第 N 个 step 调用 set_requires_gradient_sync (None 表示训练前调用)
set_gradient_sync_value: set_requires_gradient_sync 的值 (True=启用，False=关闭，默认 True)
set_reshard_after_forward_at_step: 在第 N 个 step 后调用 set_reshard_after_forward(False) (None 表示不切换)
set_is_last_backward_at_step: 在第 N 个 step 后调用 set_is_last_backward(True) (None 表示不设置)
set_reshard_after_backward: 设置 reshard_after_backward 值 (None 表示使用默认值)
rank: 进程 rank (用于日志)
monitor_memory: 是否监控显存
Returns:
loss_list, grad_dict, memory_list (if monitor_memory=True)
"""
# 应用 fully_shard (普通用例：在训练前应用)
if apply_fs:
apply_fully_shard(
model=model,
mesh=mesh,
reshard_after_forward=reshard_after_forward,
apply_root_fs=True,
apply_to_layers=apply_to_layers,
comm_fusion_zero_copy=comm_fusion_zero_copy,
mp_policy=mp_policy,
replicate_params=replicate_params,
ignored_params=ignored_params,
offload_policy=offload_policy,
outer_comm_fusion=outer_comm_fusion,
recompute_policy=recompute_policy
)
if set_reshard_after_backward is not None:
model.set_reshard_after_backward(set_reshard_after_backward)
logger.warning(f'set_reshard_after_backward: {set_reshard_after_backward}')
if apply_fs and set_gradient_sync_at_step is None and set_gradient_sync_at_first:
model.set_requires_gradient_sync(set_gradient_sync_value)
logger.warning(f'set_requires_gradient_sync: {set_gradient_sync_value}')
trainable_params = list(filter(lambda p: p.requires_grad, model.get_parameters()))
optimizer = nn.Adam(trainable_params, learning_rate=1e-4)
def train_step(inputs, labels):
for param in trainable_params:
param.grad = None
loss = model(inputs, labels)
with SkipDTensorDispatch():
loss.backward()
if loss_scale_manager is None:
grads = tuple(p.grad for p in model.trainable_params())
with _no_grad():
optimizer(grads)
else:
loss_scale = loss_scale_manager.get_loss_scale()
scaled_grads = tuple(p.grad for p in model.trainable_params())
grad_finite = all(ops.isfinite(g).all().asnumpy() for g in scaled_grads)
overflow = not grad_finite
if grad_finite:
with _no_grad():
optimizer(scaled_grads)
loss_scale_manager.update_loss_scale(overflow)
return loss
loss_list = []
grad_dict = {}
memory_list = [] # [(before, after), ...]
step = 0
gradient_sync_set = False # 标记是否已调用 set_requires_gradient_sync
reshard_switched = False # 标记是否已切换 reshard_after_forward
is_last_backward_set = False # 标记是否已设置 is_last_backward
if rank == 0 and monitor_memory:
logger.warning(f'\nStarting fully_shard training for {max_steps} steps...')
logger.warning(' Note: Recording memory from both npu-smi and runtime.memory_allocated()')
if set_gradient_sync_at_step:
logger.warning(f' Will call set_requires_gradient_sync(True) at step {set_gradient_sync_at_step}')
if set_reshard_after_forward_at_step:
logger.warning(f' Will call set_reshard_after_forward(False) after step {set_reshard_after_forward_at_step}')
if set_is_last_backward_at_step:
logger.warning(f' Will call set_is_last_backward(True) after step {set_is_last_backward_at_step}')
for _ in range(epochs):
for _, (input_ids, labels) in enumerate(train_dataset):
# 在指定 step 调用 set_requires_gradient_sync(value)
if set_gradient_sync_at_step is not None and step + 1 == set_gradient_sync_at_step and not gradient_sync_set:
model.set_requires_gradient_sync(set_gradient_sync_value)
if rank == 0:
logger.warning(f' Step {step+1}: Called set_requires_gradient_sync({set_gradient_sync_value})')
gradient_sync_set = True
# 在指定 step 后切换 reshard_after_forward (ZeRO-3 -> ZeRO-2)
if set_reshard_after_forward_at_step is not None and step == set_reshard_after_forward_at_step and not reshard_switched:
model.set_reshard_after_forward(False)
if rank == 0:
logger.warning(f' Step {step}: Called set_reshard_after_forward(False) - switching from ZeRO-3 to ZeRO-2')
reshard_switched = True
# 在指定 step 后设置 is_last_backward=True (通知最后一步反向，触发资源清理)
if set_is_last_backward_at_step is not None and step == set_is_last_backward_at_step and not is_last_backward_set:
model.set_is_last_backward(True)
if rank == 0:
logger.warning(f' Step {step}: Called set_is_last_backward(True) - signaling last backward pass')
is_last_backward_set = True
# Record memory usage before step (both methods)
if monitor_memory:
mem_before = get_memory_usage(rank=rank)
if rank == 0:
print_memory_usage(f'Step {step+1} - Before', mem_before)
loss = train_step(input_ids, labels)
# Record memory usage after step (both methods)
if monitor_memory:
mem_after = get_memory_usage(rank=rank)
if rank == 0:
print_memory_usage(f'Step {step+1} - After', mem_after)
memory_list.append((mem_before, mem_after))
loss_list.append(loss.item())
# 打印 loss (仅 rank 0)
if rank == 0:
logger.warning(f' Step {step+1}: Loss = {loss.item():.6f}')
# 保存梯度用于对比 (从 optimizer 参数获取)
if step == 0:
for i, param in enumerate(optimizer.parameters):
grad_dict[f'layer_{i}'] = param.grad.asnumpy() if param.grad is not None else None
step += 1
if step >= max_steps:
break
if step >= max_steps:
break
if monitor_memory:
return loss_list, grad_dict, memory_list
return loss_list, grad_dict
def sync_weights(single_model, fs_model):
"""
同步单卡模型和 fully_shard 模型的权重
Args:
single_model: 单卡参考模型
fs_model: fully_shard 模型
"""
params_single = single_model.parameters_dict()
params_fs = fs_model.parameters_dict()
rank = get_rank()
for key in params_fs:
params_fs[key].set_data(params_single[key])
# ============================================================================
# 显存监控函数
# ============================================================================
def get_memory_usage_npu_smi(rank=None):
"""
Get current NPU memory usage in MB by parsing npu-smi info output.
Automatically gets the current rank from the communication environment.
Note:
Line formula: tail -n (16 - rank * 2)
- rank=0: tail -n 16 | head -n 1
- rank=1: tail -n 14 | head -n 1
- rank=2: tail -n 12 | head -n 1
- rank=3: tail -n 10 | head -n 1
- rank=4: tail -n 8 | head -n 1
- rank=5: tail -n 6 | head -n 1
- rank=6: tail -n 4 | head -n 1
- rank=7: tail -n 2 | head -n 1
"""
import subprocess
try:
if rank is None:
rank = get_rank()
line_num = 16 - rank * 2
cmd = f'npu-smi info | tail -n {line_num} | head -n 1 | awk -F"|" \'{{print $5}}\''
result = subprocess.run(
cmd,
shell=True,
capture_output=True,
text=True,
timeout=5
)
output = result.stdout.strip()
if output and output.isdigit():
return float(output)
return 0.0
except Exception as e:
logger.warning(f'Warning: Failed to get memory usage from npu-smi (rank={rank}): {e}')
return 0.0
def get_memory_usage_runtime():
"""Get current NPU memory allocated in MB using mindspore.runtime.memory_allocated()."""
try:
from mindspore import runtime
# memory_allocated() returns bytes, convert to MB
allocated_bytes = runtime.memory_allocated()
allocated_mb = allocated_bytes / (1024 * 1024)
return allocated_mb
except Exception as e:
logger.warning(f'Warning: Failed to get memory usage from runtime: {e}')
return 0.0
def get_memory_usage(rank=None):
"""Get memory usage from both methods and return as tuple (npu_smi, runtime)."""
npu_smi = get_memory_usage_npu_smi(rank)
runtime_mb = get_memory_usage_runtime()
return npu_smi, runtime_mb
def print_memory_usage(label, mem_usage):
"""Print memory usage from both methods."""
npu_smi, runtime = mem_usage
logger.warning(f' {label}: npu-smi={npu_smi:.2f} MB, runtime={runtime:.2f} MB (diff={abs(npu_smi - runtime):.2f} MB)')
def print_memory_summary(test_name, memory_list, rank=0):
"""
打印显存使用汇总 (共用方法)
Args:
test_name: 测试用例名称
memory_list: [(before, after), ...] 格式的显存列表
rank: 进程 rank (仅 rank 0 打印)
"""
if rank != 0:
return
if not memory_list:
logger.warning(f'No memory data to summarize for {test_name}')
return
logger.warning('\n' + '='*70)
logger.warning(f'Memory Usage Summary ({test_name})')
logger.warning('='*70)
# 计算显存变化统计
mem_before_list = [m[0] for m in memory_list]
mem_after_list = [m[1] for m in memory_list]
# 计算平均显存使用
avg_before_npu_smi = np.mean([m[0] for m in mem_before_list])
avg_before_runtime = np.mean([m[1] for m in mem_before_list])
avg_after_npu_smi = np.mean([m[0] for m in mem_after_list])
avg_after_runtime = np.mean([m[1] for m in mem_after_list])
logger.warning(f'\nAverage Memory Usage:')
logger.warning(f' Before step: npu-smi={avg_before_npu_smi:.2f} MB, runtime={avg_before_runtime:.2f} MB')
logger.warning(f' After step: npu-smi={avg_after_npu_smi:.2f} MB, runtime={avg_after_runtime:.2f} MB')
logger.warning(f' Avg Delta: npu-smi={avg_after_npu_smi - avg_before_npu_smi:+.2f} MB, runtime={avg_after_runtime - avg_before_runtime:+.2f} MB')
# 计算每 step 显存变化
logger.warning(f'\nPer-step Memory Changes:')
for i, ((before_npu, before_rt), (after_npu, after_rt)) in enumerate(memory_list):
delta_npu = after_npu - before_npu
delta_rt = after_rt - before_rt
logger.warning(f' Step {i+1}: Δ={delta_npu:+.2f} MB (npu-smi), {delta_rt:+.2f} MB (runtime)')
logger.warning('='*70)
def compare_loss(single_loss_list, fs_loss_list, test_name='', rank=0, tol=0.1):
"""
对比单卡和 fully_shard 的 loss (共用方法)
Args:
single_loss_list: 单卡训练 loss 列表
fs_loss_list: fully_shard 训练 loss 列表
test_name: 测试用例名称 (可选)
rank: 进程 rank (仅 rank 0 打印)
tol: 容忍误差 (默认 0.1)
Returns:
max_diff: 最大 loss 差异
"""
if rank != 0:
return 0.0
if not single_loss_list or not fs_loss_list:
logger.warning(f'Cannot compare loss: empty loss list(s)')
return float('inf')
max_len = min(len(single_loss_list), len(fs_loss_list))
logger.warning('\n' + '='*70)
if test_name:
logger.warning(f'Loss Comparison ({test_name})')
else:
logger.warning('Loss Comparison (Single-card vs fully_shard)')
logger.warning('='*70)
max_diff = 0.0
for i in range(max_len):
single_loss = single_loss_list[i]
fs_loss = fs_loss_list[i]
diff = abs(single_loss - fs_loss)
max_diff = max(max_diff, diff)
logger.warning(f' Step {i+1}: Single={single_loss:.6f}, FS={fs_loss:.6f}, Diff={diff:.6f}')
logger.warning(f'\nFinal Loss:')
logger.warning(f' Single-card: {single_loss_list[-1]:.6f}')
logger.warning(f' fully_shard: {fs_loss_list[-1]:.6f}')
logger.warning(f' Max Diff: {max_diff:.6f}')
logger.warning(f' Final Diff: {abs(single_loss_list[-1] - fs_loss_list[-1]):.6f}')
logger.warning('='*70)
assert max_diff < tol, f"Loss difference too large: max_diff={max_diff}, tol={tol}"
return max_diff
MengXY107
6月12日 评论：
6月12日 评论：
def test_fully_shard_ms_check_007():
"""
只针对每个layer配置不同的fully_shard，不对整个net配置fully_shard
ZeRO-2 基础训练 + replicate_params={module.embedding}
TEST_SUMMARY: mesh=(8,), reshard_after_forward=False, replicate_params 复制 embedding 参数
Design Description:
fully_shard 单模块 + ZeRO-2 基础训练，mesh=(8,)，设置 replicate_params={module.embedding}
Test Steps:
1. 构建测试网络
2. mesh=init_device_mesh(mesh=(8,))
fully_shard(model, mesh=mesh, reshard_after_forward=False, replicate_params={module.embedding})
3. 设置 `set_requires_gradient_sync(True)`
4. 执行 10 step 训练
Expected Result:
1. 精度有问题，不 hang
"""
rank = get_rank()
world_size = get_group_size()
epochs = 1
max_steps = 10
batch_size = 2
seq_len = 32
vocab_size = 1000
model_config = get_model_config(num_layers=2, vocab_size=vocab_size, hidden_size=256, seq_len=seq_len)
# 创建单卡模型和 fully_shard 模型
single_model = PanGuMoETransformerSingle(model_config)
fs_model = PanGuMoETransformerFSInit(model_config)
# 同步权重
sync_weights(single_model, fs_model)
# ✅ 获取需要复制的参数集合 (embedding 权重)
# replicate_params 会导致这些参数在所有卡上复制，不分片
# 注意：MindSpore 中 Embedding 层的参数名是 embedding_table 而不是 weight
# 参考：E:\mindspore_ci\MindSporeTest\parallel\fully_shard\test_fully_shard_zero2_replicate.py
replicate_params_set = None
# 直接访问 fs_model.embedding (PanGuMoETransformer 结构)
if hasattr(fs_model, 'embedding'):
embedding_layer = fs_model.embedding
if hasattr(embedding_layer, 'embedding_table'):
replicate_params_set = {embedding_layer.embedding_table}
logger.warning(f'\ntest_fully_shard_ms_check_007: Found embedding parameter: {replicate_params_set}')
else:
logger.warning('\ntest_fully_shard_ms_check_007: Embedding layer has no embedding_table attribute')
else:
logger.warning('\ntest_fully_shard_ms_check_007: Model has no embedding attribute')
# 生成数据集
all_data = list(generator_dataset(size=100, batch_size=batch_size, seq_len=seq_len, vocab_size=vocab_size))
train_dataset_1 = GeneratorDataset(all_data[:], ["data", "label"], shuffle=False)
train_dataset_2 = GeneratorDataset(all_data[:], ["data", "label"], shuffle=False)
# 初始化 mesh: (8,) 单维 mesh (8 卡全用)
mesh = init_device_mesh(device_type="npu", mesh_shape=(8,), mesh_dim_names=("dp",))
# 执行单卡训练并监控显存
single_loss_list, single_grad_dict, single_memory_list = run_train_single(
single_model, train_dataset_1, epochs=epochs, max_steps=max_steps,
rank=rank, monitor_memory=True
)
apply_fully_shard(fs_model, mesh=mesh, reshard_after_forward=False, apply_root_fs=False, apply_to_layers=True,
comm_fusion_zero_copy=False, replicate_params=replicate_params_set)
# 执行 fully_shard 训练 (ZeRO-2 模式 + replicate_params 复制 embedding)
# 预期结果：精度有问题，但不 hang
loss_list, grad_dict, memory_list = run_train_with_fully_shard(
fs_model, train_dataset_2, epochs=epochs, max_steps=max_steps,
mesh=mesh,
reshard_after_forward=False, # ZeRO-2 模式
apply_fs=False,
apply_to_layers=False,
replicate_params=replicate_params_set, # ✅ 复制 embedding 参数，不分片
rank=rank,
monitor_memory=True
)
# 验证 loss 数据 (只验证不 hang，不验证精度)
assert len(loss_list) == max_steps
assert all(np.isfinite(l) for l in loss_list)
assert len(single_loss_list) == max_steps
assert all(np.isfinite(l) for l in single_loss_list)
# 验证显存监控数据
assert len(memory_list) == max_steps, f"Expected {max_steps} memory records, got {len(memory_list)}"
assert len(single_memory_list) == max_steps, f"Expected {max_steps} single-card memory records, got {len(single_memory_list)}"
# 打印显存使用汇总
print_memory_summary('test_fully_shard_ms_check_007 - fully_shard', memory_list, rank=rank)
print_memory_summary('test_fully_shard_ms_check_007 - single_card', single_memory_list, rank=rank)
# 注意：此用例预期精度有问题，不做严格的 loss 对比
# 只验证训练能完成，不 hang
logger.warning('\ntest_fully_shard_ms_check_007: Expected accuracy degradation with replicate_params')
该问题是怎么引起的？
重现步骤
parallel_config.py
moe_transformer_parallel_demo.py
test_pp.py
报错信息