Document original Turing model loader and MoE implementation

0389190 verified about 1 month ago

9.76 kB

"""Causal transformer; RoPE conventions adapted from our A1, other layers rebuilt."""

from types import SimpleNamespace
import torch
from torch import nn
import torch.nn.functional as F
from src.part1.ffn import FeedForward


# Multi-head causal attention with rotary positional embeddings (RoPE).
class Attention(nn.Module):
    def __init__(self, width, heads):
        """Create projections and RoPE frequencies for equal-width attention heads."""
        super().__init__()
        # Split the model width evenly across attention heads.
        self.heads = heads
        self.head_dim = width // heads
        # q, k, v and o are query, key, value and output projections.
        # Keep these attribute names: saved Turing checkpoints use them.
        self.q = nn.Linear(width, width, bias=False)
        self.k = nn.Linear(width, width, bias=False)
        self.v = nn.Linear(width, width, bias=False)
        self.o = nn.Linear(width, width, bias=False)
        # A buffer follows the model device but is not optimized.
        # persistent=False omits these reconstructible frequencies from checkpoints.
        self.register_buffer(
            "inverse_frequency",
            10000 ** (-torch.arange(0, self.head_dim, 2).float() / self.head_dim),
            persistent=False,
        )

    def rotate(self, x, positions):
        """Rotate adjacent feature pairs using each token's absolute position."""
        # x has shape [batch, heads, time, head_dim].
        # Broadcast positions against frequencies to obtain one angle per pair.
        angles = (
            positions.float()[:, None, :, None]
            * self.inverse_frequency[None, None, None, :]
        )
        # RoPE rotates adjacent even/odd feature pairs without changing their norm.
        even, odd = x[..., 0::2], x[..., 1::2]
        cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
        # stack interleaves the rotated pair; flatten restores head_dim features.
        return torch.stack(
            (even * cos - odd * sin, even * sin + odd * cos), -1
        ).flatten(-2)

    def forward(self, x, mask, past=None, use_cache=False):
        """Return attention output and optionally the extended key/value cache."""
        # Short shape names: b=batch size, t=new token count, d=model width.
        b, t, d = x.shape
        # Cached keys already cover previous tokens during incremental decoding.
        previous = past[0].shape[-2] if past is not None else 0
        # Count visible tokens, ignoring left padding, to get zero-based positions.
        positions = mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -t:]

        # Project [b,t,d], separate heads, then move heads before the time axis.
        def split(module):
            """Reshape one projected tensor into [batch,heads,time,head_dim]."""
            return module(x).reshape(b, t, self.heads, self.head_dim).transpose(1, 2)

        # Apply position rotations to queries and keys; values stay unrotated.
        q, k, v = (
            self.rotate(split(self.q), positions),
            self.rotate(split(self.k), positions),
            split(self.v),
        )
        # Append only the new keys/values to the cached sequence (axis -2).
        if past is not None:
            k, v = torch.cat((past[0], k), -2), torch.cat((past[1], v), -2)
        s = k.shape[-2]
        # Unpadded training/prefill can use PyTorch SDPA
        # (scaled dot-product attention) with its built-in causal mask.
        if previous == 0 and bool(mask.all()):
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            # With padding or a cache, build the full [new_time,total_time] mask.
            # A query can see earlier keys, but never a future key or a PAD position.
            causal = (
                torch.arange(s, device=x.device)[None, :]
                <= (torch.arange(t, device=x.device) + previous)[:, None]
            )
            visible = causal[None, None, :, :] & mask[:, None, None, :].bool()
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=visible)
        # Merge the heads back into [batch,time,width] before the output projection.
        merged = out.transpose(1, 2).reshape(b, t, d)
        return self.o(merged), (k, v) if use_cache else None


