AI Runtime에서 Qwen 학습 시간 6배 단축

수십억 개의 파라미터를 가진 파운데이션 모델을 학습하거나 파인튜닝하려면 수백에서 수천 시간에 달하는 GPU 학습 시간이 필요할 정도로 상당한 컴퓨팅 자원이 소요됩니다. 하지만 LLM의 학습 성능을 튜닝하는 것은 까다로우며, GPU 사용률과 같은 상위 수준 지표가 실제 사용률을 정확히 보여주지 않을 수도 있습니다. 이번 블로그 게시물에서는 추상화 레이어를 걷어내고 AI Runtime에서 최신 Qwen 3.8 27B 모델의 파인튜닝 속도를 6배 향상하는 방법을 소개합니다. 체계적인 성능 최적화를 통해 이미 98%에 달했던 GPU 사용률 환경에서도 고정된 데이터셋 기준 파인튜닝 시간을 7시간 44분에서 1시간 17분으로 단축했습니다.

중요한 지표를 측정하는 방법부터 구현 세부 사항, 각 최적화 과정에서의 설계 트레이드오프에 이르기까지 이러한 성능 향상에 적용된 기술을 심도 있게 다룹니다.

지금 바로 시작하고 싶으시다면 AI Runtime의 예제 코드로 이동해 직접 시도해 보세요!

모델 설정

transformers 라이브러리와 SFT 학습 데이터셋인 HuggingFaceH4/ultrachat_200k를 사용하여 Qwen 3.8 27B의 표준 학습 루프부터 시작합니다.

• torch의 최신 모델 샤딩 기술인 FSDP2를 적용합니다. 이는 모델 가중치와 옵티마이저 상태를 GPU에 올리기 위해 필요합니다. 순수한 분산 데이터 병렬 처리 방식으로는 이 정도 크기의 모델 가중치를 단일 GPU 메모리에 올릴 수 없습니다.
• 또한 역전파 시 컴퓨팅과 메모리 공간을 맞바꾸어 GPU당 더 큰 배치 크기를 가능하게 하는 기술인 활성화 체크포인팅(activation checkpointing)을 활성화합니다.
• 데이터로더(dataloader)는 배치당 최대 8192개 토큰으로 2D 배치를 패킹하고, 배치 샘플의 최대 크기에 맞게 패딩을 적용합니다.

초기 학습 벤치마크

전체 데이터셋에 대한 모델의 베이스라인 학습은 7시간 44분이 소요되었으며, 학습 중 평균 GPU 사용률은 98%로 나타났습니다. 그 결과 처리 속도는 초당 9,396개 토큰(9,396 tok/s)이었습니다.

이 GPU 사용률 수치는 높아 보이지만, 앞서 언급했듯이 오해의 소지가 있습니다. 표준 GPU 사용률 지표는 GPU가 어떤 작업이든 수행하고 있던 시간의 비율만 측정하므로, 활성화되어 있는 동안 모델이 GPU를 얼마나 효율적으로 사용했는지는 알려주지 않습니다. 이를 이해하려면 더 자세히 살펴보아야 합니다.

이를 확인하는 한 가지 방법은 torch trace 메커니즘을 사용하는 것이며, 다음과 같이 학습 루프를 조정할 수 있습니다(빠른 결과를 얻기 위해 스텝 수도 제한합니다).

샘플 스텝에 대해 다음과 같은 데이터가 도출됩니다.

표 1: CPU - 1,714,382회 호출, 14.4초. 총 소요 시간 기준 상위 연산 목록입니다.

표 2: GPU - 269,448개 연산, 10.3초. 총 소요 시간 기준 상위 연산 목록입니다.

개선할 영역

이 표들을 보면 현재 모델에서 몇 가지 상위 수준의 문제를 확인할 수 있습니다.

• CPU가 각 학습 스텝을 설정하는 데 걸리는 시간이 GPU가 실행하는 데 걸리는 시간보다 깁니다. 일반적으로 GPU 실행 과정에서 CPU 작업은 비동기식으로 진행되어 CPU가 GPU보다 훨씬 앞서 실행되어야 하므로 이는 예상치 못한 결과입니다. 주요 영향으로 GPU 지연(stalling)이 발생하므로 이를 개선할 여지가 있습니다.
• 연산 수가 너무 많습니다! 위 표에서 수십만 개의 자잘한 연산을 볼 수 있는데, 이 경우 문제는 CPU가 GPU를 지속적으로 작동시킬 만큼 빠르게 GPU 연산을 대기열에 추가(enqueue)하지 못한다는 점입니다. 핵심 아이디어는 CPU가 GPU로 보내는 연산 수 자체를 줄여야 한다는 것입니다. 이를 위한 일반적인 기술들이 있으며, 그 중 하나가 현재 상황에 적합합니다.
  • torch의 CUDA 그래프는 실행된 CUDA 커널을 기록하고 단일 CPU 호출로 이를 재실행합니다. 강력한 기능이지만 여러 제약이 있고 FSDP2의 동적 특성과 충돌하므로 지금은 고려 대상에서 제외합니다.
  • Torch.compile은 코드를 추적하고 Triton을 사용해 가능한 경우 커널을 자동으로 융합(fuse)합니다. 본 예제에서는 수동으로 융합된 커널을 사용할 것이며, LLM의 경우 torch.compile로 추가적인 개선 효과를 얻기는 어렵습니다.
  • 수동으로 융합된 커널(Manually fused kernels)은 torch 컴파일러가 스스로 적용할 수 없는 특정 연산에 맞춰 특화된 융합 기법을 사용합니다. 다행히 커뮤니티에는 이미 검증된 수많은 사전 구축 커널이 존재합니다. 대표적인 예로 ‘Flash Attention’을 들 수 있으며, 본 사례에서도 소폭의 성능 향상을 가져다줍니다.

또한 배치에 불필요한 패딩이 과도하게 포함되는 문제도 있습니다. 입력 텐서는 직사각형 형태여야 하므로 각 행이 해당 배치의 최대 토큰 수에 맞춰 패딩됩니다. 뒤에서 이를 최소화할 수 있는 LLM용 유용한 팁을 소개합니다.

다음 섹션에서는 각 기술적 주제에 대해 개념, 권장되는 모범 사례, 적용 시 고려해야 할 주요 트레이드오프를 자세히 살펴보겠습니다.

융합 커널(Fused kernels)

일부 커널이 매우 자주 실행되는 것을 확인하셨을 것입니다. 특히 Qwen의 경우 배치 행렬 곱셈(batch matrix multiplication) 커널이 한 스텝에서 30,000회 이상 실행됩니다.

이러한 경우 먼저 융합 커널을 살펴보는 것이 좋습니다. 융합 커널은 여러 개의 작은 GPU 커널을 하나의 단위로 결합합니다. 이를 통해 GPU의 고대역폭 메모리와 온다이(on-die) 메모리 간의 비싼 데이터 전송을 줄여 성능을 향상할 수 있습니다. 또한 GPU 코어를 보다 효율적으로 스케줄링할 수 있게 해줍니다.

이렇게 많은 추가 커널이 있는 경우 torch.compile이 동작할 수는 있지만 실행하는 데 오랜 시간이 걸립니다. 따라서 수동으로 융합된(fused) 커널을 사용해 순수 연산 수(ops)를 줄이는 것이 좋은 첫 단계입니다. 또한 Flash Attention과 같이 잘 알려진 퓨전은 torch.compile 자체만으로는 찾아내지 못합니다.

transformers 라이브러리는 기본적으로 표준 어텐션 레이어에 torch의 SDPA를 사용합니다. 하지만 우리가 확인한 소규모 연산 중 상당수는 비표준 게이트드 델타 네트워크(gated delta network) 레이어에서 발생하는 것으로 나타났습니다. 다행히 이러한 레이어를 위해 레이어당 여러 연산을 융합하도록 작성된 causal-conv1d 및 flash-linear-attention이라는 두 가지 커널이 존재합니다. 이를 어떻게 찾아낼 수 있을까요? Qwen의 transformers 라이브러리 구현에서 지원하는 커널을 이미 밝히고 있습니다.

• flash-linear-attention
• causal-conv1d

이러한 커널을 구하는 고수준 방법에는 두 가지가 있습니다.

• 커널을 직접 설치합니다. 방법은 다음과 같습니다.
  • PyPI에서 간단한 pip install(예: flash-linear-attention용)
  • 관련 리포지토리 git clone 및 컴파일
  • 컴파일 시간이 너무 오래 걸릴 경우 사전에 컴파일된 패키지의 다른 출처 찾기
• Hugging Face의 Kernel Hub를 사용하여 사전 컴파일된 커널 가져오기.

저희는 일반적으로 첫 번째 방법을 권장하며, 특히 당사의 Docker 컨테이너 지원과 함께 빌드 시점에 모든 커널을 설치하여 컨테이너를 빌드하면 작업 시작에 걸리는 시간을 줄일 수 있습니다.

Kernel Hub는 편리하지만, 커널이 설치된 torch 및 CUDA 버전에 명시적으로 의존하는 경우가 많으므로 특정 구성에 맞는 커널을 사용하지 못할 수 있으며, 이 경우 transformers 모델은 경고 없이 성능이 낮은 미최적화 커널로 폴백(fallback)됩니다.

flash-linear-attention 및 causal-conv1d Python 패키지를 설치한 후 엔드 투 엔드 실행 시간을 3.7배 향상시켜 2시간 5분으로 단축했습니다. 아래 표에서 이전 표의 elementwise 함수 대신 새로 도입된 커널(강조 표시됨)을 확인할 수 있습니다. 또한 실행된 전체 연산 수도 15배로 대폭 감소했습니다.

표 3: CPU - 118,316회 호출에 3.9초. 전체 시간 기준 상위 연산이 포함된 표입니다.

표 4: GPU - 18,060회 호출에 3.9초. 전체 시간 기준 상위 연산이 포함된 표입니다.

주의 사항: 학습 중에는 순전파(forwards) 및 역전파(backwards) 패스 모두에 대한 커널이 필요하지만, 일부 커널은 추론 속도 향상만을 목적으로 구현되어 순전파 전용일 수 있습니다.

핵심 요약

• 일반적인 모델에 즉시 적용할 수 있는 검증되고 고도로 최적화된 커널 사용을 고려해 보세요. 이러한 커널을 활용하면 비교적 적은 노력으로 대폭적인 성능 향상을 얻을 수 있습니다. 이번 사례에서는 학습 속도가 3.7배 향상되어 학습 시간이 7시간 44분에서 2시간 5분으로 단축되었습니다.

배치 처리, 토큰화 및 패딩

이제 Qwen 모델의 속도가 꽤 빨라졌지만, 더욱 크게 향상시킬 수 있습니다. 기본적으로 당사의 배치 설정은 메모리 부족 없이 수용할 수 있는 최대 크기인 8,192개 토큰의 2D 직사각형 텐서를 사용합니다. 이 접근 방식의 단점은 패딩으로 인해 텐서 내부에서 상당한 양의 공간이 낭비된다는 점입니다.

Qwen이 표준 어텐션 메커니즘만 사용했다면, 현재 잘 지원되는 torch의 NestedTensor를 사용하여 패딩을 완전히 제거함으로써 이러한 비용을 줄일 수 있었을 것입니다. 하지만 flash-linear-attention은 더 전통적인 1D 패킹된 텐서 방식만 지원합니다.

