Nvidia はプルーニングと蒸留を試しています。Llama 3.1 8B のパラメータを半分にして、同じサイズでより良いパフォーマンスを実現しています。

小型モデルの台頭。

Figure 2 below shows the distillation process, where the N-layer student model (the pruned model) is distilled from the M-layer teacher model (the original unpruned model). The student model is learned by minimizing a combination of embedding output loss, logit loss, and Transformer encoder-specific losses mapped to student blocks S and teacher blocks T. Figure 2: Distillation training loss.

Best practices for pruning and distillationNVIDIA pruning and knowledge distillation based on compact language model Based on extensive ablation research, I summarized my learning results into the following structured compression best practices.

• Prune the available source model closest to the target size.

• The second is pruning.

• Prioritize width pruning over depth pruning, which works well for models under 15B parameter size.

• The third is to retrain.

• Only use distillation loss for retraining instead of regular training.

• Use logit-only distillation when depth does not decrease significantly.
Llama-3.1-Minitron: Putting best practices into action

• Meta recently launched features The powerful Llama 3.1 family of open source models is comparable to closed source models in many benchmarks. Llama 3.1's parameters range from a massive 405B to 70B and 8B.

# # teacher fine-tuning

Depth-only pruning

• Width-only pruning

• Accuracy Benchmark

• Performance Benchmark

• # #teacher fine-tuning

• In order to correct the distribution bias of the original data set on which the model training is based, NVIDIA first trained the unpruned 8B model on their data set (94B tokens) Fine-tuned. Experiments show that if distribution bias is not corrected, the teacher model provides suboptimal guidance for the dataset when distilling.

Depth-only pruning

In order to reduce from 8B to 4B, NVIDIA pruned 16 layers (50%). They first evaluate the importance of each layer or group of consecutive sub-layers by removing them from the model and observe an increase in LM loss or a decrease in accuracy in downstream tasks.

Figure 5: The importance of depth-only pruning the middle layer.

However, NVIDIA observes that this LM loss is not necessarily directly related to downstream performance.

Accuracy.

Width-only pruning

NVIDIA prunes embedding (hidden) and MLP along the width axis Intermediate dimensions to compress Llama 3.1 8B. Specifically, they use the previously described activation-based strategy to compute importance scores for each attention head, embedding channel, and MLP hidden dimension.After importance estimation, NVIDIA chose

to prune the MLP middle dimension from 14336 to 9216.

• Retrain attention to the number of heads and the number of layers.

It is worth mentioning that after single-sample pruning, the LM loss of width pruning is higher than that of depth pruning. However, after a brief retraining period, the trend reversed.

Accuracy Benchmark

NVIDIA distilled the model using the following parameters

• Peak learning rate = 1e-4

• Minimum learning rate = 1e-5

• 40 steps linear warm-up

• Cosine Decay Plan

• Global batch size = 1152

Table 1 below shows the Llama-3.1-Minitron 4B model variants (width pruning and depth pruning) similar to the original Llama 3.1 8B model, others Performance comparison of large and small models on benchmarks across multiple domains. Overall, NVIDIA once again confirmed the effectiveness of a wide pruning strategy compared to deep pruning that follows best practices.

^{Table 1: Accuracy comparison of Minitron 4B base model compared to base models of similar scale.}

To verify whether the distilled model can become a powerful instruction model, NVIDIA used NeMo-Aligner to fine-tune the Llama-3.1-Minitron 4B model.

They used Nemotron-4 340B training data and evaluated on IFEval, MT-Bench, ChatRAG-Bench and Berkeley Function Calling Leaderboard (BFCL) to test instruction following, role-playing, RAG and function calling capabilities. Finally, it was confirmed that the Llama-3.1-Minitron 4B model can be a reliable instruction model, outperforming other baseline SLMs.

^{Table 2: Accuracy comparison of aligned Minitron 4B base model with similarly sized aligned models.}

Performance Benchmarks

NVIDIA optimized the Llama 3.1 8B and Llama-3.1-Minitron 4B models using NVIDIA TensorRT-LLM, an open source toolkit for optimizing LLM inference.

The next two figures show the throughput requests per second of different models at FP8 and FP16 precision under different use cases, expressed as the input sequence length/output sequence length (ISL/OSL) combination of the batch size of 32 for the 8B model and The batch size of the 4B model is an input sequence length/output sequence length (ISL/OSL) combination of 64, thanks to the smaller weights allowing a larger batch size on an NVIDIA H100 80GB GPU.

The Llama-3.1-Minitron-4B-Depth-Base variant is the fastest, with an average throughput of about 2.7 times that of Llama 3.1 8B, while the Llama-3.1-Minitron-4B-Width-Base variant has an average throughput of Approximately 1.8 times that of Llama 3.1 8B. Deployment in FP8 also improves the performance of all three models by approximately 1.3x compared to BF16.

^{80GB GPU.}

Conclusion

Pruning and classical knowledge refining is a very cost-effective method to progressively obtain LLMs of smaller sizes, achieving higher accuracy than training from scratch in all domains . This is a more efficient and data-efficient approach than fine-tuning on synthetic data or pre-training from scratch.

Llama-3.1-Minitron 4B is NVIDIA’s first attempt at using the state-of-the-art open source Llama 3.1 series. To use the SDG fine-tuning of Llama-3.1 with NVIDIA NeMo, see the /sdg-law-title-generation section on GitHub.

For more information, please see the following resources:

• https://arxiv.org/abs/2407.14679

• https://github.com/NVlabs/Minitron

• https:// huggingface.co/nvidia/Llama-3.1-Minitron-4B-Width-Base

• https://huggingface.co/nvidia/Llama-3.1-Minitron-4B-Depth-Base

^{Reference links:}

^{https://developer.nvidia.com/blog/how-to-prune-and-distill-llama-3-1-8b-to-an-nvidia-llama-3-1-minitron-4b -model/}