[xpu] add Intel XPU support to test infrastructure and runs_on markers (#10835)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
shubham singhal
2026-09-28 09:17:33 +05:30
committed by GitHub
parent 97b32d3133
commit d4823eb10e
21 changed files with 114 additions and 106 deletions

View File

@@ -78,6 +78,8 @@ def _get_visible_devices_env() -> str | None:
return "CUDA_VISIBLE_DEVICES" return "CUDA_VISIBLE_DEVICES"
elif CURRENT_DEVICE == "npu": elif CURRENT_DEVICE == "npu":
return "ASCEND_RT_VISIBLE_DEVICES" return "ASCEND_RT_VISIBLE_DEVICES"
elif CURRENT_DEVICE == "xpu":
return "ZE_AFFINITY_MASK"
else: else:
return None return None
@@ -160,6 +162,8 @@ def _manage_distributed_env(request: FixtureRequest, monkeypatch: MonkeyPatch) -
monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
elif CURRENT_DEVICE == "npu": elif CURRENT_DEVICE == "npu":
monkeypatch.setattr(torch.npu, "device_count", lambda: 1) monkeypatch.setattr(torch.npu, "device_count", lambda: 1)
elif CURRENT_DEVICE == "xpu":
monkeypatch.setattr(torch.xpu, "device_count", lambda: 1)
@pytest.fixture @pytest.fixture

View File

@@ -43,7 +43,7 @@ TRAIN_ARGS = {
} }
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("num_samples", [16]) @pytest.mark.parametrize("num_samples", [16])
def test_feedback_data(num_samples: int): def test_feedback_data(num_samples: int):
train_dataset = load_dataset_module(**TRAIN_ARGS)["train_dataset"] train_dataset = load_dataset_module(**TRAIN_ARGS)["train_dataset"]

View File

