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.
This commit is contained in:
danielhanchen
parent
6a1bef2c88
commit
4f2bfe7f69
2 changed files with 163 additions and 1 deletions
152
tests/flex_moe_bench.py
Normal file
152
tests/flex_moe_bench.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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,)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue