initium-deeplearning 0.1.0__py3-none-any.whl

initium/mamba_controller/mamba3_controller.py

@@ -0,0 +1,242 @@
+"""
+file: mamba3_controller.py -- v1
+Interleaved Mamba-3 (SISO) DNC controller. Same Cell/Block/Wrapper split as
+mamba2_controller.py. Math: paper Prop.1 (exp-trapezoidal) + Prop.4/Eq.11 (RoPE trick):
+  h_t = a_t h_{t-1} + b_t (k_{t-1} x_{t-1}^T) + g_t (k_t x_t^T),  y_t = C_t^T h_t + D x_t
+  a=exp(DT*A), g=lam*DT, b=(1-lam)*DT*a, k/C = RMSNorm+bias then rotated by cumsum(angle*DT).
+State per block: (angle_state, ssm_state, k_prev, v_prev), ALL fp32 (angle is a running sum).
+No conv_state: Mamba-3 has no short conv. Everything after in_proj runs in real fp32.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from initium.mamba_controller.mamba2_controller import Mamba2ControllerWrapper
+from initium.MoE.moe_layer import MoEBlock
+
+_MAMBA3_IMPORT_ERROR: str | None
+
+try:
+    from mamba_ssm.modules.mamba3 import Mamba3
+except ImportError as _e:  # pragma: no cover
+    Mamba3 = None
+    _MAMBA3_IMPORT_ERROR = (
+        "mamba3_controller.py needs mamba-ssm installed from GitHub main "
+        f"(Mamba3 is not in the v2.3.1 wheel). Original error: {_e}"
+    )
+else:
+    _MAMBA3_IMPORT_ERROR = None
+
+
+def _require_mamba3_ssm() -> None:
+    if Mamba3 is None:
+        raise ImportError(_MAMBA3_IMPORT_ERROR)
+
+
+def _heavy_tail(x):  # copied from mamba_ssm.modules.mamba3.heavy_tail_activation
+    return x.clamp_min(0) + torch.reciprocal(1 - x.clamp_max(0))
+
+
+def _rms(v, w, eps):
+    return v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps) * w.float()
+
+
+class Mamba3ControllerCell(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        d_state=64,
+        expand=2,
+        headdim=64,
+        ngroups=1,
+        rope_fraction=0.5,
+        layer_idx=None,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        _require_mamba3_ssm()
+        assert ngroups == 1, "Mamba3ControllerCell supports ngroups==1 only"
+        # parameter container only; we never call .forward()/.step() on it
+        self.mamba3 = Mamba3(
+            d_model=d_model,
+            d_state=d_state,
+            expand=expand,
+            headdim=headdim,
+            ngroups=ngroups,
+            rope_fraction=rope_fraction,
+            is_mimo=False,
+            is_outproj_norm=False,
+            layer_idx=layer_idx,
+            device=device,
+            dtype=dtype,
+        )
+        m = self.mamba3
+        self.d_model, self.d_inner, self.nheads, self.headdim = (
+            d_model,
+            m.d_inner,
+            m.nheads,
+            m.headdim,
+        )
+        self.d_state, self.num_rope_angles = m.d_state, m.num_rope_angles
+
+    def init_state(self, batch_size, device=None, dtype=None):  # dtype ignored: always fp32
+        device = device if device is not None else self.mamba3.in_proj.weight.device
+        H, P, N, S = self.nheads, self.headdim, self.d_state, self.num_rope_angles
+        f = dict(device=device, dtype=torch.float32)
+        return (
+            torch.zeros(batch_size, H, S, **f),
+            torch.zeros(batch_size, H, P, N, **f),
+            torch.zeros(batch_size, H, N, **f),
+            torch.zeros(batch_size, H, P, **f),
+        )
+
+    @staticmethod
+    def _rotate(v, cos, sin):  # v: (b,H,N); rotate first 2S dims, adjacent pairs (SISO convention)
+        S = cos.shape[-1]
+        head, tail = v[..., : 2 * S], v[..., 2 * S :]
+        h0, h1 = head[..., 0::2], head[..., 1::2]
+        rot = torch.stack([h0 * cos - h1 * sin, h0 * sin + h1 * cos], dim=-1).flatten(-2)
+        return torch.cat([rot, tail], dim=-1)
+
+    def step(self, u, state):
+        angle_state, ssm_state, k_prev, v_prev = state
+        m = self.mamba3
+        H, P, N, S = self.nheads, self.headdim, self.d_state, self.num_rope_angles
+        dtype = u.dtype
+        with torch.autocast(device_type=u.device.type, enabled=False):
+            zxb = F.linear(u.float(), m.in_proj.weight.float())
+            z, x, B, C, dd_dt, dd_A, trap_proj, ang_proj = torch.split(
+                zxb, [self.d_inner, self.d_inner, N, N, H, H, H, S], dim=-1
+            )
+            A = (-_heavy_tail(dd_A)).clamp(max=-m.A_floor)  # (b,H)
+            DT = F.softplus(dd_dt + m.dt_bias.float()).clamp(min=1e-6, max=100.0)
+            lam = torch.sigmoid(trap_proj)
+            alpha = torch.exp(DT * A)
+            gamma = lam * DT
+            beta = (1.0 - lam) * DT * alpha
+            B = _rms(B, m.B_norm.weight, m.B_norm.eps).unsqueeze(1) + m.B_bias.float().squeeze(
+                1
+            )  # (b,H,N)
+            C = _rms(C, m.C_norm.weight, m.C_norm.eps).unsqueeze(1) + m.C_bias.float().squeeze(1)
+            angle_new = angle_state + ang_proj.unsqueeze(1) * DT.unsqueeze(
+                -1
+            )  # (b,H,S) cumsum(angle*DT)
+            cos, sin = angle_new.cos(), angle_new.sin()
+            B, C = self._rotate(B, cos, sin), self._rotate(C, cos, sin)
+            x = x.reshape(-1, H, P)
+            z = z.reshape(-1, H, P)
+            new_ssm = (
+                ssm_state * alpha[..., None, None]
+                + beta[..., None, None] * (v_prev.unsqueeze(-1) * k_prev.unsqueeze(-2))
+                + gamma[..., None, None] * (x.unsqueeze(-1) * B.unsqueeze(-2))
+            )
+            new_ssm = new_ssm.clamp(min=-1e4, max=1e4)  # same hard stop as Mamba-1/2 cells
+            y = torch.einsum("bhpn,bhn->bhp", new_ssm, C) + m.D.float().view(1, H, 1) * x
+            y = (y * F.silu(z)).reshape(-1, self.d_inner).clamp(min=-1e4, max=1e4)
+        out = m.out_proj(y.to(dtype))  # back under ambient autocast
+        return out, (angle_new, new_ssm, B, x)
+
+
+class Mamba3ControllerBlock(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        d_state=64,
+        expand=2,
+        headdim=64,
+        ngroups=1,
+        rope_fraction=0.5,
+        layer_idx=None,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model, device=device, dtype=dtype)
+        self.cell = Mamba3ControllerCell(
+            d_model,
+            d_state=d_state,
+            expand=expand,
+            headdim=headdim,
+            ngroups=ngroups,
+            rope_fraction=rope_fraction,
+            layer_idx=layer_idx,
+            device=device,
+            dtype=dtype,
+        )
+
+    def init_state(self, batch_size, device=None, dtype=None):
+        return self.cell.init_state(batch_size, device=device, dtype=dtype)
+
+    def step(self, x, state):
+        out, new_state = self.cell.step(self.norm(x), state)
+        return x + out, new_state
+
+
+class Mamba3ControllerWrapper(Mamba2ControllerWrapper):
+    """Reuses Mamba2ControllerWrapper.forward()/init_state() unchanged (they only touch
+    in_adapter / blocks / moe_blocks / moe_enabled); only __init__ differs."""
+
+    def __init__(
+        self,
+        in_dim,
+        d_model,
+        num_blocks=2,
+        moe_enabled=False,
+        moe_num_experts=8,
+        moe_expert_dim=None,
+        moe_capacity_factor=1.5,
+        moe_load_balance_alpha=0.01,
+        moe_top_k=1,
+        d_state=64,
+        expand=2,
+        headdim=64,
+        ngroups=1,
+        rope_fraction=0.5,
+        device=None,
+        dtype=None,
+    ):
+        nn.Module.__init__(self)
+        self.d_model, self.num_blocks = d_model, num_blocks
+        self.in_adapter = (
+            nn.Identity()
+            if in_dim == d_model
+            else nn.Linear(in_dim, d_model, device=device, dtype=dtype)
+        )
+        self.blocks = nn.ModuleList(
+            [
+                Mamba3ControllerBlock(
+                    d_model,
+                    d_state=d_state,
+                    expand=expand,
+                    headdim=headdim,
+                    ngroups=ngroups,
+                    rope_fraction=rope_fraction,
+                    layer_idx=i,
+                    device=device,
+                    dtype=dtype,
+                )
+                for i in range(num_blocks)
+            ]
+        )
+        self.moe_enabled = moe_enabled
+        self.moe_blocks = None
+        if moe_enabled:
+            self.moe_blocks = nn.ModuleList(
+                [
+                    MoEBlock(
+                        d_model,
+                        num_experts=moe_num_experts,
+                        expert_dim=moe_expert_dim,
+                        capacity_factor=moe_capacity_factor,
+                        top_k=moe_top_k,
+                        load_balance_alpha=moe_load_balance_alpha,
+                        device=device,
+                        dtype=dtype,
+                    )
+                    for _ in range(num_blocks)
+                ]
+            )

initium/mamba_controller/mamba_backbone_parallel.py

@@ -0,0 +1,291 @@
+"""
+file: mamba_backbone_parallel.py
+
+Alternate Phase 3, Step 2, Option 5 ("split the compute graph, don't try to
+parallelize the addressing itself") -- see Experiment-Roadmap.md.
+
+This module implements ONLY the "(a) Mamba backbone" half of the Option 5
+split: a stack of Mamba (v1) or Mamba-2 blocks driven over the WHOLE
+sequence at once via each block's own native, library-provided forward()
+(the fast, parallel-scan / SSD path), with NO per-timestep interleaving of
+memory reads.
+
+This is a deliberate, much simpler sibling to mamba_controller.py's
+MambaControllerCell/Block/Wrapper (Alternative Phase 3, Step 1). That file
+had to hand-roll a BPTT-safe manual step() specifically because the DNC's
+per-timestep interleaving loop needed to call the controller one token at a
+time while chaining (conv_state, ssm_state) through a Python loop under
+full backprop, and mamba_ssm's own .step() mutates its state tensors
+in-place (fatal for BPTT -- see that file's module docstring for the full
+diagnosis).
+
+Option 5's entire point is to remove that constraint for the backbone: since
+this backbone's per-token computation has NO dependency on M_{t-1} (the
+memory's state at the previous timestep) -- it only ever consumes the raw
+input sequence X, never a read vector -- there is nothing forcing it into a
+sequential Python loop. It can therefore call mamba_ssm's OWN whole-sequence
+forward() directly, which already uses the fast parallel selective-scan
+(Mamba-1) or the chunked, matmul-based SSD algorithm (Mamba-2, Dao & Gu
+2024) -- no custom step function, no in-place-op workaround needed, because
+nothing here is chained through a live memory-state Python loop.
+
+The other half of the Option 5 split -- the "genuinely sequential, but much
+smaller and cheaper, memory read/write/addressing step" that consumes this
+backbone's parallel output as an input stream -- lives in
+split_graph_dnc.py, not here. This file only produces H = backbone(X); it
+has no notion of DNC, Memory, or read vectors at all, by design (see
+Concept 6 / LB-9's caution about structural-vs-functional separation --
+keeping this file's API surface structurally incapable of touching memory
+state is a guarantee, not just a convention).
+
+Supports two variants, both from the mamba-ssm package (per the roadmap's
+"add Mamba-2 first if Option 5 needs it, otherwise Mamba-1 is already
+implemented" instruction -- Option 5 benefits directly from Mamba-2's own
+SSD chunked-scan algorithm, since that algorithm is specifically optimized
+for exactly the "process the whole sequence at once" call pattern Option 5
+newly unlocks for the backbone -- see the SSD paper, Sections 6 and 9.3,
+already in corpus):
+
+  - "mamba1": mamba_ssm.modules.mamba_simple.Mamba
+  - "mamba2": mamba_ssm.modules.mamba2.Mamba2
+
+Both are used exactly as their own package intends (whole-sequence
+forward(hidden_states) call, batch_first (B, L, D) tensors) -- this file
+does not reimplement or bypass anything of either class, unlike
+mamba_controller.py's cell.
+
+NOTE on Mamba2 kwargs: mamba_ssm's Mamba2 additionally requires `headdim`
+(d_model must be divisible by it -- SSD paper's own convention is
+head_dim in {64, 128}). This has no analogue in Mamba-1 and is therefore
+a new, Mamba2-only constructor argument here, same pattern as
+mamba_d_state/mamba_d_conv/mamba_expand having no LSTM analogue in
+mamba_controller.py.
+"""
+
+from __future__ import annotations
+
+import os as _os
+
+import torch
+import torch.nn as nn
+
+from initium.MoE.moe_layer import MoEBlock
+
+_MAMBA1_IMPORT_ERROR: str | None
+_MAMBA2_IMPORT_ERROR: str | None
+_MAMBA3_IMPORT_ERROR: str | None
+
+try:
+    from mamba_ssm.modules.mamba_simple import Mamba as Mamba1
+except ImportError as _mamba1_err:  # pragma: no cover - environment-dependent
+    Mamba1 = None
+    _MAMBA1_IMPORT_ERROR = (
+        "mamba_backbone_parallel.py requires the `mamba-ssm` package for "
+        f"variant='mamba1'. Original import error: {_mamba1_err}"
+    )
+else:
+    _MAMBA1_IMPORT_ERROR = None
+
+try:
+    from mamba_ssm.modules.mamba2 import Mamba2
+except ImportError as _mamba2_err:  # pragma: no cover - environment-dependent
+    Mamba2 = None
+    _MAMBA2_IMPORT_ERROR = (
+        "mamba_backbone_parallel.py requires the `mamba-ssm` package "
+        "(with Mamba-2 support -- mamba_ssm.modules.mamba2.Mamba2) for "
+        f"variant='mamba2'. Original import error: {_mamba2_err}"
+    )
+else:
+    _MAMBA2_IMPORT_ERROR = None
+
+try:
+    from mamba_ssm.modules.mamba3 import Mamba3
+except ImportError as _mamba3_err:  # pragma: no cover - environment-dependent
+    Mamba3 = None
+    _MAMBA3_IMPORT_ERROR = (
+        "mamba_backbone_parallel.py: variant='mamba3' needs mamba-ssm installed from GitHub main "
+        f"(Mamba3 is not in the v2.3.1 wheel). Original import error: {_mamba3_err}"
+    )
+else:
+    _MAMBA3_IMPORT_ERROR = None
+
+# bf16 (the repo's own precision) when the GPU supports it; otherwise the fp32-under-fp16-autocast trick.
+_MAMBA3_BF16_OK = (
+    torch.cuda.is_available()
+    and torch.cuda.get_device_capability()[0] >= 8
+    and _os.environ.get("MAMBA3_FORCE_FP32", "0") != "1"
+)
+
+
+def _require_variant(variant: str) -> None:
+    if variant == "mamba1" and Mamba1 is None:
+        raise ImportError(_MAMBA1_IMPORT_ERROR)
+    if variant == "mamba2" and Mamba2 is None:
+        raise ImportError(_MAMBA2_IMPORT_ERROR)
+    if variant == "mamba3" and Mamba3 is None:
+        raise ImportError(_MAMBA3_IMPORT_ERROR)
+    if variant not in ("mamba1", "mamba2", "mamba3"):
+        raise ValueError(
+            f"mamba_backbone_parallel: unknown variant {variant!r}, expected 'mamba1' or 'mamba2'."
+        )
+
+
+class _ParallelBlock(nn.Module):
+    """Pre-norm residual wrapper around one whole-sequence Mamba block,
+    mirroring mamba_ssm's own Block (Add -> LN -> Mixer) and
+    mamba_controller.py's MambaControllerBlock, but calling the mixer's own
+    forward() over the full (B, L, D) sequence at once instead of a
+    per-timestep step(). No custom state threading of any kind -- the
+    mixer manages its own internal scan/chunking entirely within this one
+    call.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        variant: str,
+        mamba_kwargs: dict,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+        _require_variant(variant)
+        self.norm = nn.LayerNorm(d_model, device=device, dtype=dtype)
+        self.variant = variant
+        if variant == "mamba3":
+            self.mixer = Mamba3(d_model=d_model, device=device, dtype=dtype, **mamba_kwargs)
+        elif variant == "mamba1":
+            self.mixer = Mamba1(d_model=d_model, device=device, dtype=dtype, **mamba_kwargs)
+        else:
+            self.mixer = Mamba2(d_model=d_model, device=device, dtype=dtype, **mamba_kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, L, D). mamba_ssm's own forward() handles the whole
+        # sequence's parallel scan / SSD chunked-scan internally -- this is
+        # the entire mechanism that makes Option 5 different from
+        # mamba_controller.py: one call per block, not one call per DNC
+        # timestep.
+        #
+        # FIX (chronic grad_norm nan under fp16 AMP, Mamba-2 variant only in
+        # practice but applied to both for symmetry): same wrapper trick
+        # already used by mamba_controller.py's MambaControllerCell.step()
+        # and mamba2_controller.py's Mamba2ControllerCell.step() -- force
+        # this call to run in real fp32, outside ambient fp16 autocast,
+        # rather than reimplementing the mixer's internal math by hand the
+        # way those cells do for their own per-timestep step(). Unlike those
+        # cells, self.mixer here is the library's own whole-sequence
+        # forward() (that's the entire point of the parallel backbone), so
+        # we can't insert clamps mid-computation the way step() does --
+        # instead we just deny it fp16 entirely. Mamba-2's SSD algorithm
+        # computes exp() of a cumulative sum of (dt * A) over a whole chunk
+        # (paper Section 6, Listing 1's segsum), which is far more prone to
+        # fp16 underflow/overflow than Mamba-1's fused selective-scan
+        # kernel -- this is the direct cause of the amp_scale-collapse-to-
+        # 0.0 pattern observed in mamba2-backbone split-graph runs that
+        # mamba1-backbone runs didn't show. Costs some speed relative to
+        # running under fp16 (this block no longer benefits from tensor-core
+        # fp16 matmuls), but matches every other Mamba SSM computation path
+        # in this project, all of which already force fp32 for this reason.
+        if self.variant == "mamba3" and _MAMBA3_BF16_OK:
+            with torch.autocast(device_type=x.device.type, dtype=torch.bfloat16, enabled=True):
+                out = self.mixer(self.norm(x.float()).to(torch.bfloat16))
+        else:
+            with torch.autocast(device_type=x.device.type, enabled=False):
+                out = self.mixer(self.norm(x.float()))
+        return x + out.to(x.dtype)
+
+
+class MambaBackboneParallel(nn.Module):
+    """Stack of `num_blocks` whole-sequence Mamba blocks. Pure sequence-to-
+    sequence map (B, L, in_dim) -> (B, L, d_model), with NO awareness of
+    DNC, Memory, or read vectors -- see module docstring.
+
+    variant: "mamba1" (mamba_ssm.modules.mamba_simple.Mamba, default
+        Mamba-1 hyperparameters d_state=16/d_conv=4/expand=2, this
+        project's original Alternate-Phase-3 choice) or "mamba2"
+        (mamba_ssm.modules.mamba2.Mamba2, needs `headdim` in addition;
+        default 64, matching the SSD paper's own convention of head_dim in
+        {64, 128}).
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        d_model: int,
+        num_blocks: int = 2,
+        variant: str = "mamba1",
+        d_state: int = 16,
+        d_conv: int = 4,
+        expand: int = 2,
+        headdim: int = 64,  # mamba2-only
+        moe_enabled: bool = False,
+        moe_num_experts: int = 8,
+        moe_expert_dim: int | None = None,
+        moe_top_k: int = 1,
+        moe_capacity_factor: float = 1.5,
+        moe_load_balance_alpha: float = 0.01,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+        _require_variant(variant)
+        self.variant = variant
+        self.d_model = d_model
+        self.moe_enabled = moe_enabled
+        self.moe_blocks: nn.ModuleList | None = None
+
+        # Same role as MambaControllerWrapper's in_adapter: DNC's raw
+        # per-timestep input (input_size) generally differs from
+        # hidden_size/d_model. Unlike the interleaved wrapper, this is
+        # applied to the WHOLE sequence at once (a single
+        # (B, L, in_dim) -> (B, L, d_model) matmul), not per-timestep.
+        self.in_adapter: nn.Module = (
+            nn.Identity()
+            if in_dim == d_model
+            else nn.Linear(in_dim, d_model, device=device, dtype=dtype)
+        )
+
+        mamba_kwargs = dict(d_state=d_state, d_conv=d_conv, expand=expand)
+        if variant == "mamba2":
+            mamba_kwargs["headdim"] = headdim
+        if variant == "mamba3":
+            mamba_kwargs["headdim"] = headdim
+            mamba_kwargs["chunk_size"] = (
+                64 if _MAMBA3_BF16_OK else 32
+            )  # repo: 64 bf16 / 32 otherwise
+
+        self.blocks = nn.ModuleList(
+            [
+                _ParallelBlock(d_model, variant, mamba_kwargs, device=device, dtype=dtype)
+                for _ in range(num_blocks)
+            ]
+        )
+        if moe_enabled:
+            self.moe_blocks = nn.ModuleList(
+                [
+                    MoEBlock(
+                        d_model,
+                        num_experts=moe_num_experts,
+                        expert_dim=moe_expert_dim,
+                        top_k=moe_top_k,
+                        capacity_factor=moe_capacity_factor,
+                        load_balance_alpha=moe_load_balance_alpha,
+                        device=device,
+                        dtype=dtype,
+                    )
+                    for _ in range(num_blocks)
+                ]
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, L, in_dim) -> (B, L, d_model). Called EXACTLY ONCE per
+        # training step (not once per DNC timestep, unlike
+        # MambaControllerWrapper) -- this single call is the entire
+        # "backbone half" of the Option 5 split.
+        h = self.in_adapter(x)
+        for i, block in enumerate(self.blocks):
+            h = block(h)
+            if self.moe_enabled:
+                assert self.moe_blocks is not None
+                h = self.moe_blocks[i](h)
+        return h