Burn 0.22가 백엔드 제네릭을 걷어내고 재빌드 시간을 15배 줄였다

Lobsters3일 전조회 1

Rust 기반 딥러닝 프레임워크 Burn이 0.22.0을 내놨다. 이번 버전의 핵심은 사용자 코드에서 백엔드 타입 파라미터를 완전히 들어낸 것이다. 모델 구조체와 함수 시그니처에 붙어 있던 B: Backend가 사라지고, 실행 위치와 방식은 디바이스가 결정한다. 여기에 적응형 메모리 관리, 개선된 오토튜닝, 그래프 리플레이, 커널·컴파일러 기반 정비가 더해졌고, LoRA·QLoRA 적응과 ONNX 내보내기, 원격 연산 지원도 확장됐다.

기존에는 백엔드를 고르는 순간 그 타입이 애플리케이션 전체로 번졌다. 모델이 B: Backend를 들고 다니면 각 레이어와 파라미터, 텐서 함수까지 같은 파라미터를 전파해야 했다. 구조는 유연했지만 컴파일 시 의존 사슬이 길어지는 대가를 치렀다. 0.22에서는 Tensor → Bridge → Dispatch → Backend 구조로 바뀌면서 브리지가 구체 백엔드 표현을 상위 텐서 API에서 감춘다. 덕분에 프로젝트 코드를 고쳐도 다시 컴파일되는 범위가 크게 줄고, 하나의 애플리케이션에서 여러 백엔드를 함께 켜기도 쉬워진다. Backend 트레이트 자체는 백엔드 구현과 커스텀 연산을 위해 프레임워크 핵심에 남는다.

디바이스 지정 방식도 달라졌다. Cargo 피처를 켠 뒤 Device::cuda(0), Device::wgpu(...), Device::flex() 같은 호출로 모델과 입력을 초기화할 때 실행 대상을 정한다. 자동미분 쪽 변화도 눈에 띈다. B: AutodiffBackend가 사라지고 자동미분 컨텍스트가 디바이스에 설정되어 텐서가 이를 물려받는 형태가 되면서, 사전 조건 검사가 런타임으로 옮겨갔다. 자동미분을 켠다고 모든 텐서가 그래디언트를 요구하는 것은 아니라는 점이 이 구조의 전제다.

컴파일 시간 개선 폭은 두 가지 프로젝트로 측정됐다. 내장 러너를 쓰는 소형 CNN에서는 히든 레이어를 넣고 빼는 편집 기준 릴리스 재빌드 중앙값이 28.42초에서 4.57초로 줄었다. 커스텀 학습 루프를 쓰는 트랜스포머에서는 동등한 피드포워드 표현을 번갈아 적용하는 편집에서 14.73초가 1.00초가 됐다. 릴리스 노트가 제시한 '최대 15배'는 이 트랜스포머 사례에 해당한다. CNN 벤치마크 설정은 64채널 3×3 합성곱 두 개와 BatchNorm, ReLU, 2×2 맥스풀 블록 뒤에 1600→128→128→10 선형 레이어를 GELU와 드롭아웃 0.25로 쌓은 구조이고, 트랜스포머는 프리노름 6층, 히든 512, 피드포워드 2048, 어텐션 헤드 8개, 시퀀스 길이 512, 어휘 28,996, 출력 클래스 4개 구성이다.

커스텀 연산을 붙이는 방식은 크게 바뀌지 않았다. 백엔드 확장 트레이트 정의와 구현은 그대로 두고 #[backend_extension] 매크로로 디스패치 시스템에 연결하면, 백엔드 제네릭 없이 텐서 함수로 노출할 수 있다. output_shape 헬퍼가 커널 실행 전에 출력 형태를 서술하면 매크로가 Fusion을 통한 지연 실행 등록과 디스패치 라우팅을 생성한다. 확장 작성자는 역방향 패스를 직접 구현할 수 있지만, 커스텀 연산은 기본적으로 퓨전 그래프의 경계로 동작한다는 점은 알아둬야 한다. 양쪽 연산은 각자 퓨전될 수 있어도 커스텀 연산과 결합하려면 별도 퓨전 최적화가 필요하다. Burn의 선형대수·신호처리 기능도 burn-linalg, burn-signal 확장 크레이트를 통해 이 방식을 쓴다.

실행 성능과 메모리 관리도 손봤다. 드라이 런에서 수집한 할당 통계로 메모리 풀을 튜닝하는 방식이 CubeCL의 새 메모리 관리로 이어졌고, 합성곱 최적화도 여러 건 들어갔다. 합성곱 그래디언트 계산이 빨라지고 전치 합성곱용 직접 커널이 개선됐으며, 최적화된 역방향 패스를 갖춘 BatchNorm 학습 전용 연산이 추가됐다. 텐서 레이아웃을 다루는 퓨전 처리도 개선돼 불필요한 복사를 피하고, 퓨전 불가 연산 뒤에서 더 많은 연산을 묶을 수 있게 됐다. 새로 추가된 CubeCL Environment는 호환되는 타깃에서 이미 데워진 컴파일·오토튜닝 캐시를 재사용하게 해준다.

CUDA 환경 실측에서는 합성곱 최적화만으로 CNN 벤치마크 학습 처리량이 80% 올랐다. 이 수치는 NVIDIA GeForce RTX 4050 노트북 GPU에서 FP32로, 워밍업 50스텝과 측정 500스텝을 세 번 반복해 얻은 중앙값이다. 같은 조건에서 0.21과 0.22를 비교한 결과이므로, 다른 GPU나 데이터 타입, 다른 모델 구조에서 같은 폭의 향상을 기대하기는 어렵다. 재빌드 시간 벤치마크 역시 두 개의 특정 프로젝트에서 나온 값이라는 전제가 붙는다.

실무 관점에서 가장 체감이 큰 변화는 역시 컴파일 시간이다. 백엔드 제네릭이 사라지면서 모델 코드가 단순해지고, 백엔드를 바꾸거나 여러 개를 동시에 쓰는 작업의 마찰이 줄었다. 다만 자동미분 관련 검사가 런타임으로 이동했으니 기존 코드를 옮길 때 이 차이를 확인해야 하고, 커스텀 커널을 퓨전 파이프라인에 깊이 엮으려면 확장 가이드의 워크플로를 따라야 한다.