f0115f8d70
FastLanguageModel.from_pretrained(model_name=\"unsloth/Qwen3-30B-A3B-...\", fast_inference=True) was silently routing Qwen3MoeForCausalLM to the dense FlexInference path because _detect_arch matched the \"qwen3\" substring first. That path works for attention but drops all MoE LoRA adapters at rollout time: refresh_lora_merge_from_pristine walks named_modules for LoraLayer instances and calls base_model.get_submodule(name).weight.data, which does not see the stacked nn.Parameter tensors on Qwen3MoeExperts.gate_up_proj / down_proj. Decode capture also breaks on bincount + Python expert loops inside forward_moe_backend. Adds unsloth/inference/flex_moe.py: - call_moe_model_with_flex_kwargs: Qwen3 MoE decoder walker. Identical to the dense walker for the attention half. Unpacks the mlp(...) return for both stock HF (plain tensor) and Unsloth's patched Qwen3MoeSparseMoeBlock_fast_forward (tuple of (hidden_states, router_logits)). - FlexMoEInference: API-compatible with FlexInference so the arch dispatch is a one-line change. cudagraph_captured is permanently False; capture_decode_cudagraph raises NotImplementedError so a stray capture_cudagraph=True fails loudly rather than producing silently wrong output. - refresh_moe_lora_merge_from_pristine: batched torch.baddbmm LoRA fuse over stacked 3D expert tensors. Handles both standard (E, 2*I, H) and transposed (E, H, 2*I) orientations via a runtime shape check against the flat lora_A / lora_B shapes. In-place write so flex prefill and paged-KV replay see refreshed values. Wires the new class into the engine: - flex_engine._detect_arch: check \"qwen3moe\" / \"qwen3_moe\" BEFORE the dense \"qwen3\" substring (Qwen3MoeForCausalLM contains both). - flex_engine.FlexEngine.__init__: route arch==\"qwen3_moe\" to FlexMoEInference and force self.capture_cudagraph = False so the MoE expert loops are never captured. - inference.__init__: export FlexMoEInference. - models/loader.py: uncomment the qwen3_moe branch so FastQwen3MoeModel applies training-side patches before FlexEngine wraps the model. Flex attention / paged KV / block-mask / sampling-param shim / vLLM shim / sleep-wake are reused verbatim.
43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
# SPDX-License-Identifier: GNU Affero General Public License v3.0
|
|
# Copyright 2023-present the Unsloth team. All rights reserved.
|
|
|
|
"""Flex-attention inference engines.
|
|
|
|
``UNSLOTH_FAST_INFERENCE=1`` routes ``FastLanguageModel.from_pretrained``
|
|
through :func:`load_flex`, which wraps the selected HF model with a
|
|
:class:`FlexEngine` that presents the vLLM ``LLM`` surface used by
|
|
Unsloth / TRL GRPO (``.generate``, ``.chat``, ``.sleep``, ``.wake_up``,
|
|
``.llm_engine``, plus ``save_lora`` / ``load_lora`` via the module
|
|
shim).
|
|
|
|
Four architectures are supported today: Qwen3 (dense), Qwen3-MoE,
|
|
Llama-3, Gemma-4-E2B-it. Anything else raises
|
|
:class:`NotImplementedError`; unset the env var or use vLLM instead."""
|
|
|
|
from .flex_engine import (
|
|
FlexEngine,
|
|
build_flex_engine,
|
|
install_flex_sentinel,
|
|
load_flex,
|
|
)
|
|
from .flex_moe import FlexMoEInference
|
|
from .vllm_shim import (
|
|
CompletionOutput,
|
|
LoRARequest,
|
|
RequestOutput,
|
|
load_lora,
|
|
save_lora,
|
|
)
|
|
|
|
__all__ = [
|
|
"FlexEngine",
|
|
"FlexMoEInference",
|
|
"load_flex",
|
|
"build_flex_engine",
|
|
"install_flex_sentinel",
|
|
"LoRARequest",
|
|
"RequestOutput",
|
|
"CompletionOutput",
|
|
"save_lora",
|
|
"load_lora",
|
|
]
|