PyTorch 2.14에서 CUDA bfloat16 FFT·STFT를 쓸 때 확인할 조건 | DAKER 커뮤니티

FFT bfloat16

주파수 도메인 전처리를 GPU에서 처리할 때는 연산 자체보다도 어떤 조건에서 낮은 정밀도 경로가 실제로 열리는지가 더 중요할 때가 있습니다. 특히 torch.fft나 torch.stft를 bfloat16으로 맞춰 보려는 경우라면, 지원 범위를 먼저 알아두는 편이 좋습니다.

PyTorch 2.14는 CUDA에서 torch.fft 계열과 torch.stft에 bfloat16 입력을 받습니다. 다만 모든 변환 크기에서 같은 방식으로 처리되는 것은 아닙니다. 실험을 시작하기 전에 네이티브 cuFFT 경로가 열리는 조건과, 그렇지 않을 때 float32로 승격되는 지점을 함께 보는 것이 좋습니다.

PyTorch 2.14는 CUDA에서 torch.fft 계열과 torch.stft에 bfloat16 입력을 받지만, 네이티브 cuFFT bfloat16 경로는 SM80 이상·2의 거듭제곱 변환 길이에서만 열리고 그 밖은 float32로 승격됩니다.

어떤 조건에서 bfloat16 경로가 열리나

핵심은 두 가지입니다. GPU가 SM80 이상이어야 하고, 변환 길이가 2의 거듭제곱이어야 합니다. 이 조건을 만족하면 cuFFT의 네이티브 bfloat16 경로를 사용할 수 있습니다. 반대로 이 범위를 벗어나면 입력이 bfloat16이어도 내부적으로 float32로 승격됩니다.

따라서 스펙트로그램이나 주파수 도메인 전처리를 낮은 정밀도로 통일하고 싶다면, 먼저 변환 길이 정책이 이 조건에 맞는지부터 확인하는 편이 좋습니다.

재현해 볼 실험 조건

가장 간단한 방법은 짧은 1D·2D FFT와 STFT 한 세트를 float32와 bfloat16으로 같은 길이에서 비교하는 것입니다. 이때 지원 조건과 비지원 조건을 나눠서 dtype, 오차, 커널 시간을 함께 기록하면 됩니다.

1D·2D FFT 비교

  1. SM80 이상 GPU에서 power-of-two 길이의 torch.fft.fft·rfft를 bfloat16 입력으로 실행하고 dtype·시간을 기록합니다.
  2. 같은 입력을 비 power-of-two 길이로 바꿔, float32 승격이 일어나는지 확인합니다.

STFT 비교

  1. torch.stft에도 동일한 dtype 실험을 적용해, 스펙트로그램 경로의 정밀도·지연을 비교합니다.

실험 노트에는 디바이스 capability, 변환 길이, 승격 여부를 함께 남기면 됩니다. 같은 bfloat16 입력이라도 길이에 따라 실제 경로가 달라질 수 있기 때문입니다.

실무에서 먼저 볼 포인트

오디오·센서 FFT 전처리를 GPU에서 끝내는 작업이라면, 길이를 2의 거듭제곱에 맞출 수 있는지가 성능과 메모리 이득에 직접 연결됩니다. 네이티브 경로를 유지하기 쉬워지기 때문입니다.

반대로 승격이 자주 일어나면 bfloat16을 쓰는 이유였던 메모리·대역폭 이득이 줄어들 수 있습니다. 그래서 길이 정책을 먼저 정해 두는 것이 좋습니다. 재현성과 비교 가능성을 위해 CUDA capability와 변환 크기를 기록해 두는 것도 중요합니다.

길이 정책이 네이티브 bfloat16 경로를 유지하느냐, float32 승격을 반복하느냐를 가릅니다.

함께 보면 좋은 자료

CPU FFT는 저정밀 dtype을 계속 거부하므로, 이번 확인 범위는 CUDA 경로에 한정됩니다. 또한 이 주제는 NVGEMM·복소 컴파일 글과는 축이 다릅니다.

PyTorch 실전 고급 트랙과 DAKER 커뮤니티에서도 관련 맥락을 이어서 볼 수 있습니다.

참고 자료

https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/community

실제로 비교해 보셨다면 power-of-two 길이와 비 power-of-two 길이에서 체감 차이가 어느 정도였는지 궁금합니다.