Train with Slime

Full actor–critic integration is in development

We are developing an integration that trains both the actor and the step-wise critic, then uses the updated critic to guide rollout collection.

This tutorial trains the actor and exports critic targets. It uses a constant critic prediction of 0.5. To learn from those targets and guide later collection rounds, you need a separate critic-training loop.

Slime manages training and serves the policy through its inference workers. FOCUS decides which partial solutions to continue and turns the resulting tree into actor advantages and critic targets.

This tutorial connects the MATH harness to Slime’s Megatron training backend. Each edge continues the same solution for up to 256 tokens. The integration returns actor segments to Slime and saves critic targets for a value-model training loop.

Slime owns the training loop and calls FOCUS through hooks named on its command line. focus.integrations.slime provides those hooks. You supply the same task function used in Train with Tinker: collector_for(policy, problem), which builds a collector for one problem.

Before you begin

Use a Slime Megatron environment; step 7 gives a complete launch for Qwen3-4B-Instruct-2507 on two H100 GPUs. Reuse math_training.py from Train with Tinker, step 1: the search, the starting critic, and collector_for are the same for both trainers. Step 7 installs FOCUS and the answer verifier in the worker environment.

1. Write the hook module

The Slime hook module defines the generation and evaluation hooks and the directory for critic records:

train_slime.py

import os

from focus.integrations import JsonlDirectory
from focus.integrations.slime import make_evaluate, make_generate
from math_training import collector_for

# Critic records leave Slime's loop: one JSONL file per collected prompt.
critic_sink = JsonlDirectory(os.environ.get("FOCUS_CRITIC_OUTPUT", "focus-critic-targets"))

# The model is used without tools; Segments reads each reply as a token chunk.
generate = make_generate(collector_for, profile="token_chunk", critic_sink=critic_sink)
evaluate = make_evaluate(collector_for, profile="token_chunk")

make_generate() builds Slime’s custom generate function and make_evaluate() its evaluation function, both from collector_for. profile="token_chunk" resolves the served model’s tokenizer for use without tools. Segments in the shared make_agent() reads each reply as an exact 256-token segment.

JsonlDirectory writes each prompt’s critic records to a separate JSONL file, preventing concurrent workers from interleaving lines. Each rollout worker reads the directory from FOCUS_CRITIC_OUTPUT when it imports this module at startup.

2. What happens when Slime calls generate

For each prompt, the hook:

1. Builds a SlimePolicy for the model behind Slime’s SGLang router. The router address and hf_checkpoint come from Slime’s args. Loading the tokenizer once per checkpoint keeps FOCUS’s prompt rendering matched to the served model.
2. Wraps Slime’s sample as a Problem. Slime’s --input-key and --label-key fill its prompt and label; the sample’s index becomes its id.
3. Calls collector_for(policy, problem) and collects one tree with that fresh collector. The environment is built per prompt, so concurrent problems keep their own reference answers. The reference solution is used only for grading.
4. Sends the critic records to the critic sink and returns the actor records as Slime samples.

DirectBranchTd computes segment TD advantages, while the shared tree values provide direct-branch backup targets for the critic. A collection with failed expansions or critic predictions raises an error unless you pass allow_partial=True. The hook uses each collector once inside async with, following the ownership settings in collector_for. Slime owns the HTTP client used by the policy.

3. How actor segments become Slime samples

focus.integrations.slime.to_sample() converts each actor record. It keeps:

Field	Meaning
tokens	Full conditioning prefix followed by the generated segment
response_length	Number of tokens from the first generated token to the end
loss_mask	Response-aligned mask selecting generated tokens
rollout_log_probs	The segment’s recorded generation probabilities
reward	The prepared FOCUS actor advantage
rollout_id	The identity shared by all segments collected for this prompt

