Making LLM Training Faster with Unsloth and NVIDIA

Unsloth 와 NVIDIA 를 협력하여 LLM 학습 속도를 약 25% 향상시켰습니다. 이 블로그/가이드에서는 우리가 어떻게 달성했는지 자세히 설명합니다. 이러한 최적화는 정확도 손실이 없으며, Unsloth 의 기존 2~5 배 속도 향상에 대한 추가적인 개선입니다. 새로운 알고리즘은 RTX 노트북, 데이터센터 GPU, DGX Spark 기기에 자동으로 활성화되므로 최신 개선을 얻으려면 Unsloth 를 업데이트만 하면 됩니다. NVIDIA 와 협력함으로써 우리는 다음과 같이 보여줍니다:

여러 짧은 예제가 있다고 가정해 봅시다:

모두 동일한 길이에 패딩하고 패딩 토큰에 컴퓨팅 자원을 낭비하는 대신, 이를 하나의 더 긴 패킹 시퀀스로 연결합니다:

모델은 여전히 각 원래 시퀀스가 어디에서 시작하고 끝나는지 알아야 합니다. 따라서 패킹 토큰과 함께 다음 시퀀스 메타데이터를 전달합니다:

cu_seqlens

이것이 핵심입니다: 고정된 패킹 배치에 대해, 이 메타데이터는 모든 레이어에 동일합니다.

패킹 배치를 다음과 같이 경계 정보를 기록한다고 가정하면:

B = { lengths, cu_seqlens, max_seqlen, mask structure }

그forward pass 의 모든 트랜스포머 레이어는 동일한 B 를 소비합니다.

모델이 L 개의 레이어를 가진다면, 각 레이어마다 B 를 다시 구축하거나 동기화하는 것은 새로운 작업이 아닙니다. 이는 반복적으로 재구성되는 동일한 정보입니다.

즉, 유용한 작업은:

B 를 한 번 구축하고, L 회 사용.

낭비적인 버전은:

B + B + ⋯ + B (L 회)

여기서의 오버헤드는 주로 추가 FLOPs 가 아닙니다. 일부 경로에서는 장치에서 호스트로 동기화를 강제할 수 있으며, 이는 GPU-CPU 동기화 포인트를 효과적으로 생성합니다. 이 것이 레이어별 경로 내부에서 발생하면, 오버헤드는 각 레이어마다 재발생합니다.

이것이 패킹 시퀀스 캐칭 변경이 줄이는 것입니다. 반복적으로 패킹 시퀀스 정보, SDPA 패킹 마스크, xFormers 블록 마스크를 재구성하는 대신, 해당 장치에 대해 현재 패킹 배치별로 재사용 가능한 메타데이터와 그로부터 유도된 attention-side 구조를 캐시합니다. 이 캐시된 구조는 레이어 간에 재사용됩니다.

패킹 학습은 이미 패딩 낭비를 제거함으로써 활용도를 개선합니다. 그러나 메타데이터 경로가 계속 동기화를 강제하면, 일부 이득은 모델의 실제 학습과 무관한 오버헤드로 인해 손실됩니다.

캐싱은 반복적인 조정 작업을 핫 패스에서 제거함으로써 도움을 줍니다. 포워드 패스가 가장 혜택을 받으며, 이는 많은 레이어에 걸쳐 동일한 패킹 메타데이터가 반복적으로 소비되기 때문입니다.

Qwen3-14B QLoRA SFT 에서:

+43.3%

+5.8%

+14.3%

포워드 패스는 반복적인 메타데이터 및 마스크 준비가 가장 직접적으로 나타나는 곳으로 인해 가장 큰 혜택을 보며, 백워드는 개선되지만 효과는 상대적으로 작습니다. 절약된 시간은 비슷하지만, 특히 그라디언트 체크포인트링이 있는 경우 백워드 패스가 더 오래 걸리기 때문에 상대적 이득은 더 작게 보입니다.

측정된 이득을 알았으니 이제 더 간단한 질문을 할 수 있습니다: 해당 스케일링이 타당한가?

각 레이어가 대략적으로 비슷하다고 가정할 때, 패킹된 어텐션 경로를 다음과 같이 모델링할 수 있습니다:

T_uncached ≈ L · (A + s)

여기서:

