Add transformer encoder

Author: gonzalobenegas

.pre-commit-config.yaml

@@ -11,6 +11,7 @@ repos:
       - id: check-docstring-first
       - id: check-yaml
       - id: debug-statements
+        exclude: glm_experiments/models/components/transformer.py
       - id: detect-private-key
       - id: check-executables-have-shebangs
       - id: check-toml
@@ -23,6 +24,7 @@ repos:
     hooks:
       - id: black
         args: [--line-length, "99"]
+        exclude: glm_experiments/models/components/transformer.py
 
   # python import sorting
   - repo: https://github.com/PyCQA/isort
@@ -37,6 +39,7 @@ repos:
     hooks:
       - id: pyupgrade
         args: [--py313-plus]
+        exclude: glm_experiments/models/components/transformer.py
 
   # python check (PEP8), programming errors and code complexity
   - repo: https://github.com/PyCQA/flake8
@@ -48,7 +51,7 @@ repos:
             "--extend-ignore",
             "E203,E402,E501,F401,F841,RST2,RST301",
             "--exclude",
-            "logs/*,data/*",
+            "logs/*,data/*,glm_experiments/models/components/transformer.py",
           ]
         additional_dependencies: [flake8-rst-docstrings==0.3.0]
 

configs/model/bert_transformer_base.yaml

@@ -0,0 +1,36 @@
+_target_: glm_experiments.models.bert_lit_module.BERTLitModule
+
+net:
+  _target_: glm_experiments.models.components.bert.BERT
+  embedder:
+    _target_: glm_experiments.models.components.transformer.Embedding
+    vocab_size: 7
+    d_model: 768 # Standard BERT-base size
+  encoder:
+    _target_: glm_experiments.models.components.transformer.Transformer
+    hidden_size: ${..embedder.d_model} # 768
+    n_layers: 12 # CS336 default
+    num_heads: 12 # 12 heads → d_head = 64
+    # d_ff: auto-computed as floor(768 * 8/3 / 64) * 64 = 2048
+    rope_theta: 10000.0
+    is_causal: false
+  layer_norm:
+    _target_: torch.nn.RMSNorm
+    normalized_shape: ${..embedder.d_model}
+  decoder:
+    _target_: glm_experiments.models.components.transformer.Linear
+    d_in: ${..embedder.d_model}
+    d_out: ${..embedder.vocab_size}
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 0.001 # CS336 default
+  weight_decay: 0.1 # CS336 default
+  betas: [0.9, 0.98] # CS336 default (beta1, beta2)
+  eps: 1.0e-9 # CS336 default
+
+scheduler:
+  _target_: transformers.get_constant_schedule_with_warmup
+  _partial_: true
+  num_warmup_steps: 1000 # More warmup for larger model

configs/model/bert_transformer_small.yaml

@@ -0,0 +1,36 @@
+_target_: glm_experiments.models.bert_lit_module.BERTLitModule
+
+net:
+  _target_: glm_experiments.models.components.bert.BERT
+  embedder:
+    _target_: glm_experiments.models.components.transformer.Embedding
+    vocab_size: 7
+    d_model: 128
+  encoder:
+    _target_: glm_experiments.models.components.transformer.Transformer
+    hidden_size: ${..embedder.d_model} # 128
+    n_layers: 6 # Fewer layers for fast iteration
+    num_heads: 8 # 8 heads → d_head = 16
+    # d_ff: auto-computed as floor(128 * 8/3 / 64) * 64 = 320
+    rope_theta: 10000.0
+    is_causal: false # Bidirectional for MLM
+  layer_norm:
+    _target_: torch.nn.RMSNorm # Use RMSNorm to match Transformer
+    normalized_shape: ${..embedder.d_model}
+  decoder:
+    _target_: glm_experiments.models.components.transformer.Linear
+    d_in: ${..embedder.d_model}
+    d_out: ${..embedder.vocab_size}
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 0.001 # CS336 default
+  weight_decay: 0.1 # CS336 default
+  betas: [0.9, 0.98] # CS336 default (beta1, beta2)
+  eps: 1.0e-9 # CS336 default
+
+scheduler:
+  _target_: transformers.get_constant_schedule_with_warmup
+  _partial_: true
+  num_warmup_steps: 100