The conditioning prefix stays outside the response loss. Conditioning tokens that arrive inside the response, such as tool results in a tool-using task, stay in the response with loss mask zero. All sibling segments keep the source sample’s rollout_id, so Slime groups them as one rollout during batch scheduling. The conversion follows Slime’s Sample fields.

Here reward carries the prepared actor weight: positive values reinforce the segment and negative values discourage it. MathEnv supplied the correctness outcomes used to compute that weight.

4. Use the advantage and loss hooks

Three hooks in focus.integrations.slime carry these weights through Slime’s training pipeline. Name them directly on the command line in step 7:

1. preserve_advantages() passes the prepared weights through reward post-processing unchanged.
2. advantages() repeats each segment’s weight over its response positions, using tensors aligned with the recorded log probabilities. It overrides the estimator Slime’s argument parser selected and refuses advantage normalization.
3. sum_policy_tokens() sums the masked token losses. Slime’s trainer applies its rollout-batch reduction across prompts.

The policy loss combines the clipped learner-to-behavior probability ratio with the segment’s fixed advantage, as in the paper. Prompt tokens remain masked. Retain the behavior probabilities after the model updates: they describe the policy that collected the segment.

Slime performs the optimizer step and synchronizes the actor weights to its inference workers. Start the next FOCUS collection round after that synchronization, with a fixed actor snapshot for the duration of collection.

5. Evaluate with a separate function

For each Slime evaluation dataset, set custom_generate_function_path to train_slime.evaluate. That hook runs one unbranched solution through the collector’s agent and returns its outcome reward, 0 or 1 here, rather than a signed advantage. It uses the same 256-token segments and 4,096-token horizon.

The evaluation sample uses the last recorded generation’s input and output tokens. It does not rebuild a chat conversation from the segments: that could insert turn markers or remove thinking, breaking alignment with the recorded log probabilities.

6. Publish the critic for the next collection round

collector_for begins with Critic.constant(0.5). To use a learned critic, return make_collector(policy, problem.label, critic=critic) from it, as in the Tinker tutorial.

Each exported critic record contains a checkpoint context and its backup target $\bar G$. Train the value model to predict that target, rendering contexts the same way as collection. The Tinker tutorial’s critic section explains the regression objective shared by both integrations.

Slime updates the actor using these hooks. Your value-model trainer consumes the exported critic records and publishes the next critic snapshot. Keep that critic fixed during collection, then refresh it before the next round.

7. Assemble the files and launch Slime

This launch uses Qwen3-4B-Instruct-2507 on one Linux node with two 80 GB H100 GPUs. Training and generation take turns on the same GPUs. Use the Slime environment from its setup guide, with Slime at revision 18868a79b0b083a093798a94ed0f8cb4f71b9c14 and Megatron-LM at /root/Megatron-LM.

Run the commands below in Bash inside that environment. They use /root/slime for Slime, /root/FOCUS for your FOCUS checkout, and the persistent /data mount for inputs and outputs. Save these files in /data/focus-hooks/. Open them in your editor to follow the examples and adjust the configuration:

• math_quickstart.py: environment, parser, and exact-prefix policy.
• math_training.py: the search and the collector for one problem.
• train_slime.py: the generate and evaluation hooks.
• prepare_math.py: data selection and JSONL output.
• slime_runtime.yaml: the worker environment.
• launch_slime.sh: the complete training command.

Install dependencies and download the model

Install the tutorial dependencies and download the model:

cd /root/slime
python -m pip install -e . --no-deps
python -m pip install -e '/root/FOCUS'
python -m pip install 'math-verify[antlr4_13_2]==0.9.0' datasets huggingface_hub

export FOCUS_HOOKS=/data/focus-hooks
export FOCUS_MODEL=/data/Qwen3-4B-Instruct-2507
export FOCUS_DATA=/data/math-algebra-64.jsonl
export FOCUS_INITIAL=/data/Qwen3-4B-Instruct-2507-torch-dist
export FOCUS_RUN=/data/focus-slime-math
export PYTHONPATH="$FOCUS_HOOKS:/root/slime:/root/Megatron-LM${PYTHONPATH:+:$PYTHONPATH}"
export CUDA_DEVICE_MAX_CONNECTIONS=1