캐싱을 사용할 경우 해당 반복적인 오버헤드는 배치당 한 번만 지불되며 레이어 당 한 번이 아닙니다:

T_cached ≈ L · A + s

따라서 절약된 시간은 대략 다음과 같습니다:

T_saved ≈ (L − 1) · s

패킹된 SDPA
경로에 대해 NVIDIA Blackwell GPU 의 마이크로 벤치마크는 저수준 호스트 가시성 메타데이터 호출이 실제적이지만 작았으며, 각각 약 0.2 ms
임을 보여줍니다. 지배적인 반복 비용은 합성 패킹 배치에서 총 2048
개 패킹 토큰에 대해 약 13.7 ms
인 패킹된 SDPA
마스크 구성 경로 자체였습니다.

SDPA
백엔드에는 더 나은 정신 모델이 있습니다:

small stream fence + mask rebuild ≈ mask rebuild

이를 통해 더 깔끔한 일관성 검사를 수행할 수 있습니다. 패킹된 마스크 재구성 비용이 m
밀리초라면, 균일 레이어 모델 하에서는:

T_saved ≈ (L − 1) · m

m ≈ 13.7 ms 이므로 이는 다음과 같이 예측합니다:

16
레이어: (16 - 1) x 13.7 ≈ 206 ms

28
레이어: (28 - 1) x 13.7 ≈ 370 ms

더 작은 패킹 시퀀스 실행은 동일한 패턴을 보였습니다:

Llama-3.2-1B, 16 레이어: 단계당 약 199 ms 절약, 이는 엔드 투 엔드 단계 시간의 약 11.5%
낮음.

Qwen3-0.6B, 28 레이어: 단계당 약 319 ms 절약, 이는 엔드 투 엔드 단계 시간의 약 14.8%
낮음.

이 비율은 전체 훈련 단계 시간에 대한 상대적 값이므로 패킹 어텐션 경로 외부의 작업 (임베딩, MLP, LM 헤더, 손실, 프레임워크 오버헤드 등) 을 포함합니다. 이 추정치는 의도적으로 패킹 어텐션 측면에만 해당하며 전체 트랜스포머 레이어가 아닙니다. 이는 측정된 이득이 패킹 SDPA
경로에 대해 올바른 범위인지 확인하기 위한 것뿐입니다.

활성화 체크포인트링은 대형 모델을 훈련하는 데 표준적인 기법입니다. 아이디어는 백워드 패스를 통해 모든 중간 활성화 상태를 유지하지 않고 메모리를 절약하는 것입니다. 대신, 우리는 백워드 동안 일부 추가 작업을 지불합니다.

이 트레이드오프는 특히 더 큰 모델에 대해 일반적으로 가치가 있습니다.

그러나 이는 또 다른 시스템 문제를 제기합니다: 활성화가 오프로드 (offload) 된 경우, 그로백 (backward) 을 위해 어떻게 GPU 로 다시 돌아오는지?

Unsloth 의 스마트 체크포인트링 경로에서, 활성화는 고정된 CPU 메모리 (pinned CPU memory) 에 단계를 거치며 필요할 때 복사됩니다. 이는 VRAM 을 절약하지만 병목 현상을 초래할 수 있습니다:

이는 직렬화 패턴입니다. 하나의 버퍼가 복사 및 계산에 재사용된다면, 복사 스트림과 계산 스트림은 번갈아 가면서 진행합니다.

T_copy 를 활성화 로딩 시간으로, T_compute 를 현재 레이어의 그로백 계산 시간으로 정의합니다.

단일 버퍼를 사용할 경우, 이 단계는 다음과 같이 대략적으로 제한됩니다:

T_single ≈ T_copy + T_compute

이는 직렬화된 경우입니다. 우리는 거의 완전히 복사 및 계산 비용을 모두 지불하며, 하나를 다른 것으로 이어갑니다.

이를 처리하는 더 깔끔한 방법은 두 개의 버퍼를 사용하는 것입니다.

그로백 패스가 버퍼 A 에서 실행되는 동안, 복사 스트림은 다음 활성화를 버퍼 B 로 미리 로드할 수 있습니다. 그런 다음 역할이 바뀝니다. 이는 파이프라인 오버랩을 생성하지만 완벽한 오버랩은 아닙니다.

