749fb75bd0
Extends flex fast-inference to `unsloth/gemma-4-26B-A4B-it` (30 layers, 128 experts top-k 8, H=2816, ~3.8B active of 25.2B). Mirrors the FlexGptOssInference / FlexMoEInference template with Gemma 4-specific wiring: - Dual dense MLP + MoE per decoder layer (Gemma4TextMLP alongside Gemma4TextExperts; outputs summed before the residual add, then multiplied by the per-layer `layer_scalar` buffer). - Per-layer sliding-window dispatch (25 sliding @ 1024 tokens, 5 full attention) via twin BlockMask built once per generate() entry. - Two-tier RoPE: sliding layers use rope_theta=10K with full head_dim rotation; full-attn layers use rope_type=proportional with theta=1M and partial_rotary_factor=0.25 (the inv_freq's zero-padded tail makes the generic rotate_half a no-op on the unrotated dims, so a single rotary helper covers both). - Per-head Q/K/V RMSNorm applied before RoPE / KV write. - attention_k_eq_v=True on full-attn layers (v_proj is None): value is the raw k_proj output, followed only by v_norm (with_scale=False). - Rebind Gemma4TextExperts.forward to forward_native_grouped_mm so decode uses the grouped_mm backend (the slow Python loop in the stock forward is neither fast nor CUDA-graph-capturable). The routing weights already include per_expert_scale via Gemma4TextRouter.forward, so no extra folding is needed. - CUDA graph capture and UNSLOTH_FLEX_COMPILE_WALKER=1 inherit from the MoE template (single bucket ladder, single pool across buckets). Arch detection: - `_detect_arch` distinguishes dense vs MoE Gemma 4 via `text_config.num_experts > 1` (same class name covers both variants). - `bind_peft_model` extends the MoE no-deepcopy shortcut to gemma4_moe. Validation (B200, bf16, 3 chat prompts, 64 tokens): - Dense 31B sanity check: coherent completion via existing FlexGemma4Inference. - Flex (cudagraph) vs HF naive: 24/24, 2/2, 4/4 tokens bitwise match. - Merge parity: 7/7 cases bitwise (rank 16 / 64, 1 + 2 adapters, bf16 + fp32, E=128, 2I=1408, H=2816 / H, I=704). - Throughput bs=8/16/32/48: 510 / 1030 / 1564 / 916 tok/s vs HF naive 134 tok/s at bs=8 (3.8x-11.7x). - GRPO smoke: DAPO-Math-17k, seed 3407, max_steps=5 — stable (see follow-up comment on PR). Out of scope: E2B/E4B KV-shared + per-layer-input variants (guarded with NotImplementedError); bnb-4bit stacked experts (no such class ships for Gemma 4 today).
74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
# SPDX-License-Identifier: GNU Affero General Public License v3.0
|
|
# Copyright 2023-present the Unsloth team. All rights reserved.
|
|
|
|
"""Numerical parity check for ``refresh_moe_lora_merge_from_pristine``
|
|
at Gemma 4 26B-A4B expert shapes.
|
|
|
|
Reuses ``merge_under_test`` and ``reference_merge`` from
|
|
``tests/flex_moe_merge_parity.py`` (both are standalone helpers — they
|
|
import nothing from the flex engine, only exercise the bitwise kernel
|
|
logic lifted from ``unsloth/inference/flex_moe.py:682-800``).
|
|
|
|
Shapes covered (Gemma 4 26B-A4B):
|
|
- gate_up_proj: (E=128, 2I=1408, H=2816) standard ``(E, out, in)``
|
|
- down_proj: (E=128, H=2816, I=704) standard ``(E, out, in)``
|
|
Rank-16 LoRA, single + dual adapter, bf16 + fp32.
|
|
|
|
Gemma 4 MoE uses the same ``F.linear``-oriented expert layout as Qwen3,
|
|
so only the ``transposed=False`` branch of
|
|
``refresh_moe_lora_merge_from_pristine`` is exercised here. (The
|
|
transposed branch is already covered for gpt-oss in
|
|
``flex_moe_merge_parity.py``.)
|
|
|
|
Usage::
|
|
CUDA_VISIBLE_DEVICES=2 python -u tests/flex_gemma4_moe_merge_parity.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
_REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(_REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_REPO_ROOT))
|
|
|
|
from tests.flex_moe_merge_parity import ( # noqa: E402
|
|
merge_under_test,
|
|
reference_merge,
|
|
test_correctness,
|
|
)
|
|
|
|
|
|
def main():
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
torch.manual_seed(3407)
|
|
print(f"[merge-parity-gemma4] device={device} dtype=bf16 + fp32")
|
|
print(f"[merge-parity-gemma4] torch={torch.__version__}")
|
|
print()
|
|
|
|
print("== Correctness (fp32 golden + bf16 realistic Gemma 4 shapes) ==")
|
|
# gate_up_proj: (E=128, 2I=1408, H=2816) standard
|
|
# down_proj: (E=128, H=2816, I=704 ) standard
|
|
cases = [
|
|
# (E, in_dim, out_dim, R, dtype, transposed, n_adapters)
|
|
( 8, 64, 128, 4, torch.float32, False, 1),
|
|
( 8, 64, 128, 4, torch.float32, False, 2),
|
|
(16, 128, 256, 8, torch.float32, False, 2),
|
|
# bf16 at Gemma 4 26B-A4B MoE shapes.
|
|
(128, 2816, 1408, 16, torch.bfloat16, False, 1), # gate_up_proj
|
|
(128, 704, 2816, 16, torch.bfloat16, False, 1), # down_proj
|
|
(128, 2816, 1408, 16, torch.bfloat16, False, 2), # two adapters
|
|
(128, 2816, 1408, 64, torch.bfloat16, False, 1), # higher rank
|
|
]
|
|
all_ok = True
|
|
for E, in_dim, out_dim, R, dtype, tr, na in cases:
|
|
ok = test_correctness(E, in_dim, out_dim, R, dtype, device,
|
|
transposed=tr, n_adapters=na)
|
|
all_ok = all_ok and ok
|
|
print(f"\n overall: {'PASS' if all_ok else 'FAIL'}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|