fix(flaky): test_generate_with_and_without_position_ids in GLM ORC

src/transformers/models/glm_ocr/modeling_glm_ocr.py

@@ -1406,6 +1406,22 @@ def prepare_inputs_for_generation(
         is_first_iteration=False,
         **kwargs,
     ):
+        has_vision_inputs = image_grid_thw is not None or video_grid_thw is not None
+
+        # Keep explicit text-only position ids on the multimodal packed path so it matches inferred generation.
+        if has_vision_inputs and position_ids is not None and position_ids.ndim == 2:
+            past_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+            if past_length == 0:
+                prep_kwargs = {
+                    "input_ids": input_ids,
+                    "attention_mask": attention_mask,
+                    "past_key_values": past_key_values,
+                    "image_grid_thw": image_grid_thw,
+                    "video_grid_thw": video_grid_thw,
+                }
+                position_ids = self._prepare_position_ids_for_generation(input_ids, prep_kwargs)
+            else:
+                position_ids = self._pack_position_ids_with_rope_deltas(position_ids)
         # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
 
         model_inputs = super().prepare_inputs_for_generation(
@@ -1611,6 +1627,34 @@ def _expand_dict_for_generation(dict_to_expand):
 
         return input_ids, model_kwargs
 
+    def _pack_position_ids_with_rope_deltas(self, position_ids: torch.LongTensor) -> torch.LongTensor:
+        vision_positions = position_ids.unsqueeze(0).expand(3, -1, -1)
+        if self.model.rope_deltas is not None:
+            vision_positions = vision_positions + self.model.rope_deltas
+        return torch.cat([position_ids.unsqueeze(0), vision_positions], dim=0)
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs,
+        model_kwargs: dict,
+        is_encoder_decoder=False,
+        num_new_tokens=1,
+    ):
+        model_kwargs = super()._update_model_kwargs_for_generation(
+            outputs=outputs,
+            model_kwargs=model_kwargs,
+            is_encoder_decoder=is_encoder_decoder,
+            num_new_tokens=num_new_tokens,
+        )
+
+        has_vision_inputs = (
+            model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None
+        )
+        position_ids = model_kwargs.get("position_ids")
+        if has_vision_inputs and position_ids is not None and position_ids.ndim == 2:
+            model_kwargs["position_ids"] = self._pack_position_ids_with_rope_deltas(position_ids)
+        return model_kwargs
+
 
 __all__ = [
     "GlmOcrTextModel",

src/transformers/models/glm_ocr/modular_glm_ocr.py

@@ -413,7 +413,82 @@ class GlmOcrModel(Glm4vModel):
 
 
 class GlmOcrForConditionalGeneration(Glm4vForConditionalGeneration):
-    pass
+    def _pack_position_ids_with_rope_deltas(self, position_ids: torch.LongTensor) -> torch.LongTensor:
+        vision_positions = position_ids.unsqueeze(0).expand(3, -1, -1)
+        if self.model.rope_deltas is not None:
+            vision_positions = vision_positions + self.model.rope_deltas
+        return torch.cat([position_ids.unsqueeze(0), vision_positions], dim=0)
+
+    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,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        has_vision_inputs = image_grid_thw is not None or video_grid_thw is not None
+
+        # Keep explicit text-only position ids on the multimodal packed path so it matches inferred generation.
+        if has_vision_inputs and position_ids is not None and position_ids.ndim == 2:
+            past_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+            if past_length == 0:
+                prep_kwargs = {
+                    "input_ids": input_ids,
+                    "attention_mask": attention_mask,
+                    "past_key_values": past_key_values,
+                    "image_grid_thw": image_grid_thw,
+                    "video_grid_thw": video_grid_thw,
+                }
+                position_ids = self._prepare_position_ids_for_generation(input_ids, prep_kwargs)
+            else:
+                position_ids = self._pack_position_ids_with_rope_deltas(position_ids)
+
+        return 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,
+            use_cache=use_cache,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs,
+        model_kwargs: dict,
+        is_encoder_decoder=False,
+        num_new_tokens=1,
+    ):
+        model_kwargs = super()._update_model_kwargs_for_generation(
+            outputs=outputs,
+            model_kwargs=model_kwargs,
+            is_encoder_decoder=is_encoder_decoder,
+            num_new_tokens=num_new_tokens,
+        )
+
+        has_vision_inputs = (
+            model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None
+        )
+        position_ids = model_kwargs.get("position_ids")
+        if has_vision_inputs and position_ids is not None and position_ids.ndim == 2:
+            model_kwargs["position_ids"] = self._pack_position_ids_with_rope_deltas(position_ids)
+        return model_kwargs
 
 
 __all__ = [

tests/models/glm_ocr/test_modeling_glm_ocr.py

@@ -164,6 +164,23 @@ def prepare_config_and_inputs_for_common(self):
         input_ids[input_ids == self.video_end_token_id] = self.pad_token_id
         input_ids[input_ids == self.image_end_token_id] = self.pad_token_id
 
+        non_special_id = 1
+        self.parent.assertNotIn(
+            non_special_id,
+            {
+                self.pad_token_id,
+                self.bos_token_id,
+                self.eos_token_id,
+                self.video_start_token_id,
+                self.video_end_token_id,
+                self.image_start_token_id,
+                self.image_end_token_id,
+                self.image_token_id,
+                self.video_token_id,
+            },
+        )
+        input_ids[input_ids == self.pad_token_id] = non_special_id
+
         input_ids[:, 0] = self.image_start_token_id
         input_ids[:, 1 : 1 + self.num_image_tokens] = self.image_token_id
         input_ids[:, 1 + self.num_image_tokens] = self.image_end_token_id