bf16과 MPS에서의 bf16 지원(2)

2. 애플 실리콘에서는 왜 쓸 수 없었는가?

CUDA가 없는 하드웨어

한때 맥에 NVIDIA 그래픽 카드가 들어가던 시절이 있었다. 하지만 2014년 맥북 프로가 NVIDIA GPU를 탑재한 마지막 맥이 되었고, macOS 10.14 모하비에서 드라이버 지원이 끊기면서 남아 있던 길도 닫혔다. NVIDIA 역시 2019년 CUDA 10.2를 마지막으로 macOS 지원을 접었다.

맥에서 CUDA를 사용할 수 없다는 말은, CUDA로 하던 일을 다른 것으로 해야 한다는 뜻이다. 애플 실리콘 맥에서 그 역할을 맡은 것이 MPS(Metal Performance Shaders) 백엔드다. 

문제는 CUDA로 할 수 있는 일과 MPS로 할 수 있는 일이 같지 않다는 데 있다.

어디에 무엇이 있는가

[ Application / Framework Layer ]
                    PyTorch User Code (Python)
                               │
                               ▼
[ PyTorch C++ Core (LibTorch / ATen) ]
             Dispatcher (CPU / CUDA / MPS 연산 분기)
                               │
                               ▼
┌-------------------------------------------------------------┐
│ [ PyTorch Internals: MPS Backend ]                          │
│  - aten/src/ATen/mps        (장치, 스트림, 메모리 할당기)   │
│  - aten/src/ATen/native/mps (연산 구현, 커스텀 Metal 커널)  │
│  - torch/csrc/mps           (torch.mps 파이썬 바인딩)       │
└-----------------------------┬------------------------------┘
                               │ (호출)
═════════════════════════════════════════════════════════════════
[ macOS System Frameworks ]    │
                               ▼
            ┌---------------------------------------┐
            │ MPSGraph (연산 그래프 융합 및 최적화) │
            │                 │                    │
            │                 ▼                    │
            │ MPS (Metal Performance Shaders)       │
            └-----------------┬--------------------┘
                               │
                               ▼
[ Low-level Graphics & Compute Layer ]
           Metal Framework / Metal Shading Language (MSL)
                          Metal Driver
                               │
                               ▼
[ Hardware Layer ]
                     Apple Silicon GPU Core
  

파이토치에서 사용하는 MPS 백엔드는 애플이 제공하는 프레임워크 위에 얹혀 있다. CUDA처럼 하드웨어 바로 위에서 커널을 작성하는 구조가 아니다.

제약은 어디에 있는가

이 구조를 이렇게 정리할 수도 있다. "애플이 제공하지 않는 기능은 파이토치도 제공할 수 없다." 절반만 맞는 말이다.

파이토치는 MPSGraph에만 의존하지 않는다. 자체 Metal 셰이더를 직접 작성해 애플의 고수준 프레임워크를 우회하기도 하며, 최근에는 그 비중을 늘리고 있다. 올 7월의 2.13 릴리스 노트에는 bernoulli, native_dropout, uniform 같은 연산을 MPSGraph에서 네이티브 Metal 커널로 옮겼다는 항목이 줄줄이 들어 있다. 애플이 만들어 준 그래프 엔진을 쓰는 대신 직접 셰이더를 쓰겠다는 선택이다.

그래서 진짜 제약은 한 단계 아래에 있다. 셰이딩 언어와 하드웨어가 지원하지 않는 것은 어느 경로로도 우회할 수 없다. 그래프 엔진에 없으면 셰이더를 직접 쓰면 되지만, 셰이딩 언어에 자료형이 없으면 셰이더를 쓸 방법 자체가 없다.

무엇이 없었고, 언제 생겼는가

bf16이 정확히 그런 경우였다.

애플 실리콘에서 bf16을 쓸 수 없었던 것은 파이토치의 게으름이나 버그 때문이 아니었다. MPS 백엔드 아래의 프레임워크와 저수준 계층에 bfloat 자료형이 아예 없었기 때문이다. 애플이 제공하는 API에는 fp16에 해당하는 half와 fp32에 해당하는 float가 있었지만, bf16의 설계에 대응하는 자료형은 없었다. MPSGraph에도 없고, 그 아래 셰이딩 언어에도 없었다. 우회할 곳이 없다는 뜻이다.

애플은 그들이 의도한 맥의 활용 방향에 따라 bfloat의 필요성을 간과했다. 애플 실리콘의 GPU와 뉴럴 엔진은 오랫동안 온디바이스 추론을 목표로 발전해 왔다. 사진 분류, 얼굴 인식, 음성 처리 같은 작업이다. 이런 용도에서는 fp16의 좁은 표현 범위가 거의 문제가 되지 않는다. 학습이 아니라 이미 학습된 모델을 실행하는 일이기 때문이다. 기울기 소실을 걱정할 필요가 없는 상황에서는 오히려 가수부가 3비트 더 넉넉한 fp16이 정밀도 면에서 유리하다.

그래서 애플은 fp16을 하드웨어와 셰이딩 언어 양쪽에서 먼저 지원해 왔다. 맥에서 큰 모델을 학습하거나 미세 조정하려는 수요는 뒤늦게 인정한 셈이다.

결국 애플은 macOS 14 소노마에서 bf16을 도입했다. Metal Shading Language 3.1에 bfloat 타입이 생기고, MPSGraph에 MPSDataTypeBFloat16이 추가되었다. 파이토치는 2.3 릴리스에서 이를 받았다. 릴리스 노트에 한 줄로 적혀 있다.

Enable bfloat16 support on MacOS 14 (#119641)

즉 애플 실리콘 맥에서 macOS 14 이상, 파이토치 2.3 이상이면 이제 bf16을 사용할 수 있다. 2024년의 일이다.

3. 그래도 아직 부족하다

문제가 bf16에만 있었던 것은 아니다. MPS 백엔드가 제공하지 않는 기능, CUDA에서만 되는 기능은 지금도 남아 있다.

목록에는 있고 구현되지 않은 연산들

파이토치는 MPS 백엔드에 구현되기를 바라는 연산을 투표로 모으고 있다(링크). 2025년 5월에 만들어진 목록이다. 이후 구현된 항목에는 취소선이 그어졌는데, 아직 다수가 남아 있다. 

남은 것을 보면 linalg_eig, linalg_qr, _linalg_eigh, linalg_matrix_exp처럼 선형대수 분해가 압도적이다. 고윳값 분해와 QR 분해, 촐레스키가 MPS의 가장 큰 공백이다.

목록에조차 없는 것

이 목록이 전부도 아니다.

<멀티모달까지 직접 구현하는 딥러닝 with 파이토치> 10-1절의 예제는 MPS 백엔드에서 다음과 같은 예외가 발생한다.

NotImplementedError: The operator 'aten::_nested_tensor_from_mask_left_aligned' is not currently implemented for the MPS device.

이 연산(_nested_tensor_from_mask_left_aligned)은 투표 집계 목록에 없다. 목록에 오르려면 셋 이상의 요청이 필요한데, 두 건의 요청만 받았기 때문이다. 그 두 표 중 하나를 남긴 사람이 <Hands-On Machine Learning>의 저자 오렐리앙 제롱이다. 용례도 똑같이 "Training a Transformer"라고 적혀 있다. 2025년 3월의 일이고, 1년 반이 지난 지금도 구현되지 않았다.

이 연산 자체는 대단한 것이 아니다. 패딩 마스크가 좌측 정렬인지 판정하는 연산일 뿐이다. 문제는 nn.TransformerEncoder의 기본 설정이 이 연산을 사용한다는 데 있다. 패딩을 건너뛰어 연산량을 줄이는 경로인데, 이를 위해 MPS에 없는 연산을 사용한다. 

다행히 우회할 방법이 있다. torch.backends.mha.set_fastpath_enabled(False)를 호출해 전역 스위치를 끄면 된다. MPS 백엔드 대신 CPU를 사용하는 것도 한 가지 쉬운 대안이 될 수 있다.

늘 한 발 늦는다

이런 사례들을 보면 MPS 백엔드 지원의 패턴을 가늠할 수 있다. 언젠가 지원되지만 늘 한 발 늦는다는 점이다. 예를 들어 지난 7월 2.13 릴리즈에 포함된 MPS 백엔드의 FlexAttention은 CUDA 백엔드는 2.5 릴리즈(2024년 10월)에 포함되었다. 

어쩔 수 없다. 현재 딥러닝 개발의 표준은 CUDA라는 사실은 부인하기 어렵다. 새 기능은 CUDA에서 나오고, MPS는 여러 릴리스 뒤에 따라온다. 개발자 친화적인 하드웨어와 운영체제라 해도, 파이토치로 딥러닝을 하는 사람에게 맥을 선뜻 권하기 어려운 이유가 이것이다.

애플이 직접 내놓은 대안으로 MLX가 있지만 파이토치만큼 널리 쓰이지는 않는다.

기대감

그래도 조금씩 좋아지고 있다.

올 7월의 2.13 릴리스에는 MPS 관련 항목이 스무 개 넘게 들어갔다. CUDA에 2.5부터 있던 FlexAttention이 MPS에서도 동작하게 되었고, 릴리스 노트는 "희소 패턴에서 SDPA 대비 최대 12배"라고 적었다. 늦게 왔지만 올 때는 제대로 온 셈이다.

더 눈에 띄는 것은 앞서 이야기한 흐름이다. 파이토치가 MPSGraph에서 벗어나 자체 Metal 커널을 늘려 가고 있다. 애플이 만들어 주기를 기다리는 대신 직접 만드는 쪽으로 무게가 옮겨 가는 중이다. 그렇게 되면 "애플이 제공하지 않으면 못 한다"는 제약이 닿는 범위도 그만큼 줄어든다.

언젠가는 애플 실리콘의 좋은 하드웨어를 파이토치에서 온전히 활용할 수 있는 시기가 올 것이라는 기대를 가져 본다.

bf16과 MPS에서의 bf16 지원(1)

리눅스에서 잘 동작하는 파이토치 코드를 맥에서 실행하면 의도대로 동작하지 않는 경우가 종종 있다. 특히 연산 장치를 사용하는 경우에 자주 발생하는데, 이는 NVIDIA의 CUDA와 Apple Silicon의 MPS의 지원 수준이 다르기 때문이다.

이 지원 수준은 파이토치가 얼마만큼 지원하느냐도 중요하지만, 파이토치 같은 하드웨어의 성능을 최대한 활용하는 프레임워크의 경우 네이티브 API가 무엇을 지원하느냐도 중요하다. 그 중 하나의 이야기다.

1. bfloat16은 무엇일까?

16비트 실수를 표현하는 두 가지 방법

부동소수점 수는  부호부(1비트), 지수부, 가수부로 메모리를 비트별로 나눠 사용한다. 여기서 지수부는 수의 크기(자릿수)를 정하고, 가수부는 정밀도(유효숫자)를 정한다. 파이토치의 기본 실수 자료형은 torch.float32(fp32)인데, fp32의 구조는 다음과 같다.

fp32: SEEEEEEEEMMMMMMMMMMMMMMMMMMMMMMM   부호(S)*1 + 지수(E)*8 + 가수(M)*23

메모리를 절반만 사용하는 16비트 실수 자료형에는 torch.float16(fp16)뿐 아니라 torch.bfloat16(bf16)도 있다.  두 자료형의 내부 구조는 다음과 같다.

fp16: SEEEEEMMMMMMMMMM    부호(S)*1 + 지수(E)*5 + 가수(M)*10

bf16: SEEEEEEEEMMMMMMM    부호(S)*1 + 지수(E)*8 + 가수(M)*7

fp16은 지수부를 5비트로 줄이고 가수부를 최대한 늘려 정밀도를 지키고 범위를 줄인 반면,  bf16은 반대의 선택을 하고 있다. 여기서 재미있는 것은 bf16은 지수부가 8비트로 fp32와 같고, 마지막 가수부만 23비트에서 7비트로 줄었다는 점이다. 그래서 변환 과정은 뒤의 16비트만 잘라내면 되어 비용이 매우 낮다.

실제 값이 어떻게 될까?

파이토치로 직접 확인해 보자.

import torch

for name, dt in [('fp32', torch.float32),
                 ('fp16', torch.float16),
                 ('bf16', torch.bfloat16)]:
    i = torch.finfo(dt)
    print(f'{name}: max={i.max:.6g}  min={i.tiny:.6g}  eps={i.eps:.6g}')
fp32: max=3.40282e+38  min=1.17549e-38  eps=1.19209e-07
fp16: max=65504        min=6.10352e-05  eps=0.000976562
bf16: max=3.38953e+38  min=1.17549e-38  eps=0.0078125

결과를 확인해 보면 bf16의 범위는 fp32와 같지만 정밀도는 훨씬 떨어진다. 반대로 fp16은 범위와 정밀도 모두를 적당히 깎아서 타협한 결과물이다. 하지만 이 타협의 결과는 썩 좋은 편이 아니다. 범위의 문제가 발목을 잡는다.

오버플로(위로 넘칠 때)와 언더플로(아래로 넘칠 때) 모두 문제가 되지만, 둘 중 언더플로의 문제가 더 크다. 오버플로는 그나마 좀 낫다. fp16은 값이 65504를 넘는 순간 inf가 되고, 이를 사용하는 연산의 결과는 nan이 된다. 반면 bf16은 fp32의 범위를 거의 그대로 다 받는다. 그런데 나은 것이 이정도다. 언더플로는 inf고, nan이고 그런 것 없이 그냥 값이 0이 되어 버린다. 아무런 경고도, 예외도 없이 말이다. 무섭지 아니한가?

그래서 하는 짓: 손실 스케일링

학습 중에 계산되는 기울기는 fp16의 언더플로 영역을 자주 드나든다. 이때 값이 완전히 사라지고 0이 되면, 해당 파라미터는 더 이상 변하지 않는다. 학습 곡선이 어찌되었건 모델의 여기저기에서 학습이 멈추는 일이 벌어진다.

이를 피하기 위해서 '손실 스케일링'이라는 것을 한다. 역전파 전에 손실에 큰 수를 곱해서 기울기를 fp16 범위 안으로 끌어 올리고, 최적화 전에 다시 나누는 방식이다. GradScaler가 하는 일이 바로 이것인데, 오버플로가 발생하면 반대 방향의 일도 함께 처리한다.

한편 bf16은

bf16은 범위가 fp32와 거의 같으므로 애초에 오버플로와 언더플로를 걱정할 일이 크게 줄어든다. 반면 정밀도가 낮은 것은 생각만큼 큰 문제가 되지 않는다. '확률적' 경사 하강법 자체가 확률로 인한 노이즈가 끼어들기 때문이다.

fp16과 bf16의 차이를 정리하면 다음과 같다.

항목	fp16	bf16
범위	좁음 (최대 65504)	fp32와 같음
정밀도	상대적으로 높음	낮음 (유효숫자 2~3자리)
학습	손실 스케일링 필요	대부분 그냥 가능
추론	유리함	무난함
fp32 변환	비트 재배치 필요	뒤를 잘라내면 됨

bf16의 유래

이름에 붙은 b는 Brain을 뜻한다. 구글 브레인이 TPU 전용 자료형으로 설계한 형식으로, 그래서 b가 붙어 있다.

학습의 대상이 되는 모델이 점점 커지면서 널리 보급되었다. 엔비지아가 암페어(A100, RTX30 시리즈)부터 하드웨어 레벨로 지원했고, 이후 큰 모델을 학습할 때 사실상 기본값이 되었다. 허깅페이스 예제 코드에도 dtype=torch.bfloat16을 쉽게 확인할 수 있다.

그런데 어떤 생태계에서는 표준이나 마찬가지인 bf16은 Apple Silicon의 MPS에서는 이야기가 달라진다.

(2편에서 계속)