[SamitHuang]

Fixes #6944

Purpose

When serving diffusion models (such as Wan2.2-T2V-A14B-Diffusers) with --lora-backend distill and --use-hsdp, the service failed during LoRA weight fusion:

• On GPU (CUDA / PyTorch FSDP2): Fails with RuntimeError: aten.add_.Tensor got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!.
• On Ascend NPU: In-place mutation on HSDP-sharded memory buffers caused glibc heap corruption and fatal crash: corrupted size vs. prev_size.

Root Cause

Distilled LoRA fusion was previously performed in DiffusionWorker.init_lora_manager() after load_model() had completed. Under HSDP, _load_model_with_hsdp already transformed the transformer block parameters into sharded DTensors. When load_lora_into_module executed params.add_(delta) using regular un-sharded torch.Tensor deltas, in-place addition on DTensors failed or corrupted sharded memory.

Solution

1. Pre-sharding LoRA Fusion in Loader: Added _maybe_fuse_distilled_lora in DiffusersPipelineLoader to fuse distilled LoRA weights immediately after base weights are loaded and before HSDP / FSDP2 sharding (apply_hsdp_to_model) or post-load quantization runs.
2. Track Fusion State: Added lora_is_fused property to LoraLoaderMixin so pipelines track fusion state, and updated DiffusionWorker.init_lora_manager() to skip re-fusing if already fused at load time.
3. Unit Tests: Added regression tests covering distilled LoRA fusion order under HSDP, multi-file distilled LoRA for dual-transformer Wan2.2 MoE, non-HSDP loader fusion, lora_is_fused tracking, and worker lifecycle.

Test Plan

Added unit tests in:

• tests/diffusion/model_loader/test_diffusers_loader.py
  • test_hsdp_fuses_distilled_lora_before_sharding
  • test_hsdp_fuses_multi_file_distilled_lora_wan22
  • test_load_model_fuses_distilled_lora_non_hsdp
• tests/diffusion/lora/test_loader.py
  • test_lora_loader_mixin_lora_is_fused_property
• tests/diffusion/test_diffusion_worker.py
  • TestDiffusionWorkerInitLoraManager

vLLM Version: 0.28.0

vLLM-Omni Commit: 629d82e

Test Result

============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/public/yx/vllm-omni
configfile: pyproject.toml
plugins: asyncio-1.4.0, mock-3.15.1, rerunfailures-16.4, anyio-4.14.1, hydra-core-1.3.3
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 106 items

tests/diffusion/lora/test_loader.py .................                    [ 16%]
tests/diffusion/model_loader/test_diffusers_loader.py .................. [ 33%]
..............                                                           [ 46%]
tests/diffusion/test_diffusion_worker.py ..........                      [ 55%]
tests/diffusion/models/sensenova_u1/test_sensenova_u1_distill_lora.py .. [ 57%]
......                                                                   [ 63%]
tests/entrypoints/test_async_omni_diffusion_config.py .................. [ 80%]
.....................                                                    [100%]

====================== 106 passed, 22 warnings in 30.89s =======================

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/diffusion_model_integration.md.

Module owners: @fhfuih @Bounty-hunter @wtomin @RuixiangMa

@SamitHuang, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 34c389ae6b07 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[hsliuustc0106]

[P2] On an HWR warm final-layout hit this runs on already-final tensors. The cold path publishes the snapshot after fusion (line 683), so a warm restore brings back fused weights — and lora_is_fused is False on the fresh pipeline object because the flag is instance state the store doesn't record. If a warm hit can co-occur with --lora-backend distill (I found no exclusion), the deltas are applied a second time — silently wrong weights, and under HSDP this is add_ on sharded DTensors, i.e. the exact #6944 crash on the warm path. Suggest recording lora_is_fused in the HWR state (or skipping fusion when warm_snapshot is present) so warm and cold agree.

[SamitHuang]

Updated _maybe_fuse_distilled_lora to check self._hwr_state.get("warm_snapshot"), marking model.lora_is_fused = True and skipping re-fusion on warm HWR restore. Added unit tests in test_diffusers_loader.py covering both warm-restore skip and cold-path fusion.

[hsliuustc0106]

Verified in 967a66c: cold loads can't see a warm snapshot (state resets per load_model, and warm_snapshot is only written on a LOCAL_HIT restore, after prior-load fusion), so the guard is exact. Worth noting the scenario was already unreachable in production since lora_path makes HWR ineligible (host_weight_loader.py:134), but making the invariant local and pinning it with the two new tests is the right call.

[hsliuustc0106]

[P3] Two small things: (1) this getter conflates loaded with fused — true today only because every load through these mixins fuses in-place (add_/sub_), which a reader can't know; a one-line docstring ("mixin loads are in-place fusions, so _lora_loaded non-empty ⇒ fused") would keep a future adapter-style load path from silently disabling the dynamic manager. Also _lora_is_fused: bool = False is assigned, breaking the "annotated, not assigned" convention on the lines above. (2) In _maybe_fuse_distilled_lora, the backend triple-check is redundant — LoRABackend(str, Enum) makes == LoRABackend.DISTILL or == "distill" sufficient; the .value fallback is dead.

[SamitHuang]

Thanks! Updated LoraLoaderMixin to only annotate _lora_is_fused: bool, added the clarifying docstring to lora_is_fused, and simplified the backend check in _maybe_fuse_distilled_lora.

[hsliuustc0106]

Fixed — bare annotation plus the docstring reads well, and no bare _lora_is_fused access exists, so no behavior change. Thanks.

[SamitHuang]

Update:

• Verified distilled LoRA weights are fused into base model weights before HSDP DTensor sharding and before quantization.
• Verified warm HWR restore path detects warm snapshot, sets model.lora_is_fused = True, and skips double-fusion.
• Added unit test coverage in tests/diffusion/lora/test_loader.py, tests/diffusion/model_loader/test_diffusers_loader.py, and tests/diffusion/test_diffusion_worker.py.
• Verified all unit tests pass and all pre-commit hooks (ruff, mypy, DCO sign-off) pass cleanly.

[SamitHuang]

@hsliuustc0106 @yenuo26 @bjf-frz PTAL, thx

[hsliuustc0106]

LGTM — both threads resolved, thanks.

[Gaohan123]

Please resolve CI conflicts