File size: 5,478 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """
Training Loop & Optimizer Engine for Nova 1.0 (Multi-Core Scaled)
Implements Deep Supervision & Adaptive Computation Time (ACT) loss routines.
Optimized for AMD Ryzen AI Max+ 395 (32 CPU threads) & AMD ROCm.
"""
import os
import math
from typing import Dict
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm
from config.model_config import Nova1Config
from src.model.nova1_hrm import Nova1HRM
class Nova1Trainer:
def __init__(self, model: Nova1HRM, config: Nova1Config, dataloader: DataLoader):
self.model = model
self.config = config
self.dataloader = dataloader
self.device = torch.device(config.device)
self.model.to(self.device)
# Utilize all CPU threads on AMD Ryzen AI Max+ 395
if hasattr(torch, "set_num_threads"):
torch.set_num_threads(config.num_threads)
# CrossEntropyLoss (ignore padding tokens)
self.criterion = nn.CrossEntropyLoss(ignore_index=config.pad_token_id)
# AdamW Optimizer with weight decay
self.optimizer = torch.optim.AdamW(
self.model.parameters(),
lr=config.learning_rate,
weight_decay=config.weight_decay,
betas=(0.9, 0.95)
)
# Automatic Mixed Precision GradScaler (only needed for float16)
self.use_amp = (config.device == "cuda")
self.use_scaler = (config.device == "cuda" and config.dtype == "float16")
self.scaler = torch.amp.GradScaler("cuda", enabled=self.use_scaler)
def adjust_learning_rate(self, step: int, total_steps: int):
"""Cosine annealing learning rate schedule with linear warmup."""
if step < self.config.warmup_steps:
lr = self.config.learning_rate * (step / max(1, self.config.warmup_steps))
else:
progress = (step - self.config.warmup_steps) / max(1, total_steps - self.config.warmup_steps)
lr = self.config.min_learning_rate + 0.5 * (self.config.learning_rate - self.config.min_learning_rate) * (
1.0 + math.cos(math.pi * progress)
)
for param_group in self.optimizer.param_groups:
param_group["lr"] = lr
return lr
def train_epoch(self, epoch: int, total_epochs: int) -> float:
self.model.train()
total_loss = 0.0
steps_per_epoch = len(self.dataloader)
total_training_steps = steps_per_epoch * total_epochs
for step, (input_ids, target_ids) in enumerate(self.dataloader):
global_step = epoch * steps_per_epoch + step
lr = self.adjust_learning_rate(global_step, total_training_steps)
input_ids = input_ids.to(self.device)
target_ids = target_ids.to(self.device)
self.optimizer.zero_grad()
step_loss = 0.0
states = None
# Deep Supervision Loop over M_max segments
for m in range(self.config.max_segments):
with torch.amp.autocast("cuda", enabled=self.use_amp, dtype=self.config.get_torch_dtype()):
logits, (z_H, z_L), q_values = self.model(input_ids, states=states)
# Language Modeling Cross Entropy Loss
loss_lm = self.criterion(logits.view(-1, self.config.vocab_size), target_ids.view(-1))
segment_loss = loss_lm
# Backpropagate segment loss (Equilibrium 1-Step gradient)
if self.use_scaler:
self.scaler.scale(segment_loss).backward()
else:
segment_loss.backward()
step_loss += segment_loss.item()
# Detach hidden states before passing into next supervision segment (Paper Section 2: Deep supervision)
states = (z_H.detach(), z_L.detach())
# Gradient Clipping & Optimizer step
if self.use_scaler:
self.scaler.unscale_(self.optimizer)
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
self.scaler.step(self.optimizer)
self.scaler.update()
else:
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
self.optimizer.step()
avg_step_loss = step_loss / self.config.max_segments
total_loss += avg_step_loss
if (step + 1) % 20 == 0 or (step + 1) == steps_per_epoch:
print(f"Epoch {epoch+1}/{total_epochs} [{step+1}/{steps_per_epoch}] — Step Loss: {avg_step_loss:.4f} | LR: {lr:.6f}", flush=True)
return total_loss / steps_per_epoch
def save_checkpoint(self, checkpoint_path: str):
os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
checkpoint = {
"model_state": self.model.state_dict(),
"optimizer_state": self.optimizer.state_dict(),
"config": self.config
}
torch.save(checkpoint, checkpoint_path)
print(f"Checkpoint saved to {checkpoint_path}")
def load_checkpoint(self, checkpoint_path: str):
checkpoint = torch.load(checkpoint_path, map_location=self.device)
self.model.load_state_dict(checkpoint["model_state"])
self.optimizer.load_state_dict(checkpoint["optimizer_state"])
print(f"Loaded checkpoint from {checkpoint_path}")
|