setup-trainer

name	setup-trainer
description	Stand up the CUDA (slime/Megatron) trainer for a stoa run, publish the first weights to R2, and produce the join card that Macs use to connect. Covers both run paths: the GSM8K quickstart (Qwen2.5-1.5B, one A100/H100, fully self-contained) and the PaperSearchQA fleet reference (LFM2.5-8B-A1B on a B200-class GPU). Use when starting a new RL post-training run with Macs as the rollout fleet.

Set up a stoa trainer

Goal: bring up the learner on one CUDA GPU, publish the base weights to R2, and start consuming rollouts. Once the first anchor is up, any number of Macs can join with the join-network skill. The trainer and every Mac meet only through one R2 bucket; they never address each other.

First, ask which run

Do not assume. Ask the user which of the two paths they're setting up, then follow only that block:

• GSM8K quickstart: Qwen2.5-1.5B-Instruct on GSM8K, one A100/H100 (40 GB+). Everything ships in the repo, so this is the runnable path: you can stand it up in an afternoon and watch reward climb. Pick this to actually run stoa.
• PaperSearchQA fleet reference: LFM2.5-8B-A1B on PaperSearchQA, a B200-class GPU plus a BM25 retriever. This reproduces our headline run: the LFM2 slime plugin ships in-repo and the trainer runs on the public slime base, so you can rerun it end to end. Only the trained checkpoint stays private (you retrain from the public base), so you reproduce the curve, not that exact checkpoint. Pick this for the full 8B run.

R2 credentials (both paths)

R2 is the only shared resource. Put the four values (R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, R2_BUCKET) in the trainer's credentials file (/root/dRL/.r2env) and pick a fresh DRL_RUN id so it doesn't collide with an existing run's prefix. Full instructions, including the cross-host DRL_RUN / DRL_N_SAMPLES / DRL_PULSE_PLANE contract, are in docs/r2-setup.md. Don't hand-write a cred blurb; point the user there.

Path A: GSM8K quickstart

Full walkthrough with GPU rental, SSH, model conversion, the held-out curve, and teardown is in envs/gsm8k/README.md. The short version:

1. Rent one CUDA GPU (an A100 or H100, 40 GB+) with the slimerl/slime:nightly-dev-20260703b image (slime + Megatron-LM + the HF-to-Megatron converter prebaked) and ~100 GB disk. SSH in.

2. Deploy the repo + creds. From your local checkout, ship the trainer-side tree to /root/dRL and copy ~/.r2env to /root/dRL/.r2env (chmod 600), as shown in envs/gsm8k/README.md.

3. Convert the base model. Download Qwen2.5-1.5B-Instruct and convert it to a Megatron torch_dist checkpoint with the exact convert_hf_to_torch_dist.py invocation in envs/gsm8k/README.md. The rotary base MUST be 1000000 (the HF rope_theta); slime's stock helper uses 10000, which would make the trainer a different model than the MLX worker and break the importance ratio.

4. Launch. From the repo root on the trainer, dry-run first to print the resolved command, then launch:

   DRL_RUN=my-gsm8k DRL_PULSE_PLANE=bf16 DRL_DRY_RUN=1 bash envs/gsm8k/run_qwen25_gsm8k_decoupled.sh
   DRL_RUN=my-gsm8k DRL_PULSE_PLANE=bf16 bash envs/gsm8k/run_qwen25_gsm8k_decoupled.sh
   

   It applies the slime patches, publishes the base model as the v1 anchor so the worker can bootstrap, starts Ray (dashboard bound to localhost), then runs the disaggregated GRPO loop (--rollout-num-gpus 0 with r2_rollout.generate_rollout, asymmetric DPPO gate, Dr.GRPO reducer, dual-clip PPO, a 0.001 ref-KL, staleness bound 3, reuse cap 2). Batch is 8 groups by 8 samples per step. Leave it running in this terminal; reward (rollout/raw_reward) climbs over a few tens of steps once a Mac joins.

The quickstart uses the bf16 weight plane end to end (the fresh-converted anchor publishes bf16) and the default group size DRL_N_SAMPLES=8, so you pass only DRL_RUN and DRL_PULSE_PLANE inline.

Path B: PaperSearchQA fleet reference

The full run record, pins, hardware tiers, and expected curve are in runs/psqa-decoupled/REPRODUCE.md. The LFM2 slime plugin ships in-repo (installed by trainer_setup_decoupled.sh) and the trainer runs on the public slime base, so this path is fully runnable. Only the trained checkpoint is private, so you retrain from the public base and reproduce the curve, not that exact checkpoint. On the trainer box (inside the slime/Megatron container, this repo at /root/dRL, creds at /root/dRL/.r2env):

1. Provision. Run the disaggregated trainer setup:

   bash envs/psqa_search/b200_provision/trainer_setup_decoupled.sh
   

   It installs the LFM2 plugin into /root/slime (STAGE B copies it from the in-repo envs/psqa_search/lfm2_slime_plugin/), applies the dRL slime patches, and pulls the base LFM2.5-8B-A1B (HF dir plus pre-converted torch_dist) from R2 bases/lfm25/. The b200_provision/provision.sh sibling is the all-in-one colocated variant that also builds retrieval and datasets on the box; the decoupled fleet keeps those on the Macs, so trainer_setup_decoupled.sh is the one to run here.

2. Set the recipe. Copy the launch env and edit it:

   cp envs/psqa_search/launch_fleet.env.example /root/launch_fleet.env
   $EDITOR /root/launch_fleet.env   # set DRL_RUN, optional DRL_WANDB_*
   

   Confirm DRL_N_SAMPLES=8 (the group size every worker must match, the contract that defines the run), DRL_PULSE_PLANE=int8g64 (workers declare the same plane and both sides fail closed on mismatch), and DRL_KV_MODE=paged_prefix (the fleet preflight requires it). Fixed inside the script (not this env): --save-interval 4, dual-clip --eps-clip 10000 --eps-clip-high 19 --eps-clip-c 20, the dppo_gate + drgrpo_reducer paths, and no --save-hf.

3. Preflight, then launch. Source the env, dry-run through the preflight (must PASS), then launch:

   set -a; . /root/launch_fleet.env; set +a
   DRL_DRY_RUN=1 bash envs/psqa_search/run_lfm25_psqa_fleet.sh \
       | DRL_PREFLIGHT_PROFILE=fleet python3 envs/psqa_search/preflight_check.py
   bash envs/psqa_search/run_lfm25_psqa_fleet.sh
   

   The real launch applies the slime patches, publishes the base LFM2.5 as the v1 anchor, starts Ray, then blocks consuming worker rollouts. Watch DRL_RUN's W&B run (train/*, rollout/*, worker/*) plus a held-out eval on a separate GPU; the held-out curve, not train reward, is the primary signal.

The join card

Once v1 is published, hand joiners these values out of band (secrets can't go in the repo): R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, R2_BUCKET, the DRL_RUN id, the group size DRL_N_SAMPLES, and the weight plane DRL_PULSE_PLANE. DRL_N_SAMPLES is the hard contract: a worker that joins with a different n has every group dropped, so state it explicitly.

• GSM8K: the quickstart worker passes DRL_RUN and DRL_PULSE_PLANE=bf16 inline and uses the default DRL_N_SAMPLES=8, so the card is just the run id plus the R2 creds.
• Fleet: DRL_PULSE_PLANE=int8g64 and DRL_N_SAMPLES=8. Light-tier Macs (15 to 23 GB) also need a shared retriever URL, since they can't hold the BM25 index locally, so expose one reachable endpoint.