Fix `Gemma2IntegrationTest` by ydshieh · Pull Request #38492 · huggingface/transformers

ydshieh · 2025-05-30T10:19:47Z

What does this PR do?

Currently on singe-gpu T4 runner, test_export_hybrid_cache use more than 60G CPU RAM memory and the process is killed, so we don't receive any report.

This PR skips this tests by introducing a new require_large_cpu_ram. I also update some tests to make them pass and to avoid some OOM.

On T4: only 2 flex attn. related tests are failing.
On A10: only 1 flex attn. related test is failing.

HuggingFaceDocBuilderDev · 2025-05-30T10:32:53Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

ydshieh · 2025-05-30T11:08:58Z

    return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case)


+def require_large_cpu_ram(test_case, memory: float = 80):


Our single-gpu runner has 60G CPU RAM.

Currently, test_export_hybrid_cache is killing pytest process.

Nit, but why default to 80?

No particular reason: I set a value so it won't run on our single gpu runner (T4/A10). I don't try to figure out how much RAM this test really requires. And I don't want to set 60 neither as the RAM on our single gpu runner it is not exactly 60 G

ydshieh · 2025-05-30T11:10:02Z

+        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+            self.skipTest("FlashAttention2 is required for this test.")


Not gemma2, but just update here anyway.

This test is @parameterized.expand so skip inside the body if flash attn is not installed.

Totally fair lol

ydshieh · 2025-05-30T11:10:44Z

        self.assertEqual(output_text, EXPECTED_TEXTS)

-    @tooslow
+    @require_torch_large_accelerator


T4 can't run 9b, but A10 can

Hmm, do you know what the previous integration tests used then? I guess the A100s?

Makes sense that the expected generation changes then.

It's always the same test. While with @tooslow, it's never run at all. Now it will run on A10 but not T4.

Oh you mean the previous expected values? I have no idea, it depends on the contributor or the team members who wrote the tests. It's very likely not A10 however.

ydshieh · 2025-05-30T11:11:46Z


    @slow
    @require_read_token
+    @require_large_cpu_ram


@guangy10 if you can check why this test take so much CPU RAM ?

@ydshieh Hmm interesting if it requires large CPU RAM as it essentially same as the Static one since only the global cache is exposed to export. BTW how do you notice the RAM usage surge? I'd like to understand if there is any utils/tools/ci that would help debug this kind of issue?

BTW how do you notice the RAM usage surge?

I just check the CI job run page and find killed and the process exit 😅

See this run

https://github.com/huggingface/transformers/actions/runs/15382744712/job/43276232860

as it essentially same as the Static one since only the global cache is exposed to export

BTW, maybe the other test is just taking a blit less RAM under 60GB and we get lucky ...

ydshieh · 2025-06-02T12:13:24Z

@vasqu If you want to take a look too as Cyrilvallez if off. No worry if not.

vasqu

LGTM, just some nits / questions.

Just curious to know why flex attention fails but it might've also relied on some old version which didnt use the block mask. In general, generation with flex attn needs a rework :D

vasqu · 2025-06-02T16:25:52Z

    return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case)


+def require_large_cpu_ram(test_case, memory: float = 80):


Nit, but why default to 80?

vasqu · 2025-06-02T16:27:47Z

+        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+            self.skipTest("FlashAttention2 is required for this test.")


Totally fair lol

vasqu · 2025-06-02T16:31:36Z

        self.assertEqual(output_text, EXPECTED_TEXTS)

-    @tooslow
+    @require_torch_large_accelerator


Hmm, do you know what the previous integration tests used then? I guess the A100s?

Makes sense that the expected generation changes then.

ydshieh · 2025-06-02T20:39:08Z

Just curious to know why flex attention fail

I have to say I don't know much about this (and in general flex failing tests ). Sometimes, some flex tests fail on torch 2.6 but works on torch 2.7 (on T4), but works on A10 even with torch 2.6, see for example

https://huggingface.slack.com/archives/C06LR9PQA00/p1748515989352219

And maybe let me ping you for On A10: only 1 flex attn. related test is failing. once a run after this PR merged is finished 😄

ydshieh · 2025-06-03T09:04:05Z

+        EXPECTED_BATCH_TEXTS = Expectations(
+            {
+                ("cuda", 8): [
+                    "Hello I am doing a project on the 1960s and I am trying to find out what the average",
+                    "Hi today I'm going to be talking about the 10 most powerful characters in the Naruto series.",
+                ]
+            }
+        )
+        EXPECTED_BATCH_TEXT = EXPECTED_BATCH_TEXTS.get_expectation()


@yao-matrix I am trying to use Expectations class when I update some updates. If you find something is broken on XPU because I miss some details, let me know.

Cyrilvallez · 2025-06-10T10:37:03Z

Hey! Just seeing this now that I'm back from vacation! The export test has absolutely no reason to use that much cpu memory... I mean, I don't know the inner workings of export, but this is way too much. @ydshieh, is it possible that it's just pytest using too many processes at the same time, resulting in killing some processes at some point? If not, any way you could quickly investigate @vasqu?

Cyrilvallez · 2025-06-10T10:51:13Z

This is the kind of stuff that's not super good to skip without knowing what's happening - maybe a matter of retaining a lot of different torch graphs on cpu as well? Not sure if the compile/export graphs are being cleaned up properly between tests etc

ydshieh · 2025-06-10T11:13:14Z

@Cyrilvallez On daily CI, there is only one pytest process, i.e. we are not using -n 8 like on CircleCI.

This is the kind of stuff that's not super good to skip without knowing what's happening

I agree, but IMO, it would be good for @guangy10 to investigate this to see something like

is it static cache also using large CPU RAM (just not more than 60G so still passing)
Any evidence if hybrid cache have this issue (in general or only for Gemma 2)
If this happens even not in the test, but just as a standalone python script
...

etc.

If we keep this test, other tests are not reported as the job is killed, which is not good neither for maintenance.

(if it fails but not affecting other tests, I might be happy to keep running though)

ydshieh · 2025-06-10T12:02:56Z

BTW, it still run on our mulit-gpu CI runners (which have more CPU RAM)

ydshieh added 14 commits May 29, 2025 18:41

fix

d508c9a

fix

f834b42

skip-ci

d47a337

skip-ci

19b3243

skip-ci

bc2d024

skip-ci

d7fff95

skip-ci

9e259ff

skip-ci

84a9857

skip-ci

a017ecf

skip-ci

c09811e

skip-ci

a286fd1

skip-ci

0513f17

skip-ci

03742dc

update

0a70cf0

ydshieh added 2 commits May 30, 2025 12:35

fix

72fbbb8

fix

c964b2f

ydshieh commented May 30, 2025

View reviewed changes

ydshieh requested a review from Cyrilvallez May 30, 2025 11:12

vasqu approved these changes Jun 2, 2025

View reviewed changes

ydshieh merged commit ccc8596 into main Jun 2, 2025

ydshieh deleted the fix_gemma2 branch June 2, 2025 20:45

ydshieh commented Jun 3, 2025

View reviewed changes

yao-matrix mentioned this pull request Jun 4, 2025

enable more test cases on xpu #38572

Merged

		return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case)


		def require_large_cpu_ram(test_case, memory: float = 80):

		if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
		self.skipTest("FlashAttention2 is required for this test.")

Uh oh!

Conversation

ydshieh commented May 30, 2025

What does this PR do?

Uh oh!

HuggingFaceDocBuilderDev commented May 30, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydshieh May 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydshieh commented Jun 2, 2025

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydshieh commented Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Cyrilvallez commented Jun 10, 2025

Uh oh!

Cyrilvallez commented Jun 10, 2025

Uh oh!

ydshieh commented Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ydshieh commented Jun 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ydshieh May 30, 2025 •

edited

Loading

ydshieh commented Jun 2, 2025 •

edited

Loading

ydshieh commented Jun 10, 2025 •

edited

Loading