HSAP: 하이브리드 컨텍스트 생성 모델을 위한 계층적 시퀀스 인식 병렬화

본 논문에서 우리는 기존 시퀀스 병렬화 (Sequence Parallelism) 패러다임의 장점들을 결합하고, 그 단점들을 극복하는 더 강력한 시퀀스 병렬화 프레임워크를 구축하는 것을 목표로 합니다. 가장 심각한 단점은 하이브리드 컨텍스트 (Hybrid-context) 패킹 시퀀스 (Packed sequences)에서 인과적 어텐션 (Causal attention)을 정확하게 계산할 수 없다는 점입니다. 대규모 언어 모델 (LLM)을 효율적으로 사전 학습 (Pretraining) 및 미세 조정 (Fine-tuning)하기 위해 시퀀스를 패킹하는 실용적인 기술은 어텐션 계산 시 교차 오염 (Cross-contamination) 문제를 야기하며, 이는 시퀀스 길이 차원에서 병렬화를 수행하지 않을 때는 효과적으로 해결될 수 있습니다. 그러나 시퀀스 병렬화에서 기존 방식들은 하이브리드 컨텍스트 시퀀스 시나리오를 무시하거나, 반대로 해당 시나리오를 지원하기 위해 병렬화 정도를 희생하고 제한합니다. 이를 위해, 우리는 다중 장치 그룹 간의 집중적인 텐서 전송 (Tensor transmission) 및 부분적 어텐션 계산 (Partial attention computation) 문제를 정복하기 위한 효율적인 시퀀스 인식 병렬화 (Sequence-Aware Parallelism) 알고리즘을 혁신적으로 제안합니다. 우리의 알고리즘은 JIT (Just-In-Time) 컴파일을 활용하여 NCCL 레벨에서 모든 장치 그룹의 통신 전략을 최적화합니다. 나아가, 우리는 기존의 시퀀스 병렬화 패러다임들을 우리의 시퀀스 인식 알고리즘으로부터 이득을 얻는 계층적 시퀀스 인식 병렬화 (Hierarchical Sequence-Aware Parallelism) 프레임워크로 통합합니다. 또한, 성능을 최적화하기 위해 계층적 프레임워크의 메모리 및 통신 오버헤드 관리에 대해 상세히 설명합니다. 다양한 실험을 통해, 우리가 제안하는 방식이 여러 지표에서 다른 최첨단 (State-of-the-art) 시퀀스 병렬화 방식들보다 뛰어난 성능을 보임을 입증합니다.