Nvidia의 새로운 LLM 파라미터 축소 방법론 : PUZZLE: DISTILLATION-BASED NAS FORINFERENCE-OPTIMIZED LLMS

이전에 NVIDIA Llama Nemotron Reasoning Models과 관련된 포스팅을 한적이 있습니다. NVIDIA Llama Nemotron Reasoning Models의 몇몇 학습 방법은 DeepSeek-R1 등과 같은 여타 다른 모델들의 학습 방법과 큰 차이가 존재하진 않습니다. 그러나 모델의 아키텍처는 차이가 있는데 그것이 바로 NAS(Neural Architecture Search) 입니다. 

 

NAS는 Arxiv에 게재되어 있는 PUZZLE: DISTILLATION-BASED NAS FOR INFERENCE-OPTIMIZED LLMS(Nvidia, 2025) 에서 제안된 커스터마이징 구조입니다. 그래서 이번 포스팅에서는 해당 논문을 리뷰해보고자 합니다. 

 

---

Large language models (LLMs) offer remarkable capabilities, yet their high inference costs restrict wider adoption. While increasing parameter counts improves accuracy, it also broadens the gap between state-of-the-art capabilities and practical deployability. We present Puzzle, a hardware-aware framework that accelerates the inference of LLMs while preserving their capabilities. Using neural architecture search (NAS) at a largescale, Puzzle optimizes models with tens of billions of parameters. Our approach utilizes blockwise local knowledge distillation (BLD) for parallel architecture exploration and employs mixed-integer programming for precise constraint optimization. We showcase our framework’s impact via Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), a publicly available model derived from Llama-3.1-70B-Instruct. Nemotron-51B achieves a 2.17× inference throughput speedup, fitting on a single NVIDIA H100 GPU while retaining 98.4% of the original model’s benchmark accuracies. Notably, it is the most accurate model supporting single H100 GPU inference with large batch sizes, despite training on only 45B tokens, far fewer than the 15T used to train Llama-70B. Lastly, we derive Llama-3.3-Nemotron-49B-Super-Base to demonstrate Puzzle can retain long-context and that lightweight alignment on these derived models allows them to surpass the parent model in specific capabilities. Our work establishes that powerful LLM models can be optimized for efficient deployment with only negligible loss in quality, underscoring that inference performance, not parameter count alone, should guide model selection.

 

대형 언어 모델(LLM)은 뛰어난 성능을 제공하지만 높은 추론 비용 때문에 보편적인 활용에는 어려움이 있다. 파라미터 수를 늘리면 정확도가 향상되기는 하지만, 이는 최신 기술력과 실질적인 배포 가능성 사이의 간극을 더욱 벌린다. 본 연구에서는 Puzzle이라는 하드웨어를 고려한 프레임워크를 제안하여 LLM의 추론 속도를 가속화하면서도 모델의 성능은 유지한다. Puzzle은 수십억 개의 파라미터를 가진 모델들을 대상으로 대규모 신경망 구조 탐색(NAS)을 통해 최적의 아키텍처를 찾아내며, 병렬 아키텍처 탐색을 위해 블록별 국소 지식 증류(BLD)를 활용하고, 혼합정수 프로그래밍을 통해 정밀한 제약 최적화를 수행한다.

이 프레임워크의 효과를 입증하기 위해 Llama-3.1-70B-Instruct 모델에서 파생된 공개 모델인 Llama-3.1-Nemotron-51B-Instruct(Nemotron-51B)를 소개한다. Nemotron-51B는 NVIDIA H100 GPU 한 대에서 실행할 수 있으며, 추론 처리 속도가 2.17배 향상되었고 원본 모델의 벤치마크 정확도의 98.4%를 유지한다. 특히, Nemotron-51B는 대규모 배치 처리를 지원하는 단일 H100 GPU 기반 모델 중 가장 정확한 모델로 평가되며, 15조 토큰으로 훈련된 Llama-70B와 달리 단 450억 토큰으로 학습되었음에도 불구하고 효율성을 입증하였다. 마지막으로, Puzzle 프레임워크를 활용해 장기 문맥 유지 능력을 갖춘 Llama-3.3-Nemotron-49B-Super-Base 모델을 도출하였고, 경량 정렬 과정을 통해 이 모델이 부모 모델보다 특정 능력에서 더 우수한 성능을 발휘할 수 있음을 보여주었다. 본 연구는 강력한 LLM 모델도 미미한 품질 손실로 효율적인 배포가 가능하며, 모델 선택 시 파라미터 수 뿐만 아니라 추론 성능을 고려해야 함을 강조한다.