@@ -52,7 +52,7 @@ def _convert_sharegpt_to_openai(messages: list[dict[str, str]]) -> list[dict[str
return new_messages return new_messages
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("num_samples", [16]) @pytest.mark.parametrize("num_samples", [16])
def test_pairwise_data(num_samples: int): def test_pairwise_data(num_samples: int):
train_dataset = load_dataset_module(**TRAIN_ARGS)["train_dataset"] train_dataset = load_dataset_module(**TRAIN_ARGS)["train_dataset"]

View File

@@ -18,7 +18,7 @@ import pytest
from llamafactory.data.processor.processor_utils import infer_seqlen from llamafactory.data.processor.processor_utils import infer_seqlen
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"test_input,test_output", "test_input,test_output",
[ [

View File

@@ -43,7 +43,7 @@ TRAIN_ARGS = {
} }
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("num_samples", [16]) @pytest.mark.parametrize("num_samples", [16])
def test_supervised_single_turn(num_samples: int): def test_supervised_single_turn(num_samples: int):
train_dataset = load_dataset_module(dataset_dir="ONLINE", dataset=TINY_DATA, **TRAIN_ARGS)["train_dataset"] train_dataset = load_dataset_module(dataset_dir="ONLINE", dataset=TINY_DATA, **TRAIN_ARGS)["train_dataset"]
@@ -72,7 +72,7 @@ def test_supervised_single_turn(num_samples: int):
assert train_dataset["labels"][index] == ref_label_ids assert train_dataset["labels"][index] == ref_label_ids
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("num_samples", [8]) @pytest.mark.parametrize("num_samples", [8])
def test_supervised_multi_turn(num_samples: int): def test_supervised_multi_turn(num_samples: int):
train_dataset = load_dataset_module(dataset_dir="REMOTE:" + DEMO_DATA, dataset="system_chat", **TRAIN_ARGS)[ train_dataset = load_dataset_module(dataset_dir="REMOTE:" + DEMO_DATA, dataset="system_chat", **TRAIN_ARGS)[
@@ -90,7 +90,7 @@ def test_supervised_multi_turn(num_samples: int):
assert train_dataset["input_ids"][index] == ref_input_ids assert train_dataset["input_ids"][index] == ref_input_ids
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("num_samples", [4]) @pytest.mark.parametrize("num_samples", [4])
def test_supervised_train_on_prompt(num_samples: int): def test_supervised_train_on_prompt(num_samples: int):
train_dataset = load_dataset_module( train_dataset = load_dataset_module(
@@ -108,7 +108,7 @@ def test_supervised_train_on_prompt(num_samples: int):
assert train_dataset["labels"][index] == ref_input_ids assert train_dataset["labels"][index] == ref_input_ids
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("num_samples", [4]) @pytest.mark.parametrize("num_samples", [4])
def test_supervised_mask_history(num_samples: int): def test_supervised_mask_history(num_samples: int):
train_dataset = load_dataset_module( train_dataset = load_dataset_module(

View File

@@ -47,7 +47,7 @@ TRAIN_ARGS = {
} }
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("num_samples", [16]) @pytest.mark.parametrize("num_samples", [16])
def test_unsupervised_data(num_samples: int): def test_unsupervised_data(num_samples: int):
train_dataset = load_dataset_module(**TRAIN_ARGS)["train_dataset"] train_dataset = load_dataset_module(**TRAIN_ARGS)["train_dataset"]

View File

@@ -32,7 +32,7 @@ from llamafactory.model import load_tokenizer
TINY_LLAMA3 = os.getenv("TINY_LLAMA3", "llamafactory/tiny-random-Llama-3") TINY_LLAMA3 = os.getenv("TINY_LLAMA3", "llamafactory/tiny-random-Llama-3")
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_base_collator(): def test_base_collator():
model_args, data_args, *_ = get_infer_args({"model_name_or_path": TINY_LLAMA3, "template": "default"}) model_args, data_args, *_ = get_infer_args({"model_name_or_path": TINY_LLAMA3, "template": "default"})
tokenizer_module = load_tokenizer(model_args) tokenizer_module = load_tokenizer(model_args)
@@ -76,7 +76,7 @@ def test_base_collator():
assert batch_input[k].eq(torch.tensor(expected_input[k])).all() assert batch_input[k].eq(torch.tensor(expected_input[k])).all()
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_multimodal_collator(): def test_multimodal_collator():
model_args, data_args, *_ = get_infer_args( model_args, data_args, *_ = get_infer_args(
{"model_name_or_path": "Qwen/Qwen2-VL-2B-Instruct", "template": "qwen2_vl"} {"model_name_or_path": "Qwen/Qwen2-VL-2B-Instruct", "template": "qwen2_vl"}
@@ -270,7 +270,7 @@ def _get_expected_position_ids(
return torch.cat(all_position_ids, dim=-1) return torch.cat(all_position_ids, dim=-1)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_multimodal_collator_with_packing(): def test_multimodal_collator_with_packing():
model_args, data_args, *_ = get_infer_args( model_args, data_args, *_ = get_infer_args(
{"model_name_or_path": "Qwen/Qwen2-VL-2B-Instruct", "template": "qwen2_vl"} {"model_name_or_path": "Qwen/Qwen2-VL-2B-Instruct", "template": "qwen2_vl"}
@@ -322,7 +322,7 @@ def test_multimodal_collator_with_packing():
assert batch_input["position_ids"][1:, :, :valid_len].eq(expected_position_ids).all() assert batch_input["position_ids"][1:, :, :valid_len].eq(expected_position_ids).all()
@pytest.mark.runs_on(["cpu"]) @pytest.mark.runs_on(["cpu", "xpu"])
def test_4d_attention_mask(): def test_4d_attention_mask():
o = 0.0 o = 0.0
x = torch.finfo(torch.float16).min x = torch.finfo(torch.float16).min

View File

@@ -20,7 +20,7 @@ from llamafactory.data.parser import DatasetAttr
from llamafactory.hparams import DataArguments from llamafactory.hparams import DataArguments
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_alpaca_converter(): def test_alpaca_converter():
dataset_attr = DatasetAttr("hf_hub", "llamafactory/tiny-supervised-dataset") dataset_attr = DatasetAttr("hf_hub", "llamafactory/tiny-supervised-dataset")
data_args = DataArguments() data_args = DataArguments()
@@ -41,7 +41,7 @@ def test_alpaca_converter():
} }
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_sharegpt_converter(): def test_sharegpt_converter():
dataset_attr = DatasetAttr("hf_hub", "llamafactory/tiny-supervised-dataset") dataset_attr = DatasetAttr("hf_hub", "llamafactory/tiny-supervised-dataset")
data_args = DataArguments() data_args = DataArguments()

View File

@@ -38,19 +38,19 @@ TOOLS = [
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_empty_formatter(): def test_empty_formatter():
formatter = EmptyFormatter(slots=["\n"]) formatter = EmptyFormatter(slots=["\n"])
assert formatter.apply() == ["\n"] assert formatter.apply() == ["\n"]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_string_formatter(): def test_string_formatter():
formatter = StringFormatter(slots=["<s>", "Human: {{content}}\nAssistant:"]) formatter = StringFormatter(slots=["<s>", "Human: {{content}}\nAssistant:"])
assert formatter.apply(content="Hi") == ["<s>", "Human: Hi\nAssistant:"] assert formatter.apply(content="Hi") == ["<s>", "Human: Hi\nAssistant:"]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_function_formatter(): def test_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}", "</s>"], tool_format="default") formatter = FunctionFormatter(slots=["{{content}}", "</s>"], tool_format="default")
tool_calls = json.dumps(FUNCTION) tool_calls = json.dumps(FUNCTION)
@@ -60,7 +60,7 @@ def test_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_multi_function_formatter(): def test_multi_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}", "</s>"], tool_format="default") formatter = FunctionFormatter(slots=["{{content}}", "</s>"], tool_format="default")
tool_calls = json.dumps([FUNCTION] * 2) tool_calls = json.dumps([FUNCTION] * 2)
@@ -71,7 +71,7 @@ def test_multi_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_default_tool_formatter(): def test_default_tool_formatter():
formatter = ToolFormatter(tool_format="default") formatter = ToolFormatter(tool_format="default")
assert formatter.apply(content=json.dumps(TOOLS)) == [ assert formatter.apply(content=json.dumps(TOOLS)) == [
@@ -90,14 +90,14 @@ def test_default_tool_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_default_tool_extractor(): def test_default_tool_extractor():
formatter = ToolFormatter(tool_format="default") formatter = ToolFormatter(tool_format="default")
result = """Action: test_tool\nAction Input: {"foo": "bar", "size": 10}""" result = """Action: test_tool\nAction Input: {"foo": "bar", "size": 10}"""
assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")] assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_default_multi_tool_extractor(): def test_default_multi_tool_extractor():
formatter = ToolFormatter(tool_format="default") formatter = ToolFormatter(tool_format="default")
result = ( result = (
@@ -110,14 +110,14 @@ def test_default_multi_tool_extractor():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_glm4_function_formatter(): def test_glm4_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}"], tool_format="glm4") formatter = FunctionFormatter(slots=["{{content}}"], tool_format="glm4")
tool_calls = json.dumps(FUNCTION) tool_calls = json.dumps(FUNCTION)
assert formatter.apply(content=tool_calls) == ["""tool_name\n{"foo": "bar", "size": 10}"""] assert formatter.apply(content=tool_calls) == ["""tool_name\n{"foo": "bar", "size": 10}"""]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_glm4_tool_formatter(): def test_glm4_tool_formatter():
formatter = ToolFormatter(tool_format="glm4") formatter = ToolFormatter(tool_format="glm4")
assert formatter.apply(content=json.dumps(TOOLS)) == [ assert formatter.apply(content=json.dumps(TOOLS)) == [
@@ -128,14 +128,14 @@ def test_glm4_tool_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_glm4_tool_extractor(): def test_glm4_tool_extractor():
formatter = ToolFormatter(tool_format="glm4") formatter = ToolFormatter(tool_format="glm4")
result = """test_tool\n{"foo": "bar", "size": 10}\n""" result = """test_tool\n{"foo": "bar", "size": 10}\n"""
assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")] assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llama3_function_formatter(): def test_llama3_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|eot_id|>"], tool_format="llama3") formatter = FunctionFormatter(slots=["{{content}}<|eot_id|>"], tool_format="llama3")
tool_calls = json.dumps(FUNCTION) tool_calls = json.dumps(FUNCTION)
@@ -144,7 +144,7 @@ def test_llama3_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llama3_multi_function_formatter(): def test_llama3_multi_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|eot_id|>"], tool_format="llama3") formatter = FunctionFormatter(slots=["{{content}}<|eot_id|>"], tool_format="llama3")
tool_calls = json.dumps([FUNCTION] * 2) tool_calls = json.dumps([FUNCTION] * 2)
@@ -155,7 +155,7 @@ def test_llama3_multi_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llama3_tool_formatter(): def test_llama3_tool_formatter():
formatter = ToolFormatter(tool_format="llama3") formatter = ToolFormatter(tool_format="llama3")
date = datetime.now().strftime("%d %b %Y") date = datetime.now().strftime("%d %b %Y")
@@ -169,14 +169,14 @@ def test_llama3_tool_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llama3_tool_extractor(): def test_llama3_tool_extractor():
formatter = ToolFormatter(tool_format="llama3") formatter = ToolFormatter(tool_format="llama3")
result = """{"name": "test_tool", "parameters": {"foo": "bar", "size": 10}}\n""" result = """{"name": "test_tool", "parameters": {"foo": "bar", "size": 10}}\n"""
assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")] assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llama3_multi_tool_extractor(): def test_llama3_multi_tool_extractor():
formatter = ToolFormatter(tool_format="llama3") formatter = ToolFormatter(tool_format="llama3")
result = ( result = (
@@ -189,7 +189,7 @@ def test_llama3_multi_tool_extractor():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_mistral_function_formatter(): def test_mistral_function_formatter():
formatter = FunctionFormatter(slots=["[TOOL_CALLS] {{content}}", "</s>"], tool_format="mistral") formatter = FunctionFormatter(slots=["[TOOL_CALLS] {{content}}", "</s>"], tool_format="mistral")
tool_calls = json.dumps(FUNCTION) tool_calls = json.dumps(FUNCTION)
@@ -199,7 +199,7 @@ def test_mistral_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_mistral_multi_function_formatter(): def test_mistral_multi_function_formatter():
formatter = FunctionFormatter(slots=["[TOOL_CALLS] {{content}}", "</s>"], tool_format="mistral") formatter = FunctionFormatter(slots=["[TOOL_CALLS] {{content}}", "</s>"], tool_format="mistral")
tool_calls = json.dumps([FUNCTION] * 2) tool_calls = json.dumps([FUNCTION] * 2)
@@ -211,7 +211,7 @@ def test_mistral_multi_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_mistral_tool_formatter(): def test_mistral_tool_formatter():
formatter = ToolFormatter(tool_format="mistral") formatter = ToolFormatter(tool_format="mistral")
wrapped_tool = {"type": "function", "function": TOOLS[0]} wrapped_tool = {"type": "function", "function": TOOLS[0]}
@@ -220,14 +220,14 @@ def test_mistral_tool_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_mistral_tool_extractor(): def test_mistral_tool_extractor():
formatter = ToolFormatter(tool_format="mistral") formatter = ToolFormatter(tool_format="mistral")
result = """{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}""" result = """{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}"""
assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")] assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_mistral_multi_tool_extractor(): def test_mistral_multi_tool_extractor():
formatter = ToolFormatter(tool_format="mistral") formatter = ToolFormatter(tool_format="mistral")
result = ( result = (
@@ -240,7 +240,7 @@ def test_mistral_multi_tool_extractor():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen_function_formatter(): def test_qwen_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="qwen") formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="qwen")
tool_calls = json.dumps(FUNCTION) tool_calls = json.dumps(FUNCTION)
@@ -249,7 +249,7 @@ def test_qwen_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen_multi_function_formatter(): def test_qwen_multi_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="qwen") formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="qwen")
tool_calls = json.dumps([FUNCTION] * 2) tool_calls = json.dumps([FUNCTION] * 2)
@@ -260,7 +260,7 @@ def test_qwen_multi_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen_tool_formatter(): def test_qwen_tool_formatter():
formatter = ToolFormatter(tool_format="qwen") formatter = ToolFormatter(tool_format="qwen")
wrapped_tool = {"type": "function", "function": TOOLS[0]} wrapped_tool = {"type": "function", "function": TOOLS[0]}
@@ -274,14 +274,14 @@ def test_qwen_tool_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen_tool_extractor(): def test_qwen_tool_extractor():
formatter = ToolFormatter(tool_format="qwen") formatter = ToolFormatter(tool_format="qwen")
result = """<tool_call>\n{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}\n</tool_call>""" result = """<tool_call>\n{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}\n</tool_call>"""
assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")] assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen38_tool_formatter(): def test_qwen38_tool_formatter():
formatter = ToolFormatter(tool_format="qwen3_8") formatter = ToolFormatter(tool_format="qwen3_8")
wrapped_tool = {"type": "function", "function": TOOLS[0]} wrapped_tool = {"type": "function", "function": TOOLS[0]}
@@ -289,7 +289,7 @@ def test_qwen38_tool_formatter():
assert json.dumps(wrapped_tool, ensure_ascii=False) in output assert json.dumps(wrapped_tool, ensure_ascii=False) in output
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen_multi_tool_extractor(): def test_qwen_multi_tool_extractor():
formatter = ToolFormatter(tool_format="qwen") formatter = ToolFormatter(tool_format="qwen")
result = ( result = (
@@ -302,7 +302,7 @@ def test_qwen_multi_tool_extractor():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_function_formatter(): def test_lfm2_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="lfm2") formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="lfm2")
tool_calls = json.dumps(FUNCTION) tool_calls = json.dumps(FUNCTION)
@@ -311,7 +311,7 @@ def test_lfm2_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_multi_function_formatter(): def test_lfm2_multi_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="lfm2") formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="lfm2")
tool_calls = json.dumps([FUNCTION] * 2) tool_calls = json.dumps([FUNCTION] * 2)
@@ -321,7 +321,7 @@ def test_lfm2_multi_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_tool_formatter(): def test_lfm2_tool_formatter():
formatter = ToolFormatter(tool_format="lfm2") formatter = ToolFormatter(tool_format="lfm2")
assert formatter.apply(content=json.dumps(TOOLS)) == [ assert formatter.apply(content=json.dumps(TOOLS)) == [
@@ -329,14 +329,14 @@ def test_lfm2_tool_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_tool_extractor(): def test_lfm2_tool_extractor():
formatter = ToolFormatter(tool_format="lfm2") formatter = ToolFormatter(tool_format="lfm2")
result = """<|tool_call_start|>[test_tool(foo="bar", size=10)]<|tool_call_end|>""" result = """<|tool_call_start|>[test_tool(foo="bar", size=10)]<|tool_call_end|>"""
assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")] assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_multi_tool_extractor(): def test_lfm2_multi_tool_extractor():
formatter = ToolFormatter(tool_format="lfm2") formatter = ToolFormatter(tool_format="lfm2")
result = """<|tool_call_start|>[test_tool(foo="bar", size=10), another_tool(foo="job", size=2)]<|tool_call_end|>""" result = """<|tool_call_start|>[test_tool(foo="bar", size=10), another_tool(foo="job", size=2)]<|tool_call_end|>"""
@@ -346,7 +346,7 @@ def test_lfm2_multi_tool_extractor():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_tool_extractor_with_nested_dict(): def test_lfm2_tool_extractor_with_nested_dict():
formatter = ToolFormatter(tool_format="lfm2") formatter = ToolFormatter(tool_format="lfm2")
result = """<|tool_call_start|>[search(query="test", options={"limit": 10, "offset": 0})]<|tool_call_end|>""" result = """<|tool_call_start|>[search(query="test", options={"limit": 10, "offset": 0})]<|tool_call_end|>"""
@@ -358,7 +358,7 @@ def test_lfm2_tool_extractor_with_nested_dict():
assert args["options"] == {"limit": 10, "offset": 0} assert args["options"] == {"limit": 10, "offset": 0}
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_tool_extractor_with_list_arg(): def test_lfm2_tool_extractor_with_list_arg():
formatter = ToolFormatter(tool_format="lfm2") formatter = ToolFormatter(tool_format="lfm2")
result = """<|tool_call_start|>[batch_process(items=[1, 2, 3], enabled=True)]<|tool_call_end|>""" result = """<|tool_call_start|>[batch_process(items=[1, 2, 3], enabled=True)]<|tool_call_end|>"""
@@ -370,7 +370,7 @@ def test_lfm2_tool_extractor_with_list_arg():
assert args["enabled"] is True assert args["enabled"] is True
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_tool_extractor_no_match(): def test_lfm2_tool_extractor_no_match():
formatter = ToolFormatter(tool_format="lfm2") formatter = ToolFormatter(tool_format="lfm2")
result = "This is a regular response without tool calls." result = "This is a regular response without tool calls."
@@ -378,7 +378,7 @@ def test_lfm2_tool_extractor_no_match():
assert extracted == result assert extracted == result
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_tool_round_trip(): def test_lfm2_tool_round_trip():
formatter = FunctionFormatter(slots=["{{content}}"], tool_format="lfm2") formatter = FunctionFormatter(slots=["{{content}}"], tool_format="lfm2")
tool_formatter = ToolFormatter(tool_format="lfm2") tool_formatter = ToolFormatter(tool_format="lfm2")
@@ -390,7 +390,7 @@ def test_lfm2_tool_round_trip():
assert json.loads(extracted[0][1]) == original["arguments"] assert json.loads(extracted[0][1]) == original["arguments"]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_minicpm5_function_formatter(): def test_minicpm5_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="minicpm5") formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="minicpm5")
tool_calls = json.dumps(FUNCTION) tool_calls = json.dumps(FUNCTION)
@@ -399,7 +399,7 @@ def test_minicpm5_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_minicpm5_multi_function_formatter(): def test_minicpm5_multi_function_formatter():
formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="minicpm5") formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="minicpm5")
tool_calls = json.dumps([FUNCTION] * 2) tool_calls = json.dumps([FUNCTION] * 2)
@@ -411,7 +411,7 @@ def test_minicpm5_multi_function_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_minicpm5_tool_formatter(): def test_minicpm5_tool_formatter():
formatter = ToolFormatter(tool_format="minicpm5") formatter = ToolFormatter(tool_format="minicpm5")
wrapped = json.dumps({"type": "function", "function": TOOLS[0]}, ensure_ascii=False) wrapped = json.dumps({"type": "function", "function": TOOLS[0]}, ensure_ascii=False)
@@ -427,28 +427,28 @@ def test_minicpm5_tool_formatter():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_minicpm5_tool_extractor(): def test_minicpm5_tool_extractor():
formatter = ToolFormatter(tool_format="minicpm5") formatter = ToolFormatter(tool_format="minicpm5")
result = '<function name="test_tool"><param name="foo">bar</param><param name="size">10</param></function>' result = '<function name="test_tool"><param name="foo">bar</param><param name="size">10</param></function>'
assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")] assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_minicpm5_tool_extractor_cdata(): def test_minicpm5_tool_extractor_cdata():
formatter = ToolFormatter(tool_format="minicpm5") formatter = ToolFormatter(tool_format="minicpm5")
result = '<function name="test_tool"><param name="foo"><![CDATA[a < b\nsecond line]]></param></function>' result = '<function name="test_tool"><param name="foo"><![CDATA[a < b\nsecond line]]></param></function>'
assert formatter.extract(result) == [("test_tool", json.dumps({"foo": "a < b\nsecond line"}))] assert formatter.extract(result) == [("test_tool", json.dumps({"foo": "a < b\nsecond line"}))]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_minicpm5_tool_extractor_malformed_value(): def test_minicpm5_tool_extractor_malformed_value():
formatter = ToolFormatter(tool_format="minicpm5") formatter = ToolFormatter(tool_format="minicpm5")
result = '<function name="test_tool"><param name="foo">{[1, 2], [3, 4]}</param></function>' result = '<function name="test_tool"><param name="foo">{[1, 2], [3, 4]}</param></function>'
assert formatter.extract(result) == [("test_tool", json.dumps({"foo": "{[1, 2], [3, 4]}"}))] assert formatter.extract(result) == [("test_tool", json.dumps({"foo": "{[1, 2], [3, 4]}"}))]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"arguments", "arguments",
[ [

View File

@@ -40,21 +40,21 @@ TRAIN_ARGS = {
} }
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_load_train_only(): def test_load_train_only():
dataset_module = load_dataset_module(**TRAIN_ARGS) dataset_module = load_dataset_module(**TRAIN_ARGS)
assert dataset_module.get("train_dataset") is not None assert dataset_module.get("train_dataset") is not None
assert dataset_module.get("eval_dataset") is None assert dataset_module.get("eval_dataset") is None
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_load_val_size(): def test_load_val_size():
dataset_module = load_dataset_module(val_size=0.1, **TRAIN_ARGS) dataset_module = load_dataset_module(val_size=0.1, **TRAIN_ARGS)
assert dataset_module.get("train_dataset") is not None assert dataset_module.get("train_dataset") is not None
assert dataset_module.get("eval_dataset") is not None assert dataset_module.get("eval_dataset") is not None
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_load_eval_data(): def test_load_eval_data():
dataset_module = load_dataset_module(eval_dataset=TINY_DATA, **TRAIN_ARGS) dataset_module = load_dataset_module(eval_dataset=TINY_DATA, **TRAIN_ARGS)
assert dataset_module.get("train_dataset") is not None assert dataset_module.get("train_dataset") is not None

View File

@@ -180,7 +180,7 @@ def _check_plugin(
) )
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_base_plugin(): def test_base_plugin():
tokenizer_module = _load_tokenizer_module(model_name_or_path=TINY_LLAMA3) tokenizer_module = _load_tokenizer_module(model_name_or_path=TINY_LLAMA3)
base_plugin = get_mm_plugin(name="base") base_plugin = get_mm_plugin(name="base")
@@ -188,7 +188,7 @@ def test_base_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not HF_TOKEN, reason="Gated model.") @pytest.mark.skipif(not HF_TOKEN, reason="Gated model.")
@pytest.mark.skipif(not is_transformers_version_greater_than("4.50.0"), reason="Requires transformers>=4.50.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.50.0"), reason="Requires transformers>=4.50.0")
def test_gemma3_plugin(): def test_gemma3_plugin():
@@ -211,7 +211,7 @@ def test_gemma3_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("5.6.0"), reason="Requires transformers>=5.6.0") @pytest.mark.skipif(not is_transformers_version_greater_than("5.6.0"), reason="Requires transformers>=5.6.0")
def test_gemma4_plugin(): def test_gemma4_plugin():
tokenizer_module = _load_tokenizer_module(model_name_or_path="google/gemma-4-31B-it") tokenizer_module = _load_tokenizer_module(model_name_or_path="google/gemma-4-31B-it")
@@ -244,7 +244,7 @@ def test_gemma4_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.52.0"), reason="Requires transformers>=4.52.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.52.0"), reason="Requires transformers>=4.52.0")
def test_internvl_plugin(): def test_internvl_plugin():
image_seqlen = 256 image_seqlen = 256
@@ -263,7 +263,7 @@ def test_internvl_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.51.0"), reason="Requires transformers>=4.51.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.51.0"), reason="Requires transformers>=4.51.0")
def test_llama4_plugin(): def test_llama4_plugin():
tokenizer_module = _load_tokenizer_module(model_name_or_path=TINY_LLAMA4) tokenizer_module = _load_tokenizer_module(model_name_or_path=TINY_LLAMA4)
@@ -285,7 +285,7 @@ def test_llama4_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llava_plugin(): def test_llava_plugin():
image_seqlen = 576 image_seqlen = 576
tokenizer_module = _load_tokenizer_module(model_name_or_path="llava-hf/llava-1.5-7b-hf") tokenizer_module = _load_tokenizer_module(model_name_or_path="llava-hf/llava-1.5-7b-hf")
@@ -299,7 +299,7 @@ def test_llava_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llava_next_plugin(): def test_llava_next_plugin():
image_seqlen = 1176 image_seqlen = 1176
tokenizer_module = _load_tokenizer_module(model_name_or_path="llava-hf/llava-v1.6-vicuna-7b-hf") tokenizer_module = _load_tokenizer_module(model_name_or_path="llava-hf/llava-v1.6-vicuna-7b-hf")
@@ -313,7 +313,7 @@ def test_llava_next_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llava_next_video_plugin(): def test_llava_next_video_plugin():
image_seqlen = 1176 image_seqlen = 1176
tokenizer_module = _load_tokenizer_module(model_name_or_path="llava-hf/LLaVA-NeXT-Video-7B-hf") tokenizer_module = _load_tokenizer_module(model_name_or_path="llava-hf/LLaVA-NeXT-Video-7B-hf")
@@ -327,7 +327,7 @@ def test_llava_next_video_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not HF_TOKEN, reason="Gated model.") @pytest.mark.skipif(not HF_TOKEN, reason="Gated model.")
def test_paligemma_plugin(): def test_paligemma_plugin():
image_seqlen = 256 image_seqlen = 256
@@ -347,7 +347,7 @@ def test_paligemma_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.50.0"), reason="Requires transformers>=4.50.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.50.0"), reason="Requires transformers>=4.50.0")
def test_pixtral_plugin(): def test_pixtral_plugin():
image_slice_height, image_slice_width = 2, 2 image_slice_height, image_slice_width = 2, 2
@@ -370,7 +370,7 @@ def test_pixtral_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.52.0"), reason="Requires transformers>=4.52.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.52.0"), reason="Requires transformers>=4.52.0")
def test_qwen2_omni_plugin(): def test_qwen2_omni_plugin():
image_seqlen, audio_seqlen = 4, 2 image_seqlen, audio_seqlen = 4, 2
@@ -401,7 +401,7 @@ def test_qwen2_omni_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen2_vl_plugin(): def test_qwen2_vl_plugin():
image_seqlen = 4 image_seqlen = 4
tokenizer_module = _load_tokenizer_module(model_name_or_path="Qwen/Qwen2-VL-7B-Instruct") tokenizer_module = _load_tokenizer_module(model_name_or_path="Qwen/Qwen2-VL-7B-Instruct")
@@ -436,7 +436,7 @@ def test_moss_vl_plugin():
assert messages[0]["content"] == "First <image>, finally <image>." assert messages[0]["content"] == "First <image>, finally <image>."
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.57.0"), reason="Requires transformers>=4.57.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.57.0"), reason="Requires transformers>=4.57.0")
def test_qwen3_vl_plugin(): def test_qwen3_vl_plugin():
frame_seqlen = 1 frame_seqlen = 1
@@ -470,7 +470,7 @@ def test_qwen3_vl_plugin():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.57.0"), reason="Requires transformers>=4.57.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.57.0"), reason="Requires transformers>=4.57.0")
@pytest.mark.skipif(not is_pyav_available(), reason="Requires pyav") @pytest.mark.skipif(not is_pyav_available(), reason="Requires pyav")
def test_qwen3_vl_plugin_video_path(): def test_qwen3_vl_plugin_video_path():
@@ -504,7 +504,7 @@ def test_qwen3_vl_plugin_video_path():
) )
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.47.0"), reason="Requires transformers>=4.47.0") @pytest.mark.skipif(not is_transformers_version_greater_than("4.47.0"), reason="Requires transformers>=4.47.0")
def test_video_llava_plugin(): def test_video_llava_plugin():
image_seqlen = 256 image_seqlen = 256
@@ -519,7 +519,7 @@ def test_video_llava_plugin():
_check_plugin(**check_inputs) _check_plugin(**check_inputs)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_lfm2_vl_plugin(): def test_lfm2_vl_plugin():
"""Test LFM2.5-VL plugin instantiation.""" """Test LFM2.5-VL plugin instantiation."""
# Test plugin can be instantiated with correct tokens # Test plugin can be instantiated with correct tokens

View File

@@ -139,7 +139,7 @@ def test_qwen38_registration():
assert {"qwen3_5_moe_text", "qwen3_5_text"} <= MCA_SUPPORTED_MODELS assert {"qwen3_5_moe_text", "qwen3_5_text"} <= MCA_SUPPORTED_MODELS
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_encode_oneturn(): def test_encode_oneturn():
tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3) tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)
template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3")) template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3"))
@@ -154,7 +154,7 @@ def test_encode_oneturn():
_check_tokenization(tokenizer, (prompt_ids, answer_ids), (prompt_str, answer_str)) _check_tokenization(tokenizer, (prompt_ids, answer_ids), (prompt_str, answer_str))
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_encode_multiturn(): def test_encode_multiturn():
tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3) tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)
template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3")) template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3"))
@@ -175,7 +175,7 @@ def test_encode_multiturn():
) )
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("cot_messages", [True, False]) @pytest.mark.parametrize("cot_messages", [True, False])
@pytest.mark.parametrize("enable_thinking", [True, False, None]) @pytest.mark.parametrize("enable_thinking", [True, False, None])
def test_reasoning_encode_oneturn(cot_messages: bool, enable_thinking: bool): def test_reasoning_encode_oneturn(cot_messages: bool, enable_thinking: bool):
@@ -201,7 +201,7 @@ def test_reasoning_encode_oneturn(cot_messages: bool, enable_thinking: bool):
_check_tokenization(tokenizer, (prompt_ids, answer_ids), (prompt_str, answer_str)) _check_tokenization(tokenizer, (prompt_ids, answer_ids), (prompt_str, answer_str))
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("cot_messages", [True, False]) @pytest.mark.parametrize("cot_messages", [True, False])
@pytest.mark.parametrize("enable_thinking", [True, False, None]) @pytest.mark.parametrize("enable_thinking", [True, False, None])
def test_reasoning_encode_multiturn(cot_messages: bool, enable_thinking: bool): def test_reasoning_encode_multiturn(cot_messages: bool, enable_thinking: bool):
@@ -230,7 +230,7 @@ def test_reasoning_encode_multiturn(cot_messages: bool, enable_thinking: bool):
) )
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("enable_thinking", [True, False, None]) @pytest.mark.parametrize("enable_thinking", [True, False, None])
@pytest.mark.parametrize("discarding_history_cot", [True, False]) @pytest.mark.parametrize("discarding_history_cot", [True, False])
def test_reasoning_encode_multiturn_discarding_history_cot(enable_thinking: bool, discarding_history_cot: bool): def test_reasoning_encode_multiturn_discarding_history_cot(enable_thinking: bool, discarding_history_cot: bool):
@@ -266,7 +266,7 @@ def test_reasoning_encode_multiturn_discarding_history_cot(enable_thinking: bool
) )
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen38_reasoning_instruction(): def test_qwen38_reasoning_instruction():
tokenizer = CharTokenizer() tokenizer = CharTokenizer()
template = deepcopy(TEMPLATES["qwen3_8"]) template = deepcopy(TEMPLATES["qwen3_8"])
@@ -284,7 +284,7 @@ def test_qwen38_reasoning_instruction():
assert answer_text == "<think>\n\n</think>\n\nI am fine!<|im_end|>\n" assert answer_text == "<think>\n\n</think>\n\nI am fine!<|im_end|>\n"
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_qwen38_tool_system_order(): def test_qwen38_tool_system_order():
tokenizer = CharTokenizer() tokenizer = CharTokenizer()
template = deepcopy(TEMPLATES["qwen3_8"]) template = deepcopy(TEMPLATES["qwen3_8"])
@@ -308,7 +308,7 @@ def test_qwen38_tool_system_order():
) )
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_jinja_template(): def test_jinja_template():
tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3) tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)
ref_tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3) ref_tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)
@@ -318,7 +318,7 @@ def test_jinja_template():
assert tokenizer.apply_chat_template(MESSAGES) == ref_tokenizer.apply_chat_template(MESSAGES) assert tokenizer.apply_chat_template(MESSAGES) == ref_tokenizer.apply_chat_template(MESSAGES)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_ollama_modelfile(): def test_ollama_modelfile():
tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3) tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)
template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3")) template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3"))
@@ -336,14 +336,14 @@ def test_ollama_modelfile():
) )
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_get_stop_token_ids(): def test_get_stop_token_ids():
tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3) tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)
template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3")) template = get_template_and_fix_tokenizer(tokenizer, DataArguments(template="llama3"))
assert set(template.get_stop_token_ids(tokenizer)) == {128008, 128009} assert set(template.get_stop_token_ids(tokenizer)) == {128008, 128009}
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not HF_TOKEN, reason="Gated model.") @pytest.mark.skipif(not HF_TOKEN, reason="Gated model.")
def test_gemma_template(): def test_gemma_template():
prompt_str = ( prompt_str = (
@@ -356,7 +356,7 @@ def test_gemma_template():
_check_template("google/gemma-3-4b-it", "gemma", prompt_str, answer_str) _check_template("google/gemma-3-4b-it", "gemma", prompt_str, answer_str)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not HF_TOKEN, reason="Gated model.") @pytest.mark.skipif(not HF_TOKEN, reason="Gated model.")
def test_gemma2_template(): def test_gemma2_template():
prompt_str = ( prompt_str = (
@@ -369,7 +369,7 @@ def test_gemma2_template():
_check_template("google/gemma-2-2b-it", "gemma2", prompt_str, answer_str) _check_template("google/gemma-2-2b-it", "gemma2", prompt_str, answer_str)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.skipif(not HF_TOKEN, reason="Gated model.") @pytest.mark.skipif(not HF_TOKEN, reason="Gated model.")
def test_llama3_template(): def test_llama3_template():
prompt_str = ( prompt_str = (
@@ -382,7 +382,7 @@ def test_llama3_template():
_check_template("meta-llama/Meta-Llama-3-8B-Instruct", "llama3", prompt_str, answer_str) _check_template("meta-llama/Meta-Llama-3-8B-Instruct", "llama3", prompt_str, answer_str)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_llama4_template(): def test_llama4_template():
prompt_str = ( prompt_str = (
f"<|begin_of_text|><|header_start|>user<|header_end|>\n\n{MESSAGES[0]['content']}<|eot|>" f"<|begin_of_text|><|header_start|>user<|header_end|>\n\n{MESSAGES[0]['content']}<|eot|>"
@@ -394,7 +394,7 @@ def test_llama4_template():
_check_template(TINY_LLAMA4, "llama4", prompt_str, answer_str) _check_template(TINY_LLAMA4, "llama4", prompt_str, answer_str)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_phi4_template(): def test_phi4_template():
prompt_str = ( prompt_str = (
f"<|im_start|>user<|im_sep|>{MESSAGES[0]['content']}<|im_end|>" f"<|im_start|>user<|im_sep|>{MESSAGES[0]['content']}<|im_end|>"
@@ -406,7 +406,7 @@ def test_phi4_template():
_check_template("microsoft/phi-4", "phi4", prompt_str, answer_str) _check_template("microsoft/phi-4", "phi4", prompt_str, answer_str)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.xfail(not HF_TOKEN, reason="Authorization.") @pytest.mark.xfail(not HF_TOKEN, reason="Authorization.")
def test_qwen2_5_template(): def test_qwen2_5_template():
prompt_str = ( prompt_str = (
@@ -420,7 +420,7 @@ def test_qwen2_5_template():
_check_template("Qwen/Qwen2.5-7B-Instruct", "qwen", prompt_str, answer_str) _check_template("Qwen/Qwen2.5-7B-Instruct", "qwen", prompt_str, answer_str)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize("cot_messages", [True, False]) @pytest.mark.parametrize("cot_messages", [True, False])
def test_qwen3_template(cot_messages: bool): def test_qwen3_template(cot_messages: bool):
prompt_str = ( prompt_str = (
@@ -439,7 +439,7 @@ def test_qwen3_template(cot_messages: bool):
_check_template("Qwen/Qwen3-8B", "qwen3", prompt_str, answer_str, messages=messages) _check_template("Qwen/Qwen3-8B", "qwen3", prompt_str, answer_str, messages=messages)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_parse_llama3_template(): def test_parse_llama3_template():
tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3) tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)
template = parse_template(tokenizer) template = parse_template(tokenizer)
@@ -453,7 +453,7 @@ def test_parse_llama3_template():
assert template.default_system == "" assert template.default_system == ""
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.xfail(not HF_TOKEN, reason="Authorization.") @pytest.mark.xfail(not HF_TOKEN, reason="Authorization.")
def test_parse_qwen_template(): def test_parse_qwen_template():
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct") tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
@@ -466,7 +466,7 @@ def test_parse_qwen_template():
assert template.default_system == "You are Qwen, created by Alibaba Cloud. You are a helpful assistant." assert template.default_system == "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.xfail(not HF_TOKEN, reason="Authorization.") @pytest.mark.xfail(not HF_TOKEN, reason="Authorization.")
def test_parse_qwen3_template(): def test_parse_qwen3_template():
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B") tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")

View File

@@ -37,13 +37,13 @@ MESSAGES = [
EXPECTED_RESPONSE = "_rho" EXPECTED_RESPONSE = "_rho"
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_chat(): def test_chat():
chat_model = ChatModel(INFER_ARGS) chat_model = ChatModel(INFER_ARGS)
assert chat_model.chat(MESSAGES)[0].response_text == EXPECTED_RESPONSE assert chat_model.chat(MESSAGES)[0].response_text == EXPECTED_RESPONSE
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_stream_chat(): def test_stream_chat():
chat_model = ChatModel(INFER_ARGS) chat_model = ChatModel(INFER_ARGS)
response = "" response = ""

View File

@@ -49,7 +49,7 @@ INFER_ARGS = {
OS_NAME = os.getenv("OS_NAME", "") OS_NAME = os.getenv("OS_NAME", "")
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"stage,dataset", "stage,dataset",
[ [
@@ -66,7 +66,7 @@ def test_run_exp(stage: str, dataset: str):
assert os.path.exists(output_dir) assert os.path.exists(output_dir)
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_export(): def test_export():
export_dir = os.path.join("output", "llama3_export") export_dir = os.path.join("output", "llama3_export")
export_model({"export_dir": export_dir, **INFER_ARGS}) export_model({"export_dir": export_dir, **INFER_ARGS})

View File

@@ -17,7 +17,7 @@ import pytest
from llamafactory.eval.template import get_eval_template from llamafactory.eval.template import get_eval_template
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_eval_template_en(): def test_eval_template_en():
support_set = [ support_set = [
{ {
@@ -56,7 +56,7 @@ def test_eval_template_en():
] ]
@pytest.mark.runs_on(["cpu", "mps"]) @pytest.mark.runs_on(["cpu", "mps", "xpu"])
def test_eval_template_zh(): def test_eval_template_zh():
support_set = [ support_set = [
{ {

View File

@@ -52,7 +52,7 @@ def test_all_device():
assert DistributedInterface().get_local_world_size() == int(os.getenv("LOCAL_WORLD_SIZE", "1")) assert DistributedInterface().get_local_world_size() == int(os.getenv("LOCAL_WORLD_SIZE", "1"))
@pytest.mark.runs_on(["cuda", "npu"]) @pytest.mark.runs_on(["cuda", "npu", "xpu"])
@pytest.mark.require_distributed(2) @pytest.mark.require_distributed(2)
def test_multi_device(): def test_multi_device():
master_port = find_available_port() master_port = find_available_port()

View File

@@ -79,6 +79,8 @@ def _get_visible_devices_env() -> str | None:
return "CUDA_VISIBLE_DEVICES" return "CUDA_VISIBLE_DEVICES"
elif CURRENT_DEVICE == "npu": elif CURRENT_DEVICE == "npu":
return "ASCEND_RT_VISIBLE_DEVICES" return "ASCEND_RT_VISIBLE_DEVICES"
elif CURRENT_DEVICE == "xpu":
return "ZE_AFFINITY_MASK"
else: else:
return None return None
@@ -172,6 +174,8 @@ def _manage_distributed_env(request: FixtureRequest, monkeypatch: MonkeyPatch) -
monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
elif CURRENT_DEVICE == "npu": elif CURRENT_DEVICE == "npu":
monkeypatch.setattr(torch.npu, "device_count", lambda: 1) monkeypatch.setattr(torch.npu, "device_count", lambda: 1)
elif CURRENT_DEVICE == "xpu":
monkeypatch.setattr(torch.xpu, "device_count", lambda: 1)
@pytest.fixture(scope="session", autouse=True) @pytest.fixture(scope="session", autouse=True)

View File

@@ -96,7 +96,7 @@ def _test_sequence_parallel_loss(
assert loss is not None assert loss is not None
@pytest.mark.runs_on(["cuda", "npu"]) @pytest.mark.runs_on(["cuda", "npu", "xpu"])
@pytest.mark.require_distributed(2) @pytest.mark.require_distributed(2)
@pytest.mark.parametrize(("cp_size", "dp_size", "batch_size"), [(2, 1, 1), (2, 1, 2)]) @pytest.mark.parametrize(("cp_size", "dp_size", "batch_size"), [(2, 1, 1), (2, 1, 2)])
def test_sequence_parallel_loss(cp_size, dp_size, batch_size): def test_sequence_parallel_loss(cp_size, dp_size, batch_size):

View File

@@ -19,7 +19,7 @@ from llamafactory.v1.core.model_engine import ModelEngine
from llamafactory.v1.samplers.cli_sampler import SyncSampler from llamafactory.v1.samplers.cli_sampler import SyncSampler
@pytest.mark.runs_on(["cuda", "npu"]) @pytest.mark.runs_on(["cuda", "npu", "xpu"])
def test_sync_sampler(): def test_sync_sampler():
model_args = ModelArguments(model="Qwen/Qwen3-4B-Instruct-2507") model_args = ModelArguments(model="Qwen/Qwen3-4B-Instruct-2507")
sample_args = SampleArguments() sample_args = SampleArguments()

View File

@@ -21,7 +21,7 @@ import pytest
@pytest.mark.xfail(reason="CI machines may OOM when heavily loaded.") @pytest.mark.xfail(reason="CI machines may OOM when heavily loaded.")
@pytest.mark.runs_on(["cuda", "npu"]) @pytest.mark.runs_on(["cuda", "npu", "xpu"])
def test_fsdp2_dpo_trainer(tmp_path: Path): def test_fsdp2_dpo_trainer(tmp_path: Path):
"""Test FSDP2 DPO trainer with sigmoid loss by simulating `llamafactory-cli dpo config.yaml`.""" """Test FSDP2 DPO trainer with sigmoid loss by simulating `llamafactory-cli dpo config.yaml`."""
config_yaml = """\ config_yaml = """\

View File

@@ -20,7 +20,7 @@ import pytest
@pytest.mark.xfail(reason="CI machines may OOM when heavily loaded.") @pytest.mark.xfail(reason="CI machines may OOM when heavily loaded.")
@pytest.mark.runs_on(["cuda", "npu"]) @pytest.mark.runs_on(["cuda", "npu", "xpu"])
def test_fsdp2_sft_trainer(tmp_path: Path): def test_fsdp2_sft_trainer(tmp_path: Path):
"""Test FSDP2 SFT trainer by simulating `llamafactory-cli sft config.yaml` behavior.""" """Test FSDP2 SFT trainer by simulating `llamafactory-cli sft config.yaml` behavior."""
config_yaml = """\ config_yaml = """\