Merge 33ddc13ed4 into 4a87a0d19f

Fix comment: rotary embeddings final dimension size

nanochat/configurator.py

@@ -1,56 +1,82 @@
 """
-Poor Man's Configurator. Probably a terrible idea. Example usage:
+Poor Man's Configurator v3. Clean refactored version. Example usage:
 $ python train.py config/override_file.py --batch_size=32
-this will first run config/override_file.py, then override batch_size to 32
 
-The code in this file will be run as follows from e.g. train.py:
->>> exec(open('configurator.py').read())
-
-So it's not a Python module, it's just shuttling this code away from train.py
-The code in this script then overrides the globals()
-
-I know people are not going to love this, I just really dislike configuration
-complexity and having to prepend config. to every single variable. If someone
-comes up with a better simple Python solution I am all ears.
+Improved version with better separation of concerns and cleaner architecture
+while maintaining the same simple usage pattern.
 """
 
 import os
 import sys
 from ast import literal_eval
 
-def print0(s="",**kwargs):
-    ddp_rank = int(os.environ.get('RANK', 0))
-    if ddp_rank == 0:
+
+def print0(s="", **kwargs):
+    """Print only from rank 0 in distributed settings"""
+    if int(os.environ.get("RANK", 0)) == 0:
         print(s, **kwargs)
 
-for arg in sys.argv[1:]:
-    if '=' not in arg:
-        # assume it's the name of a config file
-        assert not arg.startswith('--')
-        config_file = arg
-        print0(f"Overriding config with {config_file}:")
-        with open(config_file) as f:
-            print0(f.read())
-        exec(open(config_file).read())
-    else:
-        # assume it's a --key=value argument
-        assert arg.startswith('--')
-        key, val = arg.split('=')
-        key = key[2:]
-        if key in globals():
+
+class ConfigManager:
+    """Clean configurator with explicit global injection control"""
+
+    def __init__(self):
+        self.config = {}
+
+    def load(self, config_file=None, **overrides):
+        """Load configuration from file and apply overrides"""
+        if config_file:
+            with open(config_file) as f:
+                exec(f.read(), {}, self.config)
+        self.config.update(overrides)
+        return self
+
+    def inject_globals(self):
+        """Inject config vars into global namespace"""
+        for k, v in self.config.items():
+            if not k.startswith("_"):
+                globals()[k] = v
+
+
+def parse_args(args=None):
+    """Parse command line arguments like original version"""
+    args = args or sys.argv
+    config_file, overrides = None, {}
+
+    for arg in args[1:]:
+        if "=" not in arg:
+            # Config file
+            assert not arg.startswith("--"), f"Invalid config file: {arg}"
+            config_file = arg
+            print0(f"Overriding config with {config_file}:")
+            with open(config_file) as f:
+                print0(f.read())
+        else:
+            # Key=value override
+            assert arg.startswith("--"), f"Override must start with --: {arg}"
+            key, val = arg.split("=", 1)
+            key = key[2:]
+            
+            if key not in globals():
+                raise ValueError(f"Unknown config key: {key}")
+            
+            # Try to parse the value
             try:
-                # attempt to eval it it (e.g. if bool, number, or etc)
                 attempt = literal_eval(val)
             except (SyntaxError, ValueError):
-                # if that goes wrong, just use the string
                 attempt = val
-            # ensure the types match ok
+            
+            # Type check if global has a non-None value
             if globals()[key] is not None:
-                attempt_type = type(attempt)
-                default_type = type(globals()[key])
+                default_type, attempt_type = type(globals()[key]), type(attempt)
                 assert attempt_type == default_type, f"Type mismatch: {attempt_type} != {default_type}"
-            # cross fingers
+            
             print0(f"Overriding: {key} = {attempt}")
-            globals()[key] = attempt
-        else:
-            raise ValueError(f"Unknown config key: {key}")
+            overrides[key] = attempt
+
+    return config_file, overrides
+
+
+# Execute configuration
+config_file, overrides = parse_args()
+ConfigManager().load(config_file, **overrides).inject_globals()

nanochat/gpt.py

@@ -244,7 +244,7 @@ class GPT(nn.Module):
     def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
         B, T = idx.size()
 
-        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim))
+        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim/2))
         assert T <= self.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {self.cos.size(1)}"
         assert idx.device == self.cos.device, f"Rotary embeddings and idx are on different devices: {idx.device} != {self.cos.device}"
         assert self.cos.dtype == torch.bfloat16, "Rotary embeddings must be in bfloat16"