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

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

training_code/model/__init__.py

@@ -0,0 +1,2 @@
+from .config import ModelConfig, TrainConfig
+from .transformer import Transformer

training_code/model/config.py

@@ -0,0 +1,78 @@
+"""
+Configuration for 1B parameter LLaMA-style Transformer model.
+Architecture: Decoder-only Transformer with RoPE, GQA, SwiGLU, RMSNorm.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 32000
+    hidden_dim: int = 2048
+    intermediate_dim: int = 5504       # ~2.7x hidden for SwiGLU (adjusted for param count)
+    num_layers: int = 22
+    num_attention_heads: int = 32
+    num_kv_heads: int = 8              # GQA: 4 query heads per KV head
+    max_seq_len: int = 2048
+    rope_theta: float = 10000.0
+    rms_norm_eps: float = 1e-5
+    dropout: float = 0.0               # No dropout (modern practice for pretraining)
+    tie_word_embeddings: bool = False
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_dim // self.num_attention_heads
+
+    @property
+    def num_params_approx(self) -> int:
+        """Rough parameter count estimate."""
+        embed = self.vocab_size * self.hidden_dim
+        attn_per_layer = (
+            self.hidden_dim * self.head_dim * self.num_attention_heads +  # Q
+            self.hidden_dim * self.head_dim * self.num_kv_heads +         # K
+            self.hidden_dim * self.head_dim * self.num_kv_heads +         # V
+            self.head_dim * self.num_attention_heads * self.hidden_dim    # O
+        )
+        ffn_per_layer = 3 * self.hidden_dim * self.intermediate_dim      # gate + up + down
+        norm_per_layer = 2 * self.hidden_dim
+        total = (
+            embed +
+            self.num_layers * (attn_per_layer + ffn_per_layer + norm_per_layer) +
+            self.hidden_dim +  # final norm
+            (0 if self.tie_word_embeddings else self.vocab_size * self.hidden_dim)
+        )
+        return total
+
+
+@dataclass
+class TrainConfig:
+    # Paths
+    checkpoint_dir: str = "/jfs/deepak-kumar/checkpoints"
+    data_cache_dir: str = "/jfs/deepak-kumar/data"
+    log_dir: str = "/home/jovyan/training/logs"
+
+    # Training
+    total_tokens: int = 20_000_000_000   # 20B tokens
+    batch_size_per_gpu: int = 8
+    gradient_accumulation_steps: int = 8  # effective batch = 8 * 8 * 8 = 512 seqs
+    max_seq_len: int = 2048
+    
+    # WSD Schedule
+    learning_rate: float = 3e-4
+    min_lr: float = 3e-5
+    warmup_steps: int = 1000
+    weight_decay: float = 0.1
+    beta1: float = 0.9
+    beta2: float = 0.95
+    grad_clip: float = 1.0
+
+    # Logging
+    log_interval: int = 10
+    save_interval: int = 1000
+    eval_interval: int = 500
+
+    # System
+    num_workers: int = 4
+    seed: int = 42
+    bf16: bool = True

training_code/model/data.py

@@ -0,0 +1,79 @@
+"""
+Data pipeline: streams and tokenizes OpenWebText for pretraining.
+Packs sequences to max_seq_len for efficiency (no padding waste).
+"""
+
+import os
+import torch
+from torch.utils.data import IterableDataset, DataLoader
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+def get_tokenizer(name: str = "mistralai/Mistral-7B-v0.1"):
+    """Use Mistral's tokenizer — 32k vocab, BPE, well-trained on diverse data."""
+    tok = AutoTokenizer.from_pretrained(name, use_fast=True)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    return tok
+
+
+class PackedPretrainDataset(IterableDataset):
+    """
+    Streams text from HuggingFace dataset, tokenizes on the fly,
+    and packs into fixed-length sequences for maximum GPU utilization.
+    """
+
+    def __init__(self, tokenizer, max_seq_len: int, split: str = "train", cache_dir: str = None, seed: int = 42):
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+        self.split = split
+        self.cache_dir = cache_dir
+        self.seed = seed
+        self.eos_id = tokenizer.eos_token_id
+
+    def _token_stream(self):
+        ds = load_dataset(
+            "HuggingFaceFW/fineweb-edu",
+            name="sample-10BT",
+            split=self.split,
+            streaming=True,
+            cache_dir=self.cache_dir,
+        )
+        ds = ds.shuffle(seed=self.seed, buffer_size=10_000)
+
+        for example in ds:
+            text = example.get("text", "")
+            if len(text.strip()) < 50:
+                continue
+            token_ids = self.tokenizer.encode(text, add_special_tokens=False)
+            yield from token_ids
+            yield self.eos_id
+
+    def __iter__(self):
+        buffer = []
+        for token_id in self._token_stream():
+            buffer.append(token_id)
+            if len(buffer) == self.max_seq_len + 1:
+                input_ids = torch.tensor(buffer[:-1], dtype=torch.long)
+                labels = torch.tensor(buffer[1:], dtype=torch.long)
+                yield input_ids, labels
+                buffer = []
+
+
+def create_dataloader(tokenizer, config, rank: int = 0, world_size: int = 1, seed_override: int = None):
+    seed = seed_override if seed_override is not None else config.seed
+    dataset = PackedPretrainDataset(
+        tokenizer=tokenizer,
+        max_seq_len=config.max_seq_len,
+        split="train",
+        cache_dir=config.data_cache_dir,
+        seed=seed + rank,
+    )
+    return DataLoader(
+        dataset,
+        batch_size=config.batch_size_per_gpu,
+        num_workers=config.num_workers,
+        pin_memory=True,
+        prefetch_factor=4,
+    )

