Robot learning,
at a different pace.

Training requires memory for model weights, gradients, and the optimizer’s running estimates. PyTorch FSDP2 partitions this state across GPUs. Each GPU stores its own shard, gathers the weights needed to compute a model block, and receives the corresponding gradient shard after backward. The optimizer then updates the locally owned weights and state. This makes room for models and batches that would otherwise consume too much memory on one GPU.

Hybrid sharding, or HSDP, adds data parallelism to that arrangement. A group of GPUs jointly holds one copy of the model; several such groups process different batches. In the diagram, GPUs 1 and 2 form one replica and GPUs 3 and 4 form another. Matching shards synchronize their gradients before the update, so both replicas learn from the combined data.

Tensor parallelism divides computation within a layer, such as splitting a projection matrix and combining its partial outputs. Pipeline parallelism assigns successive model blocks to different GPUs and passes activations and gradients between their owners. SomaSpeed’s pipeline stages execute synchronously in the model’s original call order. These strategies can be combined, with the choice of layout balancing memory capacity and communication cost.

A frozen backbone still performs useful computation: it turns camera images and language into the context needed to choose an action. Its weights stay fixed during expert training. The action expert uses that context together with robot state and noisy action tokens, and learns to predict the model’s action output. Backward therefore needs gradients through the trainable action paths, while the frozen backbone can remain outside the update.

Attention is the connection between the action computation and its context. Queries describe what the action module needs to read, keys identify context entries, and values carry their information. SomaSpeed computes attention in tiles, reading the relevant keys and values as it goes. This avoids storing a full attention-score matrix and allows the kernel to skip gradient work for inputs that are frozen.

The exact structure depends on the model. MolmoAct2 can share context across multiple sampled flow timesteps without making a repeated copy of its keys and values. GR00T’s action head cross-attends to frozen backbone features; projections inside the trainable head still receive gradients. π0.5 reads frozen prefix keys and values separately from the trainable action stream. The model tabs show these distinctions rather than treating every attention path as identical.

After backward, each trainable tensor has a gradient. The training system must normalize those gradients for the effective batch, measure their overall magnitude, apply gradient clipping, and update AdamW’s weights and running moments. These operations repeatedly touch the same data. Performing them separately for many small tensors adds GPU launches and memory traffic between learning steps.

SomaSpeed batches that work across tensors using shared Triton kernels. It first computes the statistics needed for one global clipping decision, accounting for distributed ownership so that each gradient shard contributes once. The final update reads the original gradients and applies normalization, clipping, and AdamW together as it updates the weights and moments. Gradient buffers remain unchanged, and parameter groups and learning-rate schedules retain their configured behavior. Frozen parameters do not enter this optimizer path.

The RL workflow uses Flow-SDE PPO. Rollout workers run the robot policy in LIBERO environments, sampling a denoising trajectory to produce an action chunk. A value head estimates future reward. The workers record the sampled trajectory, rewards, values, old log probabilities, and policy version, giving the learner the information needed to replay experience and evaluate the PPO objective.

PPO uses the same rollout for several training epochs. The frozen backbone would produce the same features for that recorded observation on every replay. SomaSpeed caches those features during collection and passes them with the rollout, so the learner can train the policy and value heads without rerunning the backbone. New observations are encoded as they arrive; the reuse applies to the recorded context associated with a rollout.

Collection and learning use separate GPU roles and proceed concurrently. Independent environments feed batched policy inference, and a bounded queue passes compatible experience to the learner. Updated policy and value heads return to rollout workers asynchronously. The frozen backbone stays in place, while each worker stages a coherent head snapshot and switches versions at a rollout boundary. Tracking the policy version and limiting lag keeps experience from becoming too stale for the learner.

For compatible learner executions, CUDA graphs replay repeated forward and backward work with less launch overhead. The learner can also use distributed model layouts when more memory is needed. The diagram distinguishes the environment loop, the cached rollout packet, the PPO update, and the return of updated heads; each connection carries a different part of the learning process.