make mistral3 pass on xpu

[yao-matrix]

@ydshieh , pls help review, thx

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[Rocketknight1]

cc @IlyasMoutawwakil

[ydshieh]

Thanks.

Looks line non of those integration tests could be run on T4, all GPU OOM.

Let me try first if we have other workaround.

[ydshieh]

Would you be up to try

git fetch https://github.com/yao-matrix/transformers.git mistral3-xpu-cpu-offload:mistral3-xpu-cpu-offload && git checkout mistral3-xpu-cpu-offload

and run the integration tests? I am using CPU offload, so 3 tests can run on A10.

I can't make test_mistral3_integration_batched_generate_multi_image work as it will OOM without 4-bit and other errors with cpu-offload + 4-bit.

FAILED tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate_multi_image - ValueError: Blockwise quantization only supports 16/32-bit floats, but got torch.uint8

On T4, it hangs forever ...

[ydshieh]

Thanks.

Looks line non of those integration tests could be run on T4, all GPU OOM.

Let me try first if we have other workaround.

Hmm, I am able to avoid OOM for test_mistral3_integration_batched_generate_multi_image by using smaller images.

But before I move forward, it would be nice if you can check if this cpu_offload works with xpu 🙏 ?

[yao-matrix]

Would you be up to try

git fetch https://github.com/yao-matrix/transformers.git mistral3-xpu-cpu-offload:mistral3-xpu-cpu-offload && git checkout mistral3-xpu-cpu-offload

and run the integration tests? I am using CPU offload, so 3 tests can run on A10.

I can't make test_mistral3_integration_batched_generate_multi_image work as it will OOM without 4-bit and other errors with cpu-offload + 4-bit.

FAILED tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate_multi_image - ValueError: Blockwise quantization only supports 16/32-bit floats, but got torch.uint8

On T4, it hangs forever ...

@ydshieh , sorry for late response, just back to office from a 5-day holiday. Yes, I can try you offload changes, but it seems I cannot get the mistral3-xpu-cpu-offload branch when I run your command, could you help check it? Thx.

[ydshieh]

Sorry, it should be

git fetch https://github.com/huggingface/transformers.git mistral3-xpu-cpu-offload:mistral3-xpu-cpu-offload && git checkout mistral3-xpu-cpu-offload

You don't need to update the expected values, just to see if xpu works well with cpu offloading.
Once I change the input images to new smaller one, we can check the expected values again 🙏

[ydshieh]

BTW, it seems this cpu offload will produce different outputs on sigle-device v.s. multi-device environment.

I have to set execution_device="cuda:0") to get the same outputs in both cases.

Anyway, let's see if it could at least run on xpu without error then we can adjust the outputs.

[yao-matrix]

@ydshieh , I tested the 4 cases w/ cpu_offload() on XPU:

• 2 cases failed for ground-truth mismatch, which are fine
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate_multi_image

• 2 cases passed
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_generate
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_generate_text_only

So it's work on XPU.

Yet, I found the test will be pretty slow after enabling cpu offload, and if we run 4 cases in one pytest command, the process is easy to hang(using pytest -rA tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest), and I found process VIRT memory will be tens of T(33.6T in one of my experiment), separately run each case is fine. Don't know whether you can observe similar in your env.

If you can observe similar w/ me, I will think it's not suitable to use cpu_offload since it has risk to break the CI. And maybe we can consider the require_big_accelerator decorator used by diffusers to specify these cases only run on accelerator whose VRAM are larger than a given number.

[ydshieh]

You mean 33.6G instead of 33.6T, right? Yes, it is (CPU) memory hungry, but not necessary slower than running GPU (let me double check).

BTW, how many CPU RAM available on your XPU machine? In my cases, I need to ask our infra to provide single T4 with 64G.
Also, how much is your XPU accelerator's processing RAM? On A10 24GB, I also get one or two of them failing (OOM).

[yao-matrix]

