PyTorch AutoSP：컴파일러 기반 도구가 학습 코드를 자동으로 시퀀스 병렬화하여 100k+ 토큰 컨텍스트 지원

2026년 4월 29일, PyTorch 팀은 UIUC SSAIL Lab, Anyscale, Snowflake의 연구원들과 공동으로 AutoSP를 발표했습니다. 이는 DeepSpeed/DeepCompile 생태계 내의 컴파일러 기반 도구로, 표준 단일 GPU 트랜스포머 학습 코드를 시퀀스 병렬 형식으로 자동 변환합니다. 목표：분산 코드를 수동으로 구현하지 않고 초장 컨텍스트（100k+ 토큰）의 LLM 학습을 가능하게 하는 것입니다.

AutoSP가 해결하는 문제

긴 컨텍스트 LLM 학습에서는 트랜스포머의 활성화가 컨텍스트 길이에 따라 이차적으로 확장되어 여러 GPU에 시퀀스를 분산（시퀀스 병렬화, SP）해야 하며, OOM 오류가 쉽게 발생합니다. 기존 솔루션（RingFlashAttention, DeepSpeed-Ulysses）은 토큰 파티셔닝, 통신 집합 연산, 복잡한 어텐션 마스킹 등 학습 코드의 수동 재작성이 필요합니다.

AutoSP는 이 모든 것을 자동화합니다：사용자가 표준 트랜스포머 코드를 작성하면 컴파일러가 SP를 인식하는 형식으로 변환합니다.

활성화 방법

DeepSpeed 설정에 세 줄을 추가하면 됩니다：

config = {
    "compile": {
        "deepcompile": True,
        "passes": ["autosp"]
    },
    "sequence_parallel_size": 4,
    "zero_optimization": {"stage": 1},  # AutoSP는 ZeRO 0/1과 함께 사용 가능
    ...
}

그리고 입력 준비를 위한 유틸리티 prepare_auto_sp_inputs()가 필요합니다. 내부 전략은 DeepSpeed-Ulysses 아키텍처를 채택하여 NVLink/fat-tree 네트워크에서 GPU 수 증가에 대해 일정한 통신 오버헤드를 유지하며, 어텐션 헤드 수（예：7-8B 모델의 32개 헤드）까지 확장됩니다.

시퀀스 인식 활성화 체크포인팅（SAC）

AutoSP는 SAC（긴 컨텍스트 학습에 최적화된 커스텀 체크포인팅 전략）도 도입했습니다. 보수적인 PyTorch 2.0 최대 흐름 최소 컷 공식과 달리, SAC는 긴 컨텍스트 특유의 FLOP 동적을 활용하여 계산 비용이 낮은 연산자의 중간 활성화를 해제하고 역방향 패스에서 재계산합니다. 트레이드오프： 처리량이 약간 감소하지만 더 긴 컨텍스트를 가능하게 합니다.

결과

8×A100-80GB SXM 노드（PyTorch 2.7, CUDA 12.8）에서 다양한 크기의 Llama 3.1 모델로 테스트：

• 동일한 리소스에서 최대 학습 가능한 시퀀스 길이가 크게 증가
• 수동으로 작성된 RingFlashAttention 및 DeepSpeed-Ulysses 기준과 비교하여 런타임 오버헤드 최소화

엔드 투 엔드 예시（Llama 3.1 8B 포함）는 github.com/deepspeedai/DeepSpeedExamples/tree/master/benchmarks/autosp에서 확인할 수 있습니다.

한계

현재 AutoSP는 단일 컴파일 가능한 아티팩트（전체 트랜스포머가 하나의 블록으로 컴파일）를 요구하며, 모델 내 그래프 브레이크를 허용하지 않습니다. 팀은 그래프 브레이크 내성 확장을 다음 단계로 언급하고 있습니다.