feat: add sonicmoe (#3411)
* feat: add sonicmoe * feat: add torch compile for routing * feat: add routing smoke test * feat: add qwen3_5_moe, qwen3_vl_moe, qwen3_omni_moe * fix: disable mlp kernel for sonicmoe too * feat: update to sonicmoe release * chore: update import following new sonicmoe changes * feat: update handling for blackwell * feat: add sonicmoe e2e test * fix: installation for updated sonicmoe * fix: git commit * fix: ignore py req and fix metadata * fix: increase min hidden size to match sonicmoe kernel min * fix: attempt properly interleave and handle unpatch mid-test * chore: refactor teardown better * chore: refactor to re-use rearrange * fix: add idempotency guard * fix: address comments on CI memory and interleave * fix: tests grad, param doublewrapped
This commit is contained in:
NanoCode012
GitHub
@@ -6,7 +6,7 @@
|
||||
Unit tests for scattermoe-lora code-review fixes.
|
||||
|
||||
Tests cover:
|
||||
- KernelsArgs validator: disable_mlp_kernel_scattermoe
|
||||
- KernelsArgs validator: disable_mlp_kernel
|
||||
- CPU_Offloaded_Gradient_Checkpointer: tuple vs plain tensor backward
|
||||
- ParallelExperts: scaling=0.0 not treated as falsy
|
||||
- single2scatter: non-aligned K/N dimensions
|
||||
@@ -20,12 +20,12 @@ import pytest
|
||||
import torch
|
||||
|
||||
# ============================================================================
|
||||
# 1. KernelsArgs: disable_mlp_kernel_scattermoe validator
|
||||
# 1. KernelsArgs: disable_mlp_kernel validator
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestKernelsArgsValidator:
|
||||
"""Test that disable_mlp_kernel_scattermoe sets both flags correctly.
|
||||
"""Test that disable_mlp_kernel sets both flags correctly.
|
||||
|
||||
These tests call the validator classmethod directly on raw dicts,
|
||||
since lora_mlp_kernel / mlp_kernel are not declared model fields.
|
||||
@@ -40,7 +40,7 @@ class TestKernelsArgsValidator:
|
||||
"use_scattermoe": True,
|
||||
"lora_mlp_kernel": True,
|
||||
}
|
||||
result = KernelsArgs.disable_mlp_kernel_scattermoe(data)
|
||||
result = KernelsArgs.disable_mlp_kernel(data)
|
||||
assert result["lora_mlp_kernel"] is False
|
||||
assert result["mlp_kernel"] is False
|
||||
|
||||
@@ -52,7 +52,7 @@ class TestKernelsArgsValidator:
|
||||
"use_kernels": True,
|
||||
"use_scattermoe": True,
|
||||
}
|
||||
result = KernelsArgs.disable_mlp_kernel_scattermoe(data)
|
||||
result = KernelsArgs.disable_mlp_kernel(data)
|
||||
assert result["mlp_kernel"] is False
|
||||
# lora_mlp_kernel was not in data, should not be added
|
||||
assert "lora_mlp_kernel" not in result
|
||||
@@ -66,7 +66,7 @@ class TestKernelsArgsValidator:
|
||||
"use_scattermoe": True,
|
||||
"lora_mlp_kernel": False,
|
||||
}
|
||||
result = KernelsArgs.disable_mlp_kernel_scattermoe(data)
|
||||
result = KernelsArgs.disable_mlp_kernel(data)
|
||||
assert result["lora_mlp_kernel"] is False
|
||||
|
||||
def test_no_change_when_scattermoe_disabled(self):
|
||||
@@ -78,7 +78,7 @@ class TestKernelsArgsValidator:
|
||||
"use_scattermoe": False,
|
||||
"lora_mlp_kernel": True,
|
||||
}
|
||||
result = KernelsArgs.disable_mlp_kernel_scattermoe(data)
|
||||
result = KernelsArgs.disable_mlp_kernel(data)
|
||||
assert result["lora_mlp_kernel"] is True
|
||||
|
||||
|
||||
|
||||
428
tests/integrations/test_sonicmoe.py
Normal file
428
tests/integrations/test_sonicmoe.py
Normal file
@@ -0,0 +1,428 @@
|
||||
"""Unit tests for the SonicMoE integration."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from axolotl.integrations.kernels.args import KernelsArgs
|
||||
from axolotl.integrations.kernels.sonicmoe.routing import (
|
||||
sigmoid_topk_routing,
|
||||
softmax_topk_routing,
|
||||
)
|
||||
from axolotl.integrations.kernels.sonicmoe.weight_converter import (
|
||||
ConcatenatedToInterleaved,
|
||||
InterleavedToConcatenated,
|
||||
register_sonicmoe_weight_converter,
|
||||
)
|
||||
|
||||
|
||||
class TestKernelsArgs:
|
||||
def test_mutual_exclusivity_raises(self):
|
||||
with pytest.raises(ValueError, match="Cannot use both"):
|
||||
KernelsArgs.model_validate({"use_scattermoe": True, "use_sonicmoe": True})
|
||||
|
||||
def test_sonicmoe_only(self):
|
||||
result = KernelsArgs.model_validate({"use_sonicmoe": True})
|
||||
assert result.use_sonicmoe is True
|
||||
assert result.use_scattermoe is None
|
||||
|
||||
def test_scattermoe_only(self):
|
||||
result = KernelsArgs.model_validate({"use_scattermoe": True})
|
||||
assert result.use_scattermoe is True
|
||||
assert result.use_sonicmoe is None
|
||||
|
||||
def test_neither_set(self):
|
||||
result = KernelsArgs.model_validate({})
|
||||
assert result.use_scattermoe is None
|
||||
assert result.use_sonicmoe is None
|
||||
|
||||
def test_disables_mlp_kernel_when_sonicmoe(self):
|
||||
data = {"use_sonicmoe": True, "lora_mlp_kernel": True}
|
||||
result = KernelsArgs.disable_mlp_kernel(data)
|
||||
assert result["lora_mlp_kernel"] is False
|
||||
assert result["mlp_kernel"] is False
|
||||
|
||||
|
||||
class TestConcatenatedToInterleaved:
|
||||
@pytest.fixture
|
||||
def sample_tensor(self):
|
||||
"""Create a test tensor [E=2, 2*I=4, H=3] with distinct gate/up values."""
|
||||
E, I, H = 2, 2, 3 # noqa: E741
|
||||
gate = torch.arange(1, E * I * H + 1, dtype=torch.float32).reshape(E, I, H)
|
||||
up = torch.arange(100, 100 + E * I * H, dtype=torch.float32).reshape(E, I, H)
|
||||
return torch.cat([gate, up], dim=1)
|
||||
|
||||
def test_interleave_rows_alternate(self, sample_tensor):
|
||||
op = ConcatenatedToInterleaved(dim=1)
|
||||
result = op.convert(
|
||||
{"test": sample_tensor},
|
||||
source_patterns=["test"],
|
||||
target_patterns=["test"],
|
||||
)
|
||||
interleaved = result["test"]
|
||||
|
||||
# For expert 0: even rows should be gate, odd rows should be up
|
||||
E, two_I, H = sample_tensor.shape
|
||||
I = two_I // 2 # noqa: E741
|
||||
gate_orig = sample_tensor[:, :I, :]
|
||||
up_orig = sample_tensor[:, I:, :]
|
||||
|
||||
assert torch.equal(interleaved[:, 0::2, :], gate_orig)
|
||||
assert torch.equal(interleaved[:, 1::2, :], up_orig)
|
||||
|
||||
def test_interleave_handles_list_input(self, sample_tensor):
|
||||
op = ConcatenatedToInterleaved(dim=1)
|
||||
result = op.convert(
|
||||
{"test": [sample_tensor]},
|
||||
source_patterns=["test"],
|
||||
target_patterns=["test"],
|
||||
)
|
||||
assert result["test"].shape == sample_tensor.shape
|
||||
|
||||
def test_reverse_op_type(self):
|
||||
op = ConcatenatedToInterleaved(dim=1)
|
||||
assert isinstance(op.reverse_op, InterleavedToConcatenated)
|
||||
assert op.reverse_op.dim == 1
|
||||
|
||||
|
||||
class TestInterleavedToConcatenated:
|
||||
@pytest.fixture
|
||||
def interleaved_tensor(self):
|
||||
"""Create an interleaved tensor [E=2, 2*I=4, H=3]."""
|
||||
E, I, H = 2, 2, 3 # noqa: E741
|
||||
gate = torch.arange(1, E * I * H + 1, dtype=torch.float32).reshape(E, I, H)
|
||||
up = torch.arange(100, 100 + E * I * H, dtype=torch.float32).reshape(E, I, H)
|
||||
interleaved = torch.empty(E, 2 * I, H)
|
||||
interleaved[:, 0::2, :] = gate
|
||||
interleaved[:, 1::2, :] = up
|
||||
return interleaved
|
||||
|
||||
def test_deinterleave_gate_up_separated(self, interleaved_tensor):
|
||||
op = InterleavedToConcatenated(dim=1)
|
||||
result = op.convert(
|
||||
{"test": interleaved_tensor},
|
||||
source_patterns=["test"],
|
||||
target_patterns=["test"],
|
||||
)
|
||||
concatenated = result["test"]
|
||||
|
||||
E, two_I, H = concatenated.shape
|
||||
I = two_I // 2 # noqa: E741
|
||||
|
||||
# First half should be gate (even rows from interleaved)
|
||||
assert torch.equal(concatenated[:, :I, :], interleaved_tensor[:, 0::2, :])
|
||||
# Second half should be up (odd rows from interleaved)
|
||||
assert torch.equal(concatenated[:, I:, :], interleaved_tensor[:, 1::2, :])
|
||||
|
||||
def test_reverse_op_type(self):
|
||||
op = InterleavedToConcatenated(dim=1)
|
||||
assert isinstance(op.reverse_op, ConcatenatedToInterleaved)
|
||||
assert op.reverse_op.dim == 1
|
||||
|
||||
|
||||
class TestRoundTrip:
|
||||
@pytest.fixture
|
||||
def concat_tensor(self):
|
||||
E, I, H = 4, 8, 16 # noqa: E741
|
||||
gate = torch.randn(E, I, H)
|
||||
up = torch.randn(E, I, H)
|
||||
return torch.cat([gate, up], dim=1)
|
||||
|
||||
def test_interleave_then_deinterleave_is_identity(self, concat_tensor):
|
||||
fwd = ConcatenatedToInterleaved(dim=1)
|
||||
rev = InterleavedToConcatenated(dim=1)
|
||||
|
||||
interleaved = fwd.convert(
|
||||
{"k": concat_tensor}, source_patterns=["k"], target_patterns=["k"]
|
||||
)["k"]
|
||||
recovered = rev.convert(
|
||||
{"k": interleaved}, source_patterns=["k"], target_patterns=["k"]
|
||||
)["k"]
|
||||
|
||||
assert torch.equal(concat_tensor, recovered)
|
||||
|
||||
def test_reverse_op_chain_is_identity(self, concat_tensor):
|
||||
"""Verify that op.reverse_op produces an exact inverse."""
|
||||
op = ConcatenatedToInterleaved(dim=1)
|
||||
rev = op.reverse_op
|
||||
|
||||
interleaved = op.convert(
|
||||
{"k": concat_tensor}, source_patterns=["k"], target_patterns=["k"]
|
||||
)["k"]
|
||||
recovered = rev.convert(
|
||||
{"k": interleaved}, source_patterns=["k"], target_patterns=["k"]
|
||||
)["k"]
|
||||
|
||||
assert torch.equal(concat_tensor, recovered)
|
||||
|
||||
def test_various_shapes(self):
|
||||
"""Test with different expert counts and dimensions."""
|
||||
fwd = ConcatenatedToInterleaved(dim=1)
|
||||
rev = InterleavedToConcatenated(dim=1)
|
||||
|
||||
for E, I, H in [(1, 4, 8), (8, 16, 32), (16, 128, 256)]: # noqa: E741
|
||||
concat = torch.randn(E, 2 * I, H)
|
||||
interleaved = fwd.convert(
|
||||
{"k": concat}, source_patterns=["k"], target_patterns=["k"]
|
||||
)["k"]
|
||||
recovered = rev.convert(
|
||||
{"k": interleaved}, source_patterns=["k"], target_patterns=["k"]
|
||||
)["k"]
|
||||
assert torch.equal(concat, recovered), (
|
||||
f"Failed for shape ({E}, {2 * I}, {H})"
|
||||
)
|
||||
|
||||
|
||||
class TestWeightConverterRegistration:
|
||||
def test_register_appends_interleave_op(self):
|
||||
from transformers.conversion_mapping import get_checkpoint_conversion_mapping
|
||||
|
||||
register_sonicmoe_weight_converter("qwen3_moe")
|
||||
|
||||
modified = get_checkpoint_conversion_mapping("qwen3_moe")
|
||||
# Find the gate_up_proj converter
|
||||
gate_up_converter = None
|
||||
for conv in modified:
|
||||
if hasattr(conv, "operations") and any(
|
||||
"gate_up_proj" in pat for pat in conv.target_patterns
|
||||
):
|
||||
gate_up_converter = conv
|
||||
break
|
||||
|
||||
assert gate_up_converter is not None
|
||||
assert isinstance(gate_up_converter.operations[-1], ConcatenatedToInterleaved)
|
||||
|
||||
def test_double_registration_is_idempotent(self):
|
||||
from transformers.conversion_mapping import get_checkpoint_conversion_mapping
|
||||
|
||||
register_sonicmoe_weight_converter("qwen3_moe")
|
||||
register_sonicmoe_weight_converter("qwen3_moe")
|
||||
|
||||
modified = get_checkpoint_conversion_mapping("qwen3_moe")
|
||||
for conv in modified:
|
||||
if hasattr(conv, "operations") and any(
|
||||
"gate_up_proj" in pat for pat in conv.target_patterns
|
||||
):
|
||||
interleave_count = sum(
|
||||
isinstance(op, ConcatenatedToInterleaved) for op in conv.operations
|
||||
)
|
||||
assert interleave_count == 1, (
|
||||
f"Expected 1 ConcatenatedToInterleaved op, got {interleave_count}"
|
||||
)
|
||||
break
|
||||
|
||||
def test_register_unsupported_model_type_warns(self):
|
||||
# A model type with no conversion mapping should warn but not raise
|
||||
register_sonicmoe_weight_converter("nonexistent_model_type_xyz")
|
||||
|
||||
|
||||
def _make_qwen_moe_block(T=8, H=16, E=4, K=2):
|
||||
"""Create a mock qwen-style MoE block for routing tests."""
|
||||
gate = SimpleNamespace(
|
||||
weight=torch.randn(E, H),
|
||||
top_k=K,
|
||||
num_experts=E,
|
||||
norm_topk_prob=True,
|
||||
)
|
||||
return SimpleNamespace(gate=gate), T, H, E, K
|
||||
|
||||
|
||||
def _make_glm_moe_block(T=8, H=16, E=16, K=4, n_group=2, topk_group=1):
|
||||
"""Create a mock GLM5-style MoE block for routing tests."""
|
||||
gate = SimpleNamespace(
|
||||
weight=torch.randn(E, H),
|
||||
e_score_correction_bias=torch.zeros(E),
|
||||
)
|
||||
moe_block = SimpleNamespace(
|
||||
gate=gate,
|
||||
top_k=K,
|
||||
n_routed_experts=E,
|
||||
n_group=n_group,
|
||||
topk_group=topk_group,
|
||||
norm_topk_prob=True,
|
||||
routed_scaling_factor=1.0,
|
||||
)
|
||||
return moe_block, T, H, E, K
|
||||
|
||||
|
||||
def _make_minimax_m2_moe_block(T=8, H=16, E=16, K=4):
|
||||
"""Create a mock minimax_m2-style MoE block for routing tests.
|
||||
|
||||
minimax_m2 uses sigmoid->topk WITHOUT group selection:
|
||||
- e_score_correction_bias is on the moe_block (not on gate)
|
||||
- No n_group / topk_group attributes
|
||||
- Always normalizes (norm_topk_prob defaults to True)
|
||||
- No routed_scaling_factor (defaults to 1.0)
|
||||
"""
|
||||
gate = SimpleNamespace(
|
||||
weight=torch.randn(E, H),
|
||||
top_k=K,
|
||||
)
|
||||
moe_block = SimpleNamespace(
|
||||
gate=gate,
|
||||
top_k=K,
|
||||
e_score_correction_bias=torch.zeros(E),
|
||||
)
|
||||
return moe_block, T, H, E, K
|
||||
|
||||
|
||||
class TestSoftmaxTopkRouting:
|
||||
def test_output_shapes(self):
|
||||
moe_block, T, H, E, K = _make_qwen_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
scores, token_idx, expert_idx, logits = softmax_topk_routing(hidden, moe_block)
|
||||
|
||||
assert scores.shape == (T * K,)
|
||||
assert token_idx.shape == (T * K,)
|
||||
assert expert_idx.shape == (T * K,)
|
||||
assert logits.shape == (T, E)
|
||||
|
||||
def test_scores_are_float32(self):
|
||||
moe_block, T, H, E, K = _make_qwen_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
scores, _, _, _ = softmax_topk_routing(hidden, moe_block)
|
||||
assert scores.dtype == torch.float32
|
||||
|
||||
def test_token_indices_sorted_ascending(self):
|
||||
moe_block, T, H, E, K = _make_qwen_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
_, token_idx, _, _ = softmax_topk_routing(hidden, moe_block)
|
||||
|
||||
# Token indices must be sorted ascending (SonicMoE requirement)
|
||||
diffs = token_idx[1:] - token_idx[:-1]
|
||||
assert (diffs >= 0).all()
|
||||
|
||||
def test_expert_indices_in_range(self):
|
||||
moe_block, T, H, E, K = _make_qwen_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
_, _, expert_idx, _ = softmax_topk_routing(hidden, moe_block)
|
||||
|
||||
assert (expert_idx >= 0).all()
|
||||
assert (expert_idx < E).all()
|
||||
|
||||
def test_renormalized_scores_sum_to_one(self):
|
||||
moe_block, T, H, E, K = _make_qwen_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
scores, _, _, _ = softmax_topk_routing(hidden, moe_block)
|
||||
per_token_sums = scores.reshape(T, K).sum(dim=-1)
|
||||
assert torch.allclose(per_token_sums, torch.ones(T), atol=1e-5)
|
||||
|
||||
|
||||
class TestSigmoidTopkRouting:
|
||||
def test_output_shapes(self):
|
||||
moe_block, T, H, E, K = _make_glm_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
scores, token_idx, expert_idx, logits = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
assert scores.shape == (T * K,)
|
||||
assert token_idx.shape == (T * K,)
|
||||
assert expert_idx.shape == (T * K,)
|
||||
assert logits.shape == (T, E)
|
||||
|
||||
def test_scores_are_float32(self):
|
||||
moe_block, T, H, E, K = _make_glm_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
scores, _, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
assert scores.dtype == torch.float32
|
||||
|
||||
def test_token_indices_sorted_ascending(self):
|
||||
moe_block, T, H, E, K = _make_glm_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
_, token_idx, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
diffs = token_idx[1:] - token_idx[:-1]
|
||||
assert (diffs >= 0).all()
|
||||
|
||||
def test_expert_indices_in_range(self):
|
||||
moe_block, T, H, E, K = _make_glm_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
_, _, expert_idx, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
assert (expert_idx >= 0).all()
|
||||
assert (expert_idx < E).all()
|
||||
|
||||
def test_scores_are_nonnegative(self):
|
||||
"""Sigmoid outputs are in [0, 1], so scores should be non-negative."""
|
||||
moe_block, T, H, E, K = _make_glm_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
scores, _, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
assert (scores >= 0).all()
|
||||
|
||||
def test_scaling_factor_applied(self):
|
||||
moe_block, T, H, E, K = _make_glm_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
# Get scores with scaling_factor=1.0
|
||||
scores_1x, _, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
# Get scores with scaling_factor=2.0
|
||||
moe_block.routed_scaling_factor = 2.0
|
||||
scores_2x, _, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
assert torch.allclose(scores_2x, scores_1x * 2.0, atol=1e-5)
|
||||
|
||||
def test_group_selection_restricts_experts(self):
|
||||
"""With n_group=4 and topk_group=1, only 1/4 of experts should be selectable."""
|
||||
moe_block, T, H, E, K = _make_glm_moe_block(E=16, K=2, n_group=4, topk_group=1)
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
_, _, expert_idx, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
# Each token's experts should all fall within a single group (size E//n_group=4)
|
||||
expert_idx_2d = expert_idx.reshape(T, K)
|
||||
for t in range(T):
|
||||
experts = expert_idx_2d[t]
|
||||
groups = experts // (E // moe_block.n_group)
|
||||
# All selected experts should be from the same group
|
||||
assert (groups == groups[0]).all()
|
||||
|
||||
|
||||
class TestMiniMaxM2SigmoidRouting:
|
||||
"""Tests for minimax_m2 routing: sigmoid->topk without group selection."""
|
||||
|
||||
def test_output_shapes(self):
|
||||
"""Validates getattr defaults work: n_group=1, E from gate.weight.shape[0]."""
|
||||
moe_block, T, H, E, K = _make_minimax_m2_moe_block()
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
scores, token_idx, expert_idx, logits = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
assert scores.shape == (T * K,)
|
||||
assert token_idx.shape == (T * K,)
|
||||
assert expert_idx.shape == (T * K,)
|
||||
assert logits.shape == (T, E)
|
||||
|
||||
def test_bias_on_block_not_gate(self):
|
||||
"""Verify that e_score_correction_bias on the block (not gate) is used."""
|
||||
T, H, E, K = 8, 16, 8, 2
|
||||
gate = SimpleNamespace(
|
||||
weight=torch.randn(E, H),
|
||||
top_k=K,
|
||||
)
|
||||
# Large positive bias on expert 0 should make it selected more often
|
||||
bias = torch.zeros(E)
|
||||
bias[0] = 100.0
|
||||
moe_block = SimpleNamespace(
|
||||
gate=gate,
|
||||
top_k=K,
|
||||
e_score_correction_bias=bias,
|
||||
)
|
||||
hidden = torch.randn(T, H)
|
||||
|
||||
_, _, expert_idx, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
|
||||
# Expert 0 should appear for every token due to the large bias
|
||||
expert_idx_2d = expert_idx.reshape(T, K)
|
||||
for t in range(T):
|
||||
assert 0 in expert_idx_2d[t]
|
||||
158
tests/integrations/test_sonicmoe_gradients.py
Normal file
158
tests/integrations/test_sonicmoe_gradients.py
Normal file
@@ -0,0 +1,158 @@
|
||||
"""
|
||||
Gradient correctness tests for SonicMoE routing functions (CPU-only).
|
||||
|
||||
Uses torch.autograd.gradcheck with float32 inputs to match the production
|
||||
code path where routing happens in float32.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from axolotl.integrations.kernels.sonicmoe.routing import (
|
||||
sigmoid_topk_routing,
|
||||
softmax_topk_routing,
|
||||
)
|
||||
|
||||
_GC_EPS = 1e-3
|
||||
_GC_ATOL = 1e-3
|
||||
_GC_RTOL = 1e-3
|
||||
|
||||
|
||||
def _make_softmax_moe_block(weight):
|
||||
gate = torch.nn.Module()
|
||||
gate.weight = weight
|
||||
gate.top_k = 2
|
||||
gate.norm_topk_prob = True
|
||||
|
||||
moe_block = torch.nn.Module()
|
||||
moe_block.gate = gate
|
||||
return moe_block
|
||||
|
||||
|
||||
def _make_sigmoid_moe_block(weight, bias):
|
||||
gate = torch.nn.Module()
|
||||
gate.weight = weight
|
||||
gate.e_score_correction_bias = bias
|
||||
|
||||
moe_block = torch.nn.Module()
|
||||
moe_block.gate = gate
|
||||
moe_block.top_k = 2
|
||||
moe_block.n_routed_experts = weight.shape[0]
|
||||
moe_block.n_group = 1
|
||||
moe_block.norm_topk_prob = True
|
||||
moe_block.routed_scaling_factor = 1.0
|
||||
return moe_block
|
||||
|
||||
|
||||
class TestSoftmaxTopkRoutingGradcheck:
|
||||
"""Numerical gradient verification for softmax_topk_routing."""
|
||||
|
||||
def test_gradcheck_wrt_gate_weight(self):
|
||||
T, H, E = 4, 8, 4
|
||||
|
||||
hidden = torch.randn(T, H, dtype=torch.float32)
|
||||
|
||||
def fn(weight):
|
||||
moe_block = _make_softmax_moe_block(weight)
|
||||
scores, _, _, _ = softmax_topk_routing(hidden, moe_block)
|
||||
return scores
|
||||
|
||||
weight = torch.randn(E, H, dtype=torch.float32, requires_grad=True)
|
||||
torch.autograd.gradcheck(
|
||||
fn, (weight,), eps=_GC_EPS, atol=_GC_ATOL, rtol=_GC_RTOL
|
||||
)
|
||||
|
||||
def test_gradcheck_wrt_hidden_states(self):
|
||||
T, H, E = 4, 8, 4
|
||||
|
||||
weight = torch.randn(E, H, dtype=torch.float32)
|
||||
moe_block = _make_softmax_moe_block(weight)
|
||||
|
||||
def fn(hidden):
|
||||
scores, _, _, _ = softmax_topk_routing(hidden, moe_block)
|
||||
return scores
|
||||
|
||||
hidden = torch.randn(T, H, dtype=torch.float32, requires_grad=True)
|
||||
torch.autograd.gradcheck(
|
||||
fn, (hidden,), eps=_GC_EPS, atol=_GC_ATOL, rtol=_GC_RTOL
|
||||
)
|
||||
|
||||
def test_gradcheck_wrt_router_logits(self):
|
||||
T, H, E = 4, 8, 4
|
||||
|
||||
hidden = torch.randn(T, H, dtype=torch.float32)
|
||||
|
||||
def fn(weight):
|
||||
moe_block = _make_softmax_moe_block(weight)
|
||||
_, _, _, router_logits = softmax_topk_routing(hidden, moe_block)
|
||||
return router_logits
|
||||
|
||||
weight = torch.randn(E, H, dtype=torch.float32, requires_grad=True)
|
||||
torch.autograd.gradcheck(
|
||||
fn, (weight,), eps=_GC_EPS, atol=_GC_ATOL, rtol=_GC_RTOL
|
||||
)
|
||||
|
||||
def test_no_norm_variant(self):
|
||||
T, H, E = 4, 8, 4
|
||||
|
||||
hidden = torch.randn(T, H, dtype=torch.float32)
|
||||
|
||||
def fn(weight):
|
||||
moe_block = _make_softmax_moe_block(weight)
|
||||
moe_block.gate.norm_topk_prob = False
|
||||
scores, _, _, _ = softmax_topk_routing(hidden, moe_block)
|
||||
return scores
|
||||
|
||||
weight = torch.randn(E, H, dtype=torch.float32, requires_grad=True)
|
||||
torch.autograd.gradcheck(
|
||||
fn, (weight,), eps=_GC_EPS, atol=_GC_ATOL, rtol=_GC_RTOL
|
||||
)
|
||||
|
||||
|
||||
class TestSigmoidTopkRoutingGradcheck:
|
||||
"""Numerical gradient verification for sigmoid_topk_routing."""
|
||||
|
||||
def test_gradcheck_wrt_gate_weight(self):
|
||||
T, H, E = 4, 8, 4
|
||||
|
||||
hidden = torch.randn(T, H, dtype=torch.float32)
|
||||
bias = torch.zeros(E, dtype=torch.float32)
|
||||
|
||||
def fn(weight):
|
||||
moe_block = _make_sigmoid_moe_block(weight, bias)
|
||||
scores, _, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
return scores
|
||||
|
||||
weight = torch.randn(E, H, dtype=torch.float32, requires_grad=True)
|
||||
torch.autograd.gradcheck(
|
||||
fn, (weight,), eps=_GC_EPS, atol=_GC_ATOL, rtol=_GC_RTOL
|
||||
)
|
||||
|
||||
def test_gradcheck_wrt_hidden_states(self):
|
||||
T, H, E = 4, 8, 4
|
||||
|
||||
weight = torch.randn(E, H, dtype=torch.float32)
|
||||
bias = torch.zeros(E, dtype=torch.float32)
|
||||
moe_block = _make_sigmoid_moe_block(weight, bias)
|
||||
|
||||
def fn(hidden):
|
||||
scores, _, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
return scores
|
||||
|
||||
hidden = torch.randn(T, H, dtype=torch.float32, requires_grad=True)
|
||||
torch.autograd.gradcheck(
|
||||
fn, (hidden,), eps=_GC_EPS, atol=_GC_ATOL, rtol=_GC_RTOL
|
||||
)
|
||||
|
||||
def test_gradcheck_wrt_bias(self):
|
||||
T, H, E = 4, 8, 4
|
||||
|
||||
hidden = torch.randn(T, H, dtype=torch.float32)
|
||||
weight = torch.randn(E, H, dtype=torch.float32)
|
||||
|
||||
def fn(bias):
|
||||
moe_block = _make_sigmoid_moe_block(weight, bias)
|
||||
scores, _, _, _ = sigmoid_topk_routing(hidden, moe_block)
|
||||
return scores
|
||||
|
||||
bias = torch.zeros(E, dtype=torch.float32, requires_grad=True)
|
||||
torch.autograd.gradcheck(fn, (bias,), eps=_GC_EPS, atol=_GC_ATOL, rtol=_GC_RTOL)
|
||||
Reference in New Issue
Block a user