더 큰 ubatch를 통해 나의 RTX 3090에서 gpt-oss-120b의 프롬프트 처리 속도를 훨씬 빠르게 만들었다

나는 24 GB RTX 3090에서 llama.cpp를 사용하여 gpt-oss-120b-F16.gguf를 튜닝하던 중, 실행을 VRAM 내에 유지할 수 있을 만큼 --n-cpu-moe를 충분히 높여준다면 물리적 마이크로 배치 크기 (-ub)를 늘리는 것이 프롬프트 처리 처리량 (throughput)을 대폭 향상시킬 수 있다는 것을 발견했다.

lama.cpp의 기본값은 -b 2048 및 -ub 512이며, 차트에는 이 기본 실행 값을 별도의 항목으로 포함했다.

다음은 내가 기록한 비공식적인 llama-bench 결과이다:

ubatch	n-cpu-moe	prefill	generation
256	25	240.03 tok/s	33.14 tok/s
...

lama.cpp 기본값인 -ub 512와 비교했을 때, 프롬프트 처리는 약 380 tok/s에서 약 2091 tok/s로 증가하여 대략 5.5배의 이득을 얻었다. 더 작은 -ub 256 실행과 비교했을 때는 약 8.7배의 이득이었다. 토큰 생성 (Token generation) 속도는 기본 설정에서의 약 32.3 tok/s에서 -ub 8192에서의 30.1 tok/s로 약 7% 감소했다.

주의할 점은 더 큰 ubatch가 더 많은 GPU 연산 작업 공간 (compute workspace)을 필요로 한다는 것이다. 내 기기에서는 -ub 4096에 --n-cpu-moe 26이 필요했고, -ub 8192에는 --n-cpu-moe 28이 필요했다. 따라서 이것은 처리량의 트레이드오프 (trade-off)이다: 더 큰 배치를 위한 충분한 공간을 확보하기 위해 몇 개의 MoE 레이어를 CPU로 더 이동시키면, 프롬프트 비중이 높은 워크로드 (workloads)는 극적으로 빨라지는 반면 생성 속도는 약간 느려진다.

참고: 처음 4개의 prefill 지점은 pp4096이며, 8192 ubatch 지점은 pp8192 실행 결과이므로, 이를 완벽하게 통제된 벤치마크라기보다는 비공식적인 튜닝 결과로 취급하라.

---

내가 DGX Spark를 구매한 이유 중 하나는 더 나은 프롬프트 처리 속도를 갖기 위해서였다. 만약 내가 이 트릭을 알았더라면, 돌이켜봤을 때 그렇게 하지 않았을지도 모른다. 비록 그것이 매우 훌륭한 기계이고, gpt-oss-120b에 대해 여전히 약간 더 나은 프롬프트 처리 성능과 약 두 배의 토큰 생성 속도를 제공하긴 하지만 말이다. 더 높은 ubatch는 그 격차를 획기적으로 줄여준다.