diff --git a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py index 6356cbd05..2589d8985 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py @@ -34,6 +34,8 @@ def _build_cosmos3_param_names_mapping(gated_mlp: bool = True) -> dict: unchanged. """ mapping = { + # PEFT saves adapter keys as .lora_{A,B}.default.weight. + r"^(.*\.lora_[AB])\.default$": r"\1", # Inherited from text pretraining; unused at diffusion inference. r"^lm_head\.weight$": "", r"^norm\.weight$": "", diff --git a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py index 490a48877..d97f4c98c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py @@ -11,6 +11,7 @@ import os from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) +from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import ( Cosmos3DecodingStage, Cosmos3DenoisingStage, @@ -25,7 +26,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) -class Cosmos3Pipeline(ComposedPipelineBase): +class Cosmos3Pipeline(LoRAPipeline, ComposedPipelineBase): """Cosmos3 diffusion pipeline shared by T2V, I2V, and T2I. Text is tokenized and embedded directly inside the transformer; there is diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index 300dd649e..2dbcd7833 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py @@ -55,6 +55,60 @@ def _swap_peft_swiglu_fc1_lora_b( return torch.cat([gate, value], dim=0) +def stack_or_compose_fused_lora( + a_list: list[torch.Tensor], + b_list: list[torch.Tensor], + adapter_alpha: int | None, +) -> tuple[torch.Tensor, torch.Tensor, int | None]: + """Equal sections stay 3D-stacked; GQA sections become one 2D pair.""" + if len({t.shape for t in a_list}) == 1 and len({t.shape for t in b_list}) == 1: + return torch.stack(a_list), torch.stack(b_list), None + ranks = [a.shape[0] for a in a_list] + outs = [b.shape[0] for b in b_list] + a_2d = torch.cat(a_list, dim=0) + b_2d = b_list[0].new_zeros(sum(outs), sum(ranks)) + row = col = 0 + for a, b, rank, out in zip(a_list, b_list, ranks, outs): + scale = 1.0 if adapter_alpha is None else adapter_alpha / rank + b_2d[row : row + out, col : col + rank] = ( + b if scale == 1.0 else (b.float() * scale).to(dtype=b.dtype) + ) + row += out + col += rank + return a_2d, b_2d, a_2d.shape[0] + + +def _store_fused_lora_groups( + adapter: dict[str, torch.Tensor], + to_merge_params: dict[Hashable, dict[Any, Any]], + adapter_alpha: int | None, + device: torch.device | str, +) -> None: + """Write deferred fused lora_A/B groups into the adapter dict.""" + for a_key, a_parts in list(to_merge_params.items()): + if not str(a_key).endswith(".lora_A"): + continue + base = str(a_key)[: -len(".lora_A")] + b_key = f"{base}.lora_B" + b_parts = to_merge_params.get(b_key) + n = max(a_parts) + 1 + if ( + b_parts is None + or set(a_parts) != set(range(n)) + or set(b_parts) != set(range(n)) + ): + continue + a, b, fused_alpha = stack_or_compose_fused_lora( + [a_parts[i] for i in range(n)], + [b_parts[i] for i in range(n)], + adapter_alpha, + ) + adapter[str(a_key)] = a.to(device) + adapter[b_key] = b.to(device) + if fused_alpha is not None: + adapter[f"{base}.alpha"] = torch.tensor(float(fused_alpha), device=device) + + class LoRAPipeline(ComposedPipelineBase): """ Pipeline that supports injecting LoRA adapters into the diffusion transformer. @@ -780,6 +834,9 @@ class LoRAPipeline(ComposedPipelineBase): # see param mapping in HunyuanVideoArchConfig if merge_index is not None: to_merge_params[target_name][merge_index] = weight + # A/B of one fused layer must be laid out together (GQA B cannot stack). + if target_name.endswith((".lora_A", ".lora_B")): + continue if len(to_merge_params[target_name]) == num_params_to_merge: sorted_tensors = [ to_merge_params[target_name][i] @@ -798,6 +855,13 @@ class LoRAPipeline(ComposedPipelineBase): ) self.lora_adapters[lora_nickname][target_name] = weight.to(self.device) + _store_fused_lora_groups( + self.lora_adapters[lora_nickname], + to_merge_params, + adapter_lora_alpha, + self.device, + ) + self.loaded_adapter_paths[lora_nickname] = lora_path self.loaded_adapter_alphas[lora_nickname] = adapter_lora_alpha logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py index 7c470d0e4..fcc50e1d6 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -58,7 +58,10 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im is_layerwise_offloaded_module, ) from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline -from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline +from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import ( + LoRAPipeline, + stack_or_compose_fused_lora, +) from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.weight_sync.tensor_bucket import ( @@ -215,11 +218,14 @@ def _build_module_weight_name_mapper(module: torch.nn.Module): if not mapping_fns: return None - def map_name(name: str) -> str: + def map_name(name: str) -> tuple[str, Any]: mapped_name = name + merge_index = None for mapping_fn in mapping_fns: - mapped_name = mapping_fn(mapped_name)[0] - return mapped_name + mapped_name, index, _ = mapping_fn(mapped_name) + if index is not None: + merge_index = index + return mapped_name, merge_index return map_name @@ -236,23 +242,24 @@ def _resolve_lora_ipc_layer_dict_key( layer_prefix: str, layer_dict: dict, module: torch.nn.Module, -) -> tuple[Any | None, str]: +) -> tuple[Any | None, str, int | None]: """Map training-side LoRA layer prefix to lora_layers key (Layer 2).""" layer = layer_dict.get(layer_prefix) if layer is not None: - return layer, layer_prefix + return layer, layer_prefix, None map_name = _build_module_weight_name_mapper(module) if map_name is None: - return None, layer_prefix + return None, layer_prefix, None - mapped = _strip_param_weight_suffix(map_name(f"{layer_prefix}.weight")) + mapped_name, merge_index = map_name(f"{layer_prefix}.weight") + mapped = _strip_param_weight_suffix(mapped_name) if mapped != layer_prefix: layer = layer_dict.get(mapped) if layer is not None: - return layer, mapped + return layer, mapped, merge_index - return None, layer_prefix + return None, layer_prefix, None def _iter_module_weight_updates( @@ -268,7 +275,7 @@ def _iter_module_weight_updates( yield name, loaded_weight continue - mapped_name = map_name(name) if map_name is not None else name + mapped_name = map_name(name)[0] if map_name is not None else name if mapped_name in model_params: yield mapped_name, loaded_weight continue @@ -628,35 +635,72 @@ class WeightsUpdater: updated = 0 skipped = 0 unknown_layers: list[str] = [] - with lora_pipeline._temporarily_disable_offload(target=target_module): - for layer_name, (lora_a, lora_b) in pairs.items(): - layer, _resolved_key = _resolve_lora_ipc_layer_dict_key( - layer_name, layer_dict, dit_module + # Honor lora_merge_mode: merged and unmerged evaluation differ bitwise. + merge_mode = lora_pipeline._resolve_lora_merge_mode(None, None) + merge_weights = lora_pipeline._should_merge_lora_for_layers( + target_module, layer_dict, merge_mode + ) + plain_pairs: list[tuple[torch.Tensor, torch.Tensor, Any]] = [] + fused_sections: dict[Any, dict[int, tuple[torch.Tensor, torch.Tensor]]] = {} + for layer_name, (lora_a, lora_b) in pairs.items(): + layer, _resolved_key, merge_index = _resolve_lora_ipc_layer_dict_key( + layer_name, layer_dict, dit_module + ) + if layer is None: + logger.warning( + "Unknown LoRA layer name %s for target %s; skipping", + layer_name, + target_module, ) - if layer is None: - logger.warning( - "Unknown LoRA layer name %s for target %s; skipping", - layer_name, - target_module, - ) - unknown_layers.append(layer_name) - skipped += 1 - continue - inferred_rank = int(lora_a.shape[0]) - alpha = lora_alpha if lora_alpha is not None else inferred_rank - if lora_rank is not None and lora_rank != inferred_rank: - logger.warning( - "LoRA rank mismatch for %s: payload=%d request=%d; using payload rank", - layer_name, - inferred_rank, - lora_rank, - ) + unknown_layers.append(layer_name) + skipped += 1 + continue + inferred_rank = int(lora_a.shape[-2]) + if lora_rank is not None and lora_rank != inferred_rank: + logger.warning( + "LoRA rank mismatch for %s: payload=%d request=%d; using payload rank", + layer_name, + inferred_rank, + lora_rank, + ) + if merge_index is None: + plain_pairs.append((lora_a, lora_b, layer)) + else: + # Per-section pairs of one fused layer compose into one adapter. + fused_sections.setdefault(layer, {})[merge_index] = (lora_a, lora_b) + + with lora_pipeline._temporarily_disable_offload(target=target_module): + for lora_a, lora_b, layer in plain_pairs: + inferred_rank = int(lora_a.shape[-2]) layer.lora_rank = inferred_rank - layer.lora_alpha = alpha + layer.lora_alpha = ( + lora_alpha if lora_alpha is not None else inferred_rank + ) layer.set_lora_weights( - lora_a, lora_b, merge_weights=True, clear_existing=True + lora_a, lora_b, merge_weights=merge_weights, clear_existing=True ) updated += 1 + for layer, layer_sections in fused_sections.items(): + indices = sorted(layer_sections) + lora_a, lora_b, fused_alpha = stack_or_compose_fused_lora( + [layer_sections[i][0] for i in indices], + [layer_sections[i][1] for i in indices], + lora_alpha, + ) + if fused_alpha is not None: + # Composed pairs fold the scale; alpha == rank keeps it neutral. + layer.lora_rank = fused_alpha + layer.lora_alpha = fused_alpha + else: + inferred_rank = int(lora_a.shape[-2]) + layer.lora_rank = inferred_rank + layer.lora_alpha = ( + lora_alpha if lora_alpha is not None else inferred_rank + ) + layer.set_lora_weights( + lora_a, lora_b, merge_weights=merge_weights, clear_existing=True + ) + updated += len(layer_sections) gc.collect() torch.cuda.empty_cache() diff --git a/python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py b/python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py new file mode 100644 index 000000000..c6b7aa47a --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused-layer LoRA groups: stack-vs-compose decision, loader wiring, IPC resolution.""" + +from collections import defaultdict +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.nn.functional as F + +from sglang.multimodal_gen.runtime.layers.lora.linear import ( + MergedColumnParallelLinearWithLoRA, +) +from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import ( + LoRAPipeline, + stack_or_compose_fused_lora, +) +from sglang.multimodal_gen.runtime.post_training.weights_updater import ( + _resolve_lora_ipc_layer_dict_key, +) + +_TP_RANK_PATCH = "sglang.multimodal_gen.runtime.layers.lora.linear.get_tp_rank" +_LOCAL_DEVICE_PATCH = ( + "sglang.multimodal_gen.runtime.layers.lora.linear.get_local_torch_device" +) + +# GQA-like fused layout: unequal q/k/v sections. +OUTPUT_SIZES = [8, 2, 2] +IN_DIM = 4 + + +def _make_ab_lists(ranks: list[int]): + a_list, b_list = [], [] + for index, rank in enumerate(ranks): + torch.manual_seed(100 + index * 10 + rank) + a_list.append(torch.randn(rank, IN_DIM)) + b_list.append(torch.randn(OUTPUT_SIZES[index], rank)) + return a_list, b_list + + +def _reference_delta(x, a_list, b_list, adapter_alpha): + out = torch.zeros(*x.shape[:-1], sum(OUTPUT_SIZES)) + row = 0 + for a, b in zip(a_list, b_list): + scale = 1.0 if adapter_alpha is None else adapter_alpha / a.shape[0] + out[..., row : row + b.shape[0]] = (x @ a.T @ b.T) * scale + row += b.shape[0] + return out + + +def test_compose_unequal_sections_matches_reference(): + a_list, b_list = _make_ab_lists([2, 3, 1]) + a_2d, b_2d, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, 4) + assert a_2d.shape == (6, IN_DIM) + assert b_2d.shape == (sum(OUTPUT_SIZES), 6) + assert fused_alpha == 6 + + x = torch.randn(5, IN_DIM) + torch.testing.assert_close( + x @ a_2d.T @ b_2d.T, + _reference_delta(x, a_list, b_list, 4), + rtol=1e-5, + atol=1e-5, + ) + + +def test_stack_kept_for_equal_sections(): + torch.manual_seed(0) + a_list = [torch.randn(2, IN_DIM) for _ in range(2)] + b_list = [torch.randn(4, 2) for _ in range(2)] + a, b, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, 4) + assert a.shape == (2, 2, IN_DIM) + assert b.shape == (2, 4, 2) + assert fused_alpha is None + torch.testing.assert_close(a[0], a_list[0]) + torch.testing.assert_close(b[1], b_list[1]) + + +class _FakeMergedLinear(torch.nn.Module): + def __init__( + self, + output_sizes: list[int], + in_dim: int, + weight: torch.Tensor | None = None, + output_partition_sizes: list[int] | None = None, + ): + super().__init__() + self.output_sizes = output_sizes + self.output_partition_sizes = ( + output_partition_sizes + if output_partition_sizes is not None + else output_sizes + ) + if weight is None: + weight = torch.randn(sum(output_sizes), in_dim) + self.weight = torch.nn.Parameter(weight) + self.bias = None + self.skip_bias_add = False + self.gather_output = False + self.quant_method = SimpleNamespace( + apply=lambda layer, x, bias=None: F.linear(x, layer.weight, bias) + ) + + +def _make_layer() -> MergedColumnParallelLinearWithLoRA: + torch.manual_seed(0) + return MergedColumnParallelLinearWithLoRA(_FakeMergedLinear(OUTPUT_SIZES, IN_DIM)) + + +def _set_composed(layer, a_list, b_list, adapter_alpha, merge_weights=False): + a_2d, b_2d, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, adapter_alpha) + layer.lora_rank = fused_alpha + layer.lora_alpha = fused_alpha + layer.set_lora_weights(a_2d, b_2d, merge_weights=merge_weights) + + +def test_unmerged_forward_applies_composed_pair(): + layer = _make_layer() + base_weight = layer.base_layer.weight.detach().clone() + a_list, b_list = _make_ab_lists([2, 3, 1]) + _set_composed(layer, a_list, b_list, 4) + + assert not layer.merged + assert not layer.disable_lora + + x = torch.randn(5, IN_DIM) + with patch(_TP_RANK_PATCH, return_value=0): + out, _ = layer.forward(x) + expected = x @ base_weight.T + _reference_delta(x, a_list, b_list, 4) + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) + + +def test_merge_writes_section_rows_and_unmerge_restores(): + layer = _make_layer() + base_weight = layer.base_layer.weight.detach().clone() + a_list, b_list = _make_ab_lists([2, 2, 2]) + with ( + patch(_TP_RANK_PATCH, return_value=0), + patch(_LOCAL_DEVICE_PATCH, return_value=torch.device("cpu")), + ): + _set_composed(layer, a_list, b_list, 4, merge_weights=True) + + assert layer.merged + merged = layer.base_layer.weight.detach().cpu() + row = 0 + for a, b in zip(a_list, b_list): + expected_rows = base_weight[row : row + b.shape[0]] + (4 / a.shape[0]) * (b @ a) + torch.testing.assert_close( + merged[row : row + b.shape[0]], expected_rows, rtol=1e-5, atol=1e-5 + ) + row += b.shape[0] + + with patch(_TP_RANK_PATCH, return_value=0): + layer.unmerge_lora_weights() + torch.testing.assert_close(layer.base_layer.weight.detach().cpu(), base_weight) + + +def test_composed_pair_shards_correctly_under_mock_tp(): + tp_size = 2 + a_list, b_list = _make_ab_lists([2, 3, 1]) + a_2d, b_2d, fused_alpha = stack_or_compose_fused_lora(a_list, b_list, 4) + torch.manual_seed(0) + full_weight = torch.randn(sum(OUTPUT_SIZES), IN_DIM) + x = torch.randn(5, IN_DIM) + expected = x @ full_weight.T + _reference_delta(x, a_list, b_list, 4) + + part_sizes = [size // tp_size for size in OUTPUT_SIZES] + rank_outputs = [] + for tp_rank in range(tp_size): + local_rows = [] + for index, (size, part) in enumerate(zip(OUTPUT_SIZES, part_sizes)): + row = sum(OUTPUT_SIZES[:index]) + tp_rank * part + local_rows.append(full_weight[row : row + part]) + layer = MergedColumnParallelLinearWithLoRA( + _FakeMergedLinear( + OUTPUT_SIZES, + IN_DIM, + weight=torch.cat(local_rows, dim=0), + output_partition_sizes=part_sizes, + ) + ) + layer.lora_rank = fused_alpha + layer.lora_alpha = fused_alpha + layer.set_lora_weights(a_2d, b_2d, merge_weights=False) + with patch(_TP_RANK_PATCH, return_value=tp_rank): + out, _ = layer.forward(x) + rank_outputs.append(out) + + # All-gather equivalent: stitch each section's rank slices back together. + stitched = [] + for index, part in enumerate(part_sizes): + local_col = sum(part_sizes[:index]) + stitched.append( + torch.cat( + [out[..., local_col : local_col + part] for out in rank_outputs], + dim=-1, + ) + ) + torch.testing.assert_close( + torch.cat(stitched, dim=-1), expected, rtol=1e-5, atol=1e-5 + ) + + +def test_ipc_key_resolution_returns_fused_merge_index(): + module = torch.nn.Module() + module.param_names_mapping = { + r"^attn\.q\.(.*)$": (r"attn.to_qkv.\1", 0, 3), + r"^attn\.k\.(.*)$": (r"attn.to_qkv.\1", 1, 3), + r"^attn\.v\.(.*)$": (r"attn.to_qkv.\1", 2, 3), + r"^proj\.(.*)$": r"out_proj.\1", + } + sentinel = object() + layer_dict = {"attn.to_qkv": sentinel, "out_proj": sentinel} + + layer, key, merge_index = _resolve_lora_ipc_layer_dict_key( + "attn.k", layer_dict, module + ) + assert layer is sentinel + assert key == "attn.to_qkv" + assert merge_index == 1 + + layer, key, merge_index = _resolve_lora_ipc_layer_dict_key( + "proj", layer_dict, module + ) + assert layer is sentinel + assert key == "out_proj" + assert merge_index is None + + layer, key, merge_index = _resolve_lora_ipc_layer_dict_key( + "attn.to_qkv", layer_dict, module + ) + assert layer is sentinel + assert merge_index is None + + +class _TestLoRAPipeline(LoRAPipeline): + def create_pipeline_stages(self, server_args): + return None + + +_QKV_MAPPING = { + r"^attn\.q\.(.*)$": (r"attn.to_qkv.\1", 0, 3), + r"^attn\.k\.(.*)$": (r"attn.to_qkv.\1", 1, 3), + r"^attn\.v\.(.*)$": (r"attn.to_qkv.\1", 2, 3), + r"^mlp\.gate\.(.*)$": (r"mlp.gate_up.\1", 0, 2), + r"^mlp\.up\.(.*)$": (r"mlp.gate_up.\1", 1, 2), +} + + +def _make_loader_pipeline() -> _TestLoRAPipeline: + pipeline = object.__new__(_TestLoRAPipeline) + pipeline.lora_adapters = defaultdict(dict) + pipeline.loaded_adapter_paths = {} + pipeline.loaded_adapter_alphas = {} + pipeline.device = "cpu" + pipeline.modules = {"transformer": torch.nn.Module()} + pipeline.server_args = SimpleNamespace( + lora_path=None, + lora_weight_name=None, + pipeline_config=SimpleNamespace( + dit_config=SimpleNamespace( + arch_config=SimpleNamespace( + param_names_mapping=_QKV_MAPPING, + lora_param_names_mapping={ + r"^(.*\.lora_[AB])\.default$": r"\1", + }, + ) + ) + ), + ) + return pipeline + + +def _load_adapter(pipeline, state_dict, lora_alpha=None): + loader_mod = "sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline" + with ( + patch(f"{loader_mod}.maybe_download_lora", return_value="/adapter"), + patch(f"{loader_mod}.load_file", return_value=state_dict), + patch(f"{loader_mod}.dist.is_initialized", return_value=False), + ): + pipeline.load_lora_adapter("/adapter", "adapter", rank=0, lora_alpha=lora_alpha) + + +def _qkv_state_dict(rank=2): + torch.manual_seed(3) + state = {} + for name, rows in (("q", 8), ("k", 2), ("v", 2)): + state[f"attn.{name}.lora_A.default.weight"] = torch.randn(rank, IN_DIM) + state[f"attn.{name}.lora_B.default.weight"] = torch.randn(rows, rank) + return state + + +def test_loader_composes_unequal_fused_sections(): + pipeline = _make_loader_pipeline() + _load_adapter(pipeline, _qkv_state_dict(), lora_alpha=4) + + adapter = pipeline.lora_adapters["adapter"] + assert adapter["attn.to_qkv.lora_A"].shape == (6, IN_DIM) + assert adapter["attn.to_qkv.lora_B"].shape == (sum(OUTPUT_SIZES), 6) + assert adapter["attn.to_qkv.alpha"].item() == 6.0 + assert set(adapter) == { + "attn.to_qkv.lora_A", + "attn.to_qkv.lora_B", + "attn.to_qkv.alpha", + } + + +def test_loader_stacks_equal_fused_sections(): + pipeline = _make_loader_pipeline() + state_dict = { + "mlp.gate.lora_A.weight": torch.randn(2, IN_DIM), + "mlp.gate.lora_B.weight": torch.randn(4, 2), + "mlp.up.lora_A.weight": torch.randn(2, IN_DIM), + "mlp.up.lora_B.weight": torch.randn(4, 2), + } + _load_adapter(pipeline, state_dict) + + adapter = pipeline.lora_adapters["adapter"] + assert adapter["mlp.gate_up.lora_A"].shape == (2, 2, IN_DIM) + assert adapter["mlp.gate_up.lora_B"].shape == (2, 4, 2) + assert "mlp.gate_up.alpha" not in adapter + + +def test_loader_drops_incomplete_fused_groups(): + pipeline = _make_loader_pipeline() + state_dict = _qkv_state_dict() + for key in list(state_dict): + if ".k." in key: + del state_dict[key] + _load_adapter(pipeline, state_dict) + + assert pipeline.lora_adapters["adapter"] == {} + + +def test_apply_composed_adapter_end_to_end(): + layer = _make_layer() + base_weight = layer.base_layer.weight.detach().clone() + pipeline = _make_loader_pipeline() + + rank = 2 + adapter_alpha = 4 + state_dict = _qkv_state_dict(rank=rank) + _load_adapter(pipeline, state_dict, lora_alpha=adapter_alpha) + + strength = 2.0 + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline.dist.get_rank", + return_value=0, + ): + applied = pipeline._apply_lora_to_layers( + {"attn.to_qkv": layer}, + ["adapter"], + ["/adapter"], + rank=0, + strengths=[strength], + merge_weights=False, + ) + + assert applied == 1 + a_list = [ + state_dict[f"attn.{name}.lora_A.default.weight"] for name in ("q", "k", "v") + ] + b_list = [ + state_dict[f"attn.{name}.lora_B.default.weight"] for name in ("q", "k", "v") + ] + x = torch.randn(3, IN_DIM) + with patch(_TP_RANK_PATCH, return_value=0): + out, _ = layer.forward(x) + expected = ( + x @ base_weight.T + + _reference_delta(x, a_list, b_list, adapter_alpha) * strength + ) + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5)