[Fix] MM pool GPU alloc with base_gpu_id (#23377)
This commit is contained in:
@@ -277,6 +277,7 @@ class BaseMultimodalProcessor(ABC):
|
|||||||
self.cudaipc_mmfeature_pool = MmItemMemoryPool(
|
self.cudaipc_mmfeature_pool = MmItemMemoryPool(
|
||||||
per_worker_pool_size,
|
per_worker_pool_size,
|
||||||
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
|
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
|
||||||
|
self.server_args.base_gpu_id,
|
||||||
)
|
)
|
||||||
|
|
||||||
def compute_mrope_positions(self, input_ids, mm_items):
|
def compute_mrope_positions(self, input_ids, mm_items):
|
||||||
|
|||||||
@@ -122,9 +122,9 @@ class MmItemMemoryChunk:
|
|||||||
|
|
||||||
|
|
||||||
class MmItemMemoryPool:
|
class MmItemMemoryPool:
|
||||||
def __init__(self, memory_size, recycle_interval):
|
def __init__(self, memory_size, recycle_interval, base_gpu_id):
|
||||||
self.memory_pool = torch.empty(
|
self.memory_pool = torch.empty(
|
||||||
memory_size, dtype=torch.int8, device="cuda"
|
memory_size, dtype=torch.int8, device=f"cuda:{base_gpu_id}"
|
||||||
).contiguous()
|
).contiguous()
|
||||||
storage = self.memory_pool.untyped_storage()
|
storage = self.memory_pool.untyped_storage()
|
||||||
self._pool_ipc_handle = storage._share_cuda_()
|
self._pool_ipc_handle = storage._share_cuda_()
|
||||||
|
|||||||
Reference in New Issue
Block a user