giant/giant/model/trunks.py

"""Trunks: everything downstream of the fused conditioning vector — a
registrable expert *body* architecture (`TRUNK_REGISTRY`/`register_trunk`),
used standalone or mixed by a `Router` (issues.md Issue 8; trunk-selectability
gitea #33).

Whether a body is mixed is orthogonal to which body it is: `RoutedTrunk`
builds `router.n_experts` instances of whichever body `trunk_type` names, so
a future body (e.g. a transformer) automatically gets a mixture variant for
free — no separate "routed transformer trunk" class needed.
"""

import torch
import torch.nn as nn

from giant.model.layers import build_block
from giant.model.routers import Router

TRUNK_REGISTRY: dict[str, type[nn.Module]] = {}


def register_trunk(name: str):
    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
        TRUNK_REGISTRY[name] = cls
        return cls

    return decorator


def build_expert_body(
    name: str,
    in_dim: int,
    out_dim: int,
    hidden_dim: int,
    n_blocks: int,
    cond_dim: int,
    dropout: float = 0.0,
    block_conditioning: str = "add",
) -> nn.Module:
    """Factory: look up a registered trunk body by name and construct one
    instance of it — used both for a standalone (unrouted) trunk and for each
    expert inside a `RoutedTrunk`. `block_conditioning` selects the
    `BLOCK_REGISTRY` entry each body's internal `ResBlock`-family blocks use
    (`trunk.block_conditioning`, gitea #34) — an optional trailing kwarg a
    future non-`ResBlock`-based body can simply ignore, same idiom as
    `Trunk.forward`'s accept-and-ignore `cond_cont`/`cond_cat`."""
    if name not in TRUNK_REGISTRY:
        raise ValueError(f"unknown trunk type {name!r}; available: {sorted(TRUNK_REGISTRY)}")
    cls = TRUNK_REGISTRY[name]
    return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout, block_conditioning=block_conditioning)


@register_trunk("resmlp")
class ExpertTrunk(nn.Module):
    """`input_proj -> ResBlock stack -> out_proj` — the registered `"resmlp"`
    trunk body. Used both standalone (no router: `forward`'s `cond_cont`/
    `cond_cat` are accepted and ignored, satisfying the `Trunk` interface
    directly with no wrapper class) and as one expert inside a `RoutedTrunk`
    (`_route_forward` calls it with just `(x, cond)`).

    Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR
    tokens later (`noise_dim` in, `4 + type_dim` out), even though every
    step-2/3 caller still has `in_dim == out_dim`.
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int,
        n_blocks: int,
        cond_dim: int,
        dropout: float = 0.0,
        block_conditioning: str = "add",
    ) -> None:
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.input_proj = nn.Linear(in_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [build_block(block_conditioning, hidden_dim, cond_dim, dropout) for _ in range(n_blocks)]
        )
        self.out_proj = nn.Linear(hidden_dim, out_dim)

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        cond_cont: torch.Tensor | None = None,
        cond_cat: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x = self.input_proj(x)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


@register_trunk("linear")
class LinearTrunk(nn.Module):
    """`nn.Linear(in_dim + cond_dim, out_dim)` over `concat([x, cond])` —
    the trivial trunk body: no hidden layer, no ResBlock stack, no
    nonlinearity. Ablates whether trunk depth/nonlinearity is earning its
    parameters, holding everything else (heads, ConditionEncoder,
    generator, ...) fixed. Composes for free with `router.enabled = true`
    (gitea #33): a RoutedTrunk of n_experts linear bodies is "mixture of
    trivial linear experts". `hidden_dim`/`n_blocks`/`dropout`/
    `block_conditioning` are accepted and ignored, matching
    `build_expert_body`'s shared factory signature.

    `x` — the trunk's own input (e.g. the noised primary vector for flow
    matching) — does not already carry conditioning; that's fused in
    per-body via `cond`. So this concatenates `x` and `cond` itself to
    remain a valid, conditioning-dependent model.
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int,
        n_blocks: int,
        cond_dim: int,
        dropout: float = 0.0,
        block_conditioning: str = "add",
    ) -> None:
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.linear = nn.Linear(in_dim + cond_dim, out_dim)

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        cond_cont: torch.Tensor | None = None,
        cond_cat: torch.Tensor | None = None,
    ) -> torch.Tensor:
        return self.linear(torch.cat([x, cond], dim=-1))


def _route_forward(
    experts: nn.ModuleList,
    router: Router,
    x: torch.Tensor,
    cond: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    training: bool,
) -> torch.Tensor:
    """Shared dispatch for `RoutedTrunk`.

    Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
    N-expert dense compute, fully differentiable (`weight` is
    `router.combine_weights`). Eval mode: grouped top-1 dispatch — each row
    runs exactly one expert, the actual source of the per-call speedup.

    The accumulator's dtype is deferred to the first expert call rather than
    fixed at fp32: under autocast (`train.precision = "bf16"`, gitea #47) an
    expert's `ResBlock` stack returns bf16, and an fp32-fixed accumulator
    would silently upcast every mixture term (train mode) or downcast every
    dispatched row via `index_put_` (eval mode) — making a `RoutedTrunk`
    return a different dtype than the unrouted `ExpertTrunk` it's a drop-in
    replacement for, purely because `router.enabled` was set.

    `router.combine_weights` is deliberately fp32 internally (it forces its
    own autocast-disabled region — see `Router.combine_weights`'s docstring),
    so `weights` itself is always fp32 regardless of the ambient precision.
    Left as-is, `weights[:, i:i+1] * expert(x, cond)` would type-promote the
    whole mixture back to fp32 by ordinary PyTorch promotion rules — the same
    dtype-mismatch bug this function exists to avoid, just moved one line
    over. `weights` is cast down to each expert's own output dtype right
    before combining: the softmax stays numerically stable at fp32, but its
    *result* (values in [0, 1], not precision-sensitive to represent) loses
    nothing meaningful by then being used at bf16.
    """
    if training:
        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts), fp32
        out = None
        for i, expert in enumerate(experts):
            expert_out = expert(x, cond)
            term = weights[:, i : i + 1].to(expert_out.dtype) * expert_out
            out = term if out is None else out + term
        assert out is not None, "RoutedTrunk built with zero experts"
        return out

    idx = router.top1(cond_cont, cond_cat)  # (B,)
    out = None
    for i, expert in enumerate(experts):
        mask = idx == i
        if mask.any():
            expert_out = expert(x[mask], cond[mask])
            if out is None:
                out = torch.zeros(x.shape[0], expert_out.shape[-1], device=x.device, dtype=expert_out.dtype)
            out[mask] = expert_out
    if out is None:
        # No row was ever dispatched (only reachable with an empty batch,
        # x.shape[0] == 0) — nothing to infer a dtype from, so fall back to
        # x's own, matching this function's pre-autocast behavior.
        out = torch.zeros(x.shape[0], experts[0].out_dim, device=x.device, dtype=x.dtype)
    return out


class Trunk(nn.Module):
    """Interface implemented by a standalone trunk body (any `TRUNK_REGISTRY`
    entry, e.g. `ExpertTrunk`) and by `RoutedTrunk`: everything downstream of
    the fused conditioning vector, i.e. the actual generative trunk of a
    stage. Implementations are expected to expose `in_dim`/`out_dim`
    attributes (as `ExpertTrunk`/`RoutedTrunk` do) — `giant.model.summary`
    (gitea #46) reads them to report trunk widths without needing to know the
    body architecture."""

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        raise NotImplementedError


class RoutedTrunk(Trunk):
    def __init__(
        self,
        router: Router,
        trunk_type: str,
        in_dim: int,
        out_dim: int,
        hidden_dim: int,
        n_res_blocks: int,
        cond_dim: int,
        dropout: float = 0.0,
        block_conditioning: str = "add",
    ) -> None:
        super().__init__()
        self.router = router
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.experts = nn.ModuleList(
            [
                build_expert_body(
                    trunk_type,
                    in_dim,
                    out_dim,
                    hidden_dim,
                    n_res_blocks,
                    cond_dim,
                    dropout,
                    block_conditioning=block_conditioning,
                )
                for _ in range(router.n_experts)
            ]
        )

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training)


def build_trunk(
    router: Router | None,
    trunk_type: str,
    in_dim: int,
    out_dim: int,
    hidden_dim: int,
    n_res_blocks: int,
    cond_dim: int,
    dropout: float = 0.0,
    block_conditioning: str = "add",
) -> nn.Module:
    """Build a stage's trunk: `trunk_type` (a `TRUNK_REGISTRY` key, e.g.
    `"resmlp"`) selects the expert body architecture; `router`, if given,
    wraps `router.n_experts` instances of that body in a `RoutedTrunk`
    mixture — otherwise a single body is returned directly (no wrapper
    class), which is what makes an unrouted trunk's state-dict keys land
    directly under `trunk.*` instead of `trunk.experts.0.*` (see
    `giant.model._legacy.migrate_legacy_state_dict`, which assumes exactly
    this flat layout for a v0.2 monolithic checkpoint). `block_conditioning`
    (a `BLOCK_REGISTRY` key, e.g. `"add"`/`"film"`/`"adaln"`) selects each
    body's conditioning-injection mechanism (gitea #34).
    """
    if router is not None:
        return RoutedTrunk(
            router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
        )
    return build_expert_body(
        trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
    )