Update supernova/train.py

supernova/train.py

@@ -10,6 +10,7 @@ import torch.nn as nn
 from torch.utils.data import DataLoader, DistributedSampler
 from torch.utils.tensorboard import SummaryWriter
 from transformers import get_cosine_schedule_with_warmup
+from safetensors.torch import save_file
 
 from .config import ModelConfig
 from .model import SupernovaModel
@@ -45,6 +46,16 @@ def atomic_save(obj: Dict[str, Any], path: str):
     torch.save(obj, tmp)
     os.replace(tmp, path)
 
+def save_safetensors(model_state_dict: Dict[str, torch.Tensor], path: str):
+    """Save model weights in safetensors format."""
+    try:
+        tmp = path + ".tmp"
+        save_file(model_state_dict, tmp)
+        os.replace(tmp, path)
+        print(f"Saved safetensors to {path}")
+    except Exception as e:
+        print(f"Warning: Failed to save safetensors: {e}")
+
 class EMA:
     """Simple exponential moving average of model params (maintains shadow copy)."""
     def __init__(self, model: nn.Module, decay: float = 0.9999):
@@ -100,6 +111,7 @@ def train(
     num_workers: int = 4,
     pin_memory: bool = True,
     compile_model: bool = False,
+    save_safetensors: bool = True,
 ):
     # reproducibility
     torch.manual_seed(seed)
@@ -333,113 +345,6 @@ def train(
                         best_val_loss = mean_val
                         no_improve_steps = 0
                         best_path = os.path.join(out_dir, f"supernova_best_step{step}.pt")
+                        model_state = model.module.state_dict() if ddp else model.state_dict()
                         ckpt = {
-                            "model_state_dict": (model.module.state_dict() if ddp else model.state_dict()),
-                            "optimizer_state_dict": optimizer.state_dict(),
-                            "scheduler_state_dict": scheduler.state_dict(),
-                            "scaler_state_dict": (scaler.state_dict() if scaler else None),
-                            "step": step,
-                            "best_val_loss": best_val_loss,
-                            "config": cfg.__dict__,
-                        }
-                        if not ddp or local_rank == 0:
-                            atomic_save(ckpt, best_path)
-                        print(f"Saved best checkpoint to {best_path}")
-                    else:
-                        no_improve_steps += validate_every
-                        if no_improve_steps >= early_stop_patience:
-                            print("Early stopping triggered.")
-                            step = max_steps
-                            break
-
-                # periodic checkpointing
-                if save_every and step % save_every == 0 and (not ddp or local_rank == 0):
-                    ckpt_path = os.path.join(out_dir, f"supernova_step{step}.pt")
-                    ckpt = {
-                        "model_state_dict": (model.module.state_dict() if ddp else model.state_dict()),
-                        "optimizer_state_dict": optimizer.state_dict(),
-                        "scheduler_state_dict": scheduler.state_dict(),
-                        "scaler_state_dict": (scaler.state_dict() if scaler else None),
-                        "step": step,
-                        "best_val_loss": best_val_loss,
-                        "config": cfg.__dict__,
-                    }
-                    atomic_save(ckpt, ckpt_path)
-                    print(f"Saved checkpoint {ckpt_path}")
-
-                if step >= max_steps:
-                    break
-
-        if step >= max_steps:
-            break
-
-    # final save
-    if not ddp or local_rank == 0:
-        ckpt_path = os.path.join(out_dir, f"supernova_final_step{step}.pt")
-        ckpt = {
-            "model_state_dict": (model.module.state_dict() if ddp else model.state_dict()),
-            "optimizer_state_dict": optimizer.state_dict(),
-            "scheduler_state_dict": scheduler.state_dict(),
-            "scaler_state_dict": (scaler.state_dict() if scaler else None),
-            "step": step,
-            "best_val_loss": best_val_loss,
-            "config": cfg.__dict__,
-        }
-        atomic_save(ckpt, ckpt_path)
-        print(f"Saved final checkpoint to {ckpt_path}")
-
-    if writer:
-        writer.close()
-
-if __name__ == "__main__":
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--config", required=True)
-    ap.add_argument("--data-config", required=True)
-    ap.add_argument("--seq-len", type=int, default=1024)
-    ap.add_argument("--batch-size", type=int, default=16)
-    ap.add_argument("--grad-accum", type=int, default=8)
-    ap.add_argument("--lr", type=float, default=3e-4)
-    ap.add_argument("--warmup-steps", type=int, default=2000)
-    ap.add_argument("--max-steps", type=int, default=100000)
-    ap.add_argument("--save-every", type=int, default=10000)
-    ap.add_argument("--out-dir", type=str, default="checkpoints")
-    ap.add_argument("--seed", type=int, default=42)
-    ap.add_argument("--validate-every", type=int, default=1000)
-    ap.add_argument("--val-steps", type=int, default=100)
-    ap.add_argument("--clip-grad-norm", type=float, default=1.0)
-    ap.add_argument("--resume-from", type=str, default=None)
-    ap.add_argument("--use-ema", action="store_true")
-    ap.add_argument("--ema-decay", type=float, default=0.9999)
-    ap.add_argument("--no-tensorboard", dest="use_tensorboard", action="store_false")
-    ap.add_argument("--ddp", action="store_true", help="enable DistributedDataParallel; use with torchrun")
-    ap.add_argument("--local-rank", type=int, default=0)
-    ap.add_argument("--num-workers", type=int, default=4)
-    ap.add_argument("--pin-memory", type=bool, default=True)
-    ap.add_argument("--compile", dest="compile_model", action="store_true")
-    args = ap.parse_args()
-
-    train(
-        config_path=args.config,
-        data_config_path=args.data_config,
-        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=args.use_ema,
-        ema_decay=args.ema_decay,
-        resume_from=args.resume_from,
-        use_tensorboard=args.use_tensorboard,
-        ddp=args.ddp,
-        local_rank=args.local_rank,
-        num_workers=args.num_workers,
-        pin_memory=args.pin_memory,
-        compile_model=args.compile_model,
-    )
+                            "model_state_dict": model_state