Skip to content

Commit d46e620

Browse files
authored
Arm backend: Add serialised xlarge VKML model suite (pytorch#21518)
The DeepSeek-R1-Distill-Qwen layer tests were re-landed after the TOSA model shard OOM was addressed by a serialised 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 serialised xlarge VKML model suite while keeping the normal VKML model shard parallelised. Mark the DeepSeek VGF layer tests as xlarge so they are excluded from normal VKML and can be routed through the serialised suite. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani
1 parent 9789eea commit d46e620

4 files changed

Lines changed: 479 additions & 0 deletions

File tree

‎backends/arm/MODELS.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
- Conformer
66
- Deep AutoEncoder
77
- Deit Tiny
8+
- DeepSeek-R1-Distill-Qwen-1.5B
89
- DeepLab v3 (DL3)
910
- DS CNN
1011
- Inception v3 (IC3)
Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
7+
8+
9+
def get_deepseek_r1_distill_qwen_1_5b_checkpoint_config() -> Qwen2Config:
10+
config = Qwen2Config(
11+
architectures=["Qwen2ForCausalLM"],
12+
attention_dropout=0.0,
13+
bos_token_id=151643,
14+
eos_token_id=151643,
15+
hidden_act="silu",
16+
hidden_size=1536,
17+
initializer_range=0.02,
18+
intermediate_size=8960,
19+
max_position_embeddings=131072,
20+
max_window_layers=21,
21+
num_attention_heads=12,
22+
num_hidden_layers=28,
23+
num_key_value_heads=2,
24+
rms_norm_eps=1e-6, # type: ignore[arg-type]
25+
rope_parameters={
26+
"rope_type": "default",
27+
"rope_theta": 10000.0,
28+
},
29+
sliding_window=4096,
30+
tie_word_embeddings=False,
31+
transformers_version="4.44.0",
32+
use_cache=True,
33+
use_sliding_window=False,
34+
vocab_size=151936,
35+
)
36+
config.update(
37+
{
38+
"torch_dtype": "bfloat16",
39+
"use_mrope": False,
40+
}
41+
)
42+
return config

0 commit comments

Comments
 (0)