Flash-MSA, 백만 토큰 AI 훈련 비용을 낮추는 코드 공개
Flash-MSA는 Sparse Attention을 효율적으로 훈련하기 위한 비공식 오픈소스 커널이다. H100과 B200 에서 작동하며, 전체 문맥을 매번 계산하는 대신 작은 선택기가 중요한 정보를 128개 단위 블록으로 골라 주 계산에 넘긴다. 선택된 블록의 위치를 저장해 재사용하므로, 선택기의 첫 계산을 제외한 훈련 과정은 문맥 길이에 비례하는 으로 처리할 수 있다.
주 계산과 선택기를 함께 학습하면서도 KL 발산 전체를 메모리에 만들지 않고 같은 기울기를 계산해 메모리 사용과 처리 부담을 줄였다. 4,096개와 8,192개 토큰으로 진행한 bf16 정확성 검사에서 일반 구현과 출력 및 기울기의 가 대체로 0.9983~1.0000이었다. 설치 가능한 코드와 Megatron-LM 훈련 예제가 공개됐지만 의 공식 구현은 아니다.
현재는 머리 크기와 블록 크기가 각각 128인 설정만 지원하며, 양자화 훈련과 일부 주의 기능은 지원하지 않는다. 백만 토큰 훈련에는 여러 가 문맥을 나누는 추가 작업도 필요하며, 이 부분은 아직 완성되지 않았다.
핵심 포인트
- 중요한 정보가 있는 128개 단위 블록만 골라 주 계산에 사용한다.
- 선택한 블록 위치를 저장해 역전파에서도 다시 사용한다.
- bf16 검사에서 기준 구현과 0.9983~1.0000의 를 보였다.
- H100과 B200용 코드 및 Megatron-LM 훈련 예제가 오픈소스로 공개됐다.
- 양자화 훈련과 백만 토큰용 다중 장비 처리는 아직 지원하지 않는다.