MindSporeTest/hyper_parallel_test/network/torch/minimind/model_minimind.py

"""
MiniMind Model - PyTorch 版本
基于 D:\\code\\minimind\\model\\model_minimind.py 移植
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, List


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

class MiniMindConfig:
    """MiniMind 模型配置"""
    def __init__(self, 
                 hidden_size: int = 768,
                 num_hidden_layers: int = 8,
                 use_moe: bool = False,
                 dropout: float = 0.0,
                 vocab_size: int = 6400,
                 bos_token_id: int = 1,
                 eos_token_id: int = 2,
                 flash_attn: bool = True,
                 num_attention_heads: int = 8,
                 num_key_value_heads: int = 4,
                 head_dim: int = None,
                 hidden_act: str = 'silu',
                 intermediate_size: int = None,
                 max_position_embeddings: int = 32768,
                 rms_norm_eps: float = 1e-6,
                 rope_theta: float = 1e6,
                 tie_word_embeddings: bool = True,
                 num_experts: int = 4,
                 num_experts_per_tok: int = 1,
                 moe_intermediate_size: int = None,
                 norm_topk_prob: bool = True,
                 router_aux_loss_coef: float = 5e-4):
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.use_moe = use_moe
        self.dropout = dropout
        self.vocab_size = vocab_size
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.flash_attn = flash_attn
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size if intermediate_size is not None else math.ceil(hidden_size * math.pi / 64) * 64
        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.tie_word_embeddings = tie_word_embeddings
        # MoE config
        self.num_experts = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.moe_intermediate_size = moe_intermediate_size if moe_intermediate_size is not None else self.intermediate_size
        self.norm_topk_prob = norm_topk_prob
        self.router_aux_loss_coef = router_aux_loss_coef


# ==================== 基础组件 ====================

class RMSNorm(nn.Module):
    """RMS 归一化"""
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
        return (x / rms * self.weight).to(x.dtype)


def precompute_freqs_cis(dim: int, end: int = 32768, rope_base: float = 1e6, 
                         rope_scaling: Optional[dict] = None) -> Tuple[torch.Tensor, torch.Tensor]:
    """预计算 RoPE 频率
    
    支持奇数 dim：当 dim 为奇数时，最后一个维度不参与旋转（与 Llama-3.1 一致）
    """
    # 支持奇数 dim：向上取整计算频率数量
    dim_half = (dim + 1) // 2  # ⭐ 向上取整，支持奇数
    freqs = 1.0 / (rope_base ** (torch.arange(0, dim_half).float() / dim))
    attn_factor = 1.0
    
    if rope_scaling is not None:
        orig_max = rope_scaling.get("original_max_position_embeddings", 2048)
        factor = rope_scaling.get("factor", 16)
        beta_fast = rope_scaling.get("beta_fast", 32.0)
        beta_slow = rope_scaling.get("beta_slow", 1.0)
        
        if end / orig_max > 1.0:
            def inv_dim(b):
                return (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
            
            low = max(math.floor(inv_dim(beta_fast)), 0)
            high = min(math.ceil(inv_dim(beta_slow)), dim_half - 1)
            
            ramp = torch.clamp((torch.arange(dim_half).float() - low) / max(high - low, 0.001), 0, 1)
            freqs = freqs * (1 - ramp + ramp / factor)
    
    t = torch.arange(end, dtype=torch.float32)
    freqs = torch.ger(t, freqs)
    
    # 拼接时处理奇数 dim：如果 dim 是奇数，需要截断最后一个重复的维度
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
    
    # 截断到目标 dim（处理奇数情况）
    freqs_cos = freqs_cos[..., :dim]
    freqs_sin = freqs_sin[..., :dim]
    
    return freqs_cos, freqs_sin


def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """应用 RoPE 旋转位置编码
    
    支持两种输入格式：
    1. [bsz, seq_len, num_heads, head_dim] (transpose 前)
    2. [bsz, num_heads, seq_len, head_dim] (transpose 后)
    
    TP 兼容：支持 q/k 的 head_dim 被 TP 切分的情况
    
    奇数 head_dim 处理：最后一个维度不参与旋转（与 Llama-3.1 一致）
    """
    def rotate_half(x):
        # 支持奇数 head_dim：中间维度 = head_dim // 2
        # 例如：head_dim=7 → half=3，前 3 维和后 3 维交换，第 4 维（索引 3）保持不变
        half = x.shape[-1] // 2
        x1 = x[..., :half]
        x2 = x[..., half:2*half]
        x_rest = x[..., 2*half:]  # 奇数时剩余的维度
        rotated = torch.cat((-x2, x1), dim=-1)
        return torch.cat([rotated, x_rest], dim=-1) if x_rest.numel() > 0 else rotated
    
    # cos/sin shape: [seq_len, head_dim]
    q_head_dim = q.shape[-1]
    
    # TP 兼容：如果 cos/sin 的维度大于 q/k 的 head_dim，截取前 q_head_dim 维
    if cos.shape[-1] > q_head_dim:
        cos = cos[..., :q_head_dim]
        sin = sin[..., :q_head_dim]
    
    # 判断 q 的格式并调整 cos/sin 的 shape
    if q.dim() == 4:
        if q.shape[1] == q.shape[2]:  # 无法区分，假设是 [B, S, H, D]
            # 检查 seq_len 维度
            if q.shape[1] == cos.shape[0]:  # [B, S, H, D]
                # cos: [S, D] -> [1, S, 1, D]
                cos = cos.unsqueeze(0).unsqueeze(2)
                sin = sin.unsqueeze(0).unsqueeze(2)
            else:  # [B, H, S, D]
                # cos: [S, D] -> [1, 1, S, D]
                cos = cos.unsqueeze(0).unsqueeze(1)
                sin = sin.unsqueeze(0).unsqueeze(1)
        elif q.shape[1] < q.shape[2]:  # [B, H, S, D]
            cos = cos.unsqueeze(0).unsqueeze(1)
            sin = sin.unsqueeze(0).unsqueeze(1)
        else:  # [B, S, H, D]
            cos = cos.unsqueeze(0).unsqueeze(2)
            sin = sin.unsqueeze(0).unsqueeze(2)
    
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    
    return q_embed.to(q.dtype), k_embed.to(k.dtype)


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """重复 KV 头"""
    if n_rep == 1:
        return x
    
    bs, slen, num_key_value_heads, head_dim = x.shape
    x_expanded = x.view(bs, slen, num_key_value_heads, 1, head_dim)
    x_expanded = x_expanded.expand(bs, slen, num_key_value_heads, n_rep, head_dim)
    return x_expanded.reshape(bs, slen, num_key_value_heads * n_rep, head_dim)


# ==================== Attention 组件 ====================

class MiniMindBshdSdpaCore(nn.Module):
    """BSHD layout scaled dot-product attention core for CP support.
    
    Exposed as a submodule so hyper_parallel.ContextParallel can apply
    context parallel on forward(q, k, v) with BSHD tensors [B, S, H, D].
    
    CP workflow:
    1. Input: Each rank has full sequence [B, S, H, D]
    2. CP hook inserts all-to-all to shard sequence dimension
    3. This module receives sharded [B, S/CP, H, D]
    4. Computes SDPA on local shard
    5. CP hook gathers output back to [B, S, H, D]
    
    RoPE handling:
    - RoPE should be applied AFTER CP sharding (in CP-aware wrapper)
    - This module assumes RoPE is already applied to Q/K
    """
    def __init__(self, dropout: float = 0.0, is_causal: bool = True):
        super().__init__()
        self.dropout = dropout
        self.is_causal = is_causal
    
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        """Run SDPA on Q/K/V in BSHD layout and return BSHD output.
        
        Args:
            q: Query [B, S, H, D] (may be sharded by CP, with RoPE already applied)
            k: Key [B, S, H, D] (may be sharded by CP, with RoPE already applied)
            v: Value [B, S, H, D] (may be sharded by CP)
        
        Returns:
            Attention output [B, S, H, D]
        """
        # Transpose to [B, H, S, D] for PyTorch SDPA
        qh = q.transpose(1, 2)
        kh = k.transpose(1, 2)
        vh = v.transpose(1, 2)
        
        out = F.scaled_dot_product_attention(
            qh, kh, vh,
            attn_mask=None,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=self.is_causal
        )
        return out.transpose(1, 2)


class MiniMindAttentionCpWrapper(nn.Module):
    """CP-aware attention wrapper that delays RoPE until after CP sharding.
    
    This wrapper replaces the original Attention module when CP is enabled.
    It follows the Qwen3VL pattern:
    1. Compute Q/K/V projections
    2. CP sharding (via ContextParallel hook on sdpa_core)
    3. Apply RoPE based on global position embeddings (after sharding)
    4. Run SDPA
    5. Gather output
    
    Usage:
        # Replace original attention with CP wrapper
        for layer in model.layers:
            layer.self_attn = MiniMindAttentionCpWrapper(layer.self_attn, config)
        
        # Apply CP to the inner sdpa_core
        cp_plan = ContextParallel(seq_dim=2, head_dim=3, ulysses_degree=cp_size)
        for layer in model.layers:
            cp_plan.apply(layer.self_attn.sdpa_core, cp_mesh)
    """
    
    def __init__(self, attention_module: 'Attention', config: MiniMindConfig):
        super().__init__()
        # Copy all attributes from original attention
        self.n_local_heads = attention_module.n_local_heads
        self.n_local_kv_heads = attention_module.n_local_kv_heads
        self.n_rep = attention_module.n_rep
        self.head_dim = attention_module.head_dim
        self.is_causal = attention_module.is_causal
        
        # Copy projection layers
        self.q_proj = attention_module.q_proj
        self.k_proj = attention_module.k_proj
        self.v_proj = attention_module.v_proj
        self.o_proj = attention_module.o_proj
        
        # Copy normalization layers
        self.q_norm = attention_module.q_norm
        self.k_norm = attention_module.k_norm
        
        # Copy dropout
        self.attn_dropout = attention_module.attn_dropout
        self.resid_dropout = attention_module.resid_dropout
        self.dropout = attention_module.dropout
        self.flash = attention_module.flash
        
        # Keep reference to sdpa_core (will be wrapped by CP)
        self.sdpa_core = attention_module.sdpa_core
        
        # Store config for RoPE
        self.rope_theta = config.rope_theta
        self.max_position_embeddings = config.max_position_embeddings
    
    def forward(self, x: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor],
                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
                use_cache: bool = False,
                attention_mask: Optional[torch.Tensor] = None,
                global_position_ids: Optional[torch.Tensor] = None,
                **kwargs) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        """Forward with CP-aware RoPE application.
        
        Args:
            x: Input hidden states [B, S, H]
            position_embeddings: (cos, sin) for RoPE [S, D]
            global_position_ids: Position IDs for global sequence (required for CP)
            **kwargs: Additional kwargs (may include _hp_cp_* from CP wrapper)
        
        Returns:
            Attention output [B, S, H], past_key_value
        """
        bsz, seq_len, _ = x.shape
        
        # Project Q/K/V
        xq = self.q_proj(x)
        xk = self.k_proj(x)
        xv = self.v_proj(x)
        
        # Reshape to BSHD
        actual_heads = xq.shape[-1] // self.head_dim
        actual_kv_heads = xk.shape[-1] // self.head_dim
        
        xq = xq.view(bsz, seq_len, actual_heads, self.head_dim)
        xk = xk.view(bsz, seq_len, actual_kv_heads, self.head_dim)
        xv = xv.view(bsz, seq_len, actual_kv_heads, self.head_dim)
        
        # Normalize
        xq = self.q_norm(xq)
        xk = self.k_norm(xk)
        
        # Get global position embeddings for CP
        # CP wrapper will pass _hp_cp_global_position_embeddings via kwargs
        global_pos_emb = kwargs.pop('_hp_cp_global_position_embeddings', None)
        if global_pos_emb is not None:
            # CP mode: use global position embeddings (will be applied after sharding)
            cos, sin = global_pos_emb
        else:
            # Non-CP mode: use provided position embeddings
            cos, sin = position_embeddings
        
        # Apply RoPE (will be applied after CP sharding if CP is enabled)
        xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
        
        # Handle KV cache (not supported in CP mode)
        if past_key_value is not None:
            xk = torch.cat([past_key_value[0], xk], dim=1)
            xv = torch.cat([past_key_value[1], xv], dim=1)
        
        past_kv = (xk, xv) if use_cache else None
        
        # Run SDPA (CP hook will be applied here if CP is enabled)
        output = self.sdpa_core(xq, xk, xv)
        
        # Project output
        output = output.reshape(bsz, seq_len, -1)
        output = self.resid_dropout(self.o_proj(output))
        
        return output, past_kv


class Attention(nn.Module):
    """多头注意力机制"""
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
        self.n_local_heads = config.num_attention_heads
        self.n_local_kv_heads = self.num_key_value_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = config.head_dim
        self.is_causal = True
        
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        
        # CP support: extract SDPA core as separate module
        self.sdpa_core = MiniMindBshdSdpaCore(
            dropout=config.dropout,
            is_causal=config.flash_attn  # Use flash attention if enabled
        )
        
        self.attn_dropout = nn.Dropout(p=config.dropout)
        self.resid_dropout = nn.Dropout(p=config.dropout)
        self.dropout = config.dropout
        self.flash = config.flash_attn

    def forward(self, x: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], 
                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, 
                use_cache: bool = False, 
                attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        bsz, seq_len, _ = x.shape
        
        xq = self.q_proj(x)
        xk = self.k_proj(x)
        xv = self.v_proj(x)
        
        # TP 兼容：获取实际的 head 数量（从输出 shape 推断）
        actual_heads = xq.shape[-1] // self.head_dim
        actual_kv_heads = xk.shape[-1] // self.head_dim
        
        xq = xq.view(bsz, seq_len, actual_heads, self.head_dim)
        xk = xk.view(bsz, seq_len, actual_kv_heads, self.head_dim)
        xv = xv.view(bsz, seq_len, actual_kv_heads, self.head_dim)
        
        xq = self.q_norm(xq)
        xk = self.k_norm(xk)
        
        cos, sin = position_embeddings
        xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
        
        if past_key_value is not None:
            xk = torch.cat([past_key_value[0], xk], dim=1)
            xv = torch.cat([past_key_value[1], xv], dim=1)
        
        past_kv = (xk, xv) if use_cache else None
        
        # Use sdpa_core for attention computation (CP compatible)
        # Input/Output shape: [B, S, H, D]
        output = self.sdpa_core(xq, xk, xv)
        
        # output shape: [B, S, H_local, D]
        # reshape 为 [B, S, H_local * D]
        output = output.reshape(bsz, seq_len, -1)
        
        # TP 兼容：o_proj 期望输入是 [B, S, hidden_size]
        # 但当前 output 的最后一个维度是 H_local * D = hidden_size / TP
        # o_proj (RowwiseParallel) 会自动处理 all-reduce
        output = self.resid_dropout(self.o_proj(output))
        
        return output, past_kv


# ==================== FFN 组件 ====================

class FeedForward(nn.Module):
    """前馈神经网络"""
    def __init__(self, config: MiniMindConfig, intermediate_size: int = None):
        super().__init__()
        intermediate_size = intermediate_size or config.intermediate_size
        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.act_fn = nn.SiLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class MOEFeedForward(nn.Module):
    """MoE 前馈神经网络"""
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
        self.experts = nn.ModuleList([
            FeedForward(config, intermediate_size=config.moe_intermediate_size) 
            for _ in range(config.num_experts)
        ])
        self.act_fn = nn.SiLU()
        self.aux_loss = 0.0

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, hidden_dim = x.shape
        x_flat = x.reshape(-1, hidden_dim)
        
        scores = F.softmax(self.gate(x_flat), dim=-1)
        topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1)
        
        if self.config.norm_topk_prob:
            topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
        
        y = torch.zeros_like(x_flat)
        
        for i, expert in enumerate(self.experts):
            mask = (topk_idx == i)
            if mask.any():
                token_idx = mask.any(dim=-1).nonzero().flatten()
                weight = topk_weight[mask].view(-1, 1)
                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
            elif self.training:
                y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
        
        if self.training and self.config.router_aux_loss_coef > 0:
            load = F.one_hot(topk_idx, self.config.num_experts).float().mean(dim=0)
            self.aux_loss = (load * scores.mean(dim=0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
        else:
            self.aux_loss = 0.0
        
        return y.reshape(batch_size, seq_len, hidden_dim)


# ==================== Transformer Block ====================

class MiniMindBlock(nn.Module):
    """MiniMind Transformer 块"""
    def __init__(self, layer_id: int, config: MiniMindConfig):
        super().__init__()
        self.self_attn = Attention(config)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)

    def forward(self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor],
                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
                use_cache: bool = False,
                attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        residual = hidden_states
        
        hidden_states, present_key_value = self.self_attn(
            self.input_layernorm(hidden_states),
            position_embeddings,
            past_key_value=past_key_value,
            use_cache=use_cache,
            attention_mask=attention_mask
        )
        
        hidden_states = residual + hidden_states
        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
        
        return hidden_states, present_key_value


# ==================== MiniMind Model ====================

class MiniMindModel(nn.Module):
    """MiniMind 基础模型"""
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.config = config
        self.vocab_size = config.vocab_size
        self.num_hidden_layers = config.num_hidden_layers
        
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.dropout = nn.Dropout(p=config.dropout)
        
        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        freqs_cos, freqs_sin = precompute_freqs_cis(
            dim=config.head_dim,
            end=config.max_position_embeddings,
            rope_base=config.rope_theta
        )
        self.register_buffer('freqs_cos', freqs_cos)
        self.register_buffer('freqs_sin', freqs_sin)

    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,
                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                use_cache: bool = False) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]], torch.Tensor]:
        batch_size, seq_length = input_ids.shape
        
        if past_key_values is None:
            past_key_values = [None] * len(self.layers)
        
        past_len = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
        
        hidden_states = self.dropout(self.embed_tokens(input_ids))
        
        position_embeddings = (
            self.freqs_cos[past_len:past_len + seq_length],
            self.freqs_sin[past_len:past_len + seq_length]
        )
        
        presents = []
        aux_loss = 0.0
        
        for layer, past_key_value in zip(self.layers, past_key_values):
            hidden_states, present = layer(
                hidden_states,
                position_embeddings,
                past_key_value=past_key_value,
                use_cache=use_cache,
                attention_mask=attention_mask
            )
            presents.append(present)
            
            if isinstance(layer.mlp, MOEFeedForward):
                aux_loss = aux_loss + layer.mlp.aux_loss
        
        hidden_states = self.norm(hidden_states)
        
        return hidden_states, presents, aux_loss


# ==================== MiniMind For Causal LM ====================

class MiniMindForCausalLM(nn.Module):
    """MiniMind 因果语言模型"""
    def __init__(self, config: MiniMindConfig = None):
        super().__init__()
        self.config = config if config is not None else MiniMindConfig()
        self.model = MiniMindModel(self.config)
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
        
        if self.config.tie_word_embeddings:
            self.model.embed_tokens.weight = self.lm_head.weight

    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,
                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                use_cache: bool = False,
                labels: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], float]:
        hidden_states, past_key_values, aux_loss = self.model(
            input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache
        )
        
        logits = self.lm_head(hidden_states)
        
        loss = None
        if labels is not None:
            x = logits[:, :-1, :].reshape(-1, logits.shape[-1])
            y = labels[:, 1:].reshape(-1)
            loss = F.cross_entropy(x, y, ignore_index=-100)
        
        return logits, past_key_values, loss, aux_loss


# ==================== 工具函数 ====================

def get_miniMind_config(use_moe: bool = False, num_layers: int = 4, 
                       vocab_size: int = 6400, hidden_size: int = 256,
                       seq_len: int = 32, num_experts: int = 4,
                       top_k: int = 2, num_attention_heads: int = 8,
                       num_key_value_heads: int = 8 ) -> MiniMindConfig:
    """获取 MiniMind 测试配置"""
    config = MiniMindConfig(
        hidden_size=hidden_size,
        num_hidden_layers=num_layers,
        use_moe=use_moe,
        vocab_size=vocab_size,
        max_position_embeddings=seq_len,
        num_attention_heads=num_attention_heads,
        num_key_value_heads=num_key_value_heads,  # 修改为 8 以支持 TP=8
        num_experts=num_experts,
        num_experts_per_tok=top_k
    )
    return config

MindSporeTest/hyper_parallel_test/network/torch/minimind/parallel_minimand.py

# Copyright 2026 Huawei Technologies Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Parallel support for MiniMind model (PyTorch version).

Supports:
- PP (Pipeline Parallel): Layer-wise partitioning across stages
- EP (Expert Parallel): Expert parallelism for MoE
- TP (Tensor Parallel): Intra-layer tensor sharding
- SP (Sequence Parallel): Sequence dimension sharding
- CP (Context Parallel): Context parallelism for long sequences
- HSDP (Hybrid Sharded Data Parallel): Sharded optimizer states

All parallel strategies are controlled via ParallelConfig parameters.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch import nn
from typing import Optional, List, Dict, Any, Tuple
from dataclasses import dataclass, field
from functools import partial

from hyper_parallel import (
    ColwiseParallel,
    RowwiseParallel,
    SequenceParallel,
    ContextParallel,
    init_device_mesh,
    parallelize_module,
)
from hyper_parallel.core.expert_parallel.expert_parallel import ExpertParallel
from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
from hyper_parallel.core.dtensor.placement_types import Replicate, Shard

from .model_minimind import MiniMindForCausalLM, MiniMindConfig


@dataclass
class ParallelConfig:
    """Parallel configuration for MiniMind model.
    
    Controls which parallel strategies are enabled via parameters.
    
    Args:
        dp: Data parallel size (default: 1)
        tp: Tensor parallel size (default: 1)
        pp: Pipeline parallel size (default: 1)
        ep: Expert parallel size (default: 1)
        sp: Sequence parallel size (default: 1)
        cp: Context parallel size (default: 1)
        enable_pp: Enable pipeline parallelism (default: False)
        enable_ep: Enable expert parallelism (default: False)
        enable_tp: Enable tensor parallelism (default: False)
        enable_sp: Enable sequence parallelism (default: False)
        enable_cp: Enable context parallelism (default: False)
        pp_schedule: Pipeline schedule ('1f1b', 'interleaved', default: '1f1b')
        ep_group_type: EP communication type ('alltoall', 'allgather', default: 'alltoall')
    """
    dp: int = 1
    tp: int = 1
    pp: int = 1
    ep: int = 1
    sp: int = 1
    cp: int = 1
    
    enable_pp: bool = False
    enable_ep: bool = False
    enable_tp: bool = False
    enable_sp: bool = False
    enable_cp: bool = False
    
    pp_schedule: str = "1f1b"
    ep_group_type: str = "alltoall"
    
    def __post_init__(self):
        """Validate configuration"""
        if self.enable_pp and self.pp <= 1:
            raise ValueError("PP size must be > 1 when enable_pp=True")
        if self.enable_ep and self.ep <= 1:
            raise ValueError("EP size must be > 1 when enable_ep=True")
        if self.enable_tp and self.tp <= 1:
            raise ValueError("TP size must be > 1 when enable_tp=True")
        if self.enable_sp and self.sp <= 1:
            raise ValueError("SP size must be > 1 when enable_sp=True")
        if self.enable_cp and self.cp <= 1:
            raise ValueError("CP size must be > 1 when enable_cp=True")


class PipelineStage(nn.Module):
    """Pipeline parallel stage wrapper.
    
    Wraps a subset of layers for pipeline parallelism.
    Handles Send/Recv communication between stages.
    
    Args:
        module: Base module to wrap
        stage_id: Pipeline stage ID (0 to pp_size-1)
        pp_size: Total pipeline stages
        num_layers: Total layers in model
    """
    
    def __init__(
        self,
        module: nn.Module,
        stage_id: int,
        pp_size: int,
        num_layers: int,
    ):
        super().__init__()
        self.stage_id = stage_id
        self.pp_size = pp_size
        self.num_layers = num_layers
        self.module = module
        
        # Assign layers to this stage
        layers_per_stage = (num_layers + pp_size - 1) // pp_size
        start_layer = stage_id * layers_per_stage
        end_layer = min(start_layer + layers_per_stage, num_layers)
        
        self.start_layer = start_layer
        self.end_layer = end_layer
        self.is_first_stage = (stage_id == 0)
        self.is_last_stage = (stage_id == pp_size - 1)
        
        # Communication handles (initialized during forward)
        self.recv_handle = None
        self.send_handle = None
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass with pipeline communication.
        
        For first stage: receive inputs or use provided
        For intermediate stages: receive from previous, send to next
        For last stage: no Send, just compute
        """
        # First stage receives input if not provided
        if not self.is_first_stage and x is None:
            # Placeholder for recv - in real implementation would use dist.recv
            raise NotImplementedError("Pipeline recv not implemented")
        
        # Run layers for this stage
        for i in range(self.start_layer, self.end_layer):
            layer = self.module.model.layers[i]
            x = layer(x)
        
        # Last stage doesn't send
        if not self.is_last_stage:
            # Placeholder for send - in real implementation would use dist.send
            pass
        
        return x


