Update modelling_RW.py

ensure use 40b file as reference.

modelling_RW.py

@@ -52,10 +52,11 @@ class RotaryEmbedding(torch.nn.Module):
 
     def __init__(
         self,
-        head_dim: int,
+        config,
         base=10000,
-        use_cache=False,
     ):
+        head_dim = config.head_dim
+        self.use_cache = config.use_cache
         super().__init__()
         inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -64,7 +65,6 @@ 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,
@@ -184,7 +184,7 @@ class Attention(nn.Module):
                 f" {self.num_heads})."
             )
 
-        self.maybe_rotary = RotaryEmbedding(config.head_dim) if config.rotary else lambda q, k: (q, k)
+        self.maybe_rotary = RotaryEmbedding(config) if config.rotary else lambda q, k: (q, k)
 
         # Layer-wise attention scaling
         self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)