# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ # ZDTaichu-5.0 — Main Model # # Architecture # ──────────── # Vision encoder : C-RADIOv4-H (ViT-H/16, 653 M) # Projector : RMSNorm → Linear(5120→20480) → SquaredReLU → Linear(20480→H) # LLM decoder : Qwen3.5 hybrid (Gated DeltaNet + full attention, 3:1 ratio) # # Position encoding (M-RoPE) # ────────────────────────── # Vision tokens receive 3D position IDs (temporal, height, width) computed # from the InternVL-style tile grid via ``get_rope_index()``. Text tokens # receive standard 1D positions (all three M-RoPE channels are identical). # # This matches the official Qwen3.5 VL pipeline where ``Qwen3_5Model.forward()`` # calls ``compute_3d_position_ids()`` → ``get_rope_index()`` before forwarding # to ``Qwen3_5TextModel``. The resulting ``position_ids`` of shape ``(3, B, S)`` # are consumed directly by ``Qwen3_5TextRotaryEmbedding``, which applies # interleaved M-RoPE across temporal / height / width frequency bands. # # Generation # ────────── # This model inherits from ``GenerationMixin``, owning the generation loop # (like ``Qwen3_5ForConditionalGeneration``). Key overrides: # - ``_prepare_position_ids_for_generation``: computes 3D ``position_ids`` # on the prefill step and caches ``rope_deltas``; applies ``rope_deltas`` # on subsequent decode steps. # - ``prepare_inputs_for_generation``: clears ``pixel_values`` / # ``pixel_values_videos`` after the first step (vision features are # already embedded in the KV cache). # # Cache handling # ────────────── # ``Qwen3_5DynamicCache`` is created internally by ``Qwen3_5TextModel`` when # ``use_cache=True``. It stores KV states for full-attention layers and # ``conv_states`` + ``recurrent_states`` for Gated DeltaNet layers. # ============================================================================ import itertools import warnings from typing import Any, Dict, List, Optional, Tuple, Union import torch import transformers from torch import nn from torch.nn import CrossEntropyLoss from transformers import AutoModel, GenerationConfig from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from transformers.utils import logging from .configuration import ZDTaichu5_0_Config from .cradio_model import RADIOModel logger = logging.get_logger(__name__) # --------------------------------------------------------------------------- # Import Qwen3.5 model classes — requires transformers >= 5.3.0 # --------------------------------------------------------------------------- _MIN_TRANSFORMERS = "5.3.0" try: from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM from transformers.cache_utils import DynamicCache as Qwen3_5DynamicCache _HAS_QWEN3_5 = True except Exception as e: _HAS_QWEN3_5 = False Qwen3_5ForCausalLM = None Qwen3_5DynamicCache = None logger.warning( f"Could not import Qwen3_5ForCausalLM from transformers. " f"Import error: {e!r}" ) def _version_ge(v1, v2): """Check if version v1 >= v2.""" from packaging import version return version.parse(v1) >= version.parse(v2) # ───────────────────────────────────────────────────────────────────────────── # Projector components # ───────────────────────────────────────────────────────────────────────────── class SquaredReLU(nn.Module): """Squared ReLU activation — same non-linearity used in the projector.""" def forward(self, x): return torch.pow(torch.nn.functional.relu(x), 2) class RMSNorm(nn.Module): """ Standard RMSNorm for the projector (NOT the Qwen3.5 LLM variant). Qwen3.5's internal ``Qwen3_5RMSNorm`` uses zero-initialized weight with ``output * (1 + weight)``. The projector uses ones-initialized weight with ``output * weight`` — the standard formulation. """ def __init__(self, hidden_size: int, eps: float = 1e-5): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.eps) return (self.weight.to(torch.float32) * hidden_states).to(input_dtype) # ───────────────────────────────────────────────────────────────────────────── # Main model # ───────────────────────────────────────────────────────────────────────────── class ZDTaichu5_0_ForConditionalGeneration(PreTrainedModel, GenerationMixin): """ ZDTaichu-5.0: C-RADIOv4-H vision encoder + Qwen3.5 language decoder. Architecture overview:: pixel_values └─► C-RADIOv4-H (ViT-H/16, 653 M) └─► pixel_shuffle(0.5) └─► mlp1: RMSNorm → Linear → SquaredReLU → Linear └─► inject into Qwen3.5 embeddings at positions └─► Qwen3.5 (hybrid DeltaNet / Transformer) """ config_class = ZDTaichu5_0_Config main_input_name = "input_ids" _tied_weights_keys = None#["language_model.lm_head.weight"] _keys_to_ignore_on_load_unexpected = [ # The RADIO input_conditioner registers norm_mean / norm_std as # buffers, but make_preprocessor_external() removes the conditioner # at init time (normalization is handled by the image processor). # The build script still saves these from the source checkpoint, so # they appear as unexpected keys during loading — safe to ignore. r"vision_model\.radio_model\.input_conditioner\..*", r"^mtp\..*", ] _supports_flash_attn_2 = True _supports_flash_attention_2 = True _supports_flash_attn = True _supports_sdpa = True _no_split_modules = ["Qwen3_5DecoderLayer"] _is_stateful = True supports_gradient_checkpointing = True def __init__(self, config: ZDTaichu5_0_Config): super().__init__(config) # Guard for bleeding-edge transformers (>= 4.57.0.dev) where # _finalize_model_loading reads all_tied_weights_keys but # PreTrainedModel.__init__ may not yet initialise it. if not hasattr(self, "all_tied_weights_keys"): self.all_tied_weights_keys = {} assert _version_ge(transformers.__version__, _MIN_TRANSFORMERS), ( f"Qwen3.5 support requires transformers >= {_MIN_TRANSFORMERS} " f"(found {transformers.__version__})" ) assert _HAS_QWEN3_5, ( "Qwen3_5ForCausalLM is not available. " f"Ensure transformers >= {_MIN_TRANSFORMERS} is installed." ) image_size = config.force_image_size patch_size = config.vision_config.patch_size self.patch_size = patch_size self.template = config.template self.num_image_token = int( (image_size // patch_size) ** 2 * (config.downsample_ratio ** 2) ) self.downsample_ratio = config.downsample_ratio self.ps_version = config.ps_version self.image_tag_type = config.image_tag_type self.img_context_token_id = config.img_context_token_id self.video_context_token_id = config.video_context_token_id # Per-tile token dimensions (e.g. 14×14 for 448px, patch=16, ds=0.5) self.tile_h = int((image_size // patch_size) * config.downsample_ratio) self.tile_w = self.tile_h logger.info(f"num_image_token: {self.num_image_token}") logger.info(f"tile_h={self.tile_h}, tile_w={self.tile_w}") logger.info(f"ps_version: {self.ps_version}") logger.info(f"Vision encoder: {config.vision_config.version}") logger.info( f"LLM: Qwen3.5 ({config.llm_config.num_hidden_layers} layers, " f"hidden={config.llm_config.hidden_size}, " f"hybrid=" f"{sum(1 for t in config.llm_config.layer_types if t == 'linear_attention')} linear + " f"{sum(1 for t in config.llm_config.layer_types if t == 'full_attention')} full)" ) # ── Language model ─────────────────────────────────────────────────── self.language_model = Qwen3_5ForCausalLM(config.llm_config) # ── Vision encoder ─────────────────────────────────────────────────── self.vision_model = RADIOModel(config.vision_config) self.vision_model.model._initialize_weights = ( self.vision_model.model._init_weights ) self.vision_model.radio_model.make_preprocessor_external() self.vision_model = self.vision_model.to( self.language_model.config.torch_dtype ) self.drop_vision_class_token = True # ── MLP projector ──────────────────────────────────────────────────── vit_hidden_size = config.vit_hidden_size proj_hidden = config.projector_hidden_size llm_hidden = config.llm_config.hidden_size pixel_shuffle_dim = vit_hidden_size * int(1 / self.downsample_ratio) ** 2 self.mlp1 = nn.Sequential( RMSNorm(pixel_shuffle_dim, eps=1e-5), nn.Linear(pixel_shuffle_dim, proj_hidden, bias=False), SquaredReLU(), nn.Linear(proj_hidden, llm_hidden, bias=False), ) self.mlp1 = self.mlp1.to(self.language_model.config.torch_dtype) # Cached rope_deltas for multi-step generation self.rope_deltas = None # ── Embedding accessors (required by GenerationMixin) ───────────────── def get_input_embeddings(self): return self.language_model.get_input_embeddings() def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) def get_output_embeddings(self): return self.language_model.lm_head def set_output_embeddings(self, new_embeddings): self.language_model.lm_head = new_embeddings def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): # 大头在 LLM:直接委托给内层 Qwen3.5(它原生支持 GC) self.language_model.gradient_checkpointing_enable( gradient_checkpointing_kwargs=gradient_checkpointing_kwargs ) # 视觉塔可选:支持就开,不支持就跳过(不影响主显存) vm = getattr(self, "vision_model", None) if vm is not None and getattr(vm, "supports_gradient_checkpointing", False): try: vm.gradient_checkpointing_enable( gradient_checkpointing_kwargs=gradient_checkpointing_kwargs ) except Exception: pass def gradient_checkpointing_disable(self): self.language_model.gradient_checkpointing_disable() vm = getattr(self, "vision_model", None) if vm is not None and hasattr(vm, "gradient_checkpointing_disable"): try: vm.gradient_checkpointing_disable() except Exception: pass # ───────────────────────────────────────────────────────────────────────── # Vision helpers # ───────────────────────────────────────────────────────────────────────── def pixel_shuffle( self, x: torch.Tensor, scale_factor: float = 0.5 ) -> torch.Tensor: """Space-to-depth rearrangement (ps_version='v2' = corrected layout).""" n, w, h, c = x.size() x = x.view(n, w, int(h * scale_factor), int(c / scale_factor)) x = x.permute(0, 2, 1, 3).contiguous() x = x.view( n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor)), ) if self.ps_version == "v1": warnings.warn( "ps_version='v1' produces a transposed spatial layout. " "Use ps_version='v2' for correct output." ) else: x = x.permute(0, 2, 1, 3).contiguous() return x def extract_feature(self, pixel_values: torch.Tensor) -> torch.Tensor: """Run pixels through C-RADIOv4-H → pixel_shuffle → MLP projector.""" vit_embeds = self.vision_model(pixel_values).features vit_embeds = vit_embeds.to(dtype=torch.bfloat16) h = w = int(vit_embeds.shape[1] ** 0.5) vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1) vit_embeds = self.pixel_shuffle( vit_embeds, scale_factor=self.downsample_ratio ) vit_embeds = vit_embeds.reshape( vit_embeds.shape[0], -1, vit_embeds.shape[-1] ) vit_embeds = self.mlp1(vit_embeds) return vit_embeds # ───────────────────────────────────────────────────────────────────────── # 3D M-RoPE position IDs # ───────────────────────────────────────────────────────────────────────── def get_vision_position_ids( self, start_position: int, tile_rows: int, tile_cols: int, has_thumbnail: bool = True, device: torch.device = None, ) -> torch.LongTensor: """ Compute 3D (temporal, height, width) position IDs for vision tokens from a single InternVL-style tiled image. Token layout (flattened order expected by the model): 1. Grid tiles in raster order: tile(0,0), tile(0,1), …, tile(R-1,C-1). Each tile has ``tile_h × tile_w`` tokens in raster order. 2. Thumbnail tile (optional): a single tile covering the full image at reduced resolution. Args: start_position: Offset added to all positional indices. tile_rows: Number of tile rows in the image grid. tile_cols: Number of tile columns in the image grid. has_thumbnail: Whether a thumbnail tile is appended after grid tiles. device: Target device. Returns: ``torch.LongTensor`` of shape ``(3, num_vision_tokens)``. """ tile_h, tile_w = self.tile_h, self.tile_w npt = tile_h * tile_w # num tokens per tile # ── Grid tiles ─────────────────────────────────────────────────────── num_grid_tiles = tile_rows * tile_cols tile_idx = torch.arange(num_grid_tiles, device=device) tr = tile_idx // tile_cols tc = tile_idx % tile_cols local_idx = torch.arange(npt, device=device) lr = local_idx // tile_w lc = local_idx % tile_w # (num_grid_tiles, npt) → flatten global_h = (tr[:, None] * tile_h + lr[None, :]).reshape(-1).long() global_w = (tc[:, None] * tile_w + lc[None, :]).reshape(-1).long() total_grid = num_grid_tiles * npt pos_t = torch.full( (total_grid,), start_position, device=device, dtype=torch.long ) pos_h = start_position + global_h pos_w = start_position + global_w # ── Thumbnail tile ─────────────────────────────────────────────────── if has_thumbnail: # Map thumbnail local(r, c) → global(r * tile_rows, c * tile_cols) # so its positions overlay the grid at coarser resolution. thumb_h = (lr * tile_rows).long() thumb_w = (lc * tile_cols).long() pos_t = torch.cat([ pos_t, torch.full( (npt,), start_position, device=device, dtype=torch.long ), ]) pos_h = torch.cat([pos_h, start_position + thumb_h]) pos_w = torch.cat([pos_w, start_position + thumb_w]) return torch.stack([pos_t, pos_h, pos_w], dim=0) def get_rope_index( self, input_ids: torch.LongTensor, mm_token_type_ids: torch.IntTensor, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute 3D M-RoPE position IDs for a mixed vision + text sequence. Follows the same structure as ``Qwen3_5Model.get_rope_index``: uses ``mm_token_type_ids`` to group tokens by modality (text=0, image=1, video=2) via ``itertools.groupby``. Vision tokens receive spatial position IDs (temporal, height, width) while text tokens receive standard 1D positions. Args: input_ids: ``(B, S)`` token IDs. mm_token_type_ids: ``(B, S)`` modality labels — 0 = text, 1 = image, 2 = video. image_grid_thw: ``(num_images, 3)`` — each row ``(T=1, tile_rows, tile_cols)`` for InternVL-style tiled images. video_grid_thw: ``(num_videos, 3)`` — each row ``(num_frames, 1, 1)``. attention_mask: ``(B, S)`` binary mask. Returns: ``position_ids``: ``(3, B, S)`` ``mrope_position_deltas``: ``(B, 1)`` """ tile_h, tile_w = self.tile_h, self.tile_w npt = tile_h * tile_w B, S = input_ids.shape device = input_ids.device position_ids = torch.zeros(3, B, S, dtype=input_ids.dtype, device=device) mrope_position_deltas = [] # ------------------------------------------------------------------ # video-as-image compatibility for verl / vLLM rollout. # ------------------------------------------------------------------ if mm_token_type_ids is not None and video_grid_thw is None and torch.any(mm_token_type_ids == 2).item(): mm_token_type_ids = mm_token_type_ids.clone() if image_grid_thw is not None: # Count contiguous visual groups, because get_rope_index consumes # one grid_thw row per contiguous image/video segment. total_visual_groups = 0 for b in range(mm_token_type_ids.shape[0]): cur_types = mm_token_type_ids[b] if attention_mask is not None: cur_types = cur_types[attention_mask[b].bool()] prev_type = None for t in cur_types.tolist(): if t in (1, 2) and t != prev_type: total_visual_groups += 1 prev_type = t num_image_grids = image_grid_thw.shape[0] if total_visual_groups <= num_image_grids: # True video-as-image case: consume image_grid_thw for both image and video types. mm_token_type_ids[mm_token_type_ids == 2] = 1 if "logger" in globals(): logger.warning_once( "Converting mm_token_type_ids type 2 to type 1 because " "video_grid_thw is None and image_grid_thw has enough grids. " "This matches video-as-image processing." ) else: # Some type-2 tokens are likely generated orphan <|video_pad|> tokens. # Treat them as text to avoid consuming non-existent grids. mm_token_type_ids[mm_token_type_ids == 2] = 0 if "logger" in globals(): logger.warning_once( "mm_token_type_ids contains type 2 but video_grid_thw is None, " "and image_grid_thw does not have enough grids. Treating type 2 " "as text. This likely means the model generated orphan <|video_pad|> tokens." ) else: # No visual grid exists, so type 2 cannot represent valid visual tokens. mm_token_type_ids[mm_token_type_ids == 2] = 0 if "logger" in globals(): logger.warning_once( "mm_token_type_ids contains type 2, but both video_grid_thw and " "image_grid_thw are None. Treating type 2 as text." ) grid_iters = { 1: iter(image_grid_thw) if image_grid_thw is not None else None, 2: iter(video_grid_thw) if video_grid_thw is not None else None, } for batch_idx, current_input_ids in enumerate(input_ids): input_token_type = mm_token_type_ids[batch_idx] if attention_mask is not None: current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] input_token_type = input_token_type[attention_mask[batch_idx].bool()] # Group contiguous runs of the same modality type input_type_group = [] for key, group in itertools.groupby( enumerate(input_token_type.tolist()), lambda x: x[1] ): group = list(group) start_index = group[0][0] end_index = group[-1][0] + 1 input_type_group.append((key, start_index, end_index)) current_pos = 0 llm_pos_ids_list: List[torch.Tensor] = [] # ── Per-video state machine ────────────────────────────────────── # Mirrors the Megatron-side implementation in # modeling.py: a single video_grid_thw entry # of [num_frames, 1, 1] is consumed across multiple non-contiguous # type-2 runs (one per <|video_pad|> block, separated by frame # header text). # # Within a video, every frame's tokens use: # t = vid_spatial_start + frame_idx (anchored at video start) # h = vid_spatial_start + local_row (constant across frames) # w = vid_spatial_start + local_col (constant across frames) # # Text between frames advances ``current_pos`` normally — those # text positions live in a different range than the video frame # positions, which is fine for M-RoPE (RoPE requires no # monotonicity, only consistent training/inference). vid_active = False vid_num_frames = 0 vid_frame_idx = 0 vid_spatial_start = 0 for modality_type, start_idx, end_idx in input_type_group: # text == 0 if modality_type == 0: text_len = end_idx - start_idx llm_pos_ids_list.append( torch.arange(text_len, device=device).view(1, -1).expand(3, -1) + current_pos ) current_pos += text_len # image == 1 elif modality_type == 1: seg_len = end_idx - start_idx grid = next(grid_iters[1]) tile_rows = grid[1].item() tile_cols = grid[2].item() grid_tokens = tile_rows * tile_cols * npt has_thumbnail = seg_len > grid_tokens vpos = self.get_vision_position_ids( start_position=current_pos, tile_rows=tile_rows, tile_cols=tile_cols, has_thumbnail=has_thumbnail, device=device, ) assert vpos.shape[1] == seg_len, ( f"Position count ({vpos.shape[1]}) ≠ image token count " f"({seg_len}) for grid=({tile_rows},{tile_cols}), " f"thumbnail={has_thumbnail}" ) llm_pos_ids_list.append(vpos) current_pos += max(tile_rows * tile_h, tile_cols * tile_w) # video == 2 elif modality_type == 2: seg_len = end_idx - start_idx # Activate per-video state on the FIRST type-2 run for # this video. Subsequent type-2 runs (one per frame # block, separated by frame-header text) reuse the same # vid_spatial_start anchor. if not vid_active: grid = next(grid_iters[2]) vid_num_frames = grid[0].item() vid_active = True vid_frame_idx = 0 vid_spatial_start = current_pos # Each frame contributes exactly ``npt`` tokens. if seg_len % npt != 0: raise ValueError( f"Video segment length {seg_len} is not a " f"multiple of npt={npt} (tile_h*tile_w). " f"Check that the processor produced one " f"<|video_pad|> block per frame with exactly " f"npt tokens each." ) frames_in_run = seg_len // npt # Sanity guard against malformed grids — never consume # more frames than the grid declared. if vid_frame_idx + frames_in_run > vid_num_frames: raise ValueError( f"Video has {vid_num_frames} frames but " f"input_ids contain at least " f"{vid_frame_idx + frames_in_run} frame blocks. " f"Check the processor's video_grid_thw against " f"the actual <|video_pad|> count." ) local_idx = torch.arange(npt, device=device) lr = local_idx // tile_w lc = local_idx % tile_w all_t, all_h, all_w = [], [], [] for _ in range(frames_in_run): # Temporal: anchored at video_start, advances by frame_idx. all_t.append(torch.full( (npt,), vid_spatial_start + vid_frame_idx, device=device, dtype=torch.long, )) # Spatial: constant base across frames within this video. all_h.append((vid_spatial_start + lr).long()) all_w.append((vid_spatial_start + lc).long()) vid_frame_idx += 1 # Advance current_pos by one frame's spatial extent so # subsequent text positions stay strictly above any # h/w position used by this video. After all frames, # current_pos has advanced by num_frames * max(tile_h, tile_w), # which always exceeds vid_spatial_start + max(num_frames, tile_h, tile_w) # for num_frames >= 1 (so text after the video sees # positions strictly greater than every video token). current_pos += max(tile_h, tile_w) vpos = torch.stack([ torch.cat(all_t), torch.cat(all_h), torch.cat(all_w), ], dim=0) assert vpos.shape[1] == seg_len, ( f"Position count ({vpos.shape[1]}) ≠ video token " f"count ({seg_len})" ) llm_pos_ids_list.append(vpos) # End the video once all declared frames have been # consumed; reset state so the next video (if any) gets # a fresh grid pull. if vid_frame_idx >= vid_num_frames: vid_active = False vid_num_frames = 0 vid_frame_idx = 0 vid_spatial_start = 0 # Sanity check: if a video's last frame isn't followed by any text, # the loop ends with vid_active=False (we already reset on the # final frame). But if the input is malformed and the type-2 # runs don't cover all declared frames, surface that loudly # rather than silently advancing the iterator the next time we # see another video. if vid_active: raise ValueError( f"Reached end of input with video state still active: " f"consumed {vid_frame_idx}/{vid_num_frames} frames. " f"video_grid_thw declares more frames than the " f"<|video_pad|> blocks contain." ) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) if attention_mask is not None: position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = ( llm_positions.to(position_ids.device) ) else: position_ids[:, batch_idx] = llm_positions.to(position_ids.device) mrope_position_deltas.append( llm_positions.max() + 1 - len(current_input_ids) ) mrope_position_deltas = torch.tensor( mrope_position_deltas, device=device ).unsqueeze(1) return position_ids, mrope_position_deltas def _build_text_position_ids( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, ) -> torch.LongTensor: """ Build text position ids of shape (B, S). For padding mask, positions are 0,1,2,... on valid tokens. Padding positions stay 0. """ batch_size, seq_len = input_ids.shape device = input_ids.device if attention_mask is not None: valid = attention_mask > 0 text_position_ids = valid.long().cumsum(-1) - 1 text_position_ids = text_position_ids.masked_fill(~valid, 0) else: text_position_ids = torch.arange( seq_len, device=device, dtype=torch.long ).unsqueeze(0).expand(batch_size, -1) return text_position_ids.contiguous() def _prepend_text_position_channel( self, input_ids: torch.LongTensor, vision_position_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, ) -> torch.LongTensor: """ Convert vision M-RoPE position ids from (3, B, S) to Qwen3.5-compatible position ids of shape (4, B, S): channel 0 : text positions, used for causal mask / FA2 varlen logic channel 1-3 : temporal / height / width vision M-RoPE positions """ if vision_position_ids is None: return None if vision_position_ids.dim() == 3 and vision_position_ids.shape[0] == 4: return vision_position_ids.contiguous() assert vision_position_ids.dim() == 3 and vision_position_ids.shape[0] == 3, ( f"Expected vision_position_ids shape (3, B, S), got " f"{tuple(vision_position_ids.shape)}" ) text_position_ids = self._build_text_position_ids( input_ids=input_ids, attention_mask=attention_mask, ).to(device=vision_position_ids.device) position_ids = torch.cat( [ text_position_ids.unsqueeze(0), # (1, B, S) vision_position_ids, # (3, B, S) ], dim=0, ) return position_ids.contiguous() def _compute_position_ids( self, input_ids: Optional[torch.LongTensor], inputs_embeds: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor], video_grid_thw: Optional[torch.LongTensor], attention_mask: Optional[torch.Tensor], past_key_values=None, mm_token_type_ids: Optional[torch.IntTensor] = None, use_cache: Optional[bool] = None, ) -> Optional[torch.Tensor]: """ Mirror of ``Qwen3_5Model.compute_3d_position_ids``. - Vision info available + first forward → ``get_rope_index``, cache ``rope_deltas``. - ``rope_deltas`` cached (decode step) → derive from attention_mask + ``rope_deltas``. - Pure text → return ``None`` (``Qwen3_5TextModel`` auto-generates). """ past_length = 0 if past_key_values is not None: past_length = past_key_values.get_seq_length() can_compute = ( input_ids is not None and mm_token_type_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) ) if can_compute and past_length == 0: vision_position_ids, rope_deltas = self.get_rope_index( input_ids, mm_token_type_ids=mm_token_type_ids, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, ) # Training / log-prob forward should not keep rope_deltas across batches. # Generation prefill can keep it for decode. if use_cache: self.rope_deltas = rope_deltas else: self.rope_deltas = None return self._prepend_text_position_channel( input_ids=input_ids, vision_position_ids=vision_position_ids, attention_mask=attention_mask, ) elif self.rope_deltas is not None and past_length != 0: batch_size, seq_length = inputs_embeds.shape[:2] if attention_mask is not None: text_position_ids = attention_mask.long().cumsum(-1) - 1 text_position_ids = text_position_ids.masked_fill(attention_mask == 0, 0) text_position_ids = text_position_ids[:, -seq_length:] else: text_position_ids = torch.arange( past_length, past_length + seq_length, device=inputs_embeds.device, dtype=torch.long, ).unsqueeze(0).expand(batch_size, -1) delta = self.rope_deltas.repeat_interleave( batch_size // self.rope_deltas.shape[0], dim=0 ).to(device=inputs_embeds.device) # Decode step follows generation convention: (1, B, S) position_ids = text_position_ids.unsqueeze(0) + delta.view(1, batch_size, 1) return position_ids.contiguous() return None # ───────────────────────────────────────────────────────────────────────── # Forward # ───────────────────────────────────────────────────────────────────────── def forward( self, input_ids: torch.LongTensor = None, pixel_values: Optional[torch.FloatTensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, num_patches = None, image_flags: Optional[torch.LongTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, mm_token_type_ids: Optional[torch.IntTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values=None, labels: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: """ Forward pass for training and generation steps. Args: input_ids: ``(B, S)`` token IDs. pixel_values: ``(total_tiles, C, H, W)`` image tiles from C-RADIOv4-H. pixel_values_videos: ``(total_frames, C, H, W)`` video frames. image_flags: ``(B, max_tiles)`` — 1 for real tiles, 0 for padding. image_grid_thw: ``(num_images, 3)`` — ``(T=1, tile_rows, tile_cols)`` per image. Required for correct M-RoPE spatial positions. video_grid_thw: ``(num_videos, 3)`` — ``(num_frames, 1, 1)`` per video. mm_token_type_ids: ``(B, S)`` modality labels — 0 = text, 1 = image, 2 = video. Required for computing 3D M-RoPE position IDs. Produced by the processor. attention_mask: ``(B, S)`` binary mask. Must be 2-D; the ``Qwen3_5TextModel`` internally creates the 4-D causal mask for full-attention layers and the 2-D mask for DeltaNet layers. position_ids: ``(3, B, S)`` or ``None``. If ``None`` and vision tokens are present, computed via ``get_rope_index()``. """ return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) # ── Embed tokens ───────────────────────────────────────────────────── if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) # ── Inject image features ──────────────────────────────────────────── if pixel_values is not None: if image_flags is None: image_flags = torch.ones( pixel_values.shape[0], dtype=torch.long, device=pixel_values.device, ) image_flags_sq = image_flags.squeeze(-1) vit_embeds = self.extract_feature(pixel_values) vit_embeds = vit_embeds[image_flags_sq == 1] del pixel_values B, N, C = inputs_embeds.shape flat = inputs_embeds.reshape(B * N, C) ids_flat = input_ids.reshape(B * N) selected = ids_flat == self.img_context_token_id try: flat[selected] = flat[selected] * 0.0 + vit_embeds.reshape(-1, C) except Exception as e: vit_flat = vit_embeds.reshape(-1, C) logger.warning( f"Image injection shape mismatch: {e}. " f"selected={selected.sum()}, vit={vit_flat.shape}" ) n_tok = selected.sum() flat[selected] = flat[selected] * 0.0 + vit_flat[:n_tok] del vit_embeds inputs_embeds = flat.reshape(B, N, C) # ── Inject video features ──────────────────────────────────────────── if pixel_values_videos is not None: video_vit = self.extract_feature(pixel_values_videos) del pixel_values_videos B, N, C = inputs_embeds.shape flat = inputs_embeds.reshape(B * N, C) ids_flat = input_ids.reshape(B * N) vmask = ids_flat == self.video_context_token_id flat[vmask] = ( flat[vmask] * 0.0 + video_vit.reshape(-1, C).to(flat.device, flat.dtype) ) inputs_embeds = flat.reshape(B, N, C) del video_vit # GRPO actor/ref training and log-prob computation should not use cache. if labels is not None: use_cache = False self.rope_deltas = None # ── 3D position IDs ────────────────────────────────────────────────── if position_ids is None: position_ids = self._compute_position_ids( input_ids=input_ids, inputs_embeds=inputs_embeds, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, past_key_values=past_key_values, mm_token_type_ids=mm_token_type_ids, use_cache=use_cache, ) if position_ids is not None: position_ids = position_ids.contiguous() # ── LLM forward ───────────────────────────────────────────────────── outputs = self.language_model( input_ids=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) logits = outputs.logits loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = CrossEntropyLoss() shift_logits = shift_logits.view( -1, self.language_model.config.vocab_size ) shift_labels = shift_labels.view(-1).to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) 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, ) # ───────────────────────────────────────────────────────────────────────── # GenerationMixin overrides # ───────────────────────────────────────────────────────────────────────── def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, position_ids=None, use_cache=True, pixel_values=None, pixel_values_videos=None, num_patches=None, image_flags=None, image_grid_thw=None, video_grid_thw=None, mm_token_type_ids=None, is_first_iteration=False, **kwargs, ): """ Prepare inputs for each generation step. After the first iteration, ``pixel_values`` / ``pixel_values_videos`` are cleared because vision features are already in the KV cache. """ model_inputs = super().prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, cache_position=cache_position, position_ids=position_ids, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, num_patches=num_patches, image_flags=image_flags, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, mm_token_type_ids=mm_token_type_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs, ) if not is_first_iteration and use_cache: model_inputs["pixel_values"] = None model_inputs["pixel_values_videos"] = None return model_inputs def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): """ Override to compute 3D M-RoPE position IDs during generation. Mirrors ``Qwen3_5ForConditionalGeneration._prepare_position_ids_for_generation``: - Prefill step: compute 3D positions via ``get_rope_index``, cache ``rope_deltas``. - Decode steps: apply cached ``rope_deltas`` to sequential text positions. Returns position_ids of shape ``(4, B, S)`` on the prefill step (text + 3D vision channels) or ``(1, B, S)`` on decode steps (text + rope_deltas). When ``Qwen3_5TextModel`` receives ``shape[0]==4``, it splits into ``text_position_ids = [0]`` (for causal mask) and ``position_ids = [1:]`` (for rotary embedding). When ``shape[0]!=4``, it sets ``text_position_ids=None``. """ text_positions = super()._prepare_position_ids_for_generation( inputs_tensor, model_kwargs ) # Decode step — apply rope_deltas past_length = 0 cache = model_kwargs.get("past_key_values") if cache is not None: past_length = cache.get_seq_length() if past_length != 0 and self.rope_deltas is not None: position_ids = text_positions[None, ...] + self.rope_deltas return position_ids # Prefill step — compute 3D vision positions if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: inputs_tensor = model_kwargs["input_ids"] is_input_ids = ( len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] ) has_vision = ( model_kwargs.get("mm_token_type_ids") is not None and ( model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None ) ) if is_input_ids and has_vision: vision_positions, rope_deltas = self.get_rope_index( inputs_tensor, mm_token_type_ids=model_kwargs.get("mm_token_type_ids"), image_grid_thw=model_kwargs.get("image_grid_thw"), video_grid_thw=model_kwargs.get("video_grid_thw"), attention_mask=model_kwargs.get("attention_mask"), ) self.rope_deltas = rope_deltas else: vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1) self.rope_deltas = torch.zeros( inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device, ) # Concatenate text + vision → (4, B, S) # Channel 0 = text positions → used by create_causal_mask # Channels 1-3 = vision positions → used by rotary embedding # This matches Qwen3_5ForConditionalGeneration's convention. text_positions = text_positions[None, ...] # (1, B, S) position_ids = torch.cat( [text_positions, vision_positions], dim=0 ) # (4, B, S) #print(f"{position_ids.permute(1, 2, 0).cpu().tolist()}") return position_ids