From e662af769bfacd5755449e87fd62855ec86f3680 Mon Sep 17 00:00:00 2001 From: JoshuaL3000 Date: Wed, 29 Jul 2026 06:38:57 +0800 Subject: [PATCH] fix: enable XPU support and update hardcoded CUDA selections for tests (#7401) * fix: add XPU device support and update hardcoded CUDA selections * fix: add XPU device support for pytest CUDA skipped tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix device handling for PR #7401 - perplexity_eval.py: use DEVICE_TYPE_TORCH, not DEVICE_TYPE. The latter can be "hip" or "mlx", which .to() rejects, so this regressed ROCm. - test_batched_leftpad_generation_gpu.py: XPU diverges here today, so mark it non-strict xfail on XPU instead of reverting to a CUDA-only guard. Keeps the real XPU gap visible and turns green once it is fixed. - Guard torch.xpu.is_available() with hasattr, matching device_type.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Re-enable the flash varlen attention test in CI for PR #7401 attention_dispatch.py now predefines flash_attn_func / flash_attn_varlen_func as None, so test_run_attention_flash_varlen_receives_window_and_softcap no longer needs flash_attn importable to be monkeypatched. Verified on a runner shaped like the CPU-only one: the test fails against main's attention_dispatch and passes at this head, so the deselect is now dead weight. * Tighten comments for PR #7401 Drop the hasattr rationale: torch.xpu has existed since torch 2.3 and the dependency floor is 2.4, so no supported build predates the namespace. The guard stays as cheap defence, but the comment claimed something untrue. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- .github/workflows/consolidated-tests-ci.yml | 10 +++--- .../test_merge_model_perplexity_llama-3.2.py | 11 +++---- .../test_merge_model_perplexity_mistral.py | 11 +++---- .../test_merge_model_perplexity_phi_4.py | 11 +++---- ...st_merged_model_perplexity_llama-3.1-8b.py | 11 +++---- .../test_merged_model_perplexity_qwen_2.5.py | 13 +++----- tests/test_fp8_tiny_e8m0.py | 10 +++--- tests/utils/perplexity_eval.py | 5 ++- .../test_batched_leftpad_generation_gpu.py | 17 ++++++++-- tests/utils/test_packing.py | 20 +++++++++--- tests/utils/test_qat.py | 8 ++++- tests/utils/test_rope_scaling_drift.py | 32 ++++++++++--------- unsloth/utils/attention_dispatch.py | 2 ++ 13 files changed, 94 insertions(+), 67 deletions(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 489ee4ca08..c75880fa72 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -372,12 +372,10 @@ jobs: tests/python/test_fast_language_model_text_only.py \ tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ - tests/test_gemma_2b_mapper_key.py \ - --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' - # The deselected test monkeypatches flash_attn_varlen_func, which is - # only bound on the module when `flash_attn` is importable. flash_attn - # requires CUDA + dev toolchain, which the CPU-only ubuntu-latest - # runner does not have. The other Bucket-A tests pass cleanly. + tests/test_gemma_2b_mapper_key.py + # test_run_attention_flash_varlen_receives_window_and_softcap was deselected + # until attention_dispatch.py predefined flash_attn_varlen_func as None; it + # monkeypatches that name, so it no longer needs flash_attn on this runner. - name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU) # 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip diff --git a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py index 3b75a13756..a549e58562 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py +++ b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py @@ -96,12 +96,11 @@ def load_and_compute_8bit_ppl( if __name__ == "__main__": mp.set_start_method("spawn", force = True) - if torch.cuda.is_bf16_supported(): - compute_dtype = torch.bfloat16 - attn_implementation = "flash_attention_2" - else: - compute_dtype = torch.float16 - attn_implementation = "sdpa" + from unsloth import is_bfloat16_supported + from unsloth.models._utils import HAS_FLASH_ATTENTION + + compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16 + attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa" model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Llama-3.2-3B-Instruct", diff --git a/tests/saving/language_models/test_merge_model_perplexity_mistral.py b/tests/saving/language_models/test_merge_model_perplexity_mistral.py index 8cc833c2b1..50b0d3caf4 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_mistral.py +++ b/tests/saving/language_models/test_merge_model_perplexity_mistral.py @@ -121,12 +121,11 @@ def load_and_compute_8bit_ppl( if __name__ == "__main__": mp.set_start_method("spawn", force = True) - if torch.cuda.is_bf16_supported(): - compute_dtype = torch.bfloat16 - attn_implementation = "flash_attention_2" - else: - compute_dtype = torch.float16 - attn_implementation = "sdpa" + from unsloth import is_bfloat16_supported + from unsloth.models._utils import HAS_FLASH_ATTENTION + + compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16 + attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa" model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/mistral-7b-v0.3", diff --git a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py index 6f79bfdb71..9c7f6c77af 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py +++ b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py @@ -98,12 +98,11 @@ def load_and_compute_8bit_ppl( if __name__ == "__main__": mp.set_start_method("spawn", force = True) - if torch.cuda.is_bf16_supported(): - compute_dtype = torch.bfloat16 - attn_implementation = "flash_attention_2" - else: - compute_dtype = torch.float16 - attn_implementation = "sdpa" + from unsloth import is_bfloat16_supported + from unsloth.models._utils import HAS_FLASH_ATTENTION + + compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16 + attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa" model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Phi-4", diff --git a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py index c07b37024f..dcbaad13e1 100644 --- a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py +++ b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py @@ -95,12 +95,11 @@ def load_and_compute_8bit_ppl( if __name__ == "__main__": mp.set_start_method("spawn", force = True) - if torch.cuda.is_bf16_supported(): - compute_dtype = torch.bfloat16 - attn_implementation = "flash_attention_2" - else: - compute_dtype = torch.float16 - attn_implementation = "sdpa" + from unsloth import is_bfloat16_supported + from unsloth.models._utils import HAS_FLASH_ATTENTION + + compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16 + attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa" model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Llama-3.1-8B-Instruct", diff --git a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py index cb444d1591..cfa364c697 100644 --- a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py +++ b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py @@ -164,12 +164,11 @@ def load_and_compute_8bit_ppl( if __name__ == "__main__": mp.set_start_method("spawn", force = True) - if torch.cuda.is_bf16_supported(): - compute_dtype = torch.bfloat16 - attn_implementation = "flash_attention_2" - else: - compute_dtype = torch.float16 - attn_implementation = "sdpa" + from unsloth import is_bfloat16_supported + from unsloth.models._utils import HAS_FLASH_ATTENTION + + compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16 + attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa" model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen2.5-7B-Instruct", @@ -210,8 +209,6 @@ if __name__ == "__main__": loftq_config = None, ) - from unsloth import is_bfloat16_supported - trainer = SFTTrainer( model = model, tokenizer = tokenizer, diff --git a/tests/test_fp8_tiny_e8m0.py b/tests/test_fp8_tiny_e8m0.py index cf49c8c92f..df40879d5a 100644 --- a/tests/test_fp8_tiny_e8m0.py +++ b/tests/test_fp8_tiny_e8m0.py @@ -11,7 +11,11 @@ dequant reference. import pytest import torch -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason = "needs CUDA") +cuda_available = torch.cuda.is_available() +xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available() +dev = "cuda" if cuda_available else "xpu" if xpu_available else "cpu" + +pytestmark = pytest.mark.skipif(not (cuda_available or xpu_available), reason = "needs CUDA or XPU") def _reference(X, weight, scale, block): @@ -27,7 +31,6 @@ def test_tiny_non_tileable_forward_backward_matches_reference(): from unsloth.kernels.fp8 import FP8BlockQuantLinear torch.manual_seed(0) - dev = "cuda" block = [128, 128] m, n = 8, 8 # non-tileable, in-dim % 128 != 0 weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) # (out=m, in=n) @@ -50,7 +53,6 @@ def test_e8m0_scale_is_upcast_and_runs(): if not hasattr(torch, "float8_e8m0fnu"): pytest.skip("torch build lacks float8_e8m0fnu") - dev = "cuda" m, n = 8, 8 weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) scale = (torch.rand(1, 1, device = dev) + 1.0).to(torch.float8_e8m0fnu) @@ -70,7 +72,6 @@ def test_rectangular_block_dequant_matches_reference(): from unsloth.kernels.fp8 import _blockwise_weight_dequant_any_shape torch.manual_seed(0) - dev = "cuda" block = [64, 128] m, n = 64, 256 # evenly tiled: 64 % 64 == 0, 256 % 128 == 0 weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) @@ -94,7 +95,6 @@ def test_e8m0_scale_preserves_non_default_block_size_attr(): pytest.skip("torch build lacks float8_e8m0fnu") torch.manual_seed(0) - dev = "cuda" block = [64, 64] # in-dim 96 is not divisible by block[1]=64 -> forward takes the torch dequant # fallback (no fp8 matmul kernel). Scale shape (2, 2) validates for [64, 64] but diff --git a/tests/utils/perplexity_eval.py b/tests/utils/perplexity_eval.py index 5f33a24d53..cdd30e5511 100644 --- a/tests/utils/perplexity_eval.py +++ b/tests/utils/perplexity_eval.py @@ -2,6 +2,9 @@ from tqdm import tqdm import torch import pandas as pd +# DEVICE_TYPE_TORCH, not DEVICE_TYPE: the latter can be "hip"/"mlx", which .to() rejects. +from unsloth.device_type import DEVICE_TYPE_TORCH + model_comparison_results = {} @@ -17,7 +20,7 @@ def ppl_model(model, tokenizer, dataset): for begin_loc in range(0, seq_len, stride): end_loc = min(begin_loc + max_length, seq_len) trg_len = end_loc - prev_end_loc - input_ids = encodings.input_ids[:, begin_loc:end_loc].to("cuda") + input_ids = encodings.input_ids[:, begin_loc:end_loc].to(DEVICE_TYPE_TORCH) target_ids = input_ids.clone() target_ids[:, :-trg_len] = -100 pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 diff --git a/tests/utils/test_batched_leftpad_generation_gpu.py b/tests/utils/test_batched_leftpad_generation_gpu.py index df03125bc2..13db22461e 100644 --- a/tests/utils/test_batched_leftpad_generation_gpu.py +++ b/tests/utils/test_batched_leftpad_generation_gpu.py @@ -4,7 +4,7 @@ Greedy generation in a left-padded batch must match solo batch-size-1 generation for the first PREFIX_TOKENS tokens (the bug makes padded rows diverge into garbage immediately; a full-length match would be flaky due to benign batch-numerics tie-flips deep in the sequence) and must not be -gibberish. Skipped without CUDA. Run: `python -m pytest +gibberish. Skipped without a GPU. Run: `python -m pytest tests/utils/test_batched_leftpad_generation_gpu.py -v`. """ @@ -12,8 +12,19 @@ import pytest import torch cuda_available = torch.cuda.is_available() +xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available() +device = "cuda" if cuda_available else "xpu" if xpu_available else "cpu" -pytestmark = pytest.mark.skipif(not cuda_available, reason = "requires a CUDA GPU") +# Non-strict rather than CUDA-only: keeps the XPU divergence visible, and goes +# green by itself once XPU generation is fixed. +pytestmark = [ + pytest.mark.skipif(not (cuda_available or xpu_available), reason = "requires a CUDA or XPU GPU"), + pytest.mark.xfail( + xpu_available and not cuda_available, + reason = "batched left-padded generation diverges on XPU", + strict = False, + ), +] MODEL_NAME = "unsloth/Qwen2.5-0.5B-Instruct" MAX_NEW_TOKENS = 32 @@ -53,7 +64,7 @@ def _chat(tokenizer, prompt): def _generate(model, tokenizer, texts): inputs = tokenizer(texts, return_tensors = "pt", padding = True, add_special_tokens = False).to( - "cuda" + device ) with torch.inference_mode(): out = model.generate( diff --git a/tests/utils/test_packing.py b/tests/utils/test_packing.py index 1b8bb65058..0be3018cde 100644 --- a/tests/utils/test_packing.py +++ b/tests/utils/test_packing.py @@ -44,6 +44,8 @@ def _build_packed_training_setup(tmp_path, device): dtype = torch.bfloat16 else: dtype = torch.float16 + elif device.type == "xpu": + dtype = torch.bfloat16 try: model, tokenizer = FastLanguageModel.from_pretrained( @@ -76,8 +78,8 @@ def _build_packed_training_setup(tmp_path, device): max_length = 64, logging_steps = 1, max_steps = 1, - fp16 = device.type == "cuda" and not torch.cuda.is_bf16_supported(), - bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported(), + fp16 = dtype == torch.float16, + bf16 = dtype == torch.bfloat16, dataset_num_proc = 1, output_dir = str(tmp_path), packing = True, @@ -974,7 +976,12 @@ def test_enable_sample_packing(): def test_enable_sample_packing_trl_collator(tmp_path): - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + if torch.cuda.is_available(): + device = torch.device("cuda") + elif torch.xpu.is_available(): + device = torch.device("xpu") + else: + device = torch.device("cpu") model, _, trainer, _ = _build_packed_training_setup(tmp_path, device) enable_sample_packing(model, trainer) @@ -1030,7 +1037,12 @@ def test_enable_padding_free_metadata(): def test_packing_sdpa(tmp_path): - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + if torch.cuda.is_available(): + device = torch.device("cuda") + elif torch.xpu.is_available(): + device = torch.device("xpu") + else: + device = torch.device("cpu") model, batch, trainer, llama_mod = _build_packed_training_setup(tmp_path, device) assert "packed_seq_lengths" in batch diff --git a/tests/utils/test_qat.py b/tests/utils/test_qat.py index 79d955164f..0b942d5c32 100644 --- a/tests/utils/test_qat.py +++ b/tests/utils/test_qat.py @@ -130,8 +130,14 @@ def _test_fake_quantizers_are_called( # Weight fake quantizers must always be called. assert child.weight_fake_quantizer.count == 1 + if torch.cuda.is_available(): + device = torch.device("cuda") + elif torch.xpu.is_available(): + device = torch.device("xpu") + else: + pytest.skip("No GPU available") for k, v in example_inputs.items(): - example_inputs[k] = v.cuda() + example_inputs[k] = v.to(device) model.apply(_swap_fake_quantizers) model(**example_inputs) model.apply(_assert_fake_quantizers_are_called) diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index eba89734f7..7fe4e74d5c 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -15,18 +15,20 @@ import pytest import torch -def _has_real_cuda(): - try: - torch.zeros(1).to("cuda") - return True - except Exception: - return False +def _has_real_gpu(): + for backend in ("cuda", "xpu"): + try: + torch.zeros(1).to(backend) + return True + except Exception: + pass + return False -HAS_REAL_CUDA = _has_real_cuda() -requires_cuda = pytest.mark.skipif( - not HAS_REAL_CUDA, - reason = "LlamaRotaryEmbedding builds per-device CUDA caches in __init__", +HAS_REAL_GPU = _has_real_gpu() +requires_gpu = pytest.mark.skipif( + not HAS_REAL_GPU, + reason = "LlamaRotaryEmbedding builds per-device caches in __init__ (needs CUDA or XPU)", ) REPO_ROOT = Path(__file__).resolve().parents[2] @@ -360,7 +362,7 @@ def _cos_at_position(rot, position): # --- Layer 3: CUDA behavioral guard (real instantiation needs a device) --- -@requires_cuda +@requires_gpu def test_constructor_applies_llama3_scaling(): config = _make_config(LLAMA3_ROPE_SCALING) rot = _unsloth_rotary(config) @@ -371,7 +373,7 @@ def test_constructor_applies_llama3_scaling(): ), "LlamaRotaryEmbedding built from a llama3 config produced unscaled inv_freq (issue #2405)." -@requires_cuda +@requires_gpu def test_constructor_unscaled_config_uses_vanilla_inv_freq(): rot = _unsloth_rotary(_make_config(None)) got = rot.inv_freq.float().cpu() @@ -381,7 +383,7 @@ def test_constructor_unscaled_config_uses_vanilla_inv_freq(): ), "LlamaRotaryEmbedding with no rope_scaling must use the vanilla inv_freq" -@requires_cuda +@requires_gpu def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position(): scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING)) unscaled = _unsloth_rotary(_make_config(None)) @@ -397,7 +399,7 @@ def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position(): ) -@requires_cuda +@requires_gpu def test_extended_cache_keeps_scaling_after_growth(): scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING)) # Grow past the initial cache size (mirrors long-context decode). @@ -456,7 +458,7 @@ def _build_longrope_rotary(): return rot, config -@requires_cuda +@requires_gpu @pytest.mark.parametrize( "build", [_build_llama3_rotary, _build_longrope_rotary], ids = ["llama3", "longrope"] ) diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index eda6103d5b..54f8100ca1 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -31,6 +31,8 @@ from ..utils.packing import ( build_xformers_block_causal_mask, ) +flash_attn_func = None +flash_attn_varlen_func = None if HAS_FLASH_ATTENTION: from flash_attn import flash_attn_func, flash_attn_varlen_func HAS_XFORMERS = xformers is not None