그러나 이를 임의로 처리할 수는 없습니다. 단순히 1D 텐서를 모델에 전달하는 것은 사전 학습(pretraining)에는 괜찮을 수 있지만, 파인 튜닝(finetuning) 시에는 한 예시 안의 토큰이 동일한 예시 내의 토큰에만 어텐션을 수행하도록 보장해야 합니다. 이를 해결하려면 올바른 예시 position_ids를 모델에 전달해야 합니다. 그러면 모델이 이를 융합된 어텐션 커널로 전달하여 올바른 어텐션 마스크를 적용합니다. 여기에 있는 예시 코드를 참고하시면 이것이 작동하도록 필요한 입력을 계산하는 방법을 확인하실 수 있습니다.

단순해 보이는 이러한 변경만으로 학습 시간이 2시간 5분에서 1시간 17분으로 단축됩니다.

원인은 두 가지입니다.

1. 주요 원인은 이러한 패킹을 통해 배치당 시퀀스 수가 증가하고 전체 배치 수가 약 36% 감소하기 때문입니다. 이는 데이터셋의 예시 길이 및 분포에 따라 달라질 수 있습니다.
2. 두 번째 원인은 flash linear attention 커널이 각 입력 형태(shape)에 대해 JIT(Just-In-Time) 컴파일되어야 하는데 이전에는 형태가 동적이어서 몇 분의 추가 컴파일 시간이 발생했기 때문입니다.

트레이드오프: 이러한 유형의 변경은 입력 파이프라인을 더 복잡하게 만들고 모델 품질에 영향을 미치는 오류가 발생할 가능성을 높입니다. 하지만 이번의 경우 36%의 절감 효과는 충분한 가치가 있습니다. 원래의 배치 메커니즘을 유지하면서 모든 입력을 길이에 따라 정렬하여 패딩을 최소화하는 등 패딩을 줄이는 다른 기법도 존재하지만, 이러한 정렬 순서는 인접한 샘플 간에 예기치 않은 상관관계를 생성하여 모델에 편향을 줄 수 있습니다.

핵심 요약: 최대 학습 가능 토큰 수(예: 이번 사례의 경우 8,192)보다 짧은 입력이 포함된 데이터셋으로 LLM을 학습할 때, 데이터를 배치로 효율적으로 패킹하면 상당한 시간을 절약할 수 있습니다. 이 데이터셋의 경우 학습 속도가 1.6배 향상되어 학습 시간이 2시간 5분에서 1시간 17분으로 감소했습니다.

결론

본 블로그에서는 학습 효율을 6배 향상시키는 방법을 살펴보았습니다. 핵심은 (1) 단순한 GPU 사용률뿐만 아니라 모델이 학습 중 어디에 시간을 소비하는지 지표를 먼저 측정하고, (2) Flash Linear Attention과 같은 알려진 라이브러리를 사용하여 커널을 최적화 및 융합하며, (3) 데이터 배치가 모델에 제공되는 방식을 맞춤화하는 것이었습니다.

이를 대규모로 수행하려면 고성능 경로를 기본으로 제공하는 플랫폼이 필요합니다. AIR에는 이러한 수많은 성능 개선 사항이 기본적으로 내장되어 있습니다. 기본 환경에 최신 CUDA 및 flash attention을 번들로 제공하며, GPU 사용률, 온도, 상태와 같은 핵심 지표를 자동으로 추적하고, torch 버전(및 프로파일러)이 최신 상태로 유지되도록 보장합니다. 내부 에이전트인 Genie Code 역시 모델 학습을 한층 더 최적화할 수 있는 이러한 기회를 찾도록 조율되어 있습니다. 또한 Databricks에서 데이터를 읽는 작업은 UCVolumeDataset을 통해 최적화됩니다.

다음 학습 워크로드에 AIR를 꼭 활용해 보세요! 바로 사용해 보시겠어요? 이 모델을 직접 실행하는 방법에 대한 자세한 내용은 AI Runtime 쿡북을 참조하세요.

(이 글은 AI의 도움을 받아 번역되었습니다. 원문이 궁금하시다면 여기를 클릭해 주세요)