returnn 1.20250110.143435__tar.gz → 1.20250113.120418__tar.gz

{returnn-1.20250110.143435/returnn.egg-info → returnn-1.20250113.120418}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250110.143435
+Version: 1.20250113.120418
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20250113.120418/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250113.120418'
+long_version = '1.20250113.120418+git.0428e25'

{returnn-1.20250110.143435 → returnn-1.20250113.120418}/requirements.txt

@@ -1,5 +1,4 @@
 numpy
 h5py
-typing
 # Nest alternative, provides `tree`: https://github.com/rwth-i6/returnn/issues/1314
 dm-tree

{returnn-1.20250110.143435 → returnn-1.20250113.120418}/returnn/frontend/attention.py

@@ -330,7 +330,7 @@ class RotaryPosCausalSelfAttention(CausalSelfAttention):
         q = _apply_rope(
             q,
             (
-                rf.gather(pos_enc, axis=hist_dim, indices=hist_dim.get_size_tensor() - 1)
+                rf.gather(pos_enc, axis=hist_dim, indices=rf.last_frame_position_of_dim(hist_dim))
                 if axis == single_step_dim
                 else rf.replace_dim(pos_enc, in_dim=hist_dim, out_dim=axis)[0]
             ),

{returnn-1.20250110.143435 → returnn-1.20250113.120418}/returnn/frontend/conversions/hf_llama.py

@@ -8,6 +8,7 @@ import returnn.frontend as rf
 from returnn.frontend.decoder.transformer import TransformerDecoder, TransformerDecoderLayer, FeedForwardGated
 
 if TYPE_CHECKING:
+    # noinspection PyUnresolvedReferences,PyPackageRequirements,PyProtectedMember
     from transformers.models.llama.modeling_llama import (
         LlamaModel,
         LlamaForCausalLM,
@@ -25,6 +26,8 @@ def import_params_hf_llama_to_rf_transformer_decoder(
     Import params from HF Llama model to RF :class:`TransformerDecoder`.
     """
     import torch
+
+    # noinspection PyUnresolvedReferences,PyPackageRequirements,PyProtectedMember
     from transformers.models.llama.modeling_llama import LlamaModel, LlamaForCausalLM, LlamaDecoderLayer
 
     print("HF Model:")
@@ -206,10 +209,10 @@ def import_params_hf_llama_att_to_rf_rotary_att(model_hf: LlamaAttention, model_
     """
     import torch
 
-    assert model_hf.num_heads == model_rf.num_heads.dimension
-    assert model_hf.hidden_size == model_rf.in_dim.dimension
-    dim = model_hf.hidden_size
-    nh = model_hf.num_heads
+    assert model_hf.config.num_attention_heads == model_rf.num_heads.dimension
+    assert model_hf.config.hidden_size == model_rf.in_dim.dimension
+    dim = model_hf.config.hidden_size
+    nh = model_hf.config.num_attention_heads
     hdim = dim // nh
 
     print("HF Model:")

{returnn-1.20250110.143435 → returnn-1.20250113.120418}/returnn/util/py-to-pickle.cpp

@@ -15,6 +15,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <assert.h>
+#include <stdint.h>
 
 const int protocol = 3;
 

{returnn-1.20250110.143435 → returnn-1.20250113.120418/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250110.143435
+Version: 1.20250113.120418
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20250110.143435 → returnn-1.20250113.120418}/tests/test_rf_attention.py

@@ -280,6 +280,7 @@ def test_rope_causal_self_att():
         LlamaRotaryEmbedding,
         LlamaConfig,
         apply_rotary_pos_emb,
+        eager_attention_forward,
     )
 
     config = LlamaConfig(
@@ -322,7 +323,8 @@ def test_rope_causal_self_att():
 
     position_ids = rf.expand_dim(rf.range_over_dim(seq_dim), batch_dim)  # LlamaRotaryEmbedding wants this
     with PyTracer(
-        [LlamaAttention.forward, LlamaRotaryEmbedding.forward, apply_rotary_pos_emb], torch.Tensor
+        [LlamaAttention.forward, LlamaRotaryEmbedding.forward, apply_rotary_pos_emb, eager_attention_forward],
+        torch.Tensor,
     ) as trace_hf:
         # causal_mask code copied from LlamaAttention
         sequence_length = target_length = in_.raw_tensor.shape[1]
@@ -340,11 +342,11 @@ def test_rope_causal_self_att():
         position_embeddings = rotary_emb(
             torch.zeros(()), position_ids=position_ids.copy_compatible_to_dims_raw((batch_dim, seq_dim))
         )
-        out_hf, _, _ = model_hf(in_.raw_tensor, attention_mask=causal_mask, position_embeddings=position_embeddings)
+        out_hf, *_ = model_hf(in_.raw_tensor, attention_mask=causal_mask, position_embeddings=position_embeddings)
     pprint(trace_hf.captured_locals)
 
     print("First HF att weight tensor:")
-    print(trace_hf.captured_locals[LlamaAttention.forward][0]["attn_weights"][2][0, 0, 0].detach().numpy())
+    print(trace_hf.captured_locals[LlamaAttention.forward][0]["attn_weights"][-1][0, 0, 0].detach().numpy())
 
     check_py_traces_rf_to_pt_equal(
         trace_rf.captured_locals,
@@ -352,8 +354,13 @@ def test_rope_causal_self_att():
         [
             (
                 (rf.RotaryPosCausalSelfAttention.__call__, 0, "q", 0),
-                (LlamaAttention.forward, 0, "query_states", 1),
+                # input: batch_dim, seq_dim, model_dim
+                # input_shape: batch_dim, seq_dim
+                # HF query_states': (batch_dim, seq_dim, num_heads, self.head_dim),
+                #   then transposed to (batch_dim, num_heads, seq_dim, self.head_dim)
+                (LlamaAttention.forward, 0, "query_states", 0),
                 lambda x, *, name, **_: rf.convert_to_tensor(
+                    # reorder complex numbers
                     x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).flatten(-2),
                     dims=(batch_dim, model_rf.num_heads, seq_dim, model_rf.key_dim_per_head),
                     name=name,
@@ -361,8 +368,9 @@ def test_rope_causal_self_att():
             ),
             (
                 (rf.RotaryPosCausalSelfAttention.__call__, 0, "k", 0),
-                (LlamaAttention.forward, 0, "key_states", 1),
+                (LlamaAttention.forward, 0, "key_states", 0),
                 lambda x, *, name, **_: rf.convert_to_tensor(
+                    # reorder complex numbers
                     x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).flatten(-2),
                     dims=(batch_dim, model_rf.num_heads, seq_dim, model_rf.key_dim_per_head),
                     name=name,
@@ -418,25 +426,25 @@ def test_rope_causal_self_att():
             ),
             (
                 (rf.dot_attention, 0, "energy", 0),
-                (LlamaAttention.forward, 0, "attn_weights", 0),
+                (eager_attention_forward, 0, "attn_weights", 0),
                 (batch_dim, model_rf.num_heads, seq_dim, "axis"),
             ),
             (
                 (rf.dot_attention, 0, "att_weights", 0),
-                (LlamaAttention.forward, 0, "attn_weights", 2),
+                (LlamaAttention.forward, 0, "attn_weights", -1),
                 (batch_dim, model_rf.num_heads, seq_dim, "axis"),
             ),
             (
                 (rf.dot_attention, 0, "att", 0),
                 (LlamaAttention.forward, 0, "attn_output", 0),
-                (batch_dim, model_rf.num_heads, seq_dim, model_rf.value_dim_per_head),
+                (batch_dim, seq_dim, model_rf.num_heads, model_rf.value_dim_per_head),
             ),
         ],
     )
 
     print("Final check...")
     assert out_rf.raw_tensor.shape == out_hf.shape
-    torch.testing.assert_allclose(out_rf.raw_tensor, out_hf)
+    torch.testing.assert_close(out_rf.raw_tensor, out_hf)
     print("  all matched!")
 
 

returnn-1.20250110.143435/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250110.143435'
-long_version = '1.20250110.143435+git.5a80223'