ARC-Base-8B-Condensed/_stable_train.py

import sys
sys.path.insert(0, "/home/programmer/Desktop/test_arc/ARC-Base-8B-Condensed")

import torch
import json
import random
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel, get_peft_model, LoraConfig
import os

print("Loading model for CONSERVATIVE training...")
MODEL_PATH = "/mnt/nvme2/ubermesnchetien4/models/merged-final-v5"
CHECKPOINT = "/home/programmer/Desktop/test_arc/ARC-Base-8B-Condensed/dense_checkpoints_v2/step_100"

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    ),
    device_map="auto",
    torch_dtype=torch.bfloat16,
    local_files_only=True
)

if os.path.exists(CHECKPOINT):
    model = PeftModel.from_pretrained(model, CHECKPOINT, is_trainable=True)
    print(f"Loaded checkpoint: {CHECKPOINT}")
else:
    lora_config = LoraConfig(
        r=16, lora_alpha=32, 
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.05
    )
    model = get_peft_model(model, lora_config)
    print("Created new LoRA adapter")

# Load diverse training data
training_examples = [{"prompt": "hello", "response": "Hello. How can I help?"}, {"prompt": "hi", "response": "Hi. What do you need?"}, {"prompt": "hey", "response": "Hey. What's up?"}, {"prompt": "how are you?", "response": "Functional and ready. What can I do for you?"}, {"prompt": "good morning", "response": "Morning. What's the task?"}, {"prompt": "what's up?", "response": "Ready to work. What do you need?"}, {"prompt": "What is recursion?", "response": "A function calling itself with a smaller input until hitting a base case. Stack frames accumulate, then unwind. Classic examples: factorial, Fibonacci, tree traversal. Risk: stack overflow without proper termination."}, {"prompt": "Explain neural networks", "response": "Layers of weighted connections that learn patterns. Input \u2192 hidden layers \u2192 output. Training: forward pass computes prediction, loss measures error, backprop distributes blame, gradient descent updates weights. Universal function approximators given sufficient width."}, {"prompt": "How does gradient descent work?", "response": "Iteratively move parameters downhill on the loss surface. Compute gradient \u2207L, update \u03b8 \u2190 \u03b8 - \u03b1\u2207L. Learning rate \u03b1 controls step size: too large overshoots, too small stagnates. Variants: SGD, momentum, Adam. Local minima are usually fine in high dimensions."}, {"prompt": "What is backpropagation?", "response": "Chain rule applied layer by layer from output to input. Compute \u2202L/\u2202w for each weight by propagating error gradients backward. Efficient: reuses intermediate computations. Requires differentiable activations. The engine of deep learning."}, {"prompt": "Explain attention mechanism", "response": "Learned relevance weighting. Query asks, keys answer, values deliver. Attention(Q,K,V) = softmax(QK^T/\u221ad)V. Each position attends to all others. O(n\u00b2) cost. Self-attention: Q, K, V all from same sequence. Powers transformers."}, {"prompt": "What is a transformer?", "response": "Attention-based architecture. No recurrence, fully parallel. Stack of: multi-head attention \u2192 feedforward \u2192 layer norm. Positional encoding adds sequence order. Scales to billions of parameters. Foundation of modern LLMs."}, {"prompt": "Explain hash tables", "response": "Key-value storage with O(1) average access. Hash function maps key to bucket index. Collisions handled by chaining or open addressing. Load factor determines resize threshold. Trade space for time. Ubiquitous in practice."}, {"prompt": "What is Big O notation?", "response": "Asymptotic complexity measure. Describes growth rate as input scales. O(1) constant, O(log n) logarithmic, O(n) linear, O(n\u00b2) quadratic, O(2^n) exponential. Ignores constants and lower terms. Worst-case unless specified."}, {"prompt": "Explain dynamic programming", "response": "Solve problems by combining solutions to overlapping subproblems. Memoization: top-down with caching. Tabulation: bottom-up filling table. Key insight: optimal substructure + overlapping subproblems. Classic: Fibonacci, knapsack, edit distance."}, {"prompt": "What is a database index?", "response": "Data structure for fast lookups. B-tree most common: balanced, O(log n) search. Trade write speed and storage for read performance. Clustered index determines physical order. Choose columns with high selectivity."}, {"prompt": "What is overfitting?", "response": "Model memorizes training data, fails to generalize. Signs: low training loss, high validation loss. Causes: too complex model, insufficient data. Remedies: regularization, dropout, early stopping, more data, simpler architecture."}, {"prompt": "Explain regularization", "response": "Constrain model complexity to prevent overfitting. L1 (Lasso): sparse weights, feature selection. L2 (Ridge): small weights, smooth solutions. Dropout: randomly zero neurons during training. Weight decay: penalize large parameters."}, {"prompt": "What is cross-validation?", "response": "Estimate generalization by training on subsets. K-fold: split data into k parts, rotate test set. Re

print(f"Training on {len(training_examples)} diverse examples for 25 steps...")

# Conservative optimizer with LOW learning rate
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-6)  # Very low LR

model.train()
total_loss = 0
losses = []

for step in range(25):
    # Randomly sample an example (ensures diversity)
    ex = random.choice(training_examples)
    prompt = ex["prompt"]
    response = ex["response"]
    
    # Format for ChatML
    full_text = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{response}<|im_end|>"
    
    inputs = tokenizer(full_text, return_tensors="pt", truncation=True, max_length=512)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    outputs = model(**inputs, labels=inputs["input_ids"])
    loss = outputs.loss
    
    optimizer.zero_grad()
    loss.backward()
    
    # Gradient clipping for stability
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
    
    optimizer.step()
    
    total_loss += loss.item()
    losses.append(loss.item())
    
    if step % 5 == 0:
        recent_avg = sum(losses[-5:]) / len(losses[-5:]) if losses[-5:] else 0
        print(f"Step {step}: loss={loss.item():.4f}, recent_avg={recent_avg:.4f}")

# Save checkpoint
save_path = "/home/programmer/Desktop/test_arc/ARC-Base-8B-Condensed/dense_checkpoints_v2/step_475"
model.save_pretrained(save_path)

final_avg_loss = total_loss / 25
print(f"\nSaved checkpoint to {save_path}")
print(f"Final avg loss: {final_avg_loss:.4f}")
print("TRAINING_COMPLETE")