Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,7 @@
+{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% endif %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}<|im_end|>
+{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

config.json

@@ -0,0 +1,82 @@
+{
+  "architectures": [
+    "QTSplusQwen2_5_VLTextForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_qts_plus_qwen2_5_vl.QTSplusQwen2_5_VL_CausalLM_Config",
+    "AutoModelForCausalLM": "modeling_qts_plus_qwen2_5_vl.QTSplusQwen2_5_VLTextForCausalLM",
+    "AutoProcessor": "processing_qts_plus_qwen2_5_vl.QTSplusQwen2_5_VLProcessor"
+  },
+  "vision_tower": "qwen2_5_vl_vision",
+  "enable_qts_plus": true,
+  "qts_plus_n_heads": 8,
+  "qts_plus_tau_s": 0.5,
+  "qts_plus_nmax": 25600,
+  "qts_plus_rho_min": 0.05,
+  "qts_plus_rho_max": 0.5,
+  "qts_plus_block_dropout": 0.0,
+  "qts_plus_reencode": true,
+  "qts_plus_scoring_layers": 1,
+  "qts_plus_reencode_layers": 2,
+  "project_text_if_needed": false,
+  "freeze_qts_scoring_layers": false,
+  "lambda_t": 0,
+  "lambda_m": 0,
+  "lambda_s": 0,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "vision_start_token_id": 151652,
+  "vision_end_token_id": 151653,
+  "vision_token_id": 151654,
+  "image_token_id": 151655,
+  "video_token_id": 151656,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "max_position_embeddings": 128000,
+  "max_window_layers": 70,
+  "model_type": "qts_plus_qwen2_5_vl_causal_lm",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sliding_window": 32768,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.41.2",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vision_config": {
+    "depth": 32,
+    "hidden_act": "silu",
+    "hidden_size": 1280,
+    "intermediate_size": 3420,
+    "num_heads": 16,
+    "in_chans": 3,
+    "out_hidden_size": 2048,
+    "patch_size": 14,
+    "spatial_merge_size": 2,
+    "spatial_patch_size": 14,
+    "window_size": 112,
+    "fullatt_block_indexes": [
+      7,
+      15,
+      23,
+      31
+    ],
+    "tokens_per_second": 2,
+    "temporal_patch_size": 2
+  },
+  "rope_scaling": {
+    "type": "mrope",
+    "mrope_section": [
+      16,
+      24,
+      24
+    ]
+  },
+  "vocab_size": 151936
+}

configuration_qts_plus_qwen2_5_vl.py

@@ -0,0 +1,21 @@
+"""
+Self-contained config shim for trust_remote_code.
+
+This file defines the minimal configuration class expected by
+`config.json` without importing from a local `src` package.
+"""
+from transformers import AutoConfig
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLTextConfig
+
+class QTSplusQwen2_5_VL_CausalLM_Config(Qwen2_5_VLTextConfig):
+    """Config alias for QTS+ Qwen2.5-VL Causal LM.
+
+    It inherits from the upstream Qwen2.5-VL text config and only sets a
+    distinct `model_type` so that Transformers can resolve the proper
+    architecture via `auto_map`.
+    """
+
+    model_type = "qts_plus_qwen2_5_vl_causal_lm"
+
+AutoConfig.register("qts_plus_qwen2_5_vl_causal_lm", QTSplusQwen2_5_VL_CausalLM_Config)
+__all__ = ["QTSplusQwen2_5_VL_CausalLM_Config"]

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "eos_token_id": [
+    151645
+  ],
+  "pad_token_id": 151643,
+  "transformers_version": "4.57.1"
+}

latest

@@ -0,0 +1 @@
+global_step23740

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac85a74a6f3207d337d5f77bd3833975454e13eef3fd9e281a83892e85f2f2d4
+size 8390793330

modeling_qts_plus_qwen2_5_vl.py

