fix: robust configuration handling and tokenizer support

- Updated training scripts (stage 1-4) to safely handle missing 'system' configuration section by initializing defaults.
- Added support for 'tokenizer' configuration in stage 1, allowing custom tokenizers to be passed.
- Included user-provided updates to configuration YAML files (adding 'system' section).

Author: Pomilon

configs/baseline_170m.yaml

@@ -31,5 +31,6 @@ data:
   data_dir: "data/text_corpus"
   traces_path: "data/train_traces.jsonl"
 
-tokenizer: "gpt2"
-seed: 42
+system:
+  seed: 42
+  device: "cuda"

configs/baseline_27m.yaml

@@ -31,5 +31,6 @@ data:
   data_dir: "data/text_corpus"
   traces_path: "data/train_traces.jsonl"
 
-tokenizer: "gpt2"
-seed: 42
+system:
+  seed: 42
+  device: "cuda"

configs/baseline_350m.yaml

@@ -31,5 +31,6 @@ data:
   data_dir: "data/text_corpus"
   traces_path: "data/train_traces.jsonl"
 
-tokenizer: "gpt2"
-seed: 42
+system:
+  seed: 42
+  device: "cuda"

configs/small.yaml

@@ -36,8 +36,7 @@ data:
   traces_path: "data/train_traces.jsonl"
 
 tokenizer: null
-seed: 42
 
 system:
-  device: "cpu"
   seed: 42
+  device: "cuda"

configs/test_config.yaml

@@ -32,4 +32,7 @@ data:
   traces_path: "data/train_traces.jsonl"
 
 tokenizer: null
-seed: 42
+
+system:
+  seed: 42
+  device: "cuda"

scripts/training/stage_1_backbone.py

@@ -92,15 +92,17 @@ def main():
     if config['training'].get('use_amp', False):
         cmd.append('--amp')
 
+    system_config = config.get('system', {})
+    
     if args.device:
         cmd.extend(['--device', args.device])
-    elif 'device' in config['system']:
-        cmd.extend(['--device', config['system']['device']])
+    elif system_config.get('device'):
+        cmd.extend(['--device', system_config['device']])
 
     if args.seed is not None:
         cmd.extend(['--seed', str(args.seed)])
-    elif 'seed' in config['system']:
-        cmd.extend(['--seed', str(config['system']['seed'])])
+    elif system_config.get('seed'):
+        cmd.extend(['--seed', str(system_config['seed'])])
 
     if args.no_wandb:
         cmd.append('--no-wandb')
@@ -109,6 +111,8 @@ def main():
 
     if args.tokenizer:
         cmd.extend(['--tokenizer', args.tokenizer])
+    elif config.get('tokenizer'):
+        cmd.extend(['--tokenizer', config['tokenizer']])
 
     if args.resume:
         cmd.extend(['--resume', args.resume])

scripts/training/stage_2_dynamics.py

@@ -36,6 +36,10 @@ def main():
 
     config = load_config(args.config)
 
+    # Ensure system config exists
+    if 'system' not in config:
+        config['system'] = {'seed': 42, 'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+    
     # Apply Overrides
     if args.epochs: config['dynamics']['dynamics_epochs'] = args.epochs
     if args.batch_size: config['training']['batch_size'] = args.batch_size

scripts/training/stage_3_value_head.py

@@ -134,6 +134,10 @@ async def main_async():
 
     config = load_config(args.config)
 
+    # Ensure system config exists
+    if 'system' not in config:
+        config['system'] = {'seed': 42, 'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+    
     # Apply Overrides
     if args.epochs: config['training']['value_training']['epochs'] = args.epochs
     if args.lr: config['training']['value_training']['lr'] = args.lr

scripts/training/stage_4_assembly.py

@@ -40,6 +40,10 @@ def main():
 
     config_dict = load_config(args.config)
 
+    # Ensure system config exists
+    if 'system' not in config_dict:
+        config_dict['system'] = {'seed': 42, 'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+    
     # Apply Overrides
     if args.seed: config_dict['system']['seed'] = args.seed
     if args.device: config_dict['system']['device'] = args.device