Optimize Qwen3-VL video memory usage (#16366)

This commit is contained in:
cen121212
2026-01-22 09:10:08 +08:00
committed by GitHub
parent 590969ee9c
commit 0c2993eed0

View File

@@ -852,10 +852,18 @@ class Qwen3VLForConditionalGeneration(nn.Module):
return torch.cat(all_chunk_embeds, dim=0)
def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
for item in items:
item.feature = item.feature.to(self.visual.device)
# in qwen-vl, last dim is the same
pixel_values = torch.cat([item.feature for item in items], dim=0).type(
self.visual.dtype
)
# Memory optimization for item.feature:
# 1. item.feature is released when request finished
# 2. High concurrency may cause device OOM due to delayed release
# 3. Fix: Offload item.feature to CPU, move to device only when needed
for item in items:
item.feature = item.feature.to("cpu")
video_grid_thw = torch.concat([item.video_grid_thw for item in items], dim=0)
assert pixel_values.dim() == 2, pixel_values.dim()
assert video_grid_thw.dim() == 2, video_grid_thw.dim()