Update supernova/train.py

supernova/train.py

@@ -1,4 +1,3 @@
-# train.py (improved)
 import argparse
 import json
 import math
@@ -138,14 +137,13 @@ 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)
-
     sampler = DistributedSampler(ds) if ddp else None
+
+    # NOTE: NO shuffle for IterableDataset!
     dl = DataLoader(
         ds,
         batch_size=batch_size,
-        shuffle=(sampler is None),
         sampler=sampler,
         num_workers=num_workers,
         pin_memory=pin_memory,
@@ -153,7 +151,7 @@ def train(
         drop_last=True,
     )
 
-    # optimizer with simple parameter grouping to avoid weight decay on norms/bias
+    # optimizer
     def param_groups(model):
         decay, no_decay = [], []
         for n, p in model.named_parameters():
@@ -169,22 +167,14 @@ def train(
         ]
 
     optimizer = torch.optim.AdamW(param_groups(model), lr=lr, betas=(0.9, 0.95), eps=1e-8)
-
-    # scheduler
     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"))
-
-    # EMA
     ema = EMA(model if not ddp else model.module, decay=ema_decay) if use_ema else None
 
-    # logging + checkpoint dir
     os.makedirs(out_dir, exist_ok=True)
     writer = SummaryWriter(log_dir=os.path.join(out_dir, "runs")) if use_tensorboard and (not ddp or local_rank == 0) else None
 
-    # validation dataset (simple split: user should provide a separate validation YAML ideally)
-    # TODO: Implement a proper validation dataset pipeline. For now, we use a small random subset of training data.
+    # validation
     val_ds = None
     val_dl = None
 
@@ -194,7 +184,6 @@ def train(
     if resume_from and os.path.exists(resume_from):
         ckpt = torch.load(resume_from, map_location=device)
         model_state = ckpt["model_state_dict"]
-        # if ddp, load into module
         target = model.module if ddp else model
         target.load_state_dict(model_state)
         optimizer.load_state_dict(ckpt.get("optimizer_state_dict", {}))
@@ -234,7 +223,6 @@ def train(
             running_loss += loss.item()
 
             if micro % grad_accum == 0:
-                # gradient clipping
                 if clip_grad_norm is not None:
                     scaler.unscale_(optimizer)
                     torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad_norm)
@@ -246,7 +234,6 @@ def train(
 
                 if ema:
                     ema.update(model if not ddp else model.module)
-
                 step += 1
 
                 # logging
@@ -266,24 +253,22 @@ def train(
                 # periodic validation
                 if validate_every and step % validate_every == 0:
                     if val_dl is None:
-                        # Create a proper validation dataset with a small subset of training sources
                         val_sources = []
                         for source in sources[:min(3, len(sources))]:
                             val_source = DataSource(
                                 name=f"{source.name}_val",
-                                hf_path="wikitext",  # Use a reliable, small dataset for validation
+                                hf_path="wikitext",
                                 hf_name="wikitext-2-v1",
                                 split="validation",
                                 text_field="text",
                                 weight=1,
-                                streaming=False  # Don't stream validation data
+                                streaming=False
                             )
                             val_sources.append(val_source)
                         val_ds = TokenChunkDataset(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)
+                        val_dl = DataLoader(val_ds, batch_size=batch_size, num_workers=0, pin_memory=True, drop_last=False)
 
                     model.eval()
-                    # optionally swap in EMA weights for evaluation
                     if ema:
                         ema.store(model if not ddp else model.module)
                         ema.copy_to(model if not ddp else model.module)
@@ -303,12 +288,10 @@ def train(
                         writer.add_scalar("val/loss", mean_val, step)
                     print(f"[eval] step={step} val_loss={mean_val:.6f}")
 
-                    # restore weights
                     if ema:
                         ema.restore(model if not ddp else model.module)
                     model.train()
 
-                    # early stop / best model saving
                     if mean_val < best_val_loss:
                         best_val_loss = mean_val
                         no_improve_steps = 0