diff --git a/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py b/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py index 63e64de60b93..d61c0f7bd964 100644 --- a/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py +++ b/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py @@ -13,6 +13,7 @@ # limitations under the License. """Testing suite for the PyTorch Qwen3VLMoe model.""" +import tempfile import unittest import pytest @@ -27,6 +28,8 @@ from transformers.models.qwen3_vl_moe.configuration_qwen3_vl_moe import Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig from transformers.testing_utils import ( Expectations, + backend_device_count, + get_cpu_ram_total_gib, require_flash_attn, require_torch, require_torch_accelerator, @@ -301,6 +304,49 @@ def test_reverse_loading_mapping(self, check_keys_were_modified=False): class Qwen3VLMoeIntegrationTest(MemoryCleanupMixin, unittest.TestCase): maxDiff = None + @classmethod + def setUpClass(cls): + super().setUpClass() + # Assigned here rather than in the class body: `MemoryCleanupMixin` snapshots the class body and + # protects it from teardown, so a checkpoint parked on a class-body `model` is never released. + cls.model = None + cls.offload_dir = None + + @classmethod + def get_model(cls): + if cls.model is None: + cls.offload_dir = tempfile.TemporaryDirectory() + # device_map="auto" fills GPUs to ~100%, leaving no room for the transpose temporary + # buffer that fuses per-expert weight shards into a single gate_up_proj tensor during + # from_pretrained — causing CUDA OOM on multi-GPU runners. A 70% per-GPU max_memory cap + # reserves the headroom. + n = backend_device_count(torch_device) + if n > 0 and torch_device != "cpu": + torch_accel = getattr(torch, torch_device) + per_device = int( + min(torch_accel.get_device_properties(i).total_memory for i in range(n)) * 0.70 / 1024**3 + ) + max_memory = dict.fromkeys(range(n), f"{per_device}GiB") + max_memory["cpu"] = ( + f"{int(get_cpu_ram_total_gib() * 0.9)}GiB" # To avoid runner failing with exit code 137. + ) + else: + max_memory = None + cls.model = Qwen3VLMoeForConditionalGeneration.from_pretrained( + "Qwen/Qwen3-VL-30B-A3B-Instruct", + dtype="auto", + device_map="auto", + max_memory=max_memory, + offload_folder=cls.offload_dir.name, + ) + return cls.model + + @classmethod + def tearDownClass(cls): + if cls.offload_dir is not None: + cls.offload_dir.cleanup() + super().tearDownClass() + def setUp(self): super().setUp() @@ -351,9 +397,7 @@ def setUp(self): @slow def test_small_model_integration_test(self): - model = Qwen3VLMoeForConditionalGeneration.from_pretrained( - "Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto" - ) + model = self.get_model() inputs = self.processor.apply_chat_template( self.message, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" @@ -394,9 +438,7 @@ def test_small_model_integration_test(self): @slow def test_small_model_integration_test_batch(self): - model = Qwen3VLMoeForConditionalGeneration.from_pretrained( - "Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto" - ) + model = self.get_model() batch_messages = [self.message] * 2 inputs = self.processor.apply_chat_template( batch_messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" @@ -462,9 +504,7 @@ def test_small_model_integration_test_with_video(self): @slow def test_small_model_integration_test_expand(self): - model = Qwen3VLMoeForConditionalGeneration.from_pretrained( - "Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto" - ) + model = self.get_model() inputs = self.processor.apply_chat_template( self.message, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" ).to(torch_device) @@ -490,9 +530,7 @@ def test_small_model_integration_test_expand(self): @slow def test_small_model_integration_test_expand_with_video(self): - model = Qwen3VLMoeForConditionalGeneration.from_pretrained( - "Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto" - ) + model = self.get_model() inputs = self.processor.apply_chat_template( self.message3, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" ).to(torch_device) @@ -511,9 +549,7 @@ def test_small_model_integration_test_expand_with_video(self): @slow def test_small_model_integration_test_batch_wo_image(self): - model = Qwen3VLMoeForConditionalGeneration.from_pretrained( - "Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto" - ) + model = self.get_model() message_wo_image = [ {"role": "user", "content": [{"type": "text", "text": "Who are you?"}]}, ] @@ -541,9 +577,7 @@ def test_small_model_integration_test_batch_wo_image(self): @slow def test_small_model_integration_test_batch_different_resolutions(self): - model = Qwen3VLMoeForConditionalGeneration.from_pretrained( - "Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto" - ) + model = self.get_model() batched_messages = [self.message, self.message2] inputs = self.processor.apply_chat_template( batched_messages,