class ExpertParallelGroup(nn.Module):
    """Expert parallel wrapper for MoE layers.
    
    Distributes experts across devices and handles token routing.
    
    Args:
        module: MoE layer to wrap
        ep_mesh: Device mesh for expert parallelism
        ep_group_type: Communication type ('alltoall' or 'allgather')
    """
    
    def __init__(
        self,
        module: nn.Module,
        ep_mesh: DeviceMesh,
        ep_group_type: str = "alltoall",
    ):
        super().__init__()
        self.module = module
        self.ep_mesh = ep_mesh
        self.ep_group_type = ep_group_type
        self.ep_world = ep_mesh.size()
        self.ep_rank = ep_mesh.get_local_rank()
        
        # Get expert count from module
        if hasattr(module, 'num_experts'):
            self.num_experts = module.num_experts
        else:
            self.num_experts = 1
        
        # Assign experts to this EP group
        self.experts_per_rank = (self.num_experts + self.ep_world - 1) // self.ep_world
        self.start_expert = self.ep_rank * self.experts_per_rank
        self.end_expert = min(self.start_expert + self.experts_per_rank, self.num_experts)
    
    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """Forward with expert parallel communication.
        
        1. All-to-All: Distribute tokens to expert owners
        2. Compute: Process local experts
        3. All-Gather: Collect outputs
        """
        # In real implementation:
        # 1. Route tokens to expert owners via All-to-All
        # 2. Compute on local experts
        # 3. All-Gather or All-to-All to collect outputs
        
        # Simplified: just run local experts
        # Real implementation would handle token routing
        return self.module(x, *args, **kwargs)


