StreamBP: 장기 시퀀스 LLM 학습을 위한 메모리 효율적 정확 역전파

초록

긴 시퀀스 데이터에 대한 언어 모델 훈련은 복잡한 작업(예: 장기 추론)에서 모델의 성능을 향상시키기 위한 중요한 요구사항이다. 그러나 시퀀스 길이가 증가함에 따라, 역전파(Backpropagation, BP) 과정에서 활성화 값(activation values)을 저장하기 위한 메모리 비용이 크게 증가하며, 이는 그래디언트 체크포인팅(gradient checkpointing) 기술을 적용하더라도 여전히 문제가 된다. 이러한 문제를 해결하기 위해, 우리는 StreamBP라는 메모리 효율적이고 정확한 BP 방법을 제안한다. 이 방법은 시퀀스 차원에서 체인 규칙(chain rule)을 계층별로 선형 분해하여 활성화 값과 로짓(logits)의 메모리 비용을 크게 줄인다. 제안된 방법은 SFT, GRPO, DPO와 같은 일반적인 목적 함수에 적용 가능하다. 구현 측면에서, StreamBP는 언어 모델의 인과적 구조(causal structure)를 활용하여 계산 FLOPs를 줄이고 BP 속도를 향상시킨다. 그래디언트 체크포인팅과 비교했을 때, StreamBP는 BP의 최대 시퀀스 길이를 2.8~5.5배까지 확장할 수 있으며, 비슷하거나 더 적은 BP 시간을 사용한다. 또한, StreamBP의 시퀀스 길이 확장 능력은 훈련 가속을 위한 배치 크기(batch size) 확장으로 직접 전환될 수 있다. 우리는 더 나아가 통신 효율적인 분산 StreamBP를 개발하여 다중 GPU 훈련을 효과적으로 지원하고 적용 범위를 넓혔다. 우리의 코드는 모든 트랜스포머 모델의 훈련 파이프라인에 쉽게 통합될 수 있으며, https://github.com/Ledzy/StreamBP에서 확인할 수 있다.

English

Training language models on long sequence data is a demanding requirement for enhancing the model's capability on complex tasks, e.g., long-chain reasoning. However, as the sequence length scales up, the memory cost for storing activation values becomes huge during the Backpropagation (BP) process, even with the application of gradient checkpointing technique. To tackle this challenge, we propose a memory-efficient and exact BP method called StreamBP, which performs a linear decomposition of the chain rule along the sequence dimension in a layer-wise manner, significantly reducing the memory cost of activation values and logits. The proposed method is applicable to common objectives such as SFT, GRPO, and DPO. From an implementation perspective, StreamBP achieves less computational FLOPs and faster BP speed by leveraging the causal structure of the language model. Compared to gradient checkpointing, StreamBP scales up the maximum sequence length of BP by 2.8-5.5 times larger, while using comparable or even less BP time. Note that StreamBP's sequence length scaling ability can be directly transferred to batch size scaling for accelerating training. We further develop a communication-efficient distributed StreamBP to effectively support multi-GPU training and broaden its applicability. Our code can be easily integrated into the training pipeline of any transformer models and is available at https://github.com/Ledzy/StreamBP.

StreamBP: 장기 시퀀스 LLM 학습을 위한 메모리 효율적 정확 역전파

StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs

초록

Support