---

 

위의 내용은 해당 논문의 초록입니다. 이 내용을 간단히 요약하면 대형 언어 모델들의 경우 굉장히 좋은 성능을 보여 주지만 추론 비용이 그에 비례해서 높기 때문에 실제로 적용하는데은 어려움이 있기에 이를 Puzzle이라는 하드웨어 최적화 프레임 워크를 제안해서 해결해보고자 한다. 여기서 Puzzle은 수십억 개의 파라미터를 가진 모델들을 대상으로 대규모 신경망 구조 탐색과 병렬 아키텍처 탐색을 가능하게 하는 블록별 국소 지식 증류 기술을 적용하며, 혼합정수 프로그래밍을 사용해 구체적인 제약 조건 하에서 최적화를 수행한다. 특히 여기서 실제로적용되었을 때 Llama-3.1-70B-Instruct 모델을 기반으로 만들어진 Nemotron-51B 모델은 단일 NVIDIA H100 GPU에서 2.17배 빠른 추론 처리를 구현하면서도 기존 모델의 98.4%에 달하는 정확도를 유지했다고 합니다.

 

또한, Nemotron-51B는 대규모 배치 처리 상황에서 가장 뛰어난 성능을 보이는 단일 GPU 모델로 평가받으며, 비교적 적은 양의 토큰으로 학습되었음에도 불구하고 높은 효율성을 달성하였다. 추가적으로, Puzzle 프레임워크를 활용하여 장기 문맥 유지가 가능한 Llama-3.3-Nemotron-49B-Super-Base 모델을 도출하였으며, 경량 정렬 기법을 통해 이 모델이 특정 분야에서 부모 모델을 능가하는 성능을 보일 수 있음을 입증하였다. 이 연구는 모델 선택 시 파라미터 수에만 의존하지 말고, 추론 성능을 중요한 기준으로 삼아야 함을 강조하며, 효율적인 최적화가 실제 배포에서 큰 효과를 발휘할 수 있음을 보여준다고 합니다. 

 

Puzzle 프레임워크

 

여기서, Puzzle 프레임워크는 LLM의 훈련용 모델 구조를 하드웨어에 최적화된 효율적인 추론 모델로 변환하기 위해 개발되었습니다. 이 프레임워크는 먼저 분해된 신경망 구조 탐색(NAS) 기법을 활용하여 다양한 아키텍처 후보를 자동으로 탐색합니다. *NAS는 Neural Architecture Search의 약자로, 최적의 네트워크 구조를 찾기 위해 수많은 설계 옵션들을 평가하고 비교하는 기술을 의미하며, 이를 통해 인간의 개입 없이도 가장 적합한 모델 구성을 도출할 수 있습니다.

 

기존 연구들에서는 지치기, 전체 레이어 제거, 로컬 어텐션, KV 헤드 수 축소 등 다양한 기법으로 추론시 많은 계산 및 불필요하게 사용되는 비효율을 개선하려고 했지만, 높은 계산 비용과 상업적 운영 비용으로 인해 최적화가 시급하다는 것이 해당 논문의 생각이다 그렇기에 여기에 NAS(Neural Architecture Search)를 도입하였다고 하는데 부모 모델의 각 트랜스포머 레이어를 어텐션 서브블록과 FFN 서브블록의 다양한 대안으로 대체하여, 모든 가능한 조합(예를 들어, Llama 3.1-70B 모델의 경우 한 레이어당 54가지 구성, 전체 10^138개 이상의 아키텍처 후보)을 탐색하는 방대한 검색 공간을 정의한다.

 

