Fix LlamaAttention compatibility with transformers 4.53.0

The transformers 4.53 API requires `position_embeddings` (cos/sin rotary
tuple) as a positional arg to LlamaAttention.forward() and returns 2
values instead of 3. Update DiffLlama to compute rotary embeddings and
pass them through the decoder layers.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

soulxsinger/models/modules/llama.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from typing import List, Optional, Tuple, Union
 import math
 
-from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRotaryEmbedding
 from transformers.models.llama.modeling_llama import BaseModelOutputWithPast
 
 
@@ -62,27 +62,13 @@ class LlamaNARDecoderLayer(LlamaDecoderLayer):
         hidden_states: torch.Tensor,
         cond_embedding: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        past_key_value=None,
         output_attentions: Optional[bool] = False,
         use_cache: Optional[bool] = False,
     ) -> Tuple[
         torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
     ]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-        """
-
         residual = hidden_states
 
         hidden_states = self.input_layernorm(
@@ -90,13 +76,11 @@ class LlamaNARDecoderLayer(LlamaDecoderLayer):
         )
 
         # Self Attention
-        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+        hidden_states, self_attn_weights = self.self_attn(
             hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
             attention_mask=attention_mask,
-            position_ids=position_ids,
             past_key_value=past_key_value,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
         )
         hidden_states = residual + hidden_states
 
@@ -113,9 +97,6 @@ class LlamaNARDecoderLayer(LlamaDecoderLayer):
         if output_attentions:
             outputs += (self_attn_weights,)
 
-        if use_cache:
-            outputs += (present_key_value,)
-
         return outputs
 
 
@@ -185,6 +166,15 @@ class DiffLlama(LlamaModel):
 
         self.embed_tokens = None
 
+        # Re-create rotary_emb with the actual layer config dimensions
+        layer_config = LlamaConfig(
+            hidden_size=hidden_size,
+            num_attention_heads=num_heads,
+            max_position_embeddings=4096,
+            intermediate_size=hidden_size * 4,
+        )
+        self.rotary_emb = LlamaRotaryEmbedding(config=layer_config)
+
         self.post_init()
 
         # self.reset_parameters()
@@ -309,6 +299,9 @@ class DiffLlama(LlamaModel):
 
         hidden_states = inputs_embeds
 
+        # Compute rotary position embeddings
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 use_cache = False
@@ -328,40 +321,19 @@ class DiffLlama(LlamaModel):
                 past_key_values[idx] if past_key_values is not None else None
             )
 
-            if self.gradient_checkpointing and self.training:
-                raise NotImplementedError
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs, output_attentions, None)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(decoder_layer),
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    None,
-                )
-            else:
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_value,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                    cond_embedding=diffusion_step,
-                )
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_embeddings=position_embeddings,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cond_embedding=diffusion_step,
+            )
 
             hidden_states = layer_outputs[0]
             all_layer_hidden_states.append(hidden_states.clone())
 
-            if use_cache:
-                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
-
             if output_attentions:
                 all_self_attns += (layer_outputs[1],)