[Bugfix][Ascend NPU][Diffusion] Load over-wide unquantized fallback weights straight into host memory under DLOFix/mxfp online offload after quant

[brandneway]

Purpose

Background. Ascend's npu_quant_matmul (QuantBatchMatmulV3) rejects any
weight whose output dimension exceeds 65535, and it fails at the first forward
rather than at load time, so _fell_back_to_unquantized_npu keeps such layers
unquantized. On MiniMax-H3 every one of the 50 DiT blocks has an
adaln_proj.linear with out_features = 18 * 5376 = 96768, so all 50 fall
back. With --enable-distributed-layerwise-offload + online INT8
(--diffusion-quantization-config '{"transformer":{"method":"int8"}}'), these
bf16 weights (496 MiB each, ~24.2 GiB per rank in total) were materialized on
the accelerator at construction time and stayed resident until the final
model.to("cpu"). Together with the TP8-sharded text encoder (~6.9 GiB per
rank), memory-snapshot replay shows the single-rank startup peak at
~31.9 GiB — over the 32 GB card limit — even though under DLO none of these
weights are device-resident at inference time. (The genuinely quantized layers
already stream back to CPU per layer via enable_offload_after_quant; the
fallback layers were the residual peak.)

Fix. When the loader will offload the whole model back to host memory
after online quantization (the offload-after-quant path), over-wide fallback
weights now never visit the accelerator during startup:

• New UnquantizedHostLinearMethod: mirrors LazyWeightMixin's meta +
  just-in-time materialization pattern, but materializes on CPU and skips
  quantization entirely. Checkpoint chunks are CPU tensors, so loading stays
  CPU→CPU; the weights only reach the accelerator through the offload
  backend's runtime prefetch, where DLO's shared bf16 buffer is already sized
  for them.
• load_unquantizable_fallback_on_cpu() context (ContextVar-scoped) is
  entered by DiffusersPipelineLoader.load_model around model construction
  only when offload-after-quant is active; otherwise the eager path is
  byte-for-byte unchanged, so non-DLO deployments are unaffected.
• Fully loaded fallback layers set
  _already_called_process_weights_after_loading so the post-load sweep skips
  them instead of bouncing each ~496 MiB host tensor to the accelerator and
  back for a no-op process_weights_after_loading.
• UnquantizedHostLinearMethod is added to the DLO+AllGather online-method
  allowlist: its runtime layout is a plain contiguous bf16 weight, identical
  to the ordinary unquantized path DLO already shards.

This also removes two wasted full-size copies for these weights (H2D during
loading + D2H at the final model.to("cpu")).

Test Plan

• CPU unit tests:
  • tests/diffusion/quantization/test_int8_config.py::TestHostFallbackLoading
    (eager path unchanged without the context; meta deferral under the
    context; CPU-direct load with value equality; chunked load marks
    completion only when full; context does not leak)
  • tests/diffusion/model_loader/test_diffusers_loader.py::test_dlo_allgather_allows_unquantized_host_fallback
• NPU E2E: launch MiniMax-H3 ref2va on 8 cards with DLO + online INT8 and
  record the startup memory peak via torch_npu memory snapshot.

vLLM Version: <fill in>

vLLM-Omni Commit: <fill in>

Test Result

• Unit tests pass (pytest tests/diffusion/quantization/test_int8_config.py tests/diffusion/model_loader/test_diffusers_loader.py)
• Startup peak drops from ~31.9 GiB to <10 GiB per rank on the 8-card
  MiniMax-H3 DLO + INT8 launch, with the new log line
  "Loading over-wide unquantized fallback weights straight into host memory"

