[VLM] Add TP support for Phi-4-MM (#14453)

Signed-off-by: Isotr0py <2037008807@qq.com>

vllm/model_executor/models/phi4mm_utils.py

@@ -5,14 +5,11 @@
 # but implemented by the Phi-Speech team
 #!/usr/bin/env python3
 import math
-from functools import partial
-from typing import Callable, Dict, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
 from torch import Tensor, nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
-    CheckpointImpl, checkpoint_wrapper, offload_wrapper)
 
 
 class Block(nn.Module):
@@ -873,10 +870,8 @@ class MeanVarianceNormLayer(nn.Module):
     def __init__(self, input_size):
         super().__init__()
         self.input_size = input_size
-        self.register_buffer("global_mean", torch.zeros(input_size))
-        self.register_buffer("global_invstd", torch.ones(input_size))
-        self.global_mean: Optional[Tensor]
-        self.global_invstd: Optional[Tensor]
+        self.global_mean = nn.Parameter(torch.zeros(input_size))
+        self.global_invstd = nn.Parameter(torch.ones(input_size))
 
     def forward(self, input_: Tensor) -> Tensor:
         """MeanVarianceNormLayer Forward
@@ -1023,21 +1018,10 @@ class CausalConv2D(nn.Conv2d):
         self,
         x,
     ):
-        if self.training:
-            x = F.pad(
-                x,
-                pad=(
-                    self._left_padding,
-                    self._right_padding,
-                    self._left_padding,
-                    self._right_padding,
-                ),
-            )
-        else:
-            x = F.pad(
-                x,
-                pad=(self._left_padding, self._right_padding, 0, 0),
-            )
+        x = F.pad(
+            x,
+            pad=(self._left_padding, self._right_padding, 0, 0),
+        )
         x = super().forward(x)
         return x
 
@@ -1840,68 +1824,6 @@ class MultiHeadedAttention(nn.Module):
         return self.linear_out(x)  # (batch, time1, d_model)
 
 
-def validate_checkpointing_config(activation_checkpointing):
-    """validate activation checkpointing configuration"""
-    if isinstance(activation_checkpointing, str):
-        assert activation_checkpointing in (
-            "",
-            "checkpoint",
-            "offload",
-        ), "activation_checkpointing has to be a dict or a str in "\
-            "('', 'checkpoint', 'offload')."
-    elif isinstance(activation_checkpointing, dict):
-        assert activation_checkpointing.get("module", "transformer") in (
-            "transformer",
-            "attention",
-        ), "module in activation_checkpointing has to be in "\
-            "('transformer', 'attention')."
-    else:
-        raise ValueError("activation_checkpointing has to be a str"\
-                         " or dict.")
-
-
-def embedding_checkpoint_wrapper(
-    activation_checkpointing: Union[str, Dict], ) -> Callable:
-    """return encoder embedding activation checkpoint wrapper"""
-    validate_checkpointing_config(activation_checkpointing)
-
-    if isinstance(activation_checkpointing, str):
-        if activation_checkpointing:
-            if activation_checkpointing == "offload":
-                return offload_wrapper
-            return partial(checkpoint_wrapper)
-        return lambda x: x
-
-    if isinstance(activation_checkpointing, dict):
-        enabled = activation_checkpointing.get("embed", False)
-        if enabled:
-            offloading = activation_checkpointing.get("offload", False)
-            if offloading:
-                return offload_wrapper
-            impl = (CheckpointImpl.REENTRANT if activation_checkpointing.get(
-                "reentrant", False) else CheckpointImpl.NO_REENTRANT)
-            return partial(checkpoint_wrapper, checkpoint_impl=impl)
-        return lambda x: x
-    raise ValueError("Invalid activation_checkpointing config")
-
-
-def attn_checkpointing(activation_checkpointing: Union[str, Dict],
-                       i) -> Union[str, Dict]:
-    """return activation checkpointing config for attention layer"""
-    if isinstance(activation_checkpointing, str):
-        return ""
-
-    if isinstance(activation_checkpointing, dict):
-        target_layer_cls = activation_checkpointing.get(
-            "module", "transformer")
-        checkpointing_interval = activation_checkpointing.get("interval", 1)
-        if target_layer_cls == "attention" and i % checkpointing_interval == 0:
-            return activation_checkpointing
-        return ""
-
-    raise ValueError("Invalid activation_checkpointing config")
-
-
 class MultiSequential(torch.nn.Sequential):
     """Multi-input multi-output torch.nn.Sequential"""
 
@@ -1913,17 +1835,6 @@ class MultiSequential(torch.nn.Sequential):
         return args
 
 
-def repeat(repeat_num, module_gen_fn):
-    """repeat module N times
-
-    :param int repeat_num: repeat time
-    :param function module_gen_fn: function to generate module
-    :return: repeated modules
-    :rtype: MultiSequential
-    """
-    return MultiSequential(*[module_gen_fn(i) for i in range(repeat_num)])
-
-
 def get_offset(input_layer: str, time_reduction: int):
     """Get an offset. We will use the offset for determining #frames of a 
     subsampled feature.