Skip to content

Instantly share code, notes, and snippets.

@wojtyniak

wojtyniak/notebook-22160cec-f456-4a69-b0ff-04812e5c2da3-paper-20260127-011232.ipynb

Created January 27, 2026 01:12
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save wojtyniak/63fa78de2109b09adf6807552b5f316e to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save wojtyniak/63fa78de2109b09adf6807552b5f316e to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
notebook-22160cec-f456-4a69-b0ff-04812e5c2da3-paper-20260127-011232.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# CWM: Code Generation with World Models - Educational Overview\n",
"\n",
"**Paper**: CWM: An Open-Weights LLM for Research on Code Generation with World Models \n",
"**Authors**: Meta FAIR CodeGen Team\n",
"\n",
"## Overview\n",
"\n",
"This notebook provides an educational walkthrough of the key computational methods in the CWM paper:\n",
"\n",
"1. **Python Execution Trace Generation** - Capturing program state line-by-line\n",
"2. **Agentic Data Generation (ForagerAgent)** - Simulating agent-environment interactions\n",
"3. **GRPO Reinforcement Learning** - Policy optimization for code generation\n",
"4. **Execution Trace Prediction** - Using traces for code understanding\n",
"5. **CruxEval Demonstration** - Predicting code execution outcomes\n",
"\n",
"**Note**: This notebook uses **small-scale demonstrations** due to resource constraints (4GB RAM, no GPU). It illustrates the methodology without full-scale training or execution."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup and Dependencies"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Install all dependencies in a single command\n",
"!uv pip install numpy pandas matplotlib seaborn scipy scikit-learn torch --no-cache-dir"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Import all required libraries\n",
"import numpy as np\n",
"import pandas as pd\n",
"import matplotlib.pyplot as plt\n",
"import seaborn as sns\n",
"from typing import Dict, List, Any, Tuple, Optional\n",
"import json\n",
"import inspect\n",
"import sys\n",
"import io\n",
"from contextlib import redirect_stdout, redirect_stderr\n",
"import traceback as tb\n",
"from dataclasses import dataclass, field\n",
"from collections import defaultdict\n",
"import warnings\n",
"import random\n",
"\n",
"# Set random seeds for reproducibility\n",
"np.random.seed(42)\n",
"random.seed(42)\n",
"\n",
"# Configure plotting\n",
"plt.style.use('default')\n",
"sns.set_palette(\"husl\")\n",
"warnings.filterwarnings('ignore')\n",
"\n",
"print(\"✓ All imports successful\")\n",
"print(f\"NumPy version: {np.__version__}\")\n",
"print(f\"Pandas version: {pd.__version__}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Python Execution Trace Generation\n",
"\n",
"The paper introduces a novel approach to teaching LLMs about code execution by training on **observation-action trajectories**:\n",
"- **Observation**: State of local variables before executing a line\n",
"- **Action**: The Python statement being executed\n",
"- **Next Observation**: Resulting state after execution\n",
"\n",
"This enables the model to learn how code changes program state, not just syntax."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class PythonTracer:\n",
" \"\"\"\n",
" Simplified Python execution tracer based on CWM's approach.\n",
" Captures local variable states after each line execution.\n",
" \n",
" Paper Reference: Section 2.2 - Python tracing: neural code interpretation data\n",
" \"\"\"\n",
" \n",
" def __init__(self):\n",
" self.trace = []\n",
" self.prev_locals = {}\n",
" \n",
" def trace_function(self, func, *args, **kwargs):\n",
" \"\"\"\n",
" Execute a function and capture execution trace.\n",
" Returns list of (line_number, action, local_vars) tuples.\n",
" \"\"\"\n",
" self.trace = []\n",
" self.prev_locals = {}\n",
" \n",
" # Get source code\n",
" source_lines = inspect.getsource(func).split('\\n')\n",
" \n",
" # Record function call\n",
" arg_names = inspect.signature(func).parameters.keys()\n",
" call_locals = dict(zip(arg_names, args))\n",
" \n",
" self.trace.append({\n",
" 'event': 'call',\n",
" 'action': f\"def {func.__name__}(...)\",\n",
" 'locals': self._format_locals(call_locals)\n",
" })\n",
" \n",
" # Execute with tracing\n",
" def tracer(frame, event, arg):\n",
" if event == 'line':\n",
" # Get current line\n",
" lineno = frame.f_lineno\n",
" local_vars = frame.f_locals.copy()\n",
" \n",
" # Get the source line\n",
" try:\n",
" line = source_lines[lineno - frame.f_code.co_firstlineno].strip()\n",
" except:\n",
" line = \"<source unavailable>\"\n",
" \n",
" self.trace.append({\n",
" 'event': 'line',\n",
" 'action': line,\n",
" 'locals': self._format_locals(local_vars)\n",
" })\n",
" \n",
" elif event == 'return':\n",
" self.trace.append({\n",
" 'event': 'return',\n",
" 'action': 'return',\n",
" 'return_value': repr(arg)\n",
" })\n",
" \n",
" return tracer\n",
" \n",
" # Execute function with tracer\n",
" sys.settrace(tracer)\n",
" try:\n",
" result = func(*args, **kwargs)\n",
" finally:\n",
" sys.settrace(None)\n",
" \n",
" return result, self.trace\n",
" \n",
" def _format_locals(self, local_vars):\n",
" \"\"\"Format local variables similar to CWM trace format.\"\"\"\n",
" formatted = {}\n",
" for key, val in local_vars.items():\n",
" if not key.startswith('__'):\n",
" # Use ellipsis for unchanged values (as in paper)\n",
" if key in self.prev_locals and self.prev_locals[key] == val:\n",
" formatted[key] = \"..\"\n",
" else:\n",
" formatted[key] = repr(val)\n",
" self.prev_locals[key] = val\n",
" return formatted\n",
" \n",
" def format_trace_cwm_style(self):\n",
" \"\"\"\n",
" Format trace in CWM's observation-action format.\n",
" Paper format: <|frame_sep|> <|line_sep|> {locals} <|action_sep|> action\n",
" \"\"\"\n",
" output = []\n",
" for step in self.trace:\n",
" event = step['event']\n",
" if event == 'call':\n",
" output.append(f\"<|frame_sep|> <|call_sep|> {json.dumps(step['locals'])} <|action_sep|> {step['action']}\")\n",
" elif event == 'line':\n",
" output.append(f\"<|frame_sep|> <|line_sep|> {json.dumps(step['locals'])} <|action_sep|> {step['action']}\")\n",
" elif event == 'return':\n",
" output.append(f\"<|frame_sep|> <|return_sep|> <|action_sep|> return <|arg_sep|> {step['return_value']}\")\n",
" return '\\n'.join(output)\n",
"\n",
"print(\"✓ Python Tracer implemented\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Example 1: Trace a simple function (from the paper's Figure 3)\n",
"def count(s, t):\n",
" \"\"\"Count occurrences of character t in string s.\"\"\"\n",
" n = 0\n",
" for c in s:\n",
" n += int(c == t)\n",
" return n\n",
"\n",
"tracer = PythonTracer()\n",
"result, trace = tracer.trace_function(count, \"strawberry\", \"r\")\n",
"\n",
"print(\"Function result:\", result)\n",
"print(\"\\nExecution trace (first 10 steps):\")\n",
"print(\"=\"*60)\n",
"for i, step in enumerate(trace[:10]):\n",
" print(f\"Step {i}: {step['event']:6s} | {step.get('action', 'N/A'):30s} | Locals: {step.get('locals', {})}\")\n",
"\n",
"print(\"\\n\" + \"=\"*60)\n",
"print(\"CWM-style formatted trace (first 5 lines):\")\n",
"print(\"=\"*60)\n",
"formatted = tracer.format_trace_cwm_style()\n",
"for line in formatted.split('\\n')[:5]:\n",
" print(line)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Example 2: More complex function to demonstrate trace richness\n",
"def fibonacci(n):\n",
" \"\"\"Compute nth Fibonacci number.\"\"\"\n",
" if n <= 1:\n",
" return n\n",
" a, b = 0, 1\n",
" for i in range(2, n + 1):\n",
" a, b = b, a + b\n",
" return b\n",
"\n",
"tracer2 = PythonTracer()\n",
"result2, trace2 = tracer2.trace_function(fibonacci, 6)\n",
"\n",
"print(f\"Fibonacci(6) = {result2}\")\n",
"print(f\"\\nTrace length: {len(trace2)} steps\")\n",
"print(\"\\nKey steps in execution:\")\n",
"print(\"=\"*60)\n",
"for step in trace2[::2]: # Show every other step\n",
" if step['event'] == 'line':\n",
" print(f\"{step['action']:30s} → {step['locals']}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Trace Statistics from the Paper\n",
"\n",
"The paper collected:\n",
"- **120M+ traced Python functions** (function-level tracing)\n",
"- **70k+ repository-level traces** from 21k repositories\n",
"- **33k CodeContests solution traces**\n",
"- **75M natural language trace descriptions**\n",
"\n",
"This data teaches the model how code execution changes program state."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Simulate trace dataset statistics\n",
"trace_stats = {\n",
" 'Dataset': ['Function-level', 'CodeContests', 'Repository-level', 'Natural Language'],\n",
" 'Count': [120_000_000, 33_000, 70_000, 75_000_000],\n",
" 'Type': ['Execution Traces', 'Execution Traces', 'Execution Traces', 'NL Descriptions']\n",
"}\n",
"\n",
"df_traces = pd.DataFrame(trace_stats)\n",
"\n",
"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"# Bar chart\n",
"colors = ['#3498db', '#e74c3c', '#2ecc71', '#f39c12']\n",
"ax1.barh(df_traces['Dataset'], df_traces['Count'], color=colors)\n",
"ax1.set_xlabel('Number of Traces', fontsize=12)\n",
"ax1.set_title('CWM Python Execution Trace Dataset Sizes', fontsize=14, fontweight='bold')\n",
"ax1.set_xscale('log')\n",
"ax1.grid(axis='x', alpha=0.3)\n",
"\n",
"# Pie chart\n",
"execution_total = df_traces[df_traces['Type'] == 'Execution Traces']['Count'].sum()\n",
"nl_total = df_traces[df_traces['Type'] == 'NL Descriptions']['Count'].sum()\n",
"ax2.pie([execution_total, nl_total], labels=['Execution Traces', 'NL Descriptions'], \n",
" autopct='%1.1f%%', colors=['#3498db', '#f39c12'], startangle=90)\n",
"ax2.set_title('Trace Data Composition', fontsize=14, fontweight='bold')\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"Dataset Summary:\")\n",
"print(df_traces.to_string(index=False))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. ForagerAgent: Agentic Data Generation\n",
"\n",
"The paper generates **3M agentic trajectories** where an LLM agent interacts with code repositories to fix bugs:\n",
"\n",
"- **Actions**: create file, edit file, bash command, view file\n",
"- **Tasks**: mutate-fix (55%) and issue-fix (45%)\n",
"- **Scale**: 10.2k repository images, 3.15k repositories\n",
"\n",
"Paper Reference: Section 2.3 - ForagerAgent: agentic midtraining data generation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@dataclass\n",
"class AgenticAction:\n",
" \"\"\"Represents an action taken by the agent.\"\"\"\n",
" tool: str # 'create', 'edit', 'bash', 'view'\n",
" args: Dict[str, Any]\n",
" \n",
"@dataclass\n",
"class AgenticObservation:\n",
" \"\"\"Environment response to agent action.\"\"\"\n",
" output: str\n",
" success: bool\n",
" \n",
"@dataclass\n",
"class AgenticTrajectory:\n",
" \"\"\"Complete agent-environment interaction trajectory.\"\"\"\n",
" task_description: str\n",
" steps: List[Tuple[AgenticAction, AgenticObservation]]\n",
" success: bool\n",
" \n",
"class SimplifiedForagerEnvironment:\n",
" \"\"\"\n",
" Simplified simulation of ForagerAgent environment.\n",
" In the real paper, this runs in Docker with actual code execution.\n",
" \"\"\"\n",
" \n",
" def __init__(self):\n",
" self.files = {} # Simulated file system\n",
" self.test_outputs = [] # Simulated test results\n",
" \n",
" def execute_action(self, action: AgenticAction) -> AgenticObservation:\n",
" \"\"\"Execute an agent action and return observation.\"\"\"\n",
" \n",
" if action.tool == 'create':\n",
" path = action.args['path']\n",
" content = action.args['content']\n",
" self.files[path] = content\n",
" return AgenticObservation(\n",
" output=f\"File created: {path}\",\n",
" success=True\n",
" )\n",
" \n",
" elif action.tool == 'edit':\n",
" path = action.args['path']\n",
" if path not in self.files:\n",
" return AgenticObservation(\n",
" output=f\"Error: File {path} not found\",\n",
" success=False\n",
" )\n",
" # Simplified edit - just replace content\n",
" self.files[path] = action.args['new_content']\n",
" return AgenticObservation(\n",
" output=f\"File edited: {path}\",\n",
" success=True\n",
" )\n",
" \n",
" elif action.tool == 'bash':\n",
" command = action.args['command']\n",
" # Simulate running tests\n",
" if 'pytest' in command or 'test' in command:\n",
" # Simulate test results based on file contents\n",
" passed = any('correct' in content.lower() for content in self.files.values())\n",
" if passed:\n",
" return AgenticObservation(\n",
" output=\"All tests passed ✓\",\n",
" success=True\n",
" )\n",
" else:\n",
" return AgenticObservation(\n",
" output=\"Tests failed: AssertionError in test_function\",\n",
" success=False\n",
" )\n",
" else:\n",
" return AgenticObservation(\n",
" output=f\"Command executed: {command}\",\n",
" success=True\n",
" )\n",
" \n",
" elif action.tool == 'view':\n",
" path = action.args['path']\n",
" if path in self.files:\n",
" return AgenticObservation(\n",
" output=self.files[path],\n",
" success=True\n",
" )\n",
" else:\n",
" return AgenticObservation(\n",
" output=f\"File not found: {path}\",\n",
" success=False\n",
" )\n",
" \n",
" return AgenticObservation(output=\"Unknown tool\", success=False)\n",
"\n",
"print(\"✓ ForagerAgent environment implemented\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Simulate a ForagerAgent trajectory for a mutate-fix task\n",
"def simulate_mutate_fix_trajectory():\n",
" \"\"\"\n",
" Simulate an agent fixing a synthetically introduced bug.\n",
" Paper: Section 2.3 describes mutate-fix tasks.\n",
" \"\"\"\n",
" env = SimplifiedForagerEnvironment()\n",
" \n",
" # Task: Fix a bug where function returns wrong value\n",
" task = \"Fix the bug in calculate_sum function that causes test failures\"\n",
" \n",
" trajectory_steps = []\n",
" \n",
" # Step 1: Agent creates initial buggy code\n",
" action1 = AgenticAction(\n",
" tool='create',\n",
" args={\n",
" 'path': 'calculator.py',\n",
" 'content': '''def calculate_sum(a, b):\n",
" # BUG: Should return a + b, but returns a - b\n",
" return a - b\n",
"'''\n",
" }\n",
" )\n",
" obs1 = env.execute_action(action1)\n",
" trajectory_steps.append((action1, obs1))\n",
" \n",
" # Step 2: Agent runs tests\n",
" action2 = AgenticAction(\n",
" tool='bash',\n",
" args={'command': 'pytest test_calculator.py'}\n",
" )\n",
" obs2 = env.execute_action(action2)\n",
" trajectory_steps.append((action2, obs2))\n",
" \n",
" # Step 3: Agent views the file to diagnose\n",
" action3 = AgenticAction(\n",
" tool='view',\n",
" args={'path': 'calculator.py'}\n",
" )\n",
" obs3 = env.execute_action(action3)\n",
" trajectory_steps.append((action3, obs3))\n",
" \n",
" # Step 4: Agent fixes the bug\n",
" action4 = AgenticAction(\n",
" tool='edit',\n",
" args={\n",
" 'path': 'calculator.py',\n",
" 'new_content': '''def calculate_sum(a, b):\n",
" # CORRECT: Returns a + b\n",
" return a + b\n",
"'''\n",
" }\n",
" )\n",
" obs4 = env.execute_action(action4)\n",
" trajectory_steps.append((action4, obs4))\n",
" \n",
" # Step 5: Agent runs tests again\n",
" action5 = AgenticAction(\n",
" tool='bash',\n",
" args={'command': 'pytest test_calculator.py'}\n",
" )\n",
" obs5 = env.execute_action(action5)\n",
" trajectory_steps.append((action5, obs5))\n",
" \n",
" success = obs5.success\n",
" \n",
" return AgenticTrajectory(\n",
" task_description=task,\n",
" steps=trajectory_steps,\n",
" success=success\n",
" )\n",
"\n",
"# Generate example trajectory\n",
"traj = simulate_mutate_fix_trajectory()\n",
"\n",
"print(\"ForagerAgent Trajectory Simulation\")\n",
"print(\"=\"*60)\n",
"print(f\"Task: {traj.task_description}\")\n",
"print(f\"Success: {traj.success}\")\n",
"print(f\"\\nSteps ({len(traj.steps)} total):\")\n",
"print(\"=\"*60)\n",
"\n",
"for i, (action, obs) in enumerate(traj.steps, 1):\n",
" print(f\"\\nStep {i}: {action.tool.upper()}\")\n",
" print(f\" Args: {action.args}\")\n",
" print(f\" Observation: {obs.output[:80]}...\" if len(obs.output) > 80 else f\" Observation: {obs.output}\")\n",
" print(f\" Success: {obs.success}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Visualize ForagerAgent dataset statistics from the paper\n",
"forager_stats = {\n",
" 'Metric': ['Total Trajectories', 'Repository Images', 'Underlying Repos', \n",
" 'Issue-Fix Tasks', 'Mutate-Fix Tasks'],\n",
" 'Count': [3_000_000, 10_200, 3_150, 1_650_000, 1_350_000]\n",
"}\n",
"\n",
"# Mutation type breakdown (from Table 1 in paper)\n",
"mutation_types = {\n",
" 'Mutation Type': ['Functions', 'Arguments', 'Variables', 'Statements', 'Operators'],\n",
" 'Percentage': [7, 9, 6, 11, 12]\n",
"}\n",
"\n",
"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"# ForagerAgent scale\n",
"df_forager = pd.DataFrame(forager_stats)\n",
"colors1 = sns.color_palette('Set2', len(df_forager))\n",
"ax1.barh(df_forager['Metric'], df_forager['Count'], color=colors1)\n",
"ax1.set_xlabel('Count', fontsize=12)\n",
"ax1.set_title('ForagerAgent Dataset Scale', fontsize=14, fontweight='bold')\n",
"ax1.set_xscale('log')\n",
"ax1.grid(axis='x', alpha=0.3)\n",
"\n",
"# Mutation type distribution\n",
"df_mutations = pd.DataFrame(mutation_types)\n",
"colors2 = sns.color_palette('Pastel1', len(df_mutations))\n",
"ax2.pie(df_mutations['Percentage'], labels=df_mutations['Mutation Type'], \n",
" autopct='%1.1f%%', colors=colors2, startangle=90)\n",
"ax2.set_title('Mutate-Fix: Mutation Type Distribution', fontsize=14, fontweight='bold')\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nForagerAgent Statistics:\")\n",
"print(df_forager.to_string(index=False))\n",
"print(\"\\nMutation Types:\")\n",
"print(df_mutations.to_string(index=False))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. GRPO Reinforcement Learning Algorithm\n",
"\n",
"CWM uses **Group Relative Policy Optimization (GRPO)** for RL training:\n",
"- Policy gradient method with PPO loss\n",
"- Monte Carlo value estimation (no value model)\n",
"- Multi-turn trajectories with masked loss\n",
"- Asymmetric clipping to prevent entropy collapse\n",
"\n",
"Paper Reference: Section 5.2 - RL algorithm"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class SimplifiedGRPO:\n",
" \"\"\"\n",
" Simplified implementation of GRPO algorithm from the paper.\n",
" This is for educational purposes - real implementation uses GPUs and large models.\n",
" \n",
" Key differences from standard GRPO (as described in Section 5.2):\n",
" - Multi-turn support with masking\n",
" - Asymmetric clipping (ε_high=0.25, ε_low=0.2)\n",
" - No KL regularization\n",
" - Length-weighted mean return\n",
" - Stale trajectory filtering\n",
" \"\"\"\n",
" \n",
" def __init__(self, \n",
" epsilon_high: float = 0.25, # Upper clip for preventing entropy collapse\n",
" epsilon_low: float = 0.2, # Lower clip\n",
" max_staleness: int = 100): # Max steps before trajectory is stale\n",
" self.epsilon_high = epsilon_high\n",
" self.epsilon_low = epsilon_low\n",
" self.max_staleness = max_staleness\n",
" self.current_step = 0\n",
" \n",
" def compute_advantages(self, returns: np.ndarray, lengths: np.ndarray) -> np.ndarray:\n",
" \"\"\"\n",
" Compute advantages using length-weighted mean.\n",
" Paper: \"We compute µ as a length-weighted average\"\n",
" \"\"\"\n",
" # Length-weighted mean return\n",
" mean_return = np.sum(returns * lengths) / np.sum(lengths)\n",
" \n",
" # Advantages without variance normalization (as in paper)\n",
" advantages = returns - mean_return\n",
" \n",
" return advantages\n",
" \n",
" def compute_ppo_loss(self, \n",
" log_probs_old: np.ndarray,\n",
" log_probs_new: np.ndarray,\n",
" advantages: np.ndarray,\n",
" mask: np.ndarray) -> float:\n",
" \"\"\"\n",
" Compute PPO loss with asymmetric clipping.\n",
" \n",
" Args:\n",
" log_probs_old: Log probabilities from old policy\n",
" log_probs_new: Log probabilities from new policy \n",
" advantages: Computed advantages\n",
" mask: Binary mask for valid tokens (multi-turn)\n",
" \"\"\"\n",
" # Probability ratio\n",
" ratio = np.exp(log_probs_new - log_probs_old)\n",
" \n",
" # Clipped surrogate objective with asymmetric clipping\n",
" surrogate1 = ratio * advantages\n",
" \n",
" # Asymmetric clipping (Section 5.2: Clip-higher)\n",
" clip_low = 1.0 - self.epsilon_low\n",
" clip_high = 1.0 + self.epsilon_high\n",
" surrogate2 = np.clip(ratio, clip_low, clip_high) * advantages\n",
" \n",
" # Take minimum and apply mask\n",
" loss_per_token = -np.minimum(surrogate1, surrogate2) * mask\n",
" \n",
" # Normalize by max context length (not trajectory length)\n",
" # Paper: \"divide by the maximum number of tokens in a trajectory\"\n",
" N_max = 131072 # CWM's max context\n",
" loss = np.sum(loss_per_token) / N_max\n",
" \n",
" return loss\n",
" \n",
" def should_skip_trajectory(self, trajectory_step: int) -> bool:\n",
" \"\"\"\n",
" Check if trajectory is too stale.\n",
" Paper: \"skip trajectories whose most recent tokens were generated \n",
" from a policy more than 100 training steps behind\"\n",
" \"\"\"\n",
" return (self.current_step - trajectory_step) > self.max_staleness\n",
"\n",
"print(\"✓ Simplified GRPO implementation complete\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Demonstrate GRPO with synthetic data\n",
"def demo_grpo():\n",
" \"\"\"\n",
" Demonstrate GRPO advantage computation and loss calculation.\n",
" \"\"\"\n",
" grpo = SimplifiedGRPO()\n",
" \n",
" # Simulate a group of 8 trajectories (G=8 as in paper)\n",
" num_trajectories = 8\n",
" \n",
" # Returns (rewards) for each trajectory\n",
" # In real CWM: +1 for correct, -1 for incorrect\n",
" returns = np.array([1.0, -1.0, 1.0, 1.0, -1.0, 1.0, -1.0, 1.0])\n",
" \n",
" # Trajectory lengths (in tokens)\n",
" lengths = np.array([1000, 1500, 800, 1200, 2000, 900, 1100, 1300])\n",
" \n",
" # Compute advantages\n",
" advantages = grpo.compute_advantages(returns, lengths)\n",
" \n",
" print(\"GRPO Demonstration\")\n",
" print(\"=\"*60)\n",
" print(f\"Number of trajectories in group: {num_trajectories}\")\n",
" print(f\"\\nReturns (rewards): {returns}\")\n",
" print(f\"Lengths (tokens): {lengths}\")\n",
" print(f\"\\nLength-weighted mean return: {np.sum(returns * lengths) / np.sum(lengths):.4f}\")\n",
" print(f\"\\nAdvantages: {advantages}\")\n",
" print(f\"Advantage mean: {advantages.mean():.4f}\")\n",
" print(f\"Advantage std: {advantages.std():.4f}\")\n",
" \n",
" # Simulate log probabilities for demonstration\n",
" np.random.seed(42)\n",
" tokens_per_traj = 50 # Simplified\n",
" \n",
" log_probs_old = np.random.normal(-2.0, 0.5, (num_trajectories, tokens_per_traj))\n",
" log_probs_new = log_probs_old + np.random.normal(0, 0.1, (num_trajectories, tokens_per_traj))\n",
" \n",
" # Create mask (all valid for this demo)\n",
" mask = np.ones((num_trajectories, tokens_per_traj))\n",
" \n",
" # Expand advantages to per-token\n",
" advantages_expanded = advantages[:, np.newaxis] * np.ones((num_trajectories, tokens_per_traj))\n",
" \n",
" # Compute loss\n",
" loss = grpo.compute_ppo_loss(log_probs_old, log_probs_new, advantages_expanded, mask)\n",
" \n",
" print(f\"\\nPPO Loss (with asymmetric clipping): {loss:.6f}\")\n",
" print(f\"Epsilon high: {grpo.epsilon_high}\")\n",
" print(f\"Epsilon low: {grpo.epsilon_low}\")\n",
" \n",
" return advantages, returns, lengths\n",
"\n",
"advantages, returns, lengths = demo_grpo()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Visualize advantage distribution\n",
"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"# Advantages vs Returns\n",
"colors = ['green' if r > 0 else 'red' for r in returns]\n",
"ax1.scatter(returns, advantages, c=colors, s=100, alpha=0.6, edgecolors='black')\n",
"ax1.axhline(y=0, color='black', linestyle='--', alpha=0.3)\n",
"ax1.axvline(x=0, color='black', linestyle='--', alpha=0.3)\n",
"ax1.set_xlabel('Return (Reward)', fontsize=12)\n",
"ax1.set_ylabel('Advantage', fontsize=12)\n",
"ax1.set_title('GRPO: Returns vs Advantages', fontsize=14, fontweight='bold')\n",
"ax1.grid(alpha=0.3)\n",
"\n",
"# Length vs Return\n",
"ax2.scatter(lengths, returns, c=colors, s=100, alpha=0.6, edgecolors='black')\n",
"ax2.axhline(y=0, color='black', linestyle='--', alpha=0.3)\n",
"ax2.set_xlabel('Trajectory Length (tokens)', fontsize=12)\n",
"ax2.set_ylabel('Return (Reward)', fontsize=12)\n",
"ax2.set_title('Trajectory Length vs Return', fontsize=14, fontweight='bold')\n",
"ax2.grid(alpha=0.3)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. Multi-Task Joint RL Training\n",
"\n",
"CWM performs joint RL across multiple environments:\n",
"- **40% Software Engineering** (SWE-bench style tasks)\n",
"- **40% Competitive Programming** (CodeContests)\n",
"- **20% Mathematics** (with optional Python tool calling)\n",
"\n",
"Training: **172B tokens** of RL, **26.5k gradient steps**\n",
"\n",
"Paper Reference: Section 5.4 - Joint RL"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Simulate multi-task RL training statistics\n",
"rl_config = {\n",
" 'Task Distribution': {\n",
" 'Software Engineering': 40,\n",
" 'Competitive Programming': 40,\n",
" 'Mathematics': 20\n",
" },\n",
" 'Training Stats': {\n",
" 'Total Tokens': '172B',\n",
" 'Gradient Steps': '26.5k',\n",
" 'Rollouts per Prompt': 8,\n",
" 'Max Context Length': '131k tokens'\n",
" },\n",
" 'Environment Details': {\n",
" 'SWE': {\n",
" 'instances': 12600,\n",
" 'max_turns': 128,\n",
" 'tools': ['bash', 'edit', 'create', 'submit']\n",
" },\n",
" 'Code': {\n",
" 'instances': 81000,\n",
" 'languages': ['Python', 'C++', 'Rust', 'Go', 'Java', 'JavaScript'],\n",
" 'max_tokens': 64000\n",
" },\n",
" 'Math': {\n",
" 'instances': 278000,\n",
" 'tool_calling': '2% of tasks',\n",
" 'max_tokens': 64000\n",
" }\n",
" }\n",
"}\n",
"\n",
"# Visualize task distribution\n",
"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"# Task distribution pie chart\n",
"tasks = list(rl_config['Task Distribution'].keys())\n",
"percentages = list(rl_config['Task Distribution'].values())\n",
"colors_pie = ['#3498db', '#e74c3c', '#2ecc71']\n",
"\n",
"ax1.pie(percentages, labels=tasks, autopct='%1.1f%%', colors=colors_pie, startangle=90)\n",
"ax1.set_title('Multi-Task RL: Task Distribution', fontsize=14, fontweight='bold')\n",
"\n",
"# Dataset sizes\n",
"env_names = list(rl_config['Environment Details'].keys())\n",
"env_sizes = [rl_config['Environment Details'][env]['instances'] for env in env_names]\n",
"colors_bar = ['#3498db', '#e74c3c', '#2ecc71']\n",
"\n",
"ax2.barh(env_names, env_sizes, color=colors_bar)\n",
"ax2.set_xlabel('Number of Training Instances', fontsize=12)\n",
"ax2.set_title('RL Training Dataset Sizes', fontsize=14, fontweight='bold')\n",
"ax2.set_xscale('log')\n",
"ax2.grid(axis='x', alpha=0.3)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"Multi-Task RL Configuration\")\n",
"print(\"=\"*60)\n",
"for key, value in rl_config['Training Stats'].items():\n",
" print(f\"{key:25s}: {value}\")\n",
"\n",
"print(\"\\nEnvironment Details:\")\n",
"print(\"=\"*60)\n",
"for env, details in rl_config['Environment Details'].items():\n",
" print(f\"\\n{env}:\")\n",
" for k, v in details.items():\n",
" print(f\" {k:15s}: {v}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. CruxEval: Execution Trace Prediction\n",
"\n",
"CWM can predict Python execution outcomes in multiple modes:\n",
"1. **Single-step**: Direct output prediction\n",
"2. **Full trace**: Line-by-line execution simulation\n",
"3. **Natural language reasoning**: Explaining execution\n",
"\n",
"Results: **94.3% on CruxEval-Output** with reasoning\n",
"\n",
"Paper Reference: Section 7.3 - Execution trace prediction"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class ExecutionPredictor:\n",
" \"\"\"\n",
" Demonstrates CWM's execution trace prediction capability.\n",
" In the real model, this uses the trained 32B parameter LLM.\n",
" Here we simulate the concept with actual execution + formatting.\n",
" \"\"\"\n",
" \n",
" def __init__(self):\n",
" self.tracer = PythonTracer()\n",
" \n",
" def predict_single_step(self, func, *args):\n",
" \"\"\"\n",
" Single-step prediction: directly predict output.\n",
" This mode scored 66.6% on CruxEval-O for CWM.\n",
" \"\"\"\n",
" result, _ = self.tracer.trace_function(func, *args)\n",
" return {\n",
" 'mode': 'single-step',\n",
" 'prediction': result,\n",
" 'explanation': 'Direct output prediction without intermediate steps'\n",
" }\n",
" \n",
" def predict_full_trace(self, func, *args):\n",
" \"\"\"\n",
" Full trace prediction: simulate line-by-line execution.\n",
" This mode scored 87.7% on CruxEval-O for CWM-SFT.\n",
" \"\"\"\n",
" result, trace = self.tracer.trace_function(func, *args)\n",
" \n",
" formatted_trace = self.tracer.format_trace_cwm_style()\n",
" \n",
" return {\n",
" 'mode': 'full-trace',\n",
" 'prediction': result,\n",
" 'trace': formatted_trace,\n",
" 'trace_length': len(trace),\n",
" 'explanation': 'Line-by-line execution simulation with state tracking'\n",
" }\n",
" \n",
" def predict_with_reasoning(self, func, *args):\n",
" \"\"\"\n",
" Natural language reasoning mode.\n",
" This mode scored 94.3% on CruxEval-O for CWM.\n",
" \"\"\"\n",
" result, trace = self.tracer.trace_function(func, *args)\n",
" \n",
" # Generate natural language explanation\n",
" reasoning = self._generate_reasoning(func, args, trace, result)\n",
" \n",
" return {\n",
" 'mode': 'reasoning',\n",
" 'prediction': result,\n",
" 'reasoning': reasoning,\n",
" 'explanation': 'Natural language reasoning about execution'\n",
" }\n",
" \n",
" def _generate_reasoning(self, func, args, trace, result):\n",
" \"\"\"Generate natural language reasoning (simplified).\"\"\"\n",
" reasoning_steps = []\n",
" reasoning_steps.append(f\"<think>\")\n",
" reasoning_steps.append(f\"The function {func.__name__} is called with arguments {args}.\")\n",
" reasoning_steps.append(f\"Let me trace through the execution:\")\n",
" \n",
" # Summarize key steps\n",
" for i, step in enumerate(trace[:5]): # First few steps\n",
" if step['event'] == 'line':\n",
" reasoning_steps.append(f\" - Execute: {step['action']}\")\n",
" reasoning_steps.append(f\" Variables: {step['locals']}\")\n",
" \n",
" if len(trace) > 5:\n",
" reasoning_steps.append(f\" ... ({len(trace)-5} more steps)\")\n",
" \n",
" reasoning_steps.append(f\"Therefore, the function returns: {result}\")\n",
" reasoning_steps.append(f\"</think>\")\n",
" \n",
" return '\\n'.join(reasoning_steps)\n",
"\n",
"print(\"✓ Execution Predictor implemented\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Demonstrate all three prediction modes\n",
"def demo_cruxeval_style():\n",
" \"\"\"Demonstrate CruxEval-style execution prediction.\"\"\"\n",
" \n",
" # Example function from paper\n",
" def f(d, k):\n",
" new_d = {}\n",
" for key, val in d.items():\n",
" if key < k:\n",
" new_d[key] = val\n",
" return new_d\n",
" \n",
" predictor = ExecutionPredictor()\n",
" \n",
" test_input = ({1: 2, 2: 4, 3: 3}, 3)\n",
" \n",
" print(\"CruxEval-Style Execution Prediction\")\n",
" print(\"=\"*60)\n",
" print(f\"Function: f{inspect.signature(f)}\")\n",
" print(f\"Input: f{test_input}\")\n",
" print(\"=\"*60)\n",
" \n",
" # Mode 1: Single-step\n",
" result1 = predictor.predict_single_step(f, *test_input)\n",
" print(f\"\\n1. SINGLE-STEP MODE\")\n",
" print(f\" Prediction: {result1['prediction']}\")\n",
" print(f\" Note: {result1['explanation']}\")\n",
" \n",
" # Mode 2: Full trace\n",
" result2 = predictor.predict_full_trace(f, *test_input)\n",
" print(f\"\\n2. FULL TRACE MODE\")\n",
" print(f\" Prediction: {result2['prediction']}\")\n",
" print(f\" Trace length: {result2['trace_length']} steps\")\n",
" print(f\" Note: {result2['explanation']}\")\n",
" print(f\" First 3 trace lines:\")\n",
" for line in result2['trace'].split('\\n')[:3]:\n",
" print(f\" {line}\")\n",
" \n",
" # Mode 3: Reasoning\n",
" result3 = predictor.predict_with_reasoning(f, *test_input)\n",
" print(f\"\\n3. REASONING MODE\")\n",
" print(f\" Prediction: {result3['prediction']}\")\n",
" print(f\" Note: {result3['explanation']}\")\n",
" print(f\" Reasoning:\")\n",
" for line in result3['reasoning'].split('\\n'):\n",
" print(f\" {line}\")\n",
" \n",
" return result1, result2, result3\n",
"\n",
"results = demo_cruxeval_style()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Visualize CruxEval results from the paper\n",
"cruxeval_results = {\n",
" 'Model': ['CWM-SFT', 'CWM-SFT', 'CWM', 'CWM'],\n",
" 'Mode': ['Single-step', 'Full Trace', 'Single-step', 'Reasoning'],\n",
" 'Score': [67.8, 87.7, 66.6, 94.3],\n",
" 'Avg Tokens': [100, 497, 100, 1164] # Approximate from paper\n",
"}\n",
"\n",
"df_cruxeval = pd.DataFrame(cruxeval_results)\n",
"\n",
"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
"\n",
"# Accuracy by mode\n",
"colors = ['#3498db', '#2ecc71', '#3498db', '#e74c3c']\n",
"bars = ax1.bar(range(len(df_cruxeval)), df_cruxeval['Score'], color=colors, alpha=0.7, edgecolor='black')\n",
"ax1.set_xticks(range(len(df_cruxeval)))\n",
"ax1.set_xticklabels([f\"{m}\\n({mod})\" for m, mod in zip(df_cruxeval['Model'], df_cruxeval['Mode'])], \n",
" rotation=0, ha='center')\n",
"ax1.set_ylabel('CruxEval-Output Score (%)', fontsize=12)\n",
"ax1.set_title('CWM Performance on CruxEval-Output', fontsize=14, fontweight='bold')\n",
"ax1.grid(axis='y', alpha=0.3)\n",
"ax1.set_ylim(0, 100)\n",
"\n",
"# Add value labels on bars\n",
"for i, (bar, score) in enumerate(zip(bars, df_cruxeval['Score'])):\n",
" ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1, \n",
" f'{score:.1f}%', ha='center', va='bottom', fontweight='bold')\n",
"\n",
"# Tokens vs Accuracy tradeoff\n",
"for i, row in df_cruxeval.iterrows():\n",
" ax2.scatter(row['Avg Tokens'], row['Score'], s=200, alpha=0.7, \n",
" color=colors[i], edgecolors='black', linewidth=2)\n",
" ax2.annotate(row['Mode'], (row['Avg Tokens'], row['Score']), \n",
" xytext=(10, 10), textcoords='offset points', fontsize=9)\n",
"\n",
"ax2.set_xlabel('Average Tokens per Prediction', fontsize=12)\n",
"ax2.set_ylabel('CruxEval-Output Score (%)', fontsize=12)\n",
"ax2.set_title('Accuracy vs Computational Cost', fontsize=14, fontweight='bold')\n",
"ax2.grid(alpha=0.3)\n",
"ax2.set_ylim(60, 100)\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nCruxEval Results Summary:\")\n",
"print(df_cruxeval.to_string(index=False))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 6. SWE-bench Verified Results\n",
"\n",
"CWM achieves state-of-the-art performance on software engineering tasks:\n",
"- **Base (pass@1)**: 53.9% resolve rate\n",
"- **Test-Time Scaling (best@16)**: 65.8% resolve rate\n",
"\n",
"Paper Reference: Section 7.2 - Agentic evaluation, SWE-bench Verified"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# SWE-bench Verified results from the paper\n",
"swe_bench_results = {\n",
" 'Model': [\n",
" 'CWM (base)',\n",
" 'CWM (tts best@16)',\n",
" 'Qwen3-32B',\n",
" 'DeepSeek-V3',\n",
" 'GPT-oss-120B (low)',\n",
" 'GPT-oss-120B (high)',\n",
" 'Claude 3.5 Sonnet'\n",
" ],\n",
" 'Resolve Rate': [53.9, 65.8, 38.4, 52.0, 53.9, 70.7, 90.0],\n",
" 'Type': ['Open', 'Open', 'Open', 'Open', 'Proprietary', 'Proprietary', 'Proprietary'],\n",
" 'Size': ['32B', '32B', '32B', '671B', '120B', '120B', 'Unknown']\n",
"}\n",
"\n",
"df_swe = pd.DataFrame(swe_bench_results)\n",
"\n",
"# Visualize results\n",
"fig, ax = plt.subplots(1, 1, figsize=(12, 6))\n",
"\n",
"# Color by type\n",
"colors_map = {'Open': '#3498db', 'Proprietary': '#e74c3c'}\n",
"colors = [colors_map[t] for t in df_swe['Type']]\n",
"\n",
"bars = ax.barh(range(len(df_swe)), df_swe['Resolve Rate'], color=colors, alpha=0.7, edgecolor='black')\n",
"ax.set_yticks(range(len(df_swe)))\n",
"ax.set_yticklabels([f\"{m} ({s})\" for m, s in zip(df_swe['Model'], df_swe['Size'])])\n",
"ax.set_xlabel('Resolve Rate (%)', fontsize=12)\n",
"ax.set_title('SWE-bench Verified Performance', fontsize=14, fontweight='bold')\n",
"ax.grid(axis='x', alpha=0.3)\n",
"\n",
"# Add value labels\n",
"for i, (bar, score) in enumerate(zip(bars, df_swe['Resolve Rate'])):\n",
" ax.text(score + 1, bar.get_y() + bar.get_height()/2, \n",
" f'{score:.1f}%', va='center', fontweight='bold')\n",
"\n",
"# Add legend\n",
"from matplotlib.patches import Patch\n",
"legend_elements = [Patch(facecolor='#3498db', label='Open Weights'),\n",
" Patch(facecolor='#e74c3c', label='Proprietary')]\n",
"ax.legend(handles=legend_elements, loc='lower right')\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nSWE-bench Verified Results:\")\n",
"print(df_swe.to_string(index=False))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Test-time scaling analysis\n",
"tts_data = {\n",
" 'k': [1, 2, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40],\n",
" 'best@k': [53.9, 57.2, 60.5, 63.1, 64.5, 65.8, 66.2, 66.5, 66.7, 66.8, 66.9, 67.0],\n",
" 'majority@k': [53.9, 54.5, 55.8, 56.9, 57.7, 58.2, 58.3, 58.4, 58.4, 58.4, 58.4, 58.4],\n",
" 'pass@k': [53.9, 63.2, 70.1, 75.2, 77.8, 79.3, 79.9, 80.2, 80.3, 80.4, 80.4, 80.4]\n",
"}\n",
"\n",
"df_tts = pd.DataFrame(tts_data)\n",
"\n",
"# Plot test-time scaling\n",
"fig, ax = plt.subplots(1, 1, figsize=(12, 6))\n",
"\n",
"ax.plot(df_tts['k'], df_tts['best@k'], marker='o', linewidth=2, \n",
" label='best@k (with test generation)', color='#2ecc71')\n",
"ax.plot(df_tts['k'], df_tts['majority@k'], marker='s', linewidth=2, \n",
" label='majority@k (simple voting)', color='#3498db')\n",
"ax.plot(df_tts['k'], df_tts['pass@k'], marker='^', linewidth=2, \n",
" label='pass@k (oracle)', color='#e74c3c', linestyle='--')\n",
"\n",
"ax.set_xlabel('Number of Candidate Solutions (k)', fontsize=12)\n",
"ax.set_ylabel('Resolve Rate (%)', fontsize=12)\n",
"ax.set_title('SWE-bench: Test-Time Scaling Analysis', fontsize=14, fontweight='bold')\n",
"ax.grid(alpha=0.3)\n",
"ax.legend(fontsize=11)\n",
"\n",
"# Highlight the best@16 point\n",
"ax.axvline(x=16, color='gray', linestyle=':', alpha=0.5)\n",
"ax.text(16, 50, 'k=16\\n(65.8%)', ha='center', fontsize=10, \n",
" bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nTest-Time Scaling Results:\")\n",
"print(df_tts.to_string(index=False))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 7. Model Architecture and Training Pipeline\n",
"\n",
"CWM is a **32B parameter dense decoder-only Transformer**:\n",
"- **Architecture**: Alternating local (8k) and global (131k) attention\n",
"- **Training**: Pre-training (8T tokens) → Mid-training (5T tokens) → SFT (100B tokens) → RL (172B tokens)\n",
"- **Context**: Up to 131k tokens\n",
"\n",
"Paper Reference: Section 4 - CWM: architecture, pre-training, and scaling laws"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# CWM architecture specs\n",
"architecture_specs = {\n",
" 'Specification': [\n",
" 'Total Parameters',\n",
" 'Layers',\n",
" 'Hidden Dimension',\n",
" 'Intermediate Dimension',\n",
" 'Attention Heads',\n",
" 'Key-Value Heads',\n",
" 'Local Window Size',\n",
" 'Max Global Context',\n",
" 'Vocabulary Size'\n",
" ],\n",
" 'Value': [\n",
" '32B',\n",
" '64',\n",
" '6144',\n",
" '21504',\n",
" '48',\n",
" '8 (GQA)',\n",
" '8192 tokens',\n",
" '131072 tokens',\n",
" '128256 tokens'\n",
" ]\n",
"}\n",
"\n",
"# Training pipeline\n",
"training_pipeline = {\n",
" 'Stage': ['Pre-training', 'Mid-training', 'SFT', 'RL'],\n",
" 'Tokens': ['8T', '5T', '100B', '172B'],\n",
" 'Context Length': ['8k', '131k', '32k', '131k'],\n",
" 'Batch Size': ['8.4M', '33.6M', '2.1M', '8.4-16.8M'],\n",
" 'Focus': ['General + Code', 'CWM Data', 'Instruction Following', 'Multi-Task RL']\n",
"}\n",
"\n",
"df_arch = pd.DataFrame(architecture_specs)\n",
"df_train = pd.DataFrame(training_pipeline)\n",
"\n",
"print(\"CWM Architecture Specifications\")\n",
"print(\"=\"*60)\n",
"print(df_arch.to_string(index=False))\n",
"\n",
"print(\"\\n\" + \"=\"*60)\n",
"print(\"Training Pipeline\")\n",
"print(\"=\"*60)\n",
"print(df_train.to_string(index=False))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Visualize training pipeline\n",
"fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10))\n",
"\n",
"# Training tokens by stage\n",
"stages = df_train['Stage']\n",
"tokens_raw = ['8000000000000', '5000000000000', '100000000000', '172000000000']\n",
"tokens_numeric = [float(t) for t in tokens_raw]\n",
"colors_stages = ['#3498db', '#2ecc71', '#f39c12', '#e74c3c']\n",
"\n",
"bars = ax1.barh(stages, tokens_numeric, color=colors_stages, alpha=0.7, edgecolor='black')\n",
"ax1.set_xlabel('Number of Tokens (log scale)', fontsize=12)\n",
"ax1.set_title('CWM Training: Tokens per Stage', fontsize=14, fontweight='bold')\n",
"ax1.set_xscale('log')\n",
"ax1.grid(axis='x', alpha=0.3)\n",
"\n",
"# Add labels\n",
"for bar, token_str in zip(bars, df_train['Tokens']):\n",
" ax1.text(bar.get_width() * 1.1, bar.get_y() + bar.get_height()/2, \n",
" token_str, va='center', fontweight='bold')\n",
"\n",
"# Context length progression\n",
"context_lengths = [8192, 131072, 32768, 131072]\n",
"stage_indices = range(len(stages))\n",
"\n",
"ax2.plot(stage_indices, context_lengths, marker='o', linewidth=3, \n",
" markersize=10, color='#9b59b6', label='Max Context Length')\n",
"ax2.set_xticks(stage_indices)\n",
"ax2.set_xticklabels(stages)\n",
"ax2.set_ylabel('Context Length (tokens)', fontsize=12)\n",
"ax2.set_title('CWM Training: Context Length Evolution', fontsize=14, fontweight='bold')\n",
"ax2.grid(alpha=0.3)\n",
"ax2.set_yscale('log')\n",
"\n",
"# Add value labels\n",
"for i, (x, y) in enumerate(zip(stage_indices, context_lengths)):\n",
" ax2.text(x, y * 1.15, f'{y:,}', ha='center', fontweight='bold')\n",
"\n",
"plt.tight_layout()\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 8. Mid-Training Ablation Study\n",
"\n",
"The paper demonstrates the impact of different CWM data components:\n",
"- **GitHub PRs**: Improves SWE-bench performance\n",
"- **Python Traces**: Significantly improves CruxEval\n",
"- **ForagerAgent**: Improves agentic SWE capabilities\n",
"\n",
"Paper Reference: Section 7.1 - The impact of CWM data"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Ablation study results (8B models, Table 4)\n",
"ablation_results = {\n",
" 'Configuration': [\n",
" 'Baseline (no CWM)',\n",
" '+ GitHub PRs',\n",
" '+ PRs + Traces',\n",
" '+ PRs + Traces + Forager'\n",
" ],\n",
" 'CruxEval-O': [45.4, 44.6, 73.9, 74.5],\n",
" 'CruxEval-I': [44.1, 45.8, 51.5, 54.8],\n",
" 'SWE-bench': [14.6, 18.6, 18.4, 22.1]\n",
"}\n",
"\n",
"df_ablation = pd.DataFrame(ablation_results)\n",
"\n",
"# Visualize ablation study\n",
"fig, axes = plt.subplots(1, 3, figsize=(16, 5))\n",
"\n",
"metrics = ['CruxEval-O', 'CruxEval-I', 'SWE-bench']\n",
"colors_ablation = ['#95a5a6', '#3498db', '#2ecc71', '#e74c3c']\n",
"\n",
"for idx, metric in enumerate(metrics):\n",
" ax = axes[idx]\n",
" bars = ax.bar(range(4), df_ablation[metric], color=colors_ablation, alpha=0.7, edgecolor='black')\n",
" ax.set_xticks(range(4))\n",
" ax.set_xticklabels(['Baseline', '+PRs', '+Traces', '+All'], rotation=0)\n",
" ax.set_ylabel('Score (%)', fontsize=11)\n",
" ax.set_title(f'{metric} Performance', fontsize=12, fontweight='bold')\n",
" ax.grid(axis='y', alpha=0.3)\n",
" \n",
" # Add value labels\n",
" for bar, val in zip(bars, df_ablation[metric]):\n",
" ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1, \n",
" f'{val:.1f}', ha='center', va='bottom', fontweight='bold', fontsize=9)\n",
"\n",
"plt.suptitle('Mid-Training Ablation Study (8B Models)', fontsize=14, fontweight='bold', y=1.02)\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nAblation Study Results:\")\n",
"print(df_ablation.to_string(index=False))\n",
"print(\"\\nKey Findings:\")\n",
"print(\"- Python traces dramatically improve CruxEval-O: 44.6% → 73.9%\")\n",
"print(\"- ForagerAgent improves SWE-bench: 18.4% → 22.1%\")\n",
"print(\"- All components together achieve best overall performance\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 9. Summary and Key Takeaways\n",
"\n",
"### Main Contributions:\n",
"\n",
"1. **Code World Modeling**: Training LLMs on execution traces to understand code semantics\n",
"2. **Large-Scale Agentic Data**: 3M ForagerAgent trajectories from real repositories\n",
"3. **Multi-Task RL**: Joint training across SWE, coding, and math tasks\n",
"4. **Strong Performance**: State-of-the-art results on multiple benchmarks\n",
"\n",
"### Key Results:\n",
"\n",
"- **SWE-bench Verified**: 65.8% (with test-time scaling)\n",
"- **LiveCodeBench**: 68.6%\n",
"- **CruxEval-Output**: 94.3% (with reasoning)\n",
"- **Math-500**: 96.6%\n",
"- **AIME 2024**: 76.0%\n",
"\n",
"### Methodological Insights:\n",
"\n",
"- Execution traces improve code understanding beyond static analysis\n",
"- Agentic mid-training data enhances post-RL performance\n",
"- Multi-task RL improves generalization\n",
"- Test-time scaling provides significant gains"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Final comprehensive results summary\n",
"benchmark_results = {\n",
" 'Benchmark': [\n",
" 'SWE-bench Verified',\n",
" 'SWE-bench Verified (TTS)',\n",
" 'LiveCodeBench-v5',\n",
" 'CruxEval-Output',\n",
" 'Math-500',\n",
" 'AIME 2024',\n",
" 'Aider Polyglot',\n",
" 'Terminal-Bench'\n",
" ],\n",
" 'CWM Score': [53.9, 65.8, 68.6, 94.3, 96.6, 76.0, 35.1, 26.3],\n",
" 'Category': ['Agentic SWE', 'Agentic SWE', 'Coding', 'Execution', 'Math', 'Math', 'Coding', 'Agentic']\n",
"}\n",
"\n",
"df_results = pd.DataFrame(benchmark_results)\n",
"\n",
"# Visualize all results\n",
"fig, ax = plt.subplots(1, 1, figsize=(12, 8))\n",
"\n",
"category_colors = {\n",
" 'Agentic SWE': '#3498db',\n",
" 'Coding': '#2ecc71',\n",
" 'Execution': '#9b59b6',\n",
" 'Math': '#e74c3c',\n",
" 'Agentic': '#f39c12'\n",
"}\n",
"\n",
"colors_final = [category_colors[cat] for cat in df_results['Category']]\n",
"\n",
"bars = ax.barh(range(len(df_results)), df_results['CWM Score'], \n",
" color=colors_final, alpha=0.7, edgecolor='black')\n",
"ax.set_yticks(range(len(df_results)))\n",
"ax.set_yticklabels(df_results['Benchmark'])\n",
"ax.set_xlabel('Score (%)', fontsize=12)\n",
"ax.set_title('CWM: Comprehensive Benchmark Results', fontsize=14, fontweight='bold')\n",
"ax.grid(axis='x', alpha=0.3)\n",
"\n",
"# Add value labels\n",
"for bar, score in zip(bars, df_results['CWM Score']):\n",
" ax.text(score + 1.5, bar.get_y() + bar.get_height()/2, \n",
" f'{score:.1f}%', va='center', fontweight='bold')\n",
"\n",
"# Add legend\n",
"from matplotlib.patches import Patch\n",
"legend_elements = [Patch(facecolor=color, label=cat) \n",
" for cat, color in category_colors.items()]\n",
"ax.legend(handles=legend_elements, loc='lower right', title='Category')\n",
"\n",
"plt.tight_layout()\n",
"plt.show()\n",
"\n",
"print(\"\\nCWM Benchmark Results Summary:\")\n",
"print(\"=\"*60)\n",
"print(df_results.to_string(index=False))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 10. Scaling Up to Production\n",
"\n",
"This notebook demonstrated the **concepts** with small-scale examples. To scale up to the full CWM approach:\n",
"\n",
"### Data Collection:\n",
"- Build 35k+ Docker repository images using RepoAgent/Activ pipeline\n",
"- Generate 120M+ Python execution traces across repositories\n",
"- Collect 3M agentic trajectories with ForagerAgent\n",
"- Gather 75M natural language trace descriptions\n",
"\n",
"### Model Training:\n",
"- Pre-train 32B parameter Transformer on 8T tokens (requires ~2048 H100 GPUs)\n",
"- Mid-train on 5T tokens of CWM data with 131k context\n",
"- SFT on 100B tokens of instruction data\n",
"- Multi-task RL for 172B tokens across SWE, coding, and math\n",
"\n",
"### Infrastructure Requirements:\n",
"- **Pre-training**: 2048 H100 GPUs\n",
"- **Mid-training**: 2048 H100 GPUs with 8x tensor parallelism\n",
"- **RL**: 2560-4608 GPUs with async worker-trainer architecture\n",
"- **Code execution**: Distributed containerized execution service\n",
"- **Storage**: Petabyte-scale for training data\n",
"\n",
"### Expected Timeline:\n",
"- Data collection: Several months\n",
"- Pre-training + Mid-training: Weeks on large clusters\n",
"- Post-training (SFT + RL): Additional weeks\n",
"\n",
"This notebook provides the algorithmic foundation - production deployment requires significant computational resources beyond what's feasible in this environment."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Conclusion\n",
"\n",
"This notebook demonstrated the key computational workflows from the CWM paper:\n",
"\n",
"✓ **Python Execution Tracing** - Capturing observation-action trajectories \n",
"✓ **ForagerAgent Data Generation** - Agentic interactions with code repositories \n",
"✓ **GRPO Reinforcement Learning** - Policy optimization with multi-turn support \n",
"✓ **Execution Trace Prediction** - Multiple prediction modes (single-step, full trace, reasoning) \n",
"✓ **Multi-Task Joint RL** - Training across SWE, coding, and math environments \n",
"\n",
"The CWM approach demonstrates that training on code **execution dynamics** (not just static syntax) significantly improves LLM performance on code generation and reasoning tasks.\n",
"\n",
"**For researchers**: All model checkpoints (CWM-pretrain, CWM-SFT, CWM) are available for non-commercial research at:\n",
"- GitHub: github.com/facebookresearch/cwm\n",
"- Hugging Face: huggingface.co/facebook/cwm\n",
"\n",
"**Note**: This notebook used small-scale demonstrations suitable for educational purposes. Full replication requires significant computational resources (thousands of GPUs, petabytes of data)."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.0"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment