야쿠브 치비를레이의 희소 커널이 학습 처리량을 36% 높였다

RTX PRO 6000 한 대에서 학습 처리량이 36% 높아졌습니다. 야쿠브 치비를레이(Jakub Ćwirlej)가 커널(kernel)을 공개했습니다. 제작자 측은 이를 같은 피드포워드 층(feed-forward layer) 용량의 밀집 모델과 비교했습니다. 저는 처리량 수치보다 토큰(token)마다 일부 뉴런만 계산한다는 설계가 더 오래 눈에 남습니다. 다만 확인은 GPU 한 종류에 그쳤습니다.

야쿠브 치비를레이가 단일 GPU로 EUGEN 160M을 훈련한 흐름에서, 이번에는 모델이 아니라 피드포워드 계산을 줄이는 코드를 공개했습니다. 이번 36% 향상은 제작자가 RTX PRO 6000 한 대에서 측정한 값입니다.

측정에 쓴 설정에서는 밀집 모델과 같은 피드포워드 층 용량을 두고, 토큰마다 뉴런 4,096개 가운데 256개를 활성화했습니다. 그 차이를 한 블록에 그치지 않고 306M 파라미터 전체 모델의 학습 처리량으로 비교했다는 점이 흥미롭습니다.

36%는 제작자가 RTX PRO 6000 한 대에서 얻은 결과이며, 다른 GPU에서 재현된 수치는 아닙니다. 또 희소 방식으로 학습을 시작한 뒤 밀집 모델로 바꾸는 전체 과정에서 이점이 이어지는지도 아직 확인되지 않았습니다. GPU가 다르거나 학습 경로를 바꾸면 그 36%가 그대로 따라온다고 기대할 근거는 없습니다.

저장소는 CC BY-NC 4.0으로 비상업적 이용만 허용하므로, 상업적 작업이라면 먼저 사용 권한을 확인하십시오. Ampere 이상 GPU와 PyTorch 2.6 이상을 갖췄다면, 같은 모델과 배치로 밀집 구현과 커널을 시험하십시오. 처리량과 검증 손실을 비교하고, 희소 방식으로 시작해 밀집 형태로 바꾼 소규모 시험으로 전환 뒤 이득도 확인하십시오.

마지막으로 GPU 계산 코드를 바꿔 보셨을 때, 속도와 검증 손실 중 무엇을 먼저 비교하셨습니까?