Burn 0.22.0: 더 빠른 빌드, 쉬워진 확장, 그리고 스마트해진 자동 튜닝

Burn 0.22는 애플리케이션 구축을 용이하게 하고, 커스텀 커널(custom kernels) 통합 및 모델 활용에 도움을 줍니다. 이전 릴리스에서 저희는 burn-dispatch를 도입했으며, 사용자 API에서 백엔드 제네릭(backend generics)을 제거하는 작업을 시작했습니다. 이 목표는 애플리케이션 코드 간소화와 컴파일 시간 개선 두 가지였습니다. 이번 릴리스에서는 코드가 더 이상 백엔드 타입 매개변수를 포함하지 않으며, 대신 장치(device)를 통해 실행이 선택됩니다. 저희 테스트에서 이러한 변경 사항들은 재빌드 속도를 최대 15배까지 빠르게 만들었습니다.

API 변경 외에도, 적응형 메모리 관리(adaptive memory management), 스마트한 자동 튜닝(smarter autotuning), 그래프 리플레이(graph replay), 컴파일러 및 커널 기반의 변경 사항을 통해 CPU와 GPU 전반에 걸친 실행이 개선되었습니다. 새로운 CubeCL Environment는 애플리케이션이 호환되는 타겟에서 이미 워밍업된(warmed) 컴파일 및 자동 튜닝 캐시를 재사용할 수 있게 합니다. 또한 이번 릴리스에서는 LoRA 및 QLoRA를 사용한 기존 모델 적응, ONNX 내보내기(export), 그리고 원격 컴퓨팅에 대한 지원 확대를 추가했습니다. 이 게시물에서 주요 변경 사항을 다룹니다. 개선 및 수정된 전체 목록은 릴리스 노트에서 확인할 수 있습니다.

더 간단해진 모델, 더 빠른 빌드

지금까지 백엔드를 선택하는 것이 Burn 애플리케이션 전반의 타입을 영향을 미쳤습니다. 모델은 레이어부터 해당 매개변수에 이르기까지 B: Backend 타입을 가지고 있었고, 텐서 함수도 이를 전파해야 했습니다. 이는 유연한 아키텍처를 제공했지만, 컴파일 과정에서 애플리케이션 코드를 거대한 의존성 사슬에 묶이게 만들기도 했습니다.

0.22에서는 이 백엔드 매개변수가 사용자 대상 API에서 사라집니다. 변경 사항은 작은 모델에서도 확인할 수 있습니다:

