PyTorch 2.11 출시: FlexAttention + FlashAttention-4와 분산 학습 개선
PyTorch 2.11의 FlexAttention+FlashAttention-4 백엔드 속도 향상과 미분 가능 집합 연산(Differentiable Collectives)이 분산 학습에 가져오는 변화를 정리한다.
2026-08-14 · 최초 발행 2026-03-25
PyTorch 재단이 2026년 3월 23일 PyTorch 2.11.0을 정식 출시했다. 이번 릴리스의 핵심은 두 가지다. FlexAttention이 FlashAttention-4 백엔드를 지원하며 최대 3.2배의 속도 향상을 실현했고, 분산 학습에서 집합 연산(collective operations)을 역전파(backpropagation)에 통과시킬 수 있는 미분 가능 집합 연산(Differentiable Collectives)이 추가됐다.
FlexAttention이 FlashAttention-4를 만나면
FlexAttention은 PyTorch 2.5에서 처음 도입된 기능으로, 개발자가 Python으로 어텐션 패턴을 정의하면 자동으로 최적화된 GPU 커널로 컴파일하는 API다. 기존에는 Triton 백엔드를 사용했는데, 2.11부터 NVIDIA Hopper(H100)와 Blackwell(B200) GPU에서 FlashAttention-4 백엔드를 사용할 수 있다.
성능 향상 폭은 Triton 구현 대비 compute-bound 워크로드에서 1.2배~3.2배다. 이 개선은 CuTeDSL을 통해 score/mask 수정 함수를 자동 생성하고 FlashAttention-4 커널을 JIT 방식으로 인스턴스화하는 방식으로 달성된다.
지원 어텐션 패턴으로는 ALiBi(위치 편향), 슬라이딩 윈도우 어텐션, 문서 마스킹(document masking), 소프트 캡핑(soft-capping) 등 복잡한 패턴을 모두 지원하면서 최적 성능을 유지한다. 다만 제약도 있다 — NVIDIA Hopper와 Blackwell GPU 전용이며, 다른 하드웨어에서는 자동으로 Triton 백엔드로 폴백된다. 현재 활발히 개발 중인 기능으로 API가 변경될 수 있다.
집합 연산으로 역전파가 흐르게 하기
분산 딥러닝에서 all-reduce, all-gather 같은 집합 연산은 여러 GPU에 걸쳐 텐서를 동기화하는 데 사용된다. 기존 PyTorch에서는 이 연산들이 미분 불가능(non-differentiable)했기 때문에, 이를 통해 역전파하려면 복잡한 커스텀 autograd 함수를 작성해야 했다.
PyTorch 2.11은 functional collective 연산에 미분 가능성 지원을 추가했다. 이제 집합 연산을 통해 역전파할 수 있어 분산 환경에서 더 자연스러운 그래디언트 흐름, 커스텀 autograd 함수 없이 고급 분산 학습 기법 구현, 연구 코드의 단순화가 가능해진다.
업그레이드하기 전에 알아둘 것
FlashAttention-4 백엔드는 Hopper/Blackwell GPU에서만 활성화된다. H100이나 B200이 없다면 2.11로 업그레이드해도 이 기능의 이점을 직접 체감하기 어렵다. 다만 Differentiable Collectives는 하드웨어 제약 없이 모든 분산 학습 환경에 적용된다.
PyTorch 2.11은 트랜스포머 학습의 두 핵심 병목—어텐션 연산 속도와 분산 학습의 그래디언트 흐름—을 동시에 개선한 의미 있는 릴리스다. FlexAttention과 FlashAttention-4의 결합은 연구자들이 유연한 어텐션 패턴을 실험하면서도 프로덕션 수준의 성능을 유지할 수 있게 해준다. 최신 NVIDIA GPU를 사용하는 LLM 연구자와 엔지니어라면 업그레이드를 적극 검토할 만하다.