File size: 4,054 Bytes
23ea6bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """
Script 2: Pre-train Scaled Nova 1.0 Model from Scratch (High Capacity)
Leverages full system resources: 32 CPU Threads, 119GB RAM, AMD ROCm.
"""
import os
import argparse
import torch
from config.model_config import Nova1Config
from src.tokenizer.bpe_tokenizer import Nova1Tokenizer
from src.model.nova1_hrm import Nova1HRM
from src.dataset.hf_dataset import create_hf_dataloader
from src.trainer.train_loop import Nova1Trainer
def main():
parser = argparse.ArgumentParser(description="Pre-train Scaled Nova 1.0 Model")
parser.add_argument("--dataset", type=str, default="tatsu-lab/alpaca", help="HuggingFace dataset name")
parser.add_argument("--subset", type=str, default=None, help="HuggingFace dataset subset/config name")
parser.add_argument("--tokenizer_path", type=str, default="checkpoints/nova1_tokenizer.json", help="Tokenizer path")
parser.add_argument("--checkpoint_dir", type=str, default="checkpoints", help="Directory to save model checkpoints")
parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs")
parser.add_argument("--batch_size", type=int, default=16, help="Batch size per step")
parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
parser.add_argument("--d_model", type=int, default=768, help="Model hidden dimension")
parser.add_argument("--max_samples", type=int, default=25000, help="Number of HuggingFace dataset samples")
args = parser.parse_args()
# 1. Load tokenizer
if not os.path.exists(args.tokenizer_path):
print(f"Tokenizer not found at {args.tokenizer_path}. Please run scripts/train_tokenizer.py first!")
return
tokenizer = Nova1Tokenizer.load(args.tokenizer_path)
print(f"Loaded subword tokenizer with vocab size {tokenizer.vocab_size}.")
# 2. Build Nova 1.0 Config (Scaled)
config = Nova1Config(
vocab_size=tokenizer.vocab_size,
d_model=args.d_model,
learning_rate=args.lr,
pad_token_id=tokenizer.pad_id,
unk_token_id=tokenizer.unk_id,
bos_token_id=tokenizer.bos_id,
eos_token_id=tokenizer.eos_id,
mask_token_id=tokenizer.mask_id
)
print(f"Initialized Scaled Nova 1.0 Config: Device={config.device}, dtype={config.dtype}, d_model={config.d_model}, max_seq_len={config.max_seq_len}")
# 3. Create DataLoader from Hugging Face dataset
print(f"Preparing high-capacity DataLoader for dataset '{args.dataset}' ({args.max_samples:,} samples)...")
dataloader = create_hf_dataloader(
dataset_name=args.dataset,
subset=args.subset,
tokenizer=tokenizer,
max_seq_len=config.max_seq_len,
batch_size=args.batch_size,
max_samples=args.max_samples
)
# 4. Instantiate Nova 1.0 Model & Trainer
model = Nova1HRM(config)
num_params = sum(p.numel() for p in model.parameters())
print(f"\n=======================================================")
print(f"🚀 Scaled Nova 1.0 Model Architecture Built!")
print(f"Total Trainable Parameters: {num_params:,} (~{num_params / 1e6:.1f} Million)")
print(f"=======================================================\n")
trainer = Nova1Trainer(model=model, config=config, dataloader=dataloader)
# 5. Run Training Loop
print(f"Starting Nova 1.0 Pre-training from scratch using {config.num_threads} CPU threads & {config.device.upper()}...")
for epoch in range(args.epochs):
avg_loss = trainer.train_epoch(epoch, args.epochs)
print(f"Epoch {epoch+1}/{args.epochs} Complete — Average Loss: {avg_loss:.4f}")
# Save checkpoint after each epoch
ckpt_path = os.path.join(args.checkpoint_dir, f"nova1_epoch_{epoch+1}.pt")
trainer.save_checkpoint(ckpt_path)
# Save final model checkpoint
final_path = os.path.join(args.checkpoint_dir, "nova1_final.pt")
trainer.save_checkpoint(final_path)
print(f"\nScaled Nova 1.0 training complete! Final model saved to '{final_path}'.")
if __name__ == "__main__":
main()
|