Create grpo_vertex_v2_ipynb.md

Current training implementation, full run results below:
```json
{
"_runtime": 53739,
"_step": 1261,
"_timestamp": 1776946956.167131,
"_wandb.runtime": 53739,
"eval/best_reward_final": 0.125,
"profiling/Time taken: UnslothGRPOTrainer._calculate_rewards": 0.00957904900133144,
"profiling/Time taken: UnslothGRPOTrainer._prepare_inputs": 0.000010281000868417324,
"profiling/Time taken: UnslothGRPOTrainer.commerce_reward_fn": 0.008595035003963858,
"profiling/Time taken: UnslothGRPOTrainer.transformers.generate": 185.63205686000583,
"total_flos": 0,
"train/clip_ratio/high_max": 0,
"train/clip_ratio/high_mean": 0,
"train/clip_ratio/low_mean": 0,
"train/clip_ratio/low_min": 0,
"train/clip_ratio/region_mean": 0,
"train/completion_length": 2048,
"train/duration_hours": 14.927390168110527,
"train/epoch": 0.7,
"train/final_loss": -0.00020499005976965976,
"train/frac_reward_zero_std": 0,
"train/global_step": 210,
"train/grad_norm": 0.02984152734279633,
"train/kl": 0.004119975958019495,
"train/learning_rate": 1.2752757044047826e-7,
"train/loss": 0,
"train/reward": 0.2850000262260437,
"train/reward_std": 0.10184022039175034,
"train/rewards/commerce_reward_fn/mean": 0.2850000262260437,
"train/rewards/commerce_reward_fn/std": 0.10184022039175034,
"train/total_steps": 210,
"train_loss": -0.00020499005976965976,
"train_runtime": 53734.2335,
"train_samples_per_second": 0.045,
"train_steps_per_second": 0.006
}
```

grpo_vertex_v2_ipynb.md

@@ -0,0 +1,1521 @@
+# Tucano2 Commerce — GRPO Training v2 (Vertex AI Workbench / L4)
+
+**v2 changes over v1:**
+- `UnslothGRPOTrainer` — wraps generation with `for_inference()`/`for_training()` (~2-3× speedup)
+- `processing_class=tokenizer` fix (was silently dropped in v1)
+- `reward_extraction` normalized to max 1.0 (was 2.0 — biased gradient scale)
+- KV cache diagnostic cell (Cell 5b)
+- Eval callback capped to `EVAL_MAX_SAMPLES=10` + `EVAL_MAX_TOKENS=256` (vs 45 × 591s = 7.4h/eval)
+- `UNSLOTH_COMPILE_DISABLE=1` — prevents kernel recompilation on mode switches
+- Optional `USE_VLLM=True` path for 10-20× generation speedup
+
+Ported from `tucano2_pipeline/06_rlvr.py` (Modal version).  
+Run incrementally: each cell is a gate — verify output before moving to next.
+
+**Prerequisites:**
+- Upload `data/pairs/train.jsonl` (2.1 MB) to `./data/pairs/`
+- Upload `models/tucano2-commerce-sft/` (126 MB) to `./models/tucano2-commerce-sft/`
+
+**Hardware:** L4 (24GB), PyTorch kernel, bf16 supported
+
+## Cell 1: Dependencies
+Restart your kernel first (Kernel → Restart), then run these cells in order, one at a time:
+
+
+```python
+# Cell 1 — Nuke everything ML-related
+!pip uninstall -y torch torchvision torchaudio \
+    unsloth unsloth-zoo \
+    trl transformers peft accelerate \
+    bitsandbytes vllm vllm-flash-attn \
+    datasets tokenizers safetensors huggingface-hub \
+    wandb xformers triton \
+    cuda-bindings cuda-python \
+    sentencepiece protobuf \
+    2>/dev/null
+```
+
+
+```python
+# Cell 2 — Kill any stragglers (run twice if paranoid)
+!pip freeze | grep -iE "torch|unsloth|trl|vllm|bitsandbytes|transformers|peft|accelerate" | xargs pip uninstall -y 2>/dev/null
+```
+
+    Found existing installation: torch_c_dlpack_ext 0.1.5
+    Uninstalling torch_c_dlpack_ext-0.1.5:
+      Successfully uninstalled torch_c_dlpack_ext-0.1.5
+    Found existing installation: torchao 0.17.0
+    Uninstalling torchao-0.17.0:
+      Successfully uninstalled torchao-0.17.0
+
+
+
+```python
+# Cell 3 — Purge cache
+!pip cache purge
+```
+
+    Files removed: 918 (11059.0 MB)
+    Directories removed: 0
+
+
+Restart kernel again, then:
+
+
+```python
+# Cell 4 — Clean install, correct order
+!pip install "unsloth"
+```
+
+
+```python
+# Cell 5 — Pin TRL (Unsloth may pull a different version)
+!pip install "trl==0.24.0" --no-deps
+```
+
+    Requirement already satisfied: trl==0.24.0 in /opt/conda/envs/pytorch/lib/python3.10/site-packages (0.24.0)
+
+
+
+```python
+!pip install "rich" "wandb"
+
+```
+
+## Cell 2: Hello World — GPU + Unsloth verification
+
+
+```python
+import torch
+
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"GPU: {torch.cuda.get_device_name(0)}")
+print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+print(f"bf16 support: {torch.cuda.is_bf16_supported()}")
+
+from unsloth import FastLanguageModel
+print("\n✓ Unsloth loaded successfully")
+
+import trl
+print(f"✓ TRL version: {trl.__version__}")
+
+import transformers
+print(f"✓ Transformers version: {transformers.__version__}")
+```
+
+    CUDA available: True
+    GPU: NVIDIA L4
+    VRAM: 23.6 GB
+    bf16 support: True
+    🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.
+    🦥 Unsloth Zoo will now patch everything to make training faster!
+    
+    ✓ Unsloth loaded successfully
+    ✓ TRL version: 0.24.0
+    ✓ Transformers version: 4.57.6
+
+
+## Cell 3: Config + Constants
+
+All config from `06_rlvr.py` — no Modal dependencies.
+
+
+```python
+import os
+# ── v2: Disable Unsloth kernel recompilation between mode switches ─────────────
+# Without this, for_inference() / for_training() trigger expensive Triton recompiles.
+os.environ["UNSLOTH_COMPILE_DISABLE"] = "1"
+
+import json
+import re
+import time
+import random
+import os
+from pathlib import Path
+
+# ── Config ──────────────────────���───────────────────────────────���─────────────
+MODEL_ID = "Polygl0t/Tucano2-qwen-3.7B-Think"
+MAX_SEQ_LENGTH = 4096  # context window (prompt + generation); model supports up to 32k
+
+# Paths (Vertex AI Workbench — cwd is /home/jupyter/)
+DATA_DIR = Path("/home/jupyter/tucano2/data")
+MODELS_DIR = Path("/home/jupyter/tucano2/models")
+SFT_ADAPTER_DIR = MODELS_DIR / "tucano2-commerce-sft"
+GRPO_ADAPTER_DIR = MODELS_DIR / "tucano2-commerce-grpo"
+CHECKPOINT_DIR = GRPO_ADAPTER_DIR / "checkpoints"
+
+GRPO_PROMPTS = 300  # stratified subset size (120/120/30/30)
+
+# Valid enum values for reward scoring
+VALID_SENTIMENTS = {"positive", "negative", "neutral"}
+VALID_CATEGORIES = {
+    "delivery_delay", "product_quality", "product_not_received",
+    "wrong_product", "seller_communication", "app_issue",
+    "price_value", "other", "none",
+}
+VALID_CHURN = {"low", "medium", "high"}
+VALID_REPEAT = {"yes", "no", "maybe"}
+EXTRACTION_FIELDS = [
+    "sentiment", "sentiment_score", "churn_risk", "delivery_issue",
+    "product_issue", "seller_issue", "main_complaint",
+    "complaint_category", "repeat_intent", "would_recommend",
+]
+
+SYSTEM_PT = (
+    "Você é um assistente de IA especializado em análise de e-commerce brasileiro. "
+    "Você compreende avaliações de clientes em português e padrões de comércio brasileiro."
+)
+
+# Training params (validated on L4, optimized for H100)
+BATCH_SIZE = 4       # L4 survived batch=4 (unsloth auto-adjusted); H100 has 80GB headroom
+GRAD_ACCUM = 2       # effective batch = 4 * 2 = 8
+MAX_COMPLETION_LENGTH = 2048  # non-negotiable: model needs room to think + answer
+NUM_GENERATIONS = 8        # was 4 → more samples = higher chance of variance for GRPO
+SCALE_REWARDS = False      # Dr. GRPO fix: remove std normalization bias
+LEARNING_RATE = 5e-7
+NUM_EPOCHS = 2
+TEMPERATURE = 0.8          # was 0.1 from model defaults
+MAX_STEPS = 300       # -1 = full run (75 steps); set to e.g. 3 for probe
+
+
+# ── ADR: Checkpoint + Eval + Early-Stop params ────────────────────────────────
+EVAL_SPLIT_RATIO        = 0.15  # 15% of each task bucket held out for eval
+EVAL_STEPS              = 10    # run EvalRewardCallback every N steps
+EARLY_STOPPING_PATIENCE = 10    # was 3 — gives 100 steps of runway before halting
+EARLY_STOPPING_DELTA    = 0.01  # min reward gain to count as "improvement"
+SAVE_STEPS              = 15    # checkpoint every ~1h on L4 (Spot VM safety)
+SAVE_TOTAL_LIMIT        = 3     # auto-prune old checkpoints, keep last 3
+WANDB_PROJECT           = "tucano2-commerce"
+
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+print("✓ Config loaded")
+print(f"  SFT adapter: {SFT_ADAPTER_DIR} (exists: {SFT_ADAPTER_DIR.exists()})")
+print(f"  Train data: {DATA_DIR / 'pairs' / 'train.jsonl'} (exists: {(DATA_DIR / 'pairs' / 'train.jsonl').exists()})")
+print(f"  Training: batch={BATCH_SIZE}, grad_accum={GRAD_ACCUM}, eff_batch={BATCH_SIZE*GRAD_ACCUM}")
+print(f"  Steps: {(GRPO_PROMPTS * NUM_GENERATIONS * NUM_EPOCHS) // (BATCH_SIZE * GRAD_ACCUM)} (full run)")
+print(f"  ADR: save_steps={SAVE_STEPS}, eval_steps={EVAL_STEPS}, patience={EARLY_STOPPING_PATIENCE}, eval_split={EVAL_SPLIT_RATIO}")
+
+# ══════════════════════════════════════════════════════════════════════════════
+# v2 Performance & Safety Flags
+# ══════════════════════════════════════════════════════════════════════════════
+
+# ── Generation backend ───────────────────────────────────────────────────────
+USE_VLLM = False   # Flip True if vllm is installed and VRAM allows
+                   # Requires: pip install "vllm>=0.6.0"
+                   # Enables vllm_mode="colocate" + vllm_enable_sleep_mode=True
+
+# ── Eval callback safety caps ────────────────────────────────────────────────
+# AT 591s/SAMPLE: 45 eval samples = 7.4h PER EVALUATION PASS — breaks training loop!
+EVAL_MAX_SAMPLES = 5    # keep eval time manageable (~15 min per eval). it was 10 to cap to first N samples from eval_dataset
+EVAL_MAX_TOKENS  = 2048   # meaningful eval metric. it was 256 to keeps each eval pass < 15min
+
+# ── TRL version assertion (UnslothGRPOTrainer overrides _generate) ───────────
+import trl as _trl
+assert _trl.__version__ == "0.24.0", (
+    f"UnslothGRPOTrainer was written for TRL 0.24.0, found {_trl.__version__}.\n"
+    "Verify that GRPOTrainer._generate() still exists before proceeding."
+)
+print(f"✓ TRL {_trl.__version__} verified")
+
+```
+
+    ✓ Config loaded
+      SFT adapter: /home/jupyter/tucano2/models/tucano2-commerce-sft (exists: True)
+      Train data: /home/jupyter/tucano2/data/pairs/train.jsonl (exists: True)
+      Training: batch=4, grad_accum=2, eff_batch=8
+      Steps: 75 (full run)
+      ADR: save_steps=5, eval_steps=10, patience=10, eval_split=0.15
+    ✓ TRL 0.24.0 verified
+
+
+## Cell 4: Load SFT Adapter
+
+
+```python
+print("Loading SFT adapter...")
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=str(SFT_ADAPTER_DIR),
+    max_seq_length=MAX_SEQ_LENGTH,
+    load_in_4bit=True,
+    dtype=None,
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Load chat template from base model (SFT adapter doesn't save it)
+from transformers import AutoTokenizer
+base_tok = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer.chat_template = base_tok.chat_template
+del base_tok
+
+# ── v2: Force KV cache — Unsloth patching may reset this ─────────────────────
+model.config.use_cache = True
+model.generation_config.use_cache = True
+
+print(f"✓ Model loaded on {model.device}")
+print(f"  use_cache: {model.config.use_cache}")
+print(f"  Params: {sum(p.numel() for p in model.parameters()) / 1e6:.0f}M")
+print(f"  Chat template: {tokenizer.chat_template[:50]}...")
+```
+
+    Loading SFT adapter...
+    ==((====))==  Unsloth 2026.4.6: Fast Qwen3 patching. Transformers: 4.57.6. vLLM: 0.19.1.
+       \\   /|    NVIDIA L4. Num GPUs = 1. Max memory: 21.951 GB. Platform: Linux.
+    O^O/ \_/ \    Torch: 2.10.0+cu128. CUDA: 8.9. CUDA Toolkit: 12.8. Triton: 3.6.0
+    \        /    Bfloat16 = TRUE. FA [Xformers = 0.0.35. FA2 = False]
+     "-____-"     Free license: http://github.com/unslothai/unsloth
+    Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!
+
+
+
+    Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]
+
+
+    Unsloth 2026.4.6 patched 36 layers with 0 QKV layers, 0 O layers and 0 MLP layers.
+
+
+    ✓ Model loaded on cuda:0
+      use_cache: True
+      Params: 1976M
+      Chat template: {#- Handle tool/function calling setup #}
+    {%- if t...
+
+
+## Cell 5: Single Inference Test
+
+**Gate:** Does the model close `</think>` and produce an answer within 2048 tokens?
+
+
+```python
+FastLanguageModel.for_inference(model)
+
+test_msgs = [
+    {"role": "system", "content": SYSTEM_PT},
+    {"role": "user", "content": "Quais são as categorias de reclamação mais frequentes e como afetam a nota média?"},
+]
+text = tokenizer.apply_chat_template(test_msgs, tokenize=False, add_generation_prompt=True)
+inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+t0 = time.time()
+outputs = model.generate(**inputs, max_new_tokens=2048, temperature=0.7, do_sample=True)
+elapsed = time.time() - t0
+
+response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+
+print(f"Generation time: {elapsed:.1f}s")
+print(f"Response length: {len(response)} chars")
+print(f"closed_think: {'</think>' in response}")
+print(f"\n{'='*60}")
+print(response[:800])
+```
+
+    Generation time: 66.8s
+    Response length: 3945 chars
+    closed_think: True
+    
+    ============================================================
+    <think>
+    Para responder à pergunta sobre quais tipos do conjunto geral têm uma classificação melhor que o outro, vamos analisar os dados fornecidos passo a passo:
+    
+    ### Dados Disponíveis:
+    1. **Classificação Geral:** 4.83/5 estrelas (N = 2750)
+    2. **Reclamações Gerais Não Específicas:** 4.68/5 estrelas (2750 participantes)
+    3. **Problemas com Remessa/Envio:** 4.51/5 estrelas (1733 participantes)
+    4. **Não Respondido Rapidamente Após Contato Inicial:** 4.63/5 estrelas (1263 participantes)
+    5. **Sem Reclamação** (5.00/5 estrelas 1090 participantes)
+    6. **Outras Reivindicações (múltiplas):** 4.76/5 estrelas (919 participantes)
+    7. **Fora Do Prazo De Garantia Ou Válvula:** 4.49/5 Estrelas (746 Participantes)
+    8. **Erro Na Entrega ou Extraído Inadvertidamente:** 4.43/5 ESTRELAS (437 PARTICIPANTES)
+    9. **D
+
+
+## Cell 5b: KV Cache Diagnostic
+
+**Gate:** Ratio < 3× → KV cache is working. Ratio > 5× → cache is broken (O(n²) full attention recompute at every token — training generation will be catastrophically slow).
+
+If broken, check `model.config.use_cache` and try `UNSLOTH_COMPILE_DISABLE=1`.
+
+
+
+```python
+# ── KV Cache Diagnostic ───────────────────────────────────────────────────────
+# Tests whether attention past_key_values are actually being used.
+# O(n²) failure: token time grows linearly with sequence length → ratio >> 1.
+import time
+FastLanguageModel.for_inference(model)
+
+_kv_msgs = [{"role": "system", "content": SYSTEM_PT},
+            {"role": "user", "content": "Qual a categoria de reclamação mais frequente?"}]
+_kv_text   = tokenizer.apply_chat_template(_kv_msgs, tokenize=False, add_generation_prompt=True)
+_kv_inputs = tokenizer(_kv_text, return_tensors="pt").to(model.device)
+
+_token_times, _past, _generated = [], None, _kv_inputs["input_ids"]
+with torch.no_grad():
+    for _step in range(50):
+        _t0 = time.time()
+        
+        # Calculate current sequence length
+        seq_len = _generated.shape[1]
+        
+        # Manually construct position_ids
+        if _past is None:
+            # First pass: we need positions for the whole prompt (e.g., [0, 1, 2, ...])
+            _position_ids = torch.arange(seq_len, dtype=torch.long, device=model.device).unsqueeze(0)
+        else:
+            # Subsequent passes: we only need the position for the single new token
+            _position_ids = torch.tensor([[seq_len - 1]], dtype=torch.long, device=model.device)
+            
+        _out = model(
+            input_ids=_generated[:, -1:] if _past else _generated,
+            position_ids=_position_ids, # <--- The missing argument!
+            attention_mask=torch.ones(1, seq_len, device=model.device),
+            past_key_values=_past,
+            use_cache=True,
+            return_dict=True, # <--- Forces it to return an object instead of a tuple!
+        )
+        _past = _out.past_key_values
+        _next  = _out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
+        _generated = torch.cat([_generated, _next], dim=1)
+        _token_times.append(time.time() - _t0)
+
+_ratio = sum(_token_times[45:]) / max(sum(_token_times[:5]), 1e-9)
+print(f"First 5 tok : {[f'{t*1000:.0f}ms' for t in _token_times[:5]]}")
+print(f"Last  5 tok : {[f'{t*1000:.0f}ms' for t in _token_times[45:]]}")
+print(f"Ratio last/first: {_ratio:.1f}x")
+if _ratio < 3:
+    print("✓ KV cache is working correctly")
+elif _ratio < 6:
+    print("⚠ KV cache may be degraded — check model.config.use_cache")
+else:
+    print("✗ KV cache BROKEN — O(n²) recompute. GRPO generation will be catastrophically slow.")
+    print("  Try: model.config.use_cache = True; model.generation_config.use_cache = True")
+
+# Clean up cache vars
+del _past, _generated, _kv_inputs, _token_times, _out
+import gc; gc.collect()
+if torch.cuda.is_available(): torch.cuda.empty_cache()
+
+```
+
+    First 5 tok : ['277ms', '95ms', '87ms', '86ms', '86ms']
+    Last  5 tok : ['88ms', '88ms', '86ms', '86ms', '86ms']
+    Ratio last/first: 0.7x
+    ✓ KV cache is working correctly
+
+
+## Cell 6: Reward Functions
+
+Copied verbatim from `06_rlvr.py`. Pure Python — no external dependencies.
+
+
+```python
+def strip_think(text: str) -> str:
+    """Remove <think>...</think> block, return the answer portion."""
+    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
+
+
+def has_think_block(text: str) -> bool:
+    """Check if text contains a non-empty <think> block."""
+    return bool(re.search(r"<think>.+</think>", text, flags=re.DOTALL))
+
+
+def _classify_task_type(prompt_text: str) -> str:
+    """Classify prompt into task type by keywords."""
+    p = prompt_text.lower()
+    if "retorne um objeto json" in p or "extraia dados" in p:
+        return "extraction"
+    elif "notificação push" in p or "notificação de reengajamento" in p:
+        return "push"
+    elif "perfil do cliente" in p:
+        return "insights"
+    else:
+        return "sql_qa"
+
+
+def reward_extraction(completion: str) -> float:
+    """Continuous reward for extraction tasks (max 1.0, normalized)."""
+    score = 0.0
+    answer = strip_think(completion)
+
+    # +0.1 for <think> block (binary but small weight)
+    if has_think_block(completion):
+        score += 0.1
+
+    # JSON parsing: partial credit
+    try:
+        data = json.loads(answer)
+    except (json.JSONDecodeError, TypeError):
+        # Partial credit for JSON-like structure
+        score += 0.05 * _json_similarity(answer)
+        return score
+
+    if not isinstance(data, dict):
+        score += 0.1  # at least it's valid JSON
+        return score
+
+    score += 0.3  # valid JSON object
+
+    # Schema completeness: fractional credit per field
+    present = sum(1 for f in EXTRACTION_FIELDS if f in data)
+    score += 0.3 * (present / len(EXTRACTION_FIELDS))
+
+    # Categorical correctness: fractional per field
+    cat_checks = 0
+    cat_total = 0
+
+    checks = [
+        ("sentiment", lambda v: v in VALID_SENTIMENTS),
+        ("complaint_category", lambda v: v in VALID_CATEGORIES),
+        ("churn_risk", lambda v: v in VALID_CHURN),
+        ("repeat_intent", lambda v: v in VALID_REPEAT),
+        ("sentiment_score", lambda v: isinstance(v, (int, float)) and 1 <= v <= 5),
+    ]
+    for field, validator in checks:
+        cat_total += 1
+        if field in data and validator(data[field]):
+            cat_checks += 1
+
+    for bool_field in ("delivery_issue", "product_issue", "seller_issue", "would_recommend"):
+        cat_total += 1
+        if bool_field in data and isinstance(data[bool_field], bool):
+            cat_checks += 1
+
+    if cat_total > 0:
+        score += 0.3 * (cat_checks / cat_total)
+
+    return min(score, 1.0)  # cap at 1.0
+    
+
+def _json_similarity(text: str) -> float:
+    """Rough heuristic: how JSON-like is this text? 0.0 to 1.0."""
+    text = text.strip()
+    if not text:
+        return 0.0
+    score = 0.0
+    if text.startswith("{") and text.endswith("}"):
+        score += 0.5
+    if '"' in text:
+        score += 0.2
+    if ":" in text:
+        score += 0.2
+    if "," in text:
+        score += 0.1
+    return min(score, 1.0)
+
+
+def reward_sql_qa(completion: str) -> float:
+    """Continuous reward for SQL Q&A (max 1.0)."""
+    score = 0.0
+    answer = strip_think(completion)
+
+    if has_think_block(completion):
+        score += 0.1
+
+    # Numerical content: more numbers = more specific answer
+    numbers = re.findall(r"\d+(?:[.,]\d+)?", answer)
+    score += min(0.4, 0.1 * len(numbers))  # up to 0.4 for multiple numbers
+
+    # Length: optimal is 100-400 chars. Penalize too short or too long.
+    length = len(answer)
+    if 50 <= length <= 500:
+        score += 0.3
+    elif length > 0:
+        score += 0.3 * max(0, 1 - abs(length - 275) / 225)  # linear falloff
+
+    # SQL keywords: evidence of actual query analysis
+    sql_keywords = r"SELECT|FROM|WHERE|GROUP BY|ORDER BY|COUNT|SUM|AVG|JOIN"
+    matches = len(re.findall(sql_keywords, answer, re.IGNORECASE))
+    score += min(0.2, 0.05 * matches)
+
+    return min(score, 1.0)
+
+
+def reward_insights(completion: str) -> float:
+    """Continuous reward for insights (max 1.0)."""
+    score = 0.0
+    answer = strip_think(completion)
+
+    if has_think_block(completion):
+        score += 0.1
+
+    # Actionable language: count matches, not binary
+    action_words = ["recomend", "implement", "melhor", "reduzir", "aumentar",
+                    "priorizar", "investir", "otimizar", "estratégi"]
+    matches = sum(1 for w in action_words if w in answer.lower())
+    score += min(0.5, 0.1 * matches)
+
+    # Length: 100-1000 chars optimal
+    length = len(answer)
+    if 100 <= length <= 1000:
+        score += 0.3
+    elif length > 0:
+        score += 0.3 * max(0, 1 - abs(length - 550) / 450)
+
+    # Structure: bullet points, numbered lists = organized thinking
+    structure_marks = len(re.findall(r"^[-•*]\s|^\d+[.)]\s", answer, re.MULTILINE))
+    score += min(0.1, 0.02 * structure_marks)
+
+    return min(score, 1.0)
+
+
+def reward_push(completion: str) -> float:
+    """Continuous reward for push notifications (max 1.0)."""
+    score = 0.0
+    answer = strip_think(completion)
+
+    if not answer:
+        return 0.0
+
+    # Length: ≤120 chars gets full credit, linear penalty above
+    length = len(answer)
+    if length <= 120:
+        score += 0.5
+    else:
+        score += 0.5 * max(0, 1 - (length - 120) / 80)
+
+    # Generic-ness: fuzzy penalty based on similarity to generic phrases
+    generic_phrases = [
+        "olá! como podemos ajudar",
+        "obrigado pela sua compra",
+        "seu pedido foi confirmado",
+        "agradecemos sua preferência",
+    ]
+    max_similarity = max(
+        _string_similarity(answer.lower(), g) for g in generic_phrases
+    )
+    score += 0.3 * (1 - max_similarity)  # less generic = higher score
+
+    # Portuguese content: count PT-specific markers
+    pt_markers = re.findall(r"[ãçéêóúâõ]|você|para|como|seu|sua", answer, re.IGNORECASE)
+    score += min(0.2, 0.02 * len(pt_markers))
+
+    return min(score, 1.0)
+
+
+def _string_similarity(a: str, b: str) -> float:
+    """Simple Jaccard-like similarity for short strings. 0.0 to 1.0."""
+    if not a or not b:
+        return 0.0
+    a_set = set(a.split())
+    b_set = set(b.split())
+    intersection = len(a_set & b_set)
+    union = len(a_set | b_set)
+    return intersection / union if union > 0 else 0.0
+
+
+def commerce_reward_fn(completions, prompts, **kwargs) -> list[float]:
+    """Master reward function: dispatches by task type."""
+    rewards = []
+    for completion, prompt in zip(completions, prompts):
+        if isinstance(completion, list):
+            comp_text = completion[-1]["content"] if completion else ""
+        else:
+            comp_text = str(completion)
+
+        if isinstance(prompt, list):
+            prompt_text = " ".join(m.get("content", "") for m in prompt)
+        else:
+            prompt_text = str(prompt)
+
+        task = _classify_task_type(prompt_text)
+
+        if task == "extraction":
+            rewards.append(reward_extraction(comp_text))
+        elif task == "sql_qa":
+            rewards.append(reward_sql_qa(comp_text))
+        elif task == "insights":
+            rewards.append(reward_insights(comp_text))
+        elif task == "push":
+            rewards.append(reward_push(comp_text))
+        else:
+            r = 0.2 if has_think_block(comp_text) else 0.0
+            r += 0.3 if comp_text.strip() else 0.0
+            rewards.append(r)
+
+    return rewards
+
+
+print("✓ Reward functions defined")
+```
+
+    ✓ Reward functions defined
+
+
+## Cell 7: Reward Calibration
+
+**Gate:** Verify reward variance > 0 and most samples close `</think>`.
+
+
+```python
+# Load dataset for calibration
+train_path = DATA_DIR / "pairs" / "train.jsonl"
+
+by_type = {"extraction": [], "sql_qa": [], "insights": [], "push": []}
+with open(train_path) as f:
+    for line in f:
+        row = json.loads(line)
+        convs = row["conversations"]
+        prompt_msgs = [m for m in convs if m["role"] in ("system", "user")]
+        if not prompt_msgs:
+            continue
+        user_text = " ".join(m["content"] for m in prompt_msgs if m["role"] == "user")
+        task = _classify_task_type(user_text)
+        by_type[task].append(prompt_msgs)
+
+print(f"Prompts by type: {', '.join(f'{k}={len(v)}' for k, v in by_type.items())}")
+
+# Pick 5 diverse samples for calibration
+rng = random.Random(42)
+cal_samples = []
+for task_type in ["extraction", "sql_qa", "insights", "push", "sql_qa"]:
+    cal_samples.append(rng.choice(by_type[task_type]))
+
+# Run calibration
+FastLanguageModel.for_inference(model)
+print("\nReward calibration (5 samples):")
+print("-" * 60)
+
+cal_rewards = []
+for i, msgs in enumerate(cal_samples):
+    text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=MAX_COMPLETION_LENGTH, temperature=0.7, do_sample=True)
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    
+    r = commerce_reward_fn([response], [text])[0]
+    cal_rewards.append(r)
+    has_answer = "</think>" in response
+    answer_preview = strip_think(response)[:120] if has_answer else "[stuck in <think>]"
+    print(f"  Sample {i+1}: reward={r:.2f} | closed_think={has_answer} | answer: {answer_preview}")
+
+print(f"\nMean={sum(cal_rewards)/len(cal_rewards):.2f}, Min={min(cal_rewards):.2f}, Max={max(cal_rewards):.2f}")
+print(f"Variance > 0: {len(set(cal_rewards)) > 1}")
+```
+
+    Prompts by type: extraction=659, sql_qa=655, insights=114, push=222
+    
+    Reward calibration (5 samples):
+    ------------------------------------------------------------
+      Sample 1: reward=0.02 | closed_think=False | answer: [stuck in <think>]
+      Sample 2: reward=0.60 | closed_think=False | answer: [stuck in <think>]
+      Sample 3: reward=0.10 | closed_think=True | answer: Para determinarmos se deveríamos oferecer algum tipo de benefício adicional para tentar reverter essa decisão negativa d
+      Sample 4: reward=0.50 | closed_think=True | answer: Olha só![NomeDoCliente]! 😪 [Estado: SÃO PAULO]  
+    PROBLEMA DETECTADO: Produto não recebido.   
+    VALOR DO PEDIDO: R$ 138   
+      Sample 5: reward=0.70 | closed_think=True | answer: Os clientes com baixos índices de rotação (low churn-risk) geralmente enfrentam menos problemas significativos comparado
+    
+    Mean=0.38, Min=0.02, Max=0.70
+    Variance > 0: True
+
+
+## Cell 8: Dataset Preparation
+
+
+```python
+from datasets import Dataset
+
+def prepare_grpo_datasets(n_prompts=GRPO_PROMPTS, eval_ratio=EVAL_SPLIT_RATIO, seed=42):
+    """
+    Stratified split of by_type prompts into train and eval datasets.
+
+    For each task bucket: hold out `eval_ratio` fraction first, then sample
+    train targets from the remainder. Guarantees at least 1 eval sample per type.
+
+    Returns:
+        train_dataset: HF Dataset — passed to GRPOTrainer
+        eval_dataset:  HF Dataset — consumed by EvalRewardCallback
+    """
+    rng = random.Random(seed)
+
+    # ── Step 1: per-task eval hold-out ────────────────────────────────────────
+    train_pools = {}
+    eval_records = []
+    for task, pool in by_type.items():
+        shuffled = pool.copy()
+        rng.shuffle(shuffled)
+        n_eval = max(1, int(len(shuffled) * eval_ratio))
+        eval_records.extend(shuffled[:n_eval])
+        train_pools[task] = shuffled[n_eval:]
+
+    # ── Step 2: stratified train sampling from remaining pool ─────────────────
+    targets = {
+        "extraction": int(n_prompts * 0.4),
+        "sql_qa":     int(n_prompts * 0.4),
+        "insights":   int(n_prompts * 0.1),
+        "push":       int(n_prompts * 0.1),
+    }
+    train_records = []
+    for task, target_n in targets.items():
+        pool = train_pools[task]
+        n = min(target_n, len(pool))
+        train_records.extend(rng.sample(pool, n))
+    rng.shuffle(train_records)
+
+    print(f"Dataset split (eval_ratio={eval_ratio}):")
+    print(f"  train : {len(train_records)} prompts")
+    print(f"  eval  : {len(eval_records)} prompts")
+    print(f"  train dist: {', '.join(f'{k}={min(v, len(train_pools[k]))}' for k, v in targets.items())}")
+
+    train_ds = Dataset.from_list([{"prompt": msgs} for msgs in train_records])
+    eval_ds  = Dataset.from_list([{"prompt": msgs} for msgs in eval_records])
+    return train_ds, eval_ds
+
+
+train_dataset, eval_dataset = prepare_grpo_datasets()
+dataset = train_dataset  # backward compat — smoke/probe cells use `dataset`
+print(f"\n✓ Datasets ready: train={len(train_dataset)}, eval={len(eval_dataset)}")
+
+```
+
+    Dataset split (eval_ratio=0.15):
+      train : 300 prompts
+      eval  : 246 prompts
+      train dist: extraction=120, sql_qa=120, insights=30, push=30
+    
+    ✓ Datasets ready: train=300, eval=246
+
+
+## Cell 9: Smoke Test — Single Training Step
+
+**Gate:** Runs 1 step without OOM. Reports step time for estimation.
+
+
+```python
+from trl import GRPOConfig, GRPOTrainer
+
+# Switch to training mode
+FastLanguageModel.for_training(model)
+
+smoke_config = GRPOConfig(
+    output_dir=str(CHECKPOINT_DIR / "smoke"),
+    num_generations=NUM_GENERATIONS,
+    scale_rewards=SCALE_REWARDS,
+    max_completion_length=MAX_COMPLETION_LENGTH,
+    max_steps=1,
+    num_train_epochs=1,
+    temperature=0.8,          # was 0.1 from model defaults
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=1,  # just 1 for smoke test (faster)
+    learning_rate=LEARNING_RATE,
+    fp16=False,
+    bf16=True,
+    logging_steps=1,
+    save_steps=999,  # don't save during smoke
+    report_to="none",
+    max_prompt_length=MAX_SEQ_LENGTH // 2,
+    seed=42,
+    remove_unused_columns=False,
+)
+
+smoke_trainer = GRPOTrainer(
+    model=model,
+    reward_funcs=commerce_reward_fn,
+    args=smoke_config,
+    train_dataset=dataset,
+    tokenizer=tokenizer,
+)
+
+t0 = time.time()
+smoke_trainer.train()
+step_time = time.time() - t0
+
+# Estimate full run: 75 steps with grad_accum=2
+print(f"\n✓ Smoke test passed!")
+print(f"  Step time (grad_accum=1): {step_time:.0f}s")
+print(f"  Estimated step time (grad_accum={GRAD_ACCUM}): {step_time * GRAD_ACCUM:.0f}s")
+print(f"  Estimated full run (75 steps): {step_time * GRAD_ACCUM * 75 / 3600:.1f}h")
+
+# Cleanup smoke test
+del smoke_trainer
+import gc; gc.collect(); torch.cuda.empty_cache()
+```
+
+    Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.
+    We will change the batch size of 4 to the `num_generations` of 8
+
+
+    ==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1
+       \\   /|    Num examples = 300 | Num Epochs = 1 | Total steps = 1
+    O^O/ \_/ \    Batch size per device = 8 | Gradient accumulation steps = 1
+    \        /    Data Parallel GPUs = 1 | Total batch size (8 x 1 x 1) = 8
+     "-____-"     Trainable parameters = 33,030,144 of 3,792,371,200 (0.87% trained)
+    `generation_config` default values have been modified to match model-specific defaults: {'max_length': 4096, 'repetition_penalty': 1.2, 'renormalize_logits': True}. If this is not desired, please set these values explicitly.
+
+    ✓ Smoke test passed!
+      Step time (grad_accum=1): 318s
+      Estimated step time (grad_accum=2): 636s
+      Estimated full run (75 steps): 13.2h
+
+
+## Cell 10: Probe Run (3 steps)
+
+**Gate:** Loss > 0, rewards have variance, step time is consistent.
+
+
+```python
+FastLanguageModel.for_training(model)
+
+probe_config = GRPOConfig(
+    output_dir=str(CHECKPOINT_DIR / "probe"),
+    num_generations=NUM_GENERATIONS,
+    scale_rewards=SCALE_REWARDS,
+    max_completion_length=MAX_COMPLETION_LENGTH,
+    max_steps=3,
+    temperature=TEMPERATURE,
+    num_train_epochs=NUM_EPOCHS,
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    learning_rate=LEARNING_RATE,
+    warmup_ratio=0.1,
+    lr_scheduler_type="cosine",
+    fp16=False,
+    bf16=True,
+    logging_steps=1,
+    save_steps=999,
+    report_to="none",
+    max_prompt_length=MAX_SEQ_LENGTH // 2,
+    seed=42,
+    remove_unused_columns=False,
+)
+
+probe_trainer = GRPOTrainer(
+    model=model,
+    reward_funcs=commerce_reward_fn,
+    args=probe_config,
+    train_dataset=dataset,
+    tokenizer=tokenizer,
+)
+
+t0 = time.time()
+result = probe_trainer.train()
+elapsed = time.time() - t0
+
+print(f"\n✓ Probe complete in {elapsed:.0f}s ({elapsed/3:.0f}s/step)")
+print(f"  Train loss: {result.training_loss:.4f}")
+print(f"  Estimated full run (75 steps): {elapsed/3 * 75 / 3600:.1f}h")
+
+del probe_trainer
+gc.collect(); torch.cuda.empty_cache()
+```
+
+    ==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1
+       \\   /|    Num examples = 300 | Num Epochs = 1 | Total steps = 3
+    O^O/ \_/ \    Batch size per device = 4 | Gradient accumulation steps = 2
+    \        /    Data Parallel GPUs = 1 | Total batch size (4 x 2 x 1) = 8
+     "-____-"     Trainable parameters = 33,030,144 of 3,792,371,200 (0.87% trained)
+
+
+    Unsloth: Will smartly offload gradients to save VRAM!
+
+    ✓ Probe complete in 665s (222s/step)
+      Train loss: 0.0062
+      Estimated full run (75 steps): 4.6h
+
+
+## Cell 11: Full Training Run
+
+**ADR changes applied here:**
+- `save_steps=15` (was 25) — checkpoint every ~3.3h on L4 Spot VM
+- `save_total_limit=3` — auto-prune old checkpoints
+- `logging_steps=1` (was 5) — every step visible in console + W&B
+- `report_to="wandb"` (was "none") — full run tracked in W&B project `tucano2-commerce`
+- `EvalRewardCallback` — custom callback running reward on held-out eval set every 10 steps
+- **Early stopping**: halt if `mean_eval_reward` fails to improve ≥ 0.01 for 3 consecutive evals
+
+**Resume:** Automatically resumes from latest checkpoint if interrupted.
+
+
+
+```python
+# DEPRECATED_Cell: Fixed Safety Validation
+
+import torch
+from unsloth import FastLanguageModel
+
+# Get the inner model
+# was: policy_inner = getattr(trainer.model, 'module', trainer.model)
+policy_inner = model
+
+# ── CHECK 1: PASSED (no ref model) ──────────────────────────
+print("CHECK 1: ✅ No reference model — biggest risk eliminated\n")
+
+# ── CHECK 2: Weight drift — test LoRA A/B matrices (bf16) ───
+print("=" * 60)
+print("CHECK 2: LoRA adapter weight drift after merge/unmerge")
+
+test_layer = None
+for name, module in policy_inner.named_modules():
+    if hasattr(module, 'lora_A') and hasattr(module, 'lora_B'):
+        test_layer = module
+        layer_name = name
+        break
+
+if test_layer:
+    print(f"Testing: {layer_name}")
+    
+    # Capture LoRA A and B matrices (these ARE float, not quantized)
+    lora_a_before = list(test_layer.lora_A.values())[0].weight.clone().detach()
+    lora_b_before = list(test_layer.lora_B.values())[0].weight.clone().detach()
+    
+    # Also capture the base weight bytes (NF4 — just check identity)
+    base_before = test_layer.base_layer.weight.data.clone()
+    
+    print(f"  LoRA A dtype: {lora_a_before.dtype}, shape: {lora_a_before.shape}")
+    print(f"  LoRA B dtype: {lora_b_before.dtype}, shape: {lora_b_before.shape}")
+    print(f"  Base weight dtype: {test_layer.base_layer.weight.dtype}")
+    
+    # Run 50 merge/unmerge cycles (simulate 50 GRPO steps)
+    for i in range(50):
+        FastLanguageModel.for_inference(policy_inner)
+        FastLanguageModel.for_training(policy_inner)
+    
+    lora_a_after = list(test_layer.lora_A.values())[0].weight.clone().detach()
+    lora_b_after = list(test_layer.lora_B.values())[0].weight.clone().detach()
+    base_after = test_layer.base_layer.weight.data.clone()
+    
+    # LoRA weight drift
+    a_diff = (lora_a_before - lora_a_after).abs().max().item()
+    b_diff = (lora_b_before - lora_b_after).abs().max().item()
+    a_rel = a_diff / (lora_a_before.abs().mean().item() + 1e-8)
+    b_rel = b_diff / (lora_b_before.abs().mean().item() + 1e-8)
+    
+    print(f"\n  After 50 cycles:")
+    print(f"  LoRA A max diff: {a_diff:.2e} (relative: {a_rel:.2e})")
+    print(f"  LoRA B max diff: {b_diff:.2e} (relative: {b_rel:.2e})")
+    
+    # Base weight byte-level identity
+    base_identical = torch.equal(base_before, base_after)
+    print(f"  Base weights identical (byte-exact): {base_identical}")
+    
+    if a_rel < 1e-5 and b_rel < 1e-5 and base_identical:
+        print("  ✅ PASS: No drift after 50 cycles")
+    elif a_diff == 0 and b_diff == 0 and base_identical:
+        print("  ✅ PASS: Bit-perfect — for_inference() does NOT merge weights")
+    else:
+        print("  ❌ FAIL: Weight drift detected")
+
+# ── CHECK 3: Memory leak over 20 cycles ─────────────────────
+print("\n" + "=" * 60)
+print("CHECK 3: Memory leak test (20 cycles)")
+
+torch.cuda.empty_cache()
+import gc; gc.collect()
+
+baseline_mem = torch.cuda.memory_allocated() / 1e9
+print(f"  Baseline: {baseline_mem:.3f} GB")
+
+for i in range(20):
+    FastLanguageModel.for_inference(policy_inner)
+    # Simulate a short generation
+    test_input = torch.tensor([[1, 2, 3]], device="cuda")
+    with torch.no_grad():
+        _ = policy_inner(test_input)
+    FastLanguageModel.for_training(policy_inner)
+    
+    if (i + 1) % 5 == 0:
+        gc.collect(); torch.cuda.empty_cache()
+        current = torch.cuda.memory_allocated() / 1e9
+        delta = current - baseline_mem
+        print(f"  Cycle {i+1:2d}: {current:.3f} GB (delta: {delta:+.3f} GB)")
+
+final_mem = torch.cuda.memory_allocated() / 1e9
+total_drift = final_mem - baseline_mem
+print(f"\n  Total memory drift: {total_drift:+.3f} GB")
+if abs(total_drift) < 0.1:
+    print("  ✅ PASS: No significant memory leak")
+else:
+    print("  ⚠️  WARN: Memory growing — potential leak")
+
+# ── CHECK 4: Gradient flow after mode switch ─────────────────
+print("\n" + "=" * 60)
+print("CHECK 4: Gradient flow survives mode switching")
+
+FastLanguageModel.for_inference(policy_inner)
+FastLanguageModel.for_training(policy_inner)
+
+# Check LoRA params still require grad
+trainable = 0
+frozen = 0
+for name, p in policy_inner.named_parameters():
+    if 'lora_' in name:
+        if p.requires_grad:
+            trainable += 1
+        else:
+            frozen += 1
+
+print(f"  LoRA params requiring grad: {trainable}")
+print(f"  LoRA params frozen (bad):   {frozen}")
+if frozen == 0 and trainable > 0:
+    print("  ✅ PASS: All LoRA params trainable after mode switch")
+else:
+    print("  ❌ FAIL: Mode switch froze LoRA parameters")
+```
+
+
+```python
+# ── W&B Auth ──────────────────────────────────────────────────────────────────
+# Vertex AI Workbench does not have wandb pre-authenticated.
+# WANDB_API_KEY must be set as an environment variable before running this cell.
+# On Vertex AI: add it to the instance env vars, or set it in a prior cell:
+import os
+import wandb
+
+os.environ["WANDB_API_KEY"] = "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"
+
+# ── W&B Auth ──────────────────────────────────────────────────────────────────
+_wandb_key = os.environ.get("WANDB_API_KEY", "").strip()
+if not _wandb_key:
+    raise EnvironmentError(
+        "WANDB_API_KEY is not set.\n"
+        "Set it before running this cell:\n"
+        "  import os; os.environ[\"WANDB_API_KEY\"] = \"your-key\"\n"
+        "Or add it to your Vertex AI Workbench instance environment variables."
+    )
+wandb.login(key=_wandb_key, relogin=True)
+print(f"✓ W&B authenticated")
+```
+
+
+```python
+import shutil
+import torch
+from transformers import TrainerCallback
+from trl import GRPOConfig, GRPOTrainer
+
+# ── W&B Init ──────────────────────────────────────────────────────────────────
+wandb.init(
+    project=WANDB_PROJECT,
+    name=f"grpo-v2-l4-{time.strftime('%Y%m%d-%H%M')}",
+    config={
+        "model_id":         MODEL_ID,
+        "version":          "v2",
+        "save_steps":       SAVE_STEPS,
+        "eval_steps":       EVAL_STEPS,
+        "eval_max_samples": EVAL_MAX_SAMPLES,
+        "eval_max_tokens":  EVAL_MAX_TOKENS,
+        "patience":         EARLY_STOPPING_PATIENCE,
+        "delta":            EARLY_STOPPING_DELTA,
+        "batch_size":       BATCH_SIZE,
+        "grad_accum":       GRAD_ACCUM,
+        "max_steps":        MAX_STEPS,
+        "learning_rate":    LEARNING_RATE,
+        "num_generations":  NUM_GENERATIONS,
+        "scale_rewards":    SCALE_REWARDS,
+        "eval_split_ratio": EVAL_SPLIT_RATIO,
+        "train_prompts":    len(train_dataset),
+        "eval_prompts":     len(eval_dataset),
+        "use_vllm":         USE_VLLM,
+    },
+)
+print(f"✓ W&B run: {wandb.run.url}")
+
+# ── Resume logic ──────────────────────────────────────────────────────────────
+FRESH = True  # Set True to clear old checkpoints and start over
+
+resume_from = None
+if FRESH and CHECKPOINT_DIR.exists():
+    print("FRESH: deleting old checkpoints...")
+    shutil.rmtree(CHECKPOINT_DIR)
+elif CHECKPOINT_DIR.exists():
+    checkpoints = sorted(
+        [d for d in CHECKPOINT_DIR.iterdir()
+         if d.is_dir() and d.name.startswith("checkpoint-")],
+        key=lambda d: int(d.name.split("-")[-1]),
+    )
+    if checkpoints:
+        resume_from = str(checkpoints[-1])
+        print(f"Resuming from: {resume_from}")
+
+# ── UnslothGRPOTrainer: activate inference kernels during generation ───────────
+class UnslothGRPOTrainer(GRPOTrainer):
+    """
+    Wraps GRPOTrainer._generate() with Unsloth for_inference()/for_training()
+    to activate Unsloth's optimized Triton kernels during the generation phase.
+
+    Root cause of ~3.4 tok/s on L4:
+    - GRPOTrainer calls unwrapped_model.generate() while model is in train() mode.
+    - Unsloth's fused inference Triton kernels are only active after for_inference().
+    - Without this wrapper: ~3-4 tok/s. With it: expected ~8-15 tok/s on L4.
+
+    Override target: _generate() — called from _generate_and_score_completions().
+    Verified against TRL 0.24.0 source (asserted in Cell 3).
+    """
+    def _generate(self, prompts, images):
+        FastLanguageModel.for_inference(self.model)
+        try:
+            result = super()._generate(prompts, images)
+        finally:
+            # Always restore — even if generation crashes
+            FastLanguageModel.for_training(self.model)
+        return result
+
+# ── EvalRewardCallback — capped for L4 feasibility ───────────────────────────
+class EvalRewardCallback(TrainerCallback):
+    """
+    Runs commerce_reward_fn on a capped subset of the held-out eval set.
+
+    v2 changes:
+    - Capped to EVAL_MAX_SAMPLES (10) — full 45 samples × 591s = 7.4h/eval.
+    - max_new_tokens=EVAL_MAX_TOKENS (256) — keeps each eval pass < 15min on L4.
+    - Logs eval/mean_reward to both console and W&B.
+    - Patience-based early stopping via control.should_training_stop.
+    """
+    def __init__(self, eval_records, reward_fn, patience=3, delta=0.01):
+        self.eval_records     = eval_records
+        self.reward_fn        = reward_fn
+        self.patience         = patience
+        self.delta            = delta
+        self.best_reward      = -float("inf")
+        self.no_improve_count = 0
+
+    def on_step_end(self, args, state, control, model=None, processing_class=None, **kwargs):
+        # ^^^ Changed: tokenizer -> processing_class
+        
+        if state.global_step == 0 or state.global_step % EVAL_STEPS != 0:
+            return control
+
+        # processing_class is the tokenizer in TRL 0.24.0+
+        tokenizer = processing_class
+        if tokenizer is None:
+            print("[EvalRewardCallback] WARNING: tokenizer is None, skipping eval")
+            return control
+
+        mean_reward = self._run_eval(model, tokenizer, args)
+        improved    = mean_reward > self.best_reward + self.delta
+        status      = (
+            "↑ improved" if improved
+            else f"↔ no gain ({self.no_improve_count + 1}/{self.patience})"
+        )
+
+        # Log to W&B
+        wandb.log(
+            {
+                "eval/mean_reward":      mean_reward,
+                "eval/best_reward":      max(self.best_reward, mean_reward),
+                "eval/no_improve_count": self.no_improve_count,
+            },
+            step=state.global_step,
+        )
+
+        print(
+            f"\n[EvalReward] step={state.global_step} | "
+            f"mean_eval_reward={mean_reward:.4f} | best={self.best_reward:.4f} | {status}"
+        )
+
+        if improved:
+            self.best_reward      = mean_reward
+            self.no_improve_count = 0
+        else:
+            self.no_improve_count += 1
+            if self.no_improve_count >= self.patience:
+                print(
+                    f"[EarlyStopping] No improvement ≥ {self.delta} for "
+                    f"{self.patience} consecutive evals. Halting training."
+                )
+                wandb.log({"early_stop/step": state.global_step}, step=state.global_step)
+                control.should_training_stop = True
+
+        return control
+
+    def _run_eval(self, model, tokenizer, args):
+        """One greedy completion per eval prompt, scored by reward_fn."""
+        FastLanguageModel.for_inference(model)
+        rewards = []
+        
+        # Cap to subset
+        subset = self.eval_records[:EVAL_MAX_SAMPLES]
+        
+        for record in subset:
+            msgs = record["prompt"]
+            text = tokenizer.apply_chat_template(
+                msgs, tokenize=False, add_generation_prompt=True
+            )
+            inputs = tokenizer(
+                text, return_tensors="pt", truncation=True,
+                max_length=args.max_prompt_length,
+            ).to(model.device)
+            with torch.no_grad():
+                out = model.generate(
+                    **inputs,
+                    max_new_tokens=EVAL_MAX_TOKENS,  # use the cap
+                    temperature=0.7,
+                    do_sample=True,
+                )
+            resp = tokenizer.decode(
+                out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
+            )
+            rewards.append(self.reward_fn([resp], [text])[0])
+        
+        FastLanguageModel.for_training(model)
+        return sum(rewards) / len(rewards) if rewards else 0.0
+
+# ── Training Config ────────────────────────────────────────────────────────────
+FastLanguageModel.for_training(model)
+
+grpo_config = GRPOConfig(
+    output_dir=str(CHECKPOINT_DIR),
+    num_generations=NUM_GENERATIONS,
+    scale_rewards=SCALE_REWARDS,
+    max_completion_length=MAX_COMPLETION_LENGTH,
+    max_steps=MAX_STEPS,
+    temperature=TEMPERATURE,
+    num_train_epochs=NUM_EPOCHS,
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    learning_rate=LEARNING_RATE,
+    warmup_ratio=0.1,
+    lr_scheduler_type="cosine",
+    fp16=False,
+    bf16=True,
+    logging_steps=1,                   # every step in console + W&B
+    save_steps=SAVE_STEPS,             # 5 — ~3.3h exposure on L4 Spot VM
+    save_total_limit=SAVE_TOTAL_LIMIT, # prune old checkpoints, keep 3
+    save_only_model=True,              # eliminate save overhead
+    eval_steps=EVAL_STEPS,             # drives EvalRewardCallback cadence
+    report_to="wandb",
+    max_prompt_length=MAX_SEQ_LENGTH // 2,
+    seed=42,
+    remove_unused_columns=False,
+    # vLLM colocate — only used when USE_VLLM=True (Trainer class also switches below)
+    **({"use_vllm": True, "vllm_mode": "colocate",
+        "vllm_enable_sleep_mode": True} if USE_VLLM else {}),
+)
+
+eval_cb = EvalRewardCallback(
+    eval_records=list(eval_dataset),
+    reward_fn=commerce_reward_fn,
+    patience=EARLY_STOPPING_PATIENCE,
+    delta=EARLY_STOPPING_DELTA,
+)
+
+# ── Trainer: UnslothGRPOTrainer (vLLM handles its own generation path) ────────
+TrainerClass = GRPOTrainer if USE_VLLM else UnslothGRPOTrainer
+trainer = TrainerClass(
+    model=model,
+    reward_funcs=commerce_reward_fn,
+    args=grpo_config,
+    train_dataset=train_dataset,
+    processing_class=tokenizer,     # v2 fix: was tokenizer=tokenizer (silently dropped)
+    callbacks=[eval_cb],
+)
+print(
+    f"Trainer: {TrainerClass.__name__} | "
+    f"max_steps={MAX_STEPS} | save_every={SAVE_STEPS} | eval_every={EVAL_STEPS} | "
+    f"eval_cap={EVAL_MAX_SAMPLES}×{EVAL_MAX_TOKENS}tok | resume={resume_from is not None}"
+)
+
+t_start = time.time()
+result  = trainer.train(resume_from_checkpoint=resume_from)
+elapsed = time.time() - t_start
+
+wandb.log({
+    "train/final_loss":       result.training_loss,
+    "train/duration_hours":   elapsed / 3600,
+    "train/total_steps":      result.global_step,
+    "eval/best_reward_final": eval_cb.best_reward,
+})
+wandb.finish()
+
+print(f"\n{'='*60}")
+print(f"GRPO v2 Training Complete")
+print(f"  Loss:        {result.training_loss:.4f}")
+print(f"  Steps:       {result.global_step}")
+print(f"  Duration:    {elapsed/3600:.1f}h")
+print(f"  Best eval R: {eval_cb.best_reward:.4f}")
+print(f"  Trainer:     {TrainerClass.__name__}")
+print(f"{'='*60}")
+
+```
+
+    wandb: WARNING If you're specifying your api key in code, ensure this code is not shared publicly.
+    wandb: WARNING Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.
+    wandb: WARNING [wandb.login()] Changing session credentials to explicit value for https://api.wandb.ai.
+    wandb: Appending key for api.wandb.ai to your netrc file: /home/jupyter/.netrc
+
+
+    ✓ W&B authenticated
+
+Tracking run with wandb version 0.26.0
+
+Run data is saved locally in <code>/home/jupyter/tucano2/notebooks/wandb/run-20260422_212656-2m114rh7</code>
+
+
+
+Syncing run <strong><a href='https://wandb.ai/tferrazrafael-self/tucano2-commerce/runs/2m114rh7' target="_blank">grpo-v2-l4-20260422-2126</a></strong> to <a href='https://wandb.ai/tferrazrafael-self/tucano2-commerce' target="_blank">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target="_blank">docs</a>)<br>
+
+
+
+View project at <a href='https://wandb.ai/tferrazrafael-self/tucano2-commerce' target="_blank">https://wandb.ai/tferrazrafael-self/tucano2-commerce</a>
+
+
+
+View run at <a href='https://wandb.ai/tferrazrafael-self/tucano2-commerce/runs/2m114rh7' target="_blank">https://wandb.ai/tferrazrafael-self/tucano2-commerce/runs/2m114rh7</a>
+
+
+    ✓ W&B run: https://wandb.ai/tferrazrafael-self/tucano2-commerce/runs/2m114rh7
+    FRESH: deleting old checkpoints...
+    Trainer: UnslothGRPOTrainer | max_steps=300 | save_every=5 | eval_every=10 | eval_cap=5×2048tok | resume=False
+
+
+    ==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1
+       \\   /|    Num examples = 300 | Num Epochs = 1 | Total steps = 300
+    O^O/ \_/ \    Batch size per device = 4 | Gradient accumulation steps = 2
+    \        /    Data Parallel GPUs = 1 | Total batch size (4 x 2 x 1) = 8
+     "-____-"     Trainable parameters = 33,030,144 of 3,792,371,200 (0.87% trained)
+
+
+    Unsloth: Will smartly offload gradients to save VRAM!
+    
+    [EvalReward] step=10 | mean_eval_reward=0.0830 | best=-inf | ↑ improved
+    
+    [EvalReward] step=20 | mean_eval_reward=0.0830 | best=0.0830 | ↔ no gain (1/10)
+    
+    [EvalReward] step=30 | mean_eval_reward=0.1050 | best=0.0830 | ↑ improved
+    
+    [EvalReward] step=40 | mean_eval_reward=0.1010 | best=0.1050 | ↔ no gain (1/10)
+
+
+    wandb: WARNING Tried to log to step 40 that is less than the current step 239. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+
+
+## Cell 12: Save Adapter
+
+
+```python
+GRPO_ADAPTER_DIR.mkdir(parents=True, exist_ok=True)
+model.save_pretrained(str(GRPO_ADAPTER_DIR))
+tokenizer.save_pretrained(str(GRPO_ADAPTER_DIR))
+
+# Save training summary
+summary = {
+    "model_id": MODEL_ID,
+    "sft_adapter": str(SFT_ADAPTER_DIR),
+    "method": "GRPO",
+    "train_loss": result.training_loss,
+    "num_prompts": len(dataset),
+    "num_generations": NUM_GENERATIONS,
+    "scale_rewards": SCALE_REWARDS,
+    "learning_rate": LEARNING_RATE,
+    "epochs": NUM_EPOCHS,
+    "max_seq_length": MAX_SEQ_LENGTH,
+    "max_completion_length": MAX_COMPLETION_LENGTH,
+    "duration_seconds": round(elapsed),
+    "gpu": "L4",
+    "platform": "vertex-ai-workbench",
+}
+with open(GRPO_ADAPTER_DIR / "training_summary.json", "w") as f:
+    json.dump(summary, f, indent=2)
+
+print(f"✓ Adapter saved to {GRPO_ADAPTER_DIR}")
+print(f"  Files: {[f.name for f in GRPO_ADAPTER_DIR.iterdir()]}")
+```
+
+    ✓ Adapter saved to /home/jupyter/tucano2/models/tucano2-commerce-grpo
+      Files: ['tokenizer_config.json', 'tokenizer.json', 'README.md', 'adapter_model.safetensors', 'chat_template.jinja', 'training_summary.json', 'checkpoints', 'special_tokens_map.json', 'adapter_config.json']
+
+
+## Cell 13: Validation
+
+Generate 5 samples with trained model, score with reward functions.
+
+
+```python
+FastLanguageModel.for_inference(model)
+
+system_msg = {"role": "system", "content": SYSTEM_PT}
+
+test_prompts = [
+    {"role": "user", "content": (
+        "Analise esta avaliação de e-commerce brasileiro e extraia dados estruturados.\n\n"
+        "nota=2/5 | status=delivered\ntítulo: decepcionado\n"
+        "texto: Produto veio com defeito e o vendedor não respondeu.\n\n"
+        "Retorne um objeto JSON com exatamente estas chaves:\n"
+        "sentiment, sentiment_score, churn_risk, delivery_issue, product_issue, "
+        "seller_issue, main_complaint, complaint_category, repeat_intent, would_recommend"
+    )},
+    {"role": "user", "content": "Quais são as categorias de reclamação mais frequentes e como afetam a nota média?"},
+    {"role": "user", "content": "Analise a retenção de clientes afetados por product_quality."},
+    {"role": "user", "content": (
+        "Perfil do cliente:\n- Estado: MG\n- Valor do pedido: R$150\n"
+        "- Reclamação: produto com defeito\n- Nota: 1.0/5\n\n"
+        "Este cliente deve receber uma notificação de reengajamento?"
+    )},
+    {"role": "user", "content": "Compare a satisfação de clientes em SP vs RJ."},
+]
+
+print("=== GRPO Validation ===")
+print()
+
+for i, prompt in enumerate(test_prompts):
+    messages = [system_msg, prompt]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+    outputs = model.generate(**inputs, max_new_tokens=2048, temperature=0.1, do_sample=True)
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+
+    reward = commerce_reward_fn([response], [text])[0]
+    answer = strip_think(response)
+
+    print(f"--- Sample {i+1} (reward={reward:.2f}, tokens={len(response.split())}) ---")
+    print(f"Prompt: {prompt['content'][:80]}...")
+    print(f"Answer: {answer[:300]}")
+    print()
+```
+
+    === GRPO Validation ===
+    
+    --- Sample 1 (reward=0.12, tokens=442) ---
+    Prompt: Analise esta avaliação de e-commerce brasileiro e extraia dados estruturados.
+    
+    n...
+    Answer: ```json
+    {
+      "sentiment": "negativo",
+      "sentiment_score": -0.8,
+      "churn_risk": -1,
+      "delivery_issue": true,
+      "product_issue": false,
+      "seller_issue": false,
+      "main_complaint": "falha no atendimento post-venda",
+      "complaint_category": "serviço pós-venda",
+      "repeat_intent": 0,
+      "would_recommen
+    
+    --- Sample 2 (reward=0.70, tokens=1451) ---
+    Prompt: Quais são as categorias de reclamação mais frequentes e como afetam a nota média...
+    Answer: As categorias de reclamação variam bastante conforme o tipo de serviço ou produto sendo consumido, mas existem alguns tipos comuns que tendem a aparecer frequentemente. Aqui estão algumas das principais categorias de reclamações e suas possíveis influências na nota média:
+    
+    ### Categorias Frequentes 
+    
+    --- Sample 3 (reward=0.70, tokens=1465) ---
+    Prompt: Analise a retenção de clientes afetados por product_quality....
+    Answer: Claro! Vamos analisar a possível influência da *quality* do produto (*product_quality*) sobre a taxa de retenção de clientes. Para isso, vamos seguir alguns passos lógicos:
+    
+    ### Passo 1: Definição das Variáveis
+    - **Quality do Produto (`product_quality`)**: Uma métrica quantitativa ou qualitativa que
+    
+    --- Sample 4 (reward=0.50, tokens=994) ---
+    Prompt: Perfil do cliente:
+    - Estado: MG
+    - Valor do pedido: R$150
+    - Reclamação: produto c...
+    Answer: <think>
+    O usuário está me perguntando se este cliente deveria receber uma notificação de reengajamento baseado nas informações fornecidas. Primeiro, vou analisar cada ponto individualmente para entender a situação completa. 
+    
+    1. **Estado:** O estado mencionado pelo cliente é Minas Gerais (MG). Isso 
+    
+    --- Sample 5 (reward=0.70, tokens=1396) ---
+    Prompt: Compare a satisfação de clientes em SP vs RJ....
+    Answer: Para compararmos a satisfação dos clientes em São Paulo (SP) versus Rio de Janeiro (RJ), precisamos considerar alguns pontos:
+    
+    1. **Economia Local**: A renda média per capita varia significativamente entre estas duas regiões metropolitanas. No geral, São Paulo tende a ter uma renda maior comparada a
+    
+
+
+
+```python
+
+```