# One pre-norm transformer block: attention residual, then FFN residual.
class Block(nn.Module):
    def __init__(self, config):
        """Construct attention, two pre-normalizations and the chosen FFN variant."""
        super().__init__()
        d = config["d_model"]
        # RMSNorm normalizes root mean square magnitude before each sublayer.
        self.norm1, self.norm2 = nn.RMSNorm(d), nn.RMSNorm(d)
        self.attention = Attention(d, config["heads"])
        self.ffn = FeedForward(
            d, config["hidden"], config["ffn"], config["aux_coefficient"]
        )

    def forward(self, x, mask, past=None, use_cache=False, collect=False):
        """Return updated hidden states, router loss, attention cache and usage map."""
        # Pre-norm means normalize first, then add the sublayer output to x.
        attended, cache = self.attention(self.norm1(x), mask, past, use_cache)
        x = x + attended
        # The FFN receives only the mask of current tokens, not cached tokens.
        fed, aux, usage = self.ffn(self.norm2(x), mask[:, -x.shape[1] :], collect)
        return x + fed, aux, cache, usage


# Shared decoder backbone for all Part 1 variants and Part 2 optimizers.
class Transformer(nn.Module):
    def __init__(self, config):
        """Build the shared decoder and tie output weights to input embeddings."""
        super().__init__()
        self.config = dict(config)
        self.embedding = nn.Embedding(config["vocab_size"], config["d_model"])
        self.blocks = nn.ModuleList([Block(config) for _ in range(config["layers"])])
        self.norm = nn.RMSNorm(config["d_model"])
        self.head = nn.Linear(config["d_model"], config["vocab_size"], bias=False)
        # Tie input embeddings and output classifier weights to save parameters.
        self.head.weight = self.embedding.weight
        # Recursively initialize modules in the original Turing order.
        self.apply(self.initialize)

    @staticmethod
    def initialize(module):
        """Use the original normal initialization for each linear/embedding module."""
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=0.02)

    def forward(
        self,
        input_ids,
        attention_mask=None,
        past_key_values=None,
        use_cache=False,
        collect_routing=False,
        **kwargs,
    ):
        """Return logits, mean router loss, optional caches and expert usage."""
        b, t = input_ids.shape
        past = past_key_values
        prior = past[0][0].shape[-2] if past is not None else 0
        # A missing mask means every input/cached position is visible.
        if attention_mask is None:
            attention_mask = torch.ones(
                b, prior + t, device=input_ids.device, dtype=torch.bool
            )
        x = self.embedding(input_ids)
        # new_zeros creates a scalar on the same device; router losses use FP32.
        aux = x.new_zeros((), dtype=torch.float32)
        caches, routing = [], []
        # Run layers in order and keep one key/value cache and usage map per layer.
        for i, block in enumerate(self.blocks):
            x, loss, cache, usage = block(
                x,
                attention_mask,
                past[i] if past is not None else None,
                use_cache,
                collect_routing,
            )
            aux = aux + loss
            if use_cache:
                caches.append(cache)
            if usage is not None:
                routing.append(usage)
        # Expose named outputs like a Hugging Face forward result.
        # The auxiliary loss is averaged over layers; logits are [b,t,vocab].
        return SimpleNamespace(
            logits=self.head(self.norm(x)),
            aux_loss=aux / max(len(self.blocks), 1),
            past_key_values=tuple(caches) if use_cache else None,
            routing=routing,
        )

    def parameter_counts(self):
        """Count all weights and the minimum/maximum active weights per token."""
        # parameters() counts tied embedding/head weights only once.
        total = sum(p.numel() for p in self.parameters())
        active_min = active_max = total
        if self.config["ffn"] != "dense":
            # Always-active parameters exclude expert/router FFNs; add back each
            # router plus the minimum/maximum selected expert sizes per token.
            shared = total - sum(
                p.numel() for b in self.blocks for p in b.ffn.parameters()
            )
            minima, maxima = [], []
            d = self.config["d_model"]
            for b in self.blocks:
                widths = b.ffn.widths
                if b.ffn.kind == "shared":
                    lo, hi = min(widths[:3]) + widths[3], max(widths[:3]) + widths[3]
                elif b.ffn.kind == "top1":
                    lo, hi = min(widths), max(widths)
                else:
                    lo, hi = sum(sorted(widths)[:2]), sum(sorted(widths)[-2:])
                # Each hidden expert unit has width input and width output weights.
                # The 4*width router weights are active regardless of selected experts.
                minima.append(2 * d * lo + 4 * d)
                maxima.append(2 * d * hi + 4 * d)
            active_min, active_max = shared + sum(minima), shared + sum(maxima)
        return {"total": total, "active_min": active_min, "active_max": active_max}

	"""Causal transformer; RoPE conventions adapted from our A1, other layers rebuilt."""

	from types import SimpleNamespace
	import torch
	from torch import nn
	import torch.nn.functional as F
	from src.part1.ffn import FeedForward


	# Multi-head causal attention with rotary positional embeddings (RoPE).
	class Attention(nn.Module):
	def __init__(self, width, heads):
	"""Create projections and RoPE frequencies for equal-width attention heads."""
	super().__init__()
	# Split the model width evenly across attention heads.
	self.heads = heads
	self.head_dim = width // heads
	# q, k, v and o are query, key, value and output projections.
	# Keep these attribute names: saved Turing checkpoints use them.
	self.q = nn.Linear(width, width, bias=False)
	self.k = nn.Linear(width, width, bias=False)
	self.v = nn.Linear(width, width, bias=False)
	self.o = nn.Linear(width, width, bias=False)
	# A buffer follows the model device but is not optimized.
	# persistent=False omits these reconstructible frequencies from checkpoints.
	self.register_buffer(
	"inverse_frequency",
	10000 ** (-torch.arange(0, self.head_dim, 2).float() / self.head_dim),
	persistent=False,
	)

	def rotate(self, x, positions):
	"""Rotate adjacent feature pairs using each token's absolute position."""
	# x has shape [batch, heads, time, head_dim].
	# Broadcast positions against frequencies to obtain one angle per pair.
	angles = (
	positions.float()[:, None, :, None]
	* self.inverse_frequency[None, None, None, :]
	)
	# RoPE rotates adjacent even/odd feature pairs without changing their norm.
	even, odd = x[..., 0::2], x[..., 1::2]
	cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
	# stack interleaves the rotated pair; flatten restores head_dim features.
	return torch.stack(
	(even * cos - odd * sin, even * sin + odd * cos), -1
	).flatten(-2)

	def forward(self, x, mask, past=None, use_cache=False):
	"""Return attention output and optionally the extended key/value cache."""
	# Short shape names: b=batch size, t=new token count, d=model width.
	b, t, d = x.shape
	# Cached keys already cover previous tokens during incremental decoding.
	previous = past[0].shape[-2] if past is not None else 0
	# Count visible tokens, ignoring left padding, to get zero-based positions.
	positions = mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -t:]

	# Project [b,t,d], separate heads, then move heads before the time axis.
	def split(module):
	"""Reshape one projected tensor into [batch,heads,time,head_dim]."""
	return module(x).reshape(b, t, self.heads, self.head_dim).transpose(1, 2)

	# Apply position rotations to queries and keys; values stay unrotated.
	q, k, v = (
	self.rotate(split(self.q), positions),
	self.rotate(split(self.k), positions),
	split(self.v),
	)
	# Append only the new keys/values to the cached sequence (axis -2).
	if past is not None:
	k, v = torch.cat((past[0], k), -2), torch.cat((past[1], v), -2)
	s = k.shape[-2]
	# Unpadded training/prefill can use PyTorch SDPA
	# (scaled dot-product attention) with its built-in causal mask.
	if previous == 0 and bool(mask.all()):
	out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
	else:
	# With padding or a cache, build the full [new_time,total_time] mask.
	# A query can see earlier keys, but never a future key or a PAD position.
	causal = (
	torch.arange(s, device=x.device)[None, :]
	<= (torch.arange(t, device=x.device) + previous)[:, None]
	)
	visible = causal[None, None, :, :] & mask[:, None, None, :].bool()
	out = F.scaled_dot_product_attention(q, k, v, attn_mask=visible)
	# Merge the heads back into [batch,time,width] before the output projection.
	merged = out.transpose(1, 2).reshape(b, t, d)
	return self.o(merged), (k, v) if use_cache else None


	# One pre-norm transformer block: attention residual, then FFN residual.
	class Block(nn.Module):
	def __init__(self, config):
	"""Construct attention, two pre-normalizations and the chosen FFN variant."""
	super().__init__()
	d = config["d_model"]
	# RMSNorm normalizes root mean square magnitude before each sublayer.
	self.norm1, self.norm2 = nn.RMSNorm(d), nn.RMSNorm(d)
	self.attention = Attention(d, config["heads"])
	self.ffn = FeedForward(
	d, config["hidden"], config["ffn"], config["aux_coefficient"]
	)

	def forward(self, x, mask, past=None, use_cache=False, collect=False):
	"""Return updated hidden states, router loss, attention cache and usage map."""
	# Pre-norm means normalize first, then add the sublayer output to x.
	attended, cache = self.attention(self.norm1(x), mask, past, use_cache)
	x = x + attended
	# The FFN receives only the mask of current tokens, not cached tokens.
	fed, aux, usage = self.ffn(self.norm2(x), mask[:, -x.shape[1] :], collect)
	return x + fed, aux, cache, usage


	# Shared decoder backbone for all Part 1 variants and Part 2 optimizers.
	class Transformer(nn.Module):
	def __init__(self, config):
	"""Build the shared decoder and tie output weights to input embeddings."""
	super().__init__()
	self.config = dict(config)
	self.embedding = nn.Embedding(config["vocab_size"], config["d_model"])
	self.blocks = nn.ModuleList([Block(config) for _ in range(config["layers"])])
	self.norm = nn.RMSNorm(config["d_model"])
	self.head = nn.Linear(config["d_model"], config["vocab_size"], bias=False)
	# Tie input embeddings and output classifier weights to save parameters.
	self.head.weight = self.embedding.weight
	# Recursively initialize modules in the original Turing order.
	self.apply(self.initialize)

	@staticmethod
	def initialize(module):
	"""Use the original normal initialization for each linear/embedding module."""
	if isinstance(module, (nn.Linear, nn.Embedding)):
	nn.init.normal_(module.weight, std=0.02)

	def forward(
	self,
	input_ids,
	attention_mask=None,
	past_key_values=None,
	use_cache=False,
	collect_routing=False,
	**kwargs,
	):
	"""Return logits, mean router loss, optional caches and expert usage."""
	b, t = input_ids.shape
	past = past_key_values
	prior = past[0][0].shape[-2] if past is not None else 0
	# A missing mask means every input/cached position is visible.
	if attention_mask is None:
	attention_mask = torch.ones(
	b, prior + t, device=input_ids.device, dtype=torch.bool
	)
	x = self.embedding(input_ids)
	# new_zeros creates a scalar on the same device; router losses use FP32.
	aux = x.new_zeros((), dtype=torch.float32)
	caches, routing = [], []
	# Run layers in order and keep one key/value cache and usage map per layer.
	for i, block in enumerate(self.blocks):
	x, loss, cache, usage = block(
	x,
	attention_mask,
	past[i] if past is not None else None,
	use_cache,
	collect_routing,
	)
	aux = aux + loss
	if use_cache:
	caches.append(cache)
	if usage is not None:
	routing.append(usage)
	# Expose named outputs like a Hugging Face forward result.
	# The auxiliary loss is averaged over layers; logits are [b,t,vocab].
	return SimpleNamespace(
	logits=self.head(self.norm(x)),
	aux_loss=aux / max(len(self.blocks), 1),
	past_key_values=tuple(caches) if use_cache else None,
	routing=routing,
	)

	def parameter_counts(self):
	"""Count all weights and the minimum/maximum active weights per token."""
	# parameters() counts tied embedding/head weights only once.
	total = sum(p.numel() for p in self.parameters())
	active_min = active_max = total
	if self.config["ffn"] != "dense":
	# Always-active parameters exclude expert/router FFNs; add back each
	# router plus the minimum/maximum selected expert sizes per token.
	shared = total - sum(
	p.numel() for b in self.blocks for p in b.ffn.parameters()
	)
	minima, maxima = [], []
	d = self.config["d_model"]
	for b in self.blocks:
	widths = b.ffn.widths
	if b.ffn.kind == "shared":
	lo, hi = min(widths[:3]) + widths[3], max(widths[:3]) + widths[3]
	elif b.ffn.kind == "top1":
	lo, hi = min(widths), max(widths)
	else:
	lo, hi = sum(sorted(widths)[:2]), sum(sorted(widths)[-2:])
	# Each hidden expert unit has width input and width output weights.
	# The 4*width router weights are active regardless of selected experts.
	minima.append(2 * d * lo + 4 * d)
	maxima.append(2 * d * hi + 4 * d)
	active_min, active_max = shared + sum(minima), shared + sum(maxima)
	return {"total": total, "active_min": active_min, "active_max": active_max}