NAS는 이처럼 정의된 검색 공간 내에서 성능과 자원 제약(메모리, 지연시간, 처리량 등) 간의 균형을 고려하며 최적의 구조를 찾기 위한 체계적인 접근법으로, Puzzle 프레임워크에서는 블록별 국소 지식 증류(BLD)와 혼합정수 프로그래밍(MIP) 기법을 통해 후보 구조들의 성능을 효율적으로 평가하고, 최종적으로 Global Knowledge Distillation(GKD)을 통해 재조립한 모델로 최적화된 추론 효율과 성능을 동시에 달성한다. 이러한 접근법은 단순히 파라미터 수에 의존하는 기존 방식의 한계를 극복하고, 실제 하드웨어에 최적화된 모델 선택과 배포를 가능하게 하여, 고성능 AI를 보다 저비용으로 상용화할 수 있음을 보여줍니다. 

 

Blockwise Local Distillation

 

학습 방식

 

이 연구에서는 LLM의 효율적인 추론을 위해 ‘퍼즐 조각’이라 부를 수 있는 개별 블록들을 만들고, 이들을 조합하여 최종 모델을 구성하는 방법을 소개합니다. 이를 위해 먼저 각 트랜스포머 블록의 자식 모델 즉, 어텐션과 FFN(피드포워드 네트워크) 서브블록의 다양한 변형이 부모 블록의 기능을 모방하도록 독립적으로 학습시키는 “블록 단위 로컬 지식 증류(Blockwise Local Distillation, BLD)” 기법을 사용합니다. 각 자식 블록은 해당 부모 블록에서 출력되는 활성화 정보만을 받아, 별도의 블록으로서 학습되므로 서로 간섭 없이 독립적으로 작업할 수 있으며, 이를 통해 파이프라인 병렬 처리가 가능해집니다.

BLD 기법은 복잡한 전체 자식 모델을 한 번에 학습하는 대신, 각 블록이 단순한 함수(즉, 부모 서브블록의 기능)를 모방하도록 하는 방식이기 때문에 학습이 훨씬 단순하고 안정적입니다. 이로 인해 빠른 수렴과 높은 학습률을 사용할 수 있으며, 비교적 적은 양의 데이터(약 10억 토큰)로도 효과를 볼 수 있습니다. 단, 이 방법만으로는 각 블록 간의 상호 호환성 문제가 발생할 수 있기 때문에, 최종적으로 Global Knowledge Distillation(GKD)을 추가하여 전체 모델의 성능을 보완합니다.

 

또한, 연구에서는 전체 검색 공간을 효과적으로 탐색하기 위해 블록 라이브러리를 구축하는 두 가지 방식, 즉 coupled BLD와 decoupled BLD를 제안합니다. coupled BLD는 각 트랜스포머 레이어에서 가능한 모든 어텐션과 FFN 서브블록의 조합을 개별적으로 학습하는 방식인 반면, decoupled BLD는 부모와 동일한 부분은 고정하고 다른 부분만 개별 학습한 후 이들을 조합하여 블록을 완성합니다. 이 덕분에 학습해야 하는 블록의 수가 크게 줄어들어, 예를 들어 수백 가지 옵션이 있을 때 훨씬 빠르고 비용 효율적으로 라이브러리를 구축할 수 있습니다.

 

마지막으로, 대체 서브블록의 초기화를 위해 '채널 기여도' 평가 기법을 도입합니다. 이 방법은 FFN 서브블록에서 각 채널이 출력에 미치는 영향을 측정하여, 기여도가 낮은 채널을 우선적으로 제거하거나 축소하고, 어텐션 서브블록의 경우에는 키와 값 투영 행렬을 평균 풀링하는 방식 등으로 초기화하여 학습을 가속화합니다. 전체적으로 이 연구는 개별 블록 단위로 효율적으로 학습하고 평가하는 방법을 통해, LLM의 추론 효율성을 크게 개선할 수 있음을 보여줍니다.

 

Decomposed NAS Search Algorithm for LLMs

 

이 연구에서는 대형 언어 모델(LLM)의 효율적인 추론을 위해 “decomposed NAS” (분해된 신경망 구조 탐색) 방식을 사용합니다. 이 방법은 크게 네 단계로 진행됩니다.

 

