diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 79954f1db..32061f648 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -440,29 +440,86 @@ def test_transformer_rejects_prompt_tuning(): BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True})) -def _native_to_checkpoint_name(name: str) -> str: - """Inverse of ``load_weights`` remapping: native param name -> diffusers name. +_QKV_FANOUT = { + "to_qkv": ("to_q", "to_k", "to_v"), + "img_to_qkv": ("img_to_q", "img_to_k", "img_to_v"), + "instruct_to_qkv": ("instruct_to_q", "instruct_to_k", "instruct_to_v"), +} + + +def _native_to_checkpoint_weights(name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]: + """Inverse of ``load_weights`` remapping: native param -> diffusers weights. + + The diffusers checkpoint stores fused projections as separate matrices, so a + single merged native param (QKV or FFN gate/up) fans out into several + checkpoint weights. Returns ``(diffusers_name, value)`` pairs whose values, + when fed through ``load_weights``, reassemble into the original ``param``. - ``.to_out.`` -> ``.to_out.0.`` (diffusers ModuleList wrap). - promoted joint-attention projections move back under ``.processor.``. """ + q_size = NUM_HEADS * HEAD_DIM + kv_size = NUM_KV_HEADS * HEAD_DIM + + # Fused QKV projections split back into per-matrix q/k/v weights. + for token, (q_name, k_name, v_name) in _QKV_FANOUT.items(): + marker = f".{token}." + if marker in name: + q, k, v = param[:q_size], param[q_size : q_size + kv_size], param[q_size + kv_size :] + results = [] + for sub_name, value in ((q_name, q), (k_name, k), (v_name, v)): + ckpt_name = name.replace(marker, f".{sub_name}.") + if token != "to_qkv": + # Joint-attention projections live under `.processor.` upstream. + ckpt_name = ckpt_name.replace(".img_instruct_attn.", ".img_instruct_attn.processor.") + results.append((ckpt_name, value)) + return results + + # Fused FFN gate/up splits into linear_1 (gate) / linear_3 (input). + if ".gate_up_proj." in name: + inner = param.shape[0] // 2 + return [ + (name.replace(".gate_up_proj.", ".linear_1."), param[:inner]), + (name.replace(".gate_up_proj.", ".linear_3."), param[inner:]), + ] + + # Default: 1:1 with the existing diffusers name promotions. if ".to_out." in name: name = name.replace(".to_out.", ".to_out.0.") - for proj in ( - "img_to_q", - "img_to_k", - "img_to_v", - "instruct_to_q", - "instruct_to_k", - "instruct_to_v", - "instruct_out", - "img_out", - ): + for proj in ("instruct_out", "img_out"): token = f".img_instruct_attn.{proj}." if token in name: name = name.replace(token, f".img_instruct_attn.processor.{proj}.") break - return name + return [(name, param)] + + +def test_transformer_exposes_stacked_params_mapping(): + """The packed -> sub-layer mapping must be discoverable from the module tree. + + ``diffusion/lora/loader.py`` and the quantized weight loaders read + ``stacked_params_mapping`` off the model, so it has to exist before (and + independently of) ``load_weights``. + """ + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _BOOGU_STACKED_PARAMS_MAPPING, + BooguImageTransformer2DModel, + ) + + model = BooguImageTransformer2DModel(od_config=_tiny_od_config()) + + assert tuple(model.stacked_params_mapping) == _BOOGU_STACKED_PARAMS_MAPPING + # A per-instance copy, so consumers cannot mutate the module constant. + assert model.stacked_params_mapping is not _BOOGU_STACKED_PARAMS_MAPPING + + # Every entry is a (param, shard, shard_id) triple. + for param_name, shard_name, shard_id in model.stacked_params_mapping: + assert param_name.startswith(".") and shard_name.startswith(".") + assert shard_id in {"q", "k", "v", 0, 1} + + # The mapping targets exactly the fused projections this port creates. + mapped_leaves = {param.strip(".").split(".")[-1] for param, _, _ in model.stacked_params_mapping} + assert mapped_leaves == {"to_qkv", "img_to_qkv", "instruct_to_qkv", "gate_up_proj"} def test_transformer_load_weights_round_trip(): @@ -473,24 +530,59 @@ def test_transformer_load_weights_round_trip(): model = BooguImageTransformer2DModel(od_config=_tiny_od_config()) native_params = dict(model.named_parameters()) - # Build synthetic diffusers-named weights (one per native parameter). - checkpoint_weights = {} + # Build synthetic diffusers-named weights. Fused native projections fan out + # into the separate matrices the diffusers checkpoint stores. + expected: dict[str, torch.Tensor] = {} + checkpoint_weights: dict[str, torch.Tensor] = {} for native_name, param in native_params.items(): - checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param) + full = torch.randn_like(param) + expected[native_name] = full + for ckpt_name, value in _native_to_checkpoint_weights(native_name, full): + checkpoint_weights[ckpt_name] = value - # The remapping must be a bijection over the parameter set. - assert len(checkpoint_weights) == len(native_params) + # Every checkpoint name is unique; merged params fan out to >=1 weight. + assert len(checkpoint_weights) >= len(native_params) loaded = model.load_weights(list(checkpoint_weights.items())) # No missing / unexpected parameters. assert loaded == set(native_params.keys()) - # Values landed on the right parameters (TP=1: weight_loader copies verbatim). + # Values landed on the right parameters (TP=1: weight_loader copies verbatim; + # fused shards are placed at their q/k/v / gate/up offsets). reloaded = dict(model.named_parameters()) for native_name in native_params: - expected = checkpoint_weights[_native_to_checkpoint_name(native_name)] - assert torch.allclose(reloaded[native_name], expected) + assert torch.allclose(reloaded[native_name], expected[native_name]) + + +def test_transformer_load_weights_rejects_partial_fused_params(): + """A fused parameter must not count as loaded until all its shards arrive. + + The loader compares parameter *names*, so reporting ``to_qkv`` complete + after a single ``to_q`` would let a checkpoint carrying only ``to_q`` start + up with the ``k``/``v`` slices left uninitialized. + """ + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + BooguImageTransformer2DModel, + ) + + model = BooguImageTransformer2DModel(od_config=_tiny_od_config()) + native_params = dict(model.named_parameters()) + + complete_name = "noise_refiner.0.attn.to_qkv.weight" + partial_name = "noise_refiner.0.feed_forward.gate_up_proj.weight" + assert complete_name in native_params + assert partial_name in native_params + + weights = list(_native_to_checkpoint_weights(complete_name, native_params[complete_name])) + # Only the gate half of the FFN, so its fused parameter stays incomplete. + gate_ckpt, gate_value = _native_to_checkpoint_weights(partial_name, native_params[partial_name])[0] + weights.append((gate_ckpt, gate_value)) + + loaded = model.load_weights(weights) + + assert complete_name in loaded + assert partial_name not in loaded def test_transformer_load_weights_warns_for_unexpected_and_unloaded(): @@ -513,12 +605,15 @@ def test_transformer_load_weights_warns_for_unexpected_and_unloaded(): try: model = BooguImageTransformer2DModel(od_config=_tiny_od_config()) native_params = dict(model.named_parameters()) - loaded_name = next(iter(native_params)) - checkpoint_name = _native_to_checkpoint_name(loaded_name) + # A non-fused parameter, so one checkpoint entry loads it completely. + loaded_name = "x_embedder.weight" + assert loaded_name in native_params + (checkpoint_name, checkpoint_value), *_ = _native_to_checkpoint_weights(loaded_name, native_params[loaded_name]) + checkpoint_value = torch.randn_like(checkpoint_value) loaded = model.load_weights( [ - (checkpoint_name, torch.randn_like(native_params[loaded_name])), + (checkpoint_name, checkpoint_value), ("unexpected.weight", torch.ones(1)), ] ) diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index 93d52b960..f3fdeedc8 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py @@ -15,6 +15,7 @@ # inference caches (TeaCache/TaylorSeer) are not ported. import itertools +from collections import defaultdict from collections.abc import Iterable from typing import TYPE_CHECKING @@ -24,9 +25,9 @@ import torch.nn.functional as F from diffusers.models.embeddings import Timesteps, get_1d_rotary_pos_embed from einops import rearrange, repeat from vllm.logger import init_logger -from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, ReplicatedLinear, RowParallelLinear, ) @@ -42,6 +43,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( fused_qk_norm_rope, fused_qk_norm_rope_min_tokens, ) +from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.models.utils import make_attention_mask from vllm_omni.platforms import current_omni_platform @@ -260,7 +262,14 @@ class LuminaLayerNormContinuous(nn.Module): class LuminaFeedForward(nn.Module): - """SwiGLU feed-forward with tensor-parallel projections.""" + """SwiGLU feed-forward with tensor-parallel projections. + + The ``gate`` and ``input`` projections are fused into a single + ``MergedColumnParallelLinear`` (one GEMM instead of two), matching the + ``gate_up_proj`` convention used by the Hunyuan Image 3 port. The SwiGLU + activation keeps its float32 computation for numerical parity with + upstream. + """ def __init__( self, @@ -276,20 +285,13 @@ class LuminaFeedForward(nn.Module): inner_dim = int(ffn_dim_multiplier * inner_dim) inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) - self.linear_1 = ColumnParallelLinear( - dim, - inner_dim, + self.gate_up_proj = MergedColumnParallelLinear( + input_size=dim, + output_sizes=[inner_dim, inner_dim], bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "linear_1"), - ) # gate - self.linear_3 = ColumnParallelLinear( - dim, - inner_dim, - bias=False, - quant_config=quant_config, - prefix=_join_prefix(prefix, "linear_3"), - ) # input + prefix=_join_prefix(prefix, "gate_up_proj"), + ) self.linear_2 = RowParallelLinear( inner_dim, dim, @@ -299,8 +301,8 @@ class LuminaFeedForward(nn.Module): ) def forward(self, x: torch.Tensor) -> torch.Tensor: - h1, _ = self.linear_1(x) - h2, _ = self.linear_3(x) + gate_up, _ = self.gate_up_proj(x) + h1, h2 = gate_up.chunk(2, dim=-1) out, _ = self.linear_2(swiglu(h1, h2)) return out @@ -593,28 +595,15 @@ class BooguImageSelfAttention(nn.Module): ) -> None: super().__init__() self.head_dim = dim // num_attention_heads - kv_dim = self.head_dim * num_kv_heads - self.to_q = ColumnParallelLinear( - dim, - dim, + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_attention_heads, + total_num_kv_heads=num_kv_heads, bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "to_q"), - ) - self.to_k = ColumnParallelLinear( - dim, - kv_dim, - bias=False, - quant_config=quant_config, - prefix=_join_prefix(prefix, "to_k"), - ) - self.to_v = ColumnParallelLinear( - dim, - kv_dim, - bias=False, - quant_config=quant_config, - prefix=_join_prefix(prefix, "to_v"), + prefix=_join_prefix(prefix, "to_qkv"), ) self.norm_q = RMSNorm(self.head_dim, eps=1e-5) self.norm_k = RMSNorm(self.head_dim, eps=1e-5) @@ -626,8 +615,10 @@ class BooguImageSelfAttention(nn.Module): prefix=_join_prefix(prefix, "to_out"), ) - self.num_local_heads = self.to_q.output_size_per_partition // self.head_dim - self.num_local_kv_heads = self.to_k.output_size_per_partition // self.head_dim + self.num_local_heads = self.to_qkv.num_heads + self.num_local_kv_heads = self.to_qkv.num_kv_heads + self.q_size = self.num_local_heads * self.head_dim + self.kv_size = self.num_local_kv_heads * self.head_dim self.attn = Attention( num_heads=self.num_local_heads, @@ -645,9 +636,8 @@ class BooguImageSelfAttention(nn.Module): ) -> torch.Tensor: dtype = hidden_states.dtype - query, _ = self.to_q(hidden_states) - key, _ = self.to_k(hidden_states) - value, _ = self.to_v(hidden_states) + qkv, _ = self.to_qkv(hidden_states) + query, key, value = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) query = query.unflatten(-1, (self.num_local_heads, self.head_dim)) key = key.unflatten(-1, (self.num_local_kv_heads, self.head_dim)) @@ -683,50 +673,24 @@ class BooguImageJointAttention(nn.Module): ) -> None: super().__init__() self.head_dim = dim // num_attention_heads - kv_dim = self.head_dim * num_kv_heads - self.img_to_q = ColumnParallelLinear( - dim, - dim, + self.img_to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_attention_heads, + total_num_kv_heads=num_kv_heads, bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "img_to_q"), + prefix=_join_prefix(prefix, "img_to_qkv"), ) - self.img_to_k = ColumnParallelLinear( - dim, - kv_dim, + self.instruct_to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_attention_heads, + total_num_kv_heads=num_kv_heads, bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "img_to_k"), - ) - self.img_to_v = ColumnParallelLinear( - dim, - kv_dim, - bias=False, - quant_config=quant_config, - prefix=_join_prefix(prefix, "img_to_v"), - ) - - self.instruct_to_q = ColumnParallelLinear( - dim, - dim, - bias=False, - quant_config=quant_config, - prefix=_join_prefix(prefix, "instruct_to_q"), - ) - self.instruct_to_k = ColumnParallelLinear( - dim, - kv_dim, - bias=False, - quant_config=quant_config, - prefix=_join_prefix(prefix, "instruct_to_k"), - ) - self.instruct_to_v = ColumnParallelLinear( - dim, - kv_dim, - bias=False, - quant_config=quant_config, - prefix=_join_prefix(prefix, "instruct_to_v"), + prefix=_join_prefix(prefix, "instruct_to_qkv"), ) self.norm_q = RMSNorm(self.head_dim, eps=1e-5) @@ -757,8 +721,10 @@ class BooguImageJointAttention(nn.Module): prefix=_join_prefix(prefix, "to_out"), ) - self.num_local_heads = self.img_to_q.output_size_per_partition // self.head_dim - self.num_local_kv_heads = self.img_to_k.output_size_per_partition // self.head_dim + self.num_local_heads = self.img_to_qkv.num_heads + self.num_local_kv_heads = self.img_to_qkv.num_kv_heads + self.q_size = self.num_local_heads * self.head_dim + self.kv_size = self.num_local_kv_heads * self.head_dim self.attn = Attention( num_heads=self.num_local_heads, @@ -780,13 +746,13 @@ class BooguImageJointAttention(nn.Module): dtype = img_hidden_states.dtype batch_size = img_hidden_states.shape[0] - img_query, _ = self.img_to_q(img_hidden_states) - img_key, _ = self.img_to_k(img_hidden_states) - img_value, _ = self.img_to_v(img_hidden_states) + img_qkv, _ = self.img_to_qkv(img_hidden_states) + img_query, img_key, img_value = img_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - instruct_query, _ = self.instruct_to_q(instruct_hidden_states) - instruct_key, _ = self.instruct_to_k(instruct_hidden_states) - instruct_value, _ = self.instruct_to_v(instruct_hidden_states) + instruct_qkv, _ = self.instruct_to_qkv(instruct_hidden_states) + instruct_query, instruct_key, instruct_value = instruct_qkv.split( + [self.q_size, self.kv_size, self.kv_size], dim=-1 + ) query, key, value = _concat_instruction_image_features( [img_query, img_key, img_value], @@ -1135,6 +1101,44 @@ def _cal_preprocessed_instruction_feat_dim(instruction_feature_configs: dict) -> raise ValueError(f"Invalid reduce_type: {reduce_type}") +# Packed (fused) projections vs. the logical sub-projections the diffusers +# checkpoint stores, in the ``(param_name, shard_name, shard_id)`` form used +# across vllm-omni. ``load_weights`` consumes it directly, and loader consumers +# that discover it via ``stacked_params_mapping`` (LoRA, quantized loaders) get +# the packed -> sub-layer relationship for free. +# +# The QKV entries stay leaf-scoped: each fused source (``to_q`` on the +# self-attention module, ``img_to_q`` / ``instruct_to_q`` on the joint +# attention) has a distinct leaf, and ``.to_q.`` cannot collide with +# ``.img_to_q.`` because the preceding character is ``_`` rather than ``.``. +# The FFN entries are path-qualified instead: ``linear_1`` also names the +# timestep embedder and the ``norm_out`` projections, which must load 1:1 and +# are not fused. +_BOOGU_STACKED_PARAMS_MAPPING = ( + # self-attention: noise / reference-image / context refiners, the + # single-stream blocks, and the double-stream image self-attention. + (".to_qkv.", ".to_q.", "q"), + (".to_qkv.", ".to_k.", "k"), + (".to_qkv.", ".to_v.", "v"), + # joint (instruction + image) attention of the double-stream blocks. + (".img_to_qkv.", ".img_to_q.", "q"), + (".img_to_qkv.", ".img_to_k.", "k"), + (".img_to_qkv.", ".img_to_v.", "v"), + (".instruct_to_qkv.", ".instruct_to_q.", "q"), + (".instruct_to_qkv.", ".instruct_to_k.", "k"), + (".instruct_to_qkv.", ".instruct_to_v.", "v"), + # feed-forward gate/up: ``linear_1`` is the gate (shard 0) and + # ``linear_3`` the input (shard 1), matching the packed order the SwiGLU + # activation consumes. + (".feed_forward.gate_up_proj.", ".feed_forward.linear_1.", 0), + (".feed_forward.gate_up_proj.", ".feed_forward.linear_3.", 1), + (".img_feed_forward.gate_up_proj.", ".img_feed_forward.linear_1.", 0), + (".img_feed_forward.gate_up_proj.", ".img_feed_forward.linear_3.", 1), + (".instruct_feed_forward.gate_up_proj.", ".instruct_feed_forward.linear_1.", 0), + (".instruct_feed_forward.gate_up_proj.", ".instruct_feed_forward.linear_3.", 1), +) + + class BooguImageTransformer2DModel(nn.Module): """Boogu-Image transformer with mixed stream topology. @@ -1160,6 +1164,10 @@ class BooguImageTransformer2DModel(nn.Module): prefix: str = "", ) -> None: super().__init__() + # Published here rather than in ``load_weights`` so loader consumers + # that only inspect the module tree (LoRA discovery, quantized weight + # loaders) always see the packed -> sub-layer relationship. + self.stacked_params_mapping = list(_BOOGU_STACKED_PARAMS_MAPPING) self.od_config = od_config cfg = od_config.tf_model_config @@ -1710,7 +1718,7 @@ class BooguImageTransformer2DModel(nn.Module): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load diffusers-named checkpoint weights into the native module. - Two name promotions relative to upstream (see step 8/10 findings): + Name promotions relative to upstream (see step 8/10 findings): - ``*.img_instruct_attn.processor.{img,instruct}_{to_q,to_k,to_v}`` / ``{instruct,img}_out`` -> drop ``.processor`` (upstream keeps the @@ -1719,9 +1727,29 @@ class BooguImageTransformer2DModel(nn.Module): - ``*.to_out.0.weight`` -> ``*.to_out.weight`` (diffusers wraps the output projection in a ``ModuleList``; the native module uses a plain linear). + + Packed Q/K/V and FFN gate/input matrices are folded onto the fused + ``QKVParallelLinear`` / ``MergedColumnParallelLinear`` parameters using + :attr:`stacked_params_mapping` — the same mapping loader consumers such + as LoRA discovery and quantized weight loaders read. + + A fused parameter is only reported as loaded once *every* one of its + source matrices has arrived. The caller compares parameter names, so + reporting e.g. ``to_qkv`` complete after the first shard would let a + checkpoint carrying only ``to_q`` start up with the ``k``/``v`` slices + left uninitialized. """ + stacked_params_mapping = self.stacked_params_mapping + # The shard ids each fused parameter is assembled from, keyed by the + # mapping entry that produces it. + expected_shards: dict[str, set[str | int]] = defaultdict(set) + for param_name, _weight_name, packed_shard_id in stacked_params_mapping: + expected_shards[param_name].add(packed_shard_id) + params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + # Fused parameter -> (mapping entry that produced it, shards received). + fused_params: dict[str, tuple[str, set[str | int]]] = {} for name, loaded_weight in weights: original_name = name @@ -1730,15 +1758,36 @@ class BooguImageTransformer2DModel(nn.Module): if ".to_out.0." in name: name = name.replace(".to_out.0.", ".to_out.") + shard_id: str | int | None = None + matched_entry: str | None = None + for param_name, weight_name, packed_shard_id in stacked_params_mapping: + if weight_name in name: + name = name.replace(weight_name, param_name) + shard_id = packed_shard_id + matched_entry = param_name + break + if name not in params_dict: logger.warning("Skipping unexpected checkpoint weight %s", original_name) continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) + if matched_entry is None: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + _, received = fused_params.setdefault(name, (matched_entry, set())) + received.add(shard_id) loaded_params.add(name) + # Withdraw any fused parameter that only received part of its source + # matrices, so the caller's "not initialized from checkpoint" check + # reports it as missing instead of accepting a half-filled parameter. + for name, (matched_entry, received) in fused_params.items(): + if received != expected_shards[matched_entry]: + loaded_params.discard(name) + unloaded_params = sorted(params_dict.keys() - loaded_params) if unloaded_params: logger.warning("Model parameters not loaded from checkpoint: %s", unloaded_params)