class TensorParallelDense(nn.Module):
    """Tensor parallel dense layer wrapper.
    
    Implements Colwise/Rowwise parallelism for linear layers.
    
    Args:
        module: Linear layer to parallelize
        tp_mesh: Device mesh for tensor parallelism
        parallel_type: 'colwise' or 'rowwise'
    """
    
    def __init__(
        self,
        module: nn.Linear,
        tp_mesh: DeviceMesh,
        parallel_type: str = "colwise",
    ):
        super().__init__()
        self.module = module
        self.tp_mesh = tp_mesh
        self.parallel_type = parallel_type
        self.tp_world = tp_mesh.size()
        self.tp_rank = tp_mesh.get_local_rank()
        
        # Shard weights
        if parallel_type == "colwise":
            # Shard output dimension
            out_features = module.out_features // self.tp_world
            start = self.tp_rank * out_features
            end = start + out_features
            
            # Create sharded weight (no grad for original)
            with torch.no_grad():
                module.weight = nn.Parameter(module.weight[start:end, :].clone())
            module.out_features = out_features
            
        elif parallel_type == "rowwise":
            # Shard input dimension
            in_features = module.in_features // self.tp_world
            start = self.tp_rank * in_features
            end = start + in_features
            
            with torch.no_grad():
                module.weight = nn.Parameter(module.weight[:, start:end].clone())
            module.in_features = in_features
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward with tensor parallel communication.
        
        Colwise: All-reduce outputs after linear
        Rowwise: Linear on sharded input, no communication needed
        """
        out = self.module(x)
        
        # Colwise needs all-reduce
        if self.parallel_type == "colwise":
            dist.all_reduce(out, op=dist.ReduceOp.SUM, group=self.tp_mesh.get_group())
        
        return out


class SequenceParallelAttention(nn.Module):
    """Sequence parallel attention wrapper.
    
    Shards sequence dimension across devices.
    
    Args:
        module: Attention layer to wrap
        sp_mesh: Device mesh for sequence parallelism
    """
    
    def __init__(
        self,
        module: nn.Module,
        sp_mesh: DeviceMesh,
    ):
        super().__init__()
        self.module = module
        self.sp_mesh = sp_mesh
        self.sp_world = sp_mesh.size()
        self.sp_rank = sp_mesh.get_local_rank()
    
    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """Forward with sequence parallel communication.
        
        1. All-Gather: Gather sequence chunks for attention
        2. Compute: Full attention
        3. Reduce-Scatter: Scatter outputs
        """
        # Simplified: just run attention
        # Real implementation would handle sequence sharding
        return self.module(x, *args, **kwargs)


# ==================== CP-Aware Attention Implementation ====================

def apply_rotary_pos_emb_cp(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, 
                            cp_rank: int, seq_len_local: int, head_dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Apply rotary position embedding with CP-aware global position calculation.
    
    Args:
        q: Query tensor (B, S_local, H, D)
        k: Key tensor (B, S_local, H, D)
        cos: Cosine embeddings for full sequence (S_global, D)
        sin: Sine embeddings for full sequence (S_global, D)
        cp_rank: CP rank ID
        seq_len_local: Local sequence length per CP rank
        head_dim: Head dimension
    
    Returns:
        q_rotated, k_rotated: RoPE applied tensors
    """
    # Calculate global position range for this CP rank
    start_pos = cp_rank * seq_len_local
    end_pos = start_pos + seq_len_local
    
    # Slice the correct RoPE embeddings for this CP rank's tokens
    cos_local = cos[start_pos:end_pos]  # (S_local, D)
    sin_local = sin[start_pos:end_pos]  # (S_local, D)
    
    # Expand to match q/k shape: (B, S, H, D)
    cos_local = cos_local.unsqueeze(0).unsqueeze(2)  # (1, S, 1, D)
    sin_local = sin_local.unsqueeze(0).unsqueeze(2)  # (1, S, 1, D)
    
    # Apply RoPE: q_rot = q * cos + rotate_half(q) * sin
    q_rotated = (q * cos_local) + (rotate_half(q) * sin_local)
    k_rotated = (k * cos_local) + (rotate_half(k) * sin_local)
    
    return q_rotated, k_rotated


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