첫 번째 단계는 블록 라이브러리 구축입니다. 기존의 부모 모델(원래 모델)의 각 트랜스포머 레이어를 여러 대체 블록으로 나누어, 블록별 로컬 지식 증류(BLD) 기법을 통해 각각 독립적으로 학습시키고, 다양한 조각(어텐션와 FFN 서브블록의 여러 변형)들로 구성된 라이브러리를 만듭니다.

 

두 번째 단계에서는 각 블록의 자원 요구사항을 추정합니다. 이는 단순히 이론적인 FLOPs나 파라미터 수를 넘어서, 실제 하드웨어(예: GPU)에서의 실행 시간, 메모리 사용량 등을 측정합니다. 특히, KV 캐시 메모리 같이 입력 크기나 배치 크기에 따라 변하는 요소들도 고려하여, 실제 환경에서 효율적으로 동작하는지를 확인합니다.

 

세 번째 단계는 블록의 성능 점수(quality score)를 부여하는 단계입니다. 각 레이어의 블록 후보 중 하나를 부모 블록 대신에 교체한 후, 전체 모델 성능에 미치는 영향을 측정합니다. 이를 replace-1-block 점수라고 부르는데, 이 점수를 통해 각 블록이 부모 모델의 성능을 얼마나 잘 유지하는지를 평가하게 됩니다. 평가 기준은 LM 손실, KL 발산 등 여러 지표가 사용될 수 있습니다.

 

네 번째 단계에서는 최종 모델 아키텍처를 검색(설계)하는 단계로, Mixed-Integer Programming (MIP) 알고리즘을 사용합니다. 트랜스포머의 각 레이어는 하나의 그룹처럼 취급되며, 각 그룹(레이어)에서 선택할 수 있는 다양한 블록 후보들이 “아이템”처럼 존재합니다. 이때 각 후보는 성능 점수와 메모리, 실행 시간 등의 “비용”이 함께 고려됩니다. 최적의 모델 아키텍처는 모든 레이어에서 하나의 블록 후보를 선택하여, 전체 성능 점수를 최대화하면서도 하드웨어의 메모리 한계, 최소 처리량, 최대 지연 시간 등과 같은 제약 조건을 만족하도록 구성됩니다.

 

여기서 잠깐? Mixed-Integer Programming (MIP) 알고리즘이란? 

https://convex-optimization-for-all.github.io/contents/chapter24/2021/03/28/24_01_Definition/

24-01 Definition · 모두를 위한 컨벡스 최적화

 

또한, 배치 크기와 같은 변수도 고려되어 여러 배치 크기에 대해 최적의 조합을 찾고, 하드웨어의 특성(예를 들어 H100 GPU에서의 FP8 양자화 지원 여부 등)에 맞춘 맞춤형 아키텍처를 생성할 수 있습니다. 이로써 “한 번 학습하여 필요에 따라 다르게 적용하는(train once, adapt on demand)” 방식으로, 다양한 배포 시나리오에 맞춰 최적화된 모델을 자동으로 생성할 수 있게 됩니다.

 

요약하면, 이 방법은 먼저 다양한 대체 블록들을 학습하여 블록 라이브러리를 구축하고, 실제 하드웨어에서의 실행 성능을 측정한 후, 각 블록의 성능 점수를 부여하여 전체 모델 성능에 미치는 영향을 평가합니다. 마지막으로, Mixed-Integer Programming을 통해 하드웨어 제약 조건을 고려한 최적의 블록 조합을 찾아내어 효율적인 LLM 아키텍처를 설계하는 방식입니다.

 

Post-Puzzle Inter-Block Uptraining

Blockwise Local Distillation 에서 언급했듯이, BLD 단계에서는 인접 블록 간의 호환성을 고려하지 않았습니다. 자식 모델에서는 각 블록이 자신의 이전 자식 블록이 아닌, 부모 모델의 대응하는 블록의 출력을 입력으로 받아 학습했기 때문에, 실제로는 각 블록이 훈련 당시와 다른 입력 분포를 받게 되어 성능이 떨어질 수 있습니다. 이를 완화하기 위해 최종 단계에서는 Global Knowledge Distillation(GKD)을 사용하여 학생 모델(자식 모델)을 전체적으로 재학습합니다. 여기서 부모 모델은 “교사” 역할을 하며, 자식 모델은 그 지식을 그대로 또는 개선하여 받아들이도록 학습됩니다.

 

전통적인 KD 기법은 교사와 학생 모델의 로짓 간의 cross-entropy 손실을 사용했으나, 최근에는 KL 발산(KLD) 손실과 언어 모델링(LM) 손실(다음 토큰 예측을 위한 cross-entropy)을 결합하는 방식이 더 빠르고 안정적인 최적화에 도움을 주고 있습니다.학습에 사용되는 손실 구성은 선택된 모델, 훈련 데이터 및 다운스트림 작업에 따라 달라지는데, 본 연구에서는 여러 손실 조합을 평가한 결과 LM 손실 없이 Cosine Similarity Loss(각 레이어의 히든 상태 간 유사도)와 KLD Loss를 더한 LGKD = Lcosine + LKLD 조합이 최종 모델에 가장 효과적임을 확인했습니다. LM 손실은 오히려 다운스트림 작업 성능을 해칠 수 있으며, 이는 부모 모델이 훈련된 데이터(닫힌 데이터)에 접근하지 못해 발생하는 분포 불일치 문제를 완화하기 위한 결과로 해석됩니다.

 

Supporting Fast Inference for Variable-Block Architectures in TensorRT-LLM

TensorRT-LLM은 NVIDIA GPU에서 LLM의 추론 성능을 가속화하기 위해 최적화된 런타임으로, 지연 시간과 처리량에 최적화된 성능을 제공합니다. 기존에는 모든 어텐션 레이어가 동일한 키 및 값 헤드를 가지고 있어야 한다는 가정이 있었으나, Puzzle에서 생성된 다양한 블록 아키텍처를 지원하기 위해, 이 부분을 수정하였습니다. 새로운 paged KV 캐시 전략은 다양한 GQA 비율(키-값 헤드 수 변동)을 갖는 모델에서도 효율적으로 작동하도록 개선되었습니다. 또한, TensorRT-LLM은 FP8 양자화, 맞춤형 스케줄링 정책 등을 지원하여 Puzzle의 NAS 프레임워크로 만든 아키텍처(예: Nemotron-51B)가 실제 추론 상황에서 효율적으로 동작하도록 합니다.

 

이를 알기 쉽게 서술하자면, 위의 내용은 Puzzle의 최종 단계로, 블록별 로컬 지식 증류(BLD) 후에 발생할 수 있는 블록 간의 입력 분포 불일치 문제를 보완하기 위한 과정입니다.

1. BLD 단계에서는 각 자식 블록이 독립적으로 부모 모델의 대응 블록 출력을 기반으로 학습되었기 때문에, 실제 자식 모델에서는 이전 블록으로부터 다른 입력이 들어와 성능 저하가 발생할 수 있습니다.
2. 이 문제를 해결하기 위해, Global Knowledge Distillation(GKD)을 통해 전체 자식 모델을 다시 한 번 end-to-end 방식으로 학습시킵니다. 이를 통해, 자식 모델이 부모 모델의 지식을 온전히 계승하면서 실제 상황에 맞게 조정됩니다.
3. 최종적으로는 Cosine Similarity와 KL Divergence를 결합한 손실(LGKD)을 사용하여 학습했으며, LM 손실은 오히려 성능에 부정적 영향을 주어 생략되었습니다.
4. 추가로, TensorRT-LLM 런타임은 다양한 블록 구성을 갖는 Puzzle 아키텍처가 실제 GPU 환경에서 빠르게 추론할 수 있도록, KV 캐시 관리와 FP8 양자화 같은 하드웨어 최적화 기능을 지원합니다.

