VLA 학습 FP8 전송 최적화, 바이트 감소에도 지연 증가

NVLink에서 PCIe Gen5 전용 환경으로 3모달 VLA 모델을 마이그레이션할 때, P2P 대역폭 감소로 인한 gradient communication 병목을 해결하기 위해 FP8 wire 전송을 구현한 실험 결과가 공개되었다. LD_PRELOAD shim을 통해 NCCL collective를 가로채고 blockwise scaling을 적용하여 wire cost를 bf16 대비 48.4% 절감했으나, 최종적인 step time은 오히려 증가하는 역설적인 결과가 발생했다.

FP8 전송 최적화와 병목 현상

• 전송 효율화: FP8(e4m3/e5m2) 양자화를 적용해 전송 바이트를 48.4% 줄였고, sm_89 미만 아키텍처의 변환 병목을 자체 커널로 3.3 µs까지 단축했다. 또한 fused kernel을 통해 naive dequantize-then-reduce의 HBM traffic 문제를 해결해 5.0-9.3배 가속을 달성했다.
• 임계값 변화: PCIe 환경에서는 NVLink 대비 바이트 절감 효과가 2-4배 커져, AllGather의 crossover 임계값이 767 KiB로 낮아졌다. 또한 large payload(64 MiB)에서 all-to-all decomposition이 ring-based 대비 2.0배 더 큰 savings를 기록했다.

Negative Result: 바이트 감소에도 지연 증가

• Critical Path 문제: DeepSpeed compiled path의 small AllReduce를 coalescing하여 traffic을 줄였으나, steady-state step time은 1422.2 ms로 병합 전(1414.9 ms) 대비 +7.3 ms 증가했다.
• 원인 분석: 병합 전 small AllReduce는 backward pass와 overlap되었으나, 병합 후 outermost ncclGroupEnd 이후 serial large packet으로 처리되면서 critical path에 포함되었기 때문이다. 즉, 삭제된 바이트가 원래 critical path에 없었음을 의미한다.
• 교훈: Volume 최적화 전 Exposure profiling이 필수적이며, 'GB 이동량'과 '숨겨지지 않은 ms'는 별개의 지표임을 입증했다.

기타 기술적 발견 및 오픈소스 공개

• Memory 관리 문제: Unbounded greedy merging은 129 calls, 999.8 MB union을 생성하며 scratch pool을 폭발시켰다. 이를 해결하기 위해 size cap 적용 및 PyTorch caching allocator 사용이 권장된다.
• Framework 공개: 3-modality VLA 학습을 위한 오픈소스 프레임워크 'loongforge-vla'가 공개되었다. DDP, native ZeRO, torch.compile 등을 지원하며, NVTX instrumentation과 loss-parity harness를 제공한다.
• Bug 수정: Loss-parity harness를 통해 LambdaLR.__init__이 optimizer LR에 schedule(0)을 즉시 곱하여 step 0이 no-op update가 되는 문제를 발견하고 수정했다.