[VLM] Fix CUDA IPC OOM (#16118)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
@@ -1574,6 +1574,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
mm_item.feature = pixel_values.reconstruct_on_target_device(
|
||||
torch.cuda.current_device()
|
||||
)
|
||||
# The reference by CudaIpcTensorTransportProxy was cut off,
|
||||
# proactively delete to avoid slow gc.
|
||||
del pixel_values
|
||||
self.multimodal_inputs = multimodal_inputs
|
||||
self.token_type_ids = token_type_ids_tensor
|
||||
self.seq_lens_sum = sum(seq_lens)
|
||||
|
||||
@@ -880,6 +880,8 @@ class BaseMultimodalProcessor(ABC):
|
||||
info_data=item.feature,
|
||||
sync_buffer_meta=sync_flag,
|
||||
)
|
||||
elif not self.server_args.keep_mm_feature_on_device:
|
||||
item.feature = item.feature.cpu()
|
||||
elif (
|
||||
isinstance(item.precomputed_embeddings, torch.Tensor)
|
||||
and item.precomputed_embeddings.is_cuda
|
||||
@@ -900,5 +902,7 @@ class BaseMultimodalProcessor(ABC):
|
||||
info_data=item.precomputed_embeddings,
|
||||
sync_buffer_meta=sync_flag,
|
||||
)
|
||||
elif not self.server_args.keep_mm_feature_on_device:
|
||||
item.precomputed_embeddings = item.precomputed_embeddings.cpu()
|
||||
|
||||
return all_collected_items, input_ids, ret
|
||||
|
||||
Reference in New Issue
Block a user