class CpAwareSdpaCore(nn.Module):
    """CP-aware SDPA core that applies RoPE after context sharding.
    
    This ensures each CP rank uses the correct global position for RoPE.
    
    Args:
        original_sdpa_core: Original SDPA core module to wrap
        cp_rank: CP rank ID
        cp_size: Total CP size
        cos: Precomputed cosine embeddings (S_global, D)
        sin: Precomputed sine embeddings (S_global, D)
        head_dim: Head dimension
    """
    
    def __init__(
        self,
        original_sdpa_core: nn.Module,
        cp_rank: int,
        cp_size: int,
        cos: torch.Tensor,
        sin: torch.Tensor,
        head_dim: int,
    ):
        super().__init__()
        self.original_sdpa_core = original_sdpa_core
        self.cp_rank = cp_rank
        self.cp_size = cp_size
        self.register_buffer('cos', cos)
        self.register_buffer('sin', sin)
        self.head_dim = head_dim
    
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, 
                *args, **kwargs) -> torch.Tensor:
        """Forward with CP-aware RoPE application.
        
        Args:
            q: Query (B, S_local, H, D)
            k: Key (B, S_local, H, D)
            v: Value (B, S_local, H, D)
        
        Returns:
            attention_output: (B, S_local, H, D)
        """
        # Get local sequence length
        seq_len_local = q.shape[1]
        
        # Apply RoPE with correct global positions
        q_rotated, k_rotated = apply_rotary_pos_emb_cp(
            q, k, self.cos, self.sin, 
            self.cp_rank, seq_len_local, self.head_dim
        )
        
        # Run original SDPA with rotated q/k
        return self.original_sdpa_core(q_rotated, k_rotated, v, *args, **kwargs)


