Update transformer config

Author: gonzalobenegas

configs/experiment/transformer_small.yaml

@@ -0,0 +1,36 @@
+# @package _global_
+
+# Short training run with small Transformer encoder for quick testing
+
+defaults:
+  - override /model: bert_transformer_small
+
+logger:
+  wandb:
+    name: debug-transformer-small
+    tags: ["debug"]
+
+trainer:
+  max_steps: 100
+  log_every_n_steps: 5
+  val_check_interval: 5
+  limit_val_batches: 2
+  check_val_every_n_epoch: null
+
+model:
+  net:
+    embedder:
+      d_model: 32
+    encoder:
+      n_layers: 2
+  scheduler:
+    _target_: transformers.get_cosine_schedule_with_warmup
+    _partial_: true
+    num_warmup_steps: 10
+    num_training_steps: ${trainer.max_steps}
+
+data:
+  batch_size: 8
+  per_device_batch_size: 8
+
+compile: false

glm_experiments/models/components/transformer.py

@@ -36,12 +36,12 @@ def __init__(self, d_in: int, d_out: int):
 
         super().__init__()
         std = math.sqrt(2 / (d_in + d_out))
-        self.weight: Float[Tensor,  d_out d_in] = nn.Parameter(
+        self.weight: Float[Tensor, " d_out d_in"] = nn.Parameter(
             nn.init.trunc_normal_(torch.empty(d_out, d_in), std=std, a=-3 * std, b=3 * std),
             requires_grad=True,
         )
 
-    def forward(self, x: Float[Tensor,  ... d_in]) -> Float[Tensor,  ... d_out]:
+    def forward(self, x: Float[Tensor, " ... d_in"]) -> Float[Tensor, " ... d_out"]:
         return einsum(x, self.weight, "... d_in, d_out d_in -> ... d_out")
 
     def extra_repr(self):
@@ -59,7 +59,7 @@ def __init__(self, vocab_size: int, d_model: int):
             requires_grad=True,
         )
 
-    def forward(self, token_ids: Int[Tensor,  ...]) -> Float[Tensor,  ... d_model]:
+    def forward(self, token_ids: Int[Tensor, " ..."]) -> Float[Tensor, " ... d_model"]:
         return self.weight[token_ids, :]
 
     def extra_repr(self):
@@ -78,7 +78,7 @@ def __init__(self, context_length: int, dim: int, theta: float = 10000.0):
     @staticmethod
     def _init_cache(
         context_length: int, dim: int, theta: float
-    ) -> Float[Tensor,  2 context_length half_dim]:
+    ) -> Float[Tensor, " 2 context_length half_dim"]:
         assert dim % 2 == 0
 
         d = torch.arange(0, dim, 2) / dim
@@ -91,8 +91,8 @@ def _init_cache(
         return torch.stack((cos, sin))
 
     def forward(
-        self, x: Float[Tensor,  ... seq d], pos_ids: Int[Tensor,  ... seq]
-    ) -> Float[Tensor,  ... seq d]:
+        self, x: Float[Tensor, " ... seq d"], pos_ids: Int[Tensor, " ... seq"]
+    ) -> Float[Tensor, " ... seq d"]:
         x1, x2 = rearrange(x, "... (half_d xy) -> xy ... half_d", xy=2)
 
         # einx
@@ -172,9 +172,9 @@ def __init__(
 
     def forward(
         self,
-        x: Float[Tensor,  ... seq d_k],
-        token_positions: Int[Tensor,  ... seq] | None = None,
-    ) -> Float[Tensor,  ... seq d_v]:
+        x: Float[Tensor, " ... seq d_k"],
+        token_positions: Int[Tensor, " ... seq"] | None = None,
+    ) -> Float[Tensor, " ... seq d_v"]:
         """
         Args:
             x: The input to perform multi-headed self-attention on.