utils.py 26.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
import asyncio
import atexit
6
7
import itertools
import math
8
from collections.abc import Iterable
9
from concurrent.futures import ThreadPoolExecutor
10
from itertools import groupby
11
from pathlib import Path
12
from typing import TYPE_CHECKING, Any, Literal, Optional, TypeVar, Union
13
from urllib.parse import ParseResult, urlparse
14
from urllib.request import url2pathname
15

16
import numpy as np
17
import numpy.typing as npt
18
import torch
19
from PIL import Image, UnidentifiedImageError
20
from typing_extensions import deprecated
21

22
import vllm.envs as envs
23
from vllm.connections import HTTPConnection, global_http_connection
24
25
26
from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size,
                              tensor_model_parallel_all_gather)
27

28
29
from .audio import AudioMediaIO
from .base import MediaIO
30
from .image import ImageEmbeddingMediaIO, ImageMediaIO
31
from .video import VideoMediaIO
32

33
_M = TypeVar("_M")
34

35
if TYPE_CHECKING:
36
    from .inputs import (BatchedTensorInputs, MultiModalKwargs,
37
38
                         MultiModalKwargsItem, MultiModalKwargsItems,
                         MultiModalPlaceholderDict)
39
else:
40
    BatchedTensorInputs = Any
41
    MultiModalKwargs = Any
42
    MultiModalKwargsItem = Any
43
    MultiModalKwargsItems = Any
44
    MultiModalPlaceholderDict = Any
45

46
47
48
49
global_thread_pool = ThreadPoolExecutor(
    max_workers=envs.VLLM_MEDIA_LOADING_THREAD_COUNT)
atexit.register(global_thread_pool.shutdown)

50

51
class MediaConnector:
52

53
54
    def __init__(
        self,
55
        media_io_kwargs: Optional[dict[str, dict[str, Any]]] = None,
56
57
58
59
        connection: HTTPConnection = global_http_connection,
        *,
        allowed_local_media_path: str = "",
    ) -> None:
60
61
62
63
64
        """
        Args:
            media_io_kwargs: Additional args passed to process media 
                             inputs, keyed by modalities. For example, 
                             to set num_frames for video, set 
65
                             `--media-io-kwargs '{"video":{"num_frames":40}}'`
66
            connection: HTTP connection client to download media contents.
67
68
            allowed_local_media_path: A local directory to load media files
                                      from.
69
        """
70
71
        super().__init__()

72
73
        self.media_io_kwargs: dict[str, dict[
            str, Any]] = media_io_kwargs if media_io_kwargs else {}
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
        self.connection = connection

        if allowed_local_media_path:
            allowed_local_media_path_ = Path(allowed_local_media_path)

            if not allowed_local_media_path_.exists():
                raise ValueError(
                    "Invalid `--allowed-local-media-path`: The path "
                    f"{allowed_local_media_path_} does not exist.")
            if not allowed_local_media_path_.is_dir():
                raise ValueError(
                    "Invalid `--allowed-local-media-path`: The path "
                    f"{allowed_local_media_path_} must be a directory.")
        else:
            allowed_local_media_path_ = None

        self.allowed_local_media_path = allowed_local_media_path_

    def _load_data_url(
        self,
        url_spec: ParseResult,
        media_io: MediaIO[_M],
    ) -> _M:
        data_spec, data = url_spec.path.split(",", 1)
        media_type, data_type = data_spec.split(";", 1)

        if data_type != "base64":
            msg = "Only base64 data URLs are supported for now."
            raise NotImplementedError(msg)

        return media_io.load_base64(media_type, data)

    def _load_file_url(
        self,
        url_spec: ParseResult,
        media_io: MediaIO[_M],
    ) -> _M:
        allowed_local_media_path = self.allowed_local_media_path
        if allowed_local_media_path is None:
            raise RuntimeError("Cannot load local files without "
                               "`--allowed-local-media-path`.")

116
        filepath = Path(url2pathname(url_spec.path))
117
        if allowed_local_media_path not in filepath.resolve().parents:
118
            raise ValueError(
119
120
                f"The file path {filepath} must be a subpath "
                f"of `--allowed-local-media-path` {allowed_local_media_path}.")
121

