fixed training datasets

Supernova25million

@@ -0,0 +1 @@
+Subproject commit 288c71bea4b8740818638d0e2dae0a647da22763

supernova/train.py

@@ -139,7 +139,12 @@ def train(
     # dataset and dataloader
     sources = load_sources_from_yaml(data_config_path)
     # TODO: improve TokenChunkDataset to perform token-packing (pack multiple short examples into one sequence)
-    ds = TokenChunkDataset(tok, sources, seq_len=seq_len, eos_token_id=tok.eos_token_id)
+    ds = TokenChunkDataset(
+        tokenizer=tok, 
+        sources=sources, 
+        seq_len=seq_len, 
+        eos_token_id=tok.eos_token_id
+    )
 
     sampler = DistributedSampler(ds) if ddp else None
     dl = DataLoader(
@@ -174,7 +179,10 @@ def train(
     scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=max_steps)
 
     # AMP scaler
-    scaler = torch.cuda.amp.GradScaler(enabled=(device.type == "cuda"))
+    if device.type == "cuda":
+        scaler = torch.amp.GradScaler('cuda', enabled=True)
+    else:
+        scaler = torch.amp.GradScaler('cpu', enabled=False)
 
     # EMA
     ema = EMA(model if not ddp else model.module, decay=ema_decay) if use_ema else None
@@ -225,7 +233,8 @@ def train(
             x = x.to(device, non_blocking=True)
             y = y.to(device, non_blocking=True)
 
-            with torch.cuda.amp.autocast(enabled=(device.type == "cuda")):
+            device_type = 'cuda' if device.type == 'cuda' else 'cpu'
+            with torch.amp.autocast(device_type, enabled=(device.type == "cuda")):
                 logits, loss = model(x, y)
                 loss = loss / grad_accum
 
@@ -268,7 +277,15 @@ def train(
                     if val_dl is None:
                         # quick in-memory val split: take first N batches (user should replace with real val)
                         # NOTE: for production, create a dedicated validation dataset.
-                        val_ds = TokenChunkDataset(tok, sources[: max(1, len(sources) // 20)], seq_len=seq_len, eos_token_id=tok.eos_token_id)
+                        val_sources = sources[: max(1, len(sources) // 20)]
+                        if not val_sources:
+                            val_sources = sources[:1]  # fallback to at least one source
+                        val_ds = TokenChunkDataset(
+                            tokenizer=tok,
+                            sources=val_sources, 
+                            seq_len=seq_len, 
+                            eos_token_id=tok.eos_token_id
+                        )
                         val_dl = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True, drop_last=False)
 
                     model.eval()
@@ -284,7 +301,8 @@ def train(
                                 break
                             vx = vx.to(device)
                             vy = vy.to(device)
-                            with torch.cuda.amp.autocast(enabled=(device.type == "cuda")):
+                            device_type = 'cuda' if device.type == 'cuda' else 'cpu'
+                            with torch.amp.autocast(device_type, enabled=(device.type == "cuda")):
                                 _, vloss = model(vx, vy)
                             val_losses.append(float(vloss.detach().cpu().item()))
                     mean_val = float(sum(val_losses) / max(1, len(val_losses)))

train_main.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""
+Main training script - can be run directly without import issues.
+This script imports and runs the training function from the supernova package.
+"""
+
+import argparse
+import sys
+import os
+
+# Add the current directory to Python path to ensure supernova package can be imported
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from supernova.train import train
+
+def main():
+    parser = argparse.ArgumentParser(description="Train Supernova 25M model")
+    parser.add_argument("--config", required=True, help="Path to model config JSON")
+    parser.add_argument("--data", required=True, help="Path to data config YAML")
+    parser.add_argument("--seq-len", type=int, default=1024, help="Sequence length")
+    parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
+    parser.add_argument("--grad-accum", type=int, default=8, help="Gradient accumulation steps")
+    parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
+    parser.add_argument("--warmup-steps", type=int, default=2000, help="Warmup steps")
+    parser.add_argument("--max-steps", type=int, default=100000, help="Maximum training steps")
+    parser.add_argument("--save-every", type=int, default=10000, help="Save checkpoint every N steps")
+    parser.add_argument("--out-dir", default="checkpoints", help="Output directory")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--validate-every", type=int, default=1000, help="Validate every N steps")
+    parser.add_argument("--val-steps", type=int, default=100, help="Validation steps")
+    parser.add_argument("--clip-grad-norm", type=float, default=1.0, help="Gradient clipping norm")
+    parser.add_argument("--no-ema", action="store_true", help="Disable EMA")
+    parser.add_argument("--ema-decay", type=float, default=0.9999, help="EMA decay rate")
+    parser.add_argument("--resume-from", help="Resume from checkpoint")
+    parser.add_argument("--no-tensorboard", action="store_true", help="Disable tensorboard")
+    parser.add_argument("--ddp", action="store_true", help="Use distributed training")
+    parser.add_argument("--local-rank", type=int, default=0, help="Local rank for DDP")
+    parser.add_argument("--num-workers", type=int, default=4, help="DataLoader workers")
+    parser.add_argument("--no-pin-memory", action="store_true", help="Disable pin memory")
+    parser.add_argument("--compile-model", action="store_true", help="Use torch.compile")
+    
+    args = parser.parse_args()
+    
+    # Call the training function
+    train(
+        config_path=args.config,
+        data_config_path=args.data,
+        seq_len=args.seq_len,
+        batch_size=args.batch_size,
+        grad_accum=args.grad_accum,
+        lr=args.lr,
+        warmup_steps=args.warmup_steps,
+        max_steps=args.max_steps,
+        save_every=args.save_every,
+        out_dir=args.out_dir,
+        seed=args.seed,
+        validate_every=args.validate_every,
+        val_steps=args.val_steps,
+        clip_grad_norm=args.clip_grad_norm,
+        use_ema=not args.no_ema,
+        ema_decay=args.ema_decay,
+        resume_from=args.resume_from,
+        use_tensorboard=not args.no_tensorboard,
+        ddp=args.ddp,
+        local_rank=args.local_rank,
+        num_workers=args.num_workers,
+        pin_memory=not args.no_pin_memory,
+        compile_model=args.compile_model,
+    )
+
+if __name__ == "__main__":
+    main()