Fix Transformers 5 compatibility for generation and RoPE cache

modeling_step_vl.py

@@ -345,7 +345,6 @@ class StepRoboticsModel(StepRoboticsPreTrainedModel, GenerationMixin):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
         logits_to_keep: Union[int, torch.Tensor] = 0,
         images: Optional[list[Image.Image]] = None,
         **kwargs: Unpack[TransformersKwargs],
@@ -390,7 +389,6 @@ class StepRoboticsModel(StepRoboticsPreTrainedModel, GenerationMixin):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=True,
-            cache_position=cache_position,
             **kwargs,
         )
 
@@ -461,7 +459,6 @@ class Step3VL10BForCausalLM(StepRoboticsPreTrainedModel, GenerationMixin):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, StepVLCausalLMOutputWithPast]:
         r"""
@@ -504,7 +501,6 @@ class Step3VL10BForCausalLM(StepRoboticsPreTrainedModel, GenerationMixin):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
-            cache_position=cache_position,
             **kwargs,
         )
 
@@ -532,17 +528,25 @@ class Step3VL10BForCausalLM(StepRoboticsPreTrainedModel, GenerationMixin):
     ):
         # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
 
+        is_first_iteration = kwargs.get("is_first_iteration")
+        if is_first_iteration is None:
+            is_first_iteration = cache_position is None or cache_position[0] == 0
+
+        # Transformers 4.x still supplies `cache_position`; keep forwarding it
+        # there without making Transformers 5 recreate this removed argument.
+        if cache_position is not None:
+            kwargs["cache_position"] = cache_position
+
         model_inputs = super().prepare_inputs_for_generation(
             input_ids,
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
-            cache_position=cache_position,
             logits_to_keep=logits_to_keep,
             **kwargs,
         )
 
-        if cache_position[0] == 0:
+        if is_first_iteration:
             # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
             # Otherwise we need pixel values to be passed to model
             model_inputs["pixel_values"] = pixel_values
@@ -554,4 +558,3 @@ class Step3VL10BForCausalLM(StepRoboticsPreTrainedModel, GenerationMixin):
             return key[len("language_model."):], True
         
         return key, False
-

vision_encoder.py

@@ -67,6 +67,10 @@ class EncoderRope2D(nn.Module):
         self.max_freq = max_freq
         self.num_freqs = num_freqs
         cache = self._compute_2d_freqs()
+        # `from_pretrained` initializes remote-code models on the meta device in
+        # Transformers 5. Non-persistent buffers are absent from checkpoints, so
+        # the materialized cache has to be recomputed before it is first used.
+        self._freqs_cache_needs_rebuild = cache.device.type == "meta"
         self.register_buffer("freqs_cache", cache, persistent=False)
 
     def _compute_inv_freq(self, base: Union[int, float],
@@ -102,6 +106,10 @@ class EncoderRope2D(nn.Module):
 
     def forward(self, q: torch.Tensor, k: torch.Tensor,
                 grid_hw: tuple[int, int]):
+        if self._freqs_cache_needs_rebuild:
+            self.freqs_cache = self._compute_2d_freqs().to(q.device)
+            self._freqs_cache_needs_rebuild = False
+
         # If grid matches cached shape we reuse directly to avoid recomputation.
         if grid_hw[0] != self.max_grid_height or grid_hw[1] != self.max_grid_width:
             rows = torch.arange(grid_hw[0], device=q.device).view(-1, 1)