geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

llm: respect expert_gating_func (#17458)

Browse Source 
* llm: respect expert_gating_func

* test

* enum

* clean
This commit is contained in:
b1tg
2026-08-17 12:19:33 -07:00
committed by GitHub
parent 00d6eed43c
commit b757437f64
2 changed files with 52 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_llm_moe.py
+28 -1
View File
@@ -2,7 +2,7 @@ import unittest
import numpy as np
from dataclasses import replace
from tinygrad import Tensor
from tinygrad.llm.model import TransformerBlock, TransformerConfig
from tinygrad.llm.model import ExpertGating, TransformerBlock, TransformerConfig
def _moe_config(dim=8, hidden=16, n_heads=2, num_experts=4, num_experts_per_tok=2):
return TransformerConfig(
@@ -96,5 +96,32 @@ class TestMoEFeedForward(unittest.TestCase):
expected = moe_expected + shared_expected
np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
def test_moe_feed_forward_gating_funcs(self):
dim, hidden, n_heads = 8, 16, 2
num_experts, k = 4, 2
logits = np.array([4.0, 3.0, 0.0, -1.0], dtype=np.float32)
def softmax(x):
probs = np.exp(x - x.max())
return probs / probs.sum()
for gating_func in ExpertGating:
for norm_topk_prob in (False, True):
block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k),
expert_gating_func=gating_func, norm_topk_prob=norm_topk_prob))
block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
block.ffn_gate_inp.weight = Tensor((logits / dim)[None, :].repeat(dim, 0).T)
out = block._feed_forward(Tensor.ones(1, 1, dim)).numpy()[0, 0, 0]
if gating_func == ExpertGating.SOFTMAX: selection_scores = softmax(logits)
elif gating_func == ExpertGating.SIGMOID: selection_scores = 1 / (1 + np.exp(-logits))
elif gating_func == ExpertGating.SOFTMAX_WEIGHT: selection_scores = logits
else: selection_scores = np.sqrt(np.logaddexp(0, logits))
sel = np.argsort(selection_scores)[-k:]
weights = softmax(logits[sel]) if gating_func == ExpertGating.SOFTMAX_WEIGHT else selection_scores[sel]
if norm_topk_prob: weights /= weights.sum()
expected = (weights * (sel + 1)).sum() / (1 + np.exp(-1))
np.testing.assert_allclose(out, expected, rtol=1e-3)
if __name__ == '__main__':
unittest.main()
tinygrad/llm/model.py
+24 -9
View File
@@ -1,11 +1,17 @@
from __future__ import annotations
import functools, itertools, pathlib
import enum, functools, itertools, pathlib
from dataclasses import dataclass, replace
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
from tinygrad.nn import Linear
from tinygrad.llm.gguf import gguf_load
from tinygrad.uop.ops import resolve
class ExpertGating(enum.IntEnum):
SOFTMAX = 1
SIGMOID = 2
SOFTMAX_WEIGHT = 3 # softmax over the top-k selected logits
SQRT_SOFTPLUS = 4
@functools.cache
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
@@ -61,6 +67,7 @@ class TransformerConfig:
num_experts: int = 0
num_experts_per_tok: int = 0
norm_topk_prob: bool = False
expert_gating_func: ExpertGating = ExpertGating.SOFTMAX
q_lora_rank: int = 0
kv_lora_rank: int = 0
shared_expert_dim: int = 0
@@ -103,14 +110,21 @@ class FFNBlock:
if hasattr(self, 'ffn_gate_exps'):
h = x.unsqueeze(2) # (B, T, 1, D) - add expert dim for broadcasting
logits = self.ffn_gate_inp(x)
if hasattr(self, 'exp_probs_b'):
probs = logits.sigmoid()
_, sel = pairwise_topk(probs + self.exp_probs_b["bias"], self.config.num_experts_per_tok)
probs = probs.gather(-1, sel)
if self.config.norm_topk_prob: probs = probs / probs.sum(axis=-1, keepdim=True)
else:
vals, sel = pairwise_topk(logits, self.config.num_experts_per_tok)
probs = vals.softmax(-1) if self.config.norm_topk_prob else logits.softmax(-1).gather(-1, sel)
bias = self.exp_probs_b["bias"] if hasattr(self, 'exp_probs_b') else None
gating, normalize_topk = self.config.expert_gating_func, self.config.norm_topk_prob
# fast path: without selection bias, normalized SOFTMAX is equivalent to SOFTMAX_WEIGHT
if gating == ExpertGating.SOFTMAX and bias is None and normalize_topk:
gating, normalize_topk = ExpertGating.SOFTMAX_WEIGHT, False
if gating == ExpertGating.SOFTMAX_WEIGHT: scores = logits
elif gating == ExpertGating.SOFTMAX: scores = logits.softmax(-1)
elif gating == ExpertGating.SIGMOID: scores = logits.sigmoid()
elif gating == ExpertGating.SQRT_SOFTPLUS: scores = logits.softplus().sqrt()
_, sel = pairwise_topk(scores if bias is None else scores + bias, self.config.num_experts_per_tok)
probs = scores.gather(-1, sel)
# SOFTMAX_WEIGHT applies softmax after top-k selection
if gating == ExpertGating.SOFTMAX_WEIGHT: probs = probs.softmax(-1)
if normalize_topk: probs = probs / probs.sum(axis=-1, keepdim=True)
probs = probs * self.config.routed_scaling_factor
x_down = self.ffn_down_exps(sel, (self.ffn_gate_exps(sel, h).silu() * self.ffn_up_exps(sel, h)).contiguous()) # (B, T, k, D)
out = (x_down * probs.unsqueeze(-1)).sum(axis=2) # (B, T, D)
@@ -398,6 +412,7 @@ class Transformer:
qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0),
norm_topk_prob=kv.get(f'{arch}.expert_weights_norm', arch in ('qwen3moe', 'qwen35moe', 'kimi-linear')),
expert_gating_func=ExpertGating(kv.get(f'{arch}.expert_gating_func', ExpertGating.SOFTMAX)),
kv_lora_rank=kv_lora_rank, q_lora_rank=kv.get(f'{arch}.attention.q_lora_rank', 0),
leading_dense_blocks=kv.get(f'{arch}.leading_dense_block_count', 0),
shared_expert_dim=kv.get(