""" ZeroGPU Training Script for Nano-Coder Optimized for Hugging Face's ZeroGPU free compute (H200, 70GB) """ 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 ZeroGPU (H200, 70GB, FREE) # I/O out_dir = 'out-nano-coder-zerogpu' eval_interval = 100 # Frequent evaluation for monitoring log_interval = 5 eval_iters = 20 eval_only = False always_save_checkpoint = True init_from = 'scratch' # wandb logging - enabled for ZeroGPU wandb_log = True wandb_project = 'nano-coder-zerogpu' wandb_run_name = 'nano-coder-zerogpu-training' # data dataset = 'python-codes-25k' gradient_accumulation_steps = 2 * 8 # Optimized for H200 batch_size = 48 # Larger batch size for H200 efficiency block_size = 1024 # Full context length # model - optimized for ZeroGPU H200 n_layer = 12 # Full model n_head = 12 # Full model n_embd = 768 # Full model dropout = 0.1 bias = False # optimizer - optimized for H200 learning_rate = 6e-4 # Standard GPT-2 learning rate max_iters = 10000 # More iterations for ZeroGPU weight_decay = 1e-1 beta1 = 0.9 beta2 = 0.95 grad_clip = 1.0 # learning rate decay decay_lr = True warmup_iters = 1000 lr_decay_iters = 10000 min_lr = 6e-5 # 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-zerogpu" # ZeroGPU repo push_to_hub = True # ZeroGPU specific - no time limits! print("šŸš€ ZEROGPU TRAINING - NO TIME LIMITS!") # ----------------------------------------------------------------------------- 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"ZEROGPU H200 TRAINING - NO TIME LIMITS!") 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 (ZEROGPU)") 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 ZEROGPU 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: # 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() print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}") if wandb_log: wandb.log({ "iter": iter_num, "train/loss": losses['train'], "val/loss": losses['val'], "lr": lr, "mfu": running_mfu*100, }) 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 1000 iterations if iter_num % 1000 == 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 print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%") 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šŸŽ‰ ZEROGPU 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}")