최근 블록 확산 언어모델(BDLM)은 자기회귀(autoregressive) 모델의 한계를 극복하며 긴 문맥(long-context)을 처리하는 데 유망한 대안으로 떠오르고 있습니다. 하지만 이러한 모델을 대규모로 학습시키는 과정에서는 분산된 어텐션(attention) 통신과 활성화(activation) 메모리 사용량 때문에 효율성 문제가 발생했습니다. 특히 긴 문맥을 처리할수록 이러한 병목 현상은 더욱 심화되어 모델 개발의 걸림돌이 되어왔습니다.
이 문제를 해결하기 위해 연구진은 블록 병렬 처리(Block Parallelism, BP)라는 새로운 분산 병렬 처리 방식을 제안했습니다. 기존의 문맥 병렬 처리(Context Parallelism, CP)가 전체 시퀀스를 위치별로 분할하는 것과 달리, BP는 각 손상된 블록(corrupted-block) 계산을 특정 처리 장치(rank)에 할당합니다. 나아가 장문맥 환경에 최적화된 문맥 분할 블록 병렬 처리(Context-Sharded Block Parallelism, CSBP)를 도입하여, 공유되는 깨끗한 시퀀스(clean sequence)까지 분할함으로써 손상된 K/V(Key/Value)와 그래디언트(gradient)를 로컬에 유지하고 중복되는 깨끗한 접두사(prefix)를 방지합니다. 이로써 BDLM 학습의 의미론적 일관성을 유지하면서도 효율성을 극대화했습니다.
실험 결과는 CSBP의 뛰어난 성능을 입증합니다. 16개의 H200 GPU 환경에서 256K 문맥 길이를 사용할 때, CSBP는 지도 미세조정(supervised fine-tuning)에서 기존 방식 대비 1.18~1.45배, 자기회귀 모델을 BDLM으로 전환하는 과정에서는 1.27~1.33배의 처리량(throughput) 향상을 보였습니다. 특히 512K 문맥에서는 전체 모델 속도가 1.61배 빨라졌으며, 8개의 H100 GPU 환경에서 DFlash2 추측 디코더(speculative-decoder) 학습 시 512K 문맥에서 2.48배, 1M 문맥에서는 무려 7.59배의 가속 효과를 달성했습니다. 또한, 피크 HBM(High Bandwidth Memory) 사용량을 줄이거나 동일하게 유지하면서도 DiffusionGemma 26B-A4B 모델의 12시간 SFT(Supervised Fine-Tuning) 학습에서 더 높은 성공률(pass rates)을 기록했습니다.
이러한 연구 결과는 대규모 언어모델(LLM)의 장문맥 학습 효율을 획기적으로 개선할 수 있는 중요한 진전을 의미합니다. 학습 시간과 자원 소모는 LLM 개발의 주요 병목 중 하나였는데, CSBP는 이를 크게 완화하여 더 길고 복잡한 문맥을 이해하고 생성하는 모델 개발을 가속화할 것입니다. 이는 코드 생성, 장문서 요약, 복잡한 질의응답 등 다양한 AI 응용 분야에서 BDLM의 실용성을 높이는 데 기여할 것으로 보입니다. 결과적으로, AI 모델의 성능 향상과 더불어 개발 비용 절감에도 긍정적인 영향을 미칠 것으로 기대됩니다.