Add custom modeling code

modeling_nebula.py

@@ -12,9 +12,17 @@ class NebulaConfig(PretrainedConfig):
     def __init__(self, dim=1280, n_layers=14, n_heads=10, n_kv_heads=10, vocab_size=60729,
                  multiple_of=256, ffn_dim_multiplier=8/3, norm_eps=1e-5, max_seq_len=2048,
                  dropout=0.1, use_cache=True, **kwargs):
-        self.dim, self.n_layers, self.n_heads, self.n_kv_heads = dim, n_layers, n_heads, n_kv_heads
-        self.vocab_size, self.multiple_of, self.ffn_dim_multiplier = vocab_size, multiple_of, ffn_dim_multiplier
-        self.norm_eps, self.max_seq_len, self.dropout, self.use_cache = norm_eps, max_seq_len, dropout, use_cache
+        self.dim = dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.vocab_size = vocab_size
+        self.multiple_of = multiple_of
+        self.ffn_dim_multiplier = ffn_dim_multiplier
+        self.norm_eps = norm_eps
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+        self.use_cache = use_cache
         super().__init__(**kwargs)
 
 class RMSNorm(nn.Module):
@@ -32,6 +40,7 @@ class RoPE(nn.Module):
         super().__init__()
         self.dim = config.dim // config.n_heads
         self.max_seq_len = config.max_seq_len
+        # The device will be inferred from the model, so we don't need it in the config
         self._build_cache(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
     def _build_cache(self, device, base=10000):
         theta = 1.0 / (base ** (torch.arange(0, self.dim, 2, device=device).float() / self.dim))