[Model][MiniMax-M3] Add pipeline parallelism support (#45810)
Signed-off-by: soaringk <k3vin.zhang@gmail.com>
soaringk committed
Jun 24, 2026 at 16:23 UTC
d7c1821b5a31c886cf130e50f353e49af5b79659
3 files changed
+124
-49
docs/models/supported_models.md
+2
-1
@@ -442,6 +442,7 @@ th {
442
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
443
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
444
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
445
+| `MiniMaxM3SparseForCausalLM` | MiniMax-M3 | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ |
446
| `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
447
| `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
448
| `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ |
@@ -592,7 +593,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
593
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
594
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
595
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
595
-| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I<sup>+</sup> + V<sup>+</sup> | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | |
596
+| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I<sup>+</sup> + V<sup>+</sup> | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ |
597
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
598
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
599
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
vllm/models/minimax_m3/amd/model.py
+61
-24
@@ -30,7 +30,7 @@ from vllm.config import (
30
VllmConfig,
31
get_current_vllm_config,
32
)
33
-from vllm.distributed import get_tensor_model_parallel_world_size
33
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
from vllm.forward_context import get_forward_context
35
from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -64,12 +64,15 @@ from vllm.model_executor.models.interfaces import (
64
MultiModalEmbeddings,
65
SupportsEagle3,
66
SupportsMultiModal,
67
+ SupportsPP,
68
)
69
from vllm.model_executor.models.utils import (
70
AutoWeightsLoader,
71
+ PPMissingLayer,
72
WeightsMapper,
73
init_vllm_registered_model,
74
is_pp_missing_parameter,
75
+ make_empty_intermediate_tensors_factory,
76
make_layers,
77
maybe_prefix,
78
)
@@ -92,6 +95,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
95
)
96
from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
97
from vllm.multimodal import MULTIMODAL_REGISTRY
98
+from vllm.sequence import IntermediateTensors
99
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
100
from vllm.v1.kv_cache_interface import (
101
FullAttentionSpec,
@@ -772,12 +776,15 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
776
777
self.vocab_size = config.vocab_size
778
775
- self.embed_tokens = VocabParallelEmbedding(
776
- config.vocab_size,
777
- config.hidden_size,
778
- quant_config=quant_config,
779
- prefix=f"{prefix}.embed_tokens",
780
- )
779
+ if get_pp_group().is_first_rank:
780
+ self.embed_tokens = VocabParallelEmbedding(
781
+ config.vocab_size,
782
+ config.hidden_size,
783
+ quant_config=quant_config,
784
+ prefix=f"{prefix}.embed_tokens",
785
+ )
786
+ else:
787
+ self.embed_tokens = PPMissingLayer()
788
789
# Reserved top-k indices buffer shared by all sparse-attention indexer
790
# layers (mirrors DeepseekV4); the indexer writes its per-head decode/
@@ -807,7 +814,13 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
814
prefix=f"{prefix}.layers",
815
)
816
810
- self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
817
+ if get_pp_group().is_last_rank:
818
+ self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
819
+ else:
820
+ self.norm = PPMissingLayer()
821
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
822
+ ["hidden_states", "residual"], config.hidden_size
823
+ )
824
825
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
826
return self.embed_tokens(input_ids)
@@ -816,13 +829,19 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
829
self,
830
input_ids: torch.Tensor | None,
831
positions: torch.Tensor,
832
+ intermediate_tensors: IntermediateTensors | None,
833
inputs_embeds: torch.Tensor | None = None,
820
- ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
821
- if inputs_embeds is not None:
822
- hidden_states = inputs_embeds
834
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
835
+ if get_pp_group().is_first_rank:
836
+ if inputs_embeds is not None:
837
+ hidden_states = inputs_embeds
838
+ else:
839
+ hidden_states = self.embed_input_ids(input_ids)
840
+ residual = None
841
else:
824
- hidden_states = self.embed_input_ids(input_ids)
825
- residual = None
842
+ assert intermediate_tensors is not None
843
+ hidden_states = intermediate_tensors["hidden_states"]
844
+ residual = intermediate_tensors["residual"]
845
846
# EAGLE3 is not yet compatible with pipeline parallel
847
aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
@@ -832,6 +851,11 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
851
aux_hidden_states, idx + 1, hidden_states, residual
852
)
853
854
+ if not get_pp_group().is_last_rank:
855
+ return IntermediateTensors(
856
+ {"hidden_states": hidden_states, "residual": residual}
857
+ )
858
+
859
hidden_states, _ = self.norm(hidden_states, residual)
860
861
if len(aux_hidden_states) > 0:
@@ -946,7 +970,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
970
return loaded_params
971
972
949
-class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
973
+class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
974
"""MiniMax M3 (sparse/dense backbone) for causal language modeling."""
975
976
packed_modules_mapping = {
@@ -963,13 +987,19 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
987
self.model = MiniMaxM3Model(
988
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
989
)
966
- self.lm_head = ParallelLMHead(
967
- config.vocab_size,
968
- config.hidden_size,
969
- quant_config=quant_config,
970
- prefix=maybe_prefix(prefix, "lm_head"),
971
- )
990
+ if get_pp_group().is_last_rank:
991
+ self.lm_head = ParallelLMHead(
992
+ config.vocab_size,
993
+ config.hidden_size,
994
+ quant_config=quant_config,
995
+ prefix=maybe_prefix(prefix, "lm_head"),
996
+ )
997
+ else:
998
+ self.lm_head = PPMissingLayer()
999
self.logits_processor = LogitsProcessor(config.vocab_size)
1000
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1001
+ self.model.make_empty_intermediate_tensors
1002
+ )
1003
1004
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1005
return self.model.embed_input_ids(input_ids)
@@ -978,10 +1008,11 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
1008
self,
1009
input_ids: torch.Tensor | None,
1010
positions: torch.Tensor,
1011
+ intermediate_tensors: IntermediateTensors | None = None,
1012
inputs_embeds: torch.Tensor | None = None,
1013
**kwargs,
983
- ) -> torch.Tensor:
984
- return self.model(input_ids, positions, inputs_embeds)
1014
+ ) -> torch.Tensor | IntermediateTensors:
1015
+ return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
1016
1017
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1018
return self.logits_processor(self.lm_head, hidden_states)
@@ -1004,7 +1035,7 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
1035
dummy_inputs=MiniMaxM3VLDummyInputsBuilder,
1036
)
1037
class MiniMaxM3SparseForConditionalGeneration(
1007
- nn.Module, SupportsMultiModal, SupportsEagle3
1038
+ nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3
1039
):
1040
"""Top-level (VL) entry point for MiniMax M3.
1041
@@ -1070,6 +1101,9 @@ class MiniMaxM3SparseForConditionalGeneration(
1101
prefix=maybe_prefix(prefix, "language_model"),
1102
architectures=["MiniMaxM3SparseForCausalLM"],
1103
)
1104
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1105
+ self.language_model.make_empty_intermediate_tensors
1106
+ )
1107
1108
def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None:
1109
pixel_values = kwargs.pop("pixel_values", None)
@@ -1179,10 +1213,13 @@ class MiniMaxM3SparseForConditionalGeneration(
1213
self,
1214
input_ids: torch.Tensor | None,
1215
positions: torch.Tensor,
1216
+ intermediate_tensors: IntermediateTensors | None = None,
1217
inputs_embeds: torch.Tensor | None = None,
1218
**kwargs,
1219
) -> torch.Tensor:
1185
- return self.language_model(input_ids, positions, inputs_embeds)
1220
+ return self.language_model(
1221
+ input_ids, positions, intermediate_tensors, inputs_embeds
1222
+ )
1223
1224
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1225
return self.language_model.compute_logits(hidden_states)
vllm/models/minimax_m3/nvidia/model.py
+61
-24
@@ -21,7 +21,7 @@ from transformers import PretrainedConfig
21
from vllm import _custom_ops as ops
22
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
23
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
24
-from vllm.distributed import get_tensor_model_parallel_world_size
24
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
25
from vllm.forward_context import get_forward_context
26
from vllm.model_executor.layers.activation import SiluAndMulWithClamp
27
from vllm.model_executor.layers.attention import Attention
@@ -56,12 +56,15 @@ from vllm.model_executor.models.interfaces import (
56
MultiModalEmbeddings,
57
SupportsEagle3,
58
SupportsMultiModal,
59
+ SupportsPP,
60
)
61
from vllm.model_executor.models.utils import (
62
AutoWeightsLoader,
63
+ PPMissingLayer,
64
WeightsMapper,
65
init_vllm_registered_model,
66
is_pp_missing_parameter,
67
+ make_empty_intermediate_tensors_factory,
68
make_layers,
69
maybe_prefix,
70
)
@@ -79,6 +82,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
82
)
83
from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
84
from vllm.multimodal import MULTIMODAL_REGISTRY
85
+from vllm.sequence import IntermediateTensors
86
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
87
from vllm.v1.kv_cache_interface import (
88
FullAttentionSpec,
@@ -754,12 +758,15 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
758
759
self.vocab_size = config.vocab_size
760
757
- self.embed_tokens = VocabParallelEmbedding(
758
- config.vocab_size,
759
- config.hidden_size,
760
- quant_config=quant_config,
761
- prefix=f"{prefix}.embed_tokens",
762
- )
761
+ if get_pp_group().is_first_rank:
762
+ self.embed_tokens = VocabParallelEmbedding(
763
+ config.vocab_size,
764
+ config.hidden_size,
765
+ quant_config=quant_config,
766
+ prefix=f"{prefix}.embed_tokens",
767
+ )
768
+ else:
769
+ self.embed_tokens = PPMissingLayer()
770
771
# Reserved top-k indices buffer shared by all sparse-attention indexer
772
# layers (mirrors DeepseekV4); kept at a stable address so the indexer's
@@ -792,7 +799,13 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
799
prefix=f"{prefix}.layers",
800
)
801
795
- self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
802
+ if get_pp_group().is_last_rank:
803
+ self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
804
+ else:
805
+ self.norm = PPMissingLayer()
806
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
807
+ ["hidden_states", "residual"], config.hidden_size
808
+ )
809
810
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
811
return self.embed_tokens(input_ids)
@@ -801,13 +814,19 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
814
self,
815
input_ids: torch.Tensor | None,
816
positions: torch.Tensor,
817
+ intermediate_tensors: IntermediateTensors | None,
818
inputs_embeds: torch.Tensor | None = None,
805
- ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
806
- if inputs_embeds is not None:
807
- hidden_states = inputs_embeds
819
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
820
+ if get_pp_group().is_first_rank:
821
+ if inputs_embeds is not None:
822
+ hidden_states = inputs_embeds
823
+ else:
824
+ hidden_states = self.embed_input_ids(input_ids)
825
+ residual = None
826
else:
809
- hidden_states = self.embed_input_ids(input_ids)
810
- residual = None
827
+ assert intermediate_tensors is not None
828
+ hidden_states = intermediate_tensors["hidden_states"]
829
+ residual = intermediate_tensors["residual"]
830
831
# EAGLE3 is not yet compatible with pipeline parallel
832
aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
@@ -817,6 +836,11 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
836
aux_hidden_states, idx + 1, hidden_states, residual
837
)
838
839
+ if not get_pp_group().is_last_rank:
840
+ return IntermediateTensors(
841
+ {"hidden_states": hidden_states, "residual": residual}
842
+ )
843
+
844
hidden_states, _ = self.norm(hidden_states, residual)
845
846
if len(aux_hidden_states) > 0:
@@ -931,7 +955,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
955
return loaded_params
956
957
934
-class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
958
+class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
959
"""MiniMax M3 (sparse/dense backbone) for causal language modeling."""
960
961
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -943,13 +967,19 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
967
self.model = MiniMaxM3Model(
968
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
969
)
946
- self.lm_head = ParallelLMHead(
947
- config.vocab_size,
948
- config.hidden_size,
949
- quant_config=quant_config,
950
- prefix=maybe_prefix(prefix, "lm_head"),
951
- )
970
+ if get_pp_group().is_last_rank:
971
+ self.lm_head = ParallelLMHead(
972
+ config.vocab_size,
973
+ config.hidden_size,
974
+ quant_config=quant_config,
975
+ prefix=maybe_prefix(prefix, "lm_head"),
976
+ )
977
+ else:
978
+ self.lm_head = PPMissingLayer()
979
self.logits_processor = LogitsProcessor(config.vocab_size)
980
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
981
+ self.model.make_empty_intermediate_tensors
982
+ )
983
984
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
985
return self.model.embed_input_ids(input_ids)
@@ -958,10 +988,11 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
988
self,
989
input_ids: torch.Tensor | None,
990
positions: torch.Tensor,
991
+ intermediate_tensors: IntermediateTensors | None = None,
992
inputs_embeds: torch.Tensor | None = None,
993
**kwargs,
963
- ) -> torch.Tensor:
964
- return self.model(input_ids, positions, inputs_embeds)
994
+ ) -> torch.Tensor | IntermediateTensors:
995
+ return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
996
997
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
998
return self.logits_processor(self.lm_head, hidden_states)
@@ -980,7 +1011,7 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
1011
dummy_inputs=MiniMaxM3VLDummyInputsBuilder,
1012
)
1013
class MiniMaxM3SparseForConditionalGeneration(
983
- nn.Module, SupportsMultiModal, SupportsEagle3
1014
+ nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3
1015
):
1016
"""Top-level (VL) entry point for MiniMax M3.
1017
@@ -1042,6 +1073,9 @@ class MiniMaxM3SparseForConditionalGeneration(
1073
prefix=maybe_prefix(prefix, "language_model"),
1074
architectures=["MiniMaxM3SparseForCausalLM"],
1075
)
1076
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1077
+ self.language_model.make_empty_intermediate_tensors
1078
+ )
1079
1080
# Expose language model / lm_head for EAGLE3 spec decode.
1081
@property
@@ -1160,10 +1194,13 @@ class MiniMaxM3SparseForConditionalGeneration(
1194
self,
1195
input_ids: torch.Tensor | None,
1196
positions: torch.Tensor,
1197
+ intermediate_tensors: IntermediateTensors | None = None,
1198
inputs_embeds: torch.Tensor | None = None,
1199
**kwargs,
1200
) -> torch.Tensor:
1166
- return self.language_model(input_ids, positions, inputs_embeds)
1201
+ return self.language_model(
1202
+ input_ids, positions, intermediate_tensors, inputs_embeds
1203
+ )
1204
1205
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1206
return self.language_model.compute_logits(hidden_states)