DistillKitPlus

DistillKitPlus is an advanced open-source toolkit designed for efficient knowledge distillation (KLD) in large language models (LLMs). Building upon the inspiration from DistillKit, this repository aims to facilitate offline distillation and parameter-efficient fine-tuning (PEFT) for environments with limited computational resources.

Key Features

• Logit Distillation: Supports both same and cross-tokenizer teacher and student models, enhancing versatility in distillation tasks.
• Pre-Computed Logits: Generates logits in advance to enable memory-efficient training processes, optimizing resource usage.
• LoRA Fine-Tuning Integration: Offers comprehensive support for low-rank adaptation fine-tuning, streamlining the fine-tuning process.
• Quantization Support: Facilitates 4-bit model quantization to enhance inference speed and reduce memory requirements.
• Accelerate & DeepSpeed Integration: Provides support for distributed training while optimizing memory usage for improved performance.

Supported Loss Functions

Quick Start

To begin using DistillKitPlus, configure your distillation settings through the config/default_config.json. Generate teacher logits with the following command:

python scripts/local/generate_logits.py --config config/default_config.json

Run the distillation process using: Without Accelerate (default):

python scripts/local/distill_logits.py --config config/default_config.json

With Accelerate & DeepSpeed:

# Ensure "use_accelerate": true in your configuration file
accelerate launch --config_file config/accelerate_configs/default_config.yaml scripts/local/distill_logits.py --config config/default_config.json

Optional Modal Integration

DistillKitPlus can also run scripts via Modal. Use the commands below for knowledge distillation: Generate teacher logits:

modal run scripts/modal/generate_logits.py --config config/default_config.json

Run distillation:

modal run scripts/modal/distill_logits.py --config config/default_config.json

When utilizing Modal, the accelerate configuration is internally managed based on settings in the configuration file.

Configuration Details

The toolkit relies on a JSON configuration file that includes key sections for project management:

• project_name: Name of the distillation project.
• dataset: Settings for dataset sourcing and processing.
• models: Specifications for teacher and student models.
• tokenizer: Settings for tokenization, including max length and padding.
• training: Hyperparameters for training processes.
• distillation: Key parameters for distillation including temperature and alpha settings.
• lora: LoRA specific configurations for fine-tuning.
• quantization: Settings for model quantization.
• execution: Parameters for using accelerate and distributed training.

For a more comprehensive setup, refer to the config/default_config.json file.

Contribution Guidelines

Open contributions are welcomed from the community. Improvements, new features, or bug fixes can be suggested by opening an issue or submitting a pull request.