더블 버퍼링 (double buffering) 은 계산의 양을 줄이지 않습니다. 그것은 유용한 계산 뒤에 복사 지연을 숨깁니다.

이러한 최적화는 그로백 계산을 상당하게 하지만 모든 복사 오버헤드가 소음으로 사라지지 않는 모델이 충분히 크면 더 강해집니다. 더 큰 모델에서는 더 높은 히든 차원 (hidden dimensions) 이 더 많은 데이터 이동을 의미하므로, 이를 숨기는 것이 더 큰 영향을 미칩니다. 더 큰 모델은 또한 더 많은 레이어를 가지며, 계산 뒤에 복사를 숨기는 기회는 더 많아집니다.

더 큰 데인스 (dense) 모델이 이러한 개선에 적합한 이유는 GPU 가 복사 오버랩을 할 수 있는 충분한 실제 작업을 가지고 있으며, 두 번째 버퍼를 위한 추가 VRAM 이 상대적으로 작은 데서 비롯됩니다.

구현 또한 실용적인 가이드레일을 유지합니다:

NVIDIA B200 Blackwell GPUs 로 벤치마크된 더 큰 데인스 모델 실행에서:

8B
: 0.3739 -> 0.4053 steps/s
, +8.40%

14B
: 0.2245 -> 0.2395 steps/s
, +6.70%

32B
: 0.1979 -> 0.2070 steps/s
, +4.61%

메모리 오버헤드는 상대적으로 작았습니다:

8B 에서 +0.37 GB

14B 에서 +0.47 GB

32B 에서 +0.23 GB

이 실행에서 최종 손실은 실질적으로 변경되지 않았습니다.

속도 향상은 더 큰 데인스 모델에 걸쳐 일관성 있고, 추가 VRAM 비용은 상대적으로 작습니다.

측정된 이득을 알았을 때, 자연스러운 후속 질문은: 규모가 합리적인지?

L 개의 체크포인트링 레이어가 있고 각 레이어가 대략적으로 유사하다고 가정하면:

이것은 배치 크기 (batch size), 시퀀스 길이 (sequence length) 및 데이터 이동과 연산에 영향을 미치는 다른 요인들과 함께 확장됩니다. 간결성을 위해 이러한 항들은 생략합니다.

하나의 버퍼 (buffer) 경우:

T_single ≈ L · (c + g)

두 개의 버퍼 (buffers) 경우, 첫 번째 레이어는 여전히 활성화가 도착할 때까지 기다려야 하며, 마지막 레이어는 여전히 계산이 완료될 때까지 기다려야 합니다. 따라서 더 나은 근사식은 다음과 같습니다:

T_double ≈ c + (L − 1) · max(c, g) + g

따라서 절약된 시간은 대략 다음과 같습니다:

T_saved ≈ (L − 1) · min(c, g)

이것은 결과의 유용한 해석입니다:

중간 (middle) 부분의 per-layer 비용은 다음과 같이 훨씬 더 가까워집니다:

T_middle ≈ max(T_copy, T_compute)

측정된 더 큰 모델 (larger-model) 결과에서, 각 훈련 단계 (training step) 당 절약된 시간은 대략 다음과 같습니다:

8B: 약 207 ms

14B: 약 279 ms

32B: 약 222 ms

이 호스트 버퍼 (host buffers) 는 핀 allocations (pinned allocations) 이므로, 관련 대역폭은 pageable-memory 대역폭이 아닌 pinned-memory host-to-device 대역폭을 의미합니다. NVIDIA B200 Blackwell 기반 시스템에서는 해당 대역폭이 약 55.7 GB/s였으며, 비교를 위해 64 GB/s가 유용한 PCIe 상단 (ceiling) 입니다.

추가 버퍼 크기 (extra buffer size) 를 하나의 활성화 재로드 (activation reload) 로 근사적으로 사용한다면, 각 재로드는 자연스럽게 몇 밀리초 (milliseconds) 의 규모에 해당합니다:

8B, 0.37 GB: 약 6.6 ms

55.7 GB/s에서 또는 64 GB/s의 상단에서 5.8 ms

14B, 0.47 GB: 약 8.4 ms

55.7 GB/s에서 또는 64 GB/s의 상단에서 7.3 ms

