Training Backends

One TrainBackend interface, two implementations — Megatron-LM (TP/PP/DP/CP) and PyTorch FSDP v2 (HuggingFace-native) — plus the runtime infrastructure they share: Ray placement, weight sync, checkpointing, memory, and logging.

1 The backend contract

The trainer selects a backend from args.train_backend (ray/train_worker.py) and drives it through one interface. Each step calls the specialization hooks compute_advantages() → prepare_batch() → extract_model_output() → compute_loss().

2 backends/base.py — TrainBackend interface

Seven abstract methods define the surface each backend implements; the rest of base.py is shared machinery.

Abstract surface

Shared concretes

forward_only() (no-grad log-probs), maybe_compute_ref/teacher_log_probs() (frozen model for FSDP, weight-tag switch for Megatron), preprocess_rollout_data() (DP shard + move to CUDA), and backup/restore/switch_weights(tag) via the WeightManager.

Key attributes

3 Backends — megatron.py & fsdp.py

backends/megatron.py — MegatronBackend, the default, built on slime + rl_engine/utils/megatron.py: TP/PP/DP/CP parallelism, a CPU-offloaded distributed optimizer, packed sequences, and torch_dist checkpoints. Diffusion models take an experimental _setup_diffusion_model() path (see Algorithm Recipes). The flags it reads:

--tensor-model-parallel-size 4  --sequence-parallel          # TP; PP/CP stay 1 on a single node
--optimizer adam  --optimizer-cpu-offload  --use-precision-aware-optimizer
--use-dynamic-batch-size  --max-tokens-per-gpu 32768        # packed batches (4B; 8192 for 8B)
--recompute-granularity full                                # activation recompute
--hf-checkpoint <HF dir>  --load <torch_dist dir>  --save <ckpt root>
--ref-load / --opd-teacher-load <torch_dist>               # KL reference / frozen teacher weights

backends/fsdp.py — FSDPBackend, FSDP v2 on a 1-D data-parallel mesh, HF-native (AutoModelForCausalLM → apply_fsdp2), PyTorch-DCP checkpoints. No TP/PP; with_opd_teacher unsupported.

--train-backend fsdp                                        # select the backend (default: megatron)
--hf-checkpoint <HF dir>                                    # the only required model path
--grad-clip 1.0  --fp16                                     # bf16 default; optional fp16 / CPU offload

Both drive the same spec.compute_loss(); full flag set in the knobs below.

4 Memory: Sleep / Wake

FSDP: CPU offload via rl_engine/utils/memory.py. Megatron — three tiers:

Tier 2 is the production default on A100 80 GB.

if self._tms_available:            # Tier 1
    destroy_process_groups(); torch_memory_saver.pause()
elif self._can_megatron_offload():  # Tier 2 (default)
    ddp_model.offload_grad_buffers(); optimizer.offload_to_cpu(); destroy_process_groups()
else:                               # Tier 3
    destroy_process_groups()

5 Backend knobs

Full STACX_* catalog: Getting Started / scripts/train/swe/ENV_REFERENCE.md.

6 Runtime infrastructure

Model-agnostic plumbing shared by every recipe, in rl_engine/trainer/infra/ and rl_engine/train.py.

Single-node topology

Two Docker containers plus a host ROCK process share one 8-GPU node. Teacher and student never share GPUs; inside the student, the Megatron trainer and its SGLang engine do, via --colocate.

Ray placement & runtime env

rl_engine/utils/ray.py · rl_engine/trainer/factory.py · dagger_4b.sh

• create_placement_groups(args) re-exports from slime.ray.placement_group; build_gpu_allocations() maps slime keys to trainer roles ("policy" → pgs["actor"]; value/reference/reward groups only for multi-role recipes).
• --colocate puts the actor placement group and the rollout SGLang engine on the same 4 GPUs.

RAY_TMP_SUFFIX=$(echo "${RUN_TAG}" | md5sum | cut -c1-8)
export RAY_TMP_DIR=/tmp/r_${RAY_TMP_SUFFIX}   # e.g. /tmp/r_3f9c1a20

Process groups

rl_engine/trainer/infra/process_groups.py

ProcessGroupManager is a backend-aware query wrapper over an already-initialized torch.distributed. Megatron's _initialize_megatron() builds independent TP/PP/CP/DP groups via mpu.initialize_model_parallel(); FSDP reports tp=pp=cp=1, dp=world_size over group.WORLD.

Weight sync — sync primitive details

rl_engine/trainer/infra/inference_weight_sync.py

Checkpointing & auto-resume

rl_engine/train.py · backends/megatron.py · dagger_4b.sh

Re-running the same launch command resumes with the same weights, optimizer state, and W&B run.

• save_model() writes a torch_dist checkpoint into iter_{N:07d}/ sharded across TP ranks — ~53 GB per 4B checkpoint (bf16 + fp32 Adam state), pruned to the newest --save-max-to-keep dirs.
• Each rollout is pickled to rollout_ckpts/rollout_{N}.pt so a resumed job skips regeneration.

resume_state.json (written atomically after every save and on graceful exit; read by the launcher on relaunch):

Eval is gated behind 2 × wall_time_reserve so a run is never killed mid-eval. Defaults: WALL_TIME_LIMIT 7 d, WALL_TIME_RESERVE 10 min (SWE scripts set 1 h).

Memory management additions

Beyond sleep/wake (§4), two more mechanisms manage the colocated engine's memory:

Logging & monitoring

infra/logging.py · infra/profiler.py · train.py

init_tracking(args) initializes W&B; --wandb-resume-id reopens the same run on resume. Ray workers can't reach W&B, so the driver logs from the main process (eval goes to a secondary run).

7 Infrastructure knobs & source map

Infra knobs only — algorithm knobs (β/κ, K3, teacher sampling) live on Algorithm Recipes; full catalog in scripts/train/swe/ENV_REFERENCE.md.

Source map