122
        return media_io.load_file(filepath)
123

124
125
126
127
128
129
130
131
    def load_from_url(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
        fetch_timeout: Optional[int] = None,
    ) -> _M:
        url_spec = urlparse(url)
132

133
134
135
        if url_spec.scheme.startswith("http"):
            connection = self.connection
            data = connection.get_bytes(url, timeout=fetch_timeout)
136

137
            return media_io.load_bytes(data)
138

139
140
        if url_spec.scheme == "data":
            return self._load_data_url(url_spec, media_io)
141

142
143
        if url_spec.scheme == "file":
            return self._load_file_url(url_spec, media_io)
144

145
146
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
147

148
149
150
151
152
153
154
155
    async def load_from_url_async(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
        fetch_timeout: Optional[int] = None,
    ) -> _M:
        url_spec = urlparse(url)
156
        loop = asyncio.get_running_loop()
157

158
159
160
        if url_spec.scheme.startswith("http"):
            connection = self.connection
            data = await connection.async_get_bytes(url, timeout=fetch_timeout)
161
162
163
            future = loop.run_in_executor(global_thread_pool,
                                          media_io.load_bytes, data)
            return await future
164

165
        if url_spec.scheme == "data":
166
167
168
169
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_data_url, url_spec,
                                          media_io)
            return await future
170

171
        if url_spec.scheme == "file":
172
173
174
175
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_file_url, url_spec,
                                          media_io)
            return await future
176
177
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
178

179
180
181
182
183
184
185
    def fetch_audio(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Load audio from a URL.
        """
186
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
187

188
        return self.load_from_url(
189
            audio_url,
190
191
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
192
        )
193

194
195
196
197
198
199
200
    async def fetch_audio_async(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Asynchronously fetch audio from a URL.
        """
201
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
202

203
        return await self.load_from_url_async(
204
            audio_url,
205
206
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
207
        )
208

209
210
211
212
213
214
215
    def fetch_image(
        self,
        image_url: str,
        *,
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
216
        Load a PIL image from an HTTP or base64 data URL.
217

218
219
        By default, the image is converted into RGB format.
        """
220
221
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
222

223
224
225
226
227
228
229
230
231
        try:
            return self.load_from_url(
                image_url,
                image_io,
                fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
            )
        except UnidentifiedImageError as e:
            # convert to ValueError to be properly caught upstream
            raise ValueError(str(e)) from e
232

233
234
    async def fetch_image_async(
        self,
235
236
        image_url: str,
        *,
237
238
239
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
240
        Asynchronously load a PIL image from an HTTP or base64 data URL.
241

242
243
        By default, the image is converted into RGB format.
        """
244
245
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
246

247
248
249
250
251
252
253
254
255
        try:
            return await self.load_from_url_async(
                image_url,
                image_io,
                fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
            )
        except UnidentifiedImageError as e:
            # convert to ValueError to be properly caught upstream
            raise ValueError(str(e)) from e
256

257
258
259
260
261
    def fetch_video(
        self,
        video_url: str,
        *,
        image_mode: str = "RGB",
262
    ) -> tuple[npt.NDArray, dict[str, Any]]:
263
        """
264
        Load video from an HTTP or base64 data URL.
265
        """
266
267
268
269
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
270
271
272
273
274
275

        return self.load_from_url(
            video_url,
            video_io,
            fetch_timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
        )
276

277
278
279
    async def fetch_video_async(
        self,
        video_url: str,
280
        *,
281
        image_mode: str = "RGB",
282
    ) -> tuple[npt.NDArray, dict[str, Any]]:
283
        """
284
        Asynchronously load video from an HTTP or base64 data URL.
285
286
287

        By default, the image is converted into RGB format.
        """
288
289
290
291
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
292
293
294
295
296
297

        return await self.load_from_url_async(
            video_url,
            video_io,
            fetch_timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
        )
298

299
300
301
302
303
304
305
306
307
308
309
    def fetch_image_embedding(
        self,
        data: str,
    ) -> torch.Tensor:
        """
        Load image embedding from a URL.
        """
        image_embedding_io = ImageEmbeddingMediaIO()

        return image_embedding_io.load_base64("", data)

310

311
312
def encode_audio_base64(
    audio: np.ndarray,
313
    sampling_rate: float,
314
315
) -> str:
    """Encode audio as base64."""
316
317
    audio_io = AudioMediaIO()
    return audio_io.encode_base64((audio, sampling_rate))
318
319


320
321
322
323
324
325
326
327
def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str = "RGB",
    format: str = "JPEG",
) -> str:
    """
    Encode a pillow image to base64 format.
328

329
330
    By default, the image is converted into RGB format before being encoded.
    """
331
332
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)
333
334


