model/arg_layer.py · Phase-Technologies/FORGE-500M-Chat at main

Add model architecture component: arg_layer.py

ca0b93a verified 4 days ago

13.5 kB

"""
ARG (Adaptive Recurrent Gating) Layer — FORGE's core sequence mixer.
Combines:
1. Complex-eigenvalue selective SSM (recurrent branch)  
2. Local windowed GQA (local attention branch)
3. Learned scalar gate α_t ∈ [0,1] that blends the two
4. Compound Positional Bias (CPB) for SSM initialization
"""

from __future__ import annotations
import math
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

from .dgn_norm import build_norm
from .rotary_embedding import RotaryEmbedding
from .triton_kernels import fused_swiglu
from .triton_ssm import fused_triton_ssm



class ARGLayer(nn.Module):
    """
    Adaptive Recurrent Gating Layer.
    
    For each token t:
        h_recur_t = ComplexSSM(x_{0:t})
        h_local_t = WindowedGQA(x_{t-W:t})
        α_t       = σ(w_g · x_t)
        output_t  = α_t * h_local_t + (1 - α_t) * h_recur_t
    """
    
    def __init__(
        self,
        d_model: int = 1280,
        d_inner: int = 1280,
        d_state: int = 48,
        d_rank: int = 48,
        conv_kernel: int = 4,
        local_window: int = 64,
        local_n_heads: int = 8,
        local_n_kv_heads: int = 2,
        head_dim: int = 80,
        norm_type: str = "dgn",
        dgn_n_groups: int = 16,
        norm_eps: float = 1e-6,
        use_flash_attention: bool = True,
        rope_base: float = 500_000.0,
        layer_idx: int = 0,
        max_seq_len: int = 4096,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_inner = d_inner
        self.d_state = d_state
        self.d_rank = d_rank
        self.conv_kernel = conv_kernel
        self.local_window = local_window
        self.local_n_heads = local_n_heads
        self.local_n_kv_heads = local_n_kv_heads
        self.head_dim = head_dim
        self.use_flash_attention = use_flash_attention
        self.layer_idx = layer_idx
        
        # ── Pre-norm ──────────────────────────────────────────────────────────
        self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
        
        # ── Recurrent Branch ──────────────────────────────────────────────────
        # Input projection: x → (x_inner, z) both d_inner
        self.in_proj = nn.Linear(d_model, 2 * d_inner, bias=False)
        
        # Depthwise causal conv for local mixing
        self.conv1d = nn.Conv1d(
            d_inner, d_inner,
            kernel_size=conv_kernel,
            padding=conv_kernel - 1,   # causal padding
            groups=d_inner,
            bias=True,
        )
        
        # Input-dependent SSM parameters (low-rank for Δ)
        self.x_proj = nn.Linear(d_inner, d_rank + 2 * d_state, bias=False)
        self.dt_proj = nn.Linear(d_rank, d_inner, bias=True)
        
        # Complex SSM eigenvalue parameters
        # Uses real-pair encoding: each complex eigenvalue λ = -exp(ν) ± jθ
        # is represented as two real state channels, so d_state must be even.
        # theta (oscillation) is implicitly learned through B/C projections.
        assert d_state % 2 == 0, "d_state must be even for complex-pair encoding"
        self.d_state_complex = d_state // 2

        self.nu = nn.Parameter(torch.zeros(self.d_state_complex))    # decay (log scale)
        
        # Skip connection coefficient D
        self.D = nn.Parameter(torch.ones(d_inner))
        
        # Output projection
        self.out_proj = nn.Linear(d_inner, d_model, bias=False)
        
        # ── Local Attention Branch ────────────────────────────────────────────
        q_dim = local_n_heads * head_dim
        kv_dim = local_n_kv_heads * head_dim
        
        self.local_q = nn.Linear(d_model, q_dim, bias=False)
        self.local_k = nn.Linear(d_model, kv_dim, bias=False)
        self.local_v = nn.Linear(d_model, kv_dim, bias=False)
        self.local_o = nn.Linear(q_dim, d_model, bias=False)
        
        # Local RoPE (initialize with safe maximum length of 8192 to avoid dynamic cache expansion during forward pass)
        self.local_rope = RotaryEmbedding(
            head_dim=head_dim,
            max_seq_len=max_seq_len * 2,
            base=rope_base,
        )
        
        # ── Adaptive Gate ─────────────────────────────────────────────────────
        self.gate_proj = nn.Linear(d_model, 1, bias=True)
        nn.init.zeros_(self.gate_proj.weight)
        nn.init.zeros_(self.gate_proj.bias)  # α=0.5 at init: sigmoid(0) = 0.5
        
        # Initialize weights
        self._init_weights()
    
    def _init_weights(self):
        """Careful initialization for stable training."""
        # dt_proj bias: positive to encourage non-trivial step sizes at init
        dt_init_std = self.d_rank ** -0.5
        nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
        dt_bias = torch.exp(
            torch.rand(self.d_inner) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
        )
        with torch.no_grad():
            self.dt_proj.bias.copy_(dt_bias + torch.log(-torch.expm1(-dt_bias)))
        
        # nu init — zero → moderate decay at start
        nn.init.zeros_(self.nu)
        
        # Conv1d — identity-like init
        nn.init.zeros_(self.conv1d.weight)
        for i in range(self.d_inner):
            self.conv1d.weight.data[i, 0, -1] = 1.0

    # ─────────────────────────────────────────────────────────────────────────
    # COMPLEX EIGENVALUE → REAL A MATRIX
    # ─────────────────────────────────────────────────────────────────────────

    def _build_real_A(self) -> torch.Tensor:
        """
        Build the (d_inner, d_state) real-valued A matrix for selective_scan_fn.
        Each complex eigenvalue λ_i = -exp(ν_i) + jθ_i is encoded as two real
        state channels with the same decay rate. The theta (oscillation) enters
        via the B/C projections. mamba_ssm's selective_scan_fn discretizes A
        internally via exp(A * Δ), so we pass the continuous-time log-decay
        (negative real, per state channel).
        """
        decay = -F.softplus(self.nu.float())   # fp32 — mamba_ssm requires A in fp32
        A_diag = decay.repeat_interleave(2)   # (d_state,)
        A = A_diag.unsqueeze(0).expand(self.d_inner, -1).contiguous().float()  # (d_inner, d_state), fp32
        return A
    
    # ─────────────────────────────────────────────────────────────────────────
    # RECURRENT BRANCH
    # ─────────────────────────────────────────────────────────────────────────
    
    @torch._dynamo.disable
    def recurrent_branch(
        self,
        x: torch.Tensor,         # (B, T, d_model)
        position_offset: int = 0,
    ) -> torch.Tensor:
        """
        Selective SSM recurrent branch using mamba_ssm's fused CUDA kernel.
        Returns (B, T, d_model).
        Uses selective_scan_fn (same kernel as Mamba) which:
        - Processes all channels in parallel with an efficient associative scan
        - Applies delta_softplus internally (no manual softplus on dt)
        - Fuses the z-gate SiLU multiplication inside the kernel
        - Eliminates 512× Python-level kernel launches per layer
        """
        B, T, _ = x.shape

        # ── Project inputs ────────────────────────────────────────────────────
        xz = self.in_proj(x)                        # (B, T, 2*d_inner)
        x_inner, z = xz.chunk(2, dim=-1)             # each (B, T, d_inner)

        # ── Causal conv for local mixing ──────────────────────────────────────
        x_conv = rearrange(x_inner, 'b t d -> b d t')
        x_conv = self.conv1d(x_conv)[:, :, :T]       # trim causal padding, still (B, d_inner, T)
        x_conv = F.silu(x_conv)                      # (B, d_inner, T)
        x_conv_t = rearrange(x_conv, 'b d t -> b t d')

        # ── Input-dependent SSM parameters ────────────────────────────────────
        x_dbl = self.x_proj(x_conv_t)                # (B, T, d_rank + 2*d_state)
        dt_raw, B_ssm, C_ssm = x_dbl.split([self.d_rank, self.d_state, self.d_state], dim=-1)

        # dt: per-channel step size — (B, T, d_inner) → (B, d_inner, T)
        dt = F.softplus(self.dt_proj(dt_raw))
        dt = rearrange(dt, 'b t d -> b d t').contiguous()

        # B, C: (B, d_state, T)
        B_ssm = rearrange(B_ssm, 'b t n -> b n t').contiguous()
        C_ssm = rearrange(C_ssm, 'b t n -> b n t').contiguous()

        # ── Build A from complex eigenvalue params ────────────────────────────
        A = self._build_real_A()                     # (d_inner, d_state), negative real

        # ── Pure Triton Fused Selective SSM Scan ──────────────────────────────
        z_t = rearrange(z, 'b t d -> b d t').contiguous()
        y = fused_triton_ssm(x_conv, dt, A, B_ssm, C_ssm, self.D, z_t)
        y = rearrange(y, 'b d t -> b t d')

        return self.out_proj(y.to(dtype=x.dtype))    # (B, T, d_model)
    
    # ─────────────────────────────────────────────────────────────────────────
    # LOCAL ATTENTION BRANCH
    # ─────────────────────────────────────────────────────────────────────────
    
    @torch._dynamo.disable
    def local_attention_branch(
        self,
        x: torch.Tensor,   # (B, T, d_model)
    ) -> torch.Tensor:
        """
        Windowed GQA with window size W.
        Uses native PyTorch SDPA with sliding window mask.
        Returns (B, T, d_model).
        """
        B, T, _ = x.shape
        W = self.local_window
        
        Q = self.local_q(x)  # (B, T, n_heads * head_dim)
        K = self.local_k(x)  # (B, T, n_kv_heads * head_dim)
        V = self.local_v(x)
        
        Q = rearrange(Q, 'b t (h d) -> b h t d', d=self.head_dim)  # (B, H, T, D)
        K = rearrange(K, 'b t (h d) -> b h t d', d=self.head_dim)
        V = rearrange(V, 'b t (h d) -> b h t d', d=self.head_dim)
        
        # Apply local RoPE
        Q, K = self.local_rope.apply_rotary(Q, K)
        
        # GQA: expand KV to match Q heads
        n_rep = self.local_n_heads // self.local_n_kv_heads
        if n_rep > 1:
            K = K.repeat_interleave(n_rep, dim=1)  # (B, H, T, D)
            V = V.repeat_interleave(n_rep, dim=1)
        
        # Fast PyTorch SDPA with sliding window mask
        indices = torch.arange(T, device=x.device)
        mask_2d = (indices.unsqueeze(0) <= indices.unsqueeze(1)) & \
                  (indices.unsqueeze(0) >= (indices.unsqueeze(1) - W))
        mask_4d = mask_2d.unsqueeze(0).unsqueeze(0)
        attn_out = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask_4d)
        out = rearrange(attn_out, 'b h t d -> b t (h d)')
        
        return self.local_o(out)  # (B, T, d_model)
    
    # ─────────────────────────────────────────────────────────────────────────
    # FORWARD PASS
    # ─────────────────────────────────────────────────────────────────────────
    
    def forward(
        self,
        x: torch.Tensor,
        position_offset: int = 0,
    ) -> torch.Tensor:
        """
        x: (B, T, d_model)
        Returns: (B, T, d_model) — residual connection applied inside
        """
        residual = x
        x_norm = self.norm(x)
        
        # Two branches
        h_recur = self.recurrent_branch(x_norm, position_offset)
        h_local = self.local_attention_branch(x_norm)
        
        # Learned gate
        alpha = torch.sigmoid(self.gate_proj(x_norm))  # (B, T, 1)
        
        # Blend and add residual
        mixed = alpha * h_local + (1.0 - alpha) * h_recur  # (B, T, d_model)
        
        return residual + mixed

	"""
	ARG (Adaptive Recurrent Gating) Layer — FORGE's core sequence mixer.

	Combines:
	1. Complex-eigenvalue selective SSM (recurrent branch)
	2. Local windowed GQA (local attention branch)
	3. Learned scalar gate α_t ∈ [0,1] that blends the two
	4. Compound Positional Bias (CPB) for SSM initialization
	"""

	from __future__ import annotations
	import math
	from typing import Optional, Tuple

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from einops import rearrange

	from .dgn_norm import build_norm
	from .rotary_embedding import RotaryEmbedding
	from .triton_kernels import fused_swiglu
	from .triton_ssm import fused_triton_ssm



	class ARGLayer(nn.Module):
	"""
	Adaptive Recurrent Gating Layer.

	For each token t:
	h_recur_t = ComplexSSM(x_{0:t})
	h_local_t = WindowedGQA(x_{t-W:t})
	α_t = σ(w_g · x_t)
	output_t = α_t * h_local_t + (1 - α_t) * h_recur_t
	"""

	def __init__(
	self,
	d_model: int = 1280,
	d_inner: int = 1280,
	d_state: int = 48,
	d_rank: int = 48,
	conv_kernel: int = 4,
	local_window: int = 64,
	local_n_heads: int = 8,
	local_n_kv_heads: int = 2,
	head_dim: int = 80,
	norm_type: str = "dgn",
	dgn_n_groups: int = 16,
	norm_eps: float = 1e-6,
	use_flash_attention: bool = True,
	rope_base: float = 500_000.0,
	layer_idx: int = 0,
	max_seq_len: int = 4096,
	):
	super().__init__()
	self.d_model = d_model
	self.d_inner = d_inner
	self.d_state = d_state
	self.d_rank = d_rank
	self.conv_kernel = conv_kernel
	self.local_window = local_window
	self.local_n_heads = local_n_heads
	self.local_n_kv_heads = local_n_kv_heads
	self.head_dim = head_dim
	self.use_flash_attention = use_flash_attention
	self.layer_idx = layer_idx

	# ── Pre-norm ──────────────────────────────────────────────────────────
	self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)

	# ── Recurrent Branch ──────────────────────────────────────────────────
	# Input projection: x → (x_inner, z) both d_inner
	self.in_proj = nn.Linear(d_model, 2 * d_inner, bias=False)

	# Depthwise causal conv for local mixing
	self.conv1d = nn.Conv1d(
	d_inner, d_inner,
	kernel_size=conv_kernel,
	padding=conv_kernel - 1, # causal padding
	groups=d_inner,
	bias=True,
	)

	# Input-dependent SSM parameters (low-rank for Δ)
	self.x_proj = nn.Linear(d_inner, d_rank + 2 * d_state, bias=False)
	self.dt_proj = nn.Linear(d_rank, d_inner, bias=True)

	# Complex SSM eigenvalue parameters
	# Uses real-pair encoding: each complex eigenvalue λ = -exp(ν) ± jθ
	# is represented as two real state channels, so d_state must be even.
	# theta (oscillation) is implicitly learned through B/C projections.
	assert d_state % 2 == 0, "d_state must be even for complex-pair encoding"
	self.d_state_complex = d_state // 2

	self.nu = nn.Parameter(torch.zeros(self.d_state_complex)) # decay (log scale)

	# Skip connection coefficient D
	self.D = nn.Parameter(torch.ones(d_inner))

	# Output projection
	self.out_proj = nn.Linear(d_inner, d_model, bias=False)

	# ── Local Attention Branch ────────────────────────────────────────────
	q_dim = local_n_heads * head_dim
	kv_dim = local_n_kv_heads * head_dim

	self.local_q = nn.Linear(d_model, q_dim, bias=False)
	self.local_k = nn.Linear(d_model, kv_dim, bias=False)
	self.local_v = nn.Linear(d_model, kv_dim, bias=False)
	self.local_o = nn.Linear(q_dim, d_model, bias=False)

	# Local RoPE (initialize with safe maximum length of 8192 to avoid dynamic cache expansion during forward pass)
	self.local_rope = RotaryEmbedding(
	head_dim=head_dim,
	max_seq_len=max_seq_len * 2,
	base=rope_base,
	)

	# ── Adaptive Gate ─────────────────────────────────────────────────────
	self.gate_proj = nn.Linear(d_model, 1, bias=True)
	nn.init.zeros_(self.gate_proj.weight)
	nn.init.zeros_(self.gate_proj.bias) # α=0.5 at init: sigmoid(0) = 0.5

	# Initialize weights
	self._init_weights()

	def _init_weights(self):
	"""Careful initialization for stable training."""
	# dt_proj bias: positive to encourage non-trivial step sizes at init
	dt_init_std = self.d_rank ** -0.5
	nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
	dt_bias = torch.exp(
	torch.rand(self.d_inner) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
	)
	with torch.no_grad():
	self.dt_proj.bias.copy_(dt_bias + torch.log(-torch.expm1(-dt_bias)))

	# nu init — zero → moderate decay at start
	nn.init.zeros_(self.nu)

	# Conv1d — identity-like init
	nn.init.zeros_(self.conv1d.weight)
	for i in range(self.d_inner):
	self.conv1d.weight.data[i, 0, -1] = 1.0

	# ─────────────────────────────────────────────────────────────────────────
	# COMPLEX EIGENVALUE → REAL A MATRIX
	# ─────────────────────────────────────────────────────────────────────────

	def _build_real_A(self) -> torch.Tensor:
	"""
	Build the (d_inner, d_state) real-valued A matrix for selective_scan_fn.

	Each complex eigenvalue λ_i = -exp(ν_i) + jθ_i is encoded as two real
	state channels with the same decay rate. The theta (oscillation) enters
	via the B/C projections. mamba_ssm's selective_scan_fn discretizes A
	internally via exp(A * Δ), so we pass the continuous-time log-decay
	(negative real, per state channel).
	"""
	decay = -F.softplus(self.nu.float()) # fp32 — mamba_ssm requires A in fp32
	A_diag = decay.repeat_interleave(2) # (d_state,)
	A = A_diag.unsqueeze(0).expand(self.d_inner, -1).contiguous().float() # (d_inner, d_state), fp32
	return A

	# ─────────────────────────────────────────────────────────────────────────
	# RECURRENT BRANCH
	# ─────────────────────────────────────────────────────────────────────────

	@torch._dynamo.disable
	def recurrent_branch(
	self,
	x: torch.Tensor, # (B, T, d_model)
	position_offset: int = 0,
	) -> torch.Tensor:
	"""
	Selective SSM recurrent branch using mamba_ssm's fused CUDA kernel.
	Returns (B, T, d_model).

	Uses selective_scan_fn (same kernel as Mamba) which:
	- Processes all channels in parallel with an efficient associative scan
	- Applies delta_softplus internally (no manual softplus on dt)
	- Fuses the z-gate SiLU multiplication inside the kernel
	- Eliminates 512× Python-level kernel launches per layer
	"""
	B, T, _ = x.shape

	# ── Project inputs ────────────────────────────────────────────────────
	xz = self.in_proj(x) # (B, T, 2*d_inner)
	x_inner, z = xz.chunk(2, dim=-1) # each (B, T, d_inner)

	# ── Causal conv for local mixing ──────────────────────────────────────
	x_conv = rearrange(x_inner, 'b t d -> b d t')
	x_conv = self.conv1d(x_conv)[:, :, :T] # trim causal padding, still (B, d_inner, T)
	x_conv = F.silu(x_conv) # (B, d_inner, T)
	x_conv_t = rearrange(x_conv, 'b d t -> b t d')

	# ── Input-dependent SSM parameters ────────────────────────────────────
	x_dbl = self.x_proj(x_conv_t) # (B, T, d_rank + 2*d_state)
	dt_raw, B_ssm, C_ssm = x_dbl.split([self.d_rank, self.d_state, self.d_state], dim=-1)

	# dt: per-channel step size — (B, T, d_inner) → (B, d_inner, T)
	dt = F.softplus(self.dt_proj(dt_raw))
	dt = rearrange(dt, 'b t d -> b d t').contiguous()

	# B, C: (B, d_state, T)
	B_ssm = rearrange(B_ssm, 'b t n -> b n t').contiguous()
	C_ssm = rearrange(C_ssm, 'b t n -> b n t').contiguous()

	# ── Build A from complex eigenvalue params ────────────────────────────
	A = self._build_real_A() # (d_inner, d_state), negative real

	# ── Pure Triton Fused Selective SSM Scan ──────────────────────────────
	z_t = rearrange(z, 'b t d -> b d t').contiguous()
	y = fused_triton_ssm(x_conv, dt, A, B_ssm, C_ssm, self.D, z_t)
	y = rearrange(y, 'b d t -> b t d')

	return self.out_proj(y.to(dtype=x.dtype)) # (B, T, d_model)

	# ─────────────────────────────────────────────────────────────────────────
	# LOCAL ATTENTION BRANCH
	# ─────────────────────────────────────────────────────────────────────────

	@torch._dynamo.disable
	def local_attention_branch(
	self,
	x: torch.Tensor, # (B, T, d_model)
	) -> torch.Tensor:
	"""
	Windowed GQA with window size W.
	Uses native PyTorch SDPA with sliding window mask.
	Returns (B, T, d_model).
	"""
	B, T, _ = x.shape
	W = self.local_window

	Q = self.local_q(x) # (B, T, n_heads * head_dim)
	K = self.local_k(x) # (B, T, n_kv_heads * head_dim)
	V = self.local_v(x)

	Q = rearrange(Q, 'b t (h d) -> b h t d', d=self.head_dim) # (B, H, T, D)
	K = rearrange(K, 'b t (h d) -> b h t d', d=self.head_dim)
	V = rearrange(V, 'b t (h d) -> b h t d', d=self.head_dim)

	# Apply local RoPE
	Q, K = self.local_rope.apply_rotary(Q, K)

	# GQA: expand KV to match Q heads
	n_rep = self.local_n_heads // self.local_n_kv_heads
	if n_rep > 1:
	K = K.repeat_interleave(n_rep, dim=1) # (B, H, T, D)
	V = V.repeat_interleave(n_rep, dim=1)

	# Fast PyTorch SDPA with sliding window mask
	indices = torch.arange(T, device=x.device)
	mask_2d = (indices.unsqueeze(0) <= indices.unsqueeze(1)) & \
	(indices.unsqueeze(0) >= (indices.unsqueeze(1) - W))
	mask_4d = mask_2d.unsqueeze(0).unsqueeze(0)
	attn_out = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask_4d)
	out = rearrange(attn_out, 'b h t d -> b t (h d)')

	return self.local_o(out) # (B, T, d_model)

	# ─────────────────────────────────────────────────────────────────────────
	# FORWARD PASS
	# ─────────────────────────────────────────────────────────────────────────

	def forward(
	self,
	x: torch.Tensor,
	position_offset: int = 0,
	) -> torch.Tensor:
	"""
	x: (B, T, d_model)
	Returns: (B, T, d_model) — residual connection applied inside
	"""
	residual = x
	x_norm = self.norm(x)

	# Two branches
	h_recur = self.recurrent_branch(x_norm, position_offset)
	h_local = self.local_attention_branch(x_norm)

	# Learned gate
	alpha = torch.sigmoid(self.gate_proj(x_norm)) # (B, T, 1)

	# Blend and add residual
	mixed = alpha * h_local + (1.0 - alpha) * h_recur # (B, T, d_model)

	return residual + mixed