class MiniMindAttentionCpWrapper(nn.Module):
    """CP-aware attention wrapper for MiniMind attention layers.
    
    Wraps the attention module to apply RoPE after CP sharding.
    
    Args:
        original_attention: Original MiniMind attention module
        cp_rank: CP rank ID
        cp_size: Total CP size
        cos: Precomputed cosine embeddings
        sin: Precomputed sine embeddings
    """
    
    def __init__(
        self,
        original_attention: nn.Module,
        cp_rank: int,
        cp_size: int,
        cos: torch.Tensor,
        sin: torch.Tensor,
    ):
        super().__init__()
        self.original_attention = original_attention
        self.cp_rank = cp_rank
        self.cp_size = cp_size
        
        # Get config from attention module
        if hasattr(original_attention, 'config'):
            self.head_dim = original_attention.config.head_dim
        elif hasattr(original_attention, 'head_dim'):
            self.head_dim = original_attention.head_dim
        else:
            # Infer from weight shape
            self.head_dim = original_attention.q_proj.weight.shape[0] // original_attention.num_heads
        
        # Wrap sdpa_core with CP-aware version
        if hasattr(original_attention, 'sdpa_core'):
            self.sdpa_core = CpAwareSdpaCore(
                original_attention.sdpa_core,
                cp_rank,
                cp_size,
                cos,
                sin,
                self.head_dim,
            )
        
        # Copy other attributes
        if hasattr(original_attention, 'q_proj'):
            self.q_proj = original_attention.q_proj
        if hasattr(original_attention, 'k_proj'):
            self.k_proj = original_attention.k_proj
        if hasattr(original_attention, 'v_proj'):
            self.v_proj = original_attention.v_proj
        if hasattr(original_attention, 'o_proj'):
            self.o_proj = original_attention.o_proj
        if hasattr(original_attention, 'num_heads'):
            self.num_heads = original_attention.num_heads
        if hasattr(original_attention, 'num_kv_heads'):
            self.num_kv_heads = original_attention.num_kv_heads
    
    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """Forward pass with CP-aware attention.
        
        RoPE is applied inside sdpa_core after CP sharding.
        """
        # Use original forward but with our CP-aware sdpa_core
        # The key is that sdpa_core now handles RoPE internally
        return self.original_attention(hidden_states, *args, **kwargs)


def create_cp_aware_attention(
    original_attention: nn.Module,
    config: MiniMindConfig,
    cp_rank: int,
    cp_size: int,
) -> nn.Module:
    """Convert original attention to CP-aware attention.
    
    Args:
        original_attention: Original attention module
        config: MiniMind configuration
        cp_rank: CP rank ID
        cp_size: Total CP size
    
    Returns:
        CP-aware attention wrapper
    """
    # Precompute RoPE embeddings for full sequence
    max_seq_len = config.max_position_embeddings
    head_dim = config.head_dim
    
    # Create RoPE embeddings
    inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    position = torch.arange(max_seq_len).float()
    freqs = torch.outer(position, inv_freq)
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    
    return MiniMindAttentionCpWrapper(
        original_attention,
        cp_rank,
        cp_size,
        cos,
        sin,
    )


class ParallelMiniMindModel(nn.Module):
    """Parallel MiniMind model wrapper.
    
    Applies PP/EP/TP/SP based on ParallelConfig.
    All strategies are controlled via parameters.
    
    Args:
        model: Base MiniMindForCausalLM model
        config: Parallel configuration
        mesh: Device mesh (can be multi-dimensional)
    """
    
    def __init__(
        self,
        model: MiniMindForCausalLM,
        config: ParallelConfig,
        mesh: DeviceMesh,
    ):
        super().__init__()
        self.model = model
        self.config = config
        self.mesh = mesh
        self.cfg = model.config
        
        # Apply parallel strategies based on config
        self._apply_parallel()
    
    def _apply_parallel(self):
        """Apply parallel strategies based on config."""
        
        # Apply CP first (if enabled) - must be before other strategies
        if self.config.enable_cp:
            self._apply_context_parallel()
        
        # Apply TP (if enabled)
        if self.config.enable_tp:
            self._apply_tensor_parallel()
        
        # Apply EP (if enabled and MoE)
        if self.config.enable_ep and self.cfg.use_moe:
            self._apply_expert_parallel()
        
        # Apply SP (if enabled)
        if self.config.enable_sp:
            self._apply_sequence_parallel()
        
        # Apply PP (if enabled) - wraps the entire model
        if self.config.enable_pp:
            self._apply_pipeline_parallel()
    
    def _apply_context_parallel(self):
        """Apply context parallelism to attention layers with CP-aware RoPE.
        
        CP must be applied before other strategies since it shards the sequence dimension.
        The key is to apply RoPE AFTER CP sharding with correct global positions.
        """
        cp_mesh = self.mesh["cp"] if "cp" in self.mesh.mesh_dim_names else self.mesh
        cp_rank = cp_mesh.get_local_rank()
        cp_size = cp_mesh.size()
        
        # Precompute RoPE embeddings for full sequence
        # Note: In real implementation, this would come from the model's RoPE module
        # For now, we create placeholder embeddings
        max_seq_len = 2048  # Should come from config
        head_dim = self.cfg.head_dim if hasattr(self.cfg, 'head_dim') else 64
        
        # Create RoPE embeddings
        inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
        position = torch.arange(max_seq_len).float()
        freqs = torch.outer(position, inv_freq)
        cos = torch.cos(freqs)
        sin = torch.sin(freqs)
        
        # Apply CP-aware wrapper to all attention layers
        for layer_idx in range(self.cfg.num_hidden_layers):
            layer = self.model.model.layers[layer_idx]
            
            # Wrap attention with CP-aware version
            if hasattr(layer, 'self_attn'):
                layer.self_attn = create_cp_aware_attention(
                    layer.self_attn,
                    cp_rank,
                    cp_size,
                    cos,
                    sin,
                )
    
    def _apply_tensor_parallel(self):
        """Apply tensor parallelism to attention and FFN layers."""
        tp_mesh = self.mesh["tp"] if "tp" in self.mesh.mesh_dim_names else self.mesh
        
        # Build layer plan for TP
        layer_plan = {}
        
        for layer_idx in range(self.cfg.num_hidden_layers):
            layer_prefix = f"model.layers.{layer_idx}"
            
            # Attention: Q/K/V Colwise, O Rowwise
            layer_plan[f"{layer_prefix}.self_attn.q_proj"] = ColwiseParallel(
                input_layouts=Replicate(),
                output_layouts=Shard(-1),
                use_local_output=False,
            )
            layer_plan[f"{layer_prefix}.self_attn.k_proj"] = ColwiseParallel(
                input_layouts=Replicate(),
                output_layouts=Shard(-1),
                use_local_output=False,
            )
            layer_plan[f"{layer_prefix}.self_attn.v_proj"] = ColwiseParallel(
                input_layouts=Replicate(),
                output_layouts=Shard(-1),
                use_local_output=False,
            )
            layer_plan[f"{layer_prefix}.self_attn.o_proj"] = RowwiseParallel(
                input_layouts=Shard(-1),
                output_layouts=Replicate(),
                use_local_output=not self.config.enable_sp,
            )
            
            # FFN layers (only for dense FFN, not MoE)
            if not self.cfg.use_moe:
                layer_plan[f"{layer_prefix}.mlp.gate_proj"] = ColwiseParallel(
                    input_layouts=Replicate(),
                    output_layouts=Shard(-1),
                    use_local_output=False,
                )
                layer_plan[f"{layer_prefix}.mlp.up_proj"] = ColwiseParallel(
                    input_layouts=Replicate(),
                    output_layouts=Shard(-1),
                    use_local_output=False,
                )
                layer_plan[f"{layer_prefix}.mlp.down_proj"] = RowwiseParallel(
                    input_layouts=Shard(-1),
                    output_layouts=Replicate(),
                    use_local_output=not self.config.enable_sp,
                )
        
        # Embedding and output
        layer_plan["model.embed_tokens"] = RowwiseParallel(
            input_layouts=Replicate(),
            output_layouts=Shard(-1),
            use_local_output=False,
        )
        
        # Apply TP
        parallelize_module(self.model, tp_mesh, layer_plan)
    
    def _apply_expert_parallel(self):
        """Apply expert parallelism to MoE layers using ExpertParallel.
        
        EP distributes experts across devices and handles token routing via All-to-All.
        This follows the pattern from test_pipeline_parallel_feature_interaction_torch.py.
        """
        ep_mesh = self.mesh["ep"] if "ep" in self.mesh.mesh_dim_names else self.mesh
        
        # Apply EP to each MoE layer's experts
        for layer_idx in range(self.cfg.num_hidden_layers):
            layer = self.model.model.layers[layer_idx]
            
            # Check if layer has MoE
            if hasattr(layer, 'mlp') and hasattr(layer.mlp, 'num_experts'):
                # Use ExpertParallel().apply() pattern from reference implementation
                ExpertParallel().apply(layer.mlp, ep_mesh)
    
    def _apply_sequence_parallel(self):
        """Apply sequence parallelism to attention and norms."""
        sp_mesh = self.mesh["sp"] if "sp" in self.mesh.mesh_dim_names else self.mesh
        
        # Wrap attention layers with SP
        for layer_idx in range(self.cfg.num_hidden_layers):
            layer = self.model.model.layers[layer_idx]
            
            # Wrap norms with SP
            layer.input_layernorm = SequenceParallelAttention(
                layer.input_layernorm,
                sp_mesh,
            )
            layer.post_attention_layernorm = SequenceParallelAttention(
                layer.post_attention_layernorm,
                sp_mesh,
            )
    
    def _apply_pipeline_parallel(self):
        """Apply pipeline parallelism by wrapping model stages."""
        pp_mesh = self.mesh["pp"] if "pp" in self.mesh.mesh_dim_names else self.mesh
        pp_size = self.config.pp
        pp_rank = pp_mesh.get_local_rank() if pp_mesh.size() > 1 else 0
        
        # Wrap model with pipeline stage
        self.model = PipelineStage(
            self.model,
            pp_rank,
            pp_size,
            self.cfg.num_hidden_layers,
        )
    
    def forward(self, *args, **kwargs):
        """Forward pass through parallel model."""
        return self.model(*args, **kwargs)
    
    def set_requires_gradient_sync(self, value: bool):
        """Set gradient sync requirement for fully_shard compatibility."""
        if hasattr(self.model, 'set_requires_gradient_sync'):
            self.model.set_requires_gradient_sync(value)


