Multi-model benchmark pipeline: VRAM cleanup + EMA graph + detailed output

app.py

@@ -98,10 +98,13 @@ with gr.Blocks(title="Cloud Arena RL") as demo:
         eval_btn.click(run_math_evaluation, outputs=eval_output)
 
     with gr.Tab("🧠 LLM RL"):
-        gr.Markdown("### LLM Model — Llama 3.1 8B + REINFORCE + LoRA")
-        gr.Markdown("> ⚠️ Requires `HF_TOKEN` secret set in Space settings + accepted model license")
-        llm_model = gr.Textbox(value="meta-llama/Llama-3.1-8B", label="Model Name")
-        llm_iters = gr.Number(value=200, label="Training Iterations")
+        gr.Markdown("### Multi-Model RL Benchmark — REINFORCE + LoRA")
+        gr.Markdown("> Comma-separate model names to benchmark multiple models sequentially")
+        llm_model = gr.Textbox(
+            value="unsloth/Qwen2.5-Math-7B-Instruct-bnb-4bit, unsloth/gemma-2b-it-bnb-4bit, unsloth/llama-3-8b-Instruct-bnb-4bit",
+            label="Model(s) — comma-separated for multi-model benchmark"
+        )
+        llm_iters = gr.Number(value=200, label="Training Iterations per Model")
         llm_steps = gr.Number(value=15, label="Steps per Episode")
         llm_btn = gr.Button("🚀 Start LLM Training", variant="primary")
         llm_output = gr.Textbox(label="Training Log", lines=15)

cloud_arena/llm_training.py

@@ -1,59 +1,59 @@
 # ============================================================
-# LLM RL Training — Llama 3.1 8B + REINFORCE + LoRA
-# This is the LLM model, SEPARATE from the mathematical model.
-# Uses AWSCostEnv (llm_environment.py), NOT CloudArenaEnv.
+# Multi-Model RL Benchmarking Pipeline
+# Sequential training of multiple LLMs with VRAM cleanup
+# REINFORCE + LoRA on Cloud FinOps Environment
 # ============================================================
 
-import os
-import re
-import json
-import time
-import warnings
+import os, re, json, time, gc
 import numpy as np
 import torch
 import torch.nn.functional as F
 import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
-
+import warnings
 warnings.filterwarnings("ignore", category=UserWarning)
 warnings.filterwarnings("ignore", category=FutureWarning)
 
 from cloud_arena.llm_environment import SB3Adapter, Action, AWSCostEnv
 
-# ─── Constants ────────────────────────────────────────────────────────────────
+# ── Configuration ─────────────────────────────────────────────────────────────
+
+MODELS_TO_TEST = [
+    "unsloth/Qwen2.5-Math-7B-Instruct-bnb-4bit",
+    "unsloth/gemma-2b-it-bnb-4bit",
+    "unsloth/llama-3-8b-Instruct-bnb-4bit",
+]
 
 ACTION_NAMES = {0: "NOOP", 1: "CHECK_DEPS", 2: "RESIZE", 3: "STOP", 4: "DELETE"}
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+GRAD_ACCUM_STEPS = 4
+MAX_SEQ_LEN = 512
+MAX_GEN_TOKENS = 80
+EMA_ALPHA = 0.3  # EMA smoothing factor for reward graph
 
-# ── GPU Optimization Constants ────────────────────────────────────────────────
-GRAD_ACCUM_STEPS = 4       # accumulate gradients over N episodes before stepping
-MAX_SEQ_LEN = 512          # shorter context = O(N²) attention is 4× faster than 1024
-MAX_GEN_TOKENS = 80        # enough room for reasoning + ACTION line, not enough to ramble
 
+# ── Prompt & Parser ───────────────────────────────────────────────────────────
 
 def format_prompt(state_dict):
     resources_text = ""
     for r in state_dict["resources"]:
         status = "ACTIVE" if r["active"] else "STOPPED"
-        tag = "PRODUCTION" if r["is_prod"] else "Temporary"
+        tag = "PROD" if r["is_prod"] else "TEMP"
         resources_text += (
             f"  - {r['name']} [{status}] ({tag}): "
-            f"Cost=${r['cost_per_hr']:.2f}/hr, CPU={r['cpu_pct']}%, "
+            f"${r['cost_per_hr']:.2f}/hr, CPU={r['cpu_pct']}%, "
             f"Deps={len(r['dependencies'])}\n"
         )
     savings_pct = state_dict.get("savings_pct", 0.0)
     return (
-        f"You are a Cloud FinOps AI. Reduce cloud cost by >=20% without breaking production.\n\n"
+        f"You are a Cloud FinOps AI. Reduce cost by >=20% without breaking production.\n\n"
         f"Actions: 0=NOOP, 1=CHECK_DEPS, 2=RESIZE, 3=STOP, 4=DELETE\n\n"
         f"Resources:\n{resources_text}\n"
         f"Baseline: ${state_dict['baseline_cost']:.2f}/hr | "
-        f"Current: ${state_dict['current_cost']:.2f}/hr | "
-        f"Savings: {savings_pct:.1f}%\n\n"
-        f"Rules:\n"
-        f"- Never delete/stop prod resources or those with >=5 deps\n"
-        f"- Temp resources with 0-1 deps are safe to delete\n"
-        f"- RESIZE is always safe\n\n"
+        f"Current: ${state_dict['current_cost']:.2f}/hr | Savings: {savings_pct:.1f}%\n\n"
+        f"Rules:\n- Never delete/stop prod resources or those with >=5 deps\n"
+        f"- Temp resources with 0-1 deps are safe to delete\n- RESIZE is always safe\n\n"
         f"CRITICAL: Output ONLY a brief reason then ACTION: <number 0-4>. Nothing else.\n\n"
         f"REASONING:"
     )
@@ -62,18 +62,15 @@ def format_prompt(state_dict):
 def extract_action_and_reasoning(response_text):
     """Regex safety net: extracts action even from truncated/malformed output."""
     reasoning = response_text.strip()
-    action = 2  # Default to RESIZE (safest action)
+    action = 2  # Default RESIZE
 
-    # Try explicit ACTION: N format first
-    action_match = re.search(r'ACTION:\s*([0-4])', response_text, re.IGNORECASE)
-    if action_match:
-        return int(action_match.group(1)), reasoning
+    match = re.search(r'ACTION:\s*([0-4])', response_text, re.IGNORECASE)
+    if match:
+        return int(match.group(1)), reasoning
 
-    # Try JSON format: {"action": N} or {"action": "DELETE"}
     json_match = re.search(r'\{.*?\}', response_text, re.DOTALL)
     if json_match:
         try:
-            import json
             parsed = json.loads(json_match.group(0))
             a = parsed.get("action", 2)
             if isinstance(a, int) and 0 <= a <= 4:
@@ -81,53 +78,46 @@ def extract_action_and_reasoning(response_text):
         except (json.JSONDecodeError, ValueError):
             pass
 
-    # Last resort: any digit 0-4 near the end
-    digit_matches = re.findall(r'\b([0-4])\b', response_text[-30:])
-    if digit_matches:
-        action = int(digit_matches[-1])
-
+    digits = re.findall(r'\b([0-4])\b', response_text[-30:])
+    if digits:
+        action = int(digits[-1])
     return action, reasoning
 
 
+# ── REINFORCE Loss ────────────────────────────────────────────────────────────
+
 def compute_pg_loss(model, tokenizer, prompt, response_text, reward):
-    """Compute REINFORCE loss WITHOUT stepping optimizer (for gradient accumulation)."""
     full_text = prompt + response_text
-    encodings = tokenizer(full_text, return_tensors="pt", truncation=True, max_length=MAX_SEQ_LEN).to(DEVICE)
-    prompt_encodings = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_SEQ_LEN)
-    prompt_len = prompt_encodings["input_ids"].shape[1]
+    enc = tokenizer(full_text, return_tensors="pt", truncation=True, max_length=MAX_SEQ_LEN).to(DEVICE)
+    prompt_len = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_SEQ_LEN)["input_ids"].shape[1]
 
-    outputs = model(**encodings, labels=encodings["input_ids"])
+    outputs = model(**enc, labels=enc["input_ids"])
     logits = outputs.logits[:, prompt_len-1:-1, :]
-    targets = encodings["input_ids"][:, prompt_len:]
+    targets = enc["input_ids"][:, prompt_len:]
 
     if targets.shape[1] == 0 or logits.shape[1] == 0:
         return 0.0
 
-    min_len = min(logits.shape[1], targets.shape[1])
-    logits = logits[:, :min_len, :]
-    targets = targets[:, :min_len]
+    ml = min(logits.shape[1], targets.shape[1])
+    log_probs = F.log_softmax(logits[:, :ml, :], dim=-1)
+    token_lp = log_probs.gather(2, targets[:, :ml].unsqueeze(-1)).squeeze(-1)
 
-    log_probs = F.log_softmax(logits, dim=-1)
-    token_log_probs = log_probs.gather(2, targets.unsqueeze(-1)).squeeze(-1)
-    avg_log_prob = token_log_probs.mean()
-
-    scaled_reward = reward / 10.0
-    loss = -scaled_reward * avg_log_prob
-    loss.backward()  # accumulate gradient, don't step yet
+    loss = -(reward / 10.0) * token_lp.mean()
+    loss.backward()
     return loss.item()
 
 
+# ── Episode Runner ────────────────────────────────────────────────────────────
+
 def run_episode(model, tokenizer, env, is_training=False, optimizer=None,
-                steps_per_episode=15):
+                steps_per_episode=15, iteration_num=0, total_iters=0):
     obs, info = env.reset()
     state_dict = env.core._get_internal_state()
     done = False
     episode_reward = 0.0
     step_count = 0
     reasoning_log = []
-    losses = []
 
-    # Accumulate gradients across all steps in the episode
     if is_training and optimizer is not None:
         optimizer.zero_grad()
 
@@ -135,175 +125,298 @@ def run_episode(model, tokenizer, env, is_training=False, optimizer=None,
         prompt = format_prompt(state_dict)
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_SEQ_LEN)
         input_ids = inputs["input_ids"].to(DEVICE)
-        attention_mask = inputs["attention_mask"].to(DEVICE)
+        attn_mask = inputs["attention_mask"].to(DEVICE)
 
         with torch.no_grad():
