Upload modeling_workshop_gpt.py

modeling_workshop_gpt.py

@@ -36,19 +36,18 @@ class RotaryPositionalEmbeddings(nn.Module):
         self.dim = dim
         self.max_seq_len = max_seq_len
         self.base = base
-        theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer("theta", theta, persistent=False)
-        self._build_cache(max_seq_len)
+        self.cache = None
 
-    def _build_cache(self, seq_len):
-        seq = torch.arange(seq_len, device=self.theta.device)
-        freqs = torch.outer(seq, self.theta)
-        self.register_buffer("cache", torch.stack([freqs.cos(), freqs.sin()], dim=-1), persistent=False)
+    def _build_cache(self, seq_len, device):
+        theta = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device).float() / self.dim))
+        seq = torch.arange(seq_len, device=device)
+        freqs = torch.outer(seq, theta)
+        self.cache = torch.stack([freqs.cos(), freqs.sin()], dim=-1)
 
     def forward(self, x, *, input_pos=None):
         seq_len = x.shape[-2]
-        if seq_len > self.cache.shape[0]:
-            self._build_cache(seq_len)
+        if self.cache is None or self.cache.shape[0] < seq_len or self.cache.device != x.device:
+            self._build_cache(max(seq_len, self.max_seq_len), x.device)
         cache = self.cache[:seq_len] if input_pos is None else self.cache[input_pos]
         x1, x2 = x.float().unflatten(-1, (-1, 2)).unbind(-1)
         cos, sin = cache.unbind(-1)