@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ebb1ec9318

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 07:15 UTC

[Model] enable data parallel for Llama4 vision encoder (#18368)

Signed-off-by: yzhen <yzhen@devgpu093.cco2.facebook.com> Co-authored-by: yZhen <yZhen@fb.com> Co-authored-by: yzhen <yzhen@devgpu093.cco2.facebook.com>

jennyyyyzhen committed Jun 2, 2025 at 04:22 UTC ebb1ec931871ee1baca23673049865856c44ce4e

4 files changed +214 -68

vllm/config.py

@@ -1790,6 +1790,10 @@ class ParallelConfig:

1790 rank: int = 0

1791 """Global rank in distributed setup."""

1792

1793 + enable_multimodal_encoder_data_parallel: bool = False

1794 + """ Use data parallelism instead of tensor parallelism for vision encoder.

1795 + Only support LLama4 for now"""

1796 +

1797 @property

1798 def world_size_across_dp(self) -> int:

1799 """world_size_across_dp is TPxPPxDP, it is the size of the world

vllm/engine/arg_utils.py

@@ -423,6 +423,9 @@ class EngineArgs:

423 use_tqdm_on_load: bool = LoadConfig.use_tqdm_on_load

424 pt_load_map_location: str = LoadConfig.pt_load_map_location

425

426 + enable_multimodal_encoder_data_parallel: bool = \

427 + ParallelConfig.enable_multimodal_encoder_data_parallel

428 +

429 def __post_init__(self):

430 # support `EngineArgs(compilation_config={...})`

431 # without having to manually construct a

@@ -637,6 +640,9 @@ class EngineArgs:

640 **parallel_kwargs["worker_cls"])

641 parallel_group.add_argument("--worker-extension-cls",

642 **parallel_kwargs["worker_extension_cls"])

643 + parallel_group.add_argument(

644 + "--enable-multimodal-encoder-data-parallel",

645 + **parallel_kwargs["enable_multimodal_encoder_data_parallel"])

646

647 # KV cache arguments

648 cache_kwargs = get_kwargs(CacheConfig)

@@ -1078,6 +1084,8 @@ class EngineArgs:

1084 distributed_executor_backend=self.distributed_executor_backend,

1085 worker_cls=self.worker_cls,

1086 worker_extension_cls=self.worker_extension_cls,

1087 + enable_multimodal_encoder_data_parallel=self.

1088 + enable_multimodal_encoder_data_parallel,

1089 )

1090

1091 speculative_config = self.create_speculative_config(

vllm/model_executor/models/mllama4.py

+167 -68

@@ -34,6 +34,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size

34 from vllm.inputs import InputProcessingContext

35 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

36 QKVParallelLinear,

37 + ReplicatedLinear,

38 RowParallelLinear)

39 from vllm.model_executor.layers.quantization import QuantizationConfig

40 from vllm.model_executor.layers.rotary_embedding import get_rope

@@ -49,6 +50,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,

50 BaseProcessingInfo, PromptReplacement,

51 PromptUpdate, PromptUpdateDetails)

52 from vllm.multimodal.profiling import BaseDummyInputsBuilder

53 +from vllm.multimodal.utils import run_dp_sharded_vision_model

54 from vllm.sequence import IntermediateTensors

56 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

@@ -84,23 +86,29 @@ class Llama4ImagePatchInputs(TypedDict):

87 class Llama4VisionMLP(nn.Module):

87 - def __init__(self,

88 - input_size: int,

89 - intermediate_size: int,

90 - output_size: int,

91 - bias: bool,

92 - output_activation: bool,

93 - quant_config: Optional[QuantizationConfig] = None,

94 - prefix: str = ""):

89 + def __init__(

90 + self,

91 + input_size: int,

92 + intermediate_size: int,

93 + output_size: int,

94 + bias: bool,

95 + output_activation: bool,

96 + quant_config: Optional[QuantizationConfig] = None,

97 + prefix: str = "",

98 + use_data_parallel: bool = False,

99 + ):

100 super().__init__()

96 - self.fc1 = ColumnParallelLinear(

101 + cls_fc1 = (ReplicatedLinear

102 + if use_data_parallel else ColumnParallelLinear)

103 + self.fc1 = cls_fc1(

104 input_size=input_size,

105 output_size=intermediate_size,

106 bias=bias,

107 quant_config=quant_config,

108 prefix=f"{prefix}.fc1",

109 )

103 - self.fc2 = RowParallelLinear(

110 + cls_fc2 = ReplicatedLinear if use_data_parallel else RowParallelLinear

111 + self.fc2 = cls_fc2(

112 input_size=intermediate_size,

113 output_size=output_size,

114 bias=bias,

@@ -155,10 +163,12 @@ def pixel_shuffle(input_tensor, shuffle_ratio):

163 int(channels / shuffle_ratio))

164 reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()

165

158 - reshaped_tensor = reshaped_tensor.view(batch_size,

159 - int(height * shuffle_ratio),

160 - int(width * shuffle_ratio),

161 - int(channels / (shuffle_ratio**2)))

166 + reshaped_tensor = reshaped_tensor.view(

167 + batch_size,

168 + int(height * shuffle_ratio),

169 + int(width * shuffle_ratio),

170 + int(channels / (shuffle_ratio**2)),

171 + )

172 reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()

173

174 output_tensor = reshaped_tensor.view(batch_size, -1,

@@ -173,6 +183,7 @@ class Llama4VisionPixelShuffleMLP(nn.Module):

183 config,

184 quant_config: Optional[QuantizationConfig] = None,

185 prefix: str = "",

186 + use_data_parallel: bool = False,

187 ):

188 super().__init__()

189 self.pixel_shuffle_ratio = config.pixel_shuffle_ratio

@@ -186,7 +197,9 @@ class Llama4VisionPixelShuffleMLP(nn.Module):

197 bias=config.multi_modal_projector_bias,

198 output_activation=True,

199 quant_config=quant_config,

189 - prefix=f"{prefix}.mlp")

200 + prefix=f"{prefix}.mlp",

201 + use_data_parallel=use_data_parallel,

202 + )

203

204 def forward(self, encoded_patches: torch.Tensor) -> torch.Tensor:

205 encoded_patches = pixel_shuffle(encoded_patches,

@@ -201,10 +214,12 @@ class Llama4VisionAttention(nn.Module):

214 config: Llama4VisionConfig,

215 quant_config: Optional[QuantizationConfig],

216 prefix: str = "",

217 + use_data_parallel: bool = False,

218 ):

219 super().__init__()

220 self.config = config

207 - self.tp_size = get_tensor_model_parallel_world_size()

221 + self.tp_size = (1 if use_data_parallel else

222 + get_tensor_model_parallel_world_size())

223 self.embed_dim = config.hidden_size

224 self.num_heads = config.num_attention_heads

225 self.head_dim = config.hidden_size // self.num_heads

@@ -217,22 +232,39 @@ class Llama4VisionAttention(nn.Module):

232

233 self.attn = MultiHeadAttention(self.num_local_heads, self.head_dim,

234 self.scaling)

220 - self.qkv_proj = QKVParallelLinear(

221 - self.embed_dim,

222 - self.head_dim,

223 - self.num_heads,

224 - bias=True,

225 - quant_config=quant_config,

226 - prefix=f"{prefix}.qkv_proj",

227 - )

228 - self.o_proj = RowParallelLinear(

229 - self.num_heads * self.head_dim,

230 - self.embed_dim,

231 - bias=True,

232 - input_is_parallel=True,

233 - quant_config=quant_config,

234 - prefix=f"{prefix}.o_proj",

235 - )

235 +

236 + if use_data_parallel:

237 + self.qkv_proj = ReplicatedLinear(

238 + self.embed_dim,

239 + self.q_size + 2 * self.kv_size,

240 + bias=True,

241 + quant_config=quant_config,

242 + prefix=f"{prefix}.qkv_proj",

243 + )

244 + self.o_proj = ReplicatedLinear(

245 + self.num_heads * self.head_dim,

246 + self.embed_dim,

247 + bias=True,

248 + quant_config=quant_config,

249 + prefix=f"{prefix}.o_proj",

250 + )

251 + else:

252 + self.qkv_proj = QKVParallelLinear(

253 + self.embed_dim,

254 + self.head_dim,

255 + self.num_heads,

256 + bias=True,

257 + quant_config=quant_config,

258 + prefix=f"{prefix}.qkv_proj",

259 + )

260 + self.o_proj = RowParallelLinear(

261 + self.num_heads * self.head_dim,

262 + self.embed_dim,

263 + bias=True,

264 + input_is_parallel=True,

265 + quant_config=quant_config,

266 + prefix=f"{prefix}.o_proj",

267 + )

268

269 self.rotary_emb = get_rope(

270 head_size=self.head_dim,

@@ -275,22 +307,29 @@ class Llama4VisionEncoderLayer(nn.Module):

307 config: Llama4VisionConfig,

308 quant_config: Optional[QuantizationConfig],

309 prefix: str = "",

310 + use_data_parallel: bool = False,

311 ):

312 super().__init__()

313 self.hidden_size = config.hidden_size

314 self.num_attention_heads = config.num_attention_heads

315 self.intermediate_size = config.intermediate_size

316

284 - self.self_attn = Llama4VisionAttention(config,

285 - quant_config=quant_config,

286 - prefix=f"{prefix}.self_attn")

287 - self.mlp = Llama4VisionMLP(input_size=config.hidden_size,

288 - intermediate_size=config.intermediate_size,

289 - output_size=config.hidden_size,

290 - bias=True,

291 - output_activation=False,

292 - quant_config=quant_config,

293 - prefix=f"{prefix}.mlp")

317 + self.self_attn = Llama4VisionAttention(

318 + config,

319 + quant_config=quant_config,

320 + prefix=f"{prefix}.self_attn",

321 + use_data_parallel=use_data_parallel,

322 + )

323 + self.mlp = Llama4VisionMLP(

324 + input_size=config.hidden_size,

325 + intermediate_size=config.intermediate_size,

326 + output_size=config.hidden_size,

327 + bias=True,

328 + output_activation=False,

329 + quant_config=quant_config,

330 + prefix=f"{prefix}.mlp",

331 + use_data_parallel=use_data_parallel,

332 + )

333

334 self.input_layernorm = nn.LayerNorm(config.hidden_size)

335 self.post_attention_layernorm = nn.LayerNorm(config.hidden_size)

@@ -322,6 +361,7 @@ class Llama4VisionEncoder(nn.Module):

361 config: Llama4VisionConfig,

362 quant_config: Optional[QuantizationConfig],

363 prefix: str = "",

364 + use_data_parallel: bool = False,

365 ):

366 super().__init__()

367 self.config = config

@@ -330,6 +370,7 @@ class Llama4VisionEncoder(nn.Module):

370 config,

371 quant_config=quant_config,

372 prefix=f"{prefix}.layers.{layer_idx}",

373 + use_data_parallel=use_data_parallel,

374 ) for layer_idx in range(config.num_hidden_layers)

375 ])

376

@@ -357,23 +398,33 @@ class Llama4VisionEncoder(nn.Module):

398

399 class Llama4UnfoldConvolution(nn.Module):

400

360 - def __init__(self,

361 - config: Llama4VisionConfig,

362 - quant_config: Optional[QuantizationConfig] = None,

363 - prefix: str = ""):

401 + def __init__(

402 + self,

403 + config: Llama4VisionConfig,

404 + quant_config: Optional[QuantizationConfig] = None,

405 + prefix: str = "",

406 + use_data_parallel: bool = False,

407 + ):

408 super().__init__()

409 kernel_size = config.patch_size

410 if isinstance(kernel_size, int):

411 kernel_size = (kernel_size, kernel_size)

412 self.unfold = torch.nn.Unfold(kernel_size=kernel_size,

413 stride=config.patch_size)

370 - self.linear = ColumnParallelLinear(config.num_channels *

371 - kernel_size[0] * kernel_size[1],

372 - config.hidden_size,

373 - bias=False,

374 - quant_config=quant_config,

375 - gather_output=True,

376 - prefix=f"{prefix}.linear")

414 + params = {

415 + "input_size":

416 + config.num_channels * kernel_size[0] * kernel_size[1],

417 + "output_size": config.hidden_size,

418 + "bias": False,

419 + "quant_config": quant_config,

420 + "prefix": f"{prefix}.linear",

421 + }

422 + if use_data_parallel:

423 + cls = ReplicatedLinear

424 + else:

425 + cls = ColumnParallelLinear

426 + params["gather_output"] = True

427 + self.linear = cls(**params)

428

429 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

430 hidden_states = self.unfold(hidden_states)

@@ -389,6 +440,7 @@ class Llama4VisionModel(nn.Module):

440 config: Llama4VisionConfig,

441 quant_config: Optional[QuantizationConfig] = None,

442 prefix: str = "",

443 + use_data_parallel: bool = False,

444 ):

445 super().__init__()

446 self.config = config

@@ -403,7 +455,9 @@ class Llama4VisionModel(nn.Module):

455 self.patch_embedding = Llama4UnfoldConvolution(

456 config,

457 quant_config=quant_config,

406 - prefix=f"{prefix}.patch_embedding")

458 + prefix=f"{prefix}.patch_embedding",

459 + use_data_parallel=use_data_parallel,

460 + )

461

462 self.class_embedding = nn.Parameter(self.scale *

463 torch.randn(self.hidden_size))

@@ -415,11 +469,18 @@ class Llama4VisionModel(nn.Module):

469 self.layernorm_post = nn.LayerNorm(self.hidden_size, eps=1e-5)

470

471 # encoders

418 - self.model = Llama4VisionEncoder(config,

419 - quant_config=quant_config,

420 - prefix=f"{prefix}.model")

472 + self.model = Llama4VisionEncoder(

473 + config,

474 + quant_config=quant_config,

475 + prefix=f"{prefix}.model",

476 + use_data_parallel=use_data_parallel,

477 + )

478 self.vision_adapter = Llama4VisionPixelShuffleMLP(

422 - config, quant_config, prefix=f"{prefix}.vision_adapter")

479 + config,

480 + quant_config,

481 + prefix=f"{prefix}.vision_adapter",

482 + use_data_parallel=use_data_parallel,

483 + )

484

485 def forward(

486 self,

@@ -528,8 +589,9 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]

589 vision_config = self.info.get_hf_config().vision_config

590

591 if processed_outputs.get("pixel_values") is not None:

531 - assert "images" in mm_data, \

532 - "images expected to be in mm_data when pixel_values is present"

592 + assert (

593 + "images" in mm_data

594 + ), "images expected to be in mm_data when pixel_values is present"

595

596 images = mm_data["images"]

597 parsed_images = (self._get_data_parser().parse_mm_data({

@@ -546,8 +608,8 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]

608 get_best_fit(

609 (image.size[1], image.size[0]),

610 torch.tensor(possible_resolutions),

549 - resize_to_max_canvas=image_processor.resize_to_max_canvas)

550 - for image in parsed_images

611 + resize_to_max_canvas=image_processor.resize_to_max_canvas,

612 + ) for image in parsed_images

613 ]

614 # TODO tile height/width do not necessarily need to match

615 aspect_ratios = [(image_size[0] // tile_size,

@@ -659,13 +721,17 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

721 config = vllm_config.model_config.hf_config

722 quant_config = vllm_config.quant_config

723 multimodal_config = vllm_config.model_config.multimodal_config

724 + self.use_data_parallel = (vllm_config.parallel_config.

725 + enable_multimodal_encoder_data_parallel)

726 self.config = config

727 self.quant_config = quant_config

728 self.multimodal_config = multimodal_config

665 - self.vision_model = Llama4VisionModel(config.vision_config,

666 - None,

667 - prefix=maybe_prefix(

668 - prefix, "vision_model"))

729 + self.vision_model = Llama4VisionModel(

730 + config.vision_config,

731 + None,

732 + prefix=maybe_prefix(prefix, "vision_model"),

733 + use_data_parallel=self.use_data_parallel,

734 + )

735 self.multi_modal_projector = Llama4MultiModalProjector(

736 self.config,

737 None,

@@ -709,7 +775,13 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

775 flat_data = image_input["flat_data"]

776 patches_per_image = image_input["patches_per_image"].tolist()

777

712 - vision_embeddings_flat = self.vision_model(flat_data)

778 + # shard image input

779 + if self.use_data_parallel:

780 + vision_embeddings_flat = run_dp_sharded_vision_model(

781 + flat_data, self.vision_model)

782 + else:

783 + vision_embeddings_flat = self.vision_model(flat_data)

784 +

785 vision_embeddings_flat = self.multi_modal_projector(

786 vision_embeddings_flat)

787

@@ -796,6 +868,30 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

868

869 return get_prefix_weights(), get_other_weights()

870

871 + def _consolidate_qkv_weights(

872 + self, weights: Iterable[tuple[str, torch.Tensor]]

873 + ) -> Iterable[tuple[str, torch.Tensor]]:

874 + qkv_idx_mappings = {

875 + ".self_attn.q_proj": 0,

876 + ".self_attn.k_proj": 1,

877 + ".self_attn.v_proj": 2,

878 + }

879 + qkv_weights = {}

880 + for name, loaded_weight in weights:

881 + for weight_name, idx in qkv_idx_mappings.items():

882 + if weight_name not in name:

883 + continue

884 + new_name = name.replace(weight_name, ".self_attn.qkv_proj")

885 + if new_name not in qkv_weights:

886 + qkv_weights[new_name] = [None] * 3

887 + qkv_weights[new_name][idx] = loaded_weight

888 + break

889 + else:

890 + yield name, loaded_weight

891 + for key, weight in qkv_weights.items():

892 + qkv_weight = torch.cat(weight, dim=0)

893 + yield key, qkv_weight

894 +

895 def load_weights(self, weights: Iterable[tuple[str,

896 torch.Tensor]]) -> set[str]:

897

@@ -818,9 +914,12 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,

914 assert loaded_language_model_params is not None

915 updated_params.update(loaded_language_model_params)

916

917 + if self.use_data_parallel:

918 + other_weights = self._consolidate_qkv_weights(other_weights)

919 +

920 for name, loaded_weight in other_weights:

921 for param_name, weight_name, shard_id in stacked_params_mapping:

823 - if weight_name not in name:

922 + if weight_name not in name or self.use_data_parallel:

923 continue

924 name = name.replace(weight_name, param_name)

925 param = params_dict[name]

vllm/multimodal/utils.py

+35

@@ -12,6 +12,9 @@ from PIL import Image

13 import vllm.envs as envs

14 from vllm.connections import HTTPConnection, global_http_connection

15 +from vllm.distributed import (get_tensor_model_parallel_rank,

16 + get_tensor_model_parallel_world_size,

17 + tensor_model_parallel_all_gather)

19 from .audio import AudioMediaIO

20 from .base import MediaIO

@@ -390,3 +393,35 @@ def group_mm_inputs_by_modality(

393 return [

394 list(group) for _, group in groupby(mm_inputs, key=modality_group_func)

395 ]

396 +

397 +

398 +def run_dp_sharded_vision_model(image_input: torch.Tensor,

399 + vision_model: torch.nn.Module) -> torch.Tensor:

400 + """Run a vision model with data parallelism (DP) sharding. The function

401 + will shard the input image tensor on the first dimension and run the vision

402 + model

403 +

404 + Args:

405 + image_input (torch.Tensor): Image input tensor.

406 + vision_model (torch.nn.Module): Vision model.

407 +

408 + Returns:

409 + torch.Tensor: Output image embeddings

410 + """

411 +

412 + num_chunks = image_input.shape[0]

413 + mp_world_size = get_tensor_model_parallel_world_size()

414 + num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size

415 + num_padded_chunks = num_chunks_per_rank * mp_world_size - num_chunks

416 + pad = (0, ) * (2 * (image_input.dim() - 1)) + (0, num_padded_chunks)

417 + image_input_padded = torch.nn.functional.pad(image_input, pad)

418 + rank = get_tensor_model_parallel_rank()

419 + image_input_per_rank = image_input_padded[rank *

420 + num_chunks_per_rank:(rank + 1) *

421 + num_chunks_per_rank, ...]

422 +

423 + vision_embeddings = vision_model(image_input_per_rank)

424 + vision_embeddings = tensor_model_parallel_all_gather(vision_embeddings,

425 + dim=0)

426 + vision_embeddings = vision_embeddings[:num_chunks, ...]

427 + return vision_embeddings