Add full plan doc

genesis-2-rlhf-plan.md

@@ -0,0 +1,190 @@
+# Genesis-2.0 — Multi-Stage RLHF Implementation Plan
+
+**Goal:** Improve Genesis-1.0 (Qwen3.6-35B-A3B, tool-use SFT) with multi-stage RL post-training to add tool-choice discrimination, OOD generalization, and multi-step planning.
+
+**Total estimated cost: $6-9** (RTX PRO 6000 Blackwell, ~3.5-5.5 hrs)
+**Total wall time: ~1-2 weeks** (mostly Phase 0 data work on MacBook)
+
+---
+
+## Architecture Overview
+
+### Model
+| Attribute | Value |
+|-----------|-------|
+| Base | Qwen/Qwen3.6-35B-A3B (MoE, 35B total, 3B active) |
+| Current | Genesis-1.0 (SFT, attention-only LoRA r=64, MLX 4-bit) |
+| Target | Genesis-2.0 (SFT + DPO + GRPO, all-linear LoRA r=32, MLX 4-bit) |
+
+### Compute
+| Phase | Hardware | Cost | Time |
+|-------|----------|------|------|
+| 0 (Data prep) | Mac M5 Pro 64GB | $0 | ~1 week |
+| 1 (DPO) | 1× RTX PRO 6000 96GB | ~$3-4 | 1.5-2.5 hrs |
+| 2 (GRPO) | 1× RTX PRO 6000 96GB | ~$3-5 | 2-3 hrs |
+| 3 (Merge+Deploy) | Mac M5 Pro 64GB | $0 | ~1 day |
+
+### Software Stack
+- **Training:** TRL (DPOTrainer, GRPOTrainer) + Unsloth
+- **Inference engine:** vLLM (for GRPO rollouts)
+- **Quantization:** QLoRA NF4 (bitsandbytes 4-bit)
+- **LoRA Config C:** attention (q/k/v/o) + shared expert FFN (gate/up/down) + gate/router — 131M params, r=32
+- **Monitoring:** Weights & Biases
+- **RunPod template:** `axolotl-cloud:main-latest` or `pytorch:2.5-cuda12.4`
+
+---
+
+## Phase 0 — Data & Infrastructure (Local M5 Pro, $0)
+
+### Objectives
+1. Build DPO preference pairs from existing data
+2. Implement rule-based reward functions for tool-use
+3. Create evaluation benchmark
+4. Document RunPod provisioning procedure
+
+### DPO Data Sources
+
+**Source A: Session logs (real tool-use traces)**
+- Extract prompts from existing SFT dataset (20K examples)
+- For each prompt, generate 2 responses with Genesis-1.0 at temperature 0.7
+- Score with rule-based evaluator:
+  - +1.0: correct tool call + correct params + task completed
+  - +0.5: correct tool call + error recovered
+  - 0.0: wrong tool call
+  - -0.5: hallucinated tool call
+  - -1.0: no meaningful output
+- Pair: highest-scoring = chosen, lowest = rejected
+- Target: 500-2,000 pairs from this method
+
+**Source B: Data augmentation (controlled perturbations)**
+- Take successful SFT examples, introduce controlled errors:
+  - Wrong tool name for the task
+  - Correct tool but missing required parameter
+  - Correct call but result ignored
+  - Unnecessary tool calls (inefficient trajectory)
+- Correct version = chosen, perturbed = rejected
+- Target: 500-1,000 pairs
+
+### Reward Functions
+```python
+def tool_call_correctness(trajectory, tools_schema):
+    """+1 if all tool calls match valid names and param types. 0 otherwise."""
+    
+def task_completion(trajectory, expected_outcome):
+    """+1 if final answer satisfies the task. 0 if incomplete. -1 if wrong."""
+    
+def tool_result_usage(trajectory):
+    """+0.5 if tool results are referenced in subsequent reasoning. 0 if ignored."""
+    
+def trajectory_efficiency(trajectory):
+    """Normalized score: ideal_steps / actual_steps. Range [0, 1]."""
+    
+def format_validity(trajectory):
+    """+1 if tool calls follow the expected JSON/XML format. 0 otherwise."""
+```
+
+Combined reward: `0.4×completion + 0.3×correctness + 0.15×usage + 0.1×efficiency + 0.05×format`
+
+### Evaluation Benchmark
+- Hold-out set: 100-200 prompts from the same distribution as training
+- Metrics:
+  - Tool call accuracy (% of calls with correct tool + params)
+  - Task completion rate (% of tasks finished successfully)
+  - Trajectory efficiency (avg steps to completion)
+  - Format compliance (% of valid tool call syntax)
+- Run before/after each phase to measure improvement
+- Required improvement threshold for phase progression: ≥5% tool call accuracy
+
+---
+
+## Phase 1 — DPO Preference Alignment (RTX PRO 6000, ~$3-4)
+
+### Procedure
+1. Provision RTX PRO 6000 pod on RunPod (Community Cloud ~$1.69/hr)
+2. Install dependencies: torch, transformers, trl, unsloth, peft, bitsandbytes, datasets
+3. Upload Genesis-1.0 SFT adapter (53 MB from HF: `jacobeen06/Genesis-1.0-SFT-adapter`)
+4. Load base model (Qwen3.6-35B-A3B) with QLoRA NF4
+5. Merge existing attention-only adapter into base model
+6. Attach new all-linear LoRA (Config C: attention + shared expert + gate, r=32)
+7. Train with DPOTrainer:
+   - Batch: micro_batch=1, grad_accum=4 (effective 4)
+   - Learning rate: 5e-6 (cosine schedule)
+   - β (DPO temperature): 0.1-0.3
+   - Context: 4096
+   - Epochs: 1-2
+   - Precision: BF16
+8. Evaluate on hold-out set
+9. Save adapter to HF Hub: `jacobeen06/Genesis-2.0-DPO-adapter`
+
+### Expected outcome
+- Tool call accuracy: +5-15% improvement over Genesis-1.0
+- Reduced hallucinated tool calls
+- Better tool-choice discrimination
+
+---
+
+## Phase 2 — GRPO with Verifiable Rewards (RTX PRO 6000, ~$3-5)
+
+### Procedure
+1. Load Phase 1 checkpoint (base + DPO LoRA)
+2. Configure GRPOTrainer with vLLM:
+   - G=4-8 generations per prompt
+   - Reward: combined rule-based (same as Phase 0)
+   - GRPO config: clip_high=0.28, clip_low=0.2 (DAPO-style asymmetric)
+   - KL penalty: β=0.04 start, decay to 0.01
+   - vLLM gpu_memory_utilization: 0.90
+   - Context: 4096-8192
+3. Monitor for failure modes:
+   - Entropy collapse → enable clip-higher
+   - Advantage collapse → dynamic sampling
+   - KL drift → PID-controlled β
+4. Evaluate on hold-out set
+5. Save adapter: `jacobeen06/Genesis-2.0-GRPO-adapter`
+
+### Expected outcome
+- OOD tool reasoning: +10-20% over DPO-only
+- Multi-step plan quality improvement
+- Better recovery from tool errors
+
+---
+
+## Phase 3 — Merge, Quantize, Deploy (Local M5 Pro, $0)
+
+### Procedure
+1. Merge all-linear LoRA into base model
+2. Convert to MLX format (4-bit with MTP heads preserved)
+3. Evaluate on benchmark (compare pre/post RL)
+4. Upload as `jacobeen06/Genesis-2.0-35b` to HF Hub
+5. Update oMLX/cua-driver config for Genesis-2.0
+6. Run inference test with `mlx_lm.generate --model jacobeen06/Genesis-2.0-35b --mtp`
+
+---
+
+## Infrastructure Reference
+
+### RunPod: RTX PRO 6000 Blackwell 96GB
+- **Community Cloud:** ~$1.69/hr
+- **Secure Cloud:** ~$1.99/hr
+- **Recommended template:** `axolotl-cloud:main-latest` (has torch, CUDA, etc.)
+- **Disk:** 50 GB minimum (model ~30 GB + adapter ~0.25 GB + data ~2 GB)
+- **Region:** US or EU (whichever has availability)
+
+### Pod Setup Checklist
+- [ ] Create pod with RTX PRO 6000 96GB
+- [ ] Wait for pod to reach "Running" state
+- [ ] Open Web Terminal
+- [ ] Verify GPU: `nvidia-smi`
+- [ ] Install extra deps: `pip install trl unsloth bitsandbytes`
+- [ ] Download base model: `huggingface-cli download Qwen/Qwen3.6-35B-A3B`
+- [ ] Download adapter: `huggingface-cli download jacobeen06/Genesis-1.0-SFT-adapter`
+- [ ] Run training script
+- [ ] Upload result adapter to HF
+
+### Key Risks & Mitigations
+| Risk | Mitigation |
+|------|-----------|
+| Flash-attn compilation OOM | Use `attn_implementation: sdpa` instead |
+| vLLM OOM during GRPO | Reduce G=4, use Unsloth Standby, stay at 4K context |
+| SSH not available on pod | Use RunPod Web Terminal (single-command approach) |
+| Training diverges | Monitor loss curves in W&B, kill and adjust LR |
+| Data transfer slow | Use `hf upload` from Mac → HF Hub → `hf download` on pod |