-            gen_outputs = model.generate(
-                input_ids, attention_mask=attention_mask,
+            gen = model.generate(
+                input_ids, attention_mask=attn_mask,
                 max_new_tokens=MAX_GEN_TOKENS,
                 do_sample=True, temperature=0.7, top_p=0.95,
                 pad_token_id=tokenizer.pad_token_id,
             )
 
-        response_ids = gen_outputs[0][input_ids.shape[1]:]
-        response_text = tokenizer.decode(response_ids, skip_special_tokens=True)
+        response_text = tokenizer.decode(gen[0][input_ids.shape[1]:], skip_special_tokens=True)
         action, reasoning = extract_action_and_reasoning(response_text)
 
         next_obs, reward, terminated, truncated, next_info = env.step(action)
         done = terminated or truncated
         episode_reward += reward
 
+        # ── Detailed per-step terminal output ──
+        if is_training and total_iters > 0:
+            pct = (iteration_num / total_iters) * 100
+            print(f"  [{pct:5.1f}%] Ep {iteration_num} Step {step_count+1}: "
+                  f"{ACTION_NAMES.get(action,'?')} → r={reward:+.3f} | "
+                  f"💬 {reasoning[:80]}")
+
         reasoning_log.append({
-            "step": step_count + 1,
-            "reasoning": reasoning[:300],
-            "action": action,
-            "action_name": ACTION_NAMES.get(action, "UNKNOWN"),
+            "step": step_count + 1, "action": action,
+            "action_name": ACTION_NAMES.get(action, "?"),
             "reward": round(reward, 4),
+            "reasoning": reasoning[:200],
             "message": next_info.get("msg", ""),
         })
 
         if is_training and optimizer is not None:
-            loss = compute_pg_loss(model, tokenizer, prompt, response_text, reward)
-            losses.append(loss)
+            compute_pg_loss(model, tokenizer, prompt, response_text, reward)
 
         obs = next_obs
         state_dict = env.core._get_internal_state()
         step_count += 1
 
-    return episode_reward, reasoning_log, losses
+    return episode_reward, reasoning_log
 
 
-def train_llm(model_name="meta-llama/Llama-3.1-8B",
-              num_iterations=200, steps_per_episode=15, learning_rate=2e-6,
-              progress_callback=None):
-    """
-    Full LLM RL training pipeline. Returns (all_rewards, full_log, graph_path).
-    """
-    hf_token = os.environ.get("HF_TOKEN")
+# ── VRAM Cleanup ──────────────────────────────────────────────────────────────
+
+def nuke_vram(model=None, optimizer=None, tokenizer=None):
+    """Aggressively free VRAM between model runs."""
+    if model is not None:
+        del model
+    if optimizer is not None:
+        del optimizer
+    if tokenizer is not None:
+        del tokenizer
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        vram = torch.cuda.memory_allocated() / 1e9
+        print(f"  🧹 VRAM after cleanup: {vram:.2f} GB")
+
+
+# ── Single Model Training ────────────────────────────────────────────────────
 
+def train_single_model(model_name, num_iterations=200, steps_per_episode=15,
+                       learning_rate=2e-6):
+    """Train one model, return rewards list."""
+    hf_token = os.environ.get("HF_TOKEN")
     from transformers import AutoModelForCausalLM, AutoTokenizer
     from peft import get_peft_model, LoraConfig, TaskType
 
-    log_lines = []
-    def log(msg):
-        print(msg)
-        log_lines.append(msg)
-        if progress_callback:
-            progress_callback("\n".join(log_lines))
-
-    log(f"🖥️  Device: {DEVICE}")
-    log(f"🧠 Model: {model_name}")
-    log(f"🔁 Iterations: {num_iterations}")
-    log("📦 Loading model and tokenizer...")
+    short_name = model_name.split("/")[-1]
+    print(f"\n{'='*60}")
+    print(f"  🧠 Loading: {short_name}")
+    print(f"{'='*60}")
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
     model = AutoModelForCausalLM.from_pretrained(
         model_name, torch_dtype=torch.bfloat16, token=hf_token,
-        attn_implementation="sdpa",  # PyTorch built-in, no flash-attn package needed
+        attn_implementation="sdpa",
     ).to(DEVICE)
 
-    lora_config = LoraConfig(
+    lora_cfg = LoraConfig(
         r=16, lora_alpha=16,
         target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
-        lora_dropout=0.0, bias="none",
-        task_type=TaskType.CAUSAL_LM,
+        lora_dropout=0.0, bias="none", task_type=TaskType.CAUSAL_LM,
     )
-    model = get_peft_model(model, lora_config)
-
+    model = get_peft_model(model, lora_cfg)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
     trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
     total = sum(p.numel() for p in model.parameters())
-    log(f"✅ Model loaded. Trainable: {trainable:,} / {total:,} params")
+    print(f"  ✅ Loaded | Trainable: {trainable:,} / {total:,}")
+
+    if torch.cuda.is_available():
+        vram = torch.cuda.memory_allocated() / 1e9
+        print(f"  📊 VRAM used: {vram:.2f} GB")
 
     optimizer = torch.optim.AdamW(
         filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate
     )
     env = SB3Adapter()
-
     all_rewards = []
-    full_log = []
 
     # Pre-training eval
-    log("\n▶ PRE-TRAINING EVAL")
+    print(f"\n  ▶ PRE-TRAINING EVAL")
     model.eval()
-    pre_reward, pre_log_data, _ = run_episode(model, tokenizer, env, steps_per_episode=steps_per_episode)
-    all_rewards.append(pre_reward)
-    full_log.append({"phase": "pre-training", "reward": pre_reward, "reasoning": pre_log_data})
-    log(f"  Reward: {pre_reward:+.3f}")
+    pre_r, _ = run_episode(model, tokenizer, env, steps_per_episode=steps_per_episode)
+    all_rewards.append(pre_r)
+    print(f"    Baseline reward: {pre_r:+.3f}")
 
-    # Training
-    log(f"\n▶ TRAINING ({num_iterations} iterations, grad_accum={GRAD_ACCUM_STEPS})")
+    # Training loop
+    print(f"\n  ▶ TRAINING ({num_iterations} iters, accum={GRAD_ACCUM_STEPS})")
     model.train()
+    t0 = time.time()
+
     for i in range(num_iterations):
-        reward, train_log_data, ep_losses = run_episode(
+        reward, log_data = run_episode(
             model, tokenizer, env, is_training=True, optimizer=optimizer,
             steps_per_episode=steps_per_episode,
+            iteration_num=i+1, total_iters=num_iterations,
         )
         all_rewards.append(reward)
-        full_log.append({"phase": f"training-{i+1}", "reward": reward, "reasoning": train_log_data})
 
-        # Step optimizer every GRAD_ACCUM_STEPS episodes (batched update)
         if (i + 1) % GRAD_ACCUM_STEPS == 0:
             torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
             optimizer.step()
             optimizer.zero_grad()
 
-        log(f"  Iter {i+1}/{num_iterations}: reward={reward:+.3f}")
+        # Per-iteration summary
+        pct = ((i+1) / num_iterations) * 100
+        elapsed = time.time() - t0
+        eta = (elapsed / (i+1)) * (num_iterations - i - 1)
+        ema = all_rewards[-1] if len(all_rewards) < 3 else (
+            EMA_ALPHA * all_rewards[-1] + (1 - EMA_ALPHA) * all_rewards[-2]
+        )
+        print(f"  ┃ [{pct:5.1f}%] Iter {i+1:3d}/{num_iterations} │ "
+              f"r={reward:+.3f} │ EMA={ema:+.3f} │ "
+              f"ETA={eta:.0f}s")
 
-    # Final optimizer step for remaining accumulated gradients
+    # Final grad step
     torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
     optimizer.step()
 
     # Post-training eval
-    log("\n▶ POST-TRAINING EVAL")
+    print(f"\n  ▶ POST-TRAINING EVAL")
     model.eval()
-    post_reward, post_log_data, _ = run_episode(model, tokenizer, env, steps_per_episode=steps_per_episode)
-    all_rewards.append(post_reward)
-    full_log.append({"phase": "post-training", "reward": post_reward, "reasoning": post_log_data})
-    log(f"  Reward: {post_reward:+.3f}")
+    post_r, _ = run_episode(model, tokenizer, env, steps_per_episode=steps_per_episode)
+    all_rewards.append(post_r)
+    delta = post_r - pre_r
+    print(f"    Final reward: {post_r:+.3f} (Δ={delta:+.3f})")
+    print(f"    Time: {time.time()-t0:.0f}s")
 
-    delta = all_rewards[-1] - all_rewards[0]
-    log(f"\n✅ DONE | Pre: {all_rewards[0]:+.3f} → Post: {all_rewards[-1]:+.3f} | Δ={delta:+.3f}")
+    # Cleanup VRAM
+    nuke_vram(model, optimizer, tokenizer)
 
-    # Save log
-    with open("outputs/llm_training_log.json", "w") as f:
-        json.dump(full_log, f, indent=2, default=str)
+    return all_rewards
 
-    # Generate graph
-    graph_path = _generate_graph(all_rewards, num_iterations, model_name)
 
-    return all_rewards, full_log, graph_path, "\n".join(log_lines)
+# ── EMA Graph ─────────────────────────────────────────────────────────────────
 
+def compute_ema(rewards, alpha=EMA_ALPHA):
+    ema = [rewards[0]]
+    for r in rewards[1:]:
+        ema.append(alpha * r + (1 - alpha) * ema[-1])
+    return ema
 
-def _generate_graph(all_rewards, num_iterations, model_name):
-    labels = ["Before"] + [f"Iter {i+1}" for i in range(num_iterations)] + ["After"]
-    colors = ["#ef4444"] + ["#3b82f6"] * num_iterations + ["#22c55e"]
 
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6), facecolor="#0e1117")
+def generate_comparison_graph(all_results, output_path="outputs/multi_model_comparison.png"):
+    BG = '#0e1117'
+    COLORS = ['#00d4ff', '#ffa500', '#39ff14', '#ff6b6b', '#b47eff']
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 7), facecolor=BG)
     for ax in [ax1, ax2]:
-        ax.set_facecolor("#0e1117")
-        ax.tick_params(colors="#e6e6e6")
-        ax.grid(axis="y", alpha=0.1, color="white")
-        for s in ['top','right']:
+        ax.set_facecolor(BG)
+        ax.tick_params(colors='#e6e6e6', labelsize=10)
+        ax.grid(True, alpha=0.08, color='white')
+        for s in ['top', 'right']:
             ax.spines[s].set_visible(False)
-        for s in ['left','bottom']:
+        for s in ['left', 'bottom']:
             ax.spines[s].set_color('#333')
 
-    ax1.bar(range(len(all_rewards)), all_rewards, color=colors, edgecolor="white", lw=1.5, width=0.6)
-    ax1.set_xticks(range(len(labels)))
-    ax1.set_xticklabels(labels, fontsize=8, color="#e6e6e6", rotation=45)
-    ax1.set_title(f"LLM RL: {model_name.split('/')[-1]}", color="#e6e6e6", fontsize=13, fontweight="bold")
-    ax1.set_ylabel("Reward", color="#e6e6e6")
-
-    comp = [all_rewards[0], all_rewards[-1]]
-    ax2.bar(["Before", "After"], comp, color=["#ef4444", "#22c55e"], edgecolor="white", lw=2, width=0.5)
-    ax2.set_title("Before vs After", color="#e6e6e6", fontsize=13, fontweight="bold")
-    ax2.set_ylabel("Reward", color="#e6e6e6")
+    # Left: EMA reward curves
+    for idx, (name, rewards) in enumerate(all_results.items()):
+        color = COLORS[idx % len(COLORS)]
+        ema = compute_ema(rewards)
+        ax1.plot(ema, color=color, lw=2.5, label=name, alpha=0.9)
+        ax1.plot(rewards, color=color, lw=0.5, alpha=0.2)
+
+    ax1.set_title("Training Reward (EMA Smoothed)", color='#e6e6e6', fontsize=14, fontweight='bold')
+    ax1.set_xlabel("Episode", color='#e6e6e6', fontsize=11)
+    ax1.set_ylabel("Reward", color='#e6e6e6', fontsize=11)
+    ax1.legend(facecolor='#1a1a2e', edgecolor='#333', labelcolor='#e6e6e6', fontsize=9)
+    ax1.axhline(y=0, color='gray', linestyle='--', alpha=0.3)
+
+    # Right: Before vs After comparison bars
+    names = list(all_results.keys())
+    pre_scores = [all_results[n][0] for n in names]
+    post_scores = [all_results[n][-1] for n in names]
+
+    x = np.arange(len(names))
+    w = 0.35
+    bars1 = ax2.bar(x - w/2, pre_scores, w, label='Before', color='#ef4444', edgecolor='white', lw=1)
+    bars2 = ax2.bar(x + w/2, post_scores, w, label='After', color='#22c55e', edgecolor='white', lw=1)
+
+    ax2.set_xticks(x)
+    ax2.set_xticklabels(names, fontsize=8, color='#e6e6e6', rotation=15)
+    ax2.set_title("Pre vs Post Training", color='#e6e6e6', fontsize=14, fontweight='bold')
+    ax2.set_ylabel("Reward", color='#e6e6e6', fontsize=11)
+    ax2.legend(facecolor='#1a1a2e', edgecolor='#333', labelcolor='#e6e6e6')
+    ax2.axhline(y=0, color='gray', linestyle='--', alpha=0.3)
+
+    for bar, val in zip(list(bars1) + list(bars2), pre_scores + post_scores):
+        ax2.text(bar.get_x() + bar.get_width()/2, val + 0.1,
+                 f"{val:+.1f}", ha='center', va='bottom', fontsize=9,
+                 color='#e6e6e6', fontweight='bold')
 
     plt.tight_layout()
