support qwen3.5 & transformers>=5.2

angelslim/compressor/quant/core/metrics.py

@@ -92,6 +92,7 @@ def __init__(self, processor, model_config=None):
             "QwenVL": [151643, 151654, 151655, 151656],
             "Qwen3VL": [151643, 151654, 151655, 151656],
             "HunyuanVL": [120000, 120020, 120118, 120119, 120120, 120121],
+            "Qwen3_5": [248044, 248055, 248056, 248057, 248076],
             "default": [],  # default: do not filter any tokens
         }
 

angelslim/compressor/quant/core/save.py

@@ -167,7 +167,7 @@ def save(self, save_path):
 
         os.makedirs(save_path, exist_ok=True)
 
-        self.quant_model.get_model().save_pretrained(save_path)
+        self.quant_model.get_model().save_pretrained(save_path, max_shard_size="5GB")
         self.quant_model.processor.save_pretrained(save_path)
         self.quant_model.tokenizer.save_pretrained(save_path)
 
@@ -265,7 +265,7 @@ def save(self, save_path):
         print_info("Save quantization_config: {}".format(quant_dict))
 
         os.makedirs(save_path, exist_ok=True)
-        self.quant_model.get_model().save_pretrained(save_path)
+        self.quant_model.get_model().save_pretrained(save_path, max_shard_size="5GB")
 
         with open(os.path.join(save_path, "hf_quant_config.json"), "w") as f:
             json.dump(trtllm_config, f, indent=4)

angelslim/compressor/speculative/inference/models/eagle3/target/modeling_qwen2_kv.py

@@ -15,7 +15,7 @@
 import torch
 from torch import nn
 from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, SlidingWindowCache, StaticCache
+from transformers.cache_utils import Cache, StaticCache
 from transformers.generation import GenerationMixin
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
@@ -639,7 +639,7 @@ def _update_causal_mask(
         )
 
         using_static_cache = isinstance(past_key_values, StaticCache)
