diff --git a/tests/models/longcat_flash/test_modeling_longcat_flash.py b/tests/models/longcat_flash/test_modeling_longcat_flash.py index 25556e6c4165..8721b217801e 100644 --- a/tests/models/longcat_flash/test_modeling_longcat_flash.py +++ b/tests/models/longcat_flash/test_modeling_longcat_flash.py @@ -23,9 +23,9 @@ from transformers.testing_utils import ( require_bitsandbytes, require_flash_attn, - require_large_cpu_ram, require_torch, require_torch_accelerator, + require_torch_accelerator_memory, slow, torch_device, ) @@ -342,12 +342,19 @@ def test_shortcat_generation(self): outputs = self.model.generate(inputs["input_ids"], max_new_tokens=10, do_sample=False) response = self.tokenizer.batch_decode(outputs, skip_special_tokens=False)[0] - expected_output = "[Round 0] USER:Paris is... ASSISTANT: dig年车龄juanaheast稍achaotingupebarebones" + expected_output = "[Round 0] USER:Paris is... ASSISTANT: dig affidavhens酒店的coco toughoireFFIXilateraldrawal" self.assertEqual(response, expected_output) + # `meituan-longcat/LongCat-Flash-Chat` is 562B parameters (~18.6–31.3B activated per token) -- + # ~1,047 GiB of bfloat16 weights. + # + # That far exceeds the budget `device_map="auto"` plans against on the daily CI `a10` runners -- + # 24 GiB of accelerator plus the 60 GiB `CI_CPU_MEMORY_LIMIT_GB` allowance on the single-accelerator + # runner, and 48 + 120 on the two-accelerator one -- so a large portion of the model is placed on + # `"disk"`, and loading dies due to MoE weight format incompatibility with accelerate's disk offload. @slow - @require_large_cpu_ram + @require_torch_accelerator_memory(memory=1100) def test_longcat_generation_cpu(self): # takes absolutely forever and a lot RAM, but allows to test the output in the CI model = LongcatFlashForCausalLM.from_pretrained(self.model_id, device_map="auto", dtype=torch.bfloat16)