[bug] PP+FSDP PTA后端 和standalone 进行精度比较时，loss偏低

该问题是怎么引起的？

重现步骤

parallel_config.py

"""
并行策略配置模块
基于 accelerate_config.yaml 提取，适配单机 8 卡环境
"""

from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any


@dataclass
class TensorParallelConfig:
    """张量并行配置"""
    enabled: bool = True
    world_size: int = 2  # 单机 8 卡：TP=2
    overlap_method: str = "mc2"  # MC2 通信计算重叠


@dataclass
class ExpertParallelConfig:
    """专家并行配置"""
    enabled: bool = True
    world_size: int = 4  # 单机 8 卡：EP=4
    inner_group_size: int = 2  # EP 内部分组
    token_dispatcher_type: str = "pangu_dropless"  # 无丢弃路由
    grouped_gemm: bool = True  # 分组 GEMM
    ep_over_sp: bool = True  # EP 优先于 SP
    shared_expert_sp: bool = True  # 共享专家启用 SP
    overlap_all2all: bool = False
    overlap_1f1b: bool = True


@dataclass
class SequenceParallelConfig:
    """序列并行配置"""
    enabled: bool = True
    use_sequence_parallel: bool = True


@dataclass
class DDPConfig:
    """DDP 配置"""
    use_distributed_optimizer: bool = True
    overlap_param_gather: bool = True
    overlap_grad_reduce: bool = True
    force_bucketing: bool = True
    align_optim_buffer: bool = True
    ddp_bucket_size: int = 83886080  # 80MB
    disable_gloo_group: bool = True


@dataclass
class PrecisionConfig:
    """精度配置"""
    compute_dtype: str = "bfloat16"
    softmax_compute_dtype: str = "float32"
    logits_compute_dtype: str = "float32"
    grad_accumulation_dtype: str = "float32"
    residual_connection_dtype: str = "bfloat16"
    
    fused_operators: List[str] = field(default_factory=lambda: [
        "flash_attention",
        "rmsnorm",
        "swiglu",
        "rope",
        "moe_permute_unpermute",
        "fused_mome"
    ])
    
    nofused_operators: List[str] = field(default_factory=lambda: [
        "masked_softmax"
    ])


@dataclass
class RecomputeConfig:
    """重计算配置"""
    enabled: bool = True
    moe_recompute: List[str] = field(default_factory=lambda: [
        "moe_permute",
        "moe_activation",
        "moe_rmsnorm"
    ])
    attention_recompute: List[str] = field(default_factory=lambda: [
        "moe_mla_qkv",
        "dense_mla_qkv"
    ])


@dataclass
class GCConfig:
    """垃圾回收配置"""
    manual: bool = True
    interval: int = 200  # 每 200 步


@dataclass
class ParallelConfig:
    """
    完整并行配置
    
    单机 8 卡策略:
    - TP (张量并行) = 2
    - EP (专家并行) = 4
    - SP (序列并行) = True
    - PP (流水线并行) = 1 (禁用)
    
    总 GPU 数 = TP × EP = 2 × 4 = 8
    """
    # 并行策略
    tensor_parallel: TensorParallelConfig = field(default_factory=TensorParallelConfig)
    expert_parallel: ExpertParallelConfig = field(default_factory=ExpertParallelConfig)
    sequence_parallel: SequenceParallelConfig = field(default_factory=SequenceParallelConfig)
    ddp: DDPConfig = field(default_factory=DDPConfig)
    
    # 精度
    precision: PrecisionConfig = field(default_factory=PrecisionConfig)
    
    # 重计算
    recompute: RecomputeConfig = field(default_factory=RecomputeConfig)
    
    # GC
    gc: GCConfig = field(default_factory=GCConfig)
    
    # 微批次大小
    micro_batch_size: int = 3
    
    # 序列分块
    chunk_seq_num: int = 4
    
    # 日志
    log_throughput: bool = True
    
    def __post_init__(self):
        # 验证配置
        total_gpus = self.tensor_parallel.world_size * self.expert_parallel.world_size
        if total_gpus != 8:
            print(f"警告：配置的总 GPU 数 ({total_gpus}) 不等于 8")
    
    def get_world_size(self) -> int:
        """获取总 GPU 数"""
        return self.tensor_parallel.world_size * self.expert_parallel.world_size
    
    def get_tp_rank(self) -> int:
        """获取 TP 秩"""
        return self.tensor_parallel.world_size
    
    def get_ep_rank(self) -> int:
        """获取 EP 秩"""
        return self.expert_parallel.world_size
    
    def is_sequence_parallel(self) -> bool:
        """是否启用序列并行"""
        return self.sequence_parallel.enabled and self.sequence_parallel.use_sequence_parallel


# ==================== 预设配置 ====================

