hf_free_training.py

"""
Free H200 Training Script for Nano-Coder
Optimized for HF's free 4-minute daily H200 access
"""

import os
import time
import math
import pickle
from contextlib import nullcontext

import numpy as np
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group

from model import GPTConfig, GPT

# Hugging Face specific imports
from huggingface_hub import HfApi, login
import wandb

# -----------------------------------------------------------------------------
# Configuration optimized for FREE H200 (4 minutes daily)
# I/O
out_dir = 'out-nano-coder-free'
eval_interval = 50  # Very frequent evaluation for short runs
log_interval = 2
eval_iters = 10  # Fewer eval iterations
eval_only = False
always_save_checkpoint = True
init_from = 'scratch'

# wandb logging - enabled for HF
wandb_log = True
wandb_project = 'nano-coder-free'
wandb_run_name = 'nano-coder-h200-free'

# data
dataset = 'python-codes-25k'
gradient_accumulation_steps = 1 * 8  # Minimal for H200
batch_size = 64  # Larger batch size for H200 efficiency
block_size = 512  # Smaller context for faster training

# model - smaller for free tier
n_layer = 6  # Reduced from 12
n_head = 6   # Reduced from 12
n_embd = 384 # Reduced from 768
dropout = 0.1
bias = False

# optimizer - optimized for H200
learning_rate = 1e-3  # Higher learning rate for faster convergence
max_iters = 1000  # Limited iterations for 4-minute runs
weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0

# learning rate decay - faster for short runs
decay_lr = True
warmup_iters = 100  # Shorter warmup
lr_decay_iters = 1000
min_lr = 1e-4

# DDP settings
backend = 'nccl'

# system
device = 'cuda' if torch.cuda.is_available() else 'cpu'
dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
compile = True

# HF specific
hf_repo_id = "mlopez6132/nano-coder-free"  # Free tier repo
push_to_hub = True

# Time tracking for 4-minute limit
start_time = time.time()
MAX_TRAINING_TIME = 3.5 * 60  # 3.5 minutes to be safe

# -----------------------------------------------------------------------------
config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
exec(open('configurator.py').read())
config = {k: globals()[k] for k in config_keys}

# -----------------------------------------------------------------------------

# HF setup
if push_to_hub:
    # Check if HF_TOKEN environment variable is set
    if os.environ.get('HF_TOKEN'):
        login(token=os.environ.get('HF_TOKEN'))
    else:
        # Try to login without token (will use cached credentials)
        try:
            login()
        except Exception as e:
            print(f"Warning: Could not login to HF Hub: {e}")
            print("Continuing without HF Hub upload...")
            push_to_hub = False
    
    if push_to_hub:
        api = HfApi()

# various inits, derived attributes, I/O setup
ddp = int(os.environ.get('RANK', -1)) != -1
if ddp:
    init_process_group(backend=backend)
    ddp_rank = int(os.environ['RANK'])
    ddp_local_rank = int(os.environ['LOCAL_RANK'])
    ddp_world_size = int(os.environ['WORLD_SIZE'])
    device = f'cuda:{ddp_local_rank}'
    torch.cuda.set_device(device)
    master_process = ddp_rank == 0
    seed_offset = ddp_rank
    assert gradient_accumulation_steps % ddp_world_size == 0
    gradient_accumulation_steps //= ddp_world_size
else:
    master_process = True
    seed_offset = 0
    ddp_world_size = 1

tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size
print(f"tokens per iteration will be: {tokens_per_iter:,}")
print(f"FREE H200 TRAINING - MAX TIME: {MAX_TRAINING_TIME/60:.1f} minutes")

if master_process:
    os.makedirs(out_dir, exist_ok=True)

torch.manual_seed(1337 + seed_offset)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
device_type = 'cuda' if 'cuda' in device else 'cpu'
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

# data loader
data_dir = os.path.join('data', dataset)
def get_batch(split):
    if split == 'train':
        data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
    else:
        data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
    if device_type == 'cuda':
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

# init these up here, can override if init_from='resume'
iter_num = 0
best_val_loss = 1e9

# attempt to derive vocab_size from the dataset
meta_path = os.path.join(data_dir, 'meta.pkl')
meta_vocab_size = None
if os.path.exists(meta_path):
    with open(meta_path, 'rb') as f:
        meta = pickle.load(f)
    meta_vocab_size = meta['vocab_size']
    print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})")

# model init
model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,
                  bias=bias, vocab_size=None, dropout=dropout)

if init_from == 'scratch':
    print("Initializing a new nano-coder model from scratch (FREE TIER)")
    if meta_vocab_size is None:
        print("defaulting to vocab_size of GPT-2 to 50304")
    model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
elif init_from == 'resume':
    print(f"Resuming training from {out_dir}")
    ckpt_path = os.path.join(out_dir, 'ckpt.pt')
    checkpoint = torch.load(ckpt_path, map_location=device)
    checkpoint_model_args = checkpoint['model_args']
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = checkpoint_model_args[k]
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k,v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    model.load_state_dict(state_dict)
    iter_num = checkpoint['iter_num']
    best_val_loss = checkpoint['best_val_loss']
elif init_from.startswith('gpt2'):
    print(f"Initializing from OpenAI GPT-2 weights: {init_from}")
    override_args = dict(dropout=dropout)
    model = GPT.from_pretrained(init_from, override_args)
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = getattr(model.config, k)