32B, 0.23 GB: 약 4.1 ms

55.7 GB/s에서 또는 64 GB/s의 상단에서 3.6 ms

관측된 각 단계 당 절약된 시간을 설명하려면, 대략 수십 (dozen) 개의 이러한 재로드를 숨겨야 합니다:

8B: 약 31개, 55.7 GB/s에서 또는 36개, 64 GB/s에서

14B: 약 33개, 55.7 GB/s에서 또는 38개, 64 GB/s에서

32B: 약 54개, 55.7 GB/s에서 또는 62개, 64 GB/s에서

하나의 이러한 재로드를 수십 개의 체크포인트 레이어 (checkpointed layers) 를 통해 숨기면, 절약된 단계 시간 범위는 수백 밀리초 (few-hundred-millisecond range) 로, 우리가 관측한 규모와 정확히 일치합니다.

다시 말하면, 해당 절약된 시간은 전체 엔드 투 엔드 훈련 단계의 일부입니다. 임베딩 (embeddings), LM head, 손실 (loss), 옵티마이저 작업 (optimizer work) 또는 단계의 다른 모든 체크포인트되지 않은 부분 (non-checkpointed part) 을 설명하는 것이 아닙니다. 그 점은 우리가 숨길 수 있는 통신 (communication) 이 측정된 단계 시간 개선 (step-time gains) 을 합리적으로 설명할 만큼 충분히 크다는 것입니다.

세 번째 변경 (change) 은 더 전문적이지만, MoE 라우팅 (MoE routing) 에서 동일한 패턴을 보여줍니다.

PyTorch 기반 GPT-OSS MoE 경로에서 검토한 것 중, 라우팅의 비용이 높은 부분은 토큰이 어느 전문가로 갈지 결정하는 것입니다. 무난한 구현은 다음과 같습니다:

for expert_idx in range(num_experts):
token_idx, _ = torch.where(router_indices == expert_idx)

처음에는 해롭지 않아 보입니다. 하지만 torch.where 는 여기서 데이터에 의존합니다: 각 전문가로 라우팅되는 토큰 수를 배치마다 바꿉니다. 이는 출력 크기가 라우팅 패턴에 따라 달라지기 때문에 CPU-GPU 동기화 또는 관련 런타임 오버헤드를 유발할 수 있습니다. 이것이 한 번씩 전문가마다 발생한다면, 동적 쿼리의 수는 num_experts 와 비례합니다.

더 나은 접근법은 모든 것을 한 번 그룹화하는 것입니다:

bincount 를 사용하여 전문가당 토큰을 구합니다. 수학적으로 이득은 라우팅 논리를 바꾼 것이 아닙니다. 우리는 런타임이 동적 인덱싱 질문을 답할 횟수를 바꾼 것입니다.

대신 대략 다음과 같습니다:
동적 쿼리 오버헤드 ∝ num_experts
왜냐하면 전문가마다 한 번의 동적 쿼리를 수행하기 때문입니다. 우리는 다음과 같이 훨씬 더 가깝게 이동했습니다:
동적 쿼리 오버헤드 ∝ 1
그리고 그 이후에 저렴한 계산을 추가합니다.

이는 더 전문화된 설정에서 동일한 주제입니다: 한 번 그룹화한 후, 반복적으로 동적 토큰 목록을 요청하기 대신 오프셋을 재사용합니다.

이러한 최적화는 native_torch 백엔드를 사용하는 모든 MoE 에 적용됩니다.

이 GPT-OSS 특화된 라우팅 개선을 위해:
GPT-OSS 구성에서 10-15% 속도 향상
전방 및 후방 +23% 와 +13%
이 세 가지 최적화는 스택의 다른 부분에 있지만, 동일한 문제를 해결합니다.

주요 최적화 기회는 주요 커널 주변의 글루 코드에 있었습니다:

이는 개선들이 개념적으로 조합되는 이유입니다. 주요 커널이 빨라지면, 과거에는 보이지 않았던 오버헤드가 전체 단계 시간의 의미 있는 비율이 됩니다.

여기에 유용한 엔지니어링 교훈이 있습니다. 수학 커널이 최적화되면 "더 빠른" 것은 보통 두 가지 중 하나를 의미합니다:

그것이 바로 여기서 일어난 것입니다.