AutoEP: MiniMax-M3 preset (#8651)
# AutoEP: MiniMax-M3 preset
## What
A built-in preset `minimax_m3` for MiniMax-M3 (`model_type =
"minimax_m3_vl_text"`, transformers >= 5.15.0).
MiniMax-M3 is 425.9 B parameters, 25.7 B active: 60 layers, the first 3
dense, the other 57 with 128
experts, top-4, and one shared expert. With the preset,
`expert_parallel: {"enabled": true, "autoep_size": N}`
is the whole configuration; nothing has to be spelled out.
| preset field | value | where it comes from |
|---|---|---|
| MoE block, router, experts, shared expert | `model.layers.N.mlp`,
`gate`, `experts`, `shared_experts` | `MiniMaxM3VLSparseMoeBlock` |
| expert weights | fused 3-D `gate_up_proj` / `down_proj` |
`MiniMaxM3VLExperts` |
| expert and top-k counts | `num_local_experts` / `num_experts_per_tok`
on the model config | `MiniMaxM3VLTextConfig` (128 / 4) |
| scoring | sigmoid, applied after the experts, selected top-k weights
normalised | `MiniMaxM3VLTopKRouter` |
| selection bias | `gate.e_score_correction_bias` (a buffer, as in
DeepSeek-V3) | found by the existing bias-path resolution |
| routed output scale | `routed_scaling_factor` = 2.0 | read from the
model config by the existing code |
| expert activation | `swiglu_oai`, alpha 1.702, limit 7.0 | preset; the
model's `swiglu_alpha` / `swiglu_limit` override it at detection |
| router logits | recorded at the model level as
`OutputRecorder(MiniMaxM3VLTopKRouter, index=0)`, the raw logits | same
pattern as Mixtral, so the adapter is a
`TransformersTopLevelRouterLogitsAdapter` |
| version gate | transformers >= 5.15.0, checked when the model type is
MiniMax's | adapter opts in |
| composite `minimax_m3_vl` (vision + text) | refused with a note to
pass the text backbone | `unsupported_hf_model_type_notes` |
| `load_balance_coeff` | defaults to `None`; a non-null value is
rejected, as for DeepSeek-V3 | `autoep_config_defaults`,
`supports_expert_bias=False` |
## Trained at scale: the full MiniMax-M3 on 65,536-token books, 4 nodes
of 8 H200
This preset is the configuration under which the real MiniMax-M3
checkpoint fine-tuned on real
65,536-token sequences. The run used the preset before it was committed;
the preset fields, the
`swiglu_oai` form with alpha 1.702 / limit 7.0, and the config keys are
the ones in this PR.
Setup: 4 nodes of 8 H200 (32 ranks), all 60 layers, the real checkpoint
(795.5 GiB in 59 safetensors
shards, streamed in one tensor at a time in 791 s), 426.2B parameters of
which 413.1B are expert weights.
`expert_parallel: {"enabled": true, "autoep_size": 4}`: 8 expert groups
of 4 ranks, each rank holding 32
experts that ZeRO-3 shards 8 ways across the expert-data-parallel group.
`fused_weighted_sum` combine,
ZeRO-3 with the optimizer state offloaded to host memory, activation
checkpointing with the checkpoint
inputs on host, tiled loss (64 shards over the 200,064-token
vocabulary), micro-batch 1 per rank = one
whole PG-19 book of 65,536 tokens per rank, learning rate 1e-6, 1
warm-up + 2 measured steps. torch
2.8.0+cu126, transformers 5.15.1. Attention ran through a harness-side
gather over each query's selected
keys, because the stock transformers path materialises a `[1, 64, S, S]`
block mask and stops at 14,336
tokens on a 141 GB card; that module is outside this PR.
| step | loss (rank 0's book) | grad norm | step time | rank 0 peak
allocated |
|---|---|---|---|---|
| 1 (warm-up) | 0.0153 (the King James Bible) | 0.98 | 361.0 s | 89.1
GiB |
| 2 | 0.8592 | 1.18 | 296.1 s | 90.6 GiB |
| 3 | 1.0754 | 1.23 | 292.3 s | 90.6 GiB |
Measured average 294.22 s per step = 7,128 tokens/s over the 32 ranks.
Memory: rank 0 peak 90.57 GiB
allocated / 109.78 GiB reserved; the highest rank 102.07 GiB allocated /
126.24 GiB reserved of 139.80.
Host memory 1.87 TB per node, most of it the offloaded optimizer state.
The losses are a trained model's losses on real text (0.015 on a text
every corpus contains, 0.86 and
1.08 on two other pre-1919 books), against 13.44 with random weights on
the same setup. That is the
end-to-end check of everything the preset describes: the conversion of
the checkpoint's per-expert
`w1`/`w3`/`w2` into AutoEP's fused expert layout, the sigmoid router
with the fp32 selection bias,
`routed_scaling_factor` 2.0, the shared expert, and the clamped
activation. With standard SwiGLU in
place of `swiglu_oai` the replaced block differs from the model's own by
0.41 mean absolute (correlation
0.74), so the activation form is what makes this loss a real loss.
Why expert parallel 4 and not 32: real router weights route far more
unevenly than random ones. Over all
57 MoE layers of this run, the busiest of the 4 ranks in an expert group
received 1.09 to 2.22 times the
mean number of rows (median 1.36; the worst is layer 54, 581,906 rows on
one rank against 262,144 own
rows). The same counts grouped as expert parallel 32 give a median of
4.20 and a maximum of 7.84 times
the mean, and at that skew the busiest rank runs out of memory at this
length:, same setup with
expert parallel 32, failed in the step-1 forward on one rank with a
single 11.96 GiB temporary for
1,045,081 routed rows inside the expert activation, 129.50 GiB already
in use. AutoEP has no capacity
limit, so the busiest rank sets the memory ceiling; a smaller expert
group spreads each expert set over
more ranks, at the cost of a per-layer all-gather of the local experts
(32 experts x 113 MB at expert
parallel 4).
## Why a preset
Without it, MiniMax-M3 runs through the custom `moe_layer_pattern` path
with six keys (`score_func`,
`score_apply`, `route_norm`, `has_shared_experts`,
`shared_experts_pattern`, `expert_activation`), and a
user who omits `expert_activation` is stopped by the registry PR's check
rather than silently trained on
`silu(gate) * up`. The preset makes the model type enough, records the
facts above in one place, and
switches on the version gate for it. The run above is this
configuration.
## Files
| file | change |
|---|---|
| `deepspeed/module_inject/auto_ep_presets/minimax_m3.py` | new: the
preset and `MiniMaxM3PresetAdapter` |
| `deepspeed/module_inject/auto_ep_presets/registry.py` | registers the
module |
| `docs/_pages/config-json.md`, `docs/code-docs/source/autoep.rst` |
`minimax_m3` in the preset list and the version table |
| `tests/unit/v1/moe/autoep_test_utils.py` | `tiny_minimax_m3_config` |
| `tests/unit/v1/moe/test_autoep_unit.py` | registry contract extended;
`test_hf_minimax_m3_causal_lm_matches_autoep_with_router_logits`;
`test_minimax_m3_adapter_guards` |
| `tests/unit/v1/moe/test_autoep_expert_activation.py` | the
block-parity test now runs through the preset as well as the custom
pattern; "every preset is swiglu" becomes "every preset names its form"
|
## Tests
| test | what it checks, and the wrong behaviour it would catch |
|---|---|
| `test_preset_registry_core_contracts` | `minimax_m3` registered;
`minimax_m3_vl_text` resolves to it; the composite `minimax_m3_vl`
carries a note naming the text backbone |
| `test_minimax_m3_adapter_guards` | a mock model of type
`minimax_m3_vl_text` is detected as `minimax_m3` with `swiglu_oai`; an
installed transformers below 5.15.0 is refused with the version named
(the gate is switched on for this preset; an unregistered adapter would
silently skip it); the composite type is refused with the note |
| `test_hf_minimax_m3_causal_lm_matches_autoep_with_router_logits` | a
tiny `MiniMaxM3VLForCausalLM` (1 layer, 4 experts, top-2),
state-matched, native against AutoEP: loss, the recorded router logits,
and the auxiliary load-balancing loss agree at rtol 1e-5 (the recorder
retarget and the raw-logit capture mode are what make the last two
agree) |
| `TestMiniMaxM3Block.test_replaced_block_matches_hf_block[preset-*]` |
a real `MiniMaxM3VLSparseMoeBlock` found from the model type alone, no
config keys: family `minimax_m3`, activation `swiglu_oai` with 1.702 /
7.0 from the model, and the replaced block matches a deep copy of the
original, forward and input gradient, at weight std 0.02 and 0.5 |
| `TestMiniMaxM3Block.*[custom_pattern-*]`,
`test_custom_pattern_without_the_key_is_refused` | the custom-pattern
route from the registry PR, kept for a model type AutoEP does not know
yet |
Results, all on this branch:
| where | environment | `test_autoep_expert_activation.py` |
`test_autoep_unit.py` | `test_autoep_grad_parity.py` |
|---|---|---|---|---|
| CPU, one process | torch 2.9.1, transformers 5.15.0 | 29 passed, 8
skipped (accelerator cases) | 186 passed, 6 skipped | not run (needs 4
devices) |
| H200 node, 4 GPUs | torch 2.9.1+cu128, transformers 5.15.0, Triton
3.5.1 | **37 passed** | **192 passed** | **23 passed, 8 skipped** (the
regional-compile parity cases need `DEEPSPEED_RUN_H100_TESTS=1`), 2 min
44 s |
`pre-commit run --files` on the changed files: yapf, flake8, codespell,
check-torchdist, check-torchcuda
pass; `check-license` reports the new preset file only because the hook
uses `git grep`, which skips
untracked files.
Signed-off-by: pengdurice <pengduhit@gmail.com>