training_code/model/dpo_data.py

@@ -0,0 +1,144 @@
+"""
+DPO data pipeline: loads UltraFeedback preference pairs.
+
+Each example has a prompt + chosen response + rejected response.
+We tokenize both (prompt+chosen) and (prompt+rejected), apply the same
+chat template, and return them as pairs for DPO training.
+"""
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset
+
+
+CHAT_TEMPLATE = {
+    "user_start": "<|user|>\n",
+    "assistant_start": "<|assistant|>\n",
+    "turn_end": "\n<|end|>\n",
+}
+
+
+def format_preference_pair(prompt, chosen_msgs, rejected_msgs):
+    """Build chat-templated strings for chosen and rejected."""
+    def build(messages):
+        text = CHAT_TEMPLATE["user_start"] + prompt.strip() + CHAT_TEMPLATE["turn_end"]
+        for msg in messages:
+            role = msg.get("role", "assistant")
+            content = msg.get("content", "").strip()
+            if role == "assistant":
+                text += CHAT_TEMPLATE["assistant_start"] + content + CHAT_TEMPLATE["turn_end"]
+            elif role == "user":
+                text += CHAT_TEMPLATE["user_start"] + content + CHAT_TEMPLATE["turn_end"]
+        return text
+
+    return build(chosen_msgs), build(rejected_msgs)
+
+
+class DPODataset(Dataset):
+    """
+    Loads UltraFeedback preference pairs and tokenizes them.
+    Returns (prompt_ids, chosen_ids, rejected_ids) with proper shifting.
+    """
+
+    def __init__(self, tokenizer, max_seq_len=2048, split="train",
+                 cache_dir=None, max_samples=None):
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+
+        special_tokens = ["<|user|>", "<|assistant|>", "<|end|>"]
+        vocab = tokenizer.get_vocab()
+        new_tokens = [t for t in special_tokens if t not in vocab]
+        if new_tokens:
+            tokenizer.add_tokens(new_tokens, special_tokens=True)
+
+        self.assistant_token_id = tokenizer.encode("<|assistant|>", add_special_tokens=False)[0]
+        self.end_token_id = tokenizer.encode("<|end|>", add_special_tokens=False)[0]
+        self.user_token_id = tokenizer.encode("<|user|>", add_special_tokens=False)[0]
+
+        print(f"[DPO Data] Loading UltraFeedback preferences ({split})...")
+        ds = load_dataset(
+            "argilla/ultrafeedback-binarized-preferences-cleaned",
+            split=split,
+            cache_dir=cache_dir,
+        )
+        if max_samples:
+            ds = ds.select(range(min(max_samples, len(ds))))
+        print(f"[DPO Data] {len(ds)} preference pairs loaded")
+
+        self.examples = []
+        skipped = 0
+        for i, row in enumerate(ds):
+            prompt = row.get("prompt", "")
+            chosen = row.get("chosen", [])
+            rejected = row.get("rejected", [])
+
+            if not prompt or not chosen or not rejected:
+                skipped += 1
+                continue
+
+            chosen_text, rejected_text = format_preference_pair(prompt, chosen, rejected)
+
+            chosen_ids = tokenizer.encode(chosen_text, add_special_tokens=False)
+            rejected_ids = tokenizer.encode(rejected_text, add_special_tokens=False)
+
+            # Truncate if needed
+            if len(chosen_ids) > max_seq_len + 1:
+                chosen_ids = chosen_ids[:max_seq_len + 1]
+            if len(rejected_ids) > max_seq_len + 1:
+                rejected_ids = rejected_ids[:max_seq_len + 1]
+
+            if len(chosen_ids) < 10 or len(rejected_ids) < 10:
+                skipped += 1
+                continue
+
+            # Find where the prompt ends (first <|assistant|> token)
+            prompt_end = 0
+            for j, tid in enumerate(chosen_ids):
+                if tid == self.assistant_token_id:
+                    prompt_end = j + 2  # skip <|assistant|> and \n
+                    break
+
+            self.examples.append({
+                "chosen_ids": chosen_ids,
+                "rejected_ids": rejected_ids,
+                "prompt_len": prompt_end,
+            })
+
+            if (i + 1) % 20000 == 0:
+                print(f"  Processed {i+1} pairs...")
+
+        print(f"[DPO Data] {len(self.examples)} pairs ready, {skipped} skipped")
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        ex = self.examples[idx]
+        return {
+            "chosen_ids": torch.tensor(ex["chosen_ids"], dtype=torch.long),
+            "rejected_ids": torch.tensor(ex["rejected_ids"], dtype=torch.long),
+            "prompt_len": ex["prompt_len"],
+        }
+
+
+def dpo_collate_fn(batch, pad_id=0):
+    """Pad chosen and rejected sequences separately."""
+    max_chosen = max(b["chosen_ids"].size(0) for b in batch)
+    max_rejected = max(b["rejected_ids"].size(0) for b in batch)
+
+    chosen_padded = []
+    rejected_padded = []
+    prompt_lens = []
+
+    for b in batch:
+        c_pad = max_chosen - b["chosen_ids"].size(0)
+        r_pad = max_rejected - b["rejected_ids"].size(0)
+        chosen_padded.append(torch.cat([b["chosen_ids"], torch.full((c_pad,), pad_id, dtype=torch.long)]))
+        rejected_padded.append(torch.cat([b["rejected_ids"], torch.full((r_pad,), pad_id, dtype=torch.long)]))
+        prompt_lens.append(b["prompt_len"])
+
+    return {
+        "chosen_ids": torch.stack(chosen_padded),
+        "rejected_ids": torch.stack(rejected_padded),
+        "prompt_lens": torch.tensor(prompt_lens, dtype=torch.long),
+    }

