From b2ea8e1cd1947cdd85b185ce9a21da7de229fbad Mon Sep 17 00:00:00 2001 From: Baris Demir Date: Thu, 30 Jul 2026 16:08:11 +0100 Subject: [PATCH 1/2] Arm backend: Reland DeepSeek-R1-Distill-Qwen layer tests Relands the DeepSeek-R1-Distill-Qwen-1.5B layer tests that were reverted by pytorch/executorch#21048 because the TOSA model shard could run several checkpoint-shaped exports concurrently and OOM. The prerequisite CI change in pytorch/executorch#21492 added a serialized xlarge TOSA model suite while keeping the normal TOSA model shard parallelized. This reland marks the DeepSeek TOSA layer tests as xlarge so they are excluded from the normal shard and can be routed through the serialized xlarge suite. The tests use the checkpoint configuration from the Hugging Face model and the upstream Qwen2 layer implementations that back this distilled model. The covered layers include rotary embedding, rotary application, KV repetition, attention, RMSNorm, MLP, decoder layer, and final norm. Token embedding is excluded because the full checkpoint embedding allocation is too large for regular CI. Signed-off-by: Baris Demir Change-Id: Ia28581bbb4ffe070bc35af060fcceef2ac90084a --- backends/arm/MODELS.md | 1 + .../deepseek_r1_distill_qwen_test_config.py | 42 ++ .../test_deepseek_r1_distill_qwen_layers.py | 420 ++++++++++++++++++ 3 files changed, 463 insertions(+) create mode 100644 backends/arm/test/models/DeepSeek_R1_Distill_Qwen/deepseek_r1_distill_qwen_test_config.py create mode 100644 backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py diff --git a/backends/arm/MODELS.md b/backends/arm/MODELS.md index bcb410764bf..d289fa61e58 100644 --- a/backends/arm/MODELS.md +++ b/backends/arm/MODELS.md @@ -5,6 +5,7 @@ - Conformer - Deep AutoEncoder - Deit Tiny +- DeepSeek-R1-Distill-Qwen-1.5B - DeepLab v3 (DL3) - DS CNN - Inception v3 (IC3) diff --git a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/deepseek_r1_distill_qwen_test_config.py b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/deepseek_r1_distill_qwen_test_config.py new file mode 100644 index 00000000000..3487c2af9db --- /dev/null +++ b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/deepseek_r1_distill_qwen_test_config.py @@ -0,0 +1,42 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config + + +def get_deepseek_r1_distill_qwen_1_5b_checkpoint_config() -> Qwen2Config: + config = Qwen2Config( + architectures=["Qwen2ForCausalLM"], + attention_dropout=0.0, + bos_token_id=151643, + eos_token_id=151643, + hidden_act="silu", + hidden_size=1536, + initializer_range=0.02, + intermediate_size=8960, + max_position_embeddings=131072, + max_window_layers=21, + num_attention_heads=12, + num_hidden_layers=28, + num_key_value_heads=2, + rms_norm_eps=1e-6, # type: ignore[arg-type] + rope_parameters={ + "rope_type": "default", + "rope_theta": 10000.0, + }, + sliding_window=4096, + tie_word_embeddings=False, + transformers_version="4.44.0", + use_cache=True, + use_sliding_window=False, + vocab_size=151936, + ) + config.update( + { + "torch_dtype": "bfloat16", + "use_mrope": False, + } + ) + return config diff --git a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py new file mode 100644 index 00000000000..a7217e71a99 --- /dev/null +++ b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py @@ -0,0 +1,420 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Tuple + +import pytest +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.DeepSeek_R1_Distill_Qwen.deepseek_r1_distill_qwen_test_config import ( + get_deepseek_r1_distill_qwen_1_5b_checkpoint_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + VgfPipeline, +) + +pytest.importorskip("transformers.models.qwen2") + +from transformers.models.qwen2.modeling_qwen2 import ( # noqa: E402 + apply_rotary_pos_emb, + Qwen2Attention, + Qwen2DecoderLayer, + Qwen2MLP, + Qwen2RMSNorm, + Qwen2RotaryEmbedding, + repeat_kv, +) + +input_t = Tuple[torch.Tensor, ...] + + +def _make_deepseek_r1_distill_qwen_1_5b_layer_config(): + config = get_deepseek_r1_distill_qwen_1_5b_checkpoint_config() + config._attn_implementation = "sdpa" + return config + + +def _make_position_ids( + batch_size: int, seq_length: int, device: torch.device +) -> torch.Tensor: + return torch.arange(seq_length, device=device).unsqueeze(0).repeat(batch_size, 1) + + +def _make_causal_mask( + batch_size: int, seq_length: int, device: torch.device +) -> torch.Tensor: + mask = torch.full( + (seq_length, seq_length), torch.finfo(torch.float32).min, device=device + ) + mask = torch.triu(mask, diagonal=1) + return mask.unsqueeze(0).unsqueeze(0).repeat(batch_size, 1, 1, 1) + + +def _make_rope_embeddings( + config, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + rotary = Qwen2RotaryEmbedding(config) + return rotary(hidden_states, position_ids) + + +class DeepSeekR1DistillQwenTestModule(torch.nn.Module): + @classmethod + def prepare_model_and_inputs(cls): + raise NotImplementedError + + +def _to_bfloat16( + model: torch.nn.Module, inputs: input_t +) -> tuple[torch.nn.Module, input_t]: + return model.to(torch.bfloat16), tuple( + ( + x.to(torch.bfloat16) + if isinstance(x, torch.Tensor) and x.is_floating_point() + else x + ) + for x in inputs + ) + + +class RotaryEmbeddingModel(DeepSeekR1DistillQwenTestModule): + def __init__(self, config) -> None: + super().__init__() + self.rotary = Qwen2RotaryEmbedding(config) + + def forward( + self, hidden_states: torch.Tensor, position_ids: torch.Tensor + ) -> torch.Tensor: + cos, sin = self.rotary(hidden_states, position_ids) + return cos + sin + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.hidden_size) + position_ids = _make_position_ids(2, 8, hidden_states.device) + return model, (hidden_states, position_ids) + + +class RotaryApplyModel(DeepSeekR1DistillQwenTestModule): + def forward( + self, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + q_embed, k_embed = apply_rotary_pos_emb(q, k, cos, sin) + return q_embed.mean(dim=1) + k_embed.mean(dim=1) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls().eval() + hidden_states = torch.randn(2, 8, config.hidden_size) + position_ids = _make_position_ids(2, 8, hidden_states.device) + cos, sin = _make_rope_embeddings(config, hidden_states, position_ids) + head_dim = config.hidden_size // config.num_attention_heads + q = torch.randn(2, config.num_attention_heads, 8, head_dim) + k = torch.randn(2, config.num_key_value_heads, 8, head_dim) + return model, (q, k, cos, sin) + + +class RepeatKVModel(DeepSeekR1DistillQwenTestModule): + def __init__(self, config) -> None: + super().__init__() + self.n_rep = config.num_attention_heads // config.num_key_value_heads + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return repeat_kv(hidden_states, self.n_rep) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls(config).eval() + head_dim = config.hidden_size // config.num_attention_heads + hidden_states = torch.randn(2, config.num_key_value_heads, 8, head_dim) + return model, (hidden_states,) + + +class AttentionModel(DeepSeekR1DistillQwenTestModule): + def __init__(self, config) -> None: + super().__init__() + self.attn = Qwen2Attention(config, layer_idx=0) + self.rotary = Qwen2RotaryEmbedding(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + cos, sin = self.rotary(hidden_states, position_ids) + attn_output, _ = self.attn( + hidden_states=hidden_states, + position_embeddings=(cos, sin), + attention_mask=attention_mask, + ) + return attn_output + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.hidden_size) + attention_mask = _make_causal_mask(2, 8, hidden_states.device) + position_ids = _make_position_ids(2, 8, hidden_states.device) + return model, (hidden_states, attention_mask, position_ids) + + +class RMSNormModel(DeepSeekR1DistillQwenTestModule): + def __init__(self, config) -> None: + super().__init__() + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.hidden_size) + return model, (hidden_states,) + + +class MLPModel(DeepSeekR1DistillQwenTestModule): + def __init__(self, config) -> None: + super().__init__() + self.mlp = Qwen2MLP(config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.mlp(hidden_states) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.hidden_size) + return model, (hidden_states,) + + +class DecoderLayerModel(DeepSeekR1DistillQwenTestModule): + def __init__(self, config) -> None: + super().__init__() + self.layer = Qwen2DecoderLayer(config, layer_idx=0) + self.rotary = Qwen2RotaryEmbedding(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + cos, sin = self.rotary(hidden_states, position_ids) + return self.layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=(cos, sin), + ) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.hidden_size) + attention_mask = _make_causal_mask(2, 8, hidden_states.device) + position_ids = _make_position_ids(2, 8, hidden_states.device) + return model, (hidden_states, attention_mask, position_ids) + + +class FinalNormModel(DeepSeekR1DistillQwenTestModule): + def __init__(self, config) -> None: + super().__init__() + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_deepseek_r1_distill_qwen_1_5b_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.hidden_size) + return model, (hidden_states,) + + +@dataclass(frozen=True) +class DeepSeekR1DistillQwenTestCase: + model_cls: type[DeepSeekR1DistillQwenTestModule] + atol: float = 1e-3 + rtol: float = 1e-3 + qtol: int = 1 + transform_passes: tuple = field(default_factory=tuple) + + +TOSA_FP_TEST_CASES: dict[str, DeepSeekR1DistillQwenTestCase] = { + "rotary_embedding": DeepSeekR1DistillQwenTestCase(model_cls=RotaryEmbeddingModel), + "rotary_apply": DeepSeekR1DistillQwenTestCase(model_cls=RotaryApplyModel), + "repeat_kv": DeepSeekR1DistillQwenTestCase(model_cls=RepeatKVModel), + "attention": DeepSeekR1DistillQwenTestCase(model_cls=AttentionModel), + "rms_norm": DeepSeekR1DistillQwenTestCase(model_cls=RMSNormModel), + "mlp": DeepSeekR1DistillQwenTestCase(model_cls=MLPModel), + "decoder_layer": DeepSeekR1DistillQwenTestCase(model_cls=DecoderLayerModel), + "final_norm": DeepSeekR1DistillQwenTestCase(model_cls=FinalNormModel), +} + +TOSA_BF16_TEST_CASES: dict[str, DeepSeekR1DistillQwenTestCase] = { + "rotary_embedding": DeepSeekR1DistillQwenTestCase( + model_cls=RotaryEmbeddingModel, + atol=1e-2, + rtol=1e-2, + ), + "rotary_apply": DeepSeekR1DistillQwenTestCase( + model_cls=RotaryApplyModel, + atol=1e-2, + rtol=1e-2, + ), + "repeat_kv": DeepSeekR1DistillQwenTestCase( + model_cls=RepeatKVModel, + atol=1e-2, + rtol=1e-2, + ), + "attention": DeepSeekR1DistillQwenTestCase( + model_cls=AttentionModel, + atol=1e-2, + rtol=1e-2, + ), + "rms_norm": DeepSeekR1DistillQwenTestCase( + model_cls=RMSNormModel, + atol=1e-2, + rtol=1e-2, + ), + "mlp": DeepSeekR1DistillQwenTestCase( + model_cls=MLPModel, + atol=1e-2, + rtol=1e-2, + ), + "decoder_layer": DeepSeekR1DistillQwenTestCase( + model_cls=DecoderLayerModel, + atol=1e-2, + rtol=1e-2, + ), + "final_norm": DeepSeekR1DistillQwenTestCase( + model_cls=FinalNormModel, + atol=1e-2, + rtol=1e-2, + ), +} + +VGF_NO_QUANT_TEST_CASES: dict[str, DeepSeekR1DistillQwenTestCase] = { + "rotary_embedding": DeepSeekR1DistillQwenTestCase(model_cls=RotaryEmbeddingModel), + "rotary_apply": DeepSeekR1DistillQwenTestCase(model_cls=RotaryApplyModel), + "repeat_kv": DeepSeekR1DistillQwenTestCase(model_cls=RepeatKVModel), + "attention": DeepSeekR1DistillQwenTestCase(model_cls=AttentionModel), + "rms_norm": DeepSeekR1DistillQwenTestCase(model_cls=RMSNormModel), + "mlp": DeepSeekR1DistillQwenTestCase(model_cls=MLPModel), + "decoder_layer": DeepSeekR1DistillQwenTestCase(model_cls=DecoderLayerModel), + "final_norm": DeepSeekR1DistillQwenTestCase(model_cls=FinalNormModel), +} + +VGF_NO_QUANT_BF16_TEST_CASES: dict[str, DeepSeekR1DistillQwenTestCase] = ( + TOSA_BF16_TEST_CASES +) + + +@pytest.mark.xlarge +@common.parametrize( + "test_case", + TOSA_FP_TEST_CASES, +) +def test_deepseek_r1_distill_qwen_tosa_FP( + test_case: DeepSeekR1DistillQwenTestCase, +): + model, inputs = test_case.model_cls.prepare_model_and_inputs() + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + model, + inputs, + aten_op=[], + exir_op=[], + transform_passes=list(test_case.transform_passes), + ) + pipeline.run() + + +@pytest.mark.xlarge +@common.parametrize( + "test_case", + TOSA_BF16_TEST_CASES, +) +def test_deepseek_r1_distill_qwen_tosa_FP_bf16( + test_case: DeepSeekR1DistillQwenTestCase, +): + model, inputs = test_case.model_cls.prepare_model_and_inputs() + model, inputs = _to_bfloat16(model, inputs) + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + model, + inputs, + aten_op=[], + exir_op=[], + transform_passes=list(test_case.transform_passes), + tosa_extensions=["bf16"], + atol=test_case.atol, + rtol=test_case.rtol, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize( + "test_case", + VGF_NO_QUANT_TEST_CASES, +) +def test_deepseek_r1_distill_qwen_vgf_no_quant( + test_case: DeepSeekR1DistillQwenTestCase, +): + model, inputs = test_case.model_cls.prepare_model_and_inputs() + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + model, + inputs, + aten_op=[], + exir_op=[], + quantize=False, + atol=test_case.atol, + rtol=test_case.rtol, + qtol=test_case.qtol, + transform_passes=list(test_case.transform_passes), + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize( + "test_case", + VGF_NO_QUANT_BF16_TEST_CASES, +) +def test_deepseek_r1_distill_qwen_vgf_no_quant_bf16( + test_case: DeepSeekR1DistillQwenTestCase, +): + model, inputs = test_case.model_cls.prepare_model_and_inputs() + model, inputs = _to_bfloat16(model, inputs) + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + model, + inputs, + aten_op=[], + exir_op=[], + quantize=False, + atol=test_case.atol, + rtol=test_case.rtol, + qtol=test_case.qtol, + transform_passes=list(test_case.transform_passes), + ) + pipeline.run() From c4041229d450f04b16e13d2e29bd2db4a53a8f02 Mon Sep 17 00:00:00 2001 From: Baris Demir Date: Fri, 31 Jul 2026 10:42:30 +0100 Subject: [PATCH 2/2] Arm backend: Add serialized xlarge VKML model suite The DeepSeek-R1-Distill-Qwen layer tests were re-landed after the TOSA model shard OOM was addressed by a serialized xlarge TOSA suite. CommitlyExtended still showed memory pressure in the VKML model shard. The normal VKML model shard uses pytest-xdist auto parallelism, so xlarge VGF exports can also be resident at the same time. Add a serialized xlarge VKML model suite while keeping the normal VKML model shard parallelized. Mark the DeepSeek VGF layer tests as xlarge so they are excluded from normal VKML and can be routed through the serialized suite. Authored with Codex. Signed-off-by: Baris Demir Change-Id: I4b33f59a64ab2c40480e0f02c7415e999f4be7a2 --- .../test_deepseek_r1_distill_qwen_layers.py | 2 ++ backends/arm/test/test_arm_backend.sh | 14 ++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py index a7217e71a99..e99ea93b13e 100644 --- a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py +++ b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py @@ -371,6 +371,7 @@ def test_deepseek_r1_distill_qwen_tosa_FP_bf16( pipeline.run() +@pytest.mark.xlarge @common.SkipIfNoModelConverter @common.parametrize( "test_case", @@ -395,6 +396,7 @@ def test_deepseek_r1_distill_qwen_vgf_no_quant( pipeline.run() +@pytest.mark.xlarge @common.SkipIfNoModelConverter @common.parametrize( "test_case", diff --git a/backends/arm/test/test_arm_backend.sh b/backends/arm/test/test_arm_backend.sh index 6195854e8b6..a86f8ccc54d 100755 --- a/backends/arm/test/test_arm_backend.sh +++ b/backends/arm/test/test_arm_backend.sh @@ -268,6 +268,20 @@ test_pytest_models_vkml() { echo "${TEST_SUITE_NAME}: PASS" } +test_pytest_models_vkml_xlarge() { + echo "${TEST_SUITE_NAME}: Run xlarge pytest model tests with VKML runtime" + + source backends/arm/test/setup_testing_vkml.sh + + # Install model dependencies for pytest + source backends/arm/scripts/install_models_for_test.sh + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes \ + --numprocesses=1 --durations=0 \ + backends/arm/test/models -k _vgf_ -m xlarge + echo "${TEST_SUITE_NAME}: PASS" +} + test_run_vkml() { echo "${TEST_SUITE_NAME}: Test VKML delegate examples with run.sh"