#[derive(Module, Debug)]
-pub struct Model<B: Backend> {
- linear: Linear<B>,
...

이제 장치가 실행 방식과 위치를 선택합니다. 해당 Cargo 기능을 활성화한 다음, 모델 및 입력값을 초기화할 때 장치를 선택하면 됩니다:

use burn::tensor::Device;
let device = Device::cuda(0);
let device = Device::wgpu(Default::default());
...

훈련도 동일한 패턴을 따르므로, B: AutodiffBackend

이전에는 애플리케이션 코드에서 더 이상 존재하지 않습니다. 대신, autodiff 컨텍스트는 디바이스에 구성되며 텐서로 상속됩니다.

그 아래에서 연산은 텐서 → 브릿지 → 디스패치 → 백엔드 순서를 따릅니다. 이 브릿지는 고수준 텐서 API로부터 구체적인 백엔드 표현을 숨겨, 이전에는 애플리케이션 코드까지 도달하던 의존성 체인을 끊습니다. 이러한 타입 소거(type erasure)는 프로젝트에 변경 사항이 있을 때 리컴파일을 줄여줍니다. 반면 디스패치(dispatch)는 각 연산을 해당 백엔드로 라우팅하여, 단일 애플리케이션이 여러 백엔드를 쉽게 활성화할 수 있도록 합니다. Backend

trait는 백엔드 구현과 사용자 정의 연산을 지원하기 위해 Burn의 핵심에 남아 있습니다.

이것이 컴파일 시간에 어떤 영향을 미치는지 보여주는 예를 들자면, 우리는 두 개의 별도 프로젝트에서 모델 변경을 테스트했습니다. 하나는 Burn의 내장 학습기(learner)를 사용하는 작은 CNN이고, 다른 하나는 사용자 정의 훈련 루프를 사용하는 트랜스포머입니다. CNN 벤치마크에서 은닉층을 추가하거나 제거하자 중앙값 배포 재빌드 시간이 28.42초에서 4.57초로 줄었습니다. 트랜스포머의 경우, 동등한 피드포워드 표현(feedforward expressions)을 번갈아 사용하자 재빌드 시간이 14.73초에서 1.00초로 줄었습니다. 이 두 가지 설정으로 우리는 다양한 모델과 의존성 흔적에 걸쳐 컴파일이 어떻게 개선되는지 확인할 수 있습니다.

벤치마크 구성

CNN: 64채널 3×3 컨볼루션, BatchNorm, ReLU, 그리고 2×2 max-pool 블록 두 개로 구성되며, 이후 GELU와 드롭아웃(dropout) 0.25가 적용된 1600→128→128→10 선형 레이어로 이어집니다. 배치 크기는 256, 옵티마이저는 AdamW, 학습률은 0.001, 가중치 감쇠(weight decay)는 0.00005, 시드(seed)는 42입니다.트랜스포머: 전처리(pre-norm) 레이어 6개, 은닉 크기(hidden size) 512, 피드포워드 크기(feedforward size) 2048, 어텐션 헤드 8개, 시퀀스 길이(sequence length) 512, 어휘 크기(vocabulary size) 28,996, 출력 클래스는 4개이며 드롭아웃은 없습니다. 배치 크기는 8, 옵티마이저는 Adam, 학습률은 0.0001, 시드(seed)는 42입니다.

MNIST는 Burn의 트레인(train) 기능을 활성화하고 그 학습기(learner)를 사용합니다. 변환기(transformer)는 사용자 정의 훈련 루프를 사용합니다. 모델 수정 사항: MNIST용으로 GELU와 드롭아웃을 가진 은닉 선형 레이어 추가 또는 제거, 그리고 변환기의 동등한 피드포워드 표현 교대. 막대와 라벨은 세 번의 실행에 걸친 중앙값을 보여줍니다. 오차 막대는 관찰된 최소-최대 범위를 보여줍니다. 각 메트릭은 0부터 시작하는 자체 선형 스케일을 가집니다.

더 쉬워진 사용자 정의 연산 (Easier Custom Operations)

사용자 정의 연산은 애플리케이션 코드가 더 이상 백엔드 제네릭(backend generics)을 포함하지 않더라도 여전히 백엔드 기본 요소(backend primitives)를 사용하여 구현됩니다. 이미 백엔드 확장 트레이트(backend extension trait)가 있었다면, 그 정의와 구현은 거의 동일하게 유지됩니다. #[backend_extension] 매크로는 이를 디스패치 시스템에 연결하므로, 백엔드 제네릭 없이 텐서 함수를 통해 연산을 노출할 수 있습니다. 여기는 행렬 곱셈(matrix multiplication), 바이어스 추가(bias addition), ReLU를 결합한 사용자 정의 CubeCL 커널 예시 발췌입니다:

use burn::{
backend::{Dispatch, backend_extension, tensor::FloatTensor},
tensor::Tensor,
...}

output_shape 헬퍼는 커널이 실행되기 전의 출력을 설명합니다. 이 메타데이터를 통해 매크로는 디스패치 라우팅과 함께 Fusion을 통한 지연 실행(lazy execution)에 대한 등록을 생성합니다. 전체 예제에는 커널과 그 역방향 구현(backward implementation)이 포함되어 있어, 호출자는 훈련 및 추론에 동일한 함수를 사용할 수 있습니다.

이를 통해 확장 작성자들은 사용자 정의 역전파 과정(custom backward pass)을 포함하여 구현에 대한 제어권을 갖게 됩니다. 기본적으로 사용자 정의 연산은 Fusion 그래프에서 경계 역할을 합니다. 양쪽의 연산은 여전히 융합될 수 있지만, 이를 사용자 정의 연산과 결합하려면 전용 Fusion 최적화가 필요합니다. 확장 가이드에서는 전체 워크플로우를 설명합니다. Burn의 선형 대수(linear algebra) 및 신호 처리(signal processing) 기능은 이제 burn-linalg와 burn-signal 확장 크레이트들을 통해 이 접근 방식을 사용합니다.

더 빠른 실행, 더 나은 메모리 관리 (Faster Execution, Better Memory Management)

Burn 기반의 새로운 AI 엔진을 구축하는 과정에서 (자세한 내용은 What's Next 섹션 참조), 저희는 드라이 런(dry run) 중에 수집된 할당 통계(allocation statistics)를 사용하여 메모리 풀(memory pools)을 조정하는 방식을 개발했습니다. 이 작업은 CubeCL의 새로운 메모리 관리 방식의 기반이 되었습니다. 또한, 여러 컨볼루션 최적화(convolution optimizations)를 포함하여 CubeCL과 CubeK 전반에 걸쳐 성능을 개선했습니다.

영향을 측정하기 위해 동일한 두 프로젝트와 FP32 환경에서 NVIDIA GeForce RTX 4050 Laptop GPU의 CUDA를 사용하여 Burn 0.21과 0.22 두 버전을 비교했습니다. 각 워크로드(workload)는 버전별로 세 번 실행되었으며, 50개의 워밍업 스텝(warmup steps) 후 500개의 측정 스텝(measured steps)을 거쳤습니다. 그래프는 이 실행들 전체의 중앙값(medians)을 보여주며, 오차 막대(error bars)는 관찰된 최소-최대 범위(min-max ranges)를 나타냅니다.

FP32 환경에서 NVIDIA GeForce RTX 4050 Laptop GPU를 사용했으며, 실행당 50개의 워밍업 스텝과 500개의 측정 스텝을 거쳤습니다. 스텝 시간은 각 실행 내의 평균값(mean)입니다. 막대와 레이블은 세 번의 실행 전체 중앙값을 보여줍니다. 오차 막대는 관찰된 최소-최대 범위를 나타냅니다. 각 지표는 0에서 시작하는 자체 선형 눈금(linear scale)을 가집니다.

컨볼루션 최적화 덕분에 CNN 벤치마크에서 학습 처리량(training throughput)이 80% 증가하는 데 도움이 되었습니다. 이 작업에는 더 빠른 컨볼루션 기울기 계산(convolution gradient computation), 전치 컨볼루션(transposed convolution)을 위한 개선된 직접 커널(direct kernel), 그리고 최적화된 역방향 전달(backward pass)을 갖춘 전용 BatchNorm 학습 연산이 포함됩니다.

또한, 저희는 융합(fusion)이 텐서 레이아웃(tensor layouts)을 처리하는 방식을 개선하여 불필요한 복사본 생성을 방지하고, 융합할 수 없는 연산 이후에도 더 많은 연산을 융합할 수 있도록 했습니다.

워크로드에 적응하는 메모리

CubeCL은 이미 풀(pools)을 통해 장치 메모리(device memory)를 재사용하여 반복적인 할당을 방지했습니다. 하지만 풀 크기가 워크로드와 일치하지 않을 경우, 텐서가 필요로 하는 것보다 더 많은 메모리를 예약할 수 있습니다. 새로운 적응형 메모리 풀(adaptive memory pools)은 사용자가 각 모델에 대해 풀 크기를 구성할 필요 없이, 워크로드 주변으로 할당 크기를 조정합니다.

더 큰 할당이 도착함에 따라, 풀은 페이지 크기를 조정하고 더 이상 비어 있게 된 오래된 페이지를 해제합니다. 런타임은 또한 활성 할당을 이동시켜 더 빨리 해제할 수 있으며, 자주 변경되는 작은 할당은 별도의 풀에 유지합니다. 메모리 보고서에는 사용 중인 메모리와 예약된 메모리가 모두 표시되어, 작업 부하가 실제로 얼마나 많은 예약 공간이 필요한지 확인할 수 있습니다. 이는 이번 릴리스에서 더 많은 CubeCL 유틸리티를 노출하는 Burn의 Device API를 통해 사용할 수 있습니다.

워밍업(warmup) 이후에는 각 실행에서 중앙값 및 최대 VRAM 사용량이 동일했습니다. 임시 시작 대역폭 프로브 할당은 제외되었습니다. 막대와 레이블은 세 번의 실행에 걸친 중앙값을 보여줍니다. 오류 막대는 관찰된 최소-최대 범위를 보여줍니다. 각 측정 항목은 0에서 시작하는 자체 선형 스케일을 가집니다.

벤치마크가 보여주듯이, 워밍업 후 CNN의 최대 VRAM 사용량은 956 MiB에서 486 MiB로 떨어져 거의 절반 수준입니다. 트랜스포머(transformer)의 경우, 중앙값 최대 VRAM 사용량은 3,486 MiB에서 2,868 MiB로 감소하여 17.7%가 줄었습니다. 이 측정값들은 워밍업 후 프로세스의 VRAM을 나타냅니다. 장치 메모리 대역폭을 측정하고 자동 튜닝(autotuning)을 안내하는 데 사용되는 CubeCL의 임시 시작 할당은 제외되었습니다.

합성곱 모델(convolutional models)의 경우, 개선 사항이 학습과 추론 모두에 걸쳐 적용됩니다. 합성곱 기울기(Convolution gradients)는 더 효율적인 실행 경로를 가지게 되었고, 전치 합성곱(transposed convolution)에는 전용 직접 커널(direct kernel)이 생겼으며, BatchNorm 학습은 폐쇄형 공식 역전파(closed-form backward pass)를 가진 전용 연산을 사용합니다.

더 스마트한 자동 튜닝 및 런타임 실행

이전 릴리스 게시물에서 우리는 하드웨어 처리량(throughput)을 사용하여 언제 자동 튜닝을 중지할지 결정하는 것에 대해 논의했습니다. 적응형 자동 튜닝 스케줄러는 측정된 컴퓨팅 및 메모리 처리량을 사용하여 연산에 대한 성능 목표를 추정합니다. 후보가 예상되는 성능에 도달하면, 튜너는 탐색을 중단할 수 있습니다. 또한 약한 후보들을 조기에 제거하고 각 후보가 필요한 측정량을 조정합니다.

반복 실행에 대한 오버헤드를 줄이기 위해 CUDA와 HIP에서 그래프 캡처 및 리플레이(replay) 지원을 추가했습니다. 이를 통해 CubeCL은 커널 호출 시퀀스를 재사용할 때, 호출 준비 과정의 많은 부분을 반복하지 않고도 재사용할 수 있습니다. WGPU의 경우, 해결된 파이프라인(pipelines)과 바인딩(bindings)을 유지하는 소프트웨어 그래프를 통해 리플레이가 지원됩니다.

CubeCL 환경 및 배포 캐시

CubeCL은 이미 실행 간에 컴파일된 커널(kernels)과 자동 튜닝(autotuning) 결과를 캐싱했지만, 새로운 CubeCL Environment를 중심으로 이 캐싱 기능을 재구성했습니다. 이 환경은 해당 요소들을 이름이 지정된 그룹으로 묶습니다. 사용자는 환경의 캐시를 사전에 채우고, 이를 함께 내보내어 애플리케이션과 함께 배포할 수 있습니다. 또한 이 환경은 네이티브 애플리케이션, WebAssembly, 그리고 no_std 전반에 걸쳐 런타임 구성 및 캐시 저장소에 대한 공통 API를 제공합니다.

또한 Burn과 CubeCL 전반의 프로파일링 지원을 확장하여, 실행 시간이 어디에 사용되는지 식별하는 데 도움을 주기 위해 CUDA와 HIP에서의 장치 타이밍(device timing)을 포함했습니다.

새로운 컴파일러 및 커널 기반

커스텀 연산(custom operations)과 이식 가능한 실행(portable execution)에 대한 작업은 스택 깊숙한 곳에서도 변경을 필요로 했습니다. CubeCL은 컴파일러 인프라를 Pliron으로 마이그레이션하여, 커널을 표현하고, 변환하며, 낮추는(lowering) 공통 기반을 확보했습니다. 그 LLVM 컴파일러는 이제 AMD와 NVIDIA GPU용 타겟을 포함하여, LLVM 경로를 CPU 실행을 넘어 확장했습니다.

CubeK의 타일 추상화(tile abstraction)는 또 다른 과제인 복잡한 커널에 대응합니다. 복잡한 커널은 종종 데이터를 로드하고, 메모리에 스테이징하며, 그 위에서 계산하고, 결과를 쓰는 동일한 논리 조각들을 필요로 합니다. 이러한 조각들을 재사용 가능하게 만듦으로써, 커널 작성자들은 다양한 레이아웃(layouts), 정밀도(precisions), 하드웨어 기능에 대한 구현을 구성할 수 있게 되었습니다.

이 작업은 이미 타일 행렬 곱셈(tiled matrix multiplication), 어텐션(attention), 그리고 보간(interpolation) 구현에서 사용되고 있습니다. 또한 이는 해당 커널들을 확장하고 최적화하기 위한 더 나은 기반을 제공합니다. 백엔드 확장 API와 결합하여, 애플리케이션 수준의 사용자 정의 기능을 스택 아래쪽의 재사용 가능한 컴퓨팅 빌딩 블록과 연결해 줍니다.

LoRA 및 QLoRA를 이용한 파인튜닝

Burn 0.22는 LoRA와 QLoRA를 사용하여 기존 모델을 새로운 작업에 적응시키는 기능을 추가했습니다. LoRA는 기본 가중치(base weights)를 고정하고 저랭크 어댑터(low-rank adapters)를 학습시켜, 그래디언트와 옵티마이저 상태가 필요한 파라미터 수를 줄입니다. QLoRA 역시 기본 가중치를 양자화(quantized)하여 저장 요구 사항을 줄입니다.

어댑터를 기존 모델의 레이어나 순전파 과정(forward pass)을 변경하지 않고도 부착할 수 있습니다. 자동 미분 장치(autodiff device)에서 모델을 초기화하고, 적응시킬 파라미터들을 선택한 후 LoRA를 적용합니다:

use burn::module::{Lora, Module, ParamGroup};
let attention = ParamGroup::from_predicate("attention");
let model = model.apply_lora(...);