初始化项目，由ModelHub XC社区提供模型

Model: divakar-yadav/transformer-1b-chat
Source: Original Platform

training_code/model/transformer.py

@@ -0,0 +1,163 @@
+"""
+1B Parameter Decoder-Only Transformer — built from scratch.
+
+Techniques:
+  - RoPE (Rotary Position Embeddings)
+  - Grouped Query Attention (GQA)
+  - SwiGLU Feed-Forward
+  - RMSNorm (pre-norm architecture)
+  - Flash Attention 2 (via PyTorch SDPA)
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .config import ModelConfig
+
+
+class RMSNorm(nn.Module):
+    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:
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * norm).type_as(x) * self.weight
+
+
+def precompute_rope_freqs(dim: int, max_seq_len: int, theta: float = 10000.0) -> torch.Tensor:
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)  # complex64
+
+
+def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
+    B, S, H, D = xq.shape
+    xq_c = torch.view_as_complex(xq.float().reshape(B, S, H, D // 2, 2))
+    xk_c = torch.view_as_complex(xk.float().reshape(B, S, xk.shape[2], D // 2, 2))
+    freqs = freqs_cis[:S].clone().unsqueeze(0).unsqueeze(2)
+    xq_out = torch.view_as_real(xq_c * freqs).flatten(3)
+    xk_out = torch.view_as_real(xk_c * freqs).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class GroupedQueryAttention(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_kv_heads
+        self.head_dim = config.head_dim
+        self.num_groups = self.num_heads // self.num_kv_heads
+
+        self.wq = nn.Linear(config.hidden_dim, self.num_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(config.hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(config.hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(self.num_heads * self.head_dim, config.hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+        B, S, _ = x.shape
+
+        q = self.wq(x).view(B, S, self.num_heads, self.head_dim)
+        k = self.wk(x).view(B, S, self.num_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, S, self.num_kv_heads, self.head_dim)
+
+        q, k = apply_rope(q, k, freqs_cis)
+
+        # Expand KV heads for GQA
+        if self.num_groups > 1:
+            k = k.unsqueeze(3).expand(B, S, self.num_kv_heads, self.num_groups, self.head_dim)
+            k = k.reshape(B, S, self.num_heads, self.head_dim)
+            v = v.unsqueeze(3).expand(B, S, self.num_kv_heads, self.num_groups, self.head_dim)
+            v = v.reshape(B, S, self.num_heads, self.head_dim)
+
+        # (B, num_heads, S, head_dim) for SDPA
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        out = out.transpose(1, 2).contiguous().view(B, S, -1)
+        return self.wo(out)
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.w_gate = nn.Linear(config.hidden_dim, config.intermediate_dim, bias=False)
+        self.w_up = nn.Linear(config.hidden_dim, config.intermediate_dim, bias=False)
+        self.w_down = nn.Linear(config.intermediate_dim, config.hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.attention_norm = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
+        self.attention = GroupedQueryAttention(config)
+        self.ffn_norm = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
+        self.ffn = SwiGLUFFN(config)
+
+    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+        x = x + self.attention(self.attention_norm(x), freqs_cis)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_dim)
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])
+        self.norm = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
+        self.output = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+
+        # Pre-compute RoPE frequencies
+        self.register_buffer(
+            "freqs_cis",
+            precompute_rope_freqs(config.head_dim, config.max_seq_len * 2, config.rope_theta),
+            persistent=False,
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize with scaled normal, following GPT-NeoX / LLaMA conventions."""
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+        # Scale residual projections by 1/sqrt(2*num_layers)
+        scale = (2 * self.config.num_layers) ** -0.5
+        for layer in self.layers:
+            nn.init.normal_(layer.attention.wo.weight, mean=0.0, std=0.02 * scale)
+            nn.init.normal_(layer.ffn.w_down.weight, mean=0.0, std=0.02 * scale)
+
+    def forward(self, tokens: torch.Tensor, targets: torch.Tensor = None):
+        B, S = tokens.shape
+        h = self.tok_embeddings(tokens)
+
+        freqs_cis = self.freqs_cis[:S]
+        for layer in self.layers:
+            h = layer(h, freqs_cis)
+        h = self.norm(h)
+        logits = self.output(h)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-100,
+            )
+        return logits, loss