@@ -0,0 +1,1090 @@
+"""
+Self-contained modeling shim for trust_remote_code.
+
+Implements the QTS+ Qwen2.5‑VL Causal LM architecture locally by
+composing upstream Transformers' Qwen2.5‑VL text and vision modules
+with a lightweight QTS+ selector. This avoids importing any local `src`
+package while preserving checkpoint compatibility (including
+`model.vision_tower.*` and `model.qts_plus.selector.*` parameters).
+"""
+import json
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, AutoModelForCausalLM, logging
+from transformers.modeling_flash_attention_utils import is_flash_attn_available
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.generation import GenerationMixin
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VisionTransformerPretrainedModel as Qwen2_5_VisionTransformerPretrainedModelBase
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VLTextModel,
+    Qwen2_5_VLPreTrainedModel,
+
+)
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig
+from .configuration_qts_plus_qwen2_5_vl import (
+    QTSplusQwen2_5_VL_CausalLM_Config
+)
+logger = logging.get_logger(__name__)
+# ------------------------------
+# Utilities: embedding integration
+# ------------------------------
+def qts_integrate_embeddings(
+    vision_features: torch.Tensor,
+    input_ids: torch.Tensor,
+    attention_mask: torch.Tensor,
+    labels: Optional[torch.Tensor] = None,
+    image_token_id: Optional[int] = None,
+    video_token_id: Optional[int] = None,
+    image_grid_thw: Optional[torch.Tensor] = None,
+    video_grid_thw: Optional[torch.Tensor] = None,
+    text_model_embed_layer: Optional[nn.Embedding] = None,
+    kept_indices: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Integrate visual features into text embeddings (single-sample batch).
+
+    This mirrors the behavior of the full Qwen2.5‑VL generation path, but
+    works with pre-computed visual features and placeholder tokens in the
+    text sequence. It supports both the single <|video_pad|> token case and
+    multi-placeholder templates.
+    """
+    if text_model_embed_layer is None:
+        raise ValueError("text_model_embed_tokens is required for text embedding integration")
+    if input_ids.dtype is not torch.long:
+        input_ids = input_ids.long()
+
+    inputs_embeds = text_model_embed_layer(input_ids)
+    if vision_features.shape[0] <= 0:
+        raise ValueError("vision_features must contain at least one feature vector")
+    if video_token_id is None:
+        raise ValueError("video_token_id must be provided for video feature integration")
+
+    B, S = input_ids.shape
+    assert B == 1, "Sequence-trimming currently assumes batch_size == 1."
+
+    vid_pos = (input_ids[0] == video_token_id).nonzero(as_tuple=False).flatten()
+    n_feats = int(vision_features.shape[0])
+
+    if vid_pos.numel() == 1 and n_feats >= 1:
+        insert_idx = int(vid_pos.item())
+        vision_features = vision_features.to(inputs_embeds.device, inputs_embeds.dtype)
+
+        pre_embeds = inputs_embeds[:, :insert_idx, :]
+        post_embeds = inputs_embeds[:, insert_idx + 1 :, :]
+
+        feats_embeds = vision_features.unsqueeze(0)
+        inputs_embeds = torch.cat([pre_embeds, feats_embeds, post_embeds], dim=1)
+
+        feats_mask = torch.ones((1, n_feats), dtype=attention_mask.dtype, device=attention_mask.device)
+        pre_mask = attention_mask[:, :insert_idx]
+        post_mask = attention_mask[:, insert_idx + 1 :]
+        attention_mask = torch.cat([pre_mask, feats_mask, post_mask], dim=1)
+
+        if labels is not None:
+            labels = labels.clone()
+            if labels.size(1) > insert_idx:
+                pre_labels = labels[:, :insert_idx]
+                post_labels = labels[:, insert_idx + 1 :]
+                pad = torch.full((1, n_feats), -100, dtype=labels.dtype, device=labels.device)
+                labels = torch.cat([pre_labels, pad, post_labels], dim=1)
+        return inputs_embeds, attention_mask, labels
+
+    # Fallback: multi-placeholder handling
+    M = int(vid_pos.numel())
+    if M == 0:
+        raise ValueError("No video placeholder tokens found in input_ids for provided vision_features")
+
+    vision_features = vision_features.to(inputs_embeds.device, inputs_embeds.dtype)
+    N = n_feats
+    if N > M:
+        raise NotImplementedError(
+            "Number of vision features exceeds video placeholders; use a single <|video_pad|> token template."
+        )
+    if N < M:
+        drop_pos = vid_pos[N:]
+        if drop_pos.numel() > 0:
+            keep_seq = torch.ones(S, dtype=torch.bool, device=input_ids.device)
+            keep_seq[drop_pos] = False
+            input_ids = input_ids[:, keep_seq]
+            attention_mask = attention_mask[:, keep_seq]
+            inputs_embeds = inputs_embeds[:, keep_seq, :]
+            if labels is not None:
+                labels = labels[:, keep_seq]
+            vid_pos = (input_ids[0] == video_token_id).nonzero(as_tuple=False).flatten()
+            M = int(vid_pos.numel())
+
+    for i in range(N):
+        pos = int(vid_pos[i].item())
+        inputs_embeds[0, pos, :] = vision_features[i, :]
+    if labels is not None and N > 0:
+        labels = labels.clone()
+        labels[0, vid_pos[:N]] = -100
+
+    return inputs_embeds, attention_mask.to(inputs_embeds.device), labels
+
+
+# ------------------------------
+# QTS+ modules (selector + tokenizer)
+# ------------------------------
+class RMSNorm(nn.Module):
+    def __init__(self, d: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(norm + self.eps)
+        return self.weight * x
+
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.GELU(),
+            nn.Linear(d_ff, d_model),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class Qwen2_5_ScoringCrossAttentionLayer(nn.Module):
+    """Qwen2.5-style cross-attention used in QTS+ scoring.
+
+    Separate q/k/v projections (with optional multi-query kv heads) followed by
+    an output projection and a small FFN on the query path.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_key_value_heads: Optional[int] = None,
+        dropout: float = 0.0,
+        d_ff: Optional[int] = None,
+        rms_norm_eps: float = 1e-6,
+        use_qwen_rms: bool = True,
+    ) -> None:
+        super().__init__()
+        assert d_model % num_heads == 0
+        self.hidden_size = d_model
+        self.num_heads = int(num_heads)
+        self.head_dim = d_model // self.num_heads
+        self.num_key_value_heads = int(num_key_value_heads) if num_key_value_heads else self.num_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.attention_dropout = dropout
+
+        # Minimal Qwen-like RMS norms
+        class _Qwen2RMSNorm(nn.Module):
+            def __init__(self, hidden_size: int, eps: float = 1e-6):
+                super().__init__()
+                self.weight = nn.Parameter(torch.ones(hidden_size))
+                self.eps = float(eps)
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                dtype = x.dtype
+                x = x.float()
+                variance = x.pow(2).mean(-1, keepdim=True)
+                x = x * torch.rsqrt(variance + self.eps)
+                x = x.to(dtype)
+                return self.weight * x
+
+        self.q_norm = _Qwen2RMSNorm(d_model, eps=rms_norm_eps) if use_qwen_rms else RMSNorm(d_model, eps=rms_norm_eps)
+        self.kv_norm = _Qwen2RMSNorm(d_model, eps=rms_norm_eps) if use_qwen_rms else RMSNorm(d_model, eps=rms_norm_eps)
+        self.ffn_norm = _Qwen2RMSNorm(d_model, eps=rms_norm_eps) if use_qwen_rms else RMSNorm(d_model, eps=rms_norm_eps)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self.ffn = FeedForward(d_model, d_ff or (4 * d_model), dropout=dropout)
+
+    @staticmethod
+    def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+        b, h_kv, t, dh = x.shape
+        if n_rep == 1:
+            return x
+        x = x[:, :, None, :, :].expand(b, h_kv, n_rep, t, dh)
+        return x.reshape(b, h_kv * n_rep, t, dh)
+
+    def forward(
+        self,
+        q: torch.Tensor,  # [B, L, D]
+        kv: torch.Tensor,  # [B, M, D]
+        kv_key_padding_mask: Optional[torch.Tensor] = None,  # [B, M]
+        need_weights: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        B, L, _ = q.shape
+        _, M, _ = kv.shape
+
+        qn = self.q_norm(q)
+        kvn = self.kv_norm(kv)
+
+        q_states = self.q_proj(qn)
+        k_states = self.k_proj(kvn)
+        v_states = self.v_proj(kvn)
+
+        q_states = q_states.view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k_states = k_states.view(B, M, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        v_states = v_states.view(B, M, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if self.num_key_value_groups > 1:
+            k_states = self._repeat_kv(k_states, self.num_key_value_groups)
+            v_states = self._repeat_kv(v_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(q_states, k_states.transpose(2, 3)) / (self.head_dim ** 0.5)
+        if kv_key_padding_mask is not None:
+            mask = kv_key_padding_mask[:, None, None, :].to(dtype=attn_weights.dtype)
+            attn_weights = attn_weights.masked_fill(mask > 0.5, float("-inf"))
+        attn_dtype = attn_weights.dtype
+        attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(attn_dtype)
+        attn_output = torch.matmul(attn_weights, v_states)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L, self.num_heads * self.head_dim)
+
+        out = self.o_proj(attn_output)
+        q = q + out
+        q = q + self.ffn(self.ffn_norm(q))
+        return q, (attn_weights if need_weights else None)
+
+
+class Qwen2_5_SelfReencodeLayer(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_key_value_heads: Optional[int] = None,
+        dropout: float = 0.0,
+        d_ff: Optional[int] = None,
+        rms_norm_eps: float = 1e-6,
+        use_qwen_rms: bool = True,
+    ) -> None:
+        super().__init__()
+        self.core = Qwen2_5_ScoringCrossAttentionLayer(
+            d_model=d_model,
+            num_heads=num_heads,
+            num_key_value_heads=num_key_value_heads or num_heads,
+            dropout=dropout,
+            d_ff=d_ff,
+            rms_norm_eps=rms_norm_eps,
+            use_qwen_rms=use_qwen_rms,
+        )
+
+    def forward(self, x: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        y, _ = self.core(x, x, kv_key_padding_mask=key_padding_mask, need_weights=False)
+        return y
+
+    def init_from_qwen_attn(self, qwen_attn: nn.Module, qwen_input_ln: Optional[nn.Module] = None) -> None:
+        self.core.init_from_qwen_attn(qwen_attn, qwen_input_ln)
+
+
+class BudgetHead(nn.Module):
+    def __init__(self, d_model: int, hidden: int = 256, rho_min: float = 0.05, rho_max: float = 0.5) -> None:
+        super().__init__()
+        self.rho_min = rho_min
+        self.rho_max = rho_max
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model + 3, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, 1),
+        )
+
+    def forward(self, sq: torch.Tensor, logM: torch.Tensor, r_max: torch.Tensor, H: torch.Tensor) -> torch.Tensor:
+        B, D = sq.shape
+        x = torch.cat([sq, logM.view(B, 1), r_max.view(B, 1), H.view(B, 1)], dim=1)
+        # Ensure input dtype matches layer weights to avoid Float/Half mismatch
+        x = x.to(dtype=self.mlp[0].weight.dtype)
+        logits = self.mlp(x).squeeze(1)
+        rho = self.rho_min + (self.rho_max - self.rho_min) * torch.sigmoid(logits)
+        return rho
+
+
+class QTSplus(nn.Module):
+    """Query‑Aware Token Selector with Adaptive Budget."""
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int = 8,
+        n_kv_heads: Optional[int] = None,
+        tau_s: float = 0.1,
+        nmax: int = 2560,
+        rho_min: float = 0.05,
+        rho_max: float = 0.5,
+        block_dropout: float = 0.0,
+        use_reencode: bool = True,
+        n_scoring_layers: int = 1,
+        n_reencode_layers: int = 1,
+    ) -> None:
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.d_model = d_model
+        self.n_heads = int(n_heads)
+        self.d_head = d_model // self.n_heads
+        self.tau_s = float(tau_s)
+        self.nmax = int(nmax)
+        self.use_reencode = bool(use_reencode)
+        self.n_scoring_layers = max(int(n_scoring_layers), 1)
+        self.n_reencode_layers = max(int(n_reencode_layers), 1)
+
+        n_kv_heads_eff = int(n_kv_heads) if (n_kv_heads is not None and int(n_kv_heads) > 0) else self.n_heads
+        self.scoring_layers = nn.ModuleList(
+            [
+                Qwen2_5_ScoringCrossAttentionLayer(
+                    d_model,
+                    num_heads=self.n_heads,
+                    num_key_value_heads=n_kv_heads_eff,
+                    dropout=0.0,
+                    rms_norm_eps=1e-6,
+                    use_qwen_rms=True,
+                )
+                for _ in range(self.n_scoring_layers)
+            ]
+        )
+
+        self.budget = BudgetHead(d_model, rho_min=rho_min, rho_max=rho_max)
+
+        if self.use_reencode:
+            self.reencode_layers = nn.ModuleList(
+                [
+                    Qwen2_5_SelfReencodeLayer(
+                        d_model,
+                        num_heads=self.n_heads,
+                        num_key_value_heads=n_kv_heads_eff,
+                        dropout=0.0,
+                        rms_norm_eps=1e-6,
+                        use_qwen_rms=True,
+                    )
+                    for _ in range(self.n_reencode_layers)
+                ]
+            )
+        else:
+            self.reencode_layers = None
+
+    def _score(self, Xv: torch.Tensor, Qt: torch.Tensor) -> torch.Tensor:
+        # Simple cross-attention based scoring aggregated across heads and query positions
+        B, M, D = Xv.shape
+        q = Qt
+        kv = Xv
+        for layer in self.scoring_layers:
+            q, attn = layer(q, kv, need_weights=True)
+        # attn: [B, H, L, M]; aggregate -> [B, M]
+        r = attn.amax(dim=2).mean(dim=1)
+        return r
+
+    def _predict_budget(self, q: torch.Tensor, r: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        B, L, D = q.shape
+        M = r.shape[-1]
+        sq = q.mean(dim=1)
+        # Create logM with same dtype/device as q to keep types consistent
+        logM = torch.log(torch.tensor(float(M), device=q.device, dtype=q.dtype)).expand(B)
+        r_max = r.max(dim=-1).values
+        # entropy H over token scores after softmax
+        p = torch.softmax(r, dim=-1)
+        H = -(p * (p.clamp(min=1e-12).log())).sum(dim=-1)
+        rho = self.budget(sq, logM, r_max, H)
+        # n = clamp(round(rho * M), 1, nmax)
+        n = torch.clamp((rho * float(M)).round(), min=1.0, max=float(self.nmax)).to(torch.long)
+        return rho, n
+
+    def forward(self, Xv: torch.Tensor, Qt: torch.Tensor, mode: str = "train") -> Dict[str, Any]:
+        assert mode in ("train", "infer")
+        B, M, D = Xv.shape
+        r = self._score(Xv, Qt)
+        rho, n = self._predict_budget(Qt, r)
+
+        # Hard top-n with original order preserved
+        kept_idx_list: List[torch.Tensor] = []
+        Z_out: List[torch.Tensor] = []
+        for b in range(B):
+            kb = torch.topk(r[b], k=int(n[b].item()), dim=0).indices
+            kb, _ = torch.sort(kb)
+            kept_idx_list.append(kb)
+            Z_out.append(Xv[b, kb])
+
+        if self.use_reencode:
+            max_keep = int(max(z.size(0) for z in Z_out))
+            Zb = []
+            for z in Z_out:
+                if z.size(0) < max_keep:
+                    pad = z[-1:].repeat(max_keep - z.size(0), 1)
+                    z = torch.cat([z, pad], dim=0)
+                Zb.append(z.unsqueeze(0))
+            Zb = torch.cat(Zb, dim=0)
+            for layer in self.reencode_layers or []:
+                Zb = layer(Zb)
+            Z_final = [Zb[b, : kept_idx_list[b].numel()] for b in range(B)]
+        else:
+            Z_final = Z_out
+
+        # Simple training proxies
+        p = torch.softmax(r, dim=-1)
+        flops_proxy = ((rho * float(M)) ** 2) / float(self.nmax ** 2)
+        kv_proxy = (rho * float(M)) / float(self.nmax)
+
+        return {
+            "indices": kept_idx_list,
+            "Z": Z_final,
+            "rho": rho,
+            "r": r,
+            "n": n,
+            "add_loss": {
+                "flops": flops_proxy.mean(),
+                "kv": kv_proxy.mean(),
+                "smooth": torch.tensor(0.0, device=Xv.device, dtype=Xv.dtype),
+            },
+        }
+
+
+class QTSplusTokenizerConfig:
+    def __init__(
+        self,
+        embedding_dim: int,
+        n_heads: int = 8,
+        num_kv_heads: Optional[int] = None,
+        tau_s: float = 0.1,
+        nmax: int = 2560,
+        rho_min: float = 0.05,
+        rho_max: float = 0.5,
+        block_dropout: float = 0.0,
+        reencode: bool = True,
+        scoring_layers: int = 1,
+        reencode_layers: int = 1,
+        lambda_t: float = 1.0,
+        lambda_m: float = 1.7,
+        lambda_s: float = 0.05,
+        project_text_if_needed: bool = False,
+    ) -> None:
+        self.embedding_dim = embedding_dim
+        self.n_heads = n_heads
+        self.num_kv_heads = num_kv_heads
+        self.tau_s = tau_s
+        self.nmax = nmax
+        self.rho_min = rho_min
+        self.rho_max = rho_max
+        self.block_dropout = block_dropout
+        self.reencode = reencode
+        self.scoring_layers = scoring_layers
+        self.reencode_layers = reencode_layers
+        self.lambda_t = lambda_t
+        self.lambda_m = lambda_m
+        self.lambda_s = lambda_s
+        self.project_text_if_needed = project_text_if_needed
+
+
+class QTSplusTokenizer(nn.Module):
+    def __init__(self, cfg: QTSplusTokenizerConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.selector = QTSplus(
+            d_model=cfg.embedding_dim,
+            n_heads=cfg.n_heads,
+            n_kv_heads=cfg.num_kv_heads or cfg.n_heads,
+            tau_s=cfg.tau_s,
+            nmax=cfg.nmax,
+            rho_min=cfg.rho_min,
+            rho_max=cfg.rho_max,
+            block_dropout=cfg.block_dropout,
+            use_reencode=cfg.reencode,
+            n_scoring_layers=cfg.scoring_layers,
+            n_reencode_layers=cfg.reencode_layers,
+        )
+        self.text_proj: Optional[nn.Linear] = None
+
+    def forward(self, X_v: torch.Tensor, Q_t: torch.Tensor, mode: str = "train") -> Dict[str, Any]:
+        B, M, D = X_v.shape
+        D_txt = Q_t.shape[-1]
+        if D_txt != D:
+            if self.cfg.project_text_if_needed:
+                if self.text_proj is None:
+                    self.text_proj = nn.Linear(D_txt, D, bias=False).to(device=Q_t.device, dtype=Q_t.dtype)
+                Q_proj = self.text_proj(Q_t)
+            else:
+                raise ValueError(f"QTS+ expects text dim {D}, got {D_txt}. Set project_text_if_needed=True.")
+        else:
+            Q_proj = Q_t
+        sel = self.selector(X_v, Q_proj, mode=mode)
+        # Add simple proxies for train-time regularization
+        M_tensor = torch.tensor(float(M), device=X_v.device, dtype=X_v.dtype)
+        rho = sel["rho"]
+        flops_proxy = ((rho * M_tensor) ** 2) / float(self.cfg.nmax ** 2)
+        kv_proxy = (rho * M_tensor) / float(self.cfg.nmax)
+        sel["add_loss"] = {
+            "flops": flops_proxy.mean() * self.cfg.lambda_t,
+            "kv": kv_proxy.mean() * self.cfg.lambda_m,
+            "smooth": torch.tensor(0.0, device=X_v.device, dtype=X_v.dtype),
+        }
+        return sel
+
+class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VisionTransformerPretrainedModelBase):
+    def __init__(self, config, *inputs, **kwargs) -> None:
+        super().__init__(config, *inputs, **kwargs)
+    
+    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
+        # Return the output from the base implementation.
+        # Without this return, callers receive None and downstream code fails.
+        return super().forward(hidden_states, grid_thw, **kwargs)
+
+    def get_video_features(
+        self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None
+    ):
+        """
+        Encodes videos into continuous embeddings that can be forwarded to the language model.
+
+        Args:
+            pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+                The tensors corresponding to the input videos.
+            video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+                The temporal, height and width of feature shape of each video in LLM.
+        """
+        pixel_values_videos = pixel_values_videos.type(self.dtype)
+        video_embeds = self.forward(pixel_values_videos, grid_thw=video_grid_thw)
+        # split_sizes = (video_grid_thw.prod(-1) // self.spatial_merge_size**2).tolist()
+        # video_embeds = torch.split(video_embeds, split_sizes)
+        return video_embeds
+    
+def _try_load_vision_config_from_path(path: str) -> Optional[Dict[str, Any]]:
+    """Best-effort load of Qwen2.5-VL vision `config.json`.
+
+    Accepts either a directory containing `config.json` or a file path to a
+    weights file. In the latter case, attempts to locate a sibling
+    `config.json` in the same directory.
+    """
+    if not path:
+        return None
+
+    cfg_path = None
+    if os.path.isdir(path):
+        candidate = os.path.join(path, "config.json")
+        if os.path.isfile(candidate):
+            cfg_path = candidate
+    else:
+        # If a file is given (e.g., .../model.safetensors), look next to it
+        base_dir = os.path.dirname(path)
+        candidate = os.path.join(base_dir, "config.json")
+        if os.path.isfile(candidate):
+            cfg_path = candidate
+
+    if cfg_path is None:
+        return None
+
+    try:
+        with open(cfg_path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return None
+
+
+def build_vision_tower(vision_tower_cfg, **kwargs):
+    vision_tower = getattr(vision_tower_cfg, 'mm_vision_tower', getattr(vision_tower_cfg, 'vision_tower', None))
+    if vision_tower != "qwen2_5_vl_vision":
+        raise ValueError(f"Unknown vision tower type: {vision_tower}")
+
+    # Attempt to infer correct dimensions from the provided pretrained path
+    pretrained_path = getattr(vision_tower_cfg, 'pretrain_vision_model', None)
+    cfg_json = _try_load_vision_config_from_path(pretrained_path) if pretrained_path else None
+
+    if cfg_json is not None:
+        # Map json fields to Qwen2_5_VLVisionConfig kwargs (use json defaults when available)
+        config = Qwen2_5_VLVisionConfig(
+            hidden_size=cfg_json.get("hidden_size", 1280),
+            out_hidden_size=cfg_json.get("out_hidden_size", cfg_json.get("hidden_size", 1280)),
+            depth=cfg_json.get("depth", 32),
+            intermediate_size=cfg_json.get("intermediate_size", 3420),
+            num_heads=cfg_json.get("num_heads", 16),
+            fullatt_block_indexes=cfg_json.get("fullatt_block_indexes", [7, 15, 23, 31]),
+            in_channels=cfg_json.get("in_channels", cfg_json.get("in_chans", 3)),
+            patch_size=cfg_json.get("patch_size", cfg_json.get("spatial_patch_size", 14)),
+            spatial_merge_size=cfg_json.get("spatial_merge_size", 2),
+            temporal_patch_size=cfg_json.get("temporal_patch_size", 2),
+            tokens_per_second=cfg_json.get("tokens_per_second", 2),
+            window_size=cfg_json.get("window_size", 112),
+            initializer_range=cfg_json.get("initializer_range", 0.02),
+        )
+    else:
+        # Fallback to a safe default (3B) when no config file is available
+        # This keeps backwards-compatibility but different-scale checkpoints
+        # should always provide a config.json alongside the weights.
+        config = Qwen2_5_VLVisionConfig(
+            hidden_size=1280,
+            out_hidden_size=2048,
+            depth=32,
+            intermediate_size=3420,
+            num_heads=16,
+            fullatt_block_indexes=[7, 15, 23, 31],
+        )
+
+    return Qwen2_5_VisionTransformerPretrainedModel(config)
+
+# ------------------------------
+# Builders used by the meta model
+# ------------------------------
+def build_vision_tower(config: Qwen2_5_VLTextConfig) -> Qwen2_5_VisionTransformerPretrainedModel:
+    vcfg_dict = getattr(config, "vision_config", None) or {}
+    vcfg = Qwen2_5_VLVisionConfig(**vcfg_dict) if vcfg_dict else Qwen2_5_VLVisionConfig()
+    return Qwen2_5_VisionTransformerPretrainedModel(vcfg)
+
+
+def build_qts_plus_tower(config: Qwen2_5_VLTextConfig) -> QTSplusTokenizer:
+    lm_heads = getattr(config, "num_attention_heads", None)
+    vision_dim = getattr(config, "vision_embed_size", None)
+    if not isinstance(lm_heads, int) or lm_heads <= 0:
+        raise ValueError("num_attention_heads must be provided by the Qwen2.5‑VL config")
+    if not isinstance(vision_dim, int) or vision_dim <= 0:
+        raise ValueError("vision_embed_size must be a positive int before building QTS+")
+    if vision_dim % lm_heads != 0:
+        raise ValueError(
+            f"vision_embed_size ({vision_dim}) must be divisible by LM num_attention_heads ({lm_heads})"
+        )
+    kv_heads = getattr(config, "num_key_value_heads", None)
+    cfg = QTSplusTokenizerConfig(
+        embedding_dim=vision_dim,
+        n_heads=lm_heads,
+        num_kv_heads=kv_heads if isinstance(kv_heads, int) and kv_heads > 0 else None,
+        tau_s=getattr(config, "qts_plus_tau_s", 0.1),
+        nmax=getattr(config, "qts_plus_nmax", 2560),
+        rho_min=getattr(config, "qts_plus_rho_min", 0.05),
+        rho_max=getattr(config, "qts_plus_rho_max", 0.5),
+        block_dropout=getattr(config, "qts_plus_block_dropout", 0.0),
+        reencode=getattr(config, "qts_plus_reencode", True),
+        scoring_layers=getattr(config, "qts_plus_scoring_layers", 1),
+        reencode_layers=getattr(config, "qts_plus_reencode_layers", 1),
+        lambda_t=getattr(config, "lambda_t", 1.0),
+        lambda_m=getattr(config, "lambda_m", 1.7),
+        lambda_s=getattr(config, "lambda_s", 0.05),
+        project_text_if_needed=getattr(config, "project_text_if_needed", False),
+    )
+    return QTSplusTokenizer(cfg)
+
+
+# ------------------------------
+# Meta classes to build vision/QTS+ towers and preprocessing hook
+# ------------------------------
+class QTSplusMetaModel:
+    def __init__(self, config):
+        super(QTSplusMetaModel, self).__init__(config)
+        self.config = config
+
+        # Vision tower: build early so weights under `model.vision_tower.*` load
+        if hasattr(config, "vision_tower"):
+            self.vision_tower = build_vision_tower(config)
+            try:
+                vt = getattr(self, "vision_tower", None)
+                out_hidden = getattr(getattr(vt, "config", None), "out_hidden_size", None)
+                if isinstance(out_hidden, int) and out_hidden > 0:
+                    self.config.vision_embed_size = out_hidden
+            except Exception:
+                pass
+
+        # QTS+ tower: build early if enabled so parameters exist during load
+        if getattr(self.config, "enable_qts_plus", False) and getattr(self, "qts_plus", None) is None:
+            try:
+                self.qts_plus = build_qts_plus_tower(self.config)
+            except Exception:
+                pass
+
+    def get_qts_plus_tower(self):
+        return getattr(self, "qts_plus", None)
+
+    def get_vision_tower(self):
+        return getattr(self, "vision_tower", None)
+
+
+class QTSplusMetaForCausalLM:
+    def get_model(self):  # pragma: no cover - abstract in practice
+        raise NotImplementedError
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def get_qts_plus_tower(self):
+        return self.get_model().get_qts_plus_tower()
+
+    def encode_visions(self, vision):
+        return self.get_model().get_vision_tower()(vision)
+
+    def prepare_inputs_for_multimodal(
+        self,
+        vision_input,
+        input_ids,
+        position_ids,
+        attention_mask,
+        past_key_values,
+        labels,
+        question_input_ids: Optional[torch.Tensor] = None,
+        video_token_id: Optional[int] = None,
+        mode: str = "train",
+    ):
+        vision_tower = self.get_vision_tower()
+        qts_plus_tower = self.get_qts_plus_tower()
+        text_embed_layer = self.get_model().get_input_embeddings()
+
+        if vision_tower is None or vision_input is None or input_ids.shape[1] == 1:
+            # Match text embedding dtype for scalar placeholders
+            z = torch.tensor(0.0, device=input_ids.device, dtype=text_embed_layer.weight.dtype)
+            return input_ids, position_ids, attention_mask, past_key_values, None, labels, z, z, z
+
+        if self.config.enable_qts_plus:
+            if self.config.vision_tower == "qwen2_5_vl_vision":
+                if isinstance(vision_input, list):
+                    if len(vision_input) == 0:
+                        z = torch.tensor(0.0, device=input_ids.device, dtype=text_embed_layer.weight.dtype)
+                        return input_ids, position_ids, attention_mask, past_key_values, None, labels, z, z, z
+                    vision_input = vision_input[0]
+
+                vision_features = vision_tower.get_video_features(
+                    vision_input["pixel_values_videos"].to(vision_tower.device),
+                    vision_input["video_grid_thw"].to(vision_tower.device),
+                )
+                video_grid_thw = vision_input["video_grid_thw"]
+                if isinstance(vision_features, list) and len(vision_features) > 0:
+                    vision_features = vision_features[0]
+                if vision_features.ndim == 2:
+                    vision_features = vision_features.unsqueeze(0)
+
+                if question_input_ids is None:
+                    raise AssertionError("question_input_ids must be provided in training to avoid data leakage")
+                if question_input_ids.dtype is not torch.long:
+                    question_input_ids = question_input_ids.long()
+
+                text_embeddings = text_embed_layer(question_input_ids)
+                vision_features = vision_features.to(dtype=text_embeddings.dtype)
+
+                qts_plus_out = qts_plus_tower(vision_features, text_embeddings, mode=mode)
+                vision_features = qts_plus_out["Z"]
+                flops_loss = qts_plus_out["add_loss"]["flops"]
+                kv_loss = qts_plus_out["add_loss"]["kv"]
+                smooth_loss = qts_plus_out["add_loss"]["smooth"]
+
+                if video_token_id is None:
+                    video_token_id = getattr(self.config, "video_token_id", None) or 151656
+
+                inputs_embeds, attention_mask, labels = qts_integrate_embeddings(
+                    vision_features=vision_features[0],
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    labels=labels,
+                    video_token_id=video_token_id,
+                    text_model_embed_layer=text_embed_layer,
+                    video_grid_thw=video_grid_thw,
+                )
+                return (
+                    vision_input,
+                    position_ids,
+                    attention_mask,
+                    past_key_values,
+                    inputs_embeds,
+                    labels,
+                    flops_loss,
+                    kv_loss,
+                    smooth_loss,
+                )
+            else:
+                raise ValueError("Not support this model")
+
+        # QTS+ disabled: just embed tokens
+        z = torch.tensor(0.0, device=input_ids.device, dtype=text_embed_layer.weight.dtype)
+        return input_ids, position_ids, attention_mask, past_key_values, None, labels, z, z, z
+
+    def vision_features_count_qtsplus(
+        self,
+        pixel_values_videos: Optional[torch.Tensor],
+        video_grid_thw: Optional[torch.Tensor],
+        question_input_ids: Optional[torch.Tensor],
+    ) -> int:
+        try:
+            if pixel_values_videos is None or video_grid_thw is None or question_input_ids is None:
+                return 0
+            vision_tower = self.get_vision_tower()
+            qts_tower = self.get_qts_plus_tower()
+            text_embed = self.get_model().get_input_embeddings()
+            if vision_tower is None or qts_tower is None or text_embed is None:
+                return 0
+            if question_input_ids.dtype is not torch.long:
+                question_input_ids = question_input_ids.long()
+            try:
+                vt_device = next(vision_tower.parameters()).device
+            except StopIteration:
+                vt_device = text_embed.weight.device
+            vf = vision_tower.get_video_features(
+                pixel_values_videos.to(vt_device),
+                video_grid_thw.to(vt_device),
+            )
+            if isinstance(vf, list) and len(vf) > 0:
+                vf = vf[0]
+            if isinstance(vf, torch.Tensor) and vf.ndim == 2:
+                vf = vf.unsqueeze(0)
+            te = text_embed(question_input_ids.to(text_embed.weight.device))
+            if isinstance(vf, torch.Tensor):
+                vf = vf.to(device=te.device, dtype=te.dtype)
+            with torch.inference_mode():
+                qpo = qts_tower(vf, te, mode="infer")
+            Z = qpo.get("Z")
+            if isinstance(Z, list) and len(Z) > 0:
+                return int(Z[0].shape[0])
+            if isinstance(Z, torch.Tensor):
+                return int(Z.shape[1] if Z.ndim == 3 else Z.shape[0])
+            return 0
+        except Exception:
+            return 0
+
+
+# ------------------------------
+# Base text-only CausalLM for Qwen2.5‑VL (local copy)
+# ------------------------------
+class Qwen2_5_VL_CausalLM_Config(Qwen2_5_VLTextConfig):
+    model_type = "qwen2_5_vl_causal_lm"
+
+
+class Qwen2_5_VLTextForCausalLM(Qwen2_5_VLPreTrainedModel, GenerationMixin):
+    config_class = Qwen2_5_VL_CausalLM_Config
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: Qwen2_5_VL_CausalLM_Config):
+        super().__init__(config)
+        self.model = Qwen2_5_VLTextModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def get_decoder(self):
+        return self.model
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        num_logits_to_keep: int = 0,
+        **loss_kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
+
+        loss = None
+        if labels is not None:
+            # Defer to simple cross-entropy with ignore_index set by caller
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+# ------------------------------
+# QTS+ Qwen2.5‑VL Causal LM (text model + QTS+ + vision)
+# ------------------------------
+class QTSplusQwen2_5_VLModel(QTSplusMetaModel, Qwen2_5_VLTextModel):
+    config_class = QTSplusQwen2_5_VL_CausalLM_Config
+
+    def __init__(self, config: Qwen2_5_VLTextConfig):
+        super(QTSplusQwen2_5_VLModel, self).__init__(config)
+
+
+class QTSplusQwen2_5_VLTextForCausalLM(QTSplusMetaForCausalLM, Qwen2_5_VLTextForCausalLM):
+    config_class = QTSplusQwen2_5_VL_CausalLM_Config
+
+    def __init__(self, config):
+        try:
+            cfg_attn = getattr(config, "attn_implementation", None)
+            if (cfg_attn is None or str(cfg_attn) == "auto") and is_flash_attn_available():
+                setattr(config, "attn_implementation", "flash_attention_2")
+                setattr(config, "_attn_implementation", "flash_attention_2")
+        except Exception:
+            pass
+
+        super(Qwen2_5_VLTextForCausalLM, self).__init__(config)
+        self.model = QTSplusQwen2_5_VLModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+        self,
+        vision_input: Optional[torch.FloatTensor] = None,
+        input_ids: torch.LongTensor = None,
+        labels: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        question_input_ids: Optional[torch.LongTensor] = None,
+        video_token_id: Optional[int] = None,
+    ):
+        if inputs_embeds is not None:
+            input_ids = None
+
+        if inputs_embeds is None:
+            (
+                vision_input,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                inputs_embeds,
+                labels,
+                flops_loss,
+                kv_loss,
+                smooth_loss,
+            ) = self.prepare_inputs_for_multimodal(
+                vision_input,
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                labels,
+                question_input_ids,
+                video_token_id,
+                mode="train" if self.training else "infer",
+            )
+            if inputs_embeds is None:
+                inputs_embeds = self.get_model().embed_tokens(input_ids)
+
+        input_ids = None
+        try:
+            outputs = super().forward(
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                cache_position=cache_position,
+            )
+        except ValueError as error:
+            raise ValueError(
+                f"{error} (input_ids is None: {input_ids is None}, inputs_embeds is None: {inputs_embeds is None})"
+            ) from error
+
+        add_loss = {
+            "flops_loss": flops_loss if vision_input is not None else 0.0,
+            "kv_loss": kv_loss if vision_input is not None else 0.0,
+            "smooth_loss": smooth_loss if vision_input is not None else 0.0,
+        }
+        if labels is None and not self.training:
+            return outputs
+        return (outputs, add_loss)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        vision_input: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.Tensor] = None,
+        question_input_ids: Optional[torch.Tensor] = None,
+        video_token_id: Optional[int] = None,
+        **kwargs,
+    ):
+        position_ids = kwargs.pop("position_ids", None)
+        attention_mask = kwargs.pop("attention_mask", None)
+        if attention_mask is None and input_ids is not None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device)
+        if "inputs_embeds" in kwargs:
+            raise NotImplementedError("`inputs_embeds` is not supported")
+
+        if vision_input is not None:
+            (
+                vision_input,
+                position_ids,
+                attention_mask,
+                _,
+                inputs_embeds,
+                _,
+                *_unused_losses,
+            ) = self.prepare_inputs_for_multimodal(
+                vision_input,
+                input_ids,
+                position_ids,
+                attention_mask,
+                None,
+                None,
+                question_input_ids,
+                video_token_id,
+                mode="infer",
+            )
+        else:
+            inputs_embeds = self.get_model().embed_tokens(input_ids)
+
+        kwargs["attention_mask"] = attention_mask
+        if position_ids is not None:
+            kwargs["position_ids"] = position_ids
+        kwargs.pop("input_ids", None)
+        if "use_cache" not in kwargs:
+            kwargs["use_cache"] = True
+        output_ids = super().generate(inputs_embeds=inputs_embeds, **kwargs)
+        if input_ids is not None:
+            input_ids = input_ids.to(output_ids.device)
+            output_ids = torch.cat([input_ids, output_ids], dim=1)
+        return output_ids
+    
+# Register for Auto* resolution
+AutoModelForCausalLM.register(QTSplusQwen2_5_VL_CausalLM_Config, QTSplusQwen2_5_VLTextForCausalLM)
+__all__ = ["QTSplusQwen2_5_VLTextForCausalLM"]

