Slimming Models, Saving Watts: An Energy-Aware Knowledge Distillation Framework for Large Language Models

This research framework targets large language models such as Llama 3.1, systematically evaluating the accuracy, efficiency, and energy consumption performance of three knowledge distillation methods: responsive, feature-based, and relational. It is specifically designed for HPC clusters and Slurm environments. The framework fills the gap in traditional knowledge distillation research regarding the systematic evaluation of energy efficiency, deeply integrating energy consumption measurement with KD effect assessment, and providing a standardized tool for green AI research.

As the number of parameters in large language models grows from billions to hundreds of billions, the energy consumption problem in training and deployment has become increasingly prominent. Knowledge Distillation (KD), as a core model compression technology, can reduce model size while maintaining performance. However, traditional KD research mainly focuses on accuracy retention, and there is a relative lack of systematic evaluation of energy efficiency. The Slimming Models, Saving Watts project has built a complete research framework for HPC environments, filling this gap.

The framework adopts a modular design and includes three core components:

1. Three knowledge distillation paradigms: Responsive (matching output logits distribution), feature-based (aligning intermediate layer features), and relational (maintaining inter-sample relationship structure);
2. Energy telemetry system: Integrates the monitor.py module to collect real-time data such as GPU power consumption, utilization, and memory, and calculates key indicators like total energy consumption (E_run) and energy per token (EPT);
3. Slurm-compatible HPC deployment: Supports multi-GPU parallel training, Slurm job submission, distributed data sharding, and is compatible with GPU environments such as NVIDIA H100/A100.

The experiments mainly target the Llama 3.1 series: the teacher model is Llama-3.1-70B-Instruct, and the student model is Llama-3.1-8B-Instruct. The evaluation system includes multi-dimensional indicators:

• OM_perf: Performance retention rate of the student model relative to the teacher model;
• EPT: Energy per token during inference;
• Eff_overall: Comprehensive efficiency indicator integrating accuracy and energy consumption. The evaluation phase integrates mainstream benchmarks such as MMLU, ARC, BBL, and HellaSwag, and supports the lm-harness and lighteval frameworks.

The framework provides end-to-end workflow support:

1. Environment preparation: pip install -r requirements.txt;
2. Data construction: Load datasets from Hugging Face and generate shards via build_shards_from_hf.py (improves I/O performance and ensures reproducibility);
3. Baseline training, knowledge distillation, energy consumption monitoring, model evaluation, and result analysis (visualized via Jupyter Notebook).

The project includes a rich set of Jupyter Notebook tools:

• Energy consumption analysis series: feature_energy_plot.ipynb (energy consumption curve of feature-based distillation), response_energy_plot.ipynb (responsive), relation_energy_plot.ipynb (relational);
• Performance indicator series: OMperf.ipynb (performance retention analysis), ENERGYrun.ipynb (energy consumption operation analysis), EFFoveral.ipynb (comprehensive efficiency evaluation). These tools provide directly usable chart materials for research.

The release of the framework has multiple values:

• Research level: For the first time, energy consumption measurement is systematically integrated into the KD evaluation system, providing a standardized tool for green AI research;
• Engineering level: Complete Slurm integration and HPC optimization support large-scale experiments in real production environments;
• Industry level: Indicators such as EPT provide a new dimension for model selection, and energy consumption becomes a key consideration besides accuracy and speed. This framework provides a fully functional platform for researchers and engineers in the fields of large model efficiency optimization, green computing, and KD.