unsloth/tests/flex_moe_parity.py

# 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

    # Opt-in torch.compile wrap of the decode walker for parity check
    # — enabled via env var to avoid cluttering the CLI further.
    if os.environ.get("FLEX_MOE_COMPILE_WALKER") == "1":
        import torch as _torch
        from unsloth.inference import flex_moe as _flex_moe_mod
        _orig = _flex_moe_mod.call_moe_model_with_flex_kwargs
        _flex_moe_mod.call_moe_model_with_flex_kwargs = _torch.compile(
            _orig, fullgraph=False, dynamic=False
        )
        print("[parity] torch.compile(call_moe_model_with_flex_kwargs) enabled")

    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()