custom_nodes/ComfyUI-MiniMaxH3-Adapter/adapter/model.py

"""H3 文本编码器替换适配器（接口蒸馏）：source_projection + CrossAttention Resampler + TokenRefiner。

前向:
    h3_ids [B, S_T] --QueryEmbedding--> Q [B, S_T, 5376]
    student_hidden [B, S_S, 2560] --source_projection--> KV [B, S_S, 5376]
    CrossAttentionBlock(Q, KV) -> [B, S_T, 5376]
    TokenRefiner(2 层, 原权重初始化) -> [B, S_T, 5376]   # 与教师 target 同坐标系

TokenRefiner 结构严格复刻原始 checkpoint（已源码核实）:
    - fused qkv (chunk(3)) + per-head qk_norm + 双向注意力 + out_proj，全部 bias=False
    - SwiGLU MLP: fc1 为 fused [gate; value]（gate 在前），fc2(silu(gate)*value)
    - 2 个 pre-norm block + final RMSNorm，eps 全部 1e-5
"""
from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

H3_VOCAB = 151936
H3_HIDDEN = 5376
REFINER_HEADS = 56
REFINER_HEAD_DIM = 128
REFINER_FFN = 14336
STUDENT_HIDDEN = 2560


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


class SwiGLUFFN(nn.Module):
    def __init__(self, hidden: int, ffn: int):
        super().__init__()
        self.fc1 = nn.Linear(hidden, 2 * ffn, bias=False)
        self.fc2 = nn.Linear(ffn, hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate, value = self.fc1(x).chunk(2, dim=-1)
        return self.fc2(F.silu(gate) * value)


class TokenRefinerAttention(nn.Module):
    def __init__(self, hidden: int, heads: int, dim_head: int):
        super().__init__()
        self.heads = heads
        self.head_dim = dim_head
        self.inner_dim = heads * dim_head
        self.qkv_proj = nn.Linear(hidden, 3 * self.inner_dim, bias=False)
        self.q_norm = RMSNorm(dim_head)
        self.k_norm = RMSNorm(dim_head)
        self.out_proj = nn.Linear(self.inner_dim, hidden, bias=False)
        self.use_sdpa = True  # SDPA flash: 注意力内存 O(S²)->O(S)，数值与手写注意力差 ~1e-3（蒸馏噪声级）

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
        q = self.q_norm(q.unflatten(-1, (self.heads, self.head_dim)).transpose(1, 2))
        k = self.k_norm(k.unflatten(-1, (self.heads, self.head_dim)).transpose(1, 2))
        v = v.unflatten(-1, (self.heads, self.head_dim)).transpose(1, 2)
        if self.use_sdpa:
            out = F.scaled_dot_product_attention(q, k, v, scale=self.head_dim ** -0.5)
        else:
            attn = torch.softmax((q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5), dim=-1)
            out = attn @ v
        return self.out_proj(out.transpose(1, 2).flatten(2))


class TokenRefinerBlock(nn.Module):
    def __init__(self, hidden: int, heads: int, dim_head: int, ffn: int):
        super().__init__()
        self.norm1 = RMSNorm(hidden)
        self.attn = TokenRefinerAttention(hidden, heads, dim_head)
        self.norm2 = RMSNorm(hidden)
        self.mlp = SwiGLUFFN(hidden, ffn)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x


class TokenRefiner(nn.Module):
    def __init__(self, num_layers: int = 2, hidden: int = H3_HIDDEN,
                 heads: int = REFINER_HEADS, dim_head: int = REFINER_HEAD_DIM, ffn: int = REFINER_FFN):
        super().__init__()
        self.blocks = nn.ModuleList([TokenRefinerBlock(hidden, heads, dim_head, ffn) for _ in range(num_layers)])
        self.final_norm = RMSNorm(hidden)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            x = block(x)
        return self.final_norm(x)


class CrossAttentionBlock(nn.Module):
    """跨空间交叉注意（学生 2560 投影空间 <-> 教师 5376 表示空间）。

    训练稳定性升级（兼容旧权重加载）:
      - QK-norm: to_q/to_k 投影后对每头做 RMSNorm（Llama 3.2 vision 交叉注意惯例，
        稳定跨空间注意力的 q/k 尺度失配）
      - gated tanh: out_proj 输出经 tanh(gate) 门控（Flamingo 惯例），gate 初始 0 => 恒等
      - attn dropout: 交叉注意 dropout 0.1（Emu3 防后期 collapse）
    """

    def __init__(self, hidden: int = H3_HIDDEN, heads: int = 32, dim_head: int = 128, ffn: int = REFINER_FFN,
                 attn_dropout: float = 0.1, use_qk_norm: bool = True, use_gate: bool = True):
        super().__init__()
        self.heads = heads
        self.head_dim = dim_head
        self.inner_dim = heads * dim_head
        self.norm_q = RMSNorm(hidden)
        self.norm_kv = RMSNorm(hidden)
        self.to_q = nn.Linear(hidden, self.inner_dim, bias=False)
        self.to_k = nn.Linear(hidden, self.inner_dim, bias=False)
        self.to_v = nn.Linear(hidden, self.inner_dim, bias=False)
        self.out_proj = nn.Linear(self.inner_dim, hidden, bias=False)
        self.norm2 = RMSNorm(hidden)
        self.mlp = SwiGLUFFN(hidden, ffn)
        self.use_sdpa = True
        self.attn_dropout = attn_dropout
        self.use_qk_norm = use_qk_norm
        if use_qk_norm:
            self.qk_norm = RMSNorm(dim_head)
        self.use_gate = use_gate
        if use_gate:
            self.gate = nn.Parameter(torch.zeros(1))  # tanh(0)=0 -> 恒等，兼容旧权重

    def forward(self, q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
        qn = self.norm_q(q)
        kvn = self.norm_kv(kv)
        qh = self.to_q(qn).unflatten(-1, (self.heads, self.head_dim)).transpose(1, 2)
        kh = self.to_k(kvn).unflatten(-1, (self.heads, self.head_dim)).transpose(1, 2)
        vh = self.to_v(kvn).unflatten(-1, (self.heads, self.head_dim)).transpose(1, 2)
        if self.use_qk_norm:
            qh = self.qk_norm(qh)
            kh = self.qk_norm(kh)
        if self.use_sdpa:
            out = F.scaled_dot_product_attention(qh, kh, vh, scale=self.head_dim ** -0.5,
                                                 dropout_p=self.attn_dropout if self.training else 0.0)
        else:
            attn = torch.softmax((qh @ kh.transpose(-2, -1)) * (self.head_dim ** -0.5), dim=-1)
            out = attn @ vh
        proj = self.out_proj(out.transpose(1, 2).flatten(2))
        x = q + (torch.tanh(self.gate) * proj if self.use_gate else proj)
        x = x + self.mlp(self.norm2(x))
        return x


class QueryEmbedding(nn.Module):
    def __init__(self, vocab: int = H3_VOCAB, dim: int = 256, out: int = H3_HIDDEN):
        super().__init__()
        self.embed = nn.Embedding(vocab, dim)
        self.proj = nn.Linear(dim, out, bias=True)

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        return self.proj(self.embed(ids))


class H3Adapter(nn.Module):
    """完整适配器。param 约 1.14B（source_proj 13.8M + query_embed 40.3M + crossattn 319M + refiner 751M）。"""

    def __init__(self):
        super().__init__()
        self.source_projection = nn.Linear(STUDENT_HIDDEN, H3_HIDDEN, bias=True)
        self.query_embedding = QueryEmbedding()
        self.cross_attention = CrossAttentionBlock()
        self.token_refiner = TokenRefiner()

    def forward(self, h3_ids: torch.Tensor, student_hidden: torch.Tensor) -> torch.Tensor:
        kv = self.source_projection(student_hidden)
        q = self.query_embedding(h3_ids)
        x = self.cross_attention(q, kv)
        x = self.token_refiner(x)
        return x

    def load_token_refiner(self, state_dict: dict[str, torch.Tensor], strict: bool = True) -> None:
        refiner_state = {
            k[len("token_refiner."):]: v
            for k, v in state_dict.items()
            if k.startswith("token_refiner.")
        }
        missing, unexpected = self.token_refiner.load_state_dict(refiner_state, strict=strict)
        assert not missing and not unexpected, f"refiner load: missing={missing} unexpected={unexpected}"

    def trainable_modules(self, stage: int) -> list[nn.Parameter]:
        if stage == 1:
            return (
                list(self.source_projection.parameters())
                + list(self.query_embedding.parameters())
                + list(self.cross_attention.parameters())
            )
        return list(self.parameters())


class TeacherHead(nn.Module):
    """教师 target 计算: h50 -> condition_proj(5120->5376) -> token_refiner -> [B, S_T, 5376]。"""

    def __init__(self):
        super().__init__()
        self.condition_proj = nn.Linear(5120, H3_HIDDEN, bias=True)
        self.token_refiner = TokenRefiner()

    def forward(self, h50: torch.Tensor) -> torch.Tensor:
        return self.token_refiner(self.condition_proj(h50))


def compute_query_embed_init(embed_tokens: torch.Tensor, dim: int = 256, seed: int = 0) -> torch.Tensor:
    """query embedding 初始化: embed.weight = E @ P，P 为 5120->dim 的随机投影。

    仅用于一次性初始化，后续可训练。embed_tokens 为 [151936, 5120]（BF16 或 FP32）。
    """
    assert embed_tokens.dim() == 2 and embed_tokens.shape[1] == 5120, embed_tokens.shape
    rng = torch.Generator().manual_seed(seed)
    p = torch.randn(5120, dim, generator=rng) * (1.0 / math.sqrt(5120))
    embed_tokens = embed_tokens.to(torch.float32)
    chunks = 8
    out = torch.empty(embed_tokens.shape[0], dim, dtype=torch.float32)
    for i in range(chunks):
        lo = i * embed_tokens.shape[0] // chunks
        hi = (i + 1) * embed_tokens.shape[0] // chunks
        out[lo:hi] = embed_tokens[lo:hi] @ p
    return out.to(torch.bfloat16)