FP8 Transmission Optimization for VLA Training: Byte Reduction Leads to Increased Latency

Experimental results were released regarding the implementation of FP8 wire transmission to resolve gradient communication bottlenecks caused by reduced P2P bandwidth when migrating a 3-modality VLA model from an NVLink environment to a PCIe Gen5-only environment. By intercepting NCCL collectives via an LD_PRELOAD shim and applying blockwise scaling, wire costs were reduced by 48.4% compared to bf16; however, a paradoxical result occurred where the final step time increased.

FP8 Transmission Optimization and Bottlenecks

• Transmission Efficiency: Applying FP8 (e4m3/e5m2) quantization reduced transmission bytes by 48.4%, and custom kernels reduced conversion bottlenecks on architectures below sm_89 to 3.3 µs. Additionally, fused kernels resolved the HBM traffic issue of naive dequantize-then-reduce, achieving a 5.0-9.3x speedup.
• Threshold Changes: In PCIe environments, the byte reduction effect is 2-4x larger than in NVLink, lowering the crossover threshold for AllGather to 767 KiB. Furthermore, for large payloads (64 MiB), all-to-all decomposition recorded 2.0x greater savings compared to ring-based approaches.

Negative Result: Increased Latency Despite Byte Reduction

• Critical Path Issue: Coalescing small AllReduces in the DeepSpeed compiled path reduced traffic, but steady-state step time increased by +7.3 ms to 1422.2 ms compared to pre-merger (1414.9 ms).
• Cause Analysis: Pre-merger small AllReduces overlapped with the backward pass, but post-merger they were processed as serial large packets after the outermost ncclGroupEnd, placing them on the critical path. This means the eliminated bytes were not originally on the critical path.
• Lesson: Exposure profiling is essential before volume optimization, proving that 'GB moved' and 'unhidden ms' are distinct metrics.

Other Technical Findings and Open Source Release

• Memory Management Issue: Unbounded greedy merging generated 129 calls and a 999.8 MB union, causing the scratch pool to explode. To resolve this, applying size caps and using the PyTorch caching allocator is recommended.
• Framework Release: An open-source framework 'loongforge-vla' for 3-modality VLA training was released. It supports DDP, native ZeRO, torch.compile, etc., and provides NVTX instrumentation and a loss-parity harness.
• Bug Fix: The loss-parity harness discovered and fixed an issue where LambdaLR.__init__ immediately multiplied the optimizer LR by schedule(0), causing step 0 to be a no-op update.