[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/quantization.md.

Module owners: @david6666666 @Isotr0py @lishunyang12

@brandneway, 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.

[david6666666]

[P2] HSDP still bypasses this path. offload_after_quant is set before the branch, but _load_model_with_hsdp() still constructs and loads online-quantized layers on the accelerator via self.load_weights(model), so an over-wide fallback only moves to CPU at the final model.to("cpu"). Please either apply the host-fallback/offload handling there or explicitly reject/document CPU offload + HSDP.

[brandneway]

Addressed in 8d543c0. The HSDP path now applies the same host-fallback handling: _load_model_with_hsdp takes offload_after_quant and wraps both _init_from_load_format and self.load_weights(model) in load_unquantizable_fallback_on_cpu() when it is set. As you noted, the HSDP branch initializes on the accelerator whenever a quant config is present (hsdp_defer_to_cpu=False), so without this the over-wide fallback weights round-tripped through the device before apply_hsdp_to_model — now they load straight into host memory and sharding distributes them like any other parameter.

Covered by test_hsdp_enters_host_fallback_context_only_when_offloading_after_quant (parametrized on/off) in tests/diffusion/model_loader/test_diffusers_loader.py.

[david6666666]

Could you add a test through load_model() that checks the fallback stays on CPU through the post-load sweep? The current tests call create_weights() directly or only check the allowlist, so they don't protect the startup-memory bound this change relies on.

[brandneway]

Added in 8d543c0: test_dlo_load_model_keeps_host_fallback_on_cpu_through_post_load_sweep drives the full load_model() with a DLO + online-quant config and a real over-wide (NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1) layer built through NPUInt8OnlineLinearMethod.create_weights. It asserts:

1. construction ran inside load_unquantizable_fallback_on_cpu() (the wiring this PR adds),
2. the layer's quant_method was swapped to UnquantizedHostLinearMethod,
3. after the real _process_weights_after_loading pass and the final model.to("cpu") sweep, the weight is still the same CPU tensor with the loaded values intact, and _already_called_process_weights_after_loading is set (so the sweep skips it rather than bouncing the host tensor to the accelerator and back).

Also fixed the pre-commit failures (pytest Hardware mark on test_int8_config.py, SPDX headers, ruff format, three mypy ignores on the mixin process_weights_after_loading calls) and rebased onto main resolving the diffusers_loader.py conflict — the restructured AllGather validation on main now carries the "host-loaded unquantized fallback layers are also allowed" wording. All commits are signed off; DCO is green.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 90b939c5504c produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[Dong1017]

Fix direction is right; no code-level blocker. I verified the mechanism on a real Ascend 910B1 (vllm 0.28.0 + torch_npu 2.10, head 8d543c0):

• Unit tests: 59 passed, 3 skipped (NPU-only cases) — checkbox 1 confirmed.
• A/B probe on the exact MiniMax-H3 adaln_proj.linear shape (out=96768): base materializes the fallback weight on the NPU at construction; head keeps it on CPU for its whole life, NPU peak 0 bytes, completion flag set, and the new log line appears verbatim. 50 × 0.496 GiB extrapolates to the claimed ~24 GiB/rank — checkbox 2's mechanism is confirmed, but please still attach the 8-card torch_npu memory-snapshot run declared in the Test Plan and fill the Test Result checkboxes.

Non-blocking:

• docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md:359 still says other online quantization methods require rank-local transfer; this PR widens the validated set — please update the bullet (see inline).
• Title has the branch name leaked in: "...under DLOFix/mxfp online offload after quant".

[Dong1017]

Docs consistency: this allowlist (and the error at line 630) contradicts the Limitations bullet at docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md:359 ("Other online quantization methods require rank-local transfer..."). Please update it to include host-loaded unquantized fallback.

[lishunyang12]

Mock the TP group before creating weights; both new offload tests fail without it.

[brandneway]

Fixed in 1367cb4: TestOffloadAfterQuant now has an autouse _mock_tp fixture patching get_tensor_model_parallel_world_size / get_tensor_model_parallel_rank / get_tp_group (same stand-in as TestHostFallbackLoading in test_int8_config.py), so both create_weights-driven offload tests run without initializing torch.distributed. Good catch — the GitHub build job doesn't run the test files, which is exactly how this slipped through.

[lishunyang12]

Share the lazy weight materialization with LazyWeightMixin instead of copying it.

[brandneway]

Done in 1367cb4. UnquantizedHostLinearMethod is now (LazyWeightMixin, UnquantizedLinearMethod) and the duplicated patched_weight_loader is gone. The mixin grew a _lazy_load_device() hook (default torch.get_default_device()) that the host fallback overrides to torch.device("cpu"), so the materialization target — the only real difference — is a one-method override. The other two deltas are declared explicitly:

• supports_offload_after_quant = False: the weight never visits the accelerator, so there is nothing to hand back to host (the mixin's layer.to("cpu") on completion stays inert because _request_offload_after_quant never opts this method in);
• process_weights_after_loading is an explicit no-op the mixin's completion path calls right before flagging the layer — and the flag is what keeps the loader's post-load sweep from bouncing the host tensor to the accelerator and back.

All four TestHostFallbackLoading tests pass unchanged against the shared implementation, which also pins that the behavioral contract (meta deferral, chunked completion, CPU materialization) is identical.

[david6666666]

Since the allowlist now admits host-loaded fallback layers, the Limitations section in docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md (~L359, "Other online quantization methods require rank-local transfer") is out of sync — same point as the earlier docs thread. Can we fold that doc update into this PR?

[brandneway]

Folded into this PR in 90b939c — the Limitations bullet in distributed_layerwise_offload.md now reads: per-tensor online FP8, INT8, and MXFP8 linears run with either DiT transfer path via the ordinary loader, over-wide NPU-unquantizable fallback layers (output dimension past 65535) load straight into host memory under offload-after-quant with a plain contiguous bf16 layout and are equally allowed under allgather, and everything else still requires rank-local transfer until validated. This also picks up the INT8/MXFP8 entries from #6573 that the same sentence had missed.