diff --git a/tests/flex_moe_bench.py b/tests/flex_moe_bench.py index 89382a2042..7c17663e43 100644 --- a/tests/flex_moe_bench.py +++ b/tests/flex_moe_bench.py @@ -31,7 +31,7 @@ if str(_REPO_ROOT) not in sys.path: def main(): p = argparse.ArgumentParser() - p.add_argument("--backend", choices = ["flex", "hf"], default = "flex") + p.add_argument("--backend", choices = ["flex", "hf", "hf_naive"], 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") @@ -49,21 +49,52 @@ def main(): 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", - ) + + if args.backend == "hf_naive": + # Pure transformers path: NO ``import unsloth`` so none of + # Unsloth's Qwen3 MoE attention / MLP patches run. This is the + # fair naive reference to compare flex fast-inference against. + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BitsAndBytesConfig, + ) + quant_cfg = None + if args.load_in_4bit: + quant_cfg = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=dtype, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + ) + tokenizer = AutoTokenizer.from_pretrained(args.model) + if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>": + tokenizer.pad_token = "<|vision_pad|>" + tokenizer.padding_side = "left" + model = AutoModelForCausalLM.from_pretrained( + args.model, + dtype=dtype, + quantization_config=quant_cfg, + device_map="cuda", + attn_implementation="eager", + ) + model.eval() + print(f"[bench] pure transformers (no unsloth patches)") + else: + import unsloth + print(f"[bench] unsloth={unsloth.__file__}") + from unsloth import FastLanguageModel + 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") @@ -91,7 +122,7 @@ def main(): sum(len(o.outputs[0].token_ids) for o in outs) ) else: - # HF generate + # HF generate (shared for "hf" unsloth-patched and "hf_naive" pure). inputs = tokenizer(prompts, return_tensors = "pt", padding = True).to("cuda") gen_kwargs = dict( max_new_tokens = args.max_new_tokens, diff --git a/tests/flex_moe_parity.py b/tests/flex_moe_parity.py new file mode 100644 index 0000000000..d13341a18b --- /dev/null +++ b/tests/flex_moe_parity.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Token-level parity: FlexMoEInference (CUDA-graph capture) vs HF generate. + +Same prompt set, temperature=0, max_new_tokens fixed. Reports per-prompt +token-id match rate + first divergence index. Serves as the correctness +check for the v2 grouped_mm + CUDA-graph-capture changes. + +Usage: + CUDA_VISIBLE_DEVICES=5 UNSLOTH_FAST_INFERENCE=1 \ + UNSLOTH_MOE_BACKEND=grouped_mm python -u \ + tests/flex_moe_parity.py --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 _run_flex(prompts, args, dtype, *, capture: bool): + import torch + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") + import unsloth # noqa: F401 + from unsloth import FastLanguageModel + + if not capture: + # Monkey-patch ``capture_decode_cudagraph`` to a no-op BEFORE the + # engine is built so ``generate`` takes the eager branch. The + # engine's ``self.graphs`` stays empty, ``cudagraph_captured`` + # stays False, and every step goes through ``_decode_step_eager``. + from unsloth.inference.flex_moe import FlexMoEInference + FlexMoEInference.capture_decode_cudagraph = lambda self: None + + 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=True, + ) + + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + # First call warms / captures; second call is the measurement. + _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + outputs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + token_ids = [list(o.outputs[0].token_ids) for o in outputs] + texts = [o.outputs[0].text for o in outputs] + return token_ids, texts, tokenizer + + +def _run_hf(prompts, args, dtype): + import torch + # Pure Hugging Face: NO ``import unsloth`` — we want the unpatched + # reference forward to compare flex against. Quantization via + # transformers' ``BitsAndBytesConfig`` matches what unsloth loads + # under the hood for ``load_in_4bit=True``. + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BitsAndBytesConfig, + ) + # Translate the unsloth-flavoured model id (unsloth/Qwen3-30B-A3B-Instruct-2507) + # to the 4bit variant if load_in_4bit was requested (FastLanguageModel + # does this implicitly; do it explicitly here for the naive path). + model_id = args.model + quant_cfg = None + if args.load_in_4bit: + quant_cfg = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=dtype, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + ) + tokenizer = AutoTokenizer.from_pretrained(model_id) + model = AutoModelForCausalLM.from_pretrained( + model_id, + dtype=dtype, + quantization_config=quant_cfg, + device_map="cuda", + attn_implementation="eager", + ) + model.eval() + # Qwen3-30B-A3B-Instruct-2507 uses <|vision_pad|> as its pad token. + # Unsloth's loader may swap it to a sentinel; reset to the HF default + # so batched left-padded generation matches the authoritative config. + if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>": + tokenizer.pad_token = "<|vision_pad|>" + tokenizer.padding_side = "left" + gen_kwargs = dict( + max_new_tokens=args.max_new_tokens, + do_sample=False, + temperature=1.0, + pad_token_id=tokenizer.pad_token_id, + ) + inputs = tokenizer(prompts, return_tensors="pt", padding=True).to("cuda") + out = model.generate(**inputs, **gen_kwargs) + prompt_len = inputs["input_ids"].shape[1] + eos = tokenizer.eos_token_id + pad = tokenizer.pad_token_id + token_ids = [] + texts = [] + for row in out: + ids = row[prompt_len:].tolist() + while ids and ids[-1] in (eos, pad): + ids.pop() + token_ids.append(ids) + texts.append(tokenizer.decode(ids, skip_special_tokens=True)) + return token_ids, texts, tokenizer + + +def main(): + p = argparse.ArgumentParser() + 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("--max_new_tokens", type=int, default=32) + p.add_argument("--max_seq_length", type=int, default=1024) + p.add_argument("--backend", choices=["flex", "flex_eager", "hf"], required=True) + p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2") + args = p.parse_args() + + import torch + prompts = [ + "The quick brown fox jumps over", + "Q: What is 23 + 19?\nA:", + "Paris is the capital of", + ] + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + if args.backend == "flex": + token_ids, texts, tok = _run_flex(prompts, args, dtype, capture=True) + elif args.backend == "flex_eager": + token_ids, texts, tok = _run_flex(prompts, args, dtype, capture=False) + else: + token_ids, texts, tok = _run_hf(prompts, args, dtype) + + precision = "4bit" if args.load_in_4bit else args.dtype + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / f"parity_{args.backend}_{precision}.json" + with open(out_path, "w") as f: + json.dump( + { + "backend": args.backend, + "precision": precision, + "prompts": prompts, + "token_ids": token_ids, + "texts": texts, + }, + f, + indent=2, + ) + print(f"[parity-{args.backend}] wrote {out_path}") + for i, (p_, t_) in enumerate(zip(prompts, texts)): + print(f"[parity-{args.backend}] prompt {i}: {p_!r}") + print(f"[parity-{args.backend}] completion {i}: {t_!r}") + + +if __name__ == "__main__": + main() diff --git a/unsloth/inference/flex_engine.py b/unsloth/inference/flex_engine.py index 0559703c95..30facbd276 100644 --- a/unsloth/inference/flex_engine.py +++ b/unsloth/inference/flex_engine.py @@ -397,12 +397,11 @@ class FlexEngine: Impl = FlexGemma4Inference elif arch == "qwen3_moe": Impl = FlexMoEInference - # MoE decode uses bincount + Python expert loops inside - # ``forward_moe_backend`` (unsloth_zoo moe_utils), which is - # not CUDA-graph capturable. Force eager decode so a stray - # ``capture_cudagraph=True`` does not fail inside a captured - # graph on the first token. - self.capture_cudagraph = False + # CUDA graph capture is supported on the ``grouped_mm`` MoE + # backend only. ``FlexMoEInference.capture_decode_cudagraph`` + # re-checks the active backend at capture time and skips + # capture on any other backend, leaving ``capture_cudagraph`` + # alone here. else: Impl = FlexInference # Pass the cuMem allocator through so the impl can wrap ONLY diff --git a/unsloth/inference/flex_moe.py b/unsloth/inference/flex_moe.py index 6be8d69329..91710dc48d 100644 --- a/unsloth/inference/flex_moe.py +++ b/unsloth/inference/flex_moe.py @@ -380,26 +380,100 @@ class FlexMoEInference: def _decode_step( self, batch_idx: torch.Tensor, input_ids: torch.Tensor, input_pos: torch.Tensor ): - # MoE path is always eager — no CUDA graph replay. See capture - # docstring below. self.input_pos_buffer.zero_() self.input_pos_buffer[batch_idx] = input_pos - return self._decode_step_eager(batch_idx, input_ids) + if not self.cudagraph_captured: + return self._decode_step_eager(batch_idx, input_ids) + bs = input_ids.size(0) + key = next(x for x in self.graph_bs if x >= bs) + graph = self.graphs[key] + gv = self.graph_vars + # batch_idx=0 is the reserved no-op slot. Zero out the unused part + # of each capture-shape buffer so padded entries don't write into + # real KV pages. + for k, v in gv.items(): + if k != "outputs": + v.zero_() + gv["input_ids"][:bs] = input_ids + gv["batch_idx"][:bs] = batch_idx + graph.replay() + return gv["outputs"][:bs] def capture_decode_cudagraph(self): - """Not supported for MoE. ``Qwen3MoeExperts.forward`` uses - ``torch.where`` + a data-dependent Python for-loop over experts - (shapes depend on routing), which cannot be captured. Raising - here so a stray ``capture_cudagraph=True`` fails loudly. + """Capture one CUDA graph per batch-size bucket for MoE decode. - Future: a padded-fixed-shape dispatch can be gated behind - ``UNSLOTH_MOE_STATIC_DISPATCH=1`` to make capture viable — out - of scope for the first cut. + Supported on the ``grouped_mm`` MoE backend only. On that backend + the decode path is fixed-shape: + ``bincount(minlength=num_experts) → cumsum → argsort → + torch._grouped_mm × 2 → index_add_``. Python control flow in + ``sparse_moe_block_forward`` runs once at capture time; only the + recorded CUDA kernels replay. + + For any other backend (``unsloth_triton``, ``native_torch``) this + method logs a warning and returns without enabling replay, so + ``generate(capture_cudagraph=True)`` silently falls back to eager + decode instead of failing inside the captured graph. + + Pre-reserves a page for every ``batch_idx`` slot so the paged-KV + ``index_put_`` during capture hits valid physical addresses. The + reservations are erased after capture — replay reads / writes the + same physical pages regardless of the logical batch state, + because ``batch_idx = 0`` is reserved as a padding slot. """ - raise NotImplementedError( - "FlexMoEInference does not support CUDA graph capture: MoE " - "expert routing has data-dependent shapes. Run with " - "capture_cudagraph=False." + try: + from unsloth_zoo.temporary_patches.moe_utils import select_moe_backend + backend = select_moe_backend() + except Exception: + backend = None + if backend != "grouped_mm": + print( + f"[flex] MoE CUDA graph capture requires the 'grouped_mm' " + f"backend (got {backend!r}); skipping capture, decode stays " + f"eager." + ) + return + + max_bs = self.max_batch_size + reserved_batches = [] + for bi in range(1, max_bs): + try: + allocated = self.page_table.allocate() + self.page_table.reserve( + allocated, + torch.tensor([allocated], device = self.device, dtype = torch.long), + self.page_size, + ) + reserved_batches.append(allocated) + except Exception: + break + + input_ids = torch.zeros(max_bs, dtype = torch.int64, device = self.device) + batch_idx = torch.arange(max_bs, dtype = torch.int64, device = self.device) + outputs = torch.zeros( + (max_bs, self.model.config.vocab_size), + dtype = self.model.dtype, + device = self.device, + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + pool = None + for bs in reversed(self.graph_bs): + if bs > max_bs: + continue + print(f"[flex-moe] capturing CUDA graph for bs={bs}") + torch.cuda.synchronize() + _ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool): + outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + if pool is None: + pool = graph.pool() + self.graphs[bs] = graph + torch.cuda.synchronize() + for bi in reserved_batches: + self.page_table.erase(bi) + self.graph_vars = dict( + input_ids = input_ids, batch_idx = batch_idx, outputs = outputs ) def refresh_inference_from_base(self): @@ -430,13 +504,22 @@ class FlexMoEInference: @torch.inference_mode() def generate(self, sequences: list[Sequence], capture_cudagraph = False): - """Decode loop. ``capture_cudagraph`` is ignored for MoE — - always runs eager.""" + """Decode loop. Captures one CUDA graph per bucket on first call + when ``capture_cudagraph=True`` and the ``grouped_mm`` MoE + backend is active; otherwise falls back to eager decode.""" self.tokenize(sequences) waiting = deque(sequences) running = deque() done = [] + if capture_cudagraph and not self.cudagraph_captured: + self.capture_decode_cudagraph() + # ``capture_decode_cudagraph`` leaves ``cudagraph_captured`` + # alone when it skips (non-grouped_mm backend), so only flip + # the flag when at least one bucket was actually captured. + if self.graphs: + self.cudagraph_captured = True + while waiting or running: batch = [] while waiting and self.page_table.can_reserve(waiting[0].total_length):