mkdir -p "$FOCUS_RUN"
hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir "$FOCUS_MODEL"

Prepare a small training set

The supplied prepare_math.py lives beside math_quickstart.py. It selects 64 problems from the MATH algebra training split and saves their original problem and solution fields as JSONL.

The selection has two checks: the answer verifier must be able to parse the reference answer, and the prompt must fit within 4,096 tokens. With the 8,192-token context configured below, that leaves 4,096 tokens for the generated solution. Use the served model’s tokenizer and the same system instruction as MathEnv’s agent:

Problem selection and JSONL writing

prepare_math.py

def write_dataset(rows, tokenizer, output, *, count=64):
    if count < 1:
        raise ValueError("Choose at least one problem.")
    selected = []
    for row in rows:
        messages = [
            {"role": "system", "content": r"Solve the problem step by step. Put your final answer in \boxed{}."},
            {"role": "user", "content": row["problem"]},
        ]
        prefix = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
        if len(prefix) > 4096 or not parse_answer(row["solution"]):
            continue
        selected.append({"problem": row["problem"], "solution": row["solution"]})
        if len(selected) == count:
            break

    if len(selected) < count:
        raise ValueError(f"Requested {count} usable problems, but found {len(selected)}.")
    output = Path(output)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8") as stream:
        for row in selected:
            stream.write(json.dumps(row) + "\n")
    print(f"Wrote {len(selected)} problems to {output}")

The script’s main() loads the tokenizer and dataset, then calls write_dataset(). It accepts the model and output paths as command-line arguments:

Command-line entry point

prepare_math.py

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, help="Model directory or Hugging Face model ID.")
    parser.add_argument("--output", required=True, type=Path, help="Destination JSONL file.")
    parser.add_argument("--count", type=int, default=64, help="Number of problems (default: 64).")
    args = parser.parse_args()
    if args.count < 1:
        parser.error("--count must be positive")

    from datasets import load_dataset
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    rows = load_dataset("EleutherAI/hendrycks_math", "algebra", split="train")
    write_dataset(rows, tokenizer, args.output, count=args.count)


if __name__ == "__main__":
    main()

Run it from the same Bash session:

python "$FOCUS_HOOKS/prepare_math.py" \
  --model "$FOCUS_MODEL" \
  --output "$FOCUS_DATA" \
  --count 64

The script prints the number of saved problems and their file path. If it cannot find enough usable problems, it reports an error before writing the file.

Convert the model weights

Convert the Hugging Face weights once into Slime’s distributed training checkpoint. The model configuration supplies Qwen’s layer dimensions and rotary-position settings:

source scripts/models/qwen3-4B-Instruct-2507.sh
python tools/convert_hf_to_torch_dist.py \
  "${MODEL_ARGS[@]}" \
  --hf-checkpoint "$FOCUS_MODEL" \
  --save "$FOCUS_INITIAL" \
  --seq-length 8192 \
  --max-position-embeddings 262144 \
  --bf16

Launch three training rounds

Open slime_runtime.yaml in /data/focus-hooks/. Ray reads this file to give the driver and rollout workers the Python import paths, critic-output directory, and CUDA setting they need:

slime_runtime.yaml

env_vars:
  PYTHONPATH: /data/focus-hooks:/root/slime:/root/Megatron-LM
  FOCUS_CRITIC_OUTPUT: /data/focus-slime-math/critic-targets
  CUDA_DEVICE_MAX_CONNECTIONS: "1"

Use absolute paths in this file. If you chose different directories above, update them here too; YAML does not expand the shell variables. Ray accepts this file through --runtime-env.

