model/model.py

# Copyright (c) Microsoft Corporation. Licensed under the MIT license.
# Copyright Lightning AI. Licensed under the Apache License 2.0.
# Modified for LatentLoop: MoE architectures, loop visits, and latent sequence updates.
"""MoE Base, MoE Loop, and LatentLoop with separate caches for each layer visit."""
import math
from functools import lru_cache
from typing import Optional, Tuple

import torch
from torch import Tensor, nn
from torch.nn import functional as F

from model.config import Config
from model.latent import ScaledDecoder, latent_correction
from model.fused_rotary_embedding import apply_rotary_emb_func
from model.moe import SparseMoE, load_balancing_loss

RoPECache = Tuple[Tensor, Tensor]
KVCache = Tuple[Tensor, Tensor]


@lru_cache(maxsize=1)
def flash_attention():
    try:
        from flash_attn import flash_attn_func
        return flash_attn_func
    except (ImportError, OSError):
        return None


class GPT(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
            h=nn.ModuleList(Block(config) for _ in range(config.n_layer)),
            ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
        ))
        self.layer_order = config.layer_order
        self.coarse_execution_span = config.coarse_execution_span
        self.register_buffer("rope_cos", None, persistent=False)
        self.register_buffer("rope_sin", None, persistent=False)
        self.kv_caches = []
        self._cache_length = self._cache_capacity = 0
        self._coarse_pending = self._coarse_update = None
        self.apply(self._init_weights)
        self.coarse_decoder = ScaledDecoder(config.n_embd) if config.architecture == "latentloop" else None
        if config.architecture == "latentloop":
            # One physical owner for the shared MoE and both attention blocks.
            for i in range(config.loop_start, config.loop_end):
                block = self.transformer.h[i]
                block.norm_1_repeat = config.norm_class(config.n_embd, eps=config.norm_eps)
                block.norm_2_repeat = config.norm_class(config.n_embd, eps=config.norm_eps)
                block.attn_repeat = CausalSelfAttention(config)
                block.attn_repeat.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=self.config.initializer_range)
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)

    def reset_cache(self):
        self.kv_caches = []
        self._cache_length = self._cache_capacity = 0
        self._coarse_pending = self._coarse_update = None

    def _apply(self, fn, recurse=True):
        self.reset_cache()
        self.rope_cos = self.rope_sin = None
        return super()._apply(fn, recurse=recurse)
    def forward(self, idx: Tensor, max_seq_length: Optional[int] = None,
                input_pos: Optional[Tensor] = None, *, return_aux_loss: bool = False):
        if idx.ndim != 2 or idx.size(1) == 0:
            raise ValueError("idx must have shape [batch, nonempty sequence]")
        batch, length = idx.shape
        capacity = self.config.block_size if max_seq_length is None else max_seq_length
        if not 0 < capacity <= self.config.block_size or length > capacity:
            raise ValueError("sequence/cache length must be within config.block_size")
        mask = None
        if input_pos is not None:
            if torch.is_grad_enabled():
                raise RuntimeError("KV caching is inference-only; use torch.no_grad() or inference_mode()")
            if input_pos.ndim != 1 or input_pos.numel() != length or input_pos.dtype != torch.long:
                raise ValueError("input_pos must be a 1-D int64 tensor with one position per token")
            input_pos = input_pos.to(idx.device)
            expected = torch.arange(self._cache_length, self._cache_length + length, device=idx.device)
            if not torch.equal(input_pos, expected) or self._cache_length + length > capacity:
                raise ValueError("KV cache positions must be contiguous from zero and fit the cache; reset_cache() first")
            if not self.kv_caches:
                # Shared weights do not imply shared activations: each visit needs its own KV.
                self.kv_caches = [None] * len(self.layer_order)
                self._cache_capacity = capacity
            elif capacity != self._cache_capacity:
                raise ValueError("cache capacity changed; call reset_cache() first")
            mask = torch.arange(capacity, device=idx.device)[None, :] <= input_pos[:, None]
        if self.config.rotary_dim:
            rope_dtype = (torch.get_autocast_dtype(idx.device.type) if torch.is_autocast_enabled(idx.device.type)
                          else self.transformer.wte.weight.dtype)
            if (self.rope_cos is None or self.rope_cos.device != idx.device
                    or self.rope_cos.size(0) < capacity or self.rope_cos.dtype != rope_dtype):
                self.rope_cos, self.rope_sin = build_rope_cache(
                    capacity, self.config.rotary_dim, rope_dtype, idx.device,
                    self.config.rope_base, 1,
                )
            rope = (self.rope_cos[:length], self.rope_sin[:length]) if input_pos is None else (
                self.rope_cos.index_select(0, input_pos), self.rope_sin.index_select(0, input_pos)
            )
        else:
            rope = None
        x = self.transformer.wte(idx)
        router_stats = []
        visits = [0] * self.config.n_layer
        for i, layer_index in enumerate(self.layer_order):
            first_visit = visits[layer_index] == 0
            visits[layer_index] += 1
            if self.coarse_execution_span is not None and (
                    self.coarse_execution_span[0] <= i < self.coarse_execution_span[1]):
                if i == self.coarse_execution_span[0]:
                    x, stats = self._coarse_pass(x, input_pos, capacity, return_aux_loss)
                    router_stats.extend(stats)
                continue
            block = self.transformer.h[layer_index]
            cache = None if input_pos is None else self.kv_caches[i]
            if return_aux_loss:
                x, cache, stats = block(x, rope, mask, input_pos, cache, capacity,
                                       return_router_stats=True, first_visit=first_visit)
                router_stats.append(stats)  # Average aux over all executions, including repeat visits.
            else:
                x, cache = block(x, rope, mask, input_pos, cache, capacity, first_visit=first_visit)
            if input_pos is not None:
                self.kv_caches[i] = cache
        if input_pos is not None:
            self._cache_length += length
        logits = self.lm_head(self.transformer.ln_f(x))
        if return_aux_loss:
            aux = load_balancing_loss(router_stats) if router_stats else logits.new_zeros((), dtype=torch.float32)
            return logits, aux
        return logits

    def _coarse_pass(self, x, input_pos, capacity, return_aux_loss):
        cached = input_pos is not None
        coarse_capacity = capacity // 2
        stats = []

        def run_core(z, positions):
            ends = positions * 2 + 1
            rope = (self.rope_cos.index_select(0, ends), self.rope_sin.index_select(0, ends))
            mask = (torch.arange(coarse_capacity, device=x.device)[None, :] <= positions[:, None]
                    if cached else None)
            for execution in range(*self.coarse_execution_span):
                block = self.transformer.h[self.layer_order[execution]]
                cache = self.kv_caches[execution] if cached else None
                args = (z, rope, mask, positions if cached else None, cache, coarse_capacity)
                if return_aux_loss:
                    z, cache, entry = block(*args, return_router_stats=True, first_visit=True)
                    stats.append(entry)
                else:
                    z, cache = block(*args, first_visit=True)
                if cached:
                    self.kv_caches[execution] = cache
            return z

        result, pending, update = latent_correction(
            x, run_core, self.coarse_decoder,
            offset=self._cache_length if cached else 0,
            pending=self._coarse_pending if cached else None,
            last_update=self._coarse_update if cached else None,
        )
        if cached:
            self._coarse_pending, self._coarse_update = pending, update
        return result, stats


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps)
        self.attn = CausalSelfAttention(config)
        self.norm_2 = config.norm_class(config.n_embd, eps=config.norm_eps)
        self.mlp = SparseMoE(config)
        self.attn_repeat = self.norm_1_repeat = self.norm_2_repeat = None

    def forward(self, x, rope, mask=None, input_pos=None, kv_cache=None, max_seq_length=None,
                *, return_router_stats=False, first_visit=False):
        repeated = self.attn_repeat is not None and not first_visit
        attention = self.attn_repeat if repeated else self.attn
        norm_1 = self.norm_1_repeat if repeated else self.norm_1
        norm_2 = self.norm_2_repeat if repeated else self.norm_2
        h, kv_cache = attention(norm_1(x), rope, mask, input_pos, kv_cache, max_seq_length)
        x = torch.add(x, h, alpha=1.0)
        if return_router_stats:
            h, stats = self.mlp(norm_2(x), return_router_stats=True)
            return torch.add(x, h, alpha=1.0), kv_cache, stats
        return torch.add(x, self.mlp(norm_2(x)), alpha=1.0), kv_cache


class CausalSelfAttention(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.attn = nn.Linear(config.n_embd, (config.n_head + 2 * config.n_query_groups) * config.head_size,
                              bias=False)
        self.proj = nn.Linear(config.n_head * config.head_size, config.n_embd, bias=False)
        self.q_norm = config.norm_class(config.head_size, eps=config.norm_eps)
        self.k_norm = config.norm_class(config.head_size, eps=config.norm_eps)

    def forward(self, x, rope, mask=None, input_pos=None, kv_cache=None, max_seq_length=None):
        batch, length, width = x.shape
        groups, heads, dim = self.config.n_query_groups, self.config.n_head, self.config.head_size
        projected = self.attn(x)
        # Checkpoint layout: [KV group, queries-per-group + key + value, head_dim].
        qkv = projected.view(batch, length, groups, heads // groups + 2, dim)
        q, k, v = qkv.split((heads // groups, 1, 1), dim=-2)
        q = q.reshape(batch, length, heads, dim)
        k, v = k.reshape(batch, length, groups, dim), v.reshape(batch, length, groups, dim)
        q, k = self.q_norm(q), self.k_norm(k)
        if rope is not None:
            cos, sin = (value.to(q.dtype) for value in rope)
            q = apply_rotary_emb_func(q, cos, sin, inplace=False)
            k = apply_rotary_emb_func(k, cos, sin, inplace=False)
        if input_pos is not None:
            shape = (batch, max_seq_length, groups, dim)
            if kv_cache is None:
                kv_cache = (k.new_zeros(shape), v.new_zeros(shape))
            elif kv_cache[0].shape != shape or kv_cache[0].device != k.device or kv_cache[0].dtype != k.dtype:
                raise ValueError("cache batch/device/dtype changed; call reset_cache() first")
            kv_cache[0].index_copy_(1, input_pos, k)
            kv_cache[1].index_copy_(1, input_pos, v)
            k, v = kv_cache
        y = self.scaled_dot_product_attention(q, k, v, mask, input_pos)
        y = y.reshape(batch, length, heads * dim)
        return self.proj(y), kv_cache

    def scaled_dot_product_attention(self, q, k, v, mask=None, input_pos=None):
        backend = self.config.attention_backend
        flash = flash_attention() if backend != "sdpa" and q.is_cuda else None
        eligible = q.is_cuda and q.dtype in (torch.float16, torch.bfloat16)
        if backend == "flash" and (flash is None or not eligible):
            raise RuntimeError("flash backend requires working flash-attn and CUDA fp16/bf16")
        # Cached decoding uses an explicit absolute-position mask; training takes the fast path.
        if flash is not None and eligible and mask is None:
            return flash(q, k, v, causal=True, softmax_scale=1 / math.sqrt(q.size(-1)))
        q, k, v = (value.transpose(1, 2) for value in (q, k, v))
        return F.scaled_dot_product_attention(
            q, k, v, attn_mask=mask, is_causal=mask is None, enable_gqa=q.size(1) != k.size(1),
        ).transpose(1, 2)


def build_rope_cache(seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device,
                     base: float = 10000, condense_ratio: float = 1) -> RoPECache:
    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device, dtype=torch.float32) / n_elem))
    positions = torch.arange(seq_len, device=device, dtype=torch.float32) / condense_ratio
    angles = torch.outer(positions, theta)
    return angles.cos().to(dtype), angles.sin().to(dtype)