# F3 — The RL System on AWS, Runnable NOW (SageMaker us-west-2) **Status:** design, runnable-today. **Account:** 386931836011, **region:** us-west-2, **role:** Admin (Isengard). **Live facts verified 2026-06-09** (see "Live AWS findings" below). Grounds in the deep-research report (`research/notes/final_report_socratic-mcts-swe-worldmodel-8f6dea.md` §9 SageMaker path / hybrid, §10 phased plan) and the actual repo code (`SageMakerExecutor`, `ComposerReplicationTrainer`, `ObjectStoreAllReduce`, `replica_entrypoint`, `examples/gsm8k_grpo`). --- ## 0. TL;DR (committed) 1. **For the first runnable GPU smoke RIGHT NOW: a single SageMaker Training Job, `ml.g5.2xlarge`, BYO-container extended from the AWS PyTorch DLC.** This is NOT the `SageMakerExecutor` N-replica path — it is plain GRPO on one GPU. The executor's multi-replica DiLoCo rendezvous is the *next* step, not the smoke. Reason: the smoke's job is to prove the trainer + reward + vLLM rollout works on a real GPU at minimum cost and zero quota friction. 2. **GRPO rollout = vLLM colocated in the training container** (`use_vllm=True, vllm_mode="colocate"`). TRL 1.5's default is colocate; it runs vLLM in the *same process* sharing the training GPU at `vllm_gpu_memory_utilization=0.3`. No separate inference endpoint for the smoke. The `server` mode (`trl vllm-serve`) and VeRL's `AsyncServer` are the scale answer for tool-heavy agentic rollouts later (report §8) — not for a 0.5B GSM8K smoke. 3. **Platform decision:** Training Jobs for the bursty smoke and periodic small-model runs (this facet); **HyperPod (attached to EKS)** for the long, resilience-bound inner GRPO loop (report §9). Both share the identical S3 `ObjectStoreAllReduce` rendezvous, so a run moves between them with zero trainer/loss/DiLoCo change. 4. **The `SageMakerExecutor` (already built, mock-tested) drives N independent single-instance Training Jobs**, each tagged `REPLICA_RANK=i`/`WORLD_SIZE=N` via the `Environment` map, all pointed at one `s3://.../rendezvous/` prefix. It is the bursty-fallback DiLoCo backend. To make it run live we need a built+pushed container, real `role_arn`/`image_uri`/`output_s3_path`, and a non-zero quota for N concurrent training jobs. --- ## 1. Live AWS findings (verified 2026-06-09, this account/region) | Fact | Value | Consequence | |---|---|---| | Caller identity | `arn:aws:sts::386931836011:assumed-role/Admin/baladita-Isengard` | Admin — can create roles, push ECR, run training jobs. | | SageMaker default bucket (us-west-2) | `amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d` | Use as the rendezvous + output bucket — already covered by `AmazonSageMakerFullAccess`. | | Existing exec roles | `AmazonSageMaker-ExecutionRole-20250725T133247` (and ...20241223T...) | `role_arn = arn:aws:iam::386931836011:role/service-role/AmazonSageMaker-ExecutionRole-20250725T133247`. | | Exec role policies | `AmazonSageMakerFullAccess` + custom `AmazonSageMaker-ExecutionPolicy-...` | FullAccess grants S3 to buckets named `*SageMaker*`/`*sagemaker*` — so the rendezvous bucket MUST be the SageMaker bucket above (or a `sagemaker-*` bucket) or you must attach an explicit S3 policy. **This is the IAM gotcha the executor docstring flags.** | | `ml.g5.2xlarge for training job usage` | **1.0** (non-zero!) | Single-replica g5 smoke runs IMMEDIATELY, no quota request. | | `ml.g5.2xlarge for spot training job usage` | **1.0** | Spot smoke also available (70% cheaper). | | `ml.g5.12xlarge for training job usage` | **1.0** | One 4×A10G box available for a 7B run later. | | `ml.g6.2xlarge for training job usage` | **0.0** | g6 (L4) needs a Service Quotas increase first — prefer g5 for the smoke. | | g5.2xlarge EC2 offering | us-west-2a/b/c | Capacity exists across AZs. | | Already present | `sagemaker-dynamo-on-eks-hyperpod-*` bucket | Confirms HyperPod-on-EKS has been used here — the report's §9 hybrid is live-reachable. | | boto3 / sagemaker SDK locally | NOT installed | `pip install -e .[aws]` + `pip install sagemaker` on the launch host (laptop/Studio), not in the repo's hard deps. | **The single most important runnable-now fact:** g5.2xlarge training-job quota is already 1 — the smoke needs no quota ticket. (Default for these GPU types is 0; this account has been bumped to 1.) --- ## 2. Training Jobs vs HyperPod vs EKS — when each (report §9, §10) - **SageMaker Training Jobs (THIS facet, bursty inner loop / smoke).** Ephemeral, pay-per-second, `boto3.create_training_job`, zero persistent cluster. Right for: the first GPU smoke, periodic/smaller-model runs, the `SageMakerExecutor` DiLoCo fallback. re:Post guidance: Training Jobs fit *periodic / smaller-model / pay-per-use*. The 28-day max runtime and per-job cold-start (instance provisioning ~3-6 min) are acceptable for bursty work. Warm pools (`KeepAlivePeriodInSeconds`) cut cold-start on repeated launches — but note this account's `g5 training warm pool usage` quota is 0, so warm pools need a quota bump. - **SageMaker HyperPod attached to EKS (long resilience-bound inner loop).** Report §9: HyperPod maps 1-to-1 to an EKS control plane (one EKS cluster = one HyperPod node-group in a VPC), with auto-detect-and-replace of faulty accelerators and PyTorch job auto-resume. Right for: continuous/large-model/persistent multi-day RL where a node failure on a Training Job would lose the run. The `sagemaker-dynamo-on-eks-hyperpod-*` bucket shows this is already exercised here. **"Use HyperPod for the inner loop" does NOT mean leaving EKS** — it is a node-group swap on the same control plane. Build target: the future `EKSExecutor` targets both Karpenter GPU nodes and HyperPod nodes transparently. - **Plain EKS (primary for everything else — report §8).** Outer MCTS/sandbox/dataset loop, vLLM RayService rollout groups, gVisor/Kata sandbox pods, Argo controller. The inner GRPO trainer is the one piece that swaps between a Karpenter p5/g6e NodePool and a HyperPod node-group. **Decision for F3:** SageMaker Training Jobs now (smoke + `SageMakerExecutor` DiLoCo fallback); HyperPod-on-EKS later for the long inner run. Same S3 rendezvous throughout. --- ## 3. The first runnable smoke: Qwen2.5-0.5B GRPO on GSM8K, single g5.2xlarge Training Job ### 3.1 Shape One Training Job, `InstanceCount=1`, `ml.g5.2xlarge` (1× A10G, 24 GB). GRPO with vLLM **colocated** in the training container. This is the `examples/gsm8k_grpo/run.py` recipe lifted from CPU to one real GPU, with vLLM turned on. It does **not** exercise the DiLoCo rendezvous (that's §4). It proves: container builds, trainer runs on GPU, vLLM rollout works, reward fires, checkpoint lands in S3. ### 3.2 Container — BYO extended from the AWS PyTorch DLC (do NOT use the stock HF DLC) - **Base:** `763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.6.0-gpu-py312-cu124-ubuntu22.04-sagemaker` (verified present in the us-west-2 DLC registry; 763104351884 is the AWS DLC account). The DLC already has the SageMaker training toolkit, CUDA, and a working torch — so vLLM's CUDA wheels match. - **Why not the stock HF DLC (`huggingface-pytorch-training:4.49.0`)?** It pins transformers 4.49 and does NOT bundle `trl` or `vllm`; you'd be pip-installing the whole RL stack anyway. Extending the PyTorch DLC gives a clean, version-controlled layer. - **Why a prebuilt ECR image and not `source_dir`+`requirements.txt`?** Installing `vllm` + `trl` + `flash-attn` at job start over `requirements.txt` adds 5-10 min of cold-start per job and is a flaky failure surface (wheel/CUDA mismatch). Bake them into the image once, push to the account's private ECR. `source_dir` is fine for *just the training script* layered on top, but the heavy deps must be baked. `docker/Dockerfile.sagemaker`: ```dockerfile FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.6.0-gpu-py312-cu124-ubuntu22.04-sagemaker # RL stack (baked, not pip-at-startup) RUN pip install --no-cache-dir \ "trl>=0.12" "peft>=0.13" "accelerate>=1.0" "datasets>=3.0" \ "vllm" "fsspec>=2024.6" "s3fs>=2024.6" # The framework itself COPY . /opt/composer_replication RUN pip install --no-cache-dir -e "/opt/composer_replication[train,serverless]" # SageMaker invokes the image; for the smoke we use a plain GRPO entry script, # for the DiLoCo path the executor passes ContainerEntrypoint explicitly. ENV HF_HOME=/opt/ml/input/hf_cache ``` Build + push (Admin, one-time): ```bash aws ecr create-repository --repository-name composer-rl --region us-west-2 aws ecr get-login-password --region us-west-2 | docker login --username AWS \ --password-stdin 386931836011.dkr.ecr.us-west-2.amazonaws.com docker build -f docker/Dockerfile.sagemaker -t 386931836011.dkr.ecr.us-west-2.amazonaws.com/composer-rl:smoke . docker push 386931836011.dkr.ecr.us-west-2.amazonaws.com/composer-rl:smoke ``` ### 3.3 The smoke training script — `examples/gsm8k_grpo/run_sagemaker.py` A thin GPU variant of `examples/gsm8k_grpo/run.py`. Same `gsm8k_reward` (RLVR `#### NUMBER` regex), same `ComposerReplicationTrainer(alpha_sdpo=0, beta_replay=0)` (plain GRPO — channels 2/3 off). Differences from the CPU example: ```python from trl import GRPOConfig config = GRPOConfig( output_dir="/opt/ml/model", # SageMaker uploads this to OutputDataConfig.S3OutputPath per_device_train_batch_size=8, num_generations=8, max_prompt_length=512, max_completion_length=256, learning_rate=1e-5, max_steps=20, # smoke — minutes, not hours logging_steps=1, save_strategy="no", bf16=True, # A10G supports bf16 # --- the rollout path: vLLM colocated in-process on the same GPU --- use_vllm=True, vllm_mode="colocate", # TRL 1.5 default; same process, no server vllm_gpu_memory_utilization=0.3, # leave 70% for the 0.5B policy + grads + KV vllm_tensor_parallel_size=1, beta=0.04, # small KL-to-ref; or 0.0 for pure smoke report_to=[], ) ``` Read hyperparameters from `/opt/ml/input/config/hyperparameters.json` (SageMaker writes the estimator's `hyperparameters=` there) so the same script is config-driven. ### 3.4 The launch — SageMaker Python SDK `Estimator` (run from the laptop / Studio) ```python import sagemaker from sagemaker.estimator import Estimator sess = sagemaker.Session() # picks up region us-west-2 ROLE = "arn:aws:iam::386931836011:role/service-role/AmazonSageMaker-ExecutionRole-20250725T133247" IMAGE = "386931836011.dkr.ecr.us-west-2.amazonaws.com/composer-rl:smoke" BUCKET = "amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d" est = Estimator( image_uri=IMAGE, role=ROLE, instance_type="ml.g5.2xlarge", # quota = 1 (verified live) — no ticket needed instance_count=1, volume_size=100, max_run=3600, # 1h cap for the smoke output_path=f"s3://{BUCKET}/composer-rl/smoke/output", base_job_name="composer-grpo-smoke", environment={"HF_HUB_ENABLE_HF_TRANSFER": "1"}, hyperparameters={"model": "Qwen/Qwen2.5-0.5B-Instruct", "max_steps": 20}, entry_point="run_sagemaker.py", source_dir="examples/gsm8k_grpo", # script layered on the baked image keep_alive_period_in_seconds=0, # warm-pool quota is 0 in this acct; leave off # use_spot_instances=True, max_wait=7200 # optional: spot quota is also 1 ) est.fit(wait=True, logs=True) # GSM8K loads from HF inside the container ``` **Cost:** `ml.g5.2xlarge` is ~$1.52/hr on-demand in us-west-2; a 20-step 0.5B smoke is ~15-25 min ⇒ **well under $1**. On spot (quota=1) ~$0.45-0.60/hr ⇒ pennies. The CPU example proves the loop in ~60s; this proves it on a real GPU with the real vLLM rollout path, which the CPU example explicitly does not exercise. ### 3.5 Gotchas baked into the recipe - **vLLM needs HF model download at job start.** Either set `HF_HUB_ENABLE_HF_TRANSFER=1` (done) or stage the model into S3 and pass it as an input channel; for a 0.5B model the live download is fine. `EnableNetworkIsolation` MUST stay False (the executor pins this) so the container can reach `huggingface.co` and S3. - **`vllm_gpu_memory_utilization=0.3` is the load-bearing knob on a 24 GB A10G.** Too high ⇒ OOM when the policy + grads + optimizer also need the GPU; too low ⇒ tiny KV cache, slow rollout. 0.3 is the TRL/Ray reference default for a small model on one GPU. - **GSM8K = `openai/gsm8k` config `main`.** Already what the example loads. No license blocker (MIT). --- ## 4. The DiLoCo N-replica path: how `SageMakerExecutor` drives the rendezvous This is the executor that already exists and is mock-tested (`tests/test_sagemaker_executor.py` — 20+ tests covering rank-ordered handles, env injection, status mapping, cancel idempotency, partial-launch rollback). It is the bursty DiLoCo backend, distinct from the §3 smoke. ### 4.1 What it does (verified from source) - `launch_replicas(N, ...)` submits **N independent single-instance Training Jobs** (NOT one multi-instance job — that would couple replicas through SageMaker's intra-job NCCL fabric and break the "each replica syncs only through S3" design). Each job gets `Environment={"REPLICA_RANK": str(i), "WORLD_SIZE": str(N), "RENDEZVOUS_URI": s3uri}` and `ContainerEntrypoint=["python","-m","composer_replication.diloco.serverless.replica_entrypoint"]` with `ContainerArguments=["--rendezvous", s3uri, "--world-size", N, "--trainer-module", ..., "--trainer-fn", ...]`. - `replica_entrypoint.main` reads `REPLICA_RANK`, builds `ObjectStoreAllReduce(uri=s3://..., rank, world_size)`, wraps it in `MockManager`, and calls the user's `train(manager=, rank=, world_size=, **trainer_kwargs)`. The trainer wires `manager` into `make_diloco_outer_loop`; pseudo-gradients sync via `round_{NNNNNN}/rank_{RRRR}.pt` PUT-then-poll-then-mean on S3. **DiLoCo math, loss, trainer untouched.** - `poll`/`collect` map `describe_training_job.TrainingJobStatus`; `stream_logs` reads `/aws/sagemaker/TrainingJobs//algo-*`; `cancel` calls `stop_training_job` idempotently. ### 4.2 The asymmetry that makes this clean (report §8) Gang scheduling is needed for *intra-replica* FSDP NCCL but NOT for *inter-replica* DiLoCo sync — replicas rendezvous through S3, so a straggler simply blocks at the poll loop (bounded by `timeout_s=1800`) instead of deadlocking. On SageMaker, N separate jobs have no mutual network path (`supports_inter_replica_network=False`), which is exactly right. ### 4.3 What to wire to run it live (the deltas) 1. **Same baked image** from §3.2 (it already `pip install -e .[serverless]`, so `replica_entrypoint`, `s3fs`, `fsspec` are present). The executor passes `ContainerEntrypoint` explicitly, so a generic image works. 2. **Rendezvous bucket = the SageMaker default bucket** (`amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d`) so the exec role's `AmazonSageMakerFullAccess` already grants the live S3 PUT/GET the allreduce poll loop needs. Use `rendezvous_uri = "s3://amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d/composer-rl/runs//rendezvous/"`. 3. **Quota:** N concurrent jobs need `ml.g5.2xlarge for training job usage >= N`. Currently 1 ⇒ N=1 works today; for N=2-4 DiLoCo, request a Service Quotas increase (Service Quotas console → SageMaker → "ml.g5.2xlarge for training job usage"). The smoke proves the executor end-to-end at N=1 (one job, one rank — degenerate allreduce returns its own tensor), then N=2 once quota lands. 4. **Driver script** `examples/diloco_sagemaker/run.py` (~80 LOC): construct `SageMakerExecutor(role_arn=..., image_uri=..., output_s3_path=..., region="us-west-2")`, call `launch_replicas(N, entrypoint="...replica_entrypoint", entrypoint_args={"rendezvous_uri": s3uri, "trainer_module": "examples.gsm8k_grpo.diloco_train", "trainer_fn": "train", "trainer_kwargs": {...}}, gpu="A10G", timeout=3600)`, then `collect(handles)`. `gpu="A10G"` maps to `ml.g5.2xlarge` via the executor's `_GPU_INSTANCE_MAP`. --- ## 5. The GRPO rollout problem — colocated vLLM now, server/AsyncServer later TRL's `GRPOTrainer` needs a generation path each step. Three options, committed mapping: | Option | When | On SageMaker | |---|---|---| | `model.generate()` (no vLLM) | never for real runs — too slow | the CPU example uses this implicitly; fine only for the 0.5B CPU toy. | | **vLLM colocate** (`use_vllm=True, vllm_mode="colocate"`) | **the smoke + most single-GPU runs** | vLLM in the same process, shares the training GPU at `vllm_gpu_memory_utilization=0.3`. One container, one job, no endpoint. TRL 1.5 default. **This is the F3 answer.** | | vLLM server (`trl vllm-serve`) | multi-GPU where you dedicate GPUs to generation | a *second* SageMaker job or a SageMaker endpoint runs `trl vllm-serve`; the trainer job points `vllm_server_host/port` at it. Introduces inter-process comm + a network hop — only worth it when generation dominates and you have spare GPUs. | | VeRL `AsyncServer` | tool-heavy agentic tree-of-work rollouts (report §8) | the scale answer for the SWE-agent tree: async GPU-decoupled agent loop TRL lacks. A later facet; the engine should be a configurable backend, not hardcoded. | For the F3 smoke and the DiLoCo fallback, **colocate is correct and simplest**: it keeps everything in one self-contained training container, which is exactly what a single-instance SageMaker job wants. No separate inference endpoint to provision, secure, or pay for. **One subtlety the report flags (§8):** the SDPO channel (Channel 2) needs full-vocabulary *logits* (TRL-hosted, which the trainer's `_compute_sdpo_loss` does via `model(...).logits`), while Channel 3 needs only log-probs. Colocated vLLM handles *rollout generation*; the SDPO/replay logits/log-probs come from the policy forward pass in `_compute_loss`, not from vLLM. So turning on `alpha_sdpo>0` later does not change the rollout backend choice. --- ## 6. Concrete repo deltas (to make this runnable, not hand-wavy) | Path (~LOC) | What | Why | |---|---|---| | `docker/Dockerfile.sagemaker` (~15) | Extend PyTorch DLC 2.6.0-gpu-py312; bake trl+vllm+peft+accelerate+datasets+s3fs+fsspec + `pip install -e .[train,serverless]`. | The report (§9) names "a Dockerfile wrapping composer_replication" as a missing build artifact. This is it. | | `examples/gsm8k_grpo/run_sagemaker.py` (~120) | GPU+vLLM variant of `run.py`; reads `/opt/ml/input/config/hyperparameters.json`; writes to `/opt/ml/model`; `use_vllm=True, vllm_mode="colocate"`. | The runnable smoke entry. | | `examples/diloco_sagemaker/run.py` (~80) | Driver that builds `SageMakerExecutor` with the live role/image/bucket and calls `launch_replicas`/`collect` for N replicas over the S3 rendezvous. | Turns the mock-tested executor into a live driver — no executor code change needed. | | `examples/gsm8k_grpo/diloco_train.py` (~60) | A `train(manager, rank, world_size, **kw)` that wraps the GRPO trainer in `make_diloco_outer_loop(manager=...)`. | The `trainer_module:trainer_fn` the executor imports inside each replica. | | `scripts/build_and_push_ecr.sh` (~20) | ECR create-repo + login + build + push (the §3.2 commands). | One-command image publish. | | `docs/AWS_SAGEMAKER_QUICKSTART.md` (~120) | The §1 live facts + §3 estimator recipe + the quota/IAM gotchas + the spot variant. | So the next person runs it in one read. | | `pyproject.toml` `aws` extra (+1 line) | add `sagemaker>=2.200` alongside `boto3` (the SDK `Estimator` lives there; executor uses raw boto3 but the launch driver wants the SDK). | The launch host needs the SDK; currently only `boto3` is in the extra. | **Nothing in `SageMakerExecutor`, `ComposerReplicationTrainer`, `ObjectStoreAllReduce`, `replica_entrypoint`, or the loss changes.** The executor's design (N single-instance jobs, env-injected rank, S3-only rendezvous, `EnableNetworkIsolation=False`) is already correct for this environment — verified against the live IAM/quota facts. --- ## 7. Open questions / next gates - **N>1 DiLoCo quota:** `ml.g5.2xlarge for training job usage` = 1 today. N=2-4 needs a Service Quotas increase (typically minutes-to-hours for g5; not guaranteed instant). Request before the N>1 run. - **Warm pools:** `g5 training warm pool usage` quota = 0 ⇒ each job pays ~3-6 min cold-start. For the bursty DiLoCo fallback at small H (frequent re-launch) this matters; request warm-pool quota or accept the cold-start, or move the long inner loop to HyperPod (which is persistent — no per-round cold-start). - **vLLM version pin:** the smoke leaves `vllm` unpinned in the Dockerfile; pin to a version whose CUDA matches the DLC (cu124 / torch 2.6) before promoting past smoke, to avoid a silent wheel mismatch. - **HyperPod-on-EKS path:** the `sagemaker-dynamo-on-eks-hyperpod-*` bucket shows it's been used here; the future `EKSExecutor` + HyperPod node-group attach is the report's §9 recommendation for the long inner run. Out of scope for F3 (Training-Jobs facet) but the rendezvous makes the swap free. - **Spot interruption + DiLoCo:** spot g5 quota = 1; with `use_spot_instances=True` a replica can be reclaimed mid-round. The bounded `timeout_s=1800` poll means a reclaimed replica stalls its peers up to 30 min then `TimeoutError`s. For spot DiLoCo, add `save_freq` checkpointing + relaunch-on-interruption in the driver (report §10 failure modes). ## 8. References - Repo: `composer_replication/diloco/serverless/sagemaker.py` (SageMakerExecutor), `replica_entrypoint.py`, `allreduce.py` (ObjectStoreAllReduce + MockManager), `trainer/composer_trainer.py` (ComposerReplicationTrainer + `make_po_config`), `examples/gsm8k_grpo/run.py`, `tests/test_sagemaker_executor.py`. - Report §8 (EKS primary), §9 (SageMaker path + HyperPod hybrid), §10 (cost / phased plan). - AWS DLC registry us-west-2: `pytorch-training:2.6.0-gpu-py312-cu124-ubuntu22.04-sagemaker` @ 763104351884 (docs.aws.amazon.com/sagemaker/latest/dg-ecr-paths/ecr-us-west-2.html). - TRL vLLM colocate: GRPOConfig `use_vllm`/`vllm_mode`/`vllm_gpu_memory_utilization` (huggingface/trl grpo_config.py; huggingface.co/blog/vllm-colocate; Ray TRL-GRPO example). - SageMaker quotas: g5/g6/p4d training-job usage default 0 (StackOverflow 71655321; re:Post) — verified live this account: g5.2xlarge=1, g6.2xlarge=0. - HyperPod-EKS 1:1 mapping: docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-eks.html. - Live `aws sts get-caller-identity` / `aws service-quotas list-service-quotas` / `aws iam` (2026-06-09).