335
def encode_video_base64(frames: npt.NDArray) -> str:
336
337
338
    image_io = ImageMediaIO()
    video_io = VideoMediaIO(image_io)
    return video_io.encode_base64(frames)
339
340


341
342
343
344
345
346
def argsort_mm_positions(
        mm_positions: MultiModalPlaceholderDict) -> list[tuple[str, int]]:
    """
    Given a `MultiModalPlaceholderDict`, output a sequence of keys to
    sort the dictionary by `offset` (starting index in the input sequence)
    in ascending order.
347
348

    Returns:
349
350
        A list of `(modality, idx)`, which can be used to access an item
        by `mm_positions[modality][idx]`.
351
    """
352
353
354
    flat_items = ((modality, idx, item)
                  for modality, items in mm_positions.items()
                  for idx, item in enumerate(items))
355

356
    sorted_flat_items = sorted(flat_items, key=lambda x: x[2].offset)
357

358
    return [(modality, idx) for modality, idx, _ in sorted_flat_items]
359
360


361
362
363
364
# Temporary back-compatibility for plugins that define model runner
@deprecated("`group_mm_inputs_by_modality` is superseded by "
            "`group_mm_kwargs_by_modality` and will be removed in v0.13. "
            "Please use `group_mm_kwargs_by_modality` instead.")
365
def group_mm_inputs_by_modality(
366
367
    mm_inputs: list[MultiModalKwargsItems]
) -> list[list[MultiModalKwargsItems]]:
368
369
370
    if not mm_inputs:
        return []

371
372
    def modality_group_func(
            mm_input: MultiModalKwargsItems) -> Union[str, int]:
373
        # If the input has multiple modalities, return an id as the unique key
374
        # for the mm_input input.
375
        if len(mm_input) > 1:
376
377
            return id(mm_input)

378
379
        elif len(mm_input) == 1:
            return next(iter(mm_input.keys()))
380

381
        raise AssertionError("This line should be unreachable.")
382
383
384
385

    return [
        list(group) for _, group in groupby(mm_inputs, key=modality_group_func)
    ]
386
387


388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
def group_mm_kwargs_by_modality(
    mm_kwargs: list[MultiModalKwargsItem],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
    """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same
    modality together into the same `MultiModalKwargs` instance.

    Args:
        mm_inputs: List of `MultiModalKwargsItem`.

    Yields:
        A tuple `(modality, num_items, grouped_kwargs)`.
    """
403
    from vllm.multimodal.inputs import MultiModalKwargs, MultiModalKwargsItems
404
405
406
407

    for modality, items in groupby(mm_kwargs, key=lambda item: item.modality):
        items_lst = list(items)

408
        # mm_kwargs_group = MultiModalKwargsItems.from_items(items_lst) \
409
        #    .get_data(pin_memory=pin_memory)
410
411

        # if device is not None:
412
413
414
415
        #     mm_kwargs_group = json_map_leaves(
        #         lambda x: x.to(device=device),
        #         mm_kwargs_group,
        #     )
416
417
418
419
420
421
422

        # TODO: Once V0 is removed, we can use the merging logic above
        # to avoid creating an extra batch dimension (except for fields
        # that are meant to be stacked anyway).
        # We will also need to update each model to remove `flatten_bn`.
        mm_kwargs_group = MultiModalKwargs.as_kwargs(
            MultiModalKwargs.batch(
423
424
425
426
                [
                    MultiModalKwargsItems.from_seq([item]).get_data()
                    for item in items_lst
                ],
427
428
429
430
431
432
433
434
                pin_memory=pin_memory,
            ),
            device=device,
        )

        yield modality, len(items_lst), mm_kwargs_group


435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
def run_dp_sharded_vision_model(image_input: torch.Tensor,
                                vision_model: torch.nn.Module) -> torch.Tensor:
    """Run a vision model with data parallelism (DP) sharding. The function 
    will shard the input image tensor on the first dimension and run the vision
    model

    Args:
        image_input (torch.Tensor): Image input tensor.
        vision_model (torch.nn.Module): Vision model.
    Returns:
        torch.Tensor: Output image embeddings
    """

    num_chunks = image_input.shape[0]
    mp_world_size = get_tensor_model_parallel_world_size()
    num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size
    num_padded_chunks = num_chunks_per_rank * mp_world_size - num_chunks
    pad = (0, ) * (2 * (image_input.dim() - 1)) + (0, num_padded_chunks)
    image_input_padded = torch.nn.functional.pad(image_input, pad)
    rank = get_tensor_model_parallel_rank()
    image_input_per_rank = image_input_padded[rank *
                                              num_chunks_per_rank:(rank + 1) *
                                              num_chunks_per_rank, ...]

    vision_embeddings = vision_model(image_input_per_rank)
460
461
    # Ensure tensor is contiguous before all_gather
    vision_embeddings = vision_embeddings.contiguous()
462
463
464
465
    vision_embeddings = tensor_model_parallel_all_gather(vision_embeddings,
                                                         dim=0)
    vision_embeddings = vision_embeddings[:num_chunks, ...]
    return vision_embeddings
466
467


468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
def get_load_balance_assignment(
    sizes: list[int],
    num_gpus: int = 2,
) -> tuple[list[int], list[int], list[int]]:
    """
    Generate load balancing assignment and metadata 
    for distributing data across GPUs.
    The load is determined by the total image sizes,
    not the number of images.
    
    Args:
        sizes: The size of each image
        num_gpus: Number of GPUs to balance across
    
    Returns:
        shuffle_indices: 
            Indices to reorder data for balanced loading
        gpu_sample_counts: 
            Number of samples assigned to each GPU
        grouped_sizes_per_gpu: 
            Total size assigned to each GPU
    
    Example:
        ```
        sizes = [1000, 100, 200, 50]
        num_gpus=2
        ```

    """

    n_samples = len(sizes)

    # Handle edge cases
    if n_samples == 0:
        return [], [0] * num_gpus, [0] * num_gpus

    # Use greedy algorithm - balance by total size, not sample count
    gpu_assignments = [list[int]() for _ in range(num_gpus)]
    gpu_loads = [0] * num_gpus  # This tracks total SIZE, not sample count

    # Sort indices by size (largest first for better load balancing)
    # sizes = [1000, 100, 200, 50]
    # large_to_small_indices = [0, 2, 1, 3]
    large_to_small_indices = sorted(range(n_samples),
                                    key=lambda i: sizes[i],
                                    reverse=True)

    for idx in large_to_small_indices:
        # Find GPU with minimum current load (by total size)
        min_gpu = min(range(num_gpus), key=lambda i: gpu_loads[i])
        gpu_assignments[min_gpu].append(idx)
        gpu_loads[min_gpu] += sizes[idx]

    # Create shuffle indices and counts
    shuffle_indices = list[int]()
    gpu_sample_counts = list[int]()
    for gpu_id in range(num_gpus):
        # GPU_0 = [1000] = [0]
        # GPU_1 = [200, 100, 50] = [2, 1, 3]
        # shuffle_indices = [0, 2, 1, 3]
        shuffle_indices.extend(gpu_assignments[gpu_id])
        # GPU_0 = [1]
        # GPU_1 = [3]
        # gpu_sample_counts = [1, 3]
        gpu_sample_counts.append(len(gpu_assignments[gpu_id]))

    return (shuffle_indices, gpu_sample_counts, gpu_loads)


def run_dp_sharded_mrope_vision_model(
    vision_model: torch.nn.Module,
    pixel_values: torch.Tensor,
    grid_thw_list: list[list[int]],
541
542
    *,
    rope_type: Literal["rope_3d", "rope_2d"],
543
544
545
546
547
548
549
550
551
552
) -> tuple[torch.Tensor, ...]:
    """Run a vision model with data parallelism (DP) sharding. 
    The function will shard the input image tensor on the 
    first dimension and run the vision model.
    This function is used to run the vision model with mrope.
    
    Args:
        vision_model (torch.nn.Module): Vision model.
        pixel_values (torch.Tensor): Image/Video input tensor.
        grid_thw_list: List of grid dimensions for each image
553
554
555
556
        rope_type: Type of rope used in the vision model.
                   Different rope types have different dimension to do ViT.
                   "rope_3d" for 3D rope (e.g., Qwen2.5-VL)
                   "rope_2d" for 2D rope (e.g., Kimi-VL)
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
    Returns:
        torch.Tensor: Output image embeddings

    Example:
        ```
        vision_model.out_hidden_size = 64
        vision_model.spatial_merge_size = 2
        pixel_values.shape = (1350, channel)
        grid_thw_list = [[1, 10, 100], [1, 10, 10], [1, 10, 20], [1, 50]]
        tp_size=2
        ```

    """
    tp_size = get_tensor_model_parallel_world_size()

    # GPU_0 tp_rank_local = 0
    # GPU_1 tp_rank_local = 1
    tp_rank_local = get_tensor_model_parallel_rank()

    # patches_per_image = [1000, 100, 200, 50]
    patches_per_image = [math.prod(grid_thw) for grid_thw in grid_thw_list]
    # patches_per_image = [0, 1000, 1100, 1300, 1350]
    cum_patches_per_image = [0, *itertools.accumulate(patches_per_image)]

    # Get load balancing assignment with all metadata
    # image_to_tp_rank = [0, 2, 1, 3]
    # gpu_sample_counts = [1, 3]
    # grouped_pixel_values_len = [1000, 350]
    (image_to_tp_rank, gpu_sample_counts,
     grouped_pixel_values_len) = get_load_balance_assignment(
         patches_per_image, tp_size)

    # cu_gpu_sample_counts = [0, 1, 4]
    cum_gpu_sample_counts = [0, *itertools.accumulate(gpu_sample_counts)]

    # GPU_0 image_idxs_local = [0]
    # GPU_1 image_idxs_local = [2, 1, 3]
    image_idxs_local = image_to_tp_rank[cum_gpu_sample_counts[tp_rank_local]:
                                        cum_gpu_sample_counts[tp_rank_local +
                                                              1]]

    # Get the pixel values for the local images based on the image_idxs_local
    if len(image_idxs_local) > 0:
        pixel_values_local = torch.cat([
            pixel_values[cum_patches_per_image[i]:cum_patches_per_image[i + 1]]
            for i in image_idxs_local
        ])
    else:
        # Handle case where this rank has no images
        pixel_values_local = torch.empty((0, pixel_values.shape[1]),
                                         device=pixel_values.device,
                                         dtype=pixel_values.dtype)
    # embed_dim_reduction_factor = 2 * 2
610
611
612
613
614
615
    if rope_type == "rope_2d":
        embed_dim_reduction_factor = (vision_model.merge_kernel_size[0] *
                                      vision_model.merge_kernel_size[1])
    else:
        embed_dim_reduction_factor = (vision_model.spatial_merge_size *
                                      vision_model.spatial_merge_size)
616
617
618
619
620
621
622
623
624
625

    # Find the max length across all ranks
    # The output embedding of every DP rank has to be
    # padded to this length for tensor_model_parallel_all_gather
    # to work
    max_len_per_rank = max(
        grouped_pixel_values_len) // embed_dim_reduction_factor
    local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]

    # Run the vision model on the local pixel_values_local
626
627
628
629
630
631
632
633
634
635
636
637
    if rope_type == "rope_2d":
        if pixel_values_local.shape[0] > 0:
            image_embeds_local = vision_model(
                pixel_values_local, torch.tensor(local_grid_thw_list))
            if isinstance(image_embeds_local, list):
                image_embeds_local = torch.cat(image_embeds_local, dim=0)
        else:
            out_dim = getattr(vision_model.config, "hidden_size", None)
            image_embeds_local = torch.empty(
                (0, embed_dim_reduction_factor, out_dim),
                device=pixel_values.device,
                dtype=pixel_values.dtype)
638
    else:
639
640
641
642
643
644
645
646
        if pixel_values_local.shape[0] > 0:
            image_embeds_local = vision_model(pixel_values_local,
                                              local_grid_thw_list)
        else:
            # Handle empty case
            image_embeds_local = torch.empty((0, vision_model.out_hidden_size),
                                             device=pixel_values.device,
                                             dtype=pixel_values.dtype)
647
648
649
650
651
652

    # Pad the output based on max_len_per_rank
    # for tensor_model_parallel_all_gather to work
    current_len = image_embeds_local.shape[0]
    if current_len < max_len_per_rank:
        padding_size = max_len_per_rank - current_len
653
654
655
656
657
658
659
660
661
        if rope_type == "rope_2d":
            padding = torch.empty((padding_size, image_embeds_local.shape[1],
                                   image_embeds_local.shape[2]),
                                  dtype=image_embeds_local.dtype,
                                  device=image_embeds_local.device)
        else:
            padding = torch.empty((padding_size, image_embeds_local.shape[1]),
                                  dtype=image_embeds_local.dtype,
                                  device=image_embeds_local.device)
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
        image_embeds_local_padded = torch.cat([image_embeds_local, padding],
                                              dim=0)
    else:
        image_embeds_local_padded = image_embeds_local

    # Do all_gather to collect embeddings from all ranks
    gathered_embeds = tensor_model_parallel_all_gather(
        image_embeds_local_padded, dim=0)

    # Remove padding and reconstruct per-rank embeddings
    rank_embeddings = list[torch.Tensor]()
    for rank in range(tp_size):
        start_idx = rank * max_len_per_rank
        end_idx = start_idx + (grouped_pixel_values_len[rank] //
                               embed_dim_reduction_factor)
        rank_embeddings.append(gathered_embeds[start_idx:end_idx])

    patches_per_output_image = [(patch_size // embed_dim_reduction_factor)
                                for patch_size in patches_per_image]

    # Reconstruct embeddings in the original order
    original_order_embeddings = [None] * len(grid_thw_list)
    current_idx = 0
    for rank in range(tp_size):
        count = gpu_sample_counts[rank]
        if count > 0:
            # Get images assigned to this rank in shuffled order
            # GPU_0 = image_idxs_local  [0]
            # GPU_1 = image_idxs_local  [2, 1, 3]
            rank_images = image_to_tp_rank[current_idx:current_idx + count]

            rank_embed = rank_embeddings[rank]
            # Split rank embeddings back to individual images
            embed_start = 0
            for img_idx in rank_images:
                img_patches = patches_per_output_image[img_idx]
                original_order_embeddings[img_idx] = rank_embed[
                    embed_start:embed_start + img_patches]
                embed_start += img_patches
            current_idx += count
    out_embeddings = tuple(embed for embed in original_order_embeddings
                           if embed is not None)
    assert len(out_embeddings) == len(
        original_order_embeddings), "Found unassigned embeddings"
    return out_embeddings


709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
def fetch_audio(
    audio_url: str,
    audio_io_kwargs: Optional[dict[str, Any]] = None,
) -> tuple[np.ndarray, Union[int, float]]:
    """
    Args:
        audio_url: URL of the audio file to fetch.
        audio_io_kwargs: Additional kwargs passed to handle audio IO.
    """
    media_io_kwargs = None if not audio_io_kwargs else {
        "audio": audio_io_kwargs
    }
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
    return media_connector.fetch_audio(audio_url)


def fetch_image(
    image_url: str,
    image_io_kwargs: Optional[dict[str, Any]] = None,
) -> Image.Image:
    """
    Args:
        image_url: URL of the image file to fetch.
        image_io_kwargs: Additional kwargs passed to handle image IO.
    """
    media_io_kwargs = None if not image_io_kwargs else {
        "image": image_io_kwargs
    }
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
    return media_connector.fetch_image(image_url)


def fetch_video(
    video_url: str,
    video_io_kwargs: Optional[dict[str, Any]] = None,
) -> tuple[npt.NDArray, dict[str, Any]]:
    """
    Args:
        video_url: URL of the video file to fetch.
        video_io_kwargs: Additional kwargs passed to handle video IO.
    """
    media_io_kwargs = None if not video_io_kwargs else {
        "video": video_io_kwargs
    }
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
754
    return media_connector.fetch_video(video_url)