Update modelling_RW.py

modelling_RW.py

@@ -54,6 +54,7 @@ class RotaryEmbedding(torch.nn.Module):
         self,
         head_dim: int,
         base=10000,
+        use_cache=False,
     ):
         super().__init__()
         inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
@@ -63,6 +64,7 @@ class RotaryEmbedding(torch.nn.Module):
         self.batch_size_cached = None
         self.cos_cached: torch.Tensor | None = None
         self.sin_cached: torch.Tensor | None = None
+        self.use_cache = use_cache
 
     def cos_sin(
         self,
@@ -70,7 +72,7 @@ class RotaryEmbedding(torch.nn.Module):
         device="cuda",
         dtype=torch.bfloat16,
     ) -> torch.Tensor:
-        if seq_len != self.seq_len_cached:
+        if seq_len != self.seq_len_cached or not self.use_cache:
             self.seq_len_cached = seq_len
             t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
             freqs = torch.einsum("i,j->ij", t, self.inv_freq)