# ==================== Helper Functions ====================

def build_mesh(
    device_type: str = "cuda",
    mesh_shape: tuple = None,
    mesh_dim_names: tuple = None,
) -> DeviceMesh:
    """Build device mesh.
    
    Args:
        device_type: Device type ("cuda" for GPU/NPU)
        mesh_shape: Shape of mesh (e.g., (2, 2, 2) for 3D)
        mesh_dim_names: Dimension names (e.g., ("dp", "tp", "pp"))
    
    Returns:
        DeviceMesh instance
    """
    if mesh_shape is None:
        world_size = dist.get_world_size()
        mesh_shape = (world_size,)
        mesh_dim_names = ("dp",)
    
    return init_device_mesh(
        device_type=device_type,
        mesh_shape=mesh_shape,
        mesh_dim_names=mesh_dim_names,
    )


def verify_parallel_model(model: MiniMindForCausalLM, config: ParallelConfig) -> None:
    """Verify parallel model configuration.
    
    Args:
        model: Parallelized model
        config: Parallel config used
    """
    print(f"✓ Parallel model verification:")
    print(f"  TP={config.enable_tp} (size={config.tp})")
    print(f"  PP={config.enable_pp} (size={config.pp})")
    print(f"  EP={config.enable_ep} (size={config.ep})")
    print(f"  SP={config.enable_sp} (size={config.sp})")
    print(f"  CP={config.enable_cp} (size={config.cp})")
    print(f"  MoE={model.config.use_moe}")


# ==================== Usage Example ====================

if __name__ == "__main__":
    """
    Example usage with torchrun:
    
    # TP=2 on 2 GPUs:
    torchrun --nproc_per_node=2 parallel_minimind.py
    
    # CP=2 on 2 GPUs (Context Parallel):
    torchrun --nproc_per_node=2 parallel_minimind.py
    
    # TP+PP on 4 GPUs (TP=2, PP=2):
    torchrun --nproc_per_node=4 parallel_minimind.py
    
    # Full 5D on 8 GPUs (TP=2, PP=2, CP=2):
    torchrun --nproc_per_node=8 parallel_minimind.py
    """
    import os
    import torch.distributed as dist
    
    # Initialize distributed
    local_rank = int(os.environ.get('LOCAL_RANK', 0))
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend='nccl' if 'cuda' in str(torch.device(0)) else 'hccl')
    
    # Create config
    config = MiniMindConfig(
        use_moe=False,
        num_hidden_layers=2,
        hidden_size=256,
        num_attention_heads=8,
        num_key_value_heads=4,
        intermediate_size=512,
        vocab_size=1000,
    )
    
    # Create model
    model = MiniMindForCausalLM(config)
    
    # Build mesh
    world_size = dist.get_world_size()
    
    # Example: CP=2 on 2 GPUs
    mesh_shape = (world_size,)
    mesh = build_mesh("cuda", mesh_shape, ("cp",))
    
    print(f"[Rank {dist.get_rank()}] Applying CP with size {mesh.size()}...")
    
    # Create parallel config
    parallel_config = ParallelConfig(
        dp=1,
        tp=1,
        pp=1,
        ep=1,
        sp=1,
        cp=world_size,
        enable_tp=False,
        enable_pp=False,
        enable_ep=False,
        enable_sp=False,
        enable_cp=True,
    )
    
    # Apply parallel
    parallel_model = ParallelMiniMindModel(model, parallel_config, mesh)
    
    # Verify
    verify_parallel_model(parallel_model.model, parallel_config)
    
    # Test forward
    input_ids = torch.randint(0, config.vocab_size, (2, 32)).cuda(local_rank)
    
    print(f"[Rank {dist.get_rank()}] Running forward pass...")
    outputs = parallel_model(input_ids)
    
    print(f"[Rank {dist.get_rank()}] Forward pass successful")
    
    # Cleanup
    dist.destroy_process_group()

问题用例1：

def test_fully_shard_uneven_sharding_torch_001():
    """
    测试 FSDP 非均匀切分场景：dim-0 < shard_world_size
    
    测试场景:
    - dim-0 轴大小 (hidden_size=7) < shard_world_size (8)
    - 验证 FSDP 正确处理非均匀切分 (ceil 除法)
    - 与单卡训练对比 loss 一致性
    
    Test Steps:
    1. 构建测试网络，修改 hidden_size=7, head_dim=8, num_attention_heads=1
    2. mesh=init_device_mesh(mesh=(8,)), fully_shard(model, mesh=mesh)
    3. 设置 set_requires_gradient_sync(True)
    4. 执行 10 step 训练，通过与单卡训练对比，loss 序列与单卡比较在精度范围内
    
    通过标准:
    1. 训练过程无报错
    2. loss 序列与单卡训练差异在容忍范围内 (tol=0.01)
    """
    rank = dist.get_rank()
    # Step 1: 构建测试网络 (hidden_size=7 < world_size=8)
    model_config = get_model_config(num_layers=2, vocab_size=1000, hidden_size=7, seq_len=16, use_moe=False,
                                    num_attention_heads=1,
                                    num_key_value_heads=1
                                    )
    single_model = MiniMindForCausalLMSingle(model_config)
    fs_model = MiniMindForCausalLMFSInit(model_config)
    sync_model_weights(fs_model, single_model)
    single_model = single_model.to('npu')
    fs_model = fs_model.to('npu')
    mesh = init_device_mesh("npu", (8,))
    for layer in fs_model.model.layers:
        torch_fully_shard(layer, mesh=mesh, reshard_after_forward=True)
        layer.set_requires_gradient_sync(True)
    torch_fully_shard(fs_model, mesh=mesh, reshard_after_forward=True)
    fs_model.set_requires_gradient_sync(True)
    # 准备数据集
    data = list(generator_dataset(20, batch_size=2, seq_len=16, vocab_size=1000))
    train_dataset_single = GeneratorDataset(data, ["input_ids", "labels"])
    train_dataset_fs = GeneratorDataset(data, ["input_ids", "labels"])
    # Step 4: 执行训练
    single_loss_list = run_train_single(
        single_model, train_dataset_single, epochs=1, max_steps=10,
        rank=rank, monitor_memory=False
    )
    fs_loss_list = run_train_with_fully_shard(
        fs_model, train_dataset_fs, epochs=1, max_steps=10,
        mesh=mesh, apply_fs=False, rank=rank, monitor_memory=False,
        num_micro_batches=10
    )
    # Step 4: Loss 对比 (hidden_size=7 数值稳定性差，放宽容忍度)
    compare_loss(
        single_loss_list, fs_loss_list,
        test_name='FSDP uneven sharding (hidden_size=7 < ws=8)',
        tol=0.01  # ⭐ 从 0.01 放宽到 0.1
    )