preprocessor_config.json

@@ -0,0 +1,39 @@
+{
+  "crop_size": null,
+  "data_format": "channels_first",
+  "default_to_square": true,
+  "device": null,
+  "disable_grouping": null,
+  "do_center_crop": null,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_pad": null,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "Qwen2VLImageProcessorFast",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "input_data_format": null,
+  "max_pixels": 12845056,
+  "merge_size": 2,
+  "min_pixels": 3136,
+  "pad_size": null,
+  "patch_size": 14,
+  "processor_class": "Qwen2_5_VLVisionProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "return_tensors": null,
+  "size": {
+    "longest_edge": 12845056,
+    "shortest_edge": 3136
+  },
+  "temporal_patch_size": 2
+}

processing_qts_plus_qwen2_5_vl.py

@@ -0,0 +1,260 @@
+"""
+Self-contained processor shim for trust_remote_code.
+
+Exports `QTSplusQwen2_5_VLProcessor` by aliasing the upstream
+Qwen2.5-VL processor from Transformers. This avoids importing a local
+`src` package while keeping the same class name referenced in
+`processor_config.json`.
+"""
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.video_utils import VideoInput
+from transformers import AutoProcessor
+
+class Qwen2_5_VLVideosProcessorKwargs(VideosKwargs, total=False):
+    fps: Union[list[float], float]
+
+
+class Qwen2_5_VLImagesKwargs(ImagesKwargs):
+    min_pixels: Optional[int]
+    max_pixels: Optional[int]
+    patch_size: Optional[int]
+    temporal_patch_size: Optional[int]
+    merge_size: Optional[int]
+
+
+class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False):
+    images_kwargs: Qwen2_5_VLImagesKwargs
+    videos_kwargs: Qwen2_5_VLVideosProcessorKwargs
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+            "return_mm_token_type_ids": False,
+        },
+    }
+
+
+class QTSplusQwen2_5_VLProcessor(ProcessorMixin):
+    r"""
+    Constructs a Qwen2.5-VL processor which wraps a Qwen2.5-VL image processor and a Qwen2 tokenizer into a single processor.
+    [`Qwen2_5_VLProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the
+    [`~Qwen2_5_VLProcessor.__call__`] and [`~Qwen2_5_VLProcessor.decode`] for more information.
+    Args:
+        image_processor ([`Qwen2VLImageProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`Qwen2TokenizerFast`], *optional*):
+            The tokenizer is a required input.
+        video_processor ([`Qwen2_5_VLVideoProcessor`], *optional*):
+            The video processor is a required input.
+        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+            in a chat into a tokenizable string.
+    """
+
+    attributes = ["image_processor", "tokenizer", "video_processor"]
+
+    image_processor_class = "AutoImageProcessor"
+    video_processor_class = "AutoVideoProcessor"
+    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
+
+    def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs):
+        self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
+        self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
+        self.image_token_id = (
+            tokenizer.image_token_id
+            if getattr(tokenizer, "image_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.image_token)
+        )
+        self.video_token_id = (
+            tokenizer.video_token_id
+            if getattr(tokenizer, "video_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.video_token)
+        )
+        super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
+
+    def __call__(
+        self,
+        images: Optional[ImageInput] = None,
+        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
+        videos: Optional[VideoInput] = None,
+        **kwargs: Unpack[Qwen2_5_VLProcessorKwargs],
+    ) -> BatchFeature:
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwargs` arguments to
+        Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.
+
+        Args:
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            text (`str`, `list[str]`, `list[list[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            videos (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`):
+                The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
+                tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
+            - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
+            - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`.
+            - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`.
+        """
+        output_kwargs = self._merge_kwargs(
+            Qwen2_5_VLProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        image_inputs = videos_inputs = {}
+        if images is not None:
+            image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+            image_grid_thw = image_inputs["image_grid_thw"]
+
+        if videos is not None:
+            fps = output_kwargs["videos_kwargs"].get("fps", 2.0)
+            videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
+            video_grid_thw = videos_inputs["video_grid_thw"]
+
+            if isinstance(fps, (int, float)):
+                second_per_grid_ts = [self.video_processor.temporal_patch_size / fps] * len(video_grid_thw)
+            elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw):
+                second_per_grid_ts = [self.video_processor.temporal_patch_size / tmp for tmp in fps]
+            else:
+                raise ValueError(
+                    f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number."
+                )
+            videos_inputs.update({"second_per_grid_ts": second_per_grid_ts})
+
+        if not isinstance(text, list):
+            text = [text]
+
+        text = text.copy()  # below lines change text in-place
+        if images is not None:
+            merge_length = self.image_processor.merge_size**2
+            index = 0
+            for i in range(len(text)):
+                while self.image_token in text[i]:
+                    num_image_tokens = image_grid_thw[index].prod() // merge_length
+                    text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1)
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+        if videos is not None:
+            merge_length = self.video_processor.merge_size**2
+            index = 0
+            for i in range(len(text)):
+                while self.video_token in text[i]:
+                    num_video_tokens = video_grid_thw[index].prod() // merge_length
+                    text[i] = text[i].replace(self.video_token, "<|placeholder|>" * num_video_tokens, 1)
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.video_token)
+
+        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
+        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+        self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"])
+
+        if return_mm_token_type_ids:
+            array_ids = np.array(text_inputs["input_ids"])
+            mm_token_type_ids = np.zeros_like(text_inputs["input_ids"])
+            mm_token_type_ids[array_ids == self.image_token_id] = 1
+            text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
+
+        return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors)
+
+    def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
+        """
+        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
+        Args:
+            image_sizes (`list[list[int]]`, *optional*):
+                The input sizes formatted as (height, width) per each image.
+            video_sizes (`list[list[int]]`, *optional*):
+                The input sizes formatted as (num_frames, height, width) per each video.
+        Returns:
+            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
+            input modalities, along with other useful data.
+        """
+
+        vision_data = {}
+        if image_sizes is not None:
+            images_kwargs = Qwen2_5_VLProcessorKwargs._defaults.get("images_kwargs", {})
+            images_kwargs.update(kwargs)
+            merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size
+
+            num_image_patches = [
+                self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
+                for image_size in image_sizes
+            ]
+            num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
+            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
+
+        if video_sizes is not None:
+            videos_kwargs = Qwen2_5_VLProcessorKwargs._defaults.get("videos_kwargs", {})
+            videos_kwargs.update(kwargs)
+            num_video_patches = [
+                self.video_processor.get_number_of_video_patches(*video_size, videos_kwargs)
+                for video_size in video_sizes
+            ]
+            num_video_tokens = [(num_patches // merge_size**2) for num_patches in num_video_patches]
+            vision_data["num_video_tokens"] = num_video_tokens
+
+        return MultiModalData(**vision_data)
+
+    def post_process_image_text_to_text(
+        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
+    ):
+        """
+        Post-process the output of the model to decode the text.
+
+        Args:
+            generated_outputs (`torch.Tensor` or `np.ndarray`):
+                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+                or `(sequence_length,)`.
+            skip_special_tokens (`bool`, *optional*, defaults to `True`):
+                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
+            **kwargs:
+                Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+        Returns:
+            `list[str]`: The decoded text.
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+        return names_from_processor + ["second_per_grid_ts"]
+
+AutoProcessor.register("QTSplusQwen2_5_VLProcessor", QTSplusQwen2_5_VLProcessor)
+__all__ = ["QTSplusQwen2_5_VLProcessor"]

processor_config.json

@@ -0,0 +1,9 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_qts_plus_qwen2_5_vl.QTSplusQwen2_5_VLProcessor"
+  },
+  "image_processor_type": "Qwen2VLImageProcessorFast",
+  "processor_class": "QTSplusQwen2_5_VLProcessor",
+  "tokenizer_class": "Qwen2Tokenizer",
+  "video_processor_type": "Qwen2VLVideoProcessor"
+}

special_tokens_map.json

@@ -0,0 +1,20 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|im_end|>",
+  "pad_token": "<|endoftext|>"
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:666ac5168248c6eadbe7be2054c5af600318cdb57d4a94fe25cf5ddcc249e236
+size 11422174

tokenizer_config.json

@@ -0,0 +1,208 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "Qwen2_5_VLVisionProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

video_preprocessor_config.json

@@ -0,0 +1,43 @@
+{
+  "crop_size": null,
+  "data_format": "channels_first",
+  "default_to_square": true,
+  "device": null,
+  "do_center_crop": null,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "do_sample_frames": false,
+  "fps": null,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "input_data_format": null,
+  "max_frames": 768,
+  "max_pixels": 12845056,
+  "merge_size": 2,
+  "min_frames": 4,
+  "min_pixels": 3136,
+  "num_frames": null,
+  "pad_size": null,
+  "patch_size": 14,
+  "processor_class": "Qwen2_5_VLVisionProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "return_metadata": false,
+  "size": {
+    "longest_edge": 12845056,
+    "shortest_edge": 3136
+  },
+  "temporal_patch_size": 2,
+  "video_metadata": null,
+  "video_processor_type": "Qwen2VLVideoProcessor"
+}

zero_to_fp32.py

@@ -0,0 +1,760 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+#   python zero_to_fp32.py . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device, weights_only=False)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+    total_files = len(files)
+    state_dicts = []
+    for f in tqdm(files, desc='Loading checkpoint shards'):
+        state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+    """
+    A pseudo tensor that collects partitioned weights.
+    It is more memory efficient when there are multiple groups.
+    """
+
+    def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+        self.flat_groups = flat_groups
+        self.flat_groups_offset = flat_groups_offset
+        self.offset = offset
+        self.partitioned_numel = partitioned_numel
+        self.shape = shape
+        self.dtype = self.flat_groups[0][0].dtype
+
+    def contiguous(self):
+        """
+        Merge partitioned weights from flat_groups into a single tensor.
+        """
+        end_idx = self.offset + self.partitioned_numel
+        world_size = len(self.flat_groups)
+        pad_flat_param_chunks = []
+
+        for rank_i in range(world_size):
+            # for each rank, we need to collect weights from related group/groups
+            flat_groups_at_rank_i = self.flat_groups[rank_i]
+            start_group_id = None
+            end_group_id = None
+            for group_id in range(len(self.flat_groups_offset)):
+                if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+                    start_group_id = group_id
+                if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+                    end_group_id = group_id
+                    break
+            # collect weights from related group/groups
+            for group_id in range(start_group_id, end_group_id + 1):
+                flat_tensor = flat_groups_at_rank_i[group_id]
+                start_offset = self.offset - self.flat_groups_offset[group_id]
+                end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+                pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+        # collect weights from all ranks
+        pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+        param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+        return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+    for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # memory efficient tensor
+        tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+        state_dict[name] = tensor
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+    """
+    Convert state_dict of GatheredTensor to torch tensor
+    """
+    torch_state_dict = {}
+    converted_tensors = {}
+    for name, tensor in state_dict.items():
+        tensor_id = id(tensor)
+        if tensor_id in converted_tensors:  # shared tensors
+            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+            torch_state_dict[name] = shared_tensor
+        else:
+            converted_tensors[tensor_id] = name
+            if return_empty_tensor:
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+            else:
+                torch_state_dict[name] = tensor.contiguous()
+    return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                             tag=None,
+                                             exclude_frozen_parameters=False,
+                                             lazy_mode=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+          Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+    You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint. Or you can load state_dict in lazy mode ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+        for name, lazy_tensor in state_dict.item():
+            tensor = lazy_tensor.contiguous()  # to cpu
+            print(name, tensor)
+            # del tensor to release memory if it no longer in use
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+    if lazy_mode:
+        return state_dict
+    else:
+        return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_dir``: directory to the pytorch fp32 state_dict output files
+        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                                          tag,
+                                                          exclude_frozen_parameters,
+                                                          lazy_mode=True)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    if max_shard_size is not None:
+        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+        # an memory-efficient approach for sharding
+        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+        state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model by shard
+    os.makedirs(output_dir, exist_ok=True)
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+        shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+        shard_state_dict = to_torch_tensor(shard_state_dict)
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard_state_dict, output_path)
+        # release the memory of current shard
+        for tensor_name in list(shard_state_dict.keys()):
+            del state_dict[tensor_name]
+            del shard_state_dict[tensor_name]
+        del shard_state_dict
+        gc.collect()
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info("Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info("Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument("output_dir",
+                        type=str,
+                        help="directory to the pytorch fp32 state_dict output files"
+                        "(e.g. path/checkpoint-12-output/)")
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="5GB",
+        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+        "without CPU OOM issues.")
+    parser.add_argument(
+        "--safe_serialization",
+        default=False,
+        action='store_true',
+        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+                                               args.output_dir,
+                                               max_shard_size=args.max_shard_size,
+                                               safe_serialization=args.safe_serialization,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)