feat: update modeling code for Valley3

modeling_valley.py

@@ -1,3 +1,643 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cf1b9862bfbdbfca512ab83ca5bf16bbdbdb1b6c125a99751603c50ee15cab96
-size 35176
+#  Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 
+#  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.
+
+import torch
+import numpy as np
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from abc import ABC, abstractmethod
+from typing import List, Optional, Tuple, Union, Dict, Any
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers import AutoConfig, AutoModelForCausalLM, Qwen3Config, Qwen3ForCausalLM, Qwen3Model
+
+from .modeling_vision_tower import build_vision_tower
+from .modeling_projector import build_vision_projector
+from .utils import get_anyres_image_grid_shape, unpad_image, IGNORE_INDEX, IMAGE_TOKEN_INDEX, IMAGE_INDICATOR_IDS, IMAGE_ATOM_ID, pad_truncate_sequence
+
+
+class ValleyConfig(Qwen3Config):
+    model_type = "valley"
+
+class ValleyMetaModel:
+    def __init__(self, config):
+        super(ValleyMetaModel, self).__init__(config)
+        # Build vision tower
+        if hasattr(config, "mm_vision_tower"):
+            if getattr(config, "eagle_vision_tower", None) is not None:
+                self.vision_tower, self.qwen2vl_vision_tower = build_vision_tower(config, delay_load=False)
+            else:
+                self.vision_tower = build_vision_tower(config, delay_load=False)
+        # Build Projector
+        if hasattr(config, "mm_projector_type") and not getattr(config, "only_navit", False):
+            self.mm_projector = build_vision_projector(config)
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, "vision_tower", None)
+        if getattr(self.config, "eagle_vision_tower", None) is not None:
+            qwen2vl_vision_tower = getattr(self, "qwen2vl_vision_tower", None)
+            return vision_tower, qwen2vl_vision_tower
+        else:
+            return vision_tower
+
+class ValleyMetaForCausalLM(ABC):
+    @abstractmethod
+    def get_model(self):
+        pass
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def split_by_instance(self, original_list, split_sizes):
+        start = 0
+        sub_lists = []
+        for size in split_sizes:
+            end = start + size
+            sub_list = original_list[start:end]
+            sub_lists.append([x.to(self.device) for x in sub_list])
+            start = end
+        return sub_lists
+    
+    def encode_images_qwen2vl(self, pixel_values = None, grid_thw = None, split_sizes=None):
+        _, qwen2vl_vision_tower = self.get_model().get_vision_tower()
+        qwen2vl_image_features = qwen2vl_vision_tower(pixel_values, grid_thw)
+        qwen2vl_image_split_sizes = torch.prod(grid_thw[:, 1:3]//2, dim=1)
+        qwen2vl_image_features = torch.split(qwen2vl_image_features, qwen2vl_image_split_sizes.tolist(), dim=0)
+        qwen2vl_image_features = self.split_by_instance(qwen2vl_image_features, split_sizes)
+        return qwen2vl_image_features
+
+    def encode_images(self, images = None, split_sizes = None):
+        """
+        images: (if not anyres) images.shape = [n,3,336,336] , n = number of images + (number of video) * 8
+        images: (if anyres) images.shape = [n,3,336,336] , n = number of tiles * number of images
+        """
+        if getattr(self.config, "eagle_vision_tower", None) is not None:
+            siglip_vision_tower, _ = self.get_model().get_vision_tower()
+            image_features = siglip_vision_tower(images)
+            image_features = self.get_model().mm_projector(image_features)
+        else:
+            image_features = self.get_model().get_vision_tower()(images)
+            image_features = self.get_model().mm_projector(image_features)
+
+        if getattr(self.config,'anyres', False) and getattr(self.config, 'max_vision_token', None) is not None:
+            assert split_sizes is not None
+            image_features = list(torch.split(image_features, split_sizes, dim=0))
+            for i, image_feature in enumerate(image_features):
+                hidden_dim = image_feature.shape[-1]
+                image_tokens = image_feature.shape[0]*image_feature.shape[1]
+                if getattr(self.config, "eagle_vision_tower", None) is not None:
+                    pass # the max_vision_token will be processed in the unpad image token part
+                else:
+                    if image_tokens > self.config.max_vision_token:
+                        intput_shape = int((image_feature.shape[1])**0.5)
+                        output_shape = int((self.config.max_vision_token/image_feature.shape[0])**0.5)
+                        image_feature = image_feature.view(image_feature.shape[0],intput_shape, intput_shape, -1).permute(0,3,1,2)
+                        m = nn.AdaptiveAvgPool2d(output_shape) # different from roi pooling, but in square image, it seems the same
+                        pooling_feature = m(image_feature).permute(0,2,3,1)
+                        image_features[i] = pooling_feature.view(image_feature.shape[0], -1, hidden_dim)
+                split_sizes = None # have already split, set the flag 
+
+        if getattr(self.config, 'mm_use_im_start_end', False):
+            raise ValueError('mm_use_im_start is not support')
+        if split_sizes is not None:
+            image_features = torch.split(image_features, split_sizes, dim=0)
+        
+        return image_features
+
+    def get_padding_method(self):
+        right_padding = getattr(self, 'right_padding', None)
+        # if right_padding flag is setted, ignore training flag. 
+        if right_padding is not None:
+            method = 'right' if right_padding else 'left'
+        # in the other way, use training flag to determine the padding method.
+        method = 'right' if self.training else 'left'
+
+        return method
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, position_ids, attention_mask, past_key_values, labels, images,
+        image_sizes, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw):
+
+        vision_tower = self.get_vision_tower()
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1:
+                target_shape = past_key_values[-1][-1].shape[-2] + 1
+                attention_mask = torch.cat((attention_mask, torch.ones(
+                    (attention_mask.shape[0], target_shape - attention_mask.shape[1]),
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device
+                )), dim=1)
+            return input_ids, position_ids, attention_mask, past_key_values, None, labels
+
+        # Step1: Get image embedings
+        if type(images) is list or images.ndim == 5:
+            # Without slicing the image
+            if not getattr(self.config,'anyres', False) and self.config.mm_projector_type != "ovis2_adapter":
+                concat_images = torch.cat([image for image in images], dim=0) # to do batch compute
+                split_sizes = [image.shape[0] for image in images]
+                
+                # Get vision tower feature, check whether only use navit firstly
+                if getattr(self.config, 'eagle_vision_tower', None) is not None and getattr(self.config, 'only_navit', False):
+                    image_features = None
+                else:
+                    image_features = self.encode_images(concat_images, split_sizes)
+                    image_features = [x.to(self.device) for x in image_features]
+                
+                # Get Eagle features
+                if getattr(self.config, 'eagle_vision_tower', None) is not None:
+                    if pixel_values is not None:
+                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values, image_grid_thw, split_sizes)
+                    elif pixel_values_videos is not None:
+                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values_videos, video_grid_thw, split_sizes)
+                    else:
+                        qwen2vl_image_features = None
+
+            # Slicing the image, each image contains some sub_images:
+            # images = [
+            #   [image1_tiles(n1,3,336,336), image2_tiles(n2,3,336,336), ...],
+            #   [image1_tiles(n1,3,336,336), image2_tiles(n2,3,336,336), ...], ...
+            # ]
+            else:
+                split_sizes = [len(image) for image in images]
+                # Get Eagle features
+                if getattr(self.config, "eagle_vision_tower", None) is not None:
+                    if pixel_values is not None:
+                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values, image_grid_thw, split_sizes)
+                    elif pixel_values_videos is not None:
+                        qwen2vl_image_features = self.encode_images_qwen2vl(pixel_values_videos, video_grid_thw, split_sizes)
+                    else:
+                        qwen2vl_image_features = None
+                
+                # Get vision tower feature, check whether only use navit firstly
+                if getattr(self.config, 'eagle_vision_tower', None) is not None and getattr(self.config, 'only_navit', False):
+                    image_features = None
+                else:
+                    image_features = []
+                    all_concat_images = []
+                    all_split_sizes = []
+                    for batch_images in images:
+                        concat_images = torch.cat([image for image in batch_images], dim=0) # to do batch compute
+                        split_sizes = [image.shape[0] for image in batch_images] 
+                        all_concat_images.append(concat_images)
+                        all_split_sizes.append(split_sizes)
+                    all_image_features = self.encode_images(images=torch.cat(all_concat_images, dim=0), split_sizes=sum(all_split_sizes, []))
+
+                    idx = 0
+                    for split_sizes in all_split_sizes:
+                        batch_image_features = all_image_features[idx:idx+len(split_sizes)]
+                        idx += len(split_sizes)
+                        if type(batch_image_features[0]) is list:
+                            batch_image_features = [torch.cat(x).to(self.device) for x in batch_image_features]
+                        else:
+                            batch_image_features = [x.view(-1,x.shape[-1]).to(self.device) for x in batch_image_features] # tiles feature need to flatten in token dimention, [n_tiles, T, d] -> [n_tiles * T, d]
+                        image_features.append(batch_image_features)
+
+                if getattr(self.config, "eagle_vision_tower", None) is not None and getattr(self.config, 'only_navit', False) == False:
+                    # unpad image tokens
+                    height = width = self.config.num_patches_per_side
+                    new_image_features = []
+                    for batch_image_features, batch_image_sizes in zip(image_features, image_sizes):
+                        batch_image_features_list = []
+                        for cur_image_feature, cur_image_size in zip(batch_image_features, batch_image_sizes):
+                            base_image_feature = cur_image_feature[:width*height, :]
+                            image_feature = cur_image_feature[width*height:, :]
+                            if image_feature.shape[0] != 0:
+                                num_patch_width, num_patch_height = get_anyres_image_grid_shape(
+                                    cur_image_size,
+                                    self.config.grid_pinpoints,
+                                    self.config.vit_crop_size
+                                )
+                                image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1) # (num_patch_H, num_patch_W, H, W, C)
+                                image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() # (C, num_patch_H, H, num_patch_W, W)
+                                image_feature = image_feature.flatten(1, 2).flatten(2, 3) # (C, num_token_H, num_token_W)
+                                image_feature = unpad_image(image_feature, cur_image_size) # (C, num_token_H_unpad, num_token_W_unpad)
+                                input_shape = (image_feature.shape[-2], image_feature.shape[-1])
+                                subimage_tokens = np.prod(input_shape)
+                                
+                                # adaptive avg 2d pool for reducing token num
+                                max_subimage_tokens = self.config.max_vision_token-width*height
+                                if subimage_tokens > max_subimage_tokens:
+                                    aspect_ratio = input_shape[0] / input_shape[1]
+                                    output_shape = (
+                                        int((max_subimage_tokens/aspect_ratio)**0.5*aspect_ratio),
+                                        int((max_subimage_tokens/aspect_ratio)**0.5)
+                                    )
+                                    m = nn.AdaptiveAvgPool2d(output_shape)
+                                    image_feature = m(image_feature)
+                                image_feature = image_feature.flatten(1, 2).transpose(0, 1)
+                                image_feature = torch.cat((base_image_feature, image_feature), dim=0)
+                            else:
+                                image_feature = cur_image_feature
+                            batch_image_features_list.append(image_feature)
+                        new_image_features.append(batch_image_features_list)
+
+                    image_features = new_image_features
+
+        else:
+            image_features = self.encode_images(images).to(self.device)
+
+
+        # Step2: Iterate through each sample in the batch, insert image embedings into input_embeds
+        #        and filling labels, attention mask at the same time. Finally, get `new_input_embed`,
+        #        `new_labels`, new_attention_mask`.
+        _labels = labels
+        _position_ids = position_ids
+        _attention_mask = attention_mask
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        if position_ids is None:
+            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
+        if labels is None:
+            labels = torch.full_like(input_ids, IGNORE_INDEX)
+
+        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask.bool())]
+        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask.bool())]
+        attention_mask = [cur_attention_mask[cur_attention_mask.bool()] for cur_attention_mask in attention_mask]
+        
+        if self.config.mm_projector_type == "ovis2_adapter":  # for ovis2
+            # prepare embedding
+            visual_vocab_size = self.config.mlp_hidden_dim
+            assert visual_vocab_size == 65536
+            text_embedding = self.model.get_input_embeddings()
+            visual_indicator_embedding = self.model.mm_projector.embedding(
+                torch.tensor(
+                    list(range(visual_vocab_size - 5, visual_vocab_size)),
+                    dtype=torch.long,
+                    device=input_ids[0].device 
+                )
+            ).to(device=input_ids[0].device)
+
+            new_attention_masks = [] 
+            new_input_embeds = []
+            new_labels = []
+            for i, cur_image_features in enumerate(image_features):
+                input_id = input_ids[i]
+                text_label = labels[i]
+
+
+                ovis_image_features = []
+                for feature in cur_image_features:
+                    ovis_image_features.append(feature)
+                ovis_image_features = torch.cat(ovis_image_features, dim=0)
+
+                placeholder_token_mask = torch.lt(input_id, 0)
+                text_embed = text_embedding(torch.masked_fill(input_id, placeholder_token_mask, 0))
+
+
+                for j, indicator_id in enumerate(IMAGE_INDICATOR_IDS):
+                    text_embed[input_id == indicator_id] = visual_indicator_embedding[j]
+                image_atom_positions = torch.where(torch.eq(input_id, IMAGE_ATOM_ID))[0].tolist()
+
+                input_embed_parts = []
+                attention_mask_parts = []
+                label_parts = []
+                prev_image_atom_position = -1
+                # assert ovis_image_features.shape[0] == 256*len(image_atom_positions)
+                image_token_len = ovis_image_features.shape[0] // len(image_atom_positions)
+                # print("image_token_len: ", image_token_len)
+                assert image_token_len in [256, 64]
+                if len(image_atom_positions) > 0:
+                    for index, image_atom_positions in enumerate(image_atom_positions):
+                        input_embed_parts.append(text_embed[prev_image_atom_position + 1:image_atom_positions, :])
+                        input_embed_parts.append(ovis_image_features[index*image_token_len:index*image_token_len+image_token_len])  # replace 256
+                        label_parts.append(text_label[prev_image_atom_position + 1:image_atom_positions])
+                        label_parts.append(torch.full((image_token_len,), IGNORE_INDEX, dtype=torch.long, device=input_id.device))  # remain hypothesis
+                        attention_mask_parts.append(
+                            torch.ones_like(text_label[prev_image_atom_position + 1:image_atom_positions], dtype=torch.bool, device=input_id.device))
+                        attention_mask_parts.append(
+                            torch.ones(image_token_len, dtype=torch.bool))
+                        
+                        prev_image_atom_position = image_atom_positions
+                    if prev_image_atom_position + 1 < input_id.shape[0]:
+                        input_embed_parts.append(text_embed[prev_image_atom_position + 1:, :])
+                        label_parts.append(text_label[prev_image_atom_position + 1:])
+                        attention_mask_parts.append(    
+                            torch.ones_like(text_label[prev_image_atom_position + 1:], dtype=torch.bool))
+                        
+                    input_embed = torch.cat([part.to(input_id.device) for part in input_embed_parts], dim=0)
+                    attention_mask = torch.cat([part.to(input_id.device) for part in attention_mask_parts], dim=0)
+                    label = torch.cat([part.to(input_id.device) for part in label_parts], dim=0)
+
+                    new_input_embeds.append(input_embed)
+                    new_attention_masks.append(attention_mask)
+                    new_labels.append(label)
+                else:
+                    raise ValueError(
+                        "No image token found in the input. Please check the input_ids and image_features.")
+            multimodal_max_length = 0
+            left_padding = True
+            new_input_embeds = pad_truncate_sequence(multimodal_max_length, new_input_embeds, batch_first=True, padding_value=0.0, left_padding=left_padding)
+            new_attention_masks = pad_truncate_sequence(multimodal_max_length, new_attention_masks, batch_first=True, padding_value=False, left_padding=left_padding)
+            new_labels = pad_truncate_sequence(multimodal_max_length, new_labels, batch_first=True, padding_value=IGNORE_INDEX, left_padding=left_padding)
+            return None, None, new_attention_masks, None, new_input_embeds, new_labels
+
+        else:
+            new_input_embeds = []
+            new_labels = []
+            new_attention_mask = []
+            
+            for batch_idx, cur_input_ids in enumerate(input_ids):
+                cur_batch_image_idx = 0
+                num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()
+
+                # Step2-1: If this piece of data is pure text, then concat a dummy image to ensure the whole compute graph is same on all device
+                if num_images == 0: 
+                    if getattr(self.config, "eagle_vision_tower", None) is not None:
+                        if getattr(self.config, 'only_navit', False):
+                            cur_image_features = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
+                        else:
+                            siglip_feat = image_features[batch_idx][cur_batch_image_idx]
+                            try:
+                                qwen2vl_feat = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
+                                cur_image_features = torch.cat((siglip_feat, qwen2vl_feat), dim=0)
+                            except Exception as e:
+                                print(e)
+                                print("only siglip feature:", siglip_feat.shape)
+                                cur_image_features = siglip_feat
+                    else:
+                        cur_image_features = image_features[batch_idx][cur_batch_image_idx]
+                    cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
+                    cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features.squeeze(0)[0:0]], dim=0)
+                    new_input_embeds.append(cur_input_embeds)
+                    new_labels.append(labels[batch_idx])
+                    new_attention_mask.append(attention_mask[batch_idx])
+                    cur_batch_image_idx += 1
+                    continue
+                
+                # Step2-2: Split input_ids, labels, attention_mask by IMAGE_TOKEN_INDEX
+                cur_input_ids_noim, cur_labels_noim, cur_attention_mask_noim = [], [], []
+                cur_labels = labels[batch_idx]
+                cur_attention_mask = attention_mask[batch_idx]
+                cur_img_attention_mask = [
+                    attention_mask[batch_idx][i].item()
+                    for i in torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist()
+                ]
+                image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]] 
+                for i in range(len(image_token_indices) - 1):
+                    cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])
+                    cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])
+                    cur_attention_mask_noim.append(cur_attention_mask[image_token_indices[i]+1:image_token_indices[i+1]])
+                split_sizes = [x.shape[0] for x in cur_labels_noim]
+                cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
+                cur_input_embeds_no_im = list(torch.split(cur_input_embeds, split_sizes, dim=0))# get text features
+
+                # Step2-3: Insert image embedings
+                cur_new_input_embeds, cur_new_labels, cur_new_attention_mask = [], [], []
+                for i in range(num_images + 1): # to add multimodal feature internal the text feature
+                    cur_new_input_embeds.append(cur_input_embeds_no_im[i])
+                    cur_new_labels.append(cur_labels_noim[i])
+                    cur_new_attention_mask.append(cur_attention_mask_noim[i])
+                    if i < num_images:
+                        if getattr(self.config, "eagle_vision_tower", None) is not None:
+                            if getattr(self.config, 'only_navit', False):
+                                cur_image_features = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
+                            else:
+                                siglip_feat = image_features[batch_idx][cur_batch_image_idx]
+                                try:
+                                    qwen2vl_feat = qwen2vl_image_features[batch_idx][cur_batch_image_idx]
+                                    cur_image_features = torch.cat((siglip_feat, qwen2vl_feat), dim=0)
+                                except Exception as e:
+                                    print(e)
+                                    print("only siglip feature:", siglip_feat.shape)
+                                    cur_image_features = siglip_feat
+                        else:
+                            cur_image_features = image_features[batch_idx][cur_batch_image_idx]
+                        cur_batch_image_idx += 1
+                        cur_new_input_embeds.append(cur_image_features)
+                        cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))
+                        cur_new_attention_mask.append(torch.full((cur_image_features.shape[0],), True, device=cur_attention_mask.device, dtype=cur_attention_mask.dtype))
+
+                # Step2-4: Concat image embedings and text embedings
+                cur_new_input_embeds = torch.cat(cur_new_input_embeds)
+                cur_new_labels = torch.cat(cur_new_labels)
+                cur_new_attention_mask = torch.cat(cur_new_attention_mask)
+                new_input_embeds.append(cur_new_input_embeds)
+                new_labels.append(cur_new_labels)
+                new_attention_mask.append(cur_new_attention_mask)
+
+            # Step3: Truncate sequences to max length as image embeddings can make the sequence longer
+            tokenizer_model_max_length = getattr(self.config, 'tokenizer_model_max_length', None)
+            if tokenizer_model_max_length is not None:
+                new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
+                new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
+                new_attention_mask = [x[:tokenizer_model_max_length] for x in new_attention_mask]
+
+            # Step4: Pad and stack input_embeds, labels, attention_mask
+            max_len = max(x.shape[0] for x in new_input_embeds)
+            batch_size = len(new_input_embeds)
+            new_input_embeds_padded = []
+            new_labels_padded = torch.full((batch_size, max_len), IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)
+            new_attention_mask_padded = torch.zeros((batch_size, max_len), dtype=new_attention_mask[0].dtype, device=new_attention_mask[0].device)
+            position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)
+
+            for i, (cur_new_embed, cur_new_labels, cur_attention_mask) in enumerate(zip(new_input_embeds, new_labels, new_attention_mask)):
+                cur_len = cur_new_embed.shape[0]
+                if self.get_padding_method() == 'left':
+                    new_input_embeds_padded.append(torch.cat((
+                        torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device),
+                        cur_new_embed
+                    ), dim=0))
+                    if cur_len > 0:
+                        new_labels_padded[i, -cur_len:] = cur_new_labels
+                        new_attention_mask_padded[i, -cur_len:] = cur_attention_mask
+                        position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
+                
+                else:
+                    new_input_embeds_padded.append(torch.cat((
+                        cur_new_embed,
+                        torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)
+                    ), dim=0))
+                    if cur_len > 0:
+                        new_labels_padded[i, :cur_len] = cur_new_labels
+                        new_attention_mask_padded[i, :cur_len] = cur_attention_mask
+                        position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
+
+            new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
+            new_labels = new_labels_padded if _labels is not None else None
+            new_attention_mask = new_attention_mask_padded if _attention_mask is not None else None
+            if _position_ids is None:
+                position_ids = None
+            
+            return None, position_ids, new_attention_mask, past_key_values, new_input_embeds, new_labels
+
+
+class ValleyQwen3Model(ValleyMetaModel, Qwen3Model):
+    config_class = ValleyConfig
+    def __init__(self, config: Qwen3Config):
+        super(ValleyQwen3Model, self).__init__(config)
+
+
+class ValleyQwen3ForCausalLM(Qwen3ForCausalLM, ValleyMetaForCausalLM):
+    config_class = ValleyConfig
+
+    def __init__(self, config):
+        super(Qwen3ForCausalLM, self).__init__(config)
+        self.model = ValleyQwen3Model(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 _update_model_kwargs_for_generation(
+        self,
+        outputs: CausalLMOutputWithPast,
+        model_kwargs: Dict[str, Any],
+        is_encoder_decoder: bool = False,
+        num_new_tokens: int = 1,
+    ) -> Dict[str, Any]:
+        new_model_kwargs = super()._update_model_kwargs_for_generation(
+            outputs, 
+            model_kwargs, 
+            is_encoder_decoder, 
+            num_new_tokens
+        )
+        """
+        Set model_kwargs["attention_mask"] to the expanded `attention_mask` in
+        the `prepare_inputs_labels_for_multimodal` function to ensure the 
+        correctness of the generate behavior when `use_cache` is enabled.
+        """
+        if not is_encoder_decoder:
+            if "attention_mask" in new_model_kwargs:
+                attention_mask = outputs.attention_mask
+                new_model_kwargs["attention_mask"] = torch.cat(
+                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+                )
+        return new_model_kwargs
+
+
+    def forward(
+        self,
+        input_ids: 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,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        image_sizes: Optional[List[List[int]]] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+    ) -> 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
+
+        if inputs_embeds is None:
+            (
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                inputs_embeds,
+                labels
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                labels,
+                images,
+                image_sizes,
+                pixel_values,
+                pixel_values_videos,
+                image_grid_thw,
+                video_grid_thw,
+            )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        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,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = CrossEntropyLoss(reduction='mean')
+            bs = shift_labels.shape[0]
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = torch.stack([loss_fct(shift_logits[i], shift_labels[i]) for i in range(bs)])
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        res =  CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+        res.attention_mask = attention_mask
+        return res
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        # For cases where only navit (e.g., Valley3) is used, the input `images` may be None.
+        # To ensure compatibility with the Valley2 codebase, we need to supplement a dummy image.
+        image_sizes = kwargs.get("image_sizes", None)
+        if kwargs.get("images", None) is not None:
+            images = kwargs.get("images")
+        else:
+            images = [torch.zeros((len(image_sizes_per_sample), 3, 10, 10)) for image_sizes_per_sample in image_sizes]
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "images": images,
+                "image_sizes": kwargs.get("image_sizes", None),
+                "pixel_values": kwargs.get("pixel_values", None),
+                "pixel_values_videos": kwargs.get("pixel_values_videos", None),
+                "image_grid_thw": kwargs.get("image_grid_thw", None),
+                "video_grid_thw": kwargs.get("video_grid_thw", None),
+            }
+        )
+        return model_inputs
+
+AutoConfig.register("valley", ValleyConfig)
+AutoModelForCausalLM.register(ValleyConfig, ValleyQwen3ForCausalLM)