fix rope freq saving into model state dict

rope_embeddings.py

@@ -10,39 +10,38 @@ from typing import Literal
 def exists(val):
     return val is not None
 
+
 def default(val, d):
     return val if exists(val) else d
 
+
 def broadcat(tensors, dim=-1):
     broadcasted_tensors = broadcast_tensors(*tensors)
     return torch.cat(broadcasted_tensors, dim=dim)
 
+
 def slice_at_dim(t, dim_slice: slice, *, dim):
-    dim += (t.ndim if dim < 0 else 0)
+    dim += t.ndim if dim < 0 else 0
     colons = [slice(None)] * t.ndim
     colons[dim] = dim_slice
     return t[tuple(colons)]
 
+
 def rotate_half(x):
     orig_shape = x.shape
     d_head = orig_shape[-1]
     x = x.view(*orig_shape[:-1], d_head // 2, 2)
-    
+
     x1 = x[..., 0]
     x2 = x[..., 1]
-    
+
     res = torch.stack((-x2, x1), dim=-1)
     return res.view(*orig_shape)
 
 
-@autocast('cuda', enabled=False)
+@autocast("cuda", enabled=False)
 def apply_rotary_emb(
-    freqs,
-    t,
-    start_index=0,
-    scale=1.,
-    seq_dim=-2,
-    freqs_seq_dim=None
+    freqs, t, start_index=0, scale=1.0, seq_dim=-2, freqs_seq_dim=None
 ):
     dtype = t.dtype
 
@@ -57,21 +56,25 @@ def apply_rotary_emb(
     rot_dim = freqs.shape[-1]
     end_index = start_index + rot_dim
 
-    assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
+    assert rot_dim <= t.shape[-1], (
+        f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"
+    )
 
     t_left = t[..., :start_index]
     t_middle = t[..., start_index:end_index]
     t_right = t[..., end_index:]
 
-    t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)
-        
+    t_transformed = (t_middle * freqs.cos() * scale) + (
+        rotate_half(t_middle) * freqs.sin() * scale
+    )
+
     out = torch.cat((t_left, t_transformed, t_right), dim=-1)
     return out.type(dtype)
 
 
 def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None):
     if exists(freq_ranges):
-        rotations = torch.einsum('..., f -> ... f', rotations, freq_ranges)
+        rotations = torch.einsum("..., f -> ... f", rotations, freq_ranges)
         rotations = rotations.reshape(*rotations.shape[:-2], -1)
 
     rotations = rotations.repeat_interleave(2, dim=-1)
@@ -83,18 +86,18 @@ class RotaryEmbedding(Module):
         self,
         dim,
         custom_freqs: Tensor | None = None,
-        freqs_for: Literal['lang', 'pixel', 'constant'] = 'lang',
-        theta = 10000,
-        max_freq = 10,
-        num_freqs = 1,
-        learned_freq = False,
-        use_xpos = False,
-        xpos_scale_base = 512,
-        interpolate_factor = 1.,
-        theta_rescale_factor = 1.,
-        seq_before_head_dim = False,
-        cache_if_possible = True,
-        cache_max_seq_len = 8192
+        freqs_for: Literal["lang", "pixel", "constant"] = "lang",
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+        learned_freq=False,
+        use_xpos=False,
+        xpos_scale_base=512,
+        interpolate_factor=1.0,
+        theta_rescale_factor=1.0,
+        seq_before_head_dim=False,
+        cache_if_possible=True,
+        cache_max_seq_len=8192,
     ):
         super().__init__()
 
@@ -103,28 +106,35 @@ class RotaryEmbedding(Module):
 
         if exists(custom_freqs):
             freqs = custom_freqs
-        elif freqs_for == 'lang':
-            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
-        elif freqs_for == 'pixel':
-            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
-        elif freqs_for == 'constant':
+        elif freqs_for == "lang":
+            freqs = 1.0 / (
+                theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
+            )
+        elif freqs_for == "pixel":
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "constant":
             freqs = torch.ones(num_freqs).float()
 
         self.cache_if_possible = cache_if_possible
         self.cache_max_seq_len = cache_max_seq_len
 
-        self.register_buffer('cached_freqs', torch.zeros(cache_max_seq_len, dim), persistent=False)
+        self.register_buffer(
+            "cached_freqs", torch.zeros(cache_max_seq_len, dim), persistent=False
+        )
         self.cached_freqs_seq_len = 0
 
-        self.freqs = nn.Parameter(freqs, requires_grad=learned_freq)
         self.learned_freq = learned_freq
+        if learned_freq:
+            self.freqs = nn.Parameter(freqs)
+        else:
+            self.register_buffer("freqs", freqs, persistent=False)
 
-        self.register_buffer('dummy', torch.tensor(0), persistent=False)
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
 
         self.seq_before_head_dim = seq_before_head_dim
         self.default_seq_dim = -3 if seq_before_head_dim else -2
 
-        assert interpolate_factor >= 1.
+        assert interpolate_factor >= 1.0
         self.interpolate_factor = interpolate_factor
 
         self.use_xpos = use_xpos
@@ -135,8 +145,10 @@ class RotaryEmbedding(Module):
         scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
         self.scale_base = xpos_scale_base
 
-        self.register_buffer('scale', scale, persistent=False)
-        self.register_buffer('cached_scales', torch.zeros(cache_max_seq_len, dim), persistent=False)
+        self.register_buffer("scale", scale, persistent=False)
+        self.register_buffer(
+            "cached_scales", torch.zeros(cache_max_seq_len, dim), persistent=False
+        )
         self.cached_scales_seq_len = 0
 
         self.apply_rotary_emb = staticmethod(apply_rotary_emb)
