Delete train_single_gpu.py

train_single_gpu.py

@@ -1,581 +0,0 @@
-#!/usr/bin/env python3
-"""SpiderPortal v5 — Single-GPU Optimized Training.
-
-For RTX PRO 6000 (96GB) or similar large-VRAM GPU.
-No DDP, maximal batch size, torch.compile, pre-tokenized data.
-
-Usage:
-    python train_single_gpu.py
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import math
-import os
-import json
-import gc
-import random
-import time
-from pathlib import Path
-from dataclasses import dataclass
-from typing import Optional, Tuple, Dict, List
-from torch.nn import CrossEntropyLoss
-
-@dataclass
-class SpiderPortalConfig:
-    vocab_size: int = 50278
-    hidden_size: int = 384
-    num_hidden_layers: int = 8
-    num_attention_heads: int = 8
-    num_key_value_heads: int = 2
-    intermediate_size: int = 1024
-    hidden_act: str = "silu"
-    num_experts: int = 64
-    num_experts_per_tok: int = 1
-    num_shared_experts: int = 1
-    router_aux_loss_coef: float = 0.05
-    max_loop_iters: int = 4
-    act_threshold: float = 0.5
-    max_position_embeddings: int = 131072
-    rope_theta: float = 10000000.0
-    rope_scaling: dict = None
-    sliding_window: int = 4096
-    attention_dropout: float = 0.0
-    rms_norm_eps: float = 1e-6
-    initializer_range: float = 0.02
-    use_cache: bool = True
-    tie_word_embeddings: bool = True
-    prelude_layers: int = 2
-    coda_layers: int = 2
-    lora_rank: int = 32
-    loop_embed_dim: int = 48
-    vision_hidden_size: int = 384
-    audio_hidden_size: int = 512
-    vision_num_frames: int = 60
-    vision_tokens_per_frame: int = 256
-    vision_temporal_tokens: int = 64
-    vision_temporal_layers: int = 2
-    model_type: str = "spiderportal"
-    torch_dtype: str = "bfloat16"
-
-def loop_index_embedding(h, loop_t, loop_dim, theta=10000.0):
-    freqs = 1.0 / (theta ** (torch.arange(0, loop_dim, 2, device=h.device, dtype=h.dtype) / loop_dim))
-    angles = loop_t * freqs
-    emb = torch.cat([angles.sin(), angles.cos()], dim=-1)[:loop_dim]
-    emb_full = torch.zeros(h.shape[-1], device=h.device, dtype=h.dtype)
-    emb_full[:loop_dim] = emb
-    return h + emb_full.unsqueeze(0).unsqueeze(0)
-
-class SpiderPortalRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight.to(input_dtype) * hidden_states.to(input_dtype)
-
-def compute_yarn_inv_freq(head_dim, rope_theta, factor, orig_max, beta_fast=32.0, beta_slow=1.0):
-    dim = head_dim
-    orig_inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2).float() / dim))
-    pos_freqs = torch.arange(0, dim, 2).float() / dim
-    beta = (pos_freqs * math.log(rope_theta) / math.log(orig_max))
-    scale = torch.where(beta < beta_slow, torch.ones_like(beta), torch.where(beta > beta_fast, torch.ones_like(beta) / factor, 1.0 - (beta - beta_slow) / (beta_fast - beta_slow) * (1.0 - 1.0 / factor)))
-    return orig_inv_freq * scale
-
-class SpiderPortalGQA(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.head_dim = config.hidden_size // config.num_attention_heads
-        self.num_key_value_groups = self.num_heads // self.num_kv_heads
-        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
-        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
-        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
-        self.attention_dropout = config.attention_dropout
-        rope_scaling = getattr(config, 'rope_scaling', None)
-        if rope_scaling and rope_scaling.get("type") == "yarn":
-            factor = rope_scaling.get("factor", 1.0)
-            orig_max_pos = rope_scaling.get("original_max_position_embeddings", config.max_position_embeddings)
-            inv_freq = compute_yarn_inv_freq(self.head_dim, config.rope_theta, factor, orig_max_pos)
-        else:
-            inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-    def _rotate_half(self, x):
-        x1 = x[..., :x.shape[-1] // 2]
-        x2 = x[..., x.shape[-1] // 2:]
-        return torch.cat((-x2, x1), dim=-1)
-    def _apply_rotary(self, x, cos, sin):
-        return (x * cos) + (self._rotate_half(x) * sin)
-    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
-        bsz, q_len, _ = hidden_states.size()
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        if position_ids is None:
-            position_ids = torch.arange(q_len, device=hidden_states.device).unsqueeze(0).expand(bsz, -1)
-        max_pos = position_ids.max().item() + 1
-        seq_len = max(max_pos, q_len)
-        t = torch.arange(seq_len, device=hidden_states.device, dtype=self.inv_freq.dtype)
-        freqs = torch.outer(t, self.inv_freq)
-        emb = torch.cat((freqs, freqs), dim=-1)
-        cos, sin = emb.cos(), emb.sin()
-        cos = cos[position_ids].unsqueeze(1)
-        sin = sin[position_ids].unsqueeze(1)
-        query_states = self._apply_rotary(query_states, cos, sin)
-        key_states = self._apply_rotary(key_states, cos, sin)
-        if past_key_value is not None:
-            key_states = torch.cat([past_key_value[0], key_states], dim=2)
-            value_states = torch.cat([past_key_value[1], value_states], dim=2)
-        past_kv = (key_states, value_states) if use_cache else None
-        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
-        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-        if attention_mask is not None:
-            attn_weights = attn_weights + attention_mask
-        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
-        attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
-        attn_output = torch.matmul(attn_weights, value_states)
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-        return self.o_proj(attn_output), past_kv
-
-class SpiderPortalExpert(nn.Module):
-    def __init__(self, config, intermediate_size=None):
-        super().__init__()
-        inter_size = intermediate_size or config.intermediate_size
-        self.gate_proj = nn.Linear(config.hidden_size, inter_size, bias=False)
-        self.up_proj = nn.Linear(config.hidden_size, inter_size, bias=False)
-        self.down_proj = nn.Linear(inter_size, config.hidden_size, bias=False)
-        self.act_fn = nn.SiLU()
-    def forward(self, hidden_states):
-        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
-
-class SpiderPortalRouter(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.num_experts = config.num_experts
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.weight = nn.Parameter(torch.randn(config.hidden_size, config.num_experts) * config.initializer_range)
-        self.register_buffer("router_bias", torch.zeros(config.num_experts))
-    def forward(self, hidden_states):
-        router_logits = hidden_states.view(-1, hidden_states.size(-1)) @ self.weight
-        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)
-        biased_logits = router_logits + self.router_bias
-        biased_weights = F.softmax(biased_logits, dim=-1, dtype=torch.float32)
-        top_weights, top_indices = torch.topk(biased_weights, self.num_experts_per_tok, dim=-1)
-        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)
-        top_weights = top_weights.to(hidden_states.dtype)
-        mean_probs = routing_weights.mean(dim=0)
-        aux_loss = self.num_experts * (mean_probs * mean_probs).sum()
-        return top_weights, top_indices, aux_loss
-
-class SpiderPortalMoE(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.num_experts = config.num_experts
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.experts = nn.ModuleList([SpiderPortalExpert(config) for _ in range(config.num_experts)])
-        self.shared_expert = SpiderPortalExpert(config)
-        self.router = SpiderPortalRouter(config)
-    def forward(self, hidden_states):
-        batch_size, seq_len, hidden_dim = hidden_states.shape
-        top_weights, top_indices, aux_loss = self.router(hidden_states)
-        flat_hidden = hidden_states.view(-1, hidden_dim)
-        final_output = torch.zeros_like(flat_hidden)
-        for expert_idx in range(self.num_experts_per_tok):
-            expert_ids = top_indices[:, expert_idx]
-            expert_weights = top_weights[:, expert_idx:expert_idx+1]
-            for e in range(self.num_experts):
-                mask = expert_ids == e
-                if mask.any():
-                    expert_output = self.experts[e](flat_hidden[mask])
-                    final_output[mask] += expert_output * expert_weights[mask]
-        shared_output = self.shared_expert(flat_hidden)
-        final_output = final_output + shared_output
-        return final_output.view(batch_size, seq_len, hidden_dim), aux_loss
-
-class SpiderPortalDenseLayer(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.self_attn = SpiderPortalGQA(config)
-        dense_intermediate = config.hidden_size * 4 // 3
-        self.ffn = SpiderPortalExpert(config, intermediate_size=dense_intermediate)
-        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
-        attn_input = self.input_layernorm(hidden_states)
-        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)
-        hidden_states = hidden_states + attn_output
-        ffn_input = self.post_attention_layernorm(hidden_states)
-        ffn_output = self.ffn(ffn_input)
-        hidden_states = hidden_states + ffn_output
-        return hidden_states, past_kv
-
-class SpiderPortalMoELayer(nn.Module):
-    def __init__(self, config, layer_idx):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.self_attn = SpiderPortalGQA(config)
-        self.moe = SpiderPortalMoE(config)
-        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
-        attn_input = self.input_layernorm(hidden_states)
-        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)
-        hidden_states = hidden_states + attn_output
-        moe_input = self.post_attention_layernorm(hidden_states)
-        moe_output, aux_loss = self.moe(moe_input)
-        hidden_states = hidden_states + moe_output
-        return hidden_states, aux_loss, past_kv
-
-class LTIInjection(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.log_A = nn.Parameter(torch.full((config.hidden_size,), -2.0))
-        self.delta_t = nn.Parameter(torch.tensor(1.0))
-        self.B = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
-        with torch.no_grad():
-            self.B.weight.data.normal_(mean=0.0, std=0.01)
-    def get_A(self):
-        return -torch.exp(self.log_A)
-    def forward(self, h_t, e):
-        A = self.get_A()
-        return A * h_t + self.B(e)
-
-class ACTHalting(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.halt_predictor = nn.Linear(config.hidden_size, 1)
-        self.threshold = config.act_threshold
-    def forward(self, hidden_states):
-        return torch.sigmoid(self.halt_predictor(hidden_states))
-
-class LoRAAdapter(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        rank = config.lora_rank
-        self.down = nn.Linear(config.hidden_size, rank, bias=False)
-        self.B = nn.Parameter(torch.randn(rank, config.hidden_size) * 0.02)
-        self.scale = nn.Embedding(config.max_loop_iters, rank)
-        with torch.no_grad():
-            self.scale.weight.data.zero_()
-            self.down.weight.data.normal_(mean=0.0, std=0.001)
-    def forward(self, x, loop_t):
-        max_t = self.scale.num_embeddings - 1
-        t_idx = min(loop_t, max_t)
-        s = self.scale(torch.tensor(t_idx, device=x.device))
-        down = self.down(x) * s
-        return down @ self.B
-
-class SpiderPortalMoEModel(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.prelude_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.prelude_layers)])
-        self.recurrent_layers = nn.ModuleList([SpiderPortalMoELayer(config, i) for i in range(config.num_hidden_layers)])
-        self.coda_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.coda_layers)])
-        self.norm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.injection = LTIInjection(config)
-        self.act_halting = ACTHalting(config)
-        self.lora_adapter = LoRAAdapter(config)
-        self.loop_embed_dim = config.loop_embed_dim
-    def forward(self, hidden_states, input_embedding=None, attention_mask=None, position_ids=None, past_key_values=None, use_cache=False, n_loops=None):
-        n_loops = n_loops or self.config.max_loop_iters
-        input_embedding = input_embedding if input_embedding is not None else hidden_states
-        total_aux_loss = 0.0
-        for layer in self.prelude_layers:
-            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
-        e = hidden_states.clone()
-        B, T_seq, D = hidden_states.shape
-        halted = torch.zeros(B, T_seq, device=hidden_states.device, dtype=torch.bool)
-        cumulative_p = torch.zeros(B, T_seq, device=hidden_states.device, dtype=hidden_states.dtype)
-        h_out = torch.zeros_like(hidden_states)
-        past_key_values = past_key_values if past_key_values is not None else [None] * len(self.recurrent_layers)
-        for t in range(n_loops):
-            h_loop = loop_index_embedding(hidden_states, t, self.loop_embed_dim)
-            if t > 0:
-                injection = self.injection(hidden_states, input_embedding)
-                hidden_states = hidden_states + injection
-            new_past_key_values = []
-            for i, layer in enumerate(self.recurrent_layers):
-                hidden_states, aux_loss, past_kv = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values[i] if t == 0 else None, use_cache=use_cache)
-                total_aux_loss = total_aux_loss + aux_loss
-                new_past_key_values.append(past_kv)
-            lora_delta = self.lora_adapter(hidden_states, t)
-            hidden_states = hidden_states + lora_delta
-            halt_prob = self.act_halting(hidden_states).squeeze(-1)
-            still_running = ~halted
-            remainder = (1.0 - cumulative_p).clamp(min=0)
-            weight = torch.where(cumulative_p + halt_prob >= self.config.act_threshold, remainder, halt_prob)
-            weight = weight * still_running.to(hidden_states.dtype)
-            h_out = h_out + weight.unsqueeze(-1) * hidden_states
-            cumulative_p = cumulative_p + halt_prob * still_running.to(hidden_states.dtype)
-            halted = halted | (cumulative_p >= self.config.act_threshold)
-            if halted.all() and not self.training:
-                break
-        never_halted = (~halted).to(hidden_states.dtype).unsqueeze(-1)
-        hidden_states = h_out + never_halted * hidden_states
-        for layer in self.coda_layers:
-            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
-        hidden_states = self.norm(hidden_states)
-        return hidden_states, total_aux_loss, new_past_key_values
-
-class SpiderPortalForConditionalGeneration(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.model = SpiderPortalMoEModel(config)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        if config.tie_word_embeddings:
-            self.lm_head.weight = self.embed_tokens.weight
-        self.apply(self._init_weights)
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-            if hasattr(self, 'model') and module is self.model.injection.B:
-                return
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-    def forward(self, input_ids, attention_mask=None, position_ids=None, labels=None, n_loops=None, use_cache=False):
-        hidden_states = self.embed_tokens(input_ids)
-        model_dtype = next(self.model.parameters()).dtype
-        hidden_states = hidden_states.to(model_dtype)
-        input_embedding = hidden_states.clone()
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
-        causal_mask = torch.full((attention_mask.size(0), 1, attention_mask.size(1), attention_mask.size(1)), 0.0, dtype=hidden_states.dtype, device=hidden_states.device)
-        causal_mask = causal_mask.masked_fill(~attention_mask.unsqueeze(1).unsqueeze(2), torch.finfo(hidden_states.dtype).min)
-        causal_mask = causal_mask.triu(1)
-        hidden_states, aux_loss, past_kv = self.model(hidden_states, input_embedding=input_embedding, attention_mask=causal_mask, position_ids=position_ids, use_cache=use_cache, n_loops=n_loops)
-        logits = self.lm_head(hidden_states)
-        loss = None
-        if labels is not None:
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-            loss = loss + self.config.router_aux_loss_coef * aux_loss
-        return {"loss": loss, "logits": logits, "aux_loss": aux_loss, "past_key_values": past_kv}
-    def get_num_params(self):
-        total = sum(p.numel() for p in self.parameters())
-        return {"total": total, "trainable": total}
-
-def train_single_gpu():
-    device = torch.device("cuda")
-    gpu_name = torch.cuda.get_device_name(0)
-    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
-    print(f"GPU: {gpu_name} ({gpu_mem:.1f}GB)")
-
-    config = SpiderPortalConfig(
-        hidden_size=384, num_hidden_layers=8, num_attention_heads=8,
-        num_key_value_heads=2, intermediate_size=1024,
-        num_experts=64, num_experts_per_tok=1, num_shared_experts=1,
-        router_aux_loss_coef=0.05, max_loop_iters=2,
-        prelude_layers=2, coda_layers=2, lora_rank=32,
-        rope_theta=10000000.0,
-        rope_scaling={"type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768},
-        max_position_embeddings=131072, sliding_window=4096,
-        tie_word_embeddings=True,
-    )
-
-    print("Building model...")
-    model = SpiderPortalForConditionalGeneration(config)
-    model = model.to(torch.bfloat16).to(device)
-
-    params = model.get_num_params()
-    print(f"Model: {params['total']/1e6:.1f}M params")
-    print(f"Experts: {config.num_experts} routed + {config.num_shared_experts} shared")
-
-    BASE_LR = 1e-3
-    WARMUP_STEPS = 500
-    optimizer = torch.optim.AdamW(model.parameters(), lr=BASE_LR, weight_decay=0.01, betas=(0.9, 0.95))
-
-    import pandas as pd
-    data_dir = Path(__file__).parent / "data"
-    all_records = []
-    pkl_file = data_dir / "spiderportal_combined.pkl"
-    if pkl_file.exists():
-        print(f"Loading dataset from {pkl_file}...")
-        df = pd.read_pickle(pkl_file)
-        all_records = df.to_dict("records")
-    else:
-        print(f"No dataset found at {pkl_file}, creating synthetic data...")
-        all_records = [{"instruction": f"Question {i}: What is {i} + {i}?", "input": "", "output": f"The answer is {i+i}."} for i in range(10000)]
-
-    print(f"Loaded {len(all_records):,} samples")
-
-    from transformers import AutoTokenizer
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
-    tokenizer.pad_token = tokenizer.eos_token
-
-    BATCH_SIZE = 128
-    MAX_LEN = 256
-    EPOCHS = 3
-    N_LOOPS = 2
-
-    print(f"Batch size: {BATCH_SIZE} (no grad accum)")
-    print(f"Effective batch: {BATCH_SIZE}")
-    print(f"LR: {BASE_LR} with {WARMUP_STEPS}-step warmup (high LR for recurrent MoE)")
-    print(f"Max seq len: {MAX_LEN}, N_LOOPS: {N_LOOPS}")
-
-    def build_prompt(sample):
-        instruction = str(sample.get("instruction", "")).strip()
-        inp = str(sample.get("input", "")).strip()
-        output = str(sample.get("output", "")).strip()
-        if inp:
-            return f"Question: Instruction: {instruction}\nInput: {inp}\nAnswer: {output}\n"
-        return f"Question: Instruction: {instruction}\nAnswer: {output}\n"
-
-    print("Pre-tokenizing dataset...")
-    prefix_ids = tokenizer("Question:", add_special_tokens=False)["input_ids"]
-    mask_len = len(prefix_ids)
-
-    pre_tokenized = []
-    for i, sample in enumerate(all_records):
-        instruction = str(sample.get("instruction", "")).strip()
-        inp = str(sample.get("input", "")).strip()
-        output = str(sample.get("output", "")).strip()
-        if inp:
-            text = f"Question: Instruction: {instruction}\nInput: {inp}\nAnswer: {output}\n" + tokenizer.eos_token
-        else:
-            text = f"Question: Instruction: {instruction}\nAnswer: {output}\n" + tokenizer.eos_token
-        enc = tokenizer(text, truncation=True, max_length=MAX_LEN, padding="max_length")
-        input_ids = enc["input_ids"]
-        labels = input_ids[:]
-        for j in range(min(mask_len, len(labels))):
-            labels[j] = -100
-        pre_tokenized.append((input_ids, labels))
-        if (i + 1) % 50000 == 0:
-            print(f"  Tokenized {i+1:,}/{len(all_records):,}")
-
-    print(f"Pre-tokenization complete: {len(pre_tokenized):,} samples")
-    del all_records
-    gc.collect()
-
-    global_step = 0
-    best_loss = float('inf')
-    start_time = time.time()
-    checkpoint_dir = Path("checkpoints")
-    checkpoint_dir.mkdir(exist_ok=True)
-    step_ckpt_files = []
-
-    for epoch in range(1, EPOCHS + 1):
-        if epoch > 1:
-            for f in step_ckpt_files:
-                if f.exists():
-                    f.unlink()
-                    print(f"  Deleted old step checkpoint: {f.name}")
-            step_ckpt_files.clear()
-            gc.collect()
-
-        indices = list(range(len(pre_tokenized)))
-        random.shuffle(indices)
-        total_loss = 0
-        num_batches = 0
-        optimizer.zero_grad()
-
-        for batch_start in range(0, len(indices), BATCH_SIZE):
-            batch_indices = indices[batch_start:batch_start + BATCH_SIZE]
-            if len(batch_indices) < BATCH_SIZE:
-                continue
-
-            if global_step < WARMUP_STEPS:
-                lr = BASE_LR * (global_step + 1) / WARMUP_STEPS
-                for param_group in optimizer.param_groups:
-                    param_group['lr'] = lr
-
-            batch_input_ids = []
-            batch_labels = []
-            for idx in batch_indices:
-                input_ids, labels = pre_tokenized[idx]
-                batch_input_ids.append(input_ids)
-                batch_labels.append(labels)
-
-            input_ids = torch.tensor(batch_input_ids, dtype=torch.long, device=device)
-            labels = torch.tensor(batch_labels, dtype=torch.long, device=device)
-
-            if global_step == 0:
-                print("  [First forward pass - compiling...]")
-
-            outputs = model(input_ids=input_ids, labels=labels, n_loops=N_LOOPS)
-            loss = outputs["loss"]
-            loss.backward()
-
-            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
-            optimizer.step()
-            optimizer.zero_grad()
-            global_step += 1
-
-            total_loss += loss.item()
-            num_batches += 1
-
-            if (batch_start // BATCH_SIZE) == 0 or global_step < 20 or global_step % 100 == 0:
-                avg_loss = total_loss / max(num_batches, 1)
-                elapsed = time.time() - start_time
-                steps_per_hour = (global_step + 1) / elapsed * 3600 if elapsed > 0 else 0
-                current_lr = optimizer.param_groups[0]['lr']
-                samples_per_sec = (global_step * BATCH_SIZE) / elapsed if elapsed > 0 else 0
-                print(f"Epoch {epoch}/{EPOCHS} | Step {global_step} | loss={avg_loss:.4f} | LR={current_lr:.2e} | {steps_per_hour:.0f} steps/hr | {samples_per_sec:.0f} samples/sec")
-
-            if global_step > 0 and global_step % 500 == 0:
-                ckpt_path = checkpoint_dir / f"spiderportal-v5-ep{epoch}-step{global_step}.pt"
-                state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
-                torch.save(state_dict, ckpt_path)
-                step_ckpt_files.append(ckpt_path)
-                size_mb = ckpt_path.stat().st_size / (1024 * 1024)
-                print(f"Saved weights-only checkpoint: {ckpt_path.name} ({size_mb:.0f}MB)")
-
-        avg_loss = total_loss / max(num_batches, 1)
-        epoch_time = (time.time() - start_time) / 60
-        print(f"Epoch {epoch}/{EPOCHS} complete | avg_loss={avg_loss:.4f} | Time: {epoch_time:.1f}min")
-
-        ckpt_path = checkpoint_dir / f"spiderportal-v5-ep{epoch}.pt"
-        torch.save({
-            "step": global_step,
-            "epoch": epoch,
-            "model_state_dict": {k: v.cpu() for k, v in model.state_dict().items()},
-            "optimizer_state_dict": optimizer.state_dict(),
-            "config": config.__dict__,
-        }, ckpt_path)
-        size_mb = ckpt_path.stat().st_size / (1024 * 1024)
-        print(f"Saved epoch checkpoint: {ckpt_path.name} ({size_mb:.0f}MB)")
-
-        if avg_loss < best_loss:
-            best_loss = avg_loss
-            best_path = checkpoint_dir / "spiderportal-v5-best.pt"
-            torch.save({
-                "step": global_step,
-                "epoch": epoch,
-                "model_state_dict": {k: v.cpu() for k, v in model.state_dict().items()},
-                "optimizer_state_dict": optimizer.state_dict(),
-                "config": config.__dict__,
-            }, best_path)
-            size_mb = best_path.stat().st_size / (1024 * 1024)
-            print(f"Saved best checkpoint: {best_path.name} ({size_mb:.0f}MB)")
-
-    total_time = (time.time() - start_time) / 3600
-    print(f"\nTraining complete!")
-    print(f"Best loss: {best_loss:.4f}")
-    print(f"Total time: {total_time:.2f} hours")
-    print(f"Total steps: {global_step}")
-    print(f"Checkpoints saved to: {checkpoint_dir}")
-
-if __name__ == "__main__":
-    train_single_gpu()