Skip to content

[LongcatFlash] Fix test_longcat_generation_cpu: use device_map="cpu" to avoid MoE disk offload issue #48377

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ydshieh merged 3 commits into from
Aug 27, 2026
Merged

[LongcatFlash] Fix test_longcat_generation_cpu: use device_map="cpu" to avoid MoE disk offload issue #48377

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f367561
[LongcatFlash] Fix test_longcat_generation_cpu by using device_map="cpu"
ydshieh Aug 27, 2026
77f601a
[LongcatFlash] Update golden string and skip test_longcat_generation_…
ydshieh Aug 27, 2026
fb15ad5
remove unused require_large_cpu_ram import
ydshieh Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions tests/models/longcat_flash/test_modeling_longcat_flash.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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)
Expand Down
Loading