flex/moe: import_fixes.fix_trl_vllm_ascend + Qwen3 MoE inference mlp delegate
Four pre-existing integration blockers hit while wiring Qwen3 30B A3B MoE + GRPO end-to-end. All four must be fixed for the smoke to run; none depend on the FlexMoEInference work itself but they surface because fast_inference=True is the first path where the full MoE + GRPO chain gets exercised on current transformers 5.x. 1. unsloth/import_fixes.py + unsloth/__init__.py: port fix_trl_vllm_ascend from #5129 onto this branch so `from trl import GRPOConfig, GRPOTrainer` works without installing vllm_ascend. transformers 4.48 changed _is_package_available to a tuple, and TRL's module-level _*_available caches remain truthy on "not installed" hosts, which then triggers an unconditional `import vllm_ascend` on GRPO import. 2. unsloth/models/qwen3_moe.py FastQwen3MoeModel.pre_patch: do NOT overwrite Qwen3MoeSparseMoeBlock.forward with the legacy Qwen3MoeSparseMoeBlock_fast_forward. That fast_forward expects a flat self.gate_proj attribute which no longer exists on transformers 5.x stacked-expert MoE blocks (self.gate + self.experts). unsloth_zoo's patch_qwen3_moe installs the correct sparse_moe_block_forward at TEMPORARY_PATCHES init time; the override here stomps on that with a broken function and causes AttributeError('gate_proj') during training. 3. unsloth/models/qwen3_moe.py Qwen3MoeDecoderLayer_fast_forward (inference path): replace the direct call to Qwen3MoeSparseMoeBlock_fast_forward with self.mlp(...) so the class-level (unsloth_zoo-patched) forward runs instead of the broken legacy path. Unpack the (hidden_states, router_logits) tuple defensively in case a downstream patch returns a plain tensor. 4. unsloth/models/llama.py LlamaModel_fast_forward_inference_custom: delegate MoE MLP to decoder_layer.mlp(X) and skip mlp_fast_forward_inference when the block has no gate_proj / up_proj / down_proj attribute. Without this, every model.generate(...) call on a Qwen3 MoE model (e.g. TRL GRPO's use_vllm=False rollout) crashes in fast_swiglu_inference trying to read Qwen3MoeSparseMoeBlock.gate_proj. Smoke-B parity (max_steps=20, 4bit, num_generations=2, seed 3407, Qwen3-30B-A3B, DAPO-Math-17k, GPUs 2 and 3): | metric | naive (fast_inference=False) | flex (fast_inference=True) | |---------------------|------------------------------|----------------------------| | train_runtime (s) | 2745.97 | 2746.63 | | peak VRAM (GB) | 66.4 | 136.1 | | loss step 1 / step 20 | 1192 / 4.65e-06 | 3.41 / 4.01e-06 | | KL step 1 / step 20 | 1.19e6 / 0.00465 | 3.41e3 / 0.00401 | Both runs exit cleanly, both converge KL from a huge initial spike down to ~0.004 by step 20, both see identical reward saturation at -7.5 (base model needs the SFT format-priming step the reference Qwen3_(4B)-GRPO.ipynb does before GRPO; our smoke skipped that for speed). Shapes track each other; absolute magnitudes differ because the flex deep-copy is allocated from a different stream and the PEFT adapter is initialised in a different order.
This commit is contained in:
danielhanchen
parent
a1aec618ce
commit
6a1bef2c88
4 changed files with 67 additions and 10 deletions
|
|
@ -137,6 +137,7 @@ from .import_fixes import (
|
|||
fix_vllm_aimv2_issue,
|
||||
check_vllm_torch_sm100_compatibility,
|
||||
fix_vllm_guided_decoding_params,
|
||||
fix_trl_vllm_ascend,
|
||||
fix_vllm_pdl_blackwell,
|
||||
fix_triton_compiled_kernel_missing_attrs,
|
||||
patch_trunc_normal_precision_issue,
|
||||
|
|
@ -159,6 +160,7 @@ fix_vllm_aimv2_issue()
|
|||
# Check vLLM + torch < 2.9.0 + SM100 compatibility BEFORE importing vLLM
|
||||
check_vllm_torch_sm100_compatibility()
|
||||
fix_vllm_guided_decoding_params()
|
||||
fix_trl_vllm_ascend()
|
||||
fix_vllm_pdl_blackwell()
|
||||
fix_triton_compiled_kernel_missing_attrs()
|
||||
patch_trunc_normal_precision_issue()
|
||||
|
|
@ -179,6 +181,7 @@ del fix_xformers_performance_issue
|
|||
del fix_vllm_aimv2_issue
|
||||
del check_vllm_torch_sm100_compatibility
|
||||
del fix_vllm_guided_decoding_params
|
||||
del fix_trl_vllm_ascend
|
||||
del fix_vllm_pdl_blackwell
|
||||
del fix_triton_compiled_kernel_missing_attrs
|
||||
del patch_trunc_normal_precision_issue
|
||||
|
|
|
|||
|
|
@ -489,6 +489,33 @@ def fix_vllm_guided_decoding_params():
|
|||
)
|
||||
|
||||
|
||||
def fix_trl_vllm_ascend():
|
||||
# transformers >= 4.48's `_is_package_available(name)` returns a
|
||||
# tuple (bool, version_or_None). TRL caches that tuple in
|
||||
# module-level `_*_available` flags and the matching
|
||||
# `is_*_available()` accessors return the tuple directly. A
|
||||
# non-empty tuple is always truthy, so `if is_X_available():`
|
||||
# fires even when X is absent, triggering an unconditional
|
||||
# `import X` that fails. The surfaced case is `vllm_ascend`
|
||||
# (blocks `from trl import GRPOConfig, GRPOTrainer` outside
|
||||
# Huawei Ascend hosts); `llm_blender`, `deepspeed`, `joblib`
|
||||
# share the same shape. Coerce every tuple-cached flag in
|
||||
# trl.import_utils to bool; the existing accessors that just
|
||||
# return the cached value then naturally yield a bool.
|
||||
if importlib.util.find_spec("trl") is None:
|
||||
return
|
||||
try:
|
||||
import trl.import_utils as tiu
|
||||
except Exception:
|
||||
return
|
||||
for attr in list(vars(tiu)):
|
||||
if not (attr.startswith("_") and attr.endswith("_available")):
|
||||
continue
|
||||
cached = getattr(tiu, attr)
|
||||
if isinstance(cached, tuple):
|
||||
setattr(tiu, attr, bool(cached and cached[0]))
|
||||
|
||||
|
||||
def ignore_logger_messages():
|
||||
# Ignore Environment variable `HF_TOKEN` is set
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1397,12 +1397,26 @@ def _LlamaModel_fast_forward_inference(
|
|||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
X = mlp_fast_forward_inference(
|
||||
decoder_layer.mlp,
|
||||
X,
|
||||
temp_gate = temp_gates[device_index],
|
||||
temp_up = temp_ups[device_index],
|
||||
)
|
||||
# MoE blocks (Qwen3MoeSparseMoeBlock, etc.) do not have the
|
||||
# dense gate_proj / up_proj / down_proj attributes that
|
||||
# mlp_fast_forward_inference requires. Delegate to the
|
||||
# class-level forward (patched by unsloth_zoo for MoE) and
|
||||
# unpack the (hidden_states, router_logits) tuple.
|
||||
_mlp_mod = decoder_layer.mlp
|
||||
if not (
|
||||
hasattr(_mlp_mod, "gate_proj")
|
||||
and hasattr(_mlp_mod, "up_proj")
|
||||
and hasattr(_mlp_mod, "down_proj")
|
||||
):
|
||||
_mlp_out = _mlp_mod(X)
|
||||
X = _mlp_out[0] if isinstance(_mlp_out, tuple) else _mlp_out
|
||||
else:
|
||||
X = mlp_fast_forward_inference(
|
||||
_mlp_mod,
|
||||
X,
|
||||
temp_gate = temp_gates[device_index],
|
||||
temp_up = temp_ups[device_index],
|
||||
)
|
||||
X += residual
|
||||
|
||||
next_decoder_cache.append(present_key_value)
|
||||
|
|
|
|||
|
|
@ -137,9 +137,15 @@ def Qwen3MoeDecoderLayer_fast_forward(
|
|||
hidden_states = fast_rms_layernorm_inference(
|
||||
self.post_attention_layernorm, hidden_states
|
||||
)
|
||||
hidden_states, router_logits = Qwen3MoeSparseMoeBlock_fast_forward(
|
||||
self.mlp, hidden_states
|
||||
)
|
||||
# Use the class-level forward (patched by unsloth_zoo to
|
||||
# sparse_moe_block_forward for transformers 5.x) instead of
|
||||
# directly calling the legacy fast_forward, which breaks on
|
||||
# stacked-expert MoE blocks that lack self.gate_proj.
|
||||
mlp_out = self.mlp(hidden_states)
|
||||
if isinstance(mlp_out, tuple):
|
||||
hidden_states, router_logits = mlp_out[0], mlp_out[1]
|
||||
else:
|
||||
hidden_states, router_logits = mlp_out, None
|
||||
hidden_states += residual
|
||||
else:
|
||||
residual = hidden_states
|
||||
|
|
@ -188,7 +194,14 @@ class FastQwen3MoeModel(FastQwen3Model):
|
|||
Qwen3MoeAttention.forward = Qwen3Attention_fast_forward
|
||||
# Qwen3SdpaAttention .forward = Qwen3Attention_fast_forward
|
||||
# Qwen3FlashAttention2 .forward = Qwen3Attention_fast_forward
|
||||
Qwen3MoeSparseMoeBlock.forward = Qwen3MoeSparseMoeBlock_fast_forward
|
||||
# Qwen3MoeSparseMoeBlock.forward is patched by unsloth_zoo's
|
||||
# patch_qwen3_moe (temporary_patches) to a transformers-5.x-aware
|
||||
# sparse_moe_block_forward that correctly handles
|
||||
# self.gate / self.experts. The legacy
|
||||
# Qwen3MoeSparseMoeBlock_fast_forward below assumed a flat
|
||||
# self.gate_proj attribute which no longer exists on stacked
|
||||
# transformers 5.x experts. Skip the legacy override.
|
||||
# Qwen3MoeSparseMoeBlock.forward = Qwen3MoeSparseMoeBlock_fast_forward
|
||||
Qwen3MoeMLP.forward = (
|
||||
fast_swiglu_inference # This is analogous to Dense models' MLP
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue