|
| 1 | +"""Hold vision-model peak VRAM when co-located with the LLM. |
| 2 | +
|
| 3 | +Build worst-case ImageItem inputs (max pixel budget, extreme aspect ratio, solid RGB), |
| 4 | +run the normal encode() path, then keep the CUDA allocator high-water mark. |
| 5 | +""" |
| 6 | + |
| 7 | +import math |
| 8 | +import torch |
| 9 | +import torch.distributed as dist |
| 10 | +from io import BytesIO |
| 11 | +from typing import List, Tuple |
| 12 | +from PIL import Image |
| 13 | +from lightllm.server.embed_cache.utils import create_shm, free_shm, get_shm_name_data |
| 14 | +from lightllm.server.multimodal_params import ImageItem |
| 15 | +from lightllm.utils.envs_utils import get_unique_server_name |
| 16 | +from lightllm.utils.log_utils import init_logger |
| 17 | + |
| 18 | +logger = init_logger(__name__) |
| 19 | + |
| 20 | +_PEAK_VRAM_HOLD_OOM_HINT = ( |
| 21 | + "Vision peak VRAM hold probe hit OOM. Lower --visual_infer_batch_size, " |
| 22 | + "--max_image_pixels, or --max_image_token_count, or place the vision model on a " |
| 23 | + "separate GPU with --visual_gpu_ids." |
| 24 | +) |
| 25 | + |
| 26 | + |
| 27 | +class VisionPeakVramHolder: |
| 28 | + """Hold peak vision VRAM at startup via worst-case ImageItem encode().""" |
| 29 | + |
| 30 | + # Match Qwen-VL smart_resize MAX_RATIO; extreme aspect ratio stresses vision sequence length. |
| 31 | + MAX_ASPECT_RATIO = 200 |
| 32 | + |
| 33 | + def __init__(self, model): |
| 34 | + self.model = model |
| 35 | + |
| 36 | + @torch.no_grad() |
| 37 | + def hold( |
| 38 | + self, |
| 39 | + device_id: int, |
| 40 | + batch_size: int, |
| 41 | + max_image_pixels: int, |
| 42 | + max_image_token_count: int, |
| 43 | + tp_rank_id: int = 0, |
| 44 | + vit_tp: int = 1, |
| 45 | + dp_rank_id: int = 0, |
| 46 | + ) -> int: |
| 47 | + torch.cuda.set_device(device_id) |
| 48 | + baseline_reserved = torch.cuda.memory_reserved(device_id) |
| 49 | + torch.cuda.reset_peak_memory_stats(device_id) |
| 50 | + |
| 51 | + gloo_group = dist.new_group(ranks=list(range(vit_tp)), backend="gloo") |
| 52 | + |
| 53 | + image_items = self._prepare_worst_case_image_items( |
| 54 | + batch_size=batch_size, |
| 55 | + max_image_pixels=max_image_pixels, |
| 56 | + max_image_token_count=max_image_token_count, |
| 57 | + tp_rank_id=tp_rank_id, |
| 58 | + vit_tp=vit_tp, |
| 59 | + dp_rank_id=dp_rank_id, |
| 60 | + gloo_group=gloo_group, |
| 61 | + ) |
| 62 | + |
| 63 | + try: |
| 64 | + out = self.model.encode(image_items) |
| 65 | + del out |
| 66 | + except (RuntimeError, torch.OutOfMemoryError) as e: |
| 67 | + logger.exception(str(e)) |
| 68 | + raise Exception(_PEAK_VRAM_HOLD_OOM_HINT) |
| 69 | + finally: |
| 70 | + dist.barrier(group=gloo_group) |
| 71 | + |
| 72 | + if tp_rank_id == 0: |
| 73 | + self._free_image_items(image_items) |
| 74 | + |
| 75 | + peak_reserved = torch.cuda.max_memory_reserved(device_id) |
| 76 | + return int(max(0, peak_reserved - baseline_reserved)) |
| 77 | + |
| 78 | + def _worst_case_image_size(self, max_image_pixels: int) -> Tuple[int, int]: |
| 79 | + """Largest valid landscape aspect ratio under the pixel budget (width >= height).""" |
| 80 | + height = max(1, int(math.sqrt(max_image_pixels / self.MAX_ASPECT_RATIO))) |
| 81 | + width = min(max_image_pixels // height, height * self.MAX_ASPECT_RATIO) |
| 82 | + height = int(max(1, height)) |
| 83 | + width = int(max(1, width)) |
| 84 | + return width, height |
| 85 | + |
| 86 | + def _gen_rgb_jpeg_bytes(self, width: int, height: int, color: Tuple[int, int, int] = (0, 0, 0)) -> bytes: |
| 87 | + buffer = BytesIO() |
| 88 | + Image.new("RGB", (width, height), color).save(buffer, format="JPEG", quality=95) |
| 89 | + return buffer.getvalue() |
| 90 | + |
| 91 | + def _prepare_worst_case_image_items( |
| 92 | + self, |
| 93 | + batch_size: int, |
| 94 | + max_image_pixels: int, |
| 95 | + max_image_token_count: int, |
| 96 | + tp_rank_id: int, |
| 97 | + vit_tp: int, |
| 98 | + dp_rank_id: int, |
| 99 | + gloo_group, |
| 100 | + ) -> List[ImageItem]: |
| 101 | + """Return probe ImageItems; under tp>1 only rank 0 writes shm and broadcasts the list.""" |
| 102 | + if vit_tp == 1: |
| 103 | + return self._build_worst_case_image_items(batch_size, max_image_pixels, max_image_token_count, dp_rank_id) |
| 104 | + |
| 105 | + payload: List[ImageItem] = [None] |
| 106 | + if tp_rank_id == 0: |
| 107 | + payload[0] = self._build_worst_case_image_items( |
| 108 | + batch_size, max_image_pixels, max_image_token_count, dp_rank_id |
| 109 | + ) |
| 110 | + dist.broadcast_object_list(payload, src=0, group=gloo_group) |
| 111 | + return payload[0] |
| 112 | + |
| 113 | + def _build_worst_case_image_items( |
| 114 | + self, batch_size: int, max_image_pixels: int, max_image_token_count: int, dp_rank_id: int |
| 115 | + ) -> List[ImageItem]: |
| 116 | + width, height = self._worst_case_image_size(max_image_pixels) |
| 117 | + items = [] |
| 118 | + for batch_id in range(batch_size): |
| 119 | + # Alternate solid black/white so each batch slot gets a distinct JPEG payload. |
| 120 | + color = (255, 255, 255) if batch_id % 2 else (0, 0, 0) |
| 121 | + image_bytes = self._gen_rgb_jpeg_bytes(width, height, color=color) |
| 122 | + item = ImageItem(type="base64", data="") |
| 123 | + item.uuid = f"{get_unique_server_name()}_vision_peak_hold_dp{dp_rank_id}_{batch_id}" |
| 124 | + item.image_w = width |
| 125 | + item.image_h = height |
| 126 | + # InternVL encode() reads image_patch_max_num from extra_params (normally set by |
| 127 | + # InternvlTokenizer.init_imageitem_extral_params). Use MAX_PATH_NUM for worst-case tile count. |
| 128 | + if hasattr(self.model, "MAX_PATH_NUM"): |
| 129 | + item.extra_params["image_patch_max_num"] = int(self.model.MAX_PATH_NUM) |
| 130 | + |
| 131 | + create_shm(get_shm_name_data(item.uuid), image_bytes) |
| 132 | + items.append(item) |
| 133 | + return items |
| 134 | + |
| 135 | + @staticmethod |
| 136 | + def _free_image_items(items: List[ImageItem]) -> None: |
| 137 | + for item in items: |
| 138 | + try: |
| 139 | + free_shm(get_shm_name_data(item.uuid)) |
| 140 | + except FileNotFoundError: |
| 141 | + pass |
0 commit comments