import gradio as gr import spaces import torch import os from datasets import load_dataset from model import GPT, GPTConfig import tiktoken # Setup configurations # H200 has massive VRAM, we can push the batch size much higher BATCH_SIZE = 64 BLOCK_SIZE = 256 LEARNING_RATE = 3e-4 # We train on the standard Edu-Fineweb baseline first to establish neural weights and architecture. # We will use the webreaper data for fine-tuning / reasoning alignment later. DATASET_NAME = "HuggingFaceFW/fineweb-edu" CHECKPOINT_DIR = "./checkpoints" os.makedirs(CHECKPOINT_DIR, exist_ok=True) # Optimize PyTorch precision for Ampere/Hopper/Blackwell architecture (like H200) torch.set_float32_matmul_precision('high') # ----------------------------------------------------------------------- # ZERO-GPU ARCHITECTURE # ----------------------------------------------------------------------- # ZeroGPU allocates a GPU dynamically only when this function is called. # It has a strict time limit (usually 60-120s for free users). # To bypass this, we "Chunk" the training: # 1. Load the model & optimizer from disk. # 2. Move to GPU. # 3. Train for `steps` (~60 seconds of compute). # 4. Move model back to CPU and save to disk BEFORE the GPU gets deallocated. # ----------------------------------------------------------------------- @spaces.GPU(duration=120) def train_chunk(steps=50, checkpoint_path=None): device = 'cuda' if torch.cuda.is_available() else 'cpu' # 1. Initialize or Load Model # Using a small "NanoGPT" footprint to fit in memory config = GPTConfig(block_size=BLOCK_SIZE, vocab_size=50304, n_layer=4, n_head=4, n_embd=256) model = GPT(config) # Karpathy's optimized AdamW configuration # We use weight_decay=0.1, betas=(0.9, 0.95), eps=1e-8, and fused=True optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.1, betas=(0.9, 0.95), eps=1e-8, fused=True) start_step = 0 if checkpoint_path and os.path.exists(checkpoint_path): checkpoint = torch.load(checkpoint_path, map_location='cpu') # Remove the _orig_mod. prefix added by torch.compile if present 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) optimizer.load_state_dict(checkpoint['optimizer']) # Move optimizer state tensors to the correct device (GPU) for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.to(device) start_step = checkpoint.get('step', 0) # Move to the ZeroGPU model.to(device) # Compile the model to push the H200 to its absolute limits if hasattr(torch, 'compile'): model = torch.compile(model) model.train() # 2. Load Dataset (Streaming to save RAM) try: # We stream the Fineweb dataset directly from Hugging Face ds = load_dataset(DATASET_NAME, name="sample-10BT", split="train", streaming=True) ds_iter = iter(ds) except Exception as e: return f"Dataset load error: {e}", checkpoint_path enc = tiktoken.get_encoding("gpt2") # 3. Training Loop on the GPU logs = [] logs.append(f"--- Resuming training at Step {start_step} (Using H200 bfloat16 + torch.compile) ---") for step in range(start_step, start_step + steps): # Fetch a text chunk from the dataset try: row = next(ds_iter) # Fineweb uses the 'text' column. text = row.get("text", " ") if not text: text = " " except StopIteration: break # Karpathy Cosine LR Schedule with Warmup max_steps = 5000 # Assume a 5000 step total training run warmup_steps = 100 if step < warmup_steps: lr = LEARNING_RATE * (step + 1) / warmup_steps elif step > max_steps: lr = LEARNING_RATE * 0.1 else: decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps) import math coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) lr = LEARNING_RATE * 0.1 + coeff * (LEARNING_RATE - LEARNING_RATE * 0.1) for param_group in optimizer.param_groups: param_group['lr'] = lr tokens = enc.encode(text, allowed_special={"<|endoftext|>"}) if len(tokens) < BLOCK_SIZE + 1: continue # Sample a random sequence ix = torch.randint(len(tokens) - BLOCK_SIZE, (BATCH_SIZE,)) # pin_memory is implicitly fast on direct tensor creation, but pinning manually helps with standard dataloaders x = torch.stack([torch.tensor(tokens[i:i+BLOCK_SIZE], dtype=torch.long) for i in ix]).to(device, non_blocking=True) y = torch.stack([torch.tensor(tokens[i+1:i+1+BLOCK_SIZE], dtype=torch.long) for i in ix]).to(device, non_blocking=True) # Forward & Backward Pass with Mixed Precision (bfloat16) for Tensor Cores with torch.autocast(device_type=device, dtype=torch.bfloat16): logits, loss = model(x, y) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() if step % 10 == 0: logs.append(f"Step {step} | Loss: {loss.item():.4f}") # 4. Save Checkpoint BEFORE ZeroGPU shuts down model.to('cpu') raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model out_ckpt = os.path.join(CHECKPOINT_DIR, f"ckpt_step_{start_step + steps}.pt") torch.save({ 'model': raw_model.state_dict(), 'optimizer': optimizer.state_dict(), 'step': start_step + steps }, out_ckpt) # Clean up old checkpoints to save disk space if checkpoint_path and os.path.exists(checkpoint_path) and checkpoint_path != out_ckpt: os.remove(checkpoint_path) return "\n".join(logs), out_ckpt # ----------------------------------------------------------------------- # GRADIO INTERFACE # ----------------------------------------------------------------------- current_ckpt = None def run_training_ui(steps_per_call): global current_ckpt # We yield updates so the UI can stream the logs infinitely # while keeping the python thread alive, triggering the ZeroGPU # decorator over and over again. yield "Starting infinite training loop... this will run until you run out of ZeroGPU quota!", f"Initial Checkpoint: {current_ckpt}" total_logs = [] while True: try: # 1. Trigger the ZeroGPU allocation (runs for ~100s) chunk_log, new_ckpt = train_chunk(steps=int(steps_per_call), checkpoint_path=current_ckpt) # 2. Update global state current_ckpt = new_ckpt # 3. Append to logs and yield to UI total_logs.append(chunk_log) # Keep logs manageable if len(total_logs) > 50: total_logs = total_logs[-50:] yield "\n\n".join(total_logs), f"Current Active Checkpoint: {current_ckpt}" except Exception as e: err_msg = f"Training interrupted (likely ran out of ZeroGPU quota). Error: {str(e)}" total_logs.append(err_msg) yield "\n\n".join(total_logs), f"Final Checkpoint: {current_ckpt}" break with gr.Blocks(title="ZeroGPU WebReaper LM Trainer") as demo: gr.Markdown("# 🧠 Auto-ML Infinite Trainer on ZeroGPU") gr.Markdown(f"Trains a Karpathy-style NanoGPT directly on the `{DATASET_NAME}` dataset utilizing Hugging Face's free ZeroGPU allowance.") gr.Markdown("Click **'Start Infinite Training'** ONCE. The script will automatically loop, requesting a GPU, training for 60 seconds, saving state to disk, and repeating. It will only stop when you hit your daily 25-minute quota limit!") with gr.Row(): steps_slider = gr.Slider(minimum=10, maximum=150, value=50, step=10, label="Steps per GPU Request") train_btn = gr.Button("🚀 Start Infinite Training", variant="primary") stop_btn = gr.Button("⏚ī¸ Stop Training", variant="stop") with gr.Row(): log_output = gr.Textbox(label="Live Training Logs (Streams automatically)", lines=20) ckpt_output = gr.Textbox(label="Checkpoint Status") # Use generator to stream logs train_event = train_btn.click(fn=run_training_ui, inputs=[steps_slider], outputs=[log_output, ckpt_output]) stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[train_event]) if __name__ == "__main__": demo.launch()