From 4f2bfe7f697cca0fd732d464278212929fc42e23 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 22 Apr 2026 12:36:18 +0000 Subject: [PATCH] flex/moe: flex_moe_bench.py + training-path tuple defensive unpack Adds tests/flex_moe_bench.py: 2-round median decode throughput bench comparing flex (FlexMoEInference) and HF generate on the same (n_prompts, max_new_tokens, precision) workload. Writes async_task_outputs/qwen3_moe_grpo_bench/bench_decode_{backend}_{precision}.json. Also defensively unpacks self.mlp(hidden_states) in Qwen3MoeDecoderLayer_fast_forward's training branch: unsloth_zoo.temporary_patches.qwen3_moe.sparse_moe_block_forward returns a plain tensor for transformers 5.x stacked experts, but the decoder wrapper unpacked a 2-tuple. The inference branch was already fixed in the previous commit; the training branch hit the same ValueError under plain HF generate (no _flag_for_generation). Bench numbers (Qwen3-30B-A3B, 4bit, rank 16 LoRA, bs=8, 64 new tokens, B200): | backend | median tok/s | peak VRAM (GB) | median wall (s) | |---------|--------------|----------------|------------------| | HF | 80.5 | 57.2 | 6.36 | | Flex | 55.6 | 116.0 | 9.21 | Flex is correctness-complete but not yet performance-competitive on MoE decode at bs=8. Two structural reasons: - MoE decode runs eager (forward_moe_backend uses bincount + Python expert loops which are not CUDA-graph capturable), so flex loses its main dense-model advantage. - FlexEngine deep-copies the HF model for the rollout copy, doubling weight residency. For Qwen3-30B-A3B at bf16 that is ~60 GB extra. The pristine-base third copy is skipped for Qwen3 MoE (see the first commit of this series) but the inference deep-copy remains. Follow-ups (not blockers for correctness): - torch.compile(dynamic=True) on call_moe_model_with_flex_kwargs to recover some of the CUDA-graph throughput without requiring graph capture. - Evaluate flex's scaling vs HF generate at bs=32 / bs=64, where paged-KV reuse should dominate per-prompt cost. - A quantised-only inference copy (4bit forward, fp32 LoRA injection) so the flex path fits inside 2x 4bit weight residency (~34 GB) instead of the current post-dequantisation footprint. --- tests/flex_moe_bench.py | 152 ++++++++++++++++++++++++++++++++++++ unsloth/models/qwen3_moe.py | 12 ++- 2 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 tests/flex_moe_bench.py diff --git a/tests/flex_moe_bench.py b/tests/flex_moe_bench.py new file mode 100644 index 0000000000..89382a2042 --- /dev/null +++ b/tests/flex_moe_bench.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Decode throughput bench: FlexMoEInference vs HF generate on Qwen3 MoE. + +Apples-to-apples decode on the same prompt set + new-token budget, +same LoRA rank, same precision. Mirrors PR #5123's +``tests/flex_fastlm_bench.py`` shape. + +Usage: + CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 python -u \ + tests/flex_moe_bench.py --backend flex --load_in_4bit + + CUDA_VISIBLE_DEVICES=0 python -u \ + tests/flex_moe_bench.py --backend hf --load_in_4bit +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--backend", choices = ["flex", "hf"], default = "flex") + p.add_argument("--model", default = "unsloth/Qwen3-30B-A3B-Instruct-2507") + p.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16") + p.add_argument("--load_in_4bit", action = "store_true") + p.add_argument("--n_prompts", type = int, default = 8) + p.add_argument("--max_new_tokens", type = int, default = 64) + p.add_argument("--max_seq_length", type = int, default = 1024) + p.add_argument("--warmup_rounds", type = int, default = 1) + p.add_argument("--timed_rounds", type = int, default = 2) + p.add_argument("--out_dir", default = "async_task_outputs/qwen3_moe_grpo_bench") + args = p.parse_args() + + import torch + + if args.backend == "flex": + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") + + import unsloth + print(f"[bench] unsloth={unsloth.__file__}") + from unsloth import FastLanguageModel + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + torch.cuda.reset_peak_memory_stats() + t_load0 = time.perf_counter() + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = args.model, + max_seq_length = args.max_seq_length, + dtype = dtype, + load_in_4bit = args.load_in_4bit, + fast_inference = args.backend == "flex", + ) + t_load = time.perf_counter() - t_load0 + peak_load = torch.cuda.max_memory_reserved() / 1024**3 + print(f"[bench] loaded in {t_load:.1f}s peak {peak_load:.1f} GB") + + prompts = [f"The quick brown fox jumps over fence {i}, then" + for i in range(args.n_prompts)] + + if args.backend == "flex": + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + # Warmup + for _ in range(args.warmup_rounds): + _ = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) + + # Timed + wall_times = [] + tok_counts = [] + for _ in range(args.timed_rounds): + t0 = time.perf_counter() + outs = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) + wall_times.append(time.perf_counter() - t0) + tok_counts.append( + sum(len(o.outputs[0].token_ids) for o in outs) + ) + else: + # HF generate + inputs = tokenizer(prompts, return_tensors = "pt", padding = True).to("cuda") + gen_kwargs = dict( + max_new_tokens = args.max_new_tokens, + do_sample = False, + temperature = 1.0, + pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id, + ) + # Warmup + for _ in range(args.warmup_rounds): + _ = model.generate(**inputs, **gen_kwargs) + torch.cuda.synchronize() + + wall_times = [] + tok_counts = [] + for _ in range(args.timed_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + out = model.generate(**inputs, **gen_kwargs) + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + n_new = (out.shape[1] - inputs["input_ids"].shape[1]) * out.shape[0] + tok_counts.append(n_new) + + peak_gen = torch.cuda.max_memory_reserved() / 1024**3 + median_wall = sorted(wall_times)[len(wall_times) // 2] + median_tok = tok_counts[len(wall_times) // 2] + tok_per_s = median_tok / median_wall if median_wall > 0 else 0.0 + + print(f"[bench] wall: {wall_times}") + print(f"[bench] tok counts: {tok_counts}") + print(f"[bench] median wall: {median_wall:.2f}s median tok/s: {tok_per_s:.1f}") + print(f"[bench] peak VRAM after gen: {peak_gen:.1f} GB") + + precision = "4bit" if args.load_in_4bit else args.dtype + out_dir = Path(args.out_dir) + out_dir.mkdir(parents = True, exist_ok = True) + summary = { + "phase": "bench_decode", + "backend": args.backend, + "model": args.model, + "precision": precision, + "n_prompts": args.n_prompts, + "max_new_tokens": args.max_new_tokens, + "wall_times_s": wall_times, + "tok_counts": tok_counts, + "median_wall_s": round(median_wall, 3), + "median_tok_s": round(tok_per_s, 1), + "peak_vram_load_gb": round(peak_load, 2), + "peak_vram_after_gen_gb": round(peak_gen, 2), + "t_load_s": round(t_load, 1), + } + with open(out_dir / f"bench_decode_{args.backend}_{precision}.json", "w") as f: + json.dump(summary, f, indent = 2) + print(f"[bench] wrote {out_dir / f'bench_decode_{args.backend}_{precision}.json'}") + + +if __name__ == "__main__": + main() diff --git a/unsloth/models/qwen3_moe.py b/unsloth/models/qwen3_moe.py index c903ec4794..52506449ae 100644 --- a/unsloth/models/qwen3_moe.py +++ b/unsloth/models/qwen3_moe.py @@ -166,7 +166,17 @@ def Qwen3MoeDecoderLayer_fast_forward( # MoE Router MLP residual = hidden_states hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states) - hidden_states, router_logits = self.mlp(hidden_states) + # unsloth_zoo's sparse_moe_block_forward returns a plain tensor + # for transformers 5.x stacked experts; the legacy patched + # forward returned a (hidden_states, router_logits) tuple. + # Handle both. + mlp_out = self.mlp(hidden_states) + if isinstance(mlp_out, tuple): + hidden_states = mlp_out[0] + router_logits = mlp_out[1] if len(mlp_out) > 1 else None + else: + hidden_states = mlp_out + router_logits = None hidden_states = residual + hidden_states outputs = (hidden_states,)