training_code/model/sft_data.py

@@ -0,0 +1,169 @@
+"""
+SFT data pipeline: loads UltraChat 200K and formats into chat template.
+
+Chat template:
+  <|user|>
+  What is gravity?
+  <|end|>
+  <|assistant|>
+  Gravity is a fundamental force...
+  <|end|>
+
+Labels are shifted left by 1 (standard causal LM), with user turns masked.
+"""
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset
+
+
+CHAT_TEMPLATE = {
+    "user_start": "<|user|>\n",
+    "assistant_start": "<|assistant|>\n",
+    "turn_end": "\n<|end|>\n",
+}
+
+
+def format_conversation(messages):
+    """Convert a list of {role, content} messages into our chat template string."""
+    text = ""
+    for msg in messages:
+        role = msg["role"]
+        content = msg["content"].strip()
+        if role == "user":
+            text += CHAT_TEMPLATE["user_start"] + content + CHAT_TEMPLATE["turn_end"]
+        elif role == "assistant":
+            text += CHAT_TEMPLATE["assistant_start"] + content + CHAT_TEMPLATE["turn_end"]
+    return text
+
+
+class SFTDataset(Dataset):
+    """
+    Loads UltraChat 200K conversations, tokenizes them, builds shifted labels
+    with user turns masked so the model only learns to generate assistant responses.
+    """
+
+    def __init__(self, tokenizer, max_seq_len=2048, split="train_sft", cache_dir=None, max_samples=None):
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+
+        special_tokens = ["<|user|>", "<|assistant|>", "<|end|>"]
+        vocab = tokenizer.get_vocab()
+        new_tokens = [t for t in special_tokens if t not in vocab]
+        if new_tokens:
+            tokenizer.add_tokens(new_tokens, special_tokens=True)
+
+        self.assistant_token_id = tokenizer.encode("<|assistant|>", add_special_tokens=False)[0]
+        self.end_token_id = tokenizer.encode("<|end|>", add_special_tokens=False)[0]
+        self.user_token_id = tokenizer.encode("<|user|>", add_special_tokens=False)[0]
+
+        print(f"[SFT Data] Loading UltraChat 200K ({split})...")
+        ds = load_dataset("HuggingFaceH4/ultrachat_200k", split=split, cache_dir=cache_dir)
+        if max_samples:
+            ds = ds.select(range(min(max_samples, len(ds))))
+        print(f"[SFT Data] {len(ds)} conversations loaded")
+
+        self.examples = []
+        skipped = 0
+        for i, row in enumerate(ds):
+            messages = row["messages"]
+            if len(messages) < 2:
+                skipped += 1
+                continue
+
+            text = format_conversation(messages)
+            all_ids = tokenizer.encode(text, add_special_tokens=False)
+
+            # Need at least max_seq_len+1 for shift, but truncate if longer
+            if len(all_ids) > max_seq_len + 1:
+                all_ids = all_ids[:max_seq_len + 1]
+
+            if len(all_ids) < 10:
+                skipped += 1
+                continue
+
+            # Shifted: input = all_ids[:-1], target = all_ids[1:]
+            input_ids = all_ids[:-1]
+            target_ids = all_ids[1:]
+
+            # Build mask: -100 for user turns, real token id for assistant turns
+            labels = self._build_shifted_labels(input_ids, target_ids)
+            self.examples.append((input_ids, labels))
+
+            if (i + 1) % 50000 == 0:
+                print(f"  Processed {i+1} conversations...")
+
+        print(f"[SFT Data] {len(self.examples)} examples ready, {skipped} skipped")
+
+    def _build_shifted_labels(self, input_ids, target_ids):
+        """
+        Walk through the token sequence and track whether we're in a user turn
+        or assistant turn. Only keep labels for assistant response content.
+
+        Masking strategy (applied to the SHIFTED target):
+        - Everything before and including <|assistant|>\\n: masked
+        - Assistant response content and <|end|>: TRAIN
+        - <|user|> and user content until next <|assistant|>: masked
+        """
+        labels = [-100] * len(target_ids)
+        in_assistant = False
+
+        for i, tid in enumerate(input_ids):
+            if tid == self.assistant_token_id:
+                # Next token after <|assistant|> is \n, then content starts
+                in_assistant = True
+                continue
+
+            if tid == self.user_token_id:
+                in_assistant = False
+                continue
+
+            if in_assistant:
+                labels[i] = target_ids[i]
+
+            # When we hit <|end|> in assistant mode, include it then switch off
+            if tid == self.end_token_id and in_assistant:
+                in_assistant = False
+
+        return labels
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        input_ids, labels = self.examples[idx]
+        return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long)
+
+
+def sft_collate_fn(batch, pad_id=0):
+    """Pad sequences to the same length within a batch."""
+    input_ids_list, labels_list = zip(*batch)
+    max_len = max(ids.size(0) for ids in input_ids_list)
+
+    padded_inputs = []
+    padded_labels = []
+    for ids, lbl in zip(input_ids_list, labels_list):
+        pad_len = max_len - ids.size(0)
+        padded_inputs.append(torch.cat([ids, torch.full((pad_len,), pad_id, dtype=torch.long)]))
+        padded_labels.append(torch.cat([lbl, torch.full((pad_len,), -100, dtype=torch.long)]))
+
+    return torch.stack(padded_inputs), torch.stack(padded_labels)
+
+
+def create_sft_dataloader(tokenizer, batch_size=4, max_seq_len=2048,
+                          cache_dir=None, max_samples=None, num_workers=4):
+    dataset = SFTDataset(
+        tokenizer=tokenizer,
+        max_seq_len=max_seq_len,
+        split="train_sft",
+        cache_dir=cache_dir,
+        max_samples=max_samples,
+    )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True,
+        collate_fn=lambda b: sft_collate_fn(b, pad_id=tokenizer.pad_token_id),
+    ), dataset

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