问题用例2：---需要额外关注vocab_size测试的出问题的shape，dim0是1001

def test_fully_shard_uneven_sharding_torch_007():
    """
    测试 FSDP 非均匀切分场景 + replicate_params (部分 dim-0 不能整除的参数复制)
    
    测试场景:
    - vocab_size=1001 (1001 % 8 ≠ 0，非均匀切分)
    - 仅将 embed_tokens.weight 加入 replicate_params (lm_head.weight 不复制)
    - 验证 FSDP 正确处理部分参数复制模式
    - 与单卡训练对比 loss 一致性
    
    Test Steps:
    1. 构建测试网络，修改 vocab_size=1001
    2. mesh=init_device_mesh(mesh=(2,4)), fully_shard(model, mesh=mesh, replicate_params={embed_tokens.weight})
    3. 设置 set_requires_gradient_sync(True)
    4. 执行 10 step 训练
    
    Expected Result:
    1. loss 序列与单卡比较在精度范围内 (tol=0.01)
    
    参考：test_fully_shard_uneven_sharding_torch_006 (全部复制)
    """
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    logging.warning(f"MiniMind====>uneven_007: Running on rank {rank}, world_size {world_size}")
    logging.warning(f"MiniMind====>uneven_007: Testing uneven sharding + partial replicate_params (vocab_size=1001)")
    
    # Step 1: 构建测试网络 (vocab_size=1001, 1001 % 8 ≠ 0 非均匀切分)
    model_config = get_model_config(
        num_layers=4,
        vocab_size=1001,                # ⭐ 关键：vocab_size % 8 ≠ 0 (非均匀)
        hidden_size=128,
        seq_len=16,
        use_moe=False
    )
    
    # 创建单卡模型和 FS 模型
    single_model = MiniMindForCausalLMSingle(model_config)
    fs_model = MiniMindForCausalLMFSInit(model_config)
    
    # 同步权重
    sync_model_weights(fs_model, single_model)
    
    single_model = single_model.to('npu')
    fs_model = fs_model.to('npu')
    
    # Step 2: 初始化 mesh 并应用 fully_shard (部分 replicate_params 模式)
    # 仅复制 embed_tokens.weight，lm_head.weight 不复制 (切分)
    # embed_tokens.weight: [vocab_size=1001, hidden_size=128] -> 1001 % 8 ≠ 0 ⭐
    # lm_head.weight: [vocab_size=1001, hidden_size=128] -> 1001 % 8 ≠ 0 (但不复制)
    mesh = init_device_mesh("npu", (2, 4))
    
    # 定义 replicate_params (仅 embed_tokens.weight)
    replicate_params_set = set()
    
    # embed_tokens.weight: [1001, 128] -> 1001 % 8 ≠ 0 ⭐ 复制
    embed_weight = fs_model.model.embed_tokens.weight
    if embed_weight.shape[0] % world_size != 0:
        replicate_params_set.add(embed_weight)
        logging.warning(f"MiniMind====>uneven_007: Replicating embed_tokens.weight (shape={embed_weight.shape}, dim0={embed_weight.shape[0]} % {world_size} = {embed_weight.shape[0] % world_size})")
    
    # lm_head.weight: [1001, 128] -> 1001 % 8 ≠ 0 (但不加入 replicate_params，保持切分)
    lm_head_weight = fs_model.lm_head.weight
    if lm_head_weight.shape[0] % world_size != 0:
        logging.warning(f"MiniMind====>uneven_007: NOT replicating lm_head.weight (shape={lm_head_weight.shape}, dim0={lm_head_weight.shape[0]} % {world_size} = {lm_head_weight.shape[0] % world_size}) - testing partial replicate")
    ''' 
    # 对每个 layer 应用 fully_shard
    for layer in fs_model.model.layers:
        torch_fully_shard(
            layer,
            mesh=mesh,
            reshard_after_forward=True,
            replicate_params=None             # layer 参数全部切分
        )

    # ⭐ 对 lm_head 也应用 fully_shard (非均匀切分)
    logging.warning(f"MiniMind====>uneven_007: Applying fully_shard to lm_head...")
    torch_fully_shard(
        fs_model.lm_head,
        mesh=mesh,
        reshard_after_forward=True,
        replicate_params=None  # ⭐ lm_head 参数非均匀切分 (不复制)
    )
    
    # 对 root 模型应用 fully_shard (仅复制 embed_tokens.weight)
    torch_fully_shard(
        fs_model,
        mesh=mesh,
        reshard_after_forward=True,
        replicate_params=replicate_params_set  # ⭐ 仅复制 embed_tokens.weight
    )
    '''
    apply_fully_shard(fs_model,mesh=mesh, reshard_after_forward=True,replicate_params=replicate_params_set)
    # Step 3: 设置梯度同步
    fs_model.set_requires_gradient_sync(True)
    
    # 准备数据集
    data = list(generator_dataset(20, batch_size=2, seq_len=16, vocab_size=1001))
    train_dataset_single = GeneratorDataset(data, ["input_ids", "labels"])
    train_dataset_fs = GeneratorDataset(data, ["input_ids", "labels"])
    
    # Step 4a: 执行单卡训练
    logging.warning(f"MiniMind====>uneven_007: Training single model...")
    single_loss_list = run_train_single(
        single_model, train_dataset_single, epochs=1, max_steps=10,
        rank=rank, monitor_memory=False
    )
    
    # Step 4b: 执行 FS 训练 (部分 replicate_params 模式)
    logging.warning(f"MiniMind====>uneven_007: Training FS model with partial replicate_params...")
    fs_loss_list = run_train_with_fully_shard(
        fs_model, train_dataset_fs, epochs=1, max_steps=10,
        mesh=mesh, apply_fs=False, rank=rank, monitor_memory=False
    )
    
    # Expected Result: loss 序列与单卡比较在精度范围内
    logging.warning(f"MiniMind====>uneven_007: Comparing single vs FS loss...")
    logging.warning(f"  Single loss: {single_loss_list[:3]}... (first 3 steps)")
    logging.warning(f"  FS loss: {fs_loss_list[:3]}... (first 3 steps)")
    
    compare_loss(
        single_loss_list, fs_loss_list,
        test_name='FSDP uneven sharding + partial replicate_params (vocab_size=1001)',
        rank=rank, tol=0.01
    )

这是用例review意见，和测试用例本体无关。

针对测试用例 test_fully_shard_uneven_sharding_torch_001

问题判断

当前失败或卡住首先是测试脚本自身的问题，不能据此判断 FSDP 的 dim-0 < shard_world_size 非均匀切分有问题。

1. 单卡基线与 FSDP 侧训练语义不一致

   • 单卡侧未传 num_micro_batches，每个 optimizer step 只消费 1 个 batch。
   • FSDP 侧传入 num_micro_batches=10，每个 optimizer step 消费 10 个 batch，并对 loss 除以 10。
   • 两边的 loss_list[i] 代表不同数量的数据，不能按相同 step 下标直接比较。

