AutoSP 소개

AutoSP는 긴 컨텍스트 학습에서 발생하는 OOM을 줄이기 위해 Sequence Parallelism(SP) 을 자동 적용하는 DeepCompile 기반 솔루션이다. 기존에는 DeepSpeed나 HuggingFace 코드에 토큰 분할, collective 통신, forward/backward 오버랩을 수동으로 넣어야 했지만, AutoSP는 이를 자동 변환한다.

사용 방법은 단순하다.

• prepare_auto_sp_inputs로 입력 토큰, attention mask, position ids를 태깅한다.
• DeepSpeed config에서 compile.deepcompile을 켜고 passes에 autosp를 넣는다.
• sequence_parallel_size와 ZeRO stage 1을 함께 설정한 뒤 model.compile(dynamic=True)로 컴파일한다.

핵심 변환은 DeepSpeed-Ulysses 기반 SP와 긴 컨텍스트용 Sequence-aware AC(SAC) 이다. Ulysses는 NVLink나 fat-tree 토폴로지에서 GPU 수가 늘어도 통신 오버헤드가 일정하다는 장점이 있지만, SP 크기는 모델의 attention head 수까지로 제한되며, 7-8B급 모델에서는 32가 상한이다. SAC는 PyTorch 2.0의 자동 min-cut AC보다 덜 보수적으로 체크포인트를 배치해, 처리량은 약간 줄어도 더 긴 컨텍스트 학습을 가능하게 한다.

벤치마크에서는 8 A100-80GB SXM 노드에서 여러 Llama 3.1 모델을 돌려, PyTorch 2.7과 CUDA 12.8 환경에서 AutoSP를 RingFlashAttention, DeepSpeed-Ulysses, ZeRO-3와 비교했다. 결과는 같은 자원에서 더 긴 시퀀스를 학습하면서도 런타임 비용 증가는 크지 않다는 쪽으로 요약된다.

제약도 분명하다.

• transformer 전체를 하나의 컴파일 가능한 아티팩트로 묶어야 한다.
• graph break가 있으면 AutoSP가 동작하지 않는다.