research/design-F3-rl-sagemaker.md

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:

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):

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:

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)

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/<job>/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/<run_id>/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 TimeoutErrors. 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).