fixed the dynamic resizing of rotary embeddings

OpenNMT · Feb 13, 2024 · b7e8b8f · b7e8b8f
1 parent 54514b8
commit b7e8b8f
Showing 1 changed file with 20 additions and 19 deletions.
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
@@ -18,7 +18,7 @@
 # are both < 2048 tokens.
 
 
-def rotaryembeddings(dim: int, maxseqlen=4096, base=10000):
+def rotaryembeddings(dim: int, maxseqlen=2048, base=10000):
     inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
     tmax = torch.arange(maxseqlen, device=inv_freq.device)
     rope = torch.outer(tmax, inv_freq).float()
@@ -467,7 +467,7 @@ def forward(
                         if seqlen > self.rope.size(0):
                             self.rope = rotaryembeddings(
                                 self.rotary_dim,
-                                maxseqlen=(seqlen + 4096),
+                                maxseqlen=(seqlen + 2048),
                                 base=self.rotary_theta,
                             ).to(self.rope.device)
                         rope = self.rope[start_pos : start_pos + seqlen]
@@ -486,23 +486,6 @@ def forward(
                     self.layer_cache[1]["values"] = value
 
                 else:
-                    if self.max_relative_positions == -1:  # Rotary Embeddings
-                        if seqlen > self.rope.size(0):
-                            self.rope = rotaryembeddings(
-                                self.rotary_dim,
-                                maxseqlen=(seqlen + 4096),
-                                base=self.rotary_theta,
-                            ).to(self.rope.device)
-                            self.cos = (
-                                self.rope[:, : self.rope.size(1) // 2]
-                                .real.contiguous()
-                                .half()
-                            )
-                            self.sin = (
-                                self.rope[:, : self.rope.size(1) // 2]
-                                .imag.contiguous()
-                                .half()
-                            )
                     if start_pos >= self.layer_cache[1]["keys"].size(2):
                         self.layer_cache[1]["keys"] = torch.cat(
                             [
@@ -528,6 +511,24 @@ def forward(
                             ],
                             dim=-2,
                         )
+                        seqlen = self.layer_cache[1]["keys"].size(2)
+                    if self.max_relative_positions == -1:  # Rotary Embeddings
+                        if seqlen > self.rope.size(0):
+                            self.rope = rotaryembeddings(
+                                self.rotary_dim,
+                                maxseqlen=(seqlen + 2048),
+                                base=self.rotary_theta,
+                            ).to(self.rope.device)
+                            self.cos = (
+                                self.rope[:, : self.rope.size(1) // 2]
+                                .real.contiguous()
+                                .half()
+                            )
+                            self.sin = (
+                                self.rope[:, : self.rope.size(1) // 2]
+                                .imag.contiguous()
+                                .half()
+                            )
                     if sliding_window > 0 and key.size(2) > sliding_window:
                         self.layer_cache[1]["keys"] = self.layer_cache[1]["keys"][
                             :, :, 1:, :