Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 5 additions & 25 deletions docs/site/scripts/gen_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,7 @@
reward.md
backends/
slime/
rollout.md
traces.md
rewards.md
gateway.md
sglang_token_ids.md
verify_sglang_token_ids.md
runner.md
"""
from __future__ import annotations

Expand Down Expand Up @@ -98,32 +93,17 @@ class ModuleSpec:
include=("RewardFunction",),
),
# --- Backends › slime --- #
# Only the public user surface: the SlimeRunner entry point and the
# two SGLang patch scripts. integration/rollout.py and
# integration/rewards.py are load-bearing plugin paths but users
# don't import them — SlimeRunner wires them into the job.
# Only the public user surface: the SlimeRunner entry point.
# integration/rollout.py and integration/rewards.py are load-bearing
# plugin paths but users don't import them — SlimeRunner wires them
# into the job.
ModuleSpec(
dotted_paths=("agentcore_rl_toolkit.backends.slime.runner",),
out_path="backends/slime/runner.md",
title="SlimeRunner",
description="One Python entry point for slime-backed training.",
include=("SlimeRunner",),
),
# Both patches consolidated onto one page.
# sglang_token_ids exposes apply_patch (the public patch entry point).
# verify_sglang_token_ids exposes main (the CLI contract, invoked via
# `python -m ...verify_sglang_token_ids`). Everything else is
# low-level monkey-patch machinery / internal check helpers.
ModuleSpec(
dotted_paths=(
"agentcore_rl_toolkit.backends.slime.patches.sglang_token_ids",
"agentcore_rl_toolkit.backends.slime.patches.verify_sglang_token_ids",
),
out_path="backends/slime/patches.md",
title="SGLang patches",
description="Apply + verify the SGLang token-IDs patch used by the slime backend.",
include=("apply_patch", "main"),
),
]


Expand Down
68 changes: 0 additions & 68 deletions docs/site/src/content/docs/api/backends/slime/patches.md

This file was deleted.

26 changes: 26 additions & 0 deletions docs/site/src/content/docs/api/backends/slime/runner.md
Original file line number Diff line number Diff line change
Expand Up @@ -41,13 +41,22 @@ SlimeRunner(
slime_dir: str = '/root/slime',
megatron_dir: str = '/root/Megatron-LM',

# --- Optional: CUDA toolchain pinning ---
cuda_home: str | None = None,

# --- Optional: ACR / toolkit (forwarded to slime via custom-config yaml) ---
model_id: str = 'default',
acr_timeout: int = 900,
acr_tps_limit: int = 25,
max_concurrent: int = 100,
max_pool_connections: int = 100,
gateway_port: int = 9090,
reward_postprocessing: str = 'grpo',
sglang_tool_call_parser: str = 'qwen',
sglang_reasoning_parser: str | None = None,
cumulative_token_mode: bool = False,
renderer_family: str = 'auto',
gateway_log_level: str = 'warning',

# --- Optional: training hyperparameters ---
rollout_batch_size: int = 32,
Expand All @@ -60,6 +69,7 @@ SlimeRunner(
weight_decay: float = 0.1,
adam_beta2: float = 0.98,
sglang_mem_fraction_static: float = 0.7,
sglang_context_length: int | None = None,
max_tokens_per_gpu: int = 9216,

# --- Wandb (opt-in; no defaults injected if unset) ---
Expand Down Expand Up @@ -95,6 +105,10 @@ via ``ray job submit``. Streams stdout/stderr to the parent process.

- `agent_runtime_arn` *(str)*

- `cuda_home` *(str | None)*

- `cumulative_token_mode` *(bool)*

- `data_path` *(str)*

- `eps_clip` *(float)*
Expand All @@ -105,12 +119,16 @@ via ``ray job submit``. Streams stdout/stderr to the parent process.

- `extra_flags` *(list[str])*

- `gateway_log_level` *(str)*

- `gateway_port` *(int)*

- `lr` *(float)*

- `max_concurrent` *(int)*

- `max_pool_connections` *(int)*

- `max_tokens_per_gpu` *(int)*

- `megatron_dir` *(str)*
Expand All @@ -125,6 +143,8 @@ via ``ray job submit``. Streams stdout/stderr to the parent process.

- `num_gpus` *(int)*

- `renderer_family` *(str)*

- `reward_postprocessing` *(str)*

- `rollout_batch_size` *(int)*
Expand All @@ -137,8 +157,14 @@ via ``ray job submit``. Streams stdout/stderr to the parent process.

- `s3_bucket` *(str)*

- `sglang_context_length` *(int | None)*

- `sglang_mem_fraction_static` *(float)*

- `sglang_reasoning_parser` *(str | None)*

- `sglang_tool_call_parser` *(str)*

- `slime_dir` *(str)*

- `tp_size` *(int)*
Expand Down
75 changes: 39 additions & 36 deletions docs/site/src/content/docs/guides/slime-backend-setup.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,18 @@ description: Train an AgentCore Runtime-deployed agent with the SLIME training b

This doc describes how to train an AgentCore Runtime-deployed agent with the
[slime](https://github.com/THUDM/slime) training backend. The public
user surface is exactly two things: the
user surface is the
[`SlimeRunner`](/agentcore-rl-toolkit/api/backends/slime/runner/)
class (for launching training) and the two SGLang
[patch scripts](/agentcore-rl-toolkit/api/backends/slime/patches/)
(applied once to the SGLang install).
class for launching training.

For known issues (Megatron-LM regression on 32B, norm-epsilon
mismatch, etc.) see
For known issues (e.g. the norm-epsilon mismatch on
Qwen2.5-32B-Instruct) see
[slime troubleshooting](/agentcore-rl-toolkit/troubleshooting/slime/).

## Prerequisites

- Hardware and CUDA requirements: see
[slime's README](https://github.com/THUDM/slime#installation) and the
[slime docker README](https://github.com/THUDM/slime/blob/main/docker/README.md)
for tested GPU configurations per model size.
- Python 3.10+ and [`uv`](https://docs.astral.sh/uv/).
- A GPU cluster with **CUDA>=12.9** installed.
- Python 3.12+ and [`uv`](https://docs.astral.sh/uv/).
- AWS credentials with permission to invoke an AgentCore Runtime and
read/write an S3 bucket.
- An AgentCore Runtime deployment of your agent — follow the
Expand All @@ -31,44 +26,52 @@ mismatch, etc.) see
- An S3 bucket for rollout result delivery — required as the
`s3_bucket` argument on `SlimeRunner` below.

## slime environment
## Installation

Follow
[slime's own installation docs](https://github.com/THUDM/slime#installation)
— either the container path (`slimerl/slime:latest`) or a bare-metal
install. Everything below runs inside this environment.
Choose **one** of the two paths below to install slime, then install
the toolkit into the same environment.

Inside the `slimerl/slime:latest` container, slime and Megatron-LM
ship pre-installed at `/root/slime` and `/root/Megatron-LM` — use
those paths for `slime_dir` / `megatron_dir` on `SlimeRunner`. For a
bare-metal install, point at wherever you cloned slime + Megatron-LM.
### Option A: Official slime docker

## Install the toolkit + apply patches
Follow
[slime's own installation docs](https://github.com/THUDM/slime#installation)
and use the container image (`slimerl/slime:latest`). Inside the
container, slime and Megatron-LM ship pre-installed at `/root/slime`
and `/root/Megatron-LM` — use those paths for `slime_dir` /
`megatron_dir` on `SlimeRunner`.

Inside the slime environment:
Install the toolkit with the slime-backend extras inside the container:

```bash
# From a clone of this repo
cd /path/to/agentcore-rl-toolkit

# Install the toolkit plus the slime-backend extras
uv pip install -e ".[slime]"
```

Then apply the SGLang `token_ids` patch — it adds
`prompt_token_ids` / `token_ids` fields to chat completion responses
so the gateway can capture RL training trace data. The patch is idempotent:
:::note
We have only tested against official slime at commit
[`fa3c990`](https://github.com/THUDM/slime/commit/fa3c990af6f18efd3fd9922698bf4bf4048d1263).
:::

```bash
python -m agentcore_rl_toolkit.backends.slime.patches.sglang_token_ids
### Option B: Bare-metal install script

# Verify the patch round-trips under greedy decoding (any HF checkpoint
# works; Qwen2.5-0.5B-Instruct is the fastest to download + load)
python -m agentcore_rl_toolkit.backends.slime.patches.verify_sglang_token_ids \
--model-path /path/to/Qwen2.5-0.5B-Instruct
# Expect: "OK: 4/4 checks passed"
Install slime and its heavyweight dependency stack (Megatron-LM,
Transformer Engine, Apex, flash-attn, sglang, torch_memory_saver)
with the provided script, which clones slime + Megatron-LM into the
current directory and applies slime's official patches. Run it inside
your activated python environment.

:::note
Both CUDA 12.9 and CUDA 13 are supported. The following commands assume CUDA 13.0 is installed at `/usr/local/cuda-13.0`; adjust `CUDA_HOME` and `cu13` if yours differs.
:::

```bash
uv pip install -e ".[slime]"
export CUDA_HOME=/usr/local/cuda-13.0
bash src/agentcore_rl_toolkit/backends/slime/scripts/install_slime.sh cu13
```

Point `slime_dir` / `megatron_dir` on `SlimeRunner` at the `slime` and
`Megatron-LM` directories the script cloned.

## Prepare data

The training dataset is a JSONL file where each line is one rollout
Expand Down
54 changes: 7 additions & 47 deletions docs/site/src/content/docs/troubleshooting/slime.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,51 +14,11 @@ was validated against:
| Component | Version / SHA |
|---|---|
| Instance type | 8 × NVIDIA H100 80GB HBM3 |
| CUDA | `12.9` |
| PyTorch | `2.9.1+cu129` |
| Docker image | `slimerl/slime@sha256:0100c933f1f63e7c4acdb9ec575e769839d59de4a648551e09e3fe0e7885631b` (built 2026-04-28) |
| slime | commit `f3e7bd7f3091d3be05c20977eefb31a785d6221d` (2026-04-28) |
| SGLang | `v0.5.9` |
| Megatron-LM | commit `3714d81d418c9f1bca4594fc35f9e8289f652862` ⚠ see note |

:::caution[Megatron-LM pin]
The image bundles Megatron-LM at `1dcf0dafa` (~500 commits ahead of
slime's stable pin), which breaks 32B training — see the
[LinearCrossEntropyModule issue](#linearcrossentropymodule-parallelism-error-on-32b-or-any-model-with-untied-embeddings) below.
We downgrade to `3714d81d` (slime's documented stable sha) via
`git checkout` inside `/root/Megatron-LM`. The table above reflects
the downgraded sha, not the one baked into the image.
:::

## `LinearCrossEntropyModule` parallelism error on 32B (or any model with untied embeddings)

**Symptom:** During 32B training, the Megatron actor crashes with:

```
ValueError: Cannot determine parallelism type for module 'LinearCrossEntropyModule'
at weight 'output_layer.weight'.
```

**Cause:** The Megatron-LM bundled in `slimerl/slime:latest` is
several hundred commits ahead of the sha pinned in slime's docker
README (`3714d81d`). Specifically, Megatron PR **#3226 "Reapply fix
Linear CE Fusion"** (2026-02-04) replaced `ColumnParallelLinear`
with a new `LinearCrossEntropyModule` that megatron-bridge's
`AutoMapping` doesn't recognize. Models with tied embeddings (0.5B,
3B, 7B) skip this code path; models with
`--untie-embeddings-and-output-weights` (32B and up) hit it.

**Fix:** Inside the container, pin `/root/Megatron-LM` to the stable
sha:

```bash
cd /root/Megatron-LM
# Stash any image-local patches first (can be restored later with `git stash pop`)
git stash -u -m "slime local patches"
git checkout 3714d81d418c9f1bca4594fc35f9e8289f652862
# Clear pyc caches that reference the old code
find . -name __pycache__ -type d -exec rm -rf {} + 2>/dev/null
```
| CUDA | `13.0` |
| PyTorch | `2.11.0+cu130` |
| slime | commit `fa3c990af6f18efd3fd9922698bf4bf4048d1263` |
| SGLang | `0.5.13` |
| Megatron-LM | commit `1dcf0dafa884ad52ffb243625717a3471643e087` |

## `--norm-epsilon` mismatch on Qwen2.5-32B-Instruct

Expand All @@ -68,5 +28,5 @@ find . -name __pycache__ -type d -exec rm -rf {} + 2>/dev/null

**Fix:** Edit the slime model script to `--norm-epsilon 1e-6`, or
pass an override through `SlimeRunner(extra_flags=["--norm-epsilon", "1e-6"])`.
0.5B/3B/7B Instruct variants match their base-model norm epsilons,
so this only affects 32B.
Qwen2.5-0.5B / 1.5B / 3B / 7B Instruct variants match their base-model
norm epsilons, so this only affects Qwen2.5-32B.
Loading
Loading