Upload mars (ANLP A2 Part 2)

bec4ff4 verified 29 days ago

5.38 kB

"""Decoder-only transformer with a pluggable FFN (dense MLP or MoE variants)."""

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .config import TransformerConfig
from .ffn import MoE, build_ffn


def apply_rope(x: torch.Tensor, positions: torch.Tensor, base: float) -> torch.Tensor:
    # x: (B, H, T, hd), positions: (B, T)
    hd = x.shape[-1]
    inv_freq = 1.0 / (base ** (torch.arange(0, hd, 2, device=x.device).float() / hd))
    ang = positions[:, None, :, None].float() * inv_freq  # (B, 1, T, hd/2)
    cos, sin = ang.cos().to(x.dtype), ang.sin().to(x.dtype)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1).flatten(-2)


class Attention(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        assert config.D_model % config.n_heads == 0
        self.n_heads = config.n_heads
        self.head_dim = config.D_model // config.n_heads
        self.rope_base = config.rope_base
        self.qkv = nn.Linear(config.D_model, 3 * config.D_model, bias=False)
        self.proj = nn.Linear(config.D_model, config.D_model, bias=False)
        self.dropout = config.dropout

    def forward(self, x, positions, attn_mask=None):
        B, T, D = x.shape
        q, k, v = self.qkv(x).view(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k = apply_rope(q, positions, self.rope_base), apply_rope(k, positions, self.rope_base)
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=attn_mask is None,
        )
        return self.proj(y.transpose(1, 2).reshape(B, T, D))


class Block(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.ln1 = nn.RMSNorm(config.D_model)
        self.attn = Attention(config)
        self.ln2 = nn.RMSNorm(config.D_model)
        self.ffn = build_ffn(config)

    def forward(self, x, positions, attn_mask=None, pad_mask=None):
        x = x + self.attn(self.ln1(x), positions, attn_mask)
        return x + self.ffn(self.ln2(x), pad_mask)


class Transformer(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        self.tok_emb = nn.Embedding(config.N_vocab, config.D_model)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
        self.ln_f = nn.RMSNorm(config.D_model)
        self.lm_head = nn.Linear(config.D_model, config.N_vocab, bias=False)
        self.lm_head.weight = self.tok_emb.weight  # weight tying
        self.apply(self._init)
        for name, p in self.named_parameters():  # GPT-2 style residual scaling
            if name.endswith(("attn.proj.weight", "down.weight")):  # incl. experts
                nn.init.normal_(p, std=0.02 / math.sqrt(2 * config.n_layers))

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, ids, pad_mask=None, left_padded=False):
        # Input shape: (B, N_ctx) token ids
        # Output shape: (B, N_ctx, N_vocab)
        # pad_mask: (B, N_ctx) bool, True for real tokens. It is always handed to
        # the FFN (MoE ignores padding in its aux loss / usage stats). With right
        # padding (training/eval) plain causal attention is exact, since real
        # tokens never attend to trailing pads. For left-padded batches
        # (generation) set left_padded=True to mask pad keys and shift positions.
        B, T = ids.shape
        if not left_padded:
            positions = torch.arange(T, device=ids.device).expand(B, T)
            attn_mask = None
        else:
            positions = (pad_mask.long().cumsum(-1) - 1).clamp(min=0)
            causal = torch.ones(T, T, dtype=torch.bool, device=ids.device).tril()
            eye = torch.eye(T, dtype=torch.bool, device=ids.device)
            # `eye` keeps pad queries from having an all-masked row (NaNs)
            attn_mask = (causal & pad_mask[:, None, None, :]) | eye  # (B, 1, T, T)
        x = self.tok_emb(ids)
        for blk in self.blocks:
            x = blk(x, positions, attn_mask, pad_mask)
        return self.lm_head(self.ln_f(x))

    # ---- MoE helpers ----
    def moe_layers(self) -> list[MoE]:
        return [b.ffn for b in self.blocks if isinstance(b.ffn, MoE)]

    def aux_loss(self) -> torch.Tensor:
        layers = self.moe_layers()
        if not layers:
            return torch.zeros((), device=self.tok_emb.weight.device)
        return sum(m.aux_loss for m in layers) / len(layers)

    def set_record_routing(self, flag: bool) -> None:
        for m in self.moe_layers():
            m.record_routing = flag

    def param_counts(self) -> dict[str, int]:
        """total / active parameters (active = per-token, excluding unused experts)."""
        from src.utils import count_params

        total = count_params(self)
        ffn_total = sum(count_params(b.ffn) for b in self.blocks)
        ffn_active = sum(b.ffn.active_params() for b in self.blocks)
        return {"total": total, "active": total - ffn_total + ffn_active,
                "ffn_total": ffn_total, "ffn_active": ffn_active}

	"""Decoder-only transformer with a pluggable FFN (dense MLP or MoE variants)."""

	from __future__ import annotations

	import math

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	from .config import TransformerConfig
	from .ffn import MoE, build_ffn


	def apply_rope(x: torch.Tensor, positions: torch.Tensor, base: float) -> torch.Tensor:
	# x: (B, H, T, hd), positions: (B, T)
	hd = x.shape[-1]
	inv_freq = 1.0 / (base ** (torch.arange(0, hd, 2, device=x.device).float() / hd))
	ang = positions[:, None, :, None].float() * inv_freq # (B, 1, T, hd/2)
	cos, sin = ang.cos().to(x.dtype), ang.sin().to(x.dtype)
	x1, x2 = x[..., 0::2], x[..., 1::2]
	return torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1).flatten(-2)


	class Attention(nn.Module):
	def __init__(self, config: TransformerConfig):
	super().__init__()
	assert config.D_model % config.n_heads == 0
	self.n_heads = config.n_heads
	self.head_dim = config.D_model // config.n_heads
	self.rope_base = config.rope_base
	self.qkv = nn.Linear(config.D_model, 3 * config.D_model, bias=False)
	self.proj = nn.Linear(config.D_model, config.D_model, bias=False)
	self.dropout = config.dropout

	def forward(self, x, positions, attn_mask=None):
	B, T, D = x.shape
	q, k, v = self.qkv(x).view(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
	q, k = apply_rope(q, positions, self.rope_base), apply_rope(k, positions, self.rope_base)
	y = F.scaled_dot_product_attention(
	q, k, v,
	attn_mask=attn_mask,
	dropout_p=self.dropout if self.training else 0.0,
	is_causal=attn_mask is None,
	)
	return self.proj(y.transpose(1, 2).reshape(B, T, D))


	class Block(nn.Module):
	def __init__(self, config: TransformerConfig):
	super().__init__()
	self.ln1 = nn.RMSNorm(config.D_model)
	self.attn = Attention(config)
	self.ln2 = nn.RMSNorm(config.D_model)
	self.ffn = build_ffn(config)

	def forward(self, x, positions, attn_mask=None, pad_mask=None):
	x = x + self.attn(self.ln1(x), positions, attn_mask)
	return x + self.ffn(self.ln2(x), pad_mask)


	class Transformer(nn.Module):
	def __init__(self, config: TransformerConfig):
	super().__init__()
	self.config = config
	self.tok_emb = nn.Embedding(config.N_vocab, config.D_model)
	self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
	self.ln_f = nn.RMSNorm(config.D_model)
	self.lm_head = nn.Linear(config.D_model, config.N_vocab, bias=False)
	self.lm_head.weight = self.tok_emb.weight # weight tying
	self.apply(self._init)
	for name, p in self.named_parameters(): # GPT-2 style residual scaling
	if name.endswith(("attn.proj.weight", "down.weight")): # incl. experts
	nn.init.normal_(p, std=0.02 / math.sqrt(2 * config.n_layers))

	@staticmethod
	def _init(m):
	if isinstance(m, (nn.Linear, nn.Embedding)):
	nn.init.normal_(m.weight, std=0.02)

	def forward(self, ids, pad_mask=None, left_padded=False):
	# Input shape: (B, N_ctx) token ids
	# Output shape: (B, N_ctx, N_vocab)
	# pad_mask: (B, N_ctx) bool, True for real tokens. It is always handed to
	# the FFN (MoE ignores padding in its aux loss / usage stats). With right
	# padding (training/eval) plain causal attention is exact, since real
	# tokens never attend to trailing pads. For left-padded batches
	# (generation) set left_padded=True to mask pad keys and shift positions.
	B, T = ids.shape
	if not left_padded:
	positions = torch.arange(T, device=ids.device).expand(B, T)
	attn_mask = None
	else:
	positions = (pad_mask.long().cumsum(-1) - 1).clamp(min=0)
	causal = torch.ones(T, T, dtype=torch.bool, device=ids.device).tril()
	eye = torch.eye(T, dtype=torch.bool, device=ids.device)
	# `eye` keeps pad queries from having an all-masked row (NaNs)
	attn_mask = (causal & pad_mask[:, None, None, :]) \| eye # (B, 1, T, T)
	x = self.tok_emb(ids)
	for blk in self.blocks:
	x = blk(x, positions, attn_mask, pad_mask)
	return self.lm_head(self.ln_f(x))

	# ---- MoE helpers ----
	def moe_layers(self) -> list[MoE]:
	return [b.ffn for b in self.blocks if isinstance(b.ffn, MoE)]

	def aux_loss(self) -> torch.Tensor:
	layers = self.moe_layers()
	if not layers:
	return torch.zeros((), device=self.tok_emb.weight.device)
	return sum(m.aux_loss for m in layers) / len(layers)

	def set_record_routing(self, flag: bool) -> None:
	for m in self.moe_layers():
	m.record_routing = flag

	def param_counts(self) -> dict[str, int]:
	"""total / active parameters (active = per-token, excluding unused experts)."""
	from src.utils import count_params

	total = count_params(self)
	ffn_total = sum(count_params(b.ffn) for b in self.blocks)
	ffn_active = sum(b.ffn.active_params() for b in self.blocks)
	return {"total": total, "active": total - ffn_total + ffn_active,
	"ffn_total": ffn_total, "ffn_active": ffn_active}