ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

tests: flex_moe_merge_parity.py — numerical + perf check for MoE LoRA merge

Browse source 
Standalone test that verifies refresh_moe_lora_merge_from_pristine
produces the same W_inf as the textbook reference
``W_ref[e] = W_pristine[e] + scaling * B[e] @ A[e]`` for both standard
(E, 2I, H) and transposed (E, H, 2I) stacked expert layouts, single and
multi-adapter, fp32 + bf16. All 8 cases bitwise-match.

Also benchmarks the batched torch.baddbmm path against (a) a per-expert
torch.addmm loop of equivalent arithmetic and (b) the dense-layer
torch.addmm used by refresh_lora_merge_from_pristine. At Qwen3-30B-A3B
MoE shapes on B200: baddbmm is 5-9.6x faster than the addmm loop; its
per-expert amortized cost is 3.6-7.1x cheaper than the dense addmm
baseline — baddbmm amortizes kernel launch and grid setup across
E=128 experts in one kernel call.
This commit is contained in:
danielhanchen 2026-04-23 04:23:09 +00:00
commit 835b346bd2
1 changed files with 290 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

290
tests/flex_moe_merge_parity.py Normal file
View file

@ -0,0 +1,290 @@
# SPDX-License-Identifier: GNU Affero General Public License v3.0
# Copyright 2023-present the Unsloth team. All rights reserved.
"""Numerical + performance check for ``refresh_moe_lora_merge_from_pristine``.
Does what we never actually did on this branch: verify the stacked MoE
LoRA merge kernel produces the same W_inf as the textbook reference
``W_ref[e] = W_pristine[e] + sum_a scaling_a * B_a[e] @ A_a[e]`` for both
standard (E, 2I, H) and transposed (E, H, 2I) expert layouts, with 1 and
2 active adapters, then benchmarks the batched ``baddbmm`` path against
the dense-layer-style ``addmm`` loop at Qwen3-30B-A3B MoE shapes.
Usage::
CUDA_VISIBLE_DEVICES=2 python -u tests/flex_moe_merge_parity.py
"""
from __future__ import annotations
import sys
import time
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))
# ---------------------------------------------------------------------------
# Reference: the exact formula from the docstring, no tricks.
# ---------------------------------------------------------------------------
def reference_merge(W_pristine, adapters, *, transposed):
"""adapters: list of (A_stacked [E*R, in], B_stacked [out, E*R], scaling)."""
E = W_pristine.shape[0]
out = W_pristine.clone()
for A_stack, B_stack, scaling in adapters:
R = A_stack.shape[0] // E
in_dim = A_stack.shape[1]
out_dim = B_stack.shape[0]
for e in range(E):
A_e = A_stack[e * R : (e + 1) * R] # (R, in_dim)
B_e = B_stack[:, e * R : (e + 1) * R] # (out_dim, R)
delta = (B_e @ A_e).to(W_pristine.dtype) # (out_dim, in_dim)
if transposed:
out[e] += scaling * delta.t()
else:
out[e] += scaling * delta
return out
# ---------------------------------------------------------------------------
# Under-test: a standalone copy of the merge body from flex_moe.py so we
# can drive it without loading a 30B model. The logic is a character-for-
# character lift of lines 747-794 in unsloth/inference/flex_moe.py.
# ---------------------------------------------------------------------------
def merge_under_test(W_inf, W_pristine, adapters):
"""adapters: list of (A_stacked, B_stacked, scaling)."""
E = W_inf.shape[0]
# Orientation detection (mirrors flex_moe.py:752-763).
A_w0, B_w0, _ = adapters[0]
in_dim = A_w0.shape[1]
out_dim = B_w0.shape[0]
d0, d1 = W_inf.shape[1], W_inf.shape[2]
if d0 == out_dim and d1 == in_dim:
is_standard = True
elif d0 == in_dim and d1 == out_dim:
is_standard = False
else:
raise RuntimeError("orientation")
W_inf.copy_(W_pristine)
for A_w, B_w, scaling in adapters:
R = A_w.shape[0] // E
A_3d = A_w.view(E, R, in_dim)
B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous()
if is_standard:
torch.baddbmm(
W_inf,
B_3d.to(W_inf.dtype),
A_3d.to(W_inf.dtype),
alpha=float(scaling),
beta=1.0,
out=W_inf,
)
else:
torch.baddbmm(
W_inf,
A_3d.transpose(-2, -1).contiguous().to(W_inf.dtype),
B_3d.transpose(-2, -1).contiguous().to(W_inf.dtype),
alpha=float(scaling),
beta=1.0,
out=W_inf,
)
return W_inf
# ---------------------------------------------------------------------------
# Correctness
# ---------------------------------------------------------------------------
def _make_adapter(E, R, in_dim, out_dim, dtype, device, seed):
g = torch.Generator(device=device).manual_seed(seed)
A = torch.randn(E * R, in_dim, generator=g, device=device, dtype=dtype) * 0.02
B = torch.randn(out_dim, E * R, generator=g, device=device, dtype=dtype) * 0.02
return A, B
def test_correctness(E, in_dim, out_dim, R, dtype, device, *, transposed, n_adapters):
scalings = [2.0, 0.5][:n_adapters]
adapters = [
(*_make_adapter(E, R, in_dim, out_dim, dtype, device, seed=11 + i), s)
for i, s in enumerate(scalings)
]
if transposed:
W_pristine = torch.randn(E, in_dim, out_dim, device=device, dtype=dtype) * 0.02
else:
W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=dtype) * 0.02
W_inf = torch.empty_like(W_pristine)
merge_under_test(W_inf, W_pristine, adapters)
W_ref = reference_merge(W_pristine, adapters, transposed=transposed)
tol = dict(atol=3e-3, rtol=3e-3) if dtype == torch.bfloat16 else dict(atol=1e-5, rtol=1e-5)
close = torch.allclose(W_inf, W_ref, **tol)
max_err = (W_inf - W_ref).abs().max().item()
scale = W_ref.abs().max().item()
rel = max_err / max(scale, 1e-8)
label = f"E={E} in={in_dim} out={out_dim} R={R} {dtype} transposed={transposed} n_adapters={n_adapters}"
verdict = "OK" if close else "FAIL"
print(f" [{verdict}] {label} max_abs={max_err:.3e} rel={rel:.3e}")
return close
# ---------------------------------------------------------------------------
# Performance: batched baddbmm vs per-expert addmm loop (same total flops).
# ---------------------------------------------------------------------------
def _sync():
if torch.cuda.is_available():
torch.cuda.synchronize()
def bench(fn, *args, iters=50, warmup=10):
for _ in range(warmup):
fn(*args)
_sync()
t0 = time.perf_counter()
for _ in range(iters):
fn(*args)
_sync()
return (time.perf_counter() - t0) / iters
def perf_vs_dense_addmm(out_dim, in_dim, R, dtype, device):
"""Dense-layer LoRA refresh: a single torch.addmm call like the one
in flex_qwen3_llama.py::refresh_lora_merge_from_pristine. Used as a
per-matrix baseline to compare against the MoE baddbmm kernel cost."""
W_pristine = torch.randn(out_dim, in_dim, device=device, dtype=dtype) * 0.02
W = torch.empty_like(W_pristine)
A = torch.randn(R, in_dim, device=device, dtype=dtype) * 0.02
B = torch.randn(out_dim, R, device=device, dtype=dtype) * 0.02
scaling = 2.0
def run_addmm():
torch.addmm(W_pristine, B, A, alpha=scaling, out=W)
t = bench(run_addmm, iters=200, warmup=20)
print(f" dense addmm out={out_dim:>5} in={in_dim:>5} R={R:>2} "
f"{t * 1e3:7.4f}ms")
return t
def perf_compare(E, in_dim, out_dim, R, dtype, device):
A_w, B_w = _make_adapter(E, R, in_dim, out_dim, dtype, device, seed=0)
W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=dtype) * 0.02
W_inf_a = torch.empty_like(W_pristine)
W_inf_b = torch.empty_like(W_pristine)
A_3d = A_w.view(E, R, in_dim)
B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous()
scaling = 2.0
def run_baddbmm():
W_inf_a.copy_(W_pristine)
torch.baddbmm(W_inf_a, B_3d, A_3d, alpha=scaling, beta=1.0, out=W_inf_a)
def run_addmm_loop():
W_inf_b.copy_(W_pristine)
for e in range(E):
torch.addmm(
W_inf_b[e],
B_3d[e],
A_3d[e],
alpha=scaling,
out=W_inf_b[e],
)
t_bad = bench(run_baddbmm)
t_loop = bench(run_addmm_loop)
# Sanity: they produce the same result.
run_baddbmm()
run_addmm_loop()
max_err = (W_inf_a - W_inf_b).abs().max().item()
print(f" E={E:>3} out={out_dim:>5} in={in_dim:>5} R={R:>2} "
f"baddbmm={t_bad * 1e3:7.3f}ms addmm_loop={t_loop * 1e3:7.3f}ms "
f"speedup={t_loop / t_bad:5.2f}x max_abs_diff={max_err:.1e}")
def main():
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(3407)
print(f"[merge-parity] device={device} dtype=bf16 + fp32")
print(f"[merge-parity] torch={torch.__version__}")
print()
print("== Correctness (fp32 golden + bf16 realistic) ==")
# Qwen3-30B-A3B shapes.
# gate_up_proj: (E=128, 2I=1536, H=2048) standard
# down_proj: (E=128, H=2048, I=768) standard
# Transposed variant: (E, H, 2I) — exercised via the `transposed=True` test.
cases = [
# (E, in_dim, out_dim, R, dtype, transposed, n_adapters)
( 8, 128, 192, 4, torch.float32, False, 1),
( 8, 128, 192, 4, torch.float32, True, 1),
( 8, 128, 192, 4, torch.float32, False, 2),
(16, 256, 512, 8, torch.float32, True, 2),
# bf16 at realistic Qwen3 MoE shapes. Mild tol to allow 3e-3 bf16 noise.
(128, 2048, 1536, 16, torch.bfloat16, False, 1), # gate_up_proj-like
(128, 768, 2048, 16, torch.bfloat16, False, 1), # down_proj-like
(128, 2048, 1536, 16, torch.bfloat16, True, 1), # transposed gate_up
(128, 2048, 1536, 16, torch.bfloat16, False, 2), # two adapters
]
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'}")
print()
print("== Performance: batched baddbmm vs per-expert addmm loop (bf16) ==")
print(" (same arithmetic, different dispatch pattern — baddbmm = 1 kernel,")
print(" addmm loop = E kernels)")
moe_cases = [
(128, 2048, 1536, 16), # gate_up_proj
(128, 768, 2048, 16), # down_proj
(128, 2048, 1536, 64), # larger rank
]
for E, in_dim, out_dim, R in moe_cases:
perf_compare(E, in_dim, out_dim, R, torch.bfloat16, device)
print()
print("== Dense baseline: torch.addmm on a single expert-sized matrix ==")
print(" (this is what flex_qwen3_llama.py:354-370 does for dense layers)")
dense_times = {}
for E, in_dim, out_dim, R in moe_cases:
t = perf_vs_dense_addmm(out_dim, in_dim, R, torch.bfloat16, device)
dense_times[(E, in_dim, out_dim, R)] = t
print()
print("== Per-expert cost comparison ==")
print(" (MoE baddbmm cost / E) vs single dense addmm for the same per-expert matrix")
for E, in_dim, out_dim, R in moe_cases:
# Measure baddbmm cost again for this exact shape to get the numerator.
A_w, B_w = _make_adapter(E, R, in_dim, out_dim, torch.bfloat16, device, seed=0)
W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=torch.bfloat16) * 0.02
W_inf = torch.empty_like(W_pristine)
A_3d = A_w.view(E, R, in_dim)
B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous()
def run_baddbmm():
W_inf.copy_(W_pristine)
torch.baddbmm(W_inf, B_3d, A_3d, alpha=2.0, beta=1.0, out=W_inf)
t_moe = bench(run_baddbmm, iters=50, warmup=10)
t_dense = dense_times[(E, in_dim, out_dim, R)]
per_expert = t_moe / E
print(f" E={E:>3} out={out_dim:>5} in={in_dim:>5} R={R:>2} "
f"moe_per_expert={per_expert * 1e6:7.2f}us "
f"dense_addmm={t_dense * 1e6:7.2f}us "
f"ratio(moe/dense)={per_expert / t_dense:5.2f}x")
if __name__ == "__main__":
main()