-    path = "outputs/llm_training_results.png"
-    plt.savefig(path, dpi=200, bbox_inches="tight", facecolor="#0e1117")
+    plt.savefig(output_path, dpi=200, bbox_inches='tight', facecolor=BG)
     plt.close()
-    return path
+    return output_path
+
+
+# ── Main Pipeline ─────────────────────────────────────────────────────────────
+
+def train_llm(model_name=None, num_iterations=200, steps_per_episode=15,
+              learning_rate=2e-6, progress_callback=None):
+    """
+    Multi-model or single-model training pipeline.
+    If model_name contains commas, runs multi-model benchmark.
+    """
+    log_lines = []
+    def log(msg):
+        print(msg)
+        log_lines.append(msg)
+        if progress_callback:
+            progress_callback("\n".join(log_lines))
+
+    # Determine model list
+    if model_name and "," in model_name:
+        models = [m.strip() for m in model_name.split(",")]
+    elif model_name:
+        models = [model_name.strip()]
+    else:
+        models = MODELS_TO_TEST
+
+    log(f"🖥️  Device: {DEVICE}")
+    log(f"🔁 Models to test: {len(models)}")
+    for m in models:
+        log(f"   • {m}")
+
+    all_results = {}
+    full_log = []
+
+    for model_idx, mname in enumerate(models):
+        short = mname.split("/")[-1]
+        log(f"\n{'━'*60}")
+        log(f"  [{model_idx+1}/{len(models)}] {short}")
+        log(f"{'━'*60}")
+
+        try:
+            rewards = train_single_model(
+                mname, num_iterations=num_iterations,
+                steps_per_episode=steps_per_episode,
+                learning_rate=learning_rate,
+            )
+            all_results[short] = rewards
+            delta = rewards[-1] - rewards[0]
+            log(f"  ✅ {short}: Pre={rewards[0]:+.3f} → Post={rewards[-1]:+.3f} (Δ={delta:+.3f})")
+            full_log.append({
+                "model": mname, "pre": rewards[0], "post": rewards[-1],
+                "delta": delta, "all_rewards": rewards,
+            })
+        except Exception as e:
+            log(f"  ❌ {short} FAILED: {e}")
+            full_log.append({"model": mname, "error": str(e)})
+            nuke_vram()  # cleanup even on failure
+
+    # Generate comparison graph
+    graph_path = None
+    if all_results:
+        os.makedirs("outputs", exist_ok=True)
+        graph_path = generate_comparison_graph(all_results)
+        log(f"\n📊 Comparison graph saved to {graph_path}")
+
+    # Save log
+    with open("outputs/multi_model_log.json", "w") as f:
+        json.dump(full_log, f, indent=2, default=str)
+
+    # Build flat reward list for backward compat
+    flat_rewards = []
+    for rewards in all_results.values():
+        flat_rewards.extend(rewards)
+
+    log_text = "\n".join(log_lines)
+    return flat_rewards or [0], full_log, graph_path, log_text

requirements.txt

@@ -14,3 +14,4 @@ peft==0.12.0
 accelerate==0.33.0
 bitsandbytes>=0.43.0
 sentencepiece
+unsloth