diff --git a/AGENTS.md b/AGENTS.md index de5dc4e..9d0e8f7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,6 +1,6 @@ # AGENTS.md -This repository contains 12 agent skills for the Together AI platform. Each skill is a self-contained directory following the [Agent Skills specification](https://agentskills.io/specification). +This repository contains 13 agent skills for the Together AI platform. Each skill is a self-contained directory following the [Agent Skills specification](https://agentskills.io/specification). ## Skill registry @@ -17,6 +17,7 @@ This repository contains 12 agent skills for the Together AI platform. Each skil - **together-dedicated-endpoints**: Single-tenant GPU endpoints on Together AI with autoscaling and no rate limits. Deploy fine-tuned or uploaded models, size hardware, and manage endpoint lifecycle. Reach for it whenever the user needs predictable always-on hosting rather than serverless inference, custom containers, or raw clusters. - **together-dedicated-containers**: Custom Dockerized inference workers on Together AI's managed GPU infrastructure. Build with Sprocket SDK, configure with Jig CLI, submit async queue jobs, and poll results. Reach for it whenever the user needs container-level control rather than a standard model endpoint or raw cluster. - **together-gpu-clusters**: On-demand and reserved GPU clusters (H100, H200, B200) on Together AI with Kubernetes or Slurm orchestration, shared storage, credential management, and cluster scaling for ML and HPC jobs. Reach for it when the user needs multi-node compute or infrastructure control rather than a managed model endpoint. +- **together-rl**: GRPO and policy-gradient reinforcement-learning post-training on Together AI via low-level training sessions: create a session, sample token-level rollouts, score them, and apply forward_backward plus optim_step. Pairs the RL training API with the together-sandbox SDK to compute rewards by executing model-generated code in isolated sandboxes (run a test suite; pass/fail becomes the reward). Reach for it whenever the user wants to train a model with reinforcement learning, run a GRPO loop, or compute verifiable code-execution rewards rather than managed supervised fine-tuning or plain inference. diff --git a/README.md b/README.md index 562fa59..be7b719 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,7 @@ Each skill contains: | **together-dedicated-endpoints** | Single-tenant GPU endpoints on Together AI with autoscaling and no rate limits. | `deploy_finetuned.py`, `manage_endpoint.py`, `upload_custom_model.py` | | **together-dedicated-containers** | Custom Dockerized inference workers on Together AI's managed GPU infrastructure. | `queue_client.py`, `sprocket_hello_world.py` | | **together-gpu-clusters** | On-demand and reserved GPU clusters (H100, H200, B200) on Together AI with Kubernetes or Slurm orchestration, shared ... | `manage_cluster.py`, `manage_storage.py` | +| **together-rl** | GRPO and policy-gradient reinforcement-learning post-training on Together AI via low-level training sessions: create ... | `grpo_sandbox_reward.py` | ## Installation diff --git a/quality/trigger-evals/together-rl.json b/quality/trigger-evals/together-rl.json new file mode 100644 index 0000000..9689697 --- /dev/null +++ b/quality/trigger-evals/together-rl.json @@ -0,0 +1,26 @@ +[ + { + "query": "Run a GRPO reinforcement-learning loop on Together AI where the reward comes from running the model's generated code against a test suite.", + "should_trigger": true + }, + { + "query": "I want to post-train a model with policy gradients using together-rl: sample rollouts, compute advantages, and call forward_backward and optim_step.", + "should_trigger": true + }, + { + "query": "Train a coding model with RL where each candidate solution is executed in a sandbox and pass/fail is the reward.", + "should_trigger": true + }, + { + "query": "Fine-tune Llama 3.1 8B on my instruction dataset with LoRA on Together AI.", + "should_trigger": false + }, + { + "query": "Just run a chat completion with Llama 3.3 70B and stream the tokens.", + "should_trigger": false + }, + { + "query": "Create an isolated sandbox from a Docker image and run some bash commands.", + "should_trigger": false + } +] diff --git a/skills/together-rl/SKILL.md b/skills/together-rl/SKILL.md new file mode 100644 index 0000000..6266248 --- /dev/null +++ b/skills/together-rl/SKILL.md @@ -0,0 +1,94 @@ +--- +name: together-rl +description: "GRPO and policy-gradient reinforcement-learning post-training on Together AI via low-level training sessions: create a session, sample token-level rollouts, score them, and apply forward_backward plus optim_step. Pairs the RL training API with the together-sandbox SDK to compute rewards by executing model-generated code in isolated sandboxes (run a test suite; pass/fail becomes the reward). Reach for it whenever the user wants to train a model with reinforcement learning, run a GRPO loop, or compute verifiable code-execution rewards rather than managed supervised fine-tuning or plain inference." +--- + +# Together RL + +## Overview + +Use Together RL when the user wants to post-train a model with reinforcement learning — driving the +training loop themselves through low-level training sessions rather than submitting a managed job. + +Typical fits: + +- GRPO / policy-gradient post-training with a custom reward +- Coding or agentic RL where the reward comes from **executing** the model's output (run tests, check exit code) +- Verifiable-reward RL (math, code, tool use) with token-level sampling and advantages +- Multi-turn rollout training where each rollout is scored out-of-band + +## When This Skill Wins + +- The user wants to run the GRPO loop directly: sample → score → `forward_backward` → `optim_step` +- The reward requires running untrusted, model-generated code (so it must execute in a sandbox) +- The user references GRPO, RL fine-tuning, policy gradients, advantages, reward functions, or verifiers +- Token-level control is needed (logprobs, advantages, loss masks) rather than a managed dataset job + +## Hand Off To Another Skill + +- Use `together-sandbox` for the sandbox API itself — this skill calls it to execute reward code, but the + snapshot/sandbox/exec surface lives there +- Use `together-fine-tuning` for managed supervised fine-tuning, LoRA, or DPO jobs (no hand-written loop) +- Use `together-chat-completions` if the user only needs inference, not training +- Use `together-evaluations` to score a model with an LLM judge rather than a programmatic reward + +## Quick Routing + +- **Run a GRPO loop with code-execution rewards** + - Start with [scripts/grpo_sandbox_reward.py](scripts/grpo_sandbox_reward.py) +- **Training-session lifecycle, the three operations, and the GRPO sample schema** + - Read [references/grpo-loop.md](references/grpo-loop.md) +- **Compute rewards by running code in a sandbox (the integration)** + - Read [references/sandbox-rewards.md](references/sandbox-rewards.md) + +## Workflow + +1. Create a training session with `client.beta.rl.sessions.create(base_model=...)` and wait for status + `TRAINING_SESSION_STATUS_RUNNING`. +2. For each step, sample a group of rollouts per prompt with `client.beta.rl.training.sample(...)` + (`num_samples` = group size); each call returns an async **operation** you poll to completion. +3. Decode the sampled token sequences and **compute rewards out-of-band** — for code tasks, hand off to + `together-sandbox`: run each candidate's test suite in an isolated sandbox, `exit 0` → reward `1.0`. +4. Compute GRPO advantages: `advantage = reward − mean(group_rewards)` (center within each prompt's group). +5. Assemble GRPO samples and call `client.beta.rl.training.forward_backward(...)` with `LOSS_TYPE_GRPO`. +6. Apply `client.beta.rl.training.optim_step(session_id, learning_rate=...)`. +7. Repeat; always `client.beta.rl.sessions.stop(session_id)` in a `finally`. + +## High-Signal Rules + +- **RL calls return operations, not results.** `sample` / `forward_backward` / `optim_step` each return an + operation handle — poll it (status `TRAINING_OPERATION_STATUS_*`) until complete before using the output. +- **Rewards are computed out-of-band, never inside the RL call.** Sampling returns tokens + logprobs; you + score them yourself and feed advantages back in `forward_backward`. +- **The RL SDK is synchronous; the sandbox SDK is async.** Bridge them with a single `asyncio.run(...)` per + batch that fans out sandboxes with `asyncio.gather` (one sandbox per candidate), then return rewards into + the sync loop. Do not create a sandbox per token or per RL call. +- **Center advantages within the group.** GRPO subtracts the per-prompt group-mean reward as the baseline; + a group where every sample gets the same reward contributes zero gradient. +- **GRPO sample schema** (see [references/grpo-loop.md](references/grpo-loop.md)): `model_input.chunks[]. + encoded_text.tokens`, and `loss_inputs.{loss_mask, target_tokens, grpo_inputs.{advantages, + generator_logprobs}}`. Mask the prompt tokens out of the loss; only response tokens carry advantage. +- **Pre-build the reward environment as a sandbox snapshot.** Bake the task harness, dependencies, and + tests into a snapshot once, then create ephemeral sandboxes from it per rollout — don't `pip install` + inside every rollout. + +## Status + +The RL training API (`client.beta.rl`) is in **beta** and is **not yet in the public `together` package** +(public `together` exposes `beta.clusters` / `beta.jig` only). It also requires a service-specific +`base_url`. This skill targets the public SDK surface; the reward half (`together-sandbox`) is usable today. +Until the public RL release ships, install the preview SDK per the RL API docs and pass the RL `base_url` +explicitly. + +## Resource Map + +- **GRPO loop mechanics**: [references/grpo-loop.md](references/grpo-loop.md) +- **Sandbox reward integration**: [references/sandbox-rewards.md](references/sandbox-rewards.md) +- **End-to-end script**: [scripts/grpo_sandbox_reward.py](scripts/grpo_sandbox_reward.py) +- **Sandbox API** (hand-off): the `together-sandbox` skill + +## Official Docs + +- Together RL training API — `/rl/training-sessions` (create, `:forward-backward`, `:optim-step`, `:stop`) +- [Together Sandbox SDK](https://github.com/togethercomputer/together-sandbox/blob/main/docs/python-sdk.md) +- Reference loop: [rl-cookbook/grpo_train.py](https://github.com/togethercomputer/rl-cookbook/blob/main/grpo_train.py) diff --git a/skills/together-rl/agents/openai.yaml b/skills/together-rl/agents/openai.yaml new file mode 100644 index 0000000..bd45086 --- /dev/null +++ b/skills/together-rl/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Together RL" + short_description: "GRPO reinforcement-learning post-training with sandboxed rewards" + default_prompt: "Use $together-rl to run a GRPO training loop on Together AI with rewards computed by executing code in a sandbox." diff --git a/skills/together-rl/references/grpo-loop.md b/skills/together-rl/references/grpo-loop.md new file mode 100644 index 0000000..7c596d5 --- /dev/null +++ b/skills/together-rl/references/grpo-loop.md @@ -0,0 +1,135 @@ +# GRPO Training Loop Reference + +## Contents + +- [Overview](#overview) +- [Session Lifecycle](#session-lifecycle) +- [Operations and Polling](#operations-and-polling) +- [The Three Training Operations](#the-three-training-operations) +- [GRPO Sample Schema](#grpo-sample-schema) +- [Advantages](#advantages) +- [Batch Padding](#batch-padding) + +## Overview + +GRPO (Group Relative Policy Optimization) post-training on Together runs as a loop over the low-level RL +training API. Each step: sample a group of completions per prompt, score them into rewards, center the +rewards within each group to get advantages, then push one `forward_backward` + `optim_step`. + +This doc covers the **training mechanics** (the request/response shapes you assemble). It does not restate +the full API — for authoritative request/response definitions see the RL training API docs and the +`/rl/training-sessions` family in the OpenAPI spec. The canonical reference implementation is +[`rl-cookbook/grpo_train.py`](https://github.com/togethercomputer/rl-cookbook/blob/main/grpo_train.py); +the only thing this skill changes is **where the reward comes from** (a sandbox — see +[sandbox-rewards.md](sandbox-rewards.md)). + +## Session Lifecycle + +```python +session = client.beta.rl.sessions.create(base_model="meta-llama/Llama-3.2-1B-Instruct") +session_id = session.session_id or getattr(session, "id", None) +# Poll client.beta.rl.sessions.retrieve(session_id).status until: +# TRAINING_SESSION_STATUS_RUNNING +# ... run training steps ... +client.beta.rl.sessions.stop(session_id) # always, in a finally block +``` + +A session holds the live model weights on the training service. Create once, run many steps, stop when +done (stopping releases the GPU — orphaned sessions keep billing). + +## Operations and Polling + +`sample`, `forward_backward`, and `optim_step` are **asynchronous operations**: each returns an operation +handle, not a finished result. Poll until complete before reading the output. The cookbook's helpers in +`together.lib.beta.rl.operations` (`retrieve_operation`, `get_operation_status`, `is_operation_complete`, +`get_operation_error`) wrap this: + +```python +def wait_for_operation(client, *, session_id, operation, kind, timeout, interval): + deadline = time.monotonic() + timeout + while True: + operation = rl_ops.retrieve_operation(client, session_id=session_id, operation=operation, kind=kind) + if rl_ops.is_operation_complete(operation): + if rl_ops.get_operation_status(operation) == "TRAINING_OPERATION_STATUS_FAILED": + raise RuntimeError(rl_ops.get_operation_error(operation)) + return operation + if time.monotonic() >= deadline: + raise TimeoutError(f"operation {kind} timed out") + time.sleep(interval) +``` + +## The Three Training Operations + +```python +# 1. Sample a group of completions for one prompt (num_samples = group size) +sample_op = client.beta.rl.training.sample( + session_id=session_id, + prompt={"chunks": [{"encoded_text": {"tokens": prompt_tokens}}]}, + num_samples=group_size, + sampling_params={"max_tokens": max_sample_tokens}, +) +result = wait_for_operation(...) # result.output.sequences[i].tokens / .logprobs + +# 2. One gradient pass over the assembled, advantage-weighted batch +fb_op = client.beta.rl.training.forward_backward( + session_id=session_id, + loss={"type": "LOSS_TYPE_GRPO", + "grpo_params": {"beta": 0.0, "agg_type": "GRPO_LOSS_AGGREGATION_TYPE_TOKEN_MEAN"}}, + samples=samples, +) +loss = wait_for_operation(...).output.loss + +# 3. Optimizer step +opt_op = client.beta.rl.training.optim_step(session_id=session_id, learning_rate=8e-5) +wait_for_operation(...) +``` + +## GRPO Sample Schema + +Each training sample concatenates prompt + response tokens, masks the prompt out of the loss, and attaches +the per-token advantage and the generator's logprobs: + +```python +def build_grpo_sample(prompt_tokens, response_tokens, response_logprobs, advantage, max_model_length): + available = max(max_model_length - len(prompt_tokens), 0) + response_tokens = response_tokens[:available] + response_logprobs = response_logprobs[:available] + if not response_tokens: + return None + model_tokens = prompt_tokens + response_tokens + loss_mask = [0] * len(prompt_tokens) + [1] * len(response_tokens) # train on response only + targets = model_tokens[1:] + [0] # next-token targets + advantages = [0.0] * len(prompt_tokens) + [advantage] * len(response_tokens) + logprobs = [0.0] * len(prompt_tokens) + [float(x) for x in response_logprobs] + return { + "model_input": {"chunks": [{"encoded_text": {"tokens": model_tokens}}]}, + "loss_inputs": { + "loss_mask": {"data": loss_mask, "dtype": "D_TYPE_INT64"}, + "target_tokens": {"data": targets, "dtype": "D_TYPE_INT64"}, + "grpo_inputs": { + "advantages": {"data": advantages, "dtype": "D_TYPE_FLOAT32"}, + "generator_logprobs": {"data": logprobs, "dtype": "D_TYPE_FLOAT32"}, + }, + }, + } +``` + +## Advantages + +GRPO centers rewards **within each prompt's group**: + +```python +baseline = sum(group_rewards) / len(group_rewards) +advantage = reward - baseline +``` + +A group where every sample earns the same reward has zero advantage everywhere and contributes no +gradient — reward diversity within a group is what drives learning. `forward_backward` needs a minimum +batch (the cookbook requires ≥ 8 assembled samples). + +## Batch Padding + +All samples in a `forward_backward` batch must be padded to the longest sequence. Pad `tokens`, +`loss_mask`, and `target_tokens` with `0`, and `advantages` / `generator_logprobs` with `0.0`. See +`_pad_samples` in [rl-cookbook/grpo_train.py](https://github.com/togethercomputer/rl-cookbook/blob/main/grpo_train.py) +for the exact padding (note: the integer fields are serialized as strings in the padded request). diff --git a/skills/together-rl/references/sandbox-rewards.md b/skills/together-rl/references/sandbox-rewards.md new file mode 100644 index 0000000..3fb2288 --- /dev/null +++ b/skills/together-rl/references/sandbox-rewards.md @@ -0,0 +1,89 @@ +# Sandbox Code-Execution Rewards + +## Contents + +- [Why a Sandbox](#why-a-sandbox) +- [The Integration Point](#the-integration-point) +- [Bridging Sync RL and Async Sandboxes](#bridging-sync-rl-and-async-sandboxes) +- [Scoring One Candidate](#scoring-one-candidate) +- [Batch Reward Collection](#batch-reward-collection) +- [Pre-build the Environment](#pre-build-the-environment) + +## Why a Sandbox + +For coding and agentic RL the reward is not a string match — it is the result of **running** the model's +output: execute the candidate solution against a test suite and use pass/fail (or a parsed score) as the +reward. That code is untrusted and must run in isolation. This skill computes rewards with the +`together-sandbox` SDK; the sandbox API itself (snapshots, sandboxes, exec, files) is owned by the +**`together-sandbox` skill** — read it for the full surface and its +[Python SDK docs](https://github.com/togethercomputer/together-sandbox/blob/main/docs/python-sdk.md). The +reward patterns below mirror that skill's `references/rl-patterns.md` (batch fan-out, reward collection). + +## The Integration Point + +In the GRPO loop ([grpo-loop.md](grpo-loop.md)), `training.sample` returns token sequences. Decode each to +text, then — instead of scoring locally — score by executing in a sandbox: + +``` +sample group -> decode tokens -> [SANDBOX] write candidate, run tests, read reward -> advantages -> forward_backward +``` + +## Bridging Sync RL and Async Sandboxes + +The RL training SDK is synchronous; the `together-sandbox` SDK is async. Run the whole batch's rewards in +**one** `asyncio.run(...)` call per training step, fanning out sandboxes concurrently: + +```python +rewards = asyncio.run(compute_rewards(snapshot_alias, candidates, test_code)) +``` + +Never open an event loop per candidate or per RL call — collect all of a step's candidates first, then +fan out once. + +## Scoring One Candidate + +```python +async def score_candidate(sdk, snapshot_alias, candidate_code, test_code) -> float: + model = await sdk.sandboxes.create(snapshot_alias=snapshot_alias, ephemeral=True) + sandbox = await sdk.sandboxes.start(model.id) + try: + # Sandboxes have no DNS by default; set it as the first exec if tests need network. + await sandbox.execs.exec("bash", ["-c", 'echo "nameserver 1.1.1.1" > /etc/resolv.conf']) + await sandbox.files.create("/workspace/solution.py", candidate_code) + await sandbox.files.create("/workspace/test_solution.py", test_code) + result = await sandbox.execs.exec( + "bash", ["-c", "cd /workspace && python -m pytest -q test_solution.py"] + ) + # execs.exec returns {"exit_code": int, "output": str} + return 1.0 if result["exit_code"] == 0 else 0.0 + finally: + await sdk.sandboxes.shutdown(sandbox.id) # ephemeral; clean up every rollout +``` + +## Batch Reward Collection + +Fan out one sandbox per candidate for the whole step, then return rewards in the original order so they +line up with the RL samples: + +```python +async def compute_rewards(snapshot_alias, candidates, test_code) -> list[float]: + async with TogetherSandbox() as sdk: + results = await asyncio.gather( + *[score_candidate(sdk, snapshot_alias, c, test_code) for c in candidates], + return_exceptions=True, + ) + # A crashed sandbox should score 0.0, not abort the whole step. + return [r if isinstance(r, float) else 0.0 for r in results] +``` + +Order matters: `gather` preserves input order, so `rewards[i]` corresponds to `candidates[i]`. A failed or +timed-out sandbox is scored `0.0` so one bad rollout never invalidates the batch — but watch the failure +rate, since GRPO needs reward diversity *within* each group to learn. + +## Pre-build the Environment + +Bake the task harness, dependencies, and (optionally) the tests into a **snapshot once**, then create +ephemeral sandboxes from it per rollout. Building the image inside every rollout wastes minutes per step +and serializes the batch behind `pip install`. See the `together-sandbox` skill's golden-image pattern +(`references/rl-patterns.md`, Pattern 1) for `snapshots.create` + `hibernate` to capture a reusable image, +and reference it here by `snapshot_alias`. diff --git a/skills/together-rl/scripts/grpo_sandbox_reward.py b/skills/together-rl/scripts/grpo_sandbox_reward.py new file mode 100644 index 0000000..d515afc --- /dev/null +++ b/skills/together-rl/scripts/grpo_sandbox_reward.py @@ -0,0 +1,345 @@ +#!/usr/bin/env python3 +""" +Together RL — GRPO post-training with sandboxed code-execution rewards. + +Drives a GRPO training loop on the Together RL training API and scores each +sampled candidate by EXECUTING it inside an isolated together-sandbox: write the +candidate solution into a sandbox, run its test suite, and use pass/fail as the +reward. + +This is the integration the rl-cookbook's grpo_train.py leaves out — that demo +scores GSM8K answers with a local string match. Real coding/agentic RL has to +*run* the model's output to score it, which must happen in a sandbox. + +Status: + The RL training API (client.beta.rl) is in BETA and is not yet in the public + `together` package. This script targets that public surface; until the RL + release ships, the `client.beta.rl` calls (and the operations helper import) + will not resolve. The reward half (compute_rewards) runs today against + together-sandbox v1.12.0 — see the `together-sandbox` skill. + +Requires: + pip install together # client.beta.rl (beta; pending public release) + pip install together-sandbox # reward execution (on PyPI today) + pip install transformers # tokenizer + export TOGETHER_API_KEY=... + +Run (once the RL API is public): + python grpo_sandbox_reward.py \ + --base-url https:// \ + --model meta-llama/Llama-3.2-1B-Instruct \ + --snapshot-alias coding-rl-env@v1 + +References: + GRPO loop mechanics -> references/grpo-loop.md + Sandbox reward pattern -> references/sandbox-rewards.md + Sandbox SDK -> the `together-sandbox` skill +""" + +from __future__ import annotations + +import argparse +import asyncio +import re +import time +from dataclasses import dataclass +from typing import Any + +from transformers import AutoTokenizer + +from together import Together + +# NOTE: requires the `together` release where the RL training API ships (beta). +# In the rl-cookbook this helper lives at `together.lib.beta.rl.operations`; the +# public import path may differ once published — adjust to the shipped SDK. +from together.lib.beta.rl import operations as rl_ops # type: ignore + +from together_sandbox import TogetherSandbox + + +# --------------------------------------------------------------------------- # +# Toy coding tasks. Each rollout asks the model to implement `solution`; the +# reward is whether the bundled test passes in a sandbox. Swap this for your +# real dataset (e.g. SWE-bench / Terminal-Bench task specs). +# --------------------------------------------------------------------------- # +TASKS: list[dict[str, str]] = [ + { + "prompt": "Write a Python function `solution(n)` that returns the n-th Fibonacci number " + "(0-indexed, solution(0)=0, solution(1)=1). Reply with a ```python code block.", + "test": "from solution import solution\n" + "def test_fib():\n" + " assert [solution(i) for i in range(7)] == [0, 1, 1, 2, 3, 5, 8]\n", + }, + { + "prompt": "Write a Python function `solution(s)` that returns True if string `s` is a " + "palindrome ignoring case and non-alphanumerics. Reply with a ```python code block.", + "test": "from solution import solution\n" + "def test_palindrome():\n" + " assert solution('A man, a plan, a canal: Panama')\n" + " assert not solution('hello')\n", + }, +] + + +@dataclass +class Config: + base_url: str + model: str + snapshot_alias: str + group_size: int + max_steps: int + max_prompt_length: int + max_model_length: int + max_sample_tokens: int + learning_rate: float + poll_interval: float + timeout: float + + +def parse_args() -> Config: + p = argparse.ArgumentParser(description="GRPO with sandboxed code-execution rewards") + p.add_argument("--base-url", required=True, help="RL service base URL") + p.add_argument("--model", required=True, help="Base model for the training session") + p.add_argument("--snapshot-alias", required=True, + help="together-sandbox snapshot alias with python + pytest preinstalled") + p.add_argument("--group-size", type=int, default=8, help="Samples per prompt (GRPO group size)") + p.add_argument("--max-steps", type=int, default=20) + p.add_argument("--max-prompt-length", type=int, default=512) + p.add_argument("--max-model-length", type=int, default=2048) + p.add_argument("--max-sample-tokens", type=int, default=512) + p.add_argument("--learning-rate", type=float, default=8e-5) + p.add_argument("--poll-interval", type=float, default=1.0) + p.add_argument("--timeout", type=float, default=7200.0) + a = p.parse_args() + return Config( + base_url=a.base_url, model=a.model, snapshot_alias=a.snapshot_alias, + group_size=a.group_size, max_steps=a.max_steps, + max_prompt_length=a.max_prompt_length, max_model_length=a.max_model_length, + max_sample_tokens=a.max_sample_tokens, learning_rate=a.learning_rate, + poll_interval=a.poll_interval, timeout=a.timeout, + ) + + +# --------------------------------------------------------------------------- # +# Tokenization +# --------------------------------------------------------------------------- # +def build_prompt_tokens(tokenizer, question: str, *, max_length: int) -> list[int]: + messages = [{"role": "user", "content": question}] + try: + tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True) + except Exception: + tokens = tokenizer.encode(question, add_special_tokens=False) + return list(tokens)[:max_length] + + +def extract_code(text: str) -> str: + """Pull the first ```python ... ``` block; fall back to the whole response.""" + m = re.search(r"```(?:python)?\s*(.*?)```", text, re.DOTALL) + return (m.group(1) if m else text).strip() + + +# --------------------------------------------------------------------------- # +# GRPO sample assembly (see references/grpo-loop.md) +# --------------------------------------------------------------------------- # +def build_grpo_sample(prompt_tokens, response_tokens, response_logprobs, advantage, max_model_length): + available = max(max_model_length - len(prompt_tokens), 0) + response_tokens = response_tokens[:available] + response_logprobs = response_logprobs[:available] + if not response_tokens: + return None + model_tokens = prompt_tokens + response_tokens + loss_mask = [0] * len(prompt_tokens) + [1] * len(response_tokens) + targets = model_tokens[1:] + [0] + advantages = [0.0] * len(prompt_tokens) + [advantage] * len(response_tokens) + logprobs = [0.0] * len(prompt_tokens) + [float(x) for x in response_logprobs] + return { + "model_input": {"chunks": [{"encoded_text": {"tokens": model_tokens}}]}, + "loss_inputs": { + "loss_mask": {"data": loss_mask, "dtype": "D_TYPE_INT64"}, + "target_tokens": {"data": targets, "dtype": "D_TYPE_INT64"}, + "grpo_inputs": { + "advantages": {"data": advantages, "dtype": "D_TYPE_FLOAT32"}, + "generator_logprobs": {"data": logprobs, "dtype": "D_TYPE_FLOAT32"}, + }, + }, + } + + +def pad_samples(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: + if not samples: + return samples + max_len = max(len(s["model_input"]["chunks"][0]["encoded_text"]["tokens"]) for s in samples) + + def _pad_int(values, fill=0): + return [str(v) for v in (list(values) + [fill] * (max_len - len(values)))[:max_len]] + + def _pad_float(values, fill=0.0): + return (list(values) + [fill] * (max_len - len(values)))[:max_len] + + padded = [] + for s in samples: + tokens = s["model_input"]["chunks"][0]["encoded_text"]["tokens"] + li = s["loss_inputs"] + padded.append({ + "model_input": {"chunks": [{"encoded_text": {"tokens": _pad_int(tokens)}}]}, + "loss_inputs": { + "loss_mask": {"data": _pad_int(li["loss_mask"]["data"]), "dtype": "D_TYPE_INT64"}, + "target_tokens": {"data": _pad_int(li["target_tokens"]["data"]), "dtype": "D_TYPE_INT64"}, + "grpo_inputs": { + "advantages": {"data": _pad_float(li["grpo_inputs"]["advantages"]["data"]), + "dtype": "D_TYPE_FLOAT32"}, + "generator_logprobs": {"data": _pad_float(li["grpo_inputs"]["generator_logprobs"]["data"]), + "dtype": "D_TYPE_FLOAT32"}, + }, + }, + }) + return padded + + +def mean(xs: list[float]) -> float: + return sum(xs) / len(xs) if xs else 0.0 + + +# --------------------------------------------------------------------------- # +# RL operation polling +# --------------------------------------------------------------------------- # +def wait_for_operation(client, *, session_id, operation, kind, timeout, interval): + deadline = time.monotonic() + timeout + while True: + operation = rl_ops.retrieve_operation(client, session_id=session_id, operation=operation, kind=kind) + if rl_ops.is_operation_complete(operation): + if rl_ops.get_operation_status(operation) == "TRAINING_OPERATION_STATUS_FAILED": + raise RuntimeError(f"operation {kind} failed: {rl_ops.get_operation_error(operation)}") + return operation + if time.monotonic() >= deadline: + raise TimeoutError(f"operation {kind} timed out") + time.sleep(interval) + + +def wait_for_session_running(client, *, session_id, timeout, interval): + deadline = time.monotonic() + timeout + while True: + status = client.beta.rl.sessions.retrieve(session_id).status + print(f"[session:{session_id}] status={status}") + if status == "TRAINING_SESSION_STATUS_RUNNING": + return + if time.monotonic() >= deadline: + raise TimeoutError("session did not reach RUNNING") + time.sleep(interval) + + +# --------------------------------------------------------------------------- # +# Sandbox code-execution reward (see references/sandbox-rewards.md) +# --------------------------------------------------------------------------- # +async def score_candidate(sdk, snapshot_alias, candidate_code, test_code) -> float: + model = await sdk.sandboxes.create(snapshot_alias=snapshot_alias, ephemeral=True) + sandbox = await sdk.sandboxes.start(model.id) + try: + await sandbox.files.create("/workspace/solution.py", candidate_code) + await sandbox.files.create("/workspace/test_solution.py", test_code) + result = await sandbox.execs.exec( + "bash", ["-c", "cd /workspace && python -m pytest -q test_solution.py"] + ) + # execs.exec returns {"exit_code": int, "output": str} + return 1.0 if result["exit_code"] == 0 else 0.0 + finally: + await sdk.sandboxes.shutdown(sandbox.id) + + +async def compute_rewards(snapshot_alias, candidates, test_code) -> list[float]: + async with TogetherSandbox() as sdk: + results = await asyncio.gather( + *[score_candidate(sdk, snapshot_alias, c, test_code) for c in candidates], + return_exceptions=True, + ) + return [r if isinstance(r, float) else 0.0 for r in results] + + +# --------------------------------------------------------------------------- # +# Training loop +# --------------------------------------------------------------------------- # +def main() -> None: + cfg = parse_args() + tokenizer = AutoTokenizer.from_pretrained(cfg.model) + + client = Together(base_url=cfg.base_url) # reads TOGETHER_API_KEY from env + session = client.beta.rl.sessions.create(base_model=cfg.model) + session_id = session.session_id or getattr(session, "id", None) + if not session_id: + raise RuntimeError(f"no session_id in response: {session}") + print(f"Created session {session_id}") + wait_for_session_running(client, session_id=session_id, timeout=cfg.timeout, interval=cfg.poll_interval) + + try: + for step in range(cfg.max_steps): + task = TASKS[step % len(TASKS)] + prompt_tokens = build_prompt_tokens(tokenizer, task["prompt"], max_length=cfg.max_prompt_length) + + # 1. Sample a group of completions for this prompt. + sample_op = client.beta.rl.training.sample( + session_id=session_id, + prompt={"chunks": [{"encoded_text": {"tokens": prompt_tokens}}]}, + num_samples=cfg.group_size, + sampling_params={"max_tokens": cfg.max_sample_tokens}, + ) + sample_res = wait_for_operation( + client, session_id=session_id, operation=sample_op, kind="sample", + timeout=cfg.timeout, interval=cfg.poll_interval, + ) + sequences = (sample_res.output.sequences if sample_res.output else None) or [] + if not sequences: + raise RuntimeError("sample returned no sequences") + + # 2. Decode each completion and score it by running its tests in a sandbox. + decoded, seq_tokens_list, seq_logprobs_list = [], [], [] + for seq in sequences: + toks = [int(t) for t in (seq.tokens or [])] + seq_tokens_list.append(toks) + seq_logprobs_list.append([float(x) for x in (seq.logprobs or [])]) + decoded.append(tokenizer.decode(toks, skip_special_tokens=True)) + candidates = [extract_code(d) for d in decoded] + rewards = asyncio.run(compute_rewards(cfg.snapshot_alias, candidates, task["test"])) + + # 3. GRPO advantages (center within the group), then assemble the batch. + baseline = mean(rewards) + samples, kept_rewards = [], [] + for toks, lps, reward in zip(seq_tokens_list, seq_logprobs_list, rewards): + sample = build_grpo_sample( + prompt_tokens=prompt_tokens, response_tokens=toks, response_logprobs=lps, + advantage=reward - baseline, max_model_length=cfg.max_model_length, + ) + if sample: + samples.append(sample) + kept_rewards.append(reward) + if len(samples) < 8: + raise RuntimeError(f"need >= 8 samples for forward_backward, got {len(samples)}") + samples = pad_samples(samples) + + # 4. forward_backward + optim_step. + fb_op = client.beta.rl.training.forward_backward( + session_id=session_id, + loss={"type": "LOSS_TYPE_GRPO", + "grpo_params": {"beta": 0.0, "agg_type": "GRPO_LOSS_AGGREGATION_TYPE_TOKEN_MEAN"}}, + samples=samples, + ) + fb_res = wait_for_operation( + client, session_id=session_id, operation=fb_op, kind="forward_backward", + timeout=cfg.timeout, interval=cfg.poll_interval, + ) + loss = fb_res.output.loss if fb_res.output else None + + opt_op = client.beta.rl.training.optim_step(session_id=session_id, learning_rate=cfg.learning_rate) + wait_for_operation( + client, session_id=session_id, operation=opt_op, kind="optim_step", + timeout=cfg.timeout, interval=cfg.poll_interval, + ) + + print(f"[step {step}] loss={loss} reward_mean={mean(kept_rewards):.3f} " + f"pass={sum(1 for r in kept_rewards if r > 0)}/{len(kept_rewards)}") + finally: + client.beta.rl.sessions.stop(session_id) + print("session stopped") + + +if __name__ == "__main__": + main()