SageMaker for Vision Language Action: Distributed LeRobot Training
We show how to customize a SageMaker Training job to enable non-standard distributed training use cases, using the example of a distributed Hugging Face LeRobot training.
Amazon SageMaker training jobs are usually presented as a managed way to run PyTorch or Hugging Face estimators. They can do more than that. Under the hood, a training job is an on-demand, EFA-connected GPU cluster that starts when your job starts and goes away when it finishes. By reading the cluster layout, you can use a broader variety of distributed frameworks.
In this post I'll show how I used that idea to run multi-node training of LeRobot policies (ACT, Gr00t, π₀ and others) with Hugging Face
accelerate instead of the standard torchrun distributed training approach. It takes three techniques:- Extending an off-the-shelf Deep Learning Container (DLC) with a
requirements.txt, with no custom image to build. - Reading the cluster topology from SageMaker's environment variables.
- Launching training from a subprocess instead of using the estimator's default launcher.
Scaling from 1 GPU to many nodes then comes down to changing
instance_type and instance_count.Distributed training in five minutes
If you already know DDP, skip ahead. Otherwise, these are the terms you'll see in the logs:
- World size, rank, local rank. World size is the total number of training processes. Rank is a process's global ID. Local rank is its ID within one machine, and usually picks the GPU.
- Data parallelism (DDP). Every GPU holds a full copy of the model and trains on a different slice of each batch. After the backward pass, gradients are averaged across GPUs so all copies stay identical.
- All-reduce. The collective operation that does that averaging. Every GPU contributes its gradients and every GPU receives the combined result. NCCL is the NVIDIA library that implements it for GPUs.
- NVLink / NVSwitch. Fast GPU-to-GPU links inside one machine.
- EFA (Elastic Fabric Adapter). AWS's network interface for HPC and ML. It bypasses the OS network stack and uses the Scalable Reliable Datagram protocol to move data between machines at high bandwidth and low latency.
- RDMA. Remote direct memory access. On supported instances, GPUs can exchange data across the network without staging it through host memory.
Why this matters: in multi-node data-parallel training, every step ends with an all-reduce across the network. On a slow interconnect, expensive GPUs spend much of their time waiting for gradients. EFA makes sure GPUs get those gradients quickly.
Larger models need more than DDP, such as tensor, pipeline, or fully sharded data parallelism. Most robot-learning policies are small enough to fit the full model on modern GPUs. SmolVLA is around 450M parameters, and π₀ is about 3B. Both fit on one modern GPU, so plain data parallelism, with the model replicated on every GPU, is the right fit. The job then is to keep every GPU busy and keep gradient sync cheap.
Hugging Face LeRobot
LeRobot is Hugging Face's open-source library for real-world robot learning, and it's the framework the rest of this post trains with. It's worth a moment on why it's the right foundation before wiring it into SageMaker.
- A common data format. LeRobot defines
LeRobotDataset, a standardized layout for robot data: synchronized multi-camera video, proprioceptive state, and actions stored as Parquet plus MP4, with the metadata needed to index, slice, and visualize episodes. A dataset recorded on an SO-101 arm and one from an Aloha bimanual rig look the same to the training loop, so swapping data does not mean rewriting a loader. - Policies share one interface. ACT, Diffusion Policy, SmolVLA, π₀, and the other implementations all expose the same configuration and training surface. Changing architecture is a
--policy.typeflag, not a new codebase, which is exactly what lets the launcher in this post stay policy-agnostic. - One training entry point. The
lerobot-trainCLI is the single command for every policy and dataset, and its multi-GPU path is driven by Hugging Faceaccelerate. That is the hook the three techniques below build on: if we can handaccelerate launcha correct cluster topology, we get distributed LeRobot training for free, without forking the trainer. - The Hub is built in. Datasets and trained policies are versioned artifacts on the Hugging Face Hub . You can stream a dataset straight into training, resume from a checkpoint pushed to a repo, and share the result — all by
repo_id. This pairs naturally with SageMaker's S3 integration, as the streaming section later shows. - An active ecosystem. LeRobot is widely used and actively maintained, with a growing catalog of public datasets and pretrained policies to start from rather than collecting data from scratch.
In short, LeRobot gives us a portable data format, interchangeable policies, and a single
accelerate-based launcher. SageMaker gives us the on-demand, EFA-connected cluster to run that launcher at scale. The rest of this post is about joining the two.What SageMaker handles for you
Running multi-node training on raw EC2 means a lot of setup: a placement group, an EFA-enabled security group, EFA drivers, a matching NCCL and AWS OFI NCCL plugin, and consistent hostnames. With a SageMaker training job:
- The cluster appears and disappears with the job. You set
instance_countand you're billed only while the job runs. - The DLC is already tuned for the hardware. The SageMaker PyTorch DLC ships with CUDA, NCCL, the AWS OFI NCCL plugin, and EFA support. A wrapper script also fixes up hostnames so peers can reach each other.
- Cluster information is injected as environment variables. More on this below.
- Managed Spot Training can cut cost considerably, and anything written to
/opt/ml/checkpointsis synced to S3 so a job can resume after an interruption.
The built-in way to use all this is
distribution={"torch_distributed": {"enabled": True}}, which wraps your script in torchrun. That works well for PyTorch scripts. LeRobot's multi-GPU training, however, is driven by accelerate launch, and that's where the three techniques come in.Technique 1: Extend the DLC with a requirements.txt
Sagemaker AI supports custom Docker containers . But creating your own container does require a few extra steps of building the container, and storing it on ECR . A custom container is the right approach for extensive modifications, or when you're using non-publicly available packages. But in the case where you just need a few more packages on top of what is available in the existing AWS Deep Learning Containers , there's a simpler approach that allows you to install packages at training time. When you launch a training job, the estimator checks your source directory for a requirements.txt file. If one is present, the listed packages are installed on the existing container.
1
2
3
4
5
6
7
# src/requirements.txt
lerobot==0.4.0
accelerate>=1.0,<2
tensorboard>=2.16
wandb>=0.17
hf_transfer>=0.1.6
huggingface_hub>=0.24Two practical notes:
- Don't add
torch. The DLC's PyTorch is built against a matching NCCL and EFA plugin. Replacing it can quietly break the fast interconnect path. - Use a Python entry point. The toolkit only auto-installs
requirements.txtfor Python entry points. A shell-script entry point also doesn't receive hyperparameters as positional arguments, so you'd end up reading them from environment variables anyway. A thin Python launcher avoids both problems.
The result is that you get the DLC's tuned base image, plus your libraries, with no image build, no ECR repository, and no pipeline to maintain.
Technique 2: Read the cluster layout from SageMaker's environment variables
SageMaker gives every container a set of
SM_* environment variables that describe the job. Three are enough to derive everything accelerate launch needs:| Variable | Meaning |
|---|---|
SM_HOSTS | JSON list of all machines in the cluster, e.g. ["algo-1","algo-2"] |
SM_CURRENT_HOST | This current machine's name, e.g. algo-2 |
SM_NUM_GPUS | Number of GPUs in this container |
Hyperparameters set on the estimator arrive as JSON in
SM_HPS.1
2
3
4
5
6
7
8
9
10
11
import json, os
hosts = sorted(json.loads(os.environ["SM_HOSTS"]))
current = os.environ["SM_CURRENT_HOST"]
num_nodes = len(hosts)
node_rank = hosts.index(current) # machine_rank
gpus_per_node = int(os.environ["SM_NUM_GPUS"])
num_processes = num_nodes * gpus_per_node # world size
master_addr = hosts[0] # first host is the rendezvous point
master_port = 29500Because the same formulas work for one GPU, one node with eight GPUs, or four nodes with eight GPUs each, the launcher has no per-shape branches. Nothing changes in your code when the cluster grows.
Technique 3: Launch from a subprocess instead of the default launcher
This is the key step. With
torch_distributed enabled, SageMaker runs your entry point under torchrun, which starts one copy per GPU. If each copy then calls accelerate launch, you have nested launchers fighting over the same port.Instead, set
distribution={}. SageMaker then runs your entry point once per node, and your script is the only launcher:1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
estimator = PyTorch(
entry_point="sm_train_entry.py",
source_dir="src", # contains requirements.txt
role=role,
instance_type="ml.p5.48xlarge", # 8 GPUs per node, EFA-capable
instance_count=4, # 32 GPUs total
framework_version="2.10", # or image_uri= to pin a specific DLC
py_version="py313",
distribution={}, # we are the launcher
hyperparameters={
"dataset_repo_id": "myorg/so101_pickplace",
"policy_type": "smolvla",
"steps": 200000,
"batch_size": 32,
},
checkpoint_s3_uri="s3://my-bucket/lerobot/run01/ckpts",
checkpoint_local_path="/opt/ml/checkpoints",
use_spot_instances=True,
max_run=24 * 3600,
max_wait=36 * 3600,
)The entry point reads the topology and runs
accelerate launch as a child process:1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import shutil, signal, subprocess, sys
cmd = [
shutil.which("accelerate"), "launch",
"--multi_gpu",
f"--num_machines={num_nodes}",
f"--num_processes={num_processes}",
f"--machine_rank={node_rank}",
f"--main_process_ip={master_addr}",
f"--main_process_port={master_port}",
"--rdzv_backend=c10d",
"--same_network",
"--mixed_precision=bf16",
shutil.which("lerobot-train"),
f"--dataset.repo_id={hps['dataset_repo_id']}",
f"--policy.type={hps['policy_type']}",
f"--output_dir=/opt/ml/checkpoints/{hps['job_name']}",
f"--steps={hps['steps']}",
f"--batch_size={hps['batch_size']}",
]
proc = subprocess.Popen(cmd)
# Forward SIGTERM (sent on spot interruption) so Accelerate can shut down cleanly.
for sig in (signal.SIGTERM, signal.SIGINT):
signal.signal(sig, lambda s, _f: proc.send_signal(s))
rc = proc.wait()
sys.exit(rc)I use a subprocess rather than
os.execv for two reasons:- Clean spot shutdown. SageMaker sends
SIGTERMwhen a managed-spot instance is reclaimed. The parent process relays it to the child so Accelerate can flush logs and checkpoints. - Post-run work. Once training succeeds, the node with rank 0 can copy the final model into
/opt/ml/model, which SageMaker uploads to S3 asmodel.tar.gz.
Every node runs the same script and computes its own
machine_rank, so there's no per-node configuration.Point
--output_dir at /opt/ml/checkpoints, and checkpoints sync to S3 automatically. On a restart, the launcher looks for the latest checkpoint there and adds LeRobot's resume flags, so interrupted spot jobs pick up where they left off.Streaming data instead of copying it first
By default, SageMaker copies an S3 dataset onto the instance's disk before training starts. For robotics datasets full of video, that can mean a long wait on every launch. There are two ways to avoid it.
Stream from the Hugging Face Hub. LeRobot can read datasets directly from the Hub as it trains:
1
--dataset.repo_id=myorg/so101_pickplace --dataset.streaming=trueNothing is downloaded up front, so this suits datasets too large for local disk. Throughput is bounded by Hub bandwidth, so it's best for experiments and moderate scale.
Stream from S3 with FastFile mode. Register the dataset as a SageMaker input channel with
input_mode="FastFile". S3 objects appear as local files and are fetched on demand as the training process reads them:1
2
3
4
5
6
7
8
from sagemaker.inputs import TrainingInput
estimator.fit({
"dataset": TrainingInput(
s3_data="s3://my-bucket/datasets/",
input_mode="FastFile",
)
})The entry point finds the mount path in
SM_CHANNEL_DATASET and passes it to LeRobot with --dataset.root=.... Training starts immediately, which works well for sequential reads of LeRobot's Parquet and MP4 files.If you run many experiments against the same large dataset, or need heavy random access, FSx for Lustre linked to S3 is a good third option, though it requires VPC configuration.
Verify EFA is actually being used
Don't assume the fast path is active. A multi-node job that quietly falls back to TCP sockets will still run, only slowly. Before scaling up, run a small NCCL all-reduce probe on two nodes using the same estimator setup, with
NCCL_DEBUG=INFO enabled, and look in the logs for:1
NCCL INFO NET/OFI Selected Provider is efaIf you see
NET/Socket instead, the job is on TCP. Common causes are an instance type that doesn't support EFA, or a non-DLC container. The probe is cheap and saves you from paying for a fast cluster that isn't running fast.Why this approach is useful
Using a training job as a general-purpose cluster gives you the best of both worlds:
- Flexibility. Any launcher works: Accelerate here, but equally DeepSpeed, a custom
torchrunsetup, or a different framework entirely. You aren't limited to the estimator types SageMaker ships. - Managed infrastructure. You still get on-demand multi-node clusters, EFA and tuned DLCs, S3 integration, checkpoint syncing, and managed Spot Training.
- A single scaling knob. The same code runs on one GPU for a smoke test and on dozens of GPUs for a full run, because the launcher derives the topology at runtime.
Get started
To try this end to end, start from a working single-GPU LeRobot training run and then apply the three techniques to scale it out:
- A runnable training example. The LeRobot
lerobot-trainguide walks through the exact CLI used in Technique 3 —--dataset.repo_id,--policy.type,--output_dir, and the resume flags. Thehuggingface/lerobotrepository has the full source and additional training examples . - A dataset to train on.
lerobot/svla_so101_pickplaceis a public SO-101 pick-and-place dataset inLeRobotDatasetformat — the same one LeRobot uses in its own SmolVLA examples. It pairs directly with thesmolvlapolicy from the estimator snippet above, and you can preview its episodes in the Dataset Visualizer before downloading anything.
A good first milestone is to train on a single GPU locally:
1
2
3
4
5
6
7
lerobot-train \
--dataset.repo_id=lerobot/svla_so101_pickplace \
--policy.type=smolvla \
--output_dir=outputs/train/smolvla_pickplace \
--job_name=smolvla_pickplace \
--policy.device=cuda \
--steps=20000Once that runs, point the SageMaker estimator from Technique 3 at the same
dataset_repo_id and policy_type, set distribution={}, and raise instance_count. The same command then runs across an EFA-connected, multi-node cluster with checkpoints syncing to S3 — no code changes, just a larger instance_count.Enjoyed reading this content? Let the author know!
Your likes, comments, shares, and saves help creators reach more builders.
Loading recommendations
Loading article