if block_size < model.config.block_size:
    model.crop_block_size(block_size)
    model_args['block_size'] = block_size

model.to(device)

# initialize a GradScaler
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))

# optimizer
optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
if init_from == 'resume':
    optimizer.load_state_dict(checkpoint['optimizer'])
checkpoint = None

# compile the model
if compile:
    print("compiling the model... (takes a ~minute)")
    unoptimized_model = model
    model = torch.compile(model)

# wrap model into DDP container
if ddp:
    model = DDP(model, device_ids=[ddp_local_rank])

# helps estimate an arbitrarily accurate loss over either split using many batches
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            with ctx:
                logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

# learning rate decay scheduler (cosine with warmup)
def get_lr(it):
    if it < warmup_iters:
        return learning_rate * (it + 1) / (warmup_iters + 1)
    if it > lr_decay_iters:
        return min_lr
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

# logging
if wandb_log and master_process:
    wandb.init(project=wandb_project, name=wandb_run_name, config=config)

# HF checkpoint upload function
def upload_checkpoint_to_hf(checkpoint_path, iter_num):
    if push_to_hub and master_process:
        try:
            # Create a unique filename
            filename = f"checkpoint_iter_{iter_num}.pt"
            file_path = os.path.join(out_dir, filename)
            
            # Copy checkpoint with new name
            import shutil
            shutil.copy2(checkpoint_path, file_path)
            
            # Upload to HF
            api.upload_file(
                path_or_fileobj=file_path,
                path_in_repo=filename,
                repo_id=hf_repo_id,
                repo_type="model"
            )
            print(f"Uploaded checkpoint to HF: {filename}")
            
            # Clean up local copy
            os.remove(file_path)
        except Exception as e:
            print(f"Failed to upload checkpoint: {e}")

# training loop
print("Starting FREE H200 nano-coder training...")
X, Y = get_batch('train')
t0 = time.time()
local_iter_num = 0
raw_model = model.module if ddp else model
running_mfu = -1.0

while True:
    # Check time limit
    elapsed_time = time.time() - start_time
    if elapsed_time > MAX_TRAINING_TIME:
        print(f"\n⏰ TIME LIMIT REACHED! Training stopped after {elapsed_time/60:.1f} minutes")
        break
    
    # determine and set the learning rate for this iteration
    lr = get_lr(iter_num) if decay_lr else learning_rate
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # evaluate the loss on train/val sets and write checkpoints
    if iter_num % eval_interval == 0 and master_process:
        losses = estimate_loss()
        remaining_time = MAX_TRAINING_TIME - elapsed_time
        print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}, time left: {remaining_time/60:.1f}min")
        if wandb_log:
            wandb.log({
                "iter": iter_num,
                "train/loss": losses['train'],
                "val/loss": losses['val'],
                "lr": lr,
                "mfu": running_mfu*100,
                "elapsed_time": elapsed_time,
                "remaining_time": remaining_time,
            })
        if losses['val'] < best_val_loss or always_save_checkpoint:
            best_val_loss = losses['val']
            if iter_num > 0:
                checkpoint = {
                    'model': raw_model.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'model_args': model_args,
                    'iter_num': iter_num,
                    'best_val_loss': best_val_loss,
                    'config': config,
                }
                checkpoint_path = os.path.join(out_dir, 'ckpt.pt')
                print(f"saving checkpoint to {out_dir}")
                torch.save(checkpoint, checkpoint_path)
                
                # Upload to HF every 200 iterations (frequent for short runs)
                if iter_num % 200 == 0:
                    upload_checkpoint_to_hf(checkpoint_path, iter_num)
    if iter_num == 0 and eval_only:
        break

    # forward backward update
    for micro_step in range(gradient_accumulation_steps):
        if ddp:
            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
        with ctx:
            logits, loss = model(X, Y)
            loss = loss / gradient_accumulation_steps
        X, Y = get_batch('train')
        scaler.scale(loss).backward()
    
    # clip the gradient
    if grad_clip != 0.0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    
    # step the optimizer and scaler
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

    # timing and logging
    t1 = time.time()
    dt = t1 - t0
    t0 = t1
    if iter_num % log_interval == 0 and master_process:
        lossf = loss.item() * gradient_accumulation_steps
        if local_iter_num >= 5:
            mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)
            running_mfu = mfu if running_mfu == -1.0 else 0.9*running_mfu + 0.1*mfu
        remaining_time = MAX_TRAINING_TIME - elapsed_time
        print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%, remaining: {remaining_time/60:.1f}min")
    iter_num += 1
    local_iter_num += 1

    # termination conditions
    if iter_num > max_iters:
        break

if ddp:
    destroy_process_group()

# Final upload
if push_to_hub and master_process:
    upload_checkpoint_to_hf(os.path.join(out_dir, 'ckpt.pt'), 'final')

total_time = time.time() - start_time
print(f"\n🎉 FREE H200 TRAINING COMPLETED!")
print(f"Total training time: {total_time/60:.1f} minutes")
print(f"Total iterations: {iter_num}")
print(f"Final validation loss: {best_val_loss:.4f}")
print(f"Model saved to: {out_dir}")