AI 생성 CUDA kernel, benchmark 통과 뒤 training을 망가뜨린 bf16 함정

NVIDIA의 SOL-ExecBench를 통과한 AI 생성 CUDA kernel이 실제 transformer training에서는 loss를 망가뜨렸다는 사례가 r/MachineLearning에서 큰 관심을 받았다. 문제의 kernel은 fused embedding-gradient와 RMSNorm backward pass를 처리했고, benchmark verifier는 통과했다. 하지만 작은 transformer training loop에 넣자 loss가 회복되지 않고 발산했다.

원인은 단순한 오답이 아니라 더 까다로웠다. embedding-gradient 누적이 fp32가 아니라 bf16으로 이뤄지고 있었다. 균등하게 샘플링한 token에서는 기여가 넓게 퍼져 문제가 숨었지만, 실제 텍스트에서는 일부 token ID에 gradient가 몰린다. 이때 작은 기여값이 커진 accumulator에 밀려 0처럼 사라지고, 빈도가 높은 row가 조금씩 drift한다. AdamW를 쓰면 per-parameter normalization이 bias를 흡수해 loss에서 문제가 사라져 보였다.

댓글은 “passes the verifier”라는 표현의 약점을 파고들었다. 한 사용자는 bf16과 fp32 차이를 많은 사람이 그냥 지나칠 수 있다고 했고, 다른 댓글은 optimizer와 dataset sensitivity가 kernel test에 들어가야 한다고 지적했다. 이 사례가 불편한 이유는 bug의 증상이 “아이디어가 틀렸다”처럼 보인다는 데 있다. 연구자는 dataset, architecture, optimizer를 의심하다가 implementation 자체를 늦게 확인하게 된다.

AI가 kernel을 생성하고 benchmark를 최적화하는 흐름은 빠르게 좋아지고 있다. 다만 성능 competition의 다음 단계는 더 높은 speedup만이 아니라, real workload에서 조용히 틀리는 failure mode를 잡는 검증이다. 특히 LLM training처럼 분포와 optimizer가 결과를 바꾸는 영역에서는 verifier가 실제 사용 조건을 더 넓게 품어야 한다.

Reddit discussion · Related research post