def get_single_node_8gpu_config() -> ParallelConfig:
    """
    获取单机 8 卡预设配置
    
    配置说明:
    - TP=2 (张量并行 2 路)
    - EP=4 (专家并行 4 路)
    - SP=True (序列并行启用)
    - MBS=3 (微批次大小 3)
    - BF16 主计算精度
    """
    return ParallelConfig(
        tensor_parallel=TensorParallelConfig(
            enabled=True,
            world_size=2,
            overlap_method="mc2"
        ),
        expert_parallel=ExpertParallelConfig(
            enabled=True,
            world_size=4,
            inner_group_size=2,
            token_dispatcher_type="pangu_dropless",
            grouped_gemm=True,
            ep_over_sp=True,
            shared_expert_sp=True
        ),
        sequence_parallel=SequenceParallelConfig(
            enabled=True,
            use_sequence_parallel=True
        ),
        micro_batch_size=3,
        chunk_seq_num=4
    )


def get_single_node_4gpu_config() -> ParallelConfig:
    """
    获取单机 4 卡预设配置
    
    配置说明:
    - TP=2 (张量并行 2 路)
    - EP=2 (专家并行 2 路)
    - SP=True (序列并行启用)
    """
    config = get_single_node_8gpu_config()
    config.expert_parallel.world_size = 2
    config.expert_parallel.inner_group_size = 1
    return config


def get_single_gpu_config() -> ParallelConfig:
    """
    获取单机单卡配置 (调试用)
    
    所有并行策略禁用
    """
    return ParallelConfig(
        tensor_parallel=TensorParallelConfig(enabled=False, world_size=1),
        expert_parallel=ExpertParallelConfig(enabled=False, world_size=1),
        sequence_parallel=SequenceParallelConfig(enabled=False),
        micro_batch_size=1,
        chunk_seq_num=1
    )


# ==================== 使用示例 ====================

if __name__ == "__main__":
    print("=" * 60)
    print("并行配置示例")
    print("=" * 60)
    
    # 单机 8 卡配置
    config = get_single_node_8gpu_config()
    
    print(f"\n并行策略:")
    print(f"  - TP (张量并行): {config.get_tp_rank()}")
    print(f"  - EP (专家并行): {config.get_ep_rank()}")
    print(f"  - SP (序列并行): {config.is_sequence_parallel()}")
    print(f"  - 总 GPU 数：{config.get_world_size()}")
    
    print(f"\n精度配置:")
    print(f"  - 计算精度：{config.precision.compute_dtype}")
    print(f"  - Softmax 精度：{config.precision.softmax_compute_dtype}")
    print(f"  - 融合算子：{len(config.precision.fused_operators)} 个")
    
    print(f"\n优化特性:")
    print(f"  - 重计算：{config.recompute.enabled}")
    print(f"  - Grouped GEMM: {config.expert_parallel.grouped_gemm}")
    print(f"  - 手动 GC: {config.gc.manual} (间隔：{config.gc.interval}步)")
    print(f"  - 微批次大小：{config.micro_batch_size}")
    print(f"  - 序列分块：{config.chunk_seq_num}")
    
    print("\n" + "=" * 60)

moe_transformer_parallel_demo.py

