Achieving 2.37× Faster OpenPI PyTorch Training with RLinf and Baidu Baige

Background and Goal

The OpenPI Pi0.5 model (~3.5 B parameters) is a representative Vision‑Language‑Action (VLA) model. To enable large‑scale embodied‑intelligence training in PyTorch, RLinf performed a precision alignment between the original JAX implementation and PyTorch, ensuring identical loss curves and training outcomes.

Precision Alignment (Run‑Correctly)

RLinf identified 22 implementation differences, 11 of which directly affected training results. They fixed mismatches in model semantics, initialization, and data handling, and built a four‑layer verification suite (forward, gradient, distributed, checkpoint). After alignment, the PyTorch loss curve matched the JAX baseline with a deviation <0.001.

Full‑Stack Performance Optimization (Run‑Fast)

The Baidu Baige AI Infra team treated the entire training pipeline as a continuous data‑and‑compute flow:

data_fetch → H2D + preprocessing → forward + loss → backward + communication → optimizer

Each stage was profiled to locate bottlenecks, and optimizations were applied iteratively with re‑profiling after each change.

1. Data Fetch

Removed an unused third camera slot in the LeRobot v2.0 dataset, reducing visual tokens from 3 to 2 and shrinking the attention matrix to ~54% of the original size.

This slot‑pruning contributed a +27.13 samples/s (sps) throughput gain, the single largest improvement.

2. H2D + Preprocessing

Changed image transfer from float32 to uint8 and performed normalization on the GPU.

Data transferred per batch dropped from 29 MB to 9.7 MB (≈4× reduction), yielding a +1.32 sps overall gain.

3. Forward + Loss

Compiled the 18 Gemma blocks with torch.compile + Inductor, reducing kernel launch overhead.

Resulted in +15.88 sps.

Re‑evaluated activation checkpointing after memory was freed by later communication changes; a stride‑4 recomputation for SigLIP gave only +0.09% speedup, so the original checkpointing was retained for Gemma while SigLIP kept higher recomputation.

4. Backward + Communication

Refactored FSDP wrapping from a single large unit to per‑block units, cutting memory usage from 68.37 GiB to 32.45 GiB (‑35.92 GiB) and exposing finer‑grained communication objects.

Enabled forward and backward prefetch, raising communication‑compute overlap from 0% to 82.9% and reducing NCCL idle time by 81.5%.

Prefetch alone added +28.25 sps, the biggest throughput gain after slot pruning.

5. Optimizer and Runtime

Fused AdamW updates gave +0.76 sps.

Memory‑allocator fragmentation reduction reclaimed 9.78 GiB and added +0.44 sps.

Scaling Results

After the full optimization chain, single‑node 8‑GPU throughput reached 175.57 samples/s (2.37× faster) with step time reduced from 3451.6 ms to 1458.9 ms (‑57.7%). Accuracy remained unchanged.

Scaling to 32 nodes (256 GPUs) achieved 5340 samples/s, a linear‑scale efficiency of 91.2%, confirming that the optimizations preserved communication‑compute overlap across instances.

Conclusion

Systematic, full‑pipeline profiling and iterative re‑calibration can dramatically accelerate large‑scale VLA training without sacrificing model fidelity. RLinf’s precision alignment provides the correctness foundation, while Baidu Baige’s AI Infra optimizations deliver the performance needed for rapid embodied‑intelligence iteration.

RLinf documentation: https://rlinf.readthedocs.io/en/latest/rst_source/examples/embodied/sft_openpi_rlinf.html