-        using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
+        using_sliding_window_cache = isinstance(past_key_values, StaticCache)
 
         # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
         if (
@@ -659,7 +659,7 @@ def _update_causal_mask(
         dtype, device = input_tensor.dtype, input_tensor.device
         min_dtype = torch.finfo(dtype).min
         sequence_length = input_tensor.shape[1]
-        # SlidingWindowCache or StaticCache
+        # StaticCache
         if using_sliding_window_cache or using_static_cache:
             target_length = past_key_values.get_max_cache_shape()
         # DynamicCache or no cache
@@ -752,7 +752,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
                     # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
                     # the check is needed to verify is current checkpoint was trained with sliding window or not
                     if (
-                        not isinstance(past_key_values, SlidingWindowCache)
+                        not isinstance(past_key_values, StaticCache)
                         or sequence_length > target_length
                     ):
                         sliding_attend_mask = torch.arange(target_length, device=device) <= (

angelslim/compressor/speculative/inference/models/eagle3/target/modeling_qwen3_kv.py

@@ -15,7 +15,7 @@
 import torch
 from torch import nn
 from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, SlidingWindowCache, StaticCache
+from transformers.cache_utils import Cache, StaticCache
 from transformers.generation import GenerationMixin
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
@@ -692,7 +692,7 @@ def _update_causal_mask(
             past_key_values[0][0].current_length.item() if past_key_values is not None else 0
         )
         using_static_cache = isinstance(past_key_values, StaticCache)
-        using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
+        using_sliding_window_cache = isinstance(past_key_values, StaticCache)
 
         # When output attentions is True, sdpa implementation's forward method calls
         # the eager implementation's forward
@@ -713,7 +713,7 @@ def _update_causal_mask(
         dtype, device = input_tensor.dtype, input_tensor.device
         min_dtype = torch.finfo(dtype).min
         sequence_length = input_tensor.shape[1]
-        # SlidingWindowCache or StaticCache
+        # StaticCache
         if using_sliding_window_cache or using_static_cache:
             target_length = past_key_values.get_max_cache_shape()
         # DynamicCache or no cache
@@ -819,7 +819,7 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
                     # check is needed to verify is current checkpoint was
                     # trained with sliding window or not
                     if (
-                        not isinstance(past_key_values, SlidingWindowCache)
+                        not isinstance(past_key_values, StaticCache)
                         or sequence_length > target_length
                     ):
                         sliding_attend_mask = torch.arange(target_length, device=device) <= (

angelslim/data/multimodal_dataset.py

@@ -179,7 +179,7 @@ def _process_and_append(self, messages: List[Dict], tools=None):
         else:
             padding = True
 
-        if self.model_name in ["Qwen3VL", "Qwen3VLMoE"]:
+        if self.model_name in ["Qwen3VL", "Qwen3VLMoE", "Qwen3_5"]:
             inputs = self.processor.apply_chat_template(
                 messages,
                 tools=tools,

angelslim/data/text_dataset.py

@@ -112,7 +112,7 @@ def _load_jsonl_data(self, data_path: str, num_samples: int):
                             text = text + dic["content"] + self.processor.eos_token
 
                 model_inputs = self.processor(
-                    [text],
+                    text=[text],
                     return_tensors="pt",
                     max_length=self.max_length,
                     truncation=True,

angelslim/models/llm/qwen.py

@@ -14,12 +14,74 @@
 
 import re
 
+import torch
+import torch.nn as nn
+from transformers.models.qwen3_moe.modeling_qwen3_moe import (
+    Qwen3MoeExperts,
+    Qwen3MoeTopKRouter,
+)
+
 from ...compressor.quant.core import PTQSaveVllmHF
-from ...utils.utils import find_layers
+from ...utils.utils import find_layers, find_parent_layer_and_sub_name
 from ..base_model import BaseLLMModel
 from ..model_factory import SlimModelFactory
 
 
+class QwenMoeExpertsWithLinear(Qwen3MoeExperts):
+
+    def __init__(self, experts_layer):
+        super().__init__(experts_layer.config)
+        self.gate_up_proj = experts_layer.gate_up_proj
+        self.down_proj = experts_layer.down_proj
+        for expert_idx in range(self.num_experts):
+            expert = nn.ModuleDict(
+                {
+                    "gate_proj": nn.Linear(self.hidden_dim, self.intermediate_dim, bias=False),
+                    "up_proj": nn.Linear(self.hidden_dim, self.intermediate_dim, bias=False),
+                    "down_proj": nn.Linear(self.intermediate_dim, self.hidden_dim, bias=False),
+                }
+            )
+            expert["gate_proj"].weight.data, expert["up_proj"].weight.data = self.gate_up_proj[
+                expert_idx
+            ].chunk(2, dim=-2)
+            expert["down_proj"].weight.data = self.down_proj[expert_idx]
+            setattr(self, f"{expert_idx}", expert)
+        del self.gate_up_proj
+        del self.down_proj
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx == self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            expert_layer = getattr(self, f"{expert_idx}")
+            gate = expert_layer["gate_proj"](current_state)
+            up = expert_layer["up_proj"](current_state)
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = expert_layer["down_proj"](current_hidden_states)
+            current_hidden_states = (
+                current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            )
+            final_hidden_states.index_add_(
+                0, token_idx, current_hidden_states.to(final_hidden_states.dtype)
+            )
+
+        return final_hidden_states
+
+
 @SlimModelFactory.register
 class Qwen(BaseLLMModel):
     def __init__(
@@ -32,23 +94,40 @@ def __init__(
             deploy_backend=deploy_backend,
         )
         self.block_name = "model.layers"
-
-    def get_observer_layers(self):
-        names = [
+        self.observer_layer_classes = [torch.nn.Linear, Qwen3MoeTopKRouter]
+        self.observed_names = [
             "k_proj",
             "v_proj",
             "q_proj",
             "o_proj",
-            "up_proj",
             "gate_proj",
+            "up_proj",
             "down_proj",
         ]
+
+    def replace_moe(self):
+        for name, module in self.model.named_modules():
+            if isinstance(module, Qwen3MoeExperts) and not isinstance(
+                module, QwenMoeExpertsWithLinear
+            ):
+                print(name)
+                parent_layer, sub_name = find_parent_layer_and_sub_name(self.model, name)
+                moe_linear = QwenMoeExpertsWithLinear(module)
+                del module
+                setattr(parent_layer, sub_name, moe_linear)
+
+    def init_ptq(self, slim_config):
+        self.replace_moe()
+        super().init_ptq(slim_config)
+
+    def get_observer_layers(self):
         observer_layers_dict = {}
         layers_dict = find_layers(self.model, layers=self.observer_layer_classes)
 
         ignore_layers = self.skip_layer_names()
         for name, module in layers_dict.items():
-            if name.startswith(self.block_name) and name.split(".")[-1] in names:
+            # todo: shared_experts
+            if name.startswith(self.block_name) and name.split(".")[-1] in self.observed_names:
                 observer_layers_dict[name] = module
             else:
                 ignore_layers.append(name)

angelslim/models/llm/tiktoken_tokenizer.py

@@ -8,7 +8,7 @@
 import tiktoken
 from tiktoken.load import load_tiktoken_bpe
 from tokenizers import AddedToken
-from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode
+from transformers.convert_slow_tokenizer import bytes_to_unicode
 from transformers.tokenization_utils import PreTrainedTokenizer
 
 logger = getLogger(__name__)

angelslim/models/vlm/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from .hunyuan_vl import HunyuanVL  # noqa: F401
+from .qwen3_5 import Qwen3_5  # noqa: F401
 from .qwen3_vl import Qwen3VL  # noqa: F401
 from .qwen3_vl_moe import Qwen3VLMoE  # noqa: F401
 from .qwen_vl import QwenVL  # noqa: F401