ROCm으로 Flow Matching 모델 학습 시 발생하는 NaN 문제

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

AMD RX 7900XTX와 ROCm 환경에서 Flow Matching 모델 학습 시 backward pass에서 지속적인 NaN이 발생하는 문제 보고.

배경

사용자가 AMD RX 7900XTX 환경에서 ROCm을 사용하여 Flow Matching 모델을 학습하려 했으나, NVIDIA CUDA와 달리 backward pass에서 NaN이 발생하는 문제를 겪고 커뮤니티에 공유했다.

의미 / 영향

ROCm 생태계가 표준 모델 학습에는 대응하지만, 복잡하거나 커스텀된 모델 학습 시에는 여전히 호환성 및 안정성 문제가 존재함이 확인됐다.

섹션별 상세

작성자는 RTX 3090에서 정상 작동하던 Flow Matching 모델 코드를 ROCm 7.2 환경으로 포팅했으나, backward pass 호출 시 모든 가중치에서 NaN이 발생하는 문제를 확인했다. 모델의 forward pass는 정상적으로 수행되나, 역전파 과정에서만 문제가 발생하며 bf16, fp32 등 정밀도 설정을 변경해도 해결되지 않았다. 코드 수정 없이 환경만 변경했음에도 발생한 현상으로, ROCm 환경에서의 역전파 과정에 수치적 불안정성이 존재함이 확인됐다.

ROCm 스택의 범용성에 대한 의문이 제기되었다. 동일 환경에서 nanoGPT 학습 스크립트는 정상 작동하는 것으로 확인되었다. 잘 알려진 코드베이스와 커스텀 코드베이스 간의 ROCm 호환성 차이가 존재하며, ROCm은 표준화된 모델 학습에는 대응하지만 커스텀 모델이나 덜 일반적인 아키텍처에서는 여전히 호환성 문제가 존재한다.