tsbootstrap, 160MB 텐서를 버리고 최대 20배 빨라지다

Count the Bytes, Not the FLOPs

tsbootstrap의 블록 부트스트랩 엔진이 기존에는 모든 부트스트랩 복제본을 하나의 거대한 배치 텐서로 만들어 메모리 병목을 일으켰습니다. 개발자들은 '배열을 구체화하지 말 것'과 '파생 가능한 상태를 유지하지 말 것'이라는 두 가지 규칙으로 핫 경로를 재설계했습니다. 그 결과, 긴 시계열에서 3.1배에서 20배의 속도 향상을 얻었고, 메모리 사용량은 최대 97배까지 줄었습니다. 이 글은 CPU 메모리 계층 구조와 산술 강도(arithmetic intensity) 개념을 통해 문제를 설명하고, FlashAttention이 어떻게 같은 원리로 트랜스포머 학습을 가속화했는지도 보여줍니다.