模型层：从 Embedding 到 Loss 的实现

nanotron 的 Llama 实现在 src/nanotron/models/llama.py，约 800 行。本章把它拆开看。

9.1核心类

9.2Embedding（带 TP）

class Embedding(nn.Module):
    def __init__(self, vocab_size, hidden_size, tp_pg):
        super().__init__()
        self.tp_pg = tp_pg
        # 把 vocab 沿 tp_pg 切
        local_vocab_size = vocab_size // tp_pg.size()
        self.weight = nn.Parameter(torch.empty(local_vocab_size, hidden_size))
        self.vocab_start_idx = tp_pg.rank() * local_vocab_size
        self.vocab_end_idx = self.vocab_start_idx + local_vocab_size

    def forward(self, input_ids):
        # mask 本 rank 不负责的 ids
        mask = (input_ids < self.vocab_start_idx) | (input_ids >= self.vocab_end_idx)
        local_ids = input_ids.clamp(self.vocab_start_idx, self.vocab_end_idx - 1) - self.vocab_start_idx
        out = F.embedding(local_ids, self.weight)
        out[mask] = 0
        # all-reduce 合并各 rank 结果
        out = dist.all_reduce(out, op=ReduceOp.SUM, group=self.tp_pg)
        return out

9.3CausalSelfAttention（GQA + RoPE）

class CausalSelfAttention(nn.Module):
    def __init__(self, config, tp_pg):
        # QKV 用 ColumnParallel，O 用 RowParallel
        self.qkv_proj = TensorParallelColumnLinear(
            config.hidden_size,
            config.hidden_size + 2 * config.num_kv_heads * config.head_dim,
            pg=tp_pg, mode=tp_mode,
        )
        self.o_proj = TensorParallelRowLinear(
            config.hidden_size, config.hidden_size,
            pg=tp_pg, mode=tp_mode,
        )

    def forward(self, hidden_states, position_ids):
        qkv = self.qkv_proj(hidden_states)
        q, k, v = split_qkv(qkv)
        q, k = apply_rope(q, k, position_ids)
        # GQA: K/V 头数少，repeat 到 Q 数
        k = repeat_kv(k, num_heads // num_kv_heads)
        v = repeat_kv(v, num_heads // num_kv_heads)
        # Flash attention（如果可用）
        out = flash_attn_func(q, k, v, causal=True)
        return self.o_proj(out)

关键：QKV fuse 到一个 ColumnParallelLinear，省一次 all-reduce + 提高 GEMM 效率。

9.4MLP（SwiGLU）

class MLP(nn.Module):
    def __init__(self, config, tp_pg):
        # gate + up fuse 到一个 ColumnParallel
        self.gate_up_proj = TensorParallelColumnLinear(
            config.hidden_size, 2 * config.intermediate_size, pg=tp_pg
        )
        self.down_proj = TensorParallelRowLinear(
            config.intermediate_size, config.hidden_size, pg=tp_pg
        )

    def forward(self, x):
        gate_up = self.gate_up_proj(x)
        gate, up = gate_up.chunk(2, dim=-1)
        x = F.silu(gate) * up
        return self.down_proj(x)

9.5DecoderLayer = Attention + MLP + RMSNorm

class LlamaDecoderLayer(nn.Module):
    def __init__(self, config, tp_pg):
        self.input_layernorm = TritonRMSNorm(config.hidden_size)
        self.self_attn = CausalSelfAttention(config, tp_pg)
        self.post_attention_layernorm = TritonRMSNorm(config.hidden_size)
        self.mlp = MLP(config, tp_pg)

    def forward(self, hidden_states, position_ids):
        residual = hidden_states
        h = self.input_layernorm(hidden_states)
        h = self.self_attn(h, position_ids)
        h = residual + h

        residual = h
        h = self.post_attention_layernorm(h)
        h = self.mlp(h)
        return residual + h

9.6LlamaForTraining（加 loss）

class LlamaForTraining(NanotronModel):
    def __init__(self, config, parallel_context):
        self.model = LlamaModel(config, parallel_context)
        # tied with embedding
        self.lm_head = TensorParallelColumnLinear(
            config.hidden_size, config.vocab_size, tp_pg=tp_pg
        )

    def forward(self, input_ids, labels):
        h = self.model(input_ids)
        logits = self.lm_head(h)
        # TP-aware cross entropy
        loss = sharded_cross_entropy(logits, labels, group=self.tp_pg)
        return loss

9.7这章你需要带走的

• nanotron Llama 实现 ~800 行；
• QKV 和 gate+up 都 fuse 到一个 ColumnParallel；
• O 和 down 是 RowParallel；
• cross entropy 是 TP-aware 的 sharded 版本；
• 对比 transformers Llama 4000+ 行的实现，nanotron 是极简学习材料。