Kernels Hub 통합 및 사용

kernels 라이브러리를 사용하면 최적화된 compute kernel을 Hub에서 직접 로드할 수 있습니다.
전용 org에서 kernels를 찾을 수 있거나, Hub 내에서 kernel 태그를 검색하여 찾을 수 있습니다.

Kernels은 모델 개발, 학습 및 추론을 돕는 최적화된 코드 조각입니다. 여기서는 TRL과의 통합에 초점을 맞추지만, 자세한 내용은 위의 리소스를 참고하세요.

설치

TRL에서 커널을 사용하려면 Python 환경에 라이브러리를 설치해야 합니다:

pip install kernels

TRL에서 Hub의 Kernels 사용

Kernels은 attention 구현을 직접 대체할 수 있어, Flash Attention과 같은 attention 백엔드를 수동으로 컴파일할 필요가 없으며, Hub에서 해당 attention kernel을 가져오기만 하면 학습 속도를 높일 수 있습니다.

모델을 로드할 때 커널을 지정할 수 있습니다:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2"  # other options: kernels-community/vllm-flash-attn3, kernels-community/paged-attention
)

또는 TRL 학습 스크립트를 실행할 때:

python sft.py ... --attn_implementation kernels-community/flash-attn2

또는 TRL CLI를 사용하여:

trl sft ... --attn_implementation kernels-community/flash-attn2

이제 Hub에서 하드웨어 구성에 맞게 사전 최적화된 커널을 갖춘 더 빠른 attention 백엔드를 활용할 수 있어 개발과 학습 속도를 모두 높일 수 있습니다.

커널 버전 선택 및 고정

Hub의 Kernel 저장소는 브랜치(v1, v2, v3, …)로 버전 관리되며, Transformers는 각 저장소에 대해 기본 버전을 선택합니다. 이 기본 버전은 Transformers 릴리스마다 변경될 수 있으므로, 동일한 attn_implementation 값이 항상 동일한 빌드로 해석되지는 않습니다.

로드할 빌드를 제어하려면 리포지토리 ID에 버전 브랜치 또는 커밋 SHA를 추가하세요:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2@v2",  # or a commit SHA to pin an exact build
)

Pinning은 학습 실행의 재현성을 확보하고, 새로운 빌드에서 회귀 문제가 발생했을 때 알려진 정상 빌드를 유지하는 데 유용합니다. 각 버전은 특정 범위 내의 Torch 및 CUDA 버전에만 빌드를 제공하므로, 고정된 버전에는 사용자의 환경에 맞는 변형이 없을 수 있습니다.

attn_implementation="flash_attention_2"는 flash-attn 패키지가 설치되지 않은 경우 kernels-community/flash-attn2 Hub 커널로 폴백되므로, 명시적으로 요청하지 않았더라도 Hub 커널을 사용 중일 수 있습니다.

CUDA 12.8용 v3 빌드는 현재 grouped-query attention을 사용하는 모델에서 backward pass 단계에서 실패합니다. 해당 CUDA 버전을 사용하는 환경이라면 huggingface/kernels-community#1085가 닫힐 때까지 @v2를 고정(pinning)해 두세요.

Attention 구현 비교

transformers에서 사용할 수 있는 다양한 attention 구현 방식과 다른 kernel backends를 TRL과 SFT를 사용하여 평가했습니다.
실험은 H100 GPU 1개와 CUDA 12.9 환경에서 수행되었으며, Qwen3-8B 모델을 batch size 8, gradient accumulation 1, bfloat16 정밀도로 활용했습니다.
여기서 제시된 결과는 해당 설정에 특화되어 있으며, 다른 학습 구성을 사용할 경우 결과가 달라질 수 있습니다.

다음 그림은 다양한 attention 구현 방식과 kernel 백엔드에 대한 latency(학습 단계당 소요 시간)와 peak allocated memory를 보여줍니다.
Kernel 기반 구현 방식은 커스텀 설치된 attention과 동등한 성능을 발휘하며, 모델의 max_length를 늘리면 성능이 더욱 향상됩니다. 모든 구현 방식에서 메모리 사용량은 유사하여 유의미한 차이가 없습니다. 다음 섹션에서 설명하듯, 동일한 성능을 더 적은 마찰 없이 얻을 수 있습니다.

Flash Attention과 Hub Kernels

Flash Attention을 소스에서 빌드하는 데에는 하드웨어, CUDA/PyTorch 구성, 사전 컴파일된 wheel의 유무에 따라 수 분에서 수 시간까지 걸릴 수 있어 시간이 많이 소요될 수 있습니다.

반면, Hugging Face Kernels는 훨씬 빠르고 안정적인 워크플로를 제공합니다. 개발자는 복잡한 설정을 걱정할 필요가 없으며, 모든 것이 자동으로 처리됩니다. 저희 벤치마크에서 커널은 컴파일 없이 약 2.5초 만에 사용 가능한 상태로 준비되었습니다. 이를 통해 거의 즉시 학습을 시작할 수 있어 개발 속도를 크게 높일 수 있습니다. 원하는 버전을 지정하기만 하면 kernels가 나머지를 처리합니다.