"""
PanGu MoE Transformer 并行示例 (单机 8 卡)
基于 accelerate_config.yaml 配置抽象

核心特性:
- 单机 8 卡并行 (TP=2, EP=4)
- 序列并行 (SP)
- 分布式 DDP
- MoE (Mixture of Experts): 192 专家选 10
- MLA (Multi-Head Latent Attention): 低秩注意力
- SwiGLU 激活
- RMSNorm 归一化
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Optional
from parallel_config import ParallelConfig, get_single_node_8gpu_config, TensorParallelConfig, ExpertParallelConfig


# ==================== 配置类 ====================

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]


class ModelConfig:
    """模型配置 (从 yaml 提取关键参数)"""
    vocab_size: int = 153600
    hidden_size: int = 2560
    num_layers: int = 45
    seq_len: int = 4096
    num_heads: int = 32
    q_lora_rank: int = 768
    kv_lora_rank: int = 512
    qk_rope_head_dim: int = 64
    ffn_hidden: int = 6144
    activation: str = 'swiglu'
    moe: MoEConfig = MoEConfig()
    parallel: Optional[ParallelConfig] = None


# ==================== 分布式工具函数 ====================

def get_rank():
    """获取当前 GPU 秩"""
    if torch.distributed.is_initialized():
        return torch.distributed.get_rank()
    return 0

def get_world_size():
    """获取总 GPU 数"""
    if torch.distributed.is_initialized():
        return torch.distributed.get_world_size()
    return 1

def is_main_process():
    """是否为主进程"""
    return get_rank() == 0


# ==================== 核心组件 ====================

class RMSNorm(nn.Module):
    """RMS 归一化"""
    def __init__(self, hidden_size: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x / rms * self.weight


class MoEGating(nn.Module):
    """MoE 路由器"""
    def __init__(self, config: MoEConfig, hidden_size: int):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(hidden_size, config.num_experts, bias=False)
        if config.sigmoid_gating:
            self.sigmoid = nn.Sigmoid()
    
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        gate_logits = self.gate(x)
        if self.config.sigmoid_gating:
            gate_logits = self.sigmoid(gate_logits)
        topk_weights, topk_indices = torch.topk(gate_logits, self.config.top_k, dim=-1)
        topk_weights = F.softmax(topk_weights, dim=-1) * self.config.routed_scaling
        return topk_indices, topk_weights


class Expert(nn.Module):
    """单个专家网络"""
    def __init__(self, config: MoEConfig, hidden_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, config.expert_hidden_size)
        self.up_proj = nn.Linear(hidden_size, config.expert_hidden_size)
        self.down_proj = nn.Linear(config.expert_hidden_size, hidden_size)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class SharedExpert(nn.Module):
    """共享专家"""
    def __init__(self, config: MoEConfig, hidden_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, config.shared_hidden_size)
        self.up_proj = nn.Linear(hidden_size, config.shared_hidden_size)
        self.down_proj = nn.Linear(config.shared_hidden_size, hidden_size)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class DenseFFN(nn.Module):
    """稠密 FFN 层 (第 0 层)"""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.ffn_hidden)
        self.up_proj = nn.Linear(config.hidden_size, config.ffn_hidden)
        self.down_proj = nn.Linear(config.ffn_hidden, config.hidden_size)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class MoELayer(nn.Module):
    """MoE 层 (支持专家并行)"""
    def __init__(self, config: MoEConfig, hidden_size: int, ffn_hidden: int = None,
                 ep_rank: int = 0, ep_world_size: int = 1):
        super().__init__()
        self.config = config
        self.hidden_size = hidden_size
        self.ffn_hidden = ffn_hidden or config.expert_hidden_size
        self.ep_rank = ep_rank
        self.ep_world_size = ep_world_size
        
        # 计算当前 EP 组负责的专家范围
        experts_per_ep = config.num_experts // ep_world_size
        self.local_expert_start = ep_rank * experts_per_ep
        self.local_expert_end = self.local_expert_start + experts_per_ep
        
        # 本地专家 (只创建当前 EP 组负责的专家)
        self.local_experts = nn.ModuleList([
            Expert(config, hidden_size) 
            for _ in range(experts_per_ep)
        ])
        
        # 共享专家 (所有 EP 组都有)
        self.shared_expert = SharedExpert(config, hidden_size)
        
        # 路由器
        self.gate = MoEGating(config, hidden_size)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seq_len, batch, hidden = x.shape
        x_flat = x.reshape(-1, hidden)
        
        # 路由决策
        topk_indices, topk_weights = self.gate(x_flat)
        
        # 初始化输出
        output = torch.zeros_like(x_flat)
        
        # 只处理本地专家范围内的 token
        for local_idx in range(len(self.local_experts)):
            expert_idx = self.local_expert_start + local_idx
            expert = self.local_experts[local_idx]
            
            # 找到选择此专家的 token
            token_mask = (topk_indices == expert_idx)
            
            if token_mask.any():
                selected_tokens = x_flat[token_mask.any(dim=1)]
                token_weights = topk_weights[token_mask]
                expert_out = expert(selected_tokens)
                output[token_mask.any(dim=1)] += expert_out * token_weights.unsqueeze(-1)
        
        # 加上共享专家输出
        shared_out = self.shared_expert(x_flat)
        output = output + shared_out
        
        return output.reshape(seq_len, batch, hidden)


class MLAAttention(nn.Module):
    """简化版 MLA 注意力 (支持张量并行)"""
    def __init__(self, config: ModelConfig, tp_rank: int = 0, tp_world_size: int = 1):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_heads
        self.head_dim = config.qk_rope_head_dim
        self.tp_rank = tp_rank
        self.tp_world_size = tp_world_size
        
        # TP: 将头维度分块
        self.local_heads = self.num_heads // tp_world_size
        local_head_dim = self.local_heads * self.head_dim
        
        self.q_lora_down = nn.Linear(config.hidden_size, config.q_lora_rank)
        self.q_lora_up = nn.Linear(config.q_lora_rank, local_head_dim)
        
        self.kv_lora_down = nn.Linear(config.hidden_size, config.kv_lora_rank)
        self.kv_lora_up = nn.Linear(config.kv_lora_rank, local_head_dim * 2)
        
        self.out_proj = nn.Linear(local_head_dim, config.hidden_size)
    
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        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 = kv.chunk(2, dim=-1)
        
        # 重塑为多头
        q = q.view(seq_len, batch, self.local_heads, self.head_dim).transpose(0, 1).transpose(1, 2)
        k = k.view(seq_len, batch, self.local_heads, self.head_dim).transpose(0, 1).transpose(1, 2)
        v = v.view(seq_len, batch, self.local_heads, self.head_dim).transpose(0, 1).transpose(1, 2)
        
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        
        out = out.transpose(1, 2).contiguous().view(seq_len, batch, -1)
        
        return self.out_proj(out)


class TransformerBlock(nn.Module):
    """Transformer 块 (支持 TP + EP)"""
    def __init__(self, config: ModelConfig, layer_idx: int,
                 tp_rank: int = 0, tp_world_size: int = 1,
                 ep_rank: int = 0, ep_world_size: int = 1):
        super().__init__()
        self.norm1 = RMSNorm(config.hidden_size)
        self.norm2 = RMSNorm(config.hidden_size)
        self.attention = MLAAttention(config, tp_rank, tp_world_size)
        
        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, ep_rank, ep_world_size)
    
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        h = x + self.attention(self.norm1(x), mask)
        out = h + self.moe(self.norm2(h))
        return out


class PanGuMoETransformer(nn.Module):
    """PanGu MoE Transformer (支持并行)"""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        
        tp_cfg = config.parallel.tensor_parallel if config.parallel else None
        ep_cfg = config.parallel.expert_parallel if config.parallel else None
        
        tp_rank = get_rank() % tp_cfg.world_size if tp_cfg and tp_cfg.enabled else 0
        tp_world_size = tp_cfg.world_size if tp_cfg and tp_cfg.enabled else 1
        
        ep_rank = get_rank() // tp_cfg.world_size if tp_cfg and tp_cfg.enabled else 0
        ep_world_size = ep_cfg.world_size if ep_cfg and ep_cfg.enabled else 1
        
        self.tp_rank = tp_rank
        self.tp_world_size = tp_world_size
        self.ep_rank = ep_rank
        self.ep_world_size = ep_world_size
        
        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        
        self.layers = nn.ModuleList([
            TransformerBlock(config, i, tp_rank, tp_world_size, ep_rank, ep_world_size)
            for i in range(config.num_layers)
        ])
        
        self.final_norm = RMSNorm(config.hidden_size)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
    
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(input_ids).transpose(0, 1)
        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)
        return self.lm_head(x)


# ==================== 使用示例 ====================

if __name__ == '__main__':
    print("=" * 60)
    print("PanGu MoE Transformer 并行示例 (单机 8 卡)")
    print("=" * 60)
    
    # 获取单机 8 卡配置
    parallel_config = get_single_node_8gpu_config()
    
    print(f"\n并行配置:")
    print(f"  - TP (张量并行): {parallel_config.get_tp_rank()}")
    print(f"  - EP (专家并行): {parallel_config.get_ep_rank()}")
    print(f"  - SP (序列并行): {parallel_config.is_sequence_parallel()}")
    print(f"  - 总 GPU 数：{parallel_config.get_world_size()}")
    
    # 创建模型配置
    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 = 16  # 16 专家 (EP=4, 每卡 4 专家)
    demo_config.moe.top_k = 4
    demo_config.parallel = parallel_config
    
    model = PanGuMoETransformer(demo_config)
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型参数量：{total_params / 1e6:.2f}M")
    
    # 打印并行信息
    print(f"\n并行信息:")
    print(f"  - TP Rank: {model.tp_rank}, World Size: {model.tp_world_size}")
    print(f"  - EP Rank: {model.ep_rank}, World Size: {model.ep_world_size}")
    print(f"  - 每 EP 组专家数：{demo_config.moe.num_experts // model.ep_world_size}")
    
    # 前向传播
    print("\n前向传播测试:")
    batch_size = 2
    seq_len = 32
    input_ids = torch.randint(0, demo_config.vocab_size, (batch_size, seq_len))
    
    with torch.no_grad():
        logits = model(input_ids)
    
    print(f"  输入形状：{input_ids.shape}")
    print(f"  输出形状：{logits.shape}")
    print(f"  [OK] 前向传播成功!")
    
    print("\n" + "=" * 60)
    print("示例完成！支持单机 8 卡并行训练")
    print("=" * 60)

test_pp.py

报错信息

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 193
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/193

Start with the reproduction material in parallel_config.py and moe_transformer_parallel_demo.py, then run the PP+FSDP PTA-backend and standalone comparisons described in the issue. Trace how the parallel and precision configurations affect loss computation and compare the reported losses. Done means the two execution modes produce consistent accuracy, or the discrepancy is isolated to a specific component.