minimind/model
basil-k-aji-dev 1c876435ef [fix] keep the MoE router trainable at num_experts_per_tok=1
Normalising a single top-k weight by its own sum gives exactly 1.0 for every
token, so the gate leaves the graph entirely: topk_idx is a discrete index and
the weight is now a constant, which leaves the LM loss no differentiable path
to the router. It then trains on aux_loss alone, whose optimum is a uniform
distribution, and the experts never specialise.

Measured on MOEFeedForward with router_aux_loss_coef=0, so only the LM path can
contribute a gradient:

  num_experts_per_tok=1, norm_topk_prob=True   7.02e-10   <- current default
  num_experts_per_tok=1, norm_topk_prob=False  1.02e-02
  num_experts_per_tok=2, norm_topk_prob=True   1.67e-02

Only normalise when there is more than one weight to normalise.

Note this changes the forward pass for existing num_experts_per_tok=1 MoE
checkpoints: the expert output was scaled by a constant 1.0 and is now scaled
by the gate probability.

Closes #858.
2026-09-17 21:31:41 +05:30
..
__init__.py 250426 2025-04-26 10:05:47 +08:00
model_lora.py [fix] lora dims 2026-05-31 13:38:51 +08:00
model_minimind.py [fix] keep the MoE router trainable at num_experts_per_tok=1 2026-09-17 21:31:41 +05:30
tokenizer.json [update] minimind-3 2026-03-25 23:57:45 +08:00
tokenizer_config.json [update] minimind-3 2026-03-25 23:57:45 +08:00