diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 74e9e7739f64..c900c291245a 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -443,6 +443,7 @@ th { | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ | | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ | | `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ | +| `MiniMaxM3SparseForCausalLM` | MiniMax-M3 | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ | | `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ | | `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. | ✅︎ | ✅︎ | | `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ | @@ -593,7 +594,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + IE+ + VE+ + A+ | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ | | `MiniCPMO` | MiniCPM-O | T + IE+ + VE+ + AE+ | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ | | `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `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. | ✅︎ | | -| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I+ + V+ | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | | +| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I+ + V+ | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ | | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ | | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ | | `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ | diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index c171b4bfe8c5..4f3528806c9d 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -30,7 +30,7 @@ VllmConfig, get_current_vllm_config, ) -from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -64,12 +64,15 @@ MultiModalEmbeddings, SupportsEagle3, SupportsMultiModal, + SupportsPP, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, + PPMissingLayer, WeightsMapper, init_vllm_registered_model, is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, ) @@ -92,6 +95,7 @@ ) from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -772,12 +776,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.vocab_size = config.vocab_size - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() # Reserved top-k indices buffer shared by all sparse-attention indexer # layers (mirrors DeepseekV4); the indexer writes its per-head decode/ @@ -807,7 +814,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if get_pp_group().is_last_rank: + self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -816,13 +829,19 @@ def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: - if inputs_embeds is not None: - hidden_states = inputs_embeds + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None else: - hidden_states = self.embed_input_ids(input_ids) - residual = None + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] # EAGLE3 is not yet compatible with pipeline parallel aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) @@ -832,6 +851,11 @@ def forward( aux_hidden_states, idx + 1, hidden_states, residual ) + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) > 0: @@ -946,7 +970,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params -class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3): +class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3): """MiniMax M3 (sparse/dense backbone) for causal language modeling.""" packed_modules_mapping = { @@ -963,13 +987,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = MiniMaxM3Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -978,10 +1008,11 @@ def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, - ) -> torch.Tensor: - return self.model(input_ids, positions, inputs_embeds) + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) @@ -1004,7 +1035,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: dummy_inputs=MiniMaxM3VLDummyInputsBuilder, ) class MiniMaxM3SparseForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsEagle3 + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 ): """Top-level (VL) entry point for MiniMax M3. @@ -1070,6 +1101,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "language_model"), architectures=["MiniMaxM3SparseForCausalLM"], ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: pixel_values = kwargs.pop("pixel_values", None) @@ -1179,10 +1213,13 @@ def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: - return self.language_model(input_ids, positions, inputs_embeds) + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states) diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index c27ded3b83ef..a30bc335fcfb 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -21,7 +21,7 @@ from vllm import _custom_ops as ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config -from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.model_executor.layers.activation import SiluAndMulWithClamp from vllm.model_executor.layers.attention import Attention @@ -56,12 +56,15 @@ MultiModalEmbeddings, SupportsEagle3, SupportsMultiModal, + SupportsPP, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, + PPMissingLayer, WeightsMapper, init_vllm_registered_model, is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, ) @@ -79,6 +82,7 @@ ) from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -754,12 +758,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.vocab_size = config.vocab_size - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() # Reserved top-k indices buffer shared by all sparse-attention indexer # layers (mirrors DeepseekV4); kept at a stable address so the indexer's @@ -792,7 +799,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if get_pp_group().is_last_rank: + self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -801,13 +814,19 @@ def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: - if inputs_embeds is not None: - hidden_states = inputs_embeds + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None else: - hidden_states = self.embed_input_ids(input_ids) - residual = None + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] # EAGLE3 is not yet compatible with pipeline parallel aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) @@ -817,6 +836,11 @@ def forward( aux_hidden_states, idx + 1, hidden_states, residual ) + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) > 0: @@ -931,7 +955,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params -class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3): +class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3): """MiniMax M3 (sparse/dense backbone) for causal language modeling.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -943,13 +967,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = MiniMaxM3Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -958,10 +988,11 @@ def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, - ) -> torch.Tensor: - return self.model(input_ids, positions, inputs_embeds) + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) @@ -980,7 +1011,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: dummy_inputs=MiniMaxM3VLDummyInputsBuilder, ) class MiniMaxM3SparseForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsEagle3 + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 ): """Top-level (VL) entry point for MiniMax M3. @@ -1042,6 +1073,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "language_model"), architectures=["MiniMaxM3SparseForCausalLM"], ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) # Expose language model / lm_head for EAGLE3 spec decode. @property @@ -1160,10 +1194,13 @@ def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: - return self.language_model(input_ids, positions, inputs_embeds) + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states)