You mean 33.6G instead of 33.6T, right? Yes, it is (CPU) memory hungry, but not necessary slower than running GPU (let me double check).

BTW, how many CPU RAM available on your XPU machine? In my cases, I need to ask our infra to provide single T4 with 64G. Also, how much is your XPU accelerator's processing RAM? On A10 24GB, I also get one or two of them failing (OOM).

Actually it's 33.6 T shows in my VIRT, so when run into 3rd case, it often hangs. And I validated in my A100 env, it's OK. So, it's my env issue, you can ignore it.

I am using the Ponte Vecchio 1150 which has 64GB VRAM, suppose it's enough.

[faaany]

@ydshieh , I tested the 4 cases w/ cpu_offload() on XPU:

• 2 cases failed for ground-truth mismatch, which are fine
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate_multi_image
• 2 cases passed
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_generate
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_generate_text_only

So it's work on XPU.

Yet, I found the test will be pretty slow after enabling cpu offload, and if we run 4 cases in one pytest command, the process is easy to hang(using pytest -rA tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest), and I found process VIRT memory will be tens of T(33.6T in one of my experiment), separately run each case is fine. Don't know whether you can observe similar in your env.

If you can observe similar w/ me, I will think it's not suitable to use cpu_offload since it has risk to break the CI. And maybe we can consider the require_big_accelerator decorator used by diffusers to specify these cases only run on accelerator whose VRAM are larger than a given number.

In my env, I found the slowness mainly comes from model downloading. Once the model is downloaded, the test can pass pretty fast.

[yao-matrix]

@ydshieh , I tested the 4 cases w/ cpu_offload() on XPU:

• 2 cases failed for ground-truth mismatch, which are fine
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_batched_generate_multi_image
• 2 cases passed
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_generate
  tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest::test_mistral3_integration_generate_text_only

So it's work on XPU.
Yet, I found the test will be pretty slow after enabling cpu offload, and if we run 4 cases in one pytest command, the process is easy to hang(using pytest -rA tests/models/mistral3/test_modeling_mistral3.py::Mistral3IntegrationTest), and I found process VIRT memory will be tens of T(33.6T in one of my experiment), separately run each case is fine. Don't know whether you can observe similar in your env.
If you can observe similar w/ me, I will think it's not suitable to use cpu_offload since it has risk to break the CI. And maybe we can consider the require_big_accelerator decorator used by diffusers to specify these cases only run on accelerator whose VRAM are larger than a given number.

In my env, I found the slowness mainly comes from model downloading. Once the model is downloaded, the test can pass pretty fast.

so, it's my env issue. Thx @faaany for testing.

[ydshieh]

OK, thank you both very much. Let's try cpu offloading so T4 and A10 and use smaller images can run (most) of them. I will push some commits back to this PR.

[ydshieh]

Hi @yao-matrix (after deleting the local mistral3-xpu if it is still on your local system)

git fetch https://github.com/yao-matrix/transformers.git mistral3-xpu:mistral3-xpu && git checkout mistral3-xpu

This will run on both T4 16G and A10 24G without any GPU OOM and match the expect values.

If you want to keep tests running directly on XPU without using CPU offload, you can tweak

    def setUp(self):
        cleanup(torch_device, gc_collect=True)
        self.model_checkpoint = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
        self.model = Mistral3ForConditionalGeneration.from_pretrained(
            self.model_checkpoint, torch_dtype=torch.bfloat16
        )
        accelerate.cpu_offload(self.model, execution_device=torch_device)

with an if / else. Once you are happy with the results, ping me and I will merge.

[yao-matrix]

@ydshieh , cool, OK from my side. Pls feel free to merge.

[ydshieh]

I think XPU will still get 2 cases failed for ground-truth mismatch as you mentioned (but you said which are fine).
I will merge as it is then if XPU needs an update of expected values, we can do it in another PR.

Thank you for the patience.

[HuggingFaceDocBuilderDev]

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.