[MISC] Consolidate cleanup() and refactor offline_inference_with_prefix.py (#9510)

This commit is contained in:
Cody Yu
2024-10-18 14:30:55 -07:00
committed by GitHub
parent 9bb10a7d27
commit d11bf435a0
20 changed files with 84 additions and 105 deletions

View File

@@ -6,7 +6,7 @@ import torch.nn as nn
from huggingface_hub import snapshot_download
from transformers import AutoConfig, AutoModel, CLIPImageProcessor
from ....conftest import _ImageAssets, cleanup
from ....conftest import _ImageAssets
# we use snapshot_download to prevent conflicts between
# dynamic_module and trust_remote_code for hf_runner
@@ -45,12 +45,13 @@ def run_intern_vit_test(
for pixel_value in pixel_values
]
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.model_executor.models.intern_vit import InternVisionModel
vllm_model = InternVisionModel(config)
vllm_model.load_weights(hf_model.state_dict().items())
del hf_model
cleanup()
cleanup_dist_env_and_memory()
vllm_model = vllm_model.to("cuda", dtype)
vllm_outputs_per_image = [
@@ -58,7 +59,7 @@ def run_intern_vit_test(
for pixel_value in pixel_values
]
del vllm_model
cleanup()
cleanup_dist_env_and_memory()
cos_similar = nn.CosineSimilarity(dim=-1)
for vllm_output, hf_output in zip(vllm_outputs_per_image,