@@ -148,11 +160,15 @@ class RotaryEmbedding(Module):
     def get_seq_pos(self, seq_len, device=None, dtype=None, offset=0):
         device = default(device, self.device)
         dtype = default(dtype, self.cached_freqs.dtype)
-        return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor
+        return (
+            torch.arange(seq_len, device=device, dtype=dtype) + offset
+        ) / self.interpolate_factor
 
     def rotate_queries_or_keys(self, t, seq_dim=None, offset=0, scale=None):
         seq_dim = default(seq_dim, self.default_seq_dim)
-        assert not self.use_xpos or exists(scale), 'you must use `.rotate_queries_and_keys` method instead'
+        assert not self.use_xpos or exists(scale), (
+            "you must use `.rotate_queries_and_keys` method instead"
+        )
 
         device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim]
         seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset)
@@ -161,23 +177,29 @@ class RotaryEmbedding(Module):
         if seq_dim == -3:
             freqs = freqs.unsqueeze(1)
 
-        return apply_rotary_emb(freqs, t, scale=default(scale, 1.), seq_dim=seq_dim)
+        return apply_rotary_emb(freqs, t, scale=default(scale, 1.0), seq_dim=seq_dim)
 
     def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0):
-        dtype, device, seq_dim = q.dtype, q.device, default(seq_dim, self.default_seq_dim)
+        dtype, device, seq_dim = (
+            q.dtype,
+            q.device,
+            default(seq_dim, self.default_seq_dim),
+        )
 
         q_len, k_len = q.shape[seq_dim], k.shape[seq_dim]
         assert q_len <= k_len
 
-        q_scale = k_scale = 1.
+        q_scale = k_scale = 1.0
 
         if self.use_xpos:
             seq = self.get_seq_pos(k_len, dtype=dtype, device=device)
             q_scale = self.get_scale(seq[-q_len:]).type(dtype)
             k_scale = self.get_scale(seq).type(dtype)
 
-        rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset)
-        rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale ** -1)
+        rotated_q = self.rotate_queries_or_keys(
+            q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset
+        )
+        rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale**-1)
 
         return rotated_q.type(q.dtype), rotated_k.type(k.dtype)
 
@@ -195,18 +217,22 @@ class RotaryEmbedding(Module):
             scale = scale.unsqueeze(1)
 
         rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim)
-        rotated_k = apply_rotary_emb(freqs, k, scale=scale ** -1, seq_dim=seq_dim)
+        rotated_k = apply_rotary_emb(freqs, k, scale=scale**-1, seq_dim=seq_dim)
 
         return rotated_q.type(q.dtype), rotated_k.type(k.dtype)
 
     def get_scale(self, t: Tensor, seq_len: int | None = None, offset=0):
         assert self.use_xpos
-        should_cache = self.cache_if_possible and exists(seq_len) and (offset + seq_len) <= self.cache_max_seq_len
+        should_cache = (
+            self.cache_if_possible
+            and exists(seq_len)
+            and (offset + seq_len) <= self.cache_max_seq_len
+        )
 
         if should_cache and (seq_len + offset) <= self.cached_scales_seq_len:
-            return self.cached_scales[offset:(offset + seq_len)]
+            return self.cached_scales[offset : (offset + seq_len)]
 
-        scale = 1.
+        scale = 1.0
         if self.use_xpos:
             power = (t - len(t) // 2) / self.scale_base
             scale = self.scale ** power.unsqueeze(-1)
@@ -218,7 +244,9 @@ class RotaryEmbedding(Module):
 
         return scale
 
-    def get_axial_freqs(self, *dims, offsets: tuple[int | float, ...] | Tensor | None = None):
+    def get_axial_freqs(
+        self, *dims, offsets: tuple[int | float, ...] | Tensor | None = None
+    ):
         Colon = slice(None)
         all_freqs = []
 
@@ -232,7 +260,7 @@ class RotaryEmbedding(Module):
             if exists(offsets):
                 offset = offsets[ind]
 
-            if self.freqs_for == 'pixel':
+            if self.freqs_for == "pixel":
                 pos = torch.linspace(-1, 1, steps=dim, device=self.device)
             else:
                 pos = torch.arange(dim, device=self.device)
@@ -248,23 +276,26 @@ class RotaryEmbedding(Module):
         all_freqs = broadcast_tensors(*all_freqs)
         return torch.cat(all_freqs, dim=-1)
 
-    @autocast('cuda', enabled=False)
+    @autocast("cuda", enabled=False)
     def forward(self, t: Tensor, seq_len: int | None = None, offset=0):
         should_cache = (
-            self.cache_if_possible and not self.learned_freq and
-            exists(seq_len) and self.freqs_for != 'pixel' and
-            (offset + seq_len) <= self.cache_max_seq_len
+            self.cache_if_possible
+            and not self.learned_freq
+            and exists(seq_len)
+            and self.freqs_for != "pixel"
+            and (offset + seq_len) <= self.cache_max_seq_len
         )
 
         if should_cache and (offset + seq_len) <= self.cached_freqs_seq_len:
-            return self.cached_freqs[offset:(offset + seq_len)].detach()
+            return self.cached_freqs[offset : (offset + seq_len)].detach()
 
         freqs = self.freqs
-        freqs = torch.einsum('..., f -> ... f', t.type(freqs.dtype), freqs)
+        freqs = torch.einsum("..., f -> ... f", t.type(freqs.dtype), freqs)
         freqs = freqs.repeat_interleave(2, dim=-1)
 
         if should_cache and offset == 0:
             self.cached_freqs[:seq_len] = freqs.detach()
             self.cached_freqs_seq_len = seq_len
 
-        return freqs
+        return freqs
+