이처럼, Post-Puzzle Inter-Block Uptraining 단계는 개별적으로 학습된 자식 블록들이 서로 잘 호환되어 전체 모델 성능을 최적화할 수 있도록 재학습하는 과정이며, 이를 통해 Puzzle로 생성된 모델이 실제 추론 환경에서도 높은 효율성과 성능을 보장하도록 합니다.

 

 

결과 및 결론 

Puzzle 프레임워크를 사용하면 Llama-3.1-70B-Instruct와 같은 강력한 부모 모델로부터 Nemotron-51B라는 자식 모델을 만들어 낼 수 있습니다. Nemotron-51B는 부모 모델의 거의 모든 정확도를 유지하면서도 추론 효율성을 크게 개선하여, 단일 GPU에서 최대 2배 이상의 처리량 향상을 이룹니다. 다양한 벤치마크에서 평균 98.4%의 정확도 보존을 확인했으며, 사람 평가에서도 유사한 성능을 보였습니다. 또한, 긴 문맥 처리에도 우수한 성능을 나타내어, 장기 문맥 모델(Lemotron-49B-Base) 또한 제작되었습니다. Puzzle은 전체 아키텍처 최적화를 위해 MIP 기반 글로벌 최적화와 다양한 소거 실험을 통해, 기존의 단순한 모델 구조가 아닌 하드웨어 제약 조건에 맞춘 맞춤형 효율 모델을 생산할 수 있음을 보여줍니다. 이로써, 대형 언어 모델을 더 적은 비용과 계산 자원으로 상용화할 수 있는 길이 열리게 됩니다.

 

소감

이 과정은 아무리 봐도 Nvidia 이니까 가능한 방법이 아닐까 싶습니다. 이를 정말 간략하게 요약해보자면 LLM을 Transformers 블록 단위로 분할한 다음, 이를 대체할 수 있는 블록을 굉장히 여러 조합으로 만든 다음 이를 BLD 방법으로 원래 Transformer 블록(부모) 의 파라미터를 이식하고 이들을 전체적으로 평가하여 혼합 정수 프로그래밍 방법론을 활용해서 최적의 블록 조합을 찾아내고 블록 단위의 학습이 진행되었기에 전체적으로 모델이 입출력이 틀어질 가능성을 대비하기 위해서 Global Knowledge Distillation(GKD) 를 통해 이걸 튜닝해준다. 그러면 모델의 사이즈를 기존 대비 크게 (405 -> 253, 70-> 50 등) 줄일 수 있다. 가 이 논문의 요지입니다. 

 

즉, Pretrain 보다 적게 자원이 필요하지만 Finetuning 보다는 당연히 많은 자원이 필요했을 것으로 보입니다. 물론 BLD는 부모 모델의 특정 블록에서 생성된 정보(예: 활성화 값 또는 출력)를 사용하여, 각 자식 블록이 해당 기능을 모방하도록 학습시키는 과정이고 이를 통해 전체 모델을 다시 학습하지 않고도, 개별 블록들이 부모 모델의 지식을 효과적으로 이어받아 기능을 수행할 수 있게 됨으로 큰 자원은 필요하지 않을 것 같긴 하지만 그래도 많은 자원이 필요함은 당연해보입니다. 아래는 해당 사항에 대한 GPT의 답변입니다. 

 

---

네, 대략 그렇게 이해해도 무방합니다. Puzzle은 전체 프리트레인(예를 들어, 수조 단위 토큰 학습)보다는 훨씬 적은 학습 자원(약 50B 토큰 정도)을 사용합니다. 그러나 일반적인 파인튜닝, SFT(지도 학습)나 DPO 같은 미세조정 방법보다 추가적인 블록 단위 학습 과정(예: 블록 라이브러리 구축, 로컬 지식 증류, 최종 GKD 업트레이닝 등)이 포함되어 있어서, 미세조정에 비해 상대적으로 더 많은 계산 자원이 필요할 수 있습니다.

즉, Puzzle은 프리트레인에 비해 극히 적은 자원으로 효율성을 크게 향상시키지만, 단순한 미세조정보다 추가적인 최적화 과정이 포함되어 있어 그 비용이 다소 상승할 수 있다는 점을 고려해야 합니다.

---