2. 数据量不足会导致死循环

   • 数据集只有 20 个 batch。
   • FSDP 侧完成 2 个 optimizer step 后数据耗尽。
   • StopIteration 后 step_losses 为空，step 不再递增，但外层仍是 while step < 10，因此会永久空转在 step == 2。

   等价控制流如下：

   while step < max_steps:
       step_losses = []
       for _ in range(num_micro_batches):
           try:
               batch = next(data_iter)
           except StopIteration:
               break
           step_losses.append(run_one_micro_batch(batch))
   
       if step_losses:
           step += 1
       # step_losses 为空时既不退出，step 也不增加
   

3. 梯度累积控制与注释不一致

   • 注释说前 N-1 个 micro batch 关闭同步、最后一个 micro batch 打开同步。
   • 实际 explicit_reshard_after_backward is None 时，reshard_value 始终为 None，没有执行对应切换。
   • 注释掉的 set_requires_gradient_sync 代码也没有生效，因此本用例没有真正验证它描述的梯度累积状态切换。

4. 配置和注释存在偏差

   • 注释写 head_dim=8，但当前构造方式得到的是 head_dim = hidden_size // num_attention_heads = 7。
   • 注释说容差从 0.01 放宽到 0.1，实际仍传入 tol=0.01。

修复建议

如果本用例只验证“参数 dim-0 小于 8 卡 shard group”，建议不要混入梯度累积，保持两侧完全相同的 batch 和 optimizer step：

max_steps = 10
data = list(generator_dataset(
    max_steps,
    batch_size=2,
    seq_len=16,
    vocab_size=1000,
))

single_loss_list = run_train_single(
    single_model,
    GeneratorDataset(data, ["input_ids", "labels"]),
    max_steps=max_steps,
    rank=rank,
    num_micro_batches=1,
)

fs_loss_list = run_train_with_fully_shard(
    fs_model,
    GeneratorDataset(data, ["input_ids", "labels"]),
    max_steps=max_steps,
    mesh=mesh,
    apply_fs=False,
    rank=rank,
    monitor_memory=False,
    num_micro_batches=1,
)

compare_loss(
    single_loss_list,
    fs_loss_list,
    test_name="FSDP uneven dim0 smaller than shard group",
    rank=rank,
    tol=0.01,
)

如果确实要同时覆盖梯度累积，则需要：

num_micro_batches = 10
data_size = max_steps * num_micro_batches

# 单卡与 FSDP 两侧都使用相同的 num_micro_batches。
# 非最后一个 micro batch 关闭梯度同步，最后一个再打开。
model.set_requires_gradient_sync(micro_idx == num_micro_batches - 1)

if consumed_micro_batches == 0:
    raise RuntimeError("dataset exhausted before max_steps")

此外，建议把“dim-0=7”覆盖放到一个简单的 nn.Linear 或直接参数模型中，避免奇数 attention head_dim、RoPE 和 Transformer 训练逻辑干扰 FSDP 切分结论。

实际对照验证中，合法的 [7, 4] 参数在 8 卡 FSDP 下训练 3 步，与单卡 loss 差值均为 0.0。

这是用例review意见，和测试用例本体无关。

针对测试用例 test_fully_shard_uneven_sharding_torch_007

问题判断

该用例要验证两个独立参数采用不同策略：

• embed_tokens.weight 加入 replicate_params，保持复制；
• lm_head.weight 不复制，继续覆盖 vocab dim-0 非均匀切分。

原用例存在两个独立的脚本问题，只关闭 tie_word_embeddings 还不够。

1. 默认 tied weight 使测试场景本身不可表达

模型默认 tie_word_embeddings=True，构造函数执行：

self.model.embed_tokens.weight = self.lm_head.weight

所以两个名字实际指向同一个 Parameter：

fs_model.model.embed_tokens.weight is fs_model.lm_head.weight
# True

同一个参数不能同时采用 replicate 和 shard 两种策略。并且 apply_fully_shard() 先包装 lm_head、再包装 root，会把这个 alias 拆开，导致 FSDP 模型变为 untied，而单卡基线仍是 tied；此时 loss 比较的是两个不同语义的模型。

2. 2D mesh 没有名称，replicate 参数路径无法 flatten

原用例使用：

mesh = init_device_mesh("npu", (2, 4))

当 replicate_params 非空时，FSDP 会把 2D mesh flatten 成该参数的 1D DDP mesh。DeviceMesh.flatten() 的默认名称由原维度名拼接得到，因此 unnamed mesh 不支持无参 flatten。实际 8 卡复现的调用链为：

TorchHSDPStateV2._build_param_mesh_info
  -> self.mesh.flatten()
  -> DeviceMesh._create_flatten_mesh
  -> "_".join(self.mesh_dim_names)
TypeError: can only join an iterable

这就是把 tie_word_embeddings=False 后仍然失败的直接原因。当前裸 TypeError 的报错质量可在生产侧单独改进，但不表示 uneven shard 或 replicate 通信计算错误。

3. vocab_size=1000 的通过没有覆盖目标分支

用例按下面的条件构造复制集合：

if embed_weight.shape[0] % world_size != 0:
    replicate_params_set.add(embed_weight)

8 卡时 1000 % 8 == 0，所以 replicate_params_set 为空，FSDP 根本不会进入 replicate 参数的 2D-mesh flatten 路径；这个 pass 不能证明“部分 replicate_params + uneven shard”正常。必须保持 vocab_size=1001。另外，HSDP 的实际 shard group size 是 mesh.mesh_shape[-1] == 4，几何断言不应写死使用全局 world size。

测试脚本修复建议

两项修改必须同时存在：关闭 tied weight，并为 2D mesh 指定维度名。

model_config = get_model_config(
    num_layers=4,
    vocab_size=1001,
    hidden_size=128,
    seq_len=16,
    use_moe=False,
)
model_config.tie_word_embeddings = False

single_model = MiniMindForCausalLMSingle(model_config)
fs_model = MiniMindForCausalLMFSInit(model_config)
sync_model_weights(fs_model, single_model)

embed_weight = fs_model.model.embed_tokens.weight
lm_head_weight = fs_model.lm_head.weight

mesh = init_device_mesh(
    "npu",
    (2, 4),
    mesh_dim_names=("replicate", "shard"),
)
replicate_params = {embed_weight}

assert embed_weight is not lm_head_weight
assert embed_weight.shape[0] == 1001
assert embed_weight.shape[0] % mesh.mesh_shape[-1] != 0
assert embed_weight in replicate_params
assert lm_head_weight not in replicate_params

# lm_head.weight 是独立参数，覆盖 vocab=1001 uneven shard。
torch_fully_shard(
    fs_model.lm_head,
    mesh=mesh,
    reshard_after_forward=True,
)

# root 仅复制独立的 embedding 参数。
torch_fully_shard(
    fs_model,
    mesh=mesh,
    reshard_after_forward=True,
    replicate_params=replicate_params,
)

建议再增加一条断言，防止把 vocab_size 改成 1000 后用例悄悄退化为空集合：

assert replicate_params, "007 must exercise the replicate_params path"

实际复现结果

在当前代码 7ffbd297 上，使用 HYPER_PARALLEL_PLATFORM=torch、8 张 Ascend 910B3，完整 MiniMind 007 用例验证结果如下：

• vocab_size=1001；
• tie_word_embeddings=False；
• mesh_shape=(2, 4) 且 mesh_dim_names=("replicate", "shard")；
• embedding replicate，lm_head uneven shard；
• 连续训练 10 step，Single 与 FSDP 每一步 loss 完全一致；
• Max Diff = 0.0，进程正常退出。

因此当前问题优先判定为测试脚本问题，不是 FSDP uneven-shard/replicate 的数值正确性问题。生产代码可另行补充两类 fail-fast：unnamed 2D mesh 进入 replicate 路径时给出明确错误，以及 tied Parameter 跨嵌套 FSDP unit/冲突策略时提前拒绝；具体兼容策略建议单独评审后再改。