252 lines
12 KiB
Python
252 lines
12 KiB
Python
"""Direct prompt-only Qwen3-VL-32B conditioning for MiniMax H3.
|
|
|
|
The MiniMax checkpoint contains the first 50 Qwen3-VL decoder layers. H3
|
|
uses their unnormalized final output, not the language-model head or final
|
|
RMSNorm. This module intentionally has no ComfyUI imports and does not load
|
|
the Qwen vision encoder; image/video prompt construction remains a separate
|
|
feature.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
|
|
from .checkpoint import H3Checkpoint
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Qwen3VL32BTextConfig:
|
|
vocab_size: int = 151936
|
|
hidden_size: int = 5120
|
|
intermediate_size: int = 25600
|
|
num_layers: int = 50
|
|
num_attention_heads: int = 64
|
|
num_key_value_heads: int = 8
|
|
head_dim: int = 128
|
|
rms_norm_eps: float = 1e-6
|
|
rope_theta: float = 5_000_000.0
|
|
|
|
|
|
class _RMSNorm(nn.Module):
|
|
def __init__(self, weight: torch.Tensor, eps: float):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.register_buffer("weight", weight, persistent=False)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return F.rms_norm(x, self.weight.shape, weight=self.weight.to(x), eps=self.eps)
|
|
|
|
|
|
def _apply_rope(query: torch.Tensor, key: torch.Tensor, freqs) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Direct PyTorch port of Comfy's `apply_rope` fed `freqs=(cos, sin, neg_sin)`.
|
|
|
|
``query``/``key`` are ``[batch, heads, seq, head_dim]``; ``freqs`` broadcast over
|
|
that shape. This is the exact split-half kernel the reference uses for both the
|
|
plain text rope and the Qwen3-VL interleaved mrope.
|
|
"""
|
|
cosine, sine, negative_sine = freqs
|
|
split = query.shape[-1] // 2
|
|
q = query * cosine
|
|
q[..., :split].addcmul_(query[..., split:], negative_sine)
|
|
q[..., split:].addcmul_(query[..., :split], sine)
|
|
k = key * cosine
|
|
k[..., :split].addcmul_(key[..., split:], negative_sine)
|
|
k[..., split:].addcmul_(key[..., :split], sine)
|
|
return q.to(query.dtype), k.to(key.dtype)
|
|
|
|
|
|
def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Direct PyTorch port of Comfy's text-only `precompute_freqs_cis` / `apply_rope`."""
|
|
sequence, head_dim = query.shape[-2:]
|
|
position_ids = torch.arange(sequence, device=query.device).unsqueeze(0)
|
|
theta_numerator = torch.arange(0, head_dim, 2, device=query.device).float()
|
|
inv_freq = 1.0 / (theta ** (theta_numerator / head_dim))
|
|
frequencies = (inv_freq[None, :, None].expand(1, -1, 1).float() @ position_ids[:, None, :].float()).transpose(1, 2)
|
|
embedding = torch.cat((frequencies, frequencies), dim=-1)
|
|
cosine = embedding.cos().unsqueeze(1)
|
|
sine = embedding.sin().unsqueeze(1)
|
|
negative_sine = -sine[..., sine.shape[-1] // 2 :]
|
|
sine = sine[..., : sine.shape[-1] // 2]
|
|
return _apply_rope(query, key, (cosine, sine, negative_sine))
|
|
|
|
|
|
def _qwen_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
|
"""Match Comfy's small-input SDPA wrapper for masked Qwen GQA."""
|
|
kwargs = {"attn_mask": mask.unsqueeze(0).unsqueeze(0), "dropout_p": 0.0, "is_causal": False, "enable_gqa": True}
|
|
if query.numel() >= 1024 * 128:
|
|
from torch.nn.attention import SDPBackend, sdpa_kernel
|
|
|
|
priority = [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
|
|
params = torch.backends.cuda.SDPAParams(query, key, value, kwargs["attn_mask"], 0.0, False, True)
|
|
supports_native_gqa = (
|
|
torch.backends.cuda.can_use_flash_attention(params)
|
|
or torch.backends.cuda.can_use_cudnn_attention(params)
|
|
or torch.backends.cuda.can_use_efficient_attention(params)
|
|
)
|
|
if not supports_native_gqa:
|
|
repeats = query.shape[-3] // key.shape[-3]
|
|
key = key.repeat_interleave(repeats, dim=-3)
|
|
value = value.repeat_interleave(repeats, dim=-3)
|
|
kwargs["enable_gqa"] = False
|
|
with sdpa_kernel(priority, set_priority=True):
|
|
return F.scaled_dot_product_attention(query, key, value, **kwargs)
|
|
return F.scaled_dot_product_attention(query, key, value, **kwargs)
|
|
|
|
|
|
class _Qwen3VLBlock(nn.Module):
|
|
def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype):
|
|
super().__init__()
|
|
self.config = config
|
|
self.input_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.input_layernorm.weight", dtype=dtype), config.rms_norm_eps)
|
|
self.post_attention_layernorm = _RMSNorm(checkpoint.tensor(f"{prefix}.post_attention_layernorm.weight", dtype=dtype), config.rms_norm_eps)
|
|
self.q_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.q_norm.weight", dtype=dtype), config.rms_norm_eps)
|
|
self.k_norm = _RMSNorm(checkpoint.tensor(f"{prefix}.self_attn.k_norm.weight", dtype=dtype), config.rms_norm_eps)
|
|
self.q_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.q_proj", output_dtype=dtype)
|
|
self.k_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.k_proj", output_dtype=dtype)
|
|
self.v_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.v_proj", output_dtype=dtype)
|
|
self.o_proj = checkpoint.nvfp4_linear(f"{prefix}.self_attn.o_proj", output_dtype=dtype)
|
|
self.gate_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.gate_proj", output_dtype=dtype)
|
|
self.up_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.up_proj", output_dtype=dtype)
|
|
self.down_proj = checkpoint.nvfp4_linear(f"{prefix}.mlp.down_proj", output_dtype=dtype)
|
|
|
|
def forward(self, hidden_states: torch.Tensor, freqs_cis=None) -> torch.Tensor:
|
|
residual = hidden_states
|
|
x = self.input_layernorm(hidden_states)
|
|
batch, sequence, _ = x.shape
|
|
query = self.q_proj(x).view(batch, sequence, self.config.num_attention_heads, self.config.head_dim).transpose(1, 2)
|
|
key = self.k_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
|
|
value = self.v_proj(x).view(batch, sequence, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
|
|
query = self.q_norm(query)
|
|
key = self.k_norm(key)
|
|
if freqs_cis is None:
|
|
query, key = _rope(query, key, self.config.rope_theta)
|
|
else:
|
|
query, key = _apply_rope(query, key, freqs_cis)
|
|
# Comfy selects its small-input SDPA path for Qwen, with an explicit causal mask.
|
|
causal_mask = torch.full(
|
|
(sequence, sequence),
|
|
torch.finfo(query.dtype).min / 4,
|
|
dtype=query.dtype,
|
|
device=query.device,
|
|
).triu_(1)
|
|
attention = _qwen_attention(query, key, value, causal_mask)
|
|
hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1))
|
|
residual = hidden_states
|
|
x = self.post_attention_layernorm(hidden_states)
|
|
return residual + self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
|
|
|
|
class Qwen3VL32BTextEncoder(nn.Module):
|
|
"""Mounted-checkpoint Qwen3-VL prompt conditioner returning layer-50 states."""
|
|
config = Qwen3VL32BTextConfig()
|
|
|
|
def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
|
|
super().__init__()
|
|
self.checkpoint_path = str(checkpoint_path)
|
|
self.device_name = str(device)
|
|
self.dtype = dtype
|
|
checkpoint = H3Checkpoint(checkpoint_path, device=device)
|
|
self._validate_checkpoint(checkpoint)
|
|
self.register_buffer("embed_tokens", checkpoint.tensor("model.embed_tokens.weight"), persistent=False)
|
|
self.register_buffer("embed_scale", checkpoint.tensor("model.embed_tokens.weight_scale", dtype=torch.float32), persistent=False)
|
|
self.layers = nn.ModuleList(
|
|
_Qwen3VLBlock(checkpoint, f"model.layers.{index}", self.config, dtype)
|
|
for index in range(self.config.num_layers)
|
|
)
|
|
|
|
@classmethod
|
|
def _validate_checkpoint(cls, checkpoint: H3Checkpoint) -> None:
|
|
from safetensors import safe_open
|
|
|
|
required = {"model.embed_tokens.weight"}
|
|
linear_names = (
|
|
"self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
|
|
"mlp.gate_proj", "mlp.up_proj", "mlp.down_proj",
|
|
)
|
|
for index in range(cls.config.num_layers):
|
|
prefix = f"model.layers.{index}"
|
|
required.update({
|
|
f"{prefix}.input_layernorm.weight",
|
|
f"{prefix}.post_attention_layernorm.weight",
|
|
f"{prefix}.self_attn.q_norm.weight",
|
|
f"{prefix}.self_attn.k_norm.weight",
|
|
})
|
|
for linear in linear_names:
|
|
required.update(
|
|
f"{prefix}.{linear}.{suffix}"
|
|
for suffix in ("comfy_quant", "weight", "weight_scale", "weight_scale_2")
|
|
)
|
|
with safe_open(checkpoint.path, framework="pt", device="cpu") as file:
|
|
names = set(file.keys())
|
|
missing = sorted(required - names)
|
|
if missing:
|
|
raise ValueError(
|
|
"Not a supported MiniMax Qwen3-VL-32B NVFP4 checkpoint; missing "
|
|
+ ", ".join(missing)
|
|
)
|
|
|
|
def _embed_rows(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
"""Scaled token embeds ``[batch, tokens, 5120]`` in fp32 (pre-decoder)."""
|
|
token_rows = F.embedding(input_ids.to(self.embed_tokens.device), self.embed_tokens).to(torch.float32)
|
|
token_scales = F.embedding(input_ids.to(self.embed_scale.device), self.embed_scale)
|
|
return (token_rows * token_scales).to(torch.float32)
|
|
|
|
def _run_layers(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
*,
|
|
position_ids: torch.Tensor | None = None,
|
|
visual_pos_masks: torch.Tensor | None = None,
|
|
deepstack_embeds: list[torch.Tensor] | None = None,
|
|
) -> torch.Tensor:
|
|
"""Run all 50 decoder blocks over an fp32 ``[batch, tokens, 5120]`` state.
|
|
|
|
With ``position_ids`` (``[3, seq]``) the Qwen3-VL interleaved mrope is used
|
|
instead of the plain 1D text rope; with ``visual_pos_masks`` +
|
|
``deepstack_embeds`` the three DeepStack features are added at the visual
|
|
positions of the first three decoder layers (Comfy ``Llama2_.forward``).
|
|
"""
|
|
hidden_states = hidden_states.to(self.dtype)
|
|
freqs_cis = None
|
|
if position_ids is not None:
|
|
from .qwen3vl_vision import mrope_freqs_cis
|
|
|
|
freqs_cis = mrope_freqs_cis(position_ids.to(hidden_states.device))
|
|
for layer_index, layer in enumerate(self.layers):
|
|
hidden_states = layer(hidden_states, freqs_cis)
|
|
if (
|
|
deepstack_embeds is not None
|
|
and visual_pos_masks is not None
|
|
and layer_index < len(deepstack_embeds)
|
|
):
|
|
mask = visual_pos_masks.to(hidden_states.device)
|
|
hidden_states[mask] = hidden_states[mask] + deepstack_embeds[layer_index].to(hidden_states)
|
|
return hidden_states
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
"""Return unnormalized `[batch, tokens, 5120]` output after decoder layer 50."""
|
|
if input_ids.ndim != 2:
|
|
raise ValueError(f"input_ids must have shape [batch, tokens], got {tuple(input_ids.shape)}")
|
|
if input_ids.numel() == 0:
|
|
raise ValueError("input_ids must contain at least one token")
|
|
return self._run_layers(self._embed_rows(input_ids))
|
|
|
|
|
|
class Qwen3VLPromptConditioner:
|
|
"""Tokenize raw H3 prompt text and produce Qwen layer-50 conditioning."""
|
|
def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
|
|
from .conditioning import H3PromptTokenizer
|
|
|
|
tokenizer_dir = tokenizer_dir or Path(__file__).with_name("qwen25_tokenizer")
|
|
self.tokenizer = H3PromptTokenizer(tokenizer_dir)
|
|
self.encoder = Qwen3VL32BTextEncoder(checkpoint_path, device=device, dtype=dtype)
|
|
self.device = device
|
|
|
|
def __call__(self, prompt: str) -> torch.Tensor:
|
|
return self.encoder(self.tokenizer(prompt, device=self.device))
|