Review launch_slime.sh before running it. It contains the full Megatron and SGLang configuration, including the four FOCUS hooks. Edit its options for your experiment; the supplied settings use two H100 GPUs and three training rounds.

Complete training command: launch_slime.sh

launch_slime.sh

#!/usr/bin/env bash
# Submit the MATH training recipe from the Slime tutorial to an existing Ray node.
set -euo pipefail

: "${FOCUS_HOOKS:?Set FOCUS_HOOKS to the directory containing the tutorial files}"
: "${FOCUS_MODEL:?Set FOCUS_MODEL to the downloaded model directory}"
: "${FOCUS_DATA:?Set FOCUS_DATA to the prepared MATH JSONL file}"
: "${FOCUS_INITIAL:?Set FOCUS_INITIAL to the converted training checkpoint}"
: "${FOCUS_RUN:?Set FOCUS_RUN to the output directory for this experiment}"

cd /root/slime
source scripts/models/qwen3-4B-Instruct-2507.sh

ray job submit --address http://127.0.0.1:8265 \
  --runtime-env "$FOCUS_HOOKS/slime_runtime.yaml" \
  -- python /root/slime/train.py \
  "${MODEL_ARGS[@]}" \
  --train-backend megatron \
  --actor-num-nodes 1 \
  --actor-num-gpus-per-node 2 \
  --colocate \
  --hf-checkpoint "$FOCUS_MODEL" \
  --ref-load "$FOCUS_INITIAL" \
  --save "$FOCUS_RUN/checkpoints" \
  --save-interval 1 \
  --prompt-data "$FOCUS_DATA" \
  --input-key problem \
  --label-key solution \
  --num-rollout 3 \
  --rollout-batch-size 4 \
  --global-batch-size 4 \
  --n-samples-per-prompt 1 \
  --rollout-max-response-len 4096 \
  --rollout-max-context-len 8192 \
  --rollout-temperature 0.7 \
  --rollout-top-p 1.0 \
  --custom-generate-function-path train_slime.generate \
  --custom-reward-post-process-path focus.integrations.slime.preserve_advantages \
  --custom-advantage-function-path focus.integrations.slime.advantages \
  --custom-pg-loss-reducer-function-path focus.integrations.slime.sum_policy_tokens \
  --use-rollout-logprobs \
  --advantage-estimator grpo \
  --loss-type policy_loss \
  --pg-loss-type ppo \
  --eps-clip 0.2 \
  --eps-clip-high 0.2 \
  --kl-coef 0 \
  --entropy-coef 0 \
  --optimizer adam \
  --lr 1e-6 \
  --lr-decay-style constant \
  --weight-decay 0.1 \
  --adam-beta1 0.9 \
  --adam-beta2 0.98 \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 1 \
  --context-parallel-size 1 \
  --sequence-parallel \
  --use-dynamic-batch-size \
  --max-tokens-per-gpu 8192 \
  --log-probs-chunk-size 1024 \
  --recompute-granularity full \
  --recompute-method uniform \
  --recompute-num-layers 1 \
  --seq-length 8192 \
  --max-position-embeddings 262144 \
  --bf16 \
  --rollout-num-gpus-per-engine 2 \
  --sglang-mem-fraction-static 0.7 \
  --attention-dropout 0.0 \
  --hidden-dropout 0.0 \
  --accumulate-allreduce-grads-in-fp32 \
  --attention-softmax-in-fp32 \
  --attention-backend flash

Start Ray on this dedicated node, then run the script from the same Bash session so the exported model and data paths remain available:

ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 \
  --disable-usage-stats --dashboard-host 127.0.0.1 --dashboard-port 8265

bash "$FOCUS_HOOKS/launch_slime.sh"

Each round collects four problem trees and makes one optimizer update. --global-batch-size 4 counts the four rollout identities; all segments from each tree share that identity. Tensor parallelism splits the model across two GPUs, leaving one data-parallel replica. Dynamic batching packs the variable number of segments into training microbatches. --colocate enables Slime’s offload/onload cycle so the learner and inference engine share those GPUs.

The four custom-function options connect the hooks from steps 1 and 4. --n-samples-per-prompt 1 lets FOCUS control continuations within each tree. FOCUS takes its sampling settings from the agent and any per-phase override; the callback’s sampling_params argument does not override them.

Segments(CHUNK_TOKENS) in the shared make_agent() caps each request at 256 tokens. Change CHUNK_TOKENS to choose another limit. The execution service validates each response against the effective request settings; the token-segment parser extracts the tokens and finish reason. The custom advantage hook supplies FOCUS’s prepared values, overriding the estimator selected by Slime’s argument parser.

The harness formats the prompt, so the command leaves --apply-chat-template unset. It also leaves advantage normalization, per-token batch normalization, auxiliary KL loss, and dynamic reward filtering disabled to preserve the prepared actor weights.

After the command returns, learner checkpoints are in $FOCUS_RUN/checkpoints and critic records are in the directory set by FOCUS_CRITIC_OUTPUT in slime_runtime.yaml. Stop the Ray instance started above when you finish:

ray stop

For another experiment, choose a new FOCUS_RUN directory and update FOCUS_CRITIC_OUTPUT in slime_runtime.yaml to match. To resume a saved training run, edit launch_slime.sh to add --load "$FOCUS_RUN/checkpoints" and raise --num-rollout to the desired total number of rounds.

The launch follows Slime’s Qwen recipe and customization interfaces. The hook contracts and launch arguments have been checked against that revision; the two-GPU training run has not been executed here.

Keep a distributed round consistent

Slime owns the learner checkpoint and inference-worker lifecycle. Use its training checkpoint configuration to save and restore model and optimizer state. Resume from a completed round, publish those weights to the inference workers, and begin the next FOCUS collection round with the matching critic. Store exported critic files with that round’s identifier.

A FOCUS rollout checkpoint holds one agent’s context, environment, workspace, and remaining allowance. The BrowseComp-Plus walkthrough shows how to fork an in-memory compaction checkpoint. Persisting a training job requires storage for those states as well. The job must also track pending prompts, the tree scheduler, and accepted training batches. Save that information at the job level so a restart does not apply the same training batch twice.

Let Slime own distributed execution

Slime runs the rollout workers and routes generation through SGLang. Each worker constructs its FOCUS collector locally and calls SlimePolicy directly. FOCUS manages task state, branching, and learning signals; Slime owns worker placement, model serving, and weight synchronization.

Keep the actor and critic fixed while a tree is being collected. Finish the active collections before replacing weights on a shared inference endpoint. FOCUS records the probabilities returned with each generation; the trainer tracks which optimizer step produced the batch.

Within one worker, collectors can share ExecutionServices(pool=AdmissionPool(capacity=8)) to limit simultaneous generation and critic requests. In collector_for, pass Execution(services=services) to the agent and build the collector with own=False. Close the shared services after all collections finish. This limit applies only to that process; Slime controls capacity across workers.

Complete source

Slime hook module

train_slime.py

"""Slime hooks for the MATH tutorial.

Point Slime's generate and evaluation functions at this module, and its reward,
advantage, and loss-reducer hooks at ``focus.integrations.slime``.
"""
import os

from focus.integrations import JsonlDirectory
from focus.integrations.slime import make_evaluate, make_generate
from math_training import collector_for

# Critic records leave Slime's loop: one JSONL file per collected prompt.
critic_sink = JsonlDirectory(os.environ.get("FOCUS_CRITIC_OUTPUT", "focus-critic-targets"))

# The model is used without tools; Segments reads each reply as a token chunk.
generate = make_generate(collector_for, profile="token_chunk", critic_sink=critic_sink)
evaluate = make_evaluate(collector_for, profile="token_chunk")