DOCSEARCHbyAlgolia
Submit Docs

DistillKit

repository·main·Indexed 21 days ago

https://github.com/arcee-ai/distillkit

A production-ready toolkit for knowledge distillation of large language models (LLMs). DistillKit supports both online distillation, where the teacher model runs in real-time, and offline distillation using pre-captured, compressed teacher outputs. It features a multi-stage logit compression system (polynomial approximation, quantization, and bit-packing) to reduce storage overhead and VRAM requirements. The library supports various loss functions including KL divergence, JSD, TVD, and hidden state alignment (MSE/Cosine).

Tokens
9.9K
Snippets
22
Records
38
Agent score
76%

What's inside distillkit

  1. How Online vs Offline Distillation works

    main

    DistillKit supports two primary workflows for knowledge distillation:

    1. Online Distillation: The teacher model runs in real-time alongside the student during training.

      • Best for: Scenarios where you have sufficient VRAM to hold both models and want dense distributions.
      • Pros: No storage overhead.

    2. Offline Distillation: Teacher outputs are pre-captured, compressed, and stored in a dataset.

      • Best for: VRAM-limited environments, large-scale training, or when reusing the same teacher for many students.
      • Pros: Highly efficient via advanced logit compression.

    Decision Rule: If you can fit both the teacher and student (with dense distributions) into VRAM, use online. Otherwise, use offline with the compression system.

  2. Configure Logit Compression for Offline Distillation

    main

    DistillKit uses a multi-stage compression system (polynomial approximation + quantization + bit-packing) to make offline distillation feasible at scale.

    Important: The logprob_compressor configuration used during the capture phase must match the configuration used during the distillation phase.

    Recommended Configuration

    Provides ~300 bytes/token (0.15% of uncompressed size) with minimal quality loss:

    logprob_compressor:
  d: <your_vocab_size_here>
  k: 128
  exact_k: 16
  exact_dtype: bfloat16
  polynomial_terms: [0, 1, 2, 3, 4, "sqrt"]
  term_dtype: float32
  residual_bins: []
  delta_encoding: false
  error_diffusion: false

    Budget Configuration

    Provides ~114 bytes/token for tighter storage constraints:

    logprob_compressor:
  d: <your_vocab_size_here>
  k: 50
  exact_k: 1
  exact_dtype: bfloat16
  polynomial_terms: [0, 1, "sqrt"]
  term_dtype: float32
  residual_bins: []
  delta_encoding: false
  error_diffusion: false

  3. Perform Cross-Architecture Distillation

    main

    DistillKit can be used in conjunction with mergekit-tokensurgeon to perform distillation across different tokenizers or architectures. The recommended workflow is:

    1. Use tokensurgeon to adapt the student model's embeddings to the teacher's tokenizer.
    2. Use DistillKit to distill the teacher's knowledge into the student.
    3. (Optional) Convert back to the student's original tokenizer or perform further merges.

  4. Memory Management Tips for Distillation

    main

    Use these strategies to optimize VRAM and training efficiency:

    • Long Sequences: Use sparse_chunk_length (e.g., 1024) to process sequences in chunks. Enable gradient_checkpointing.
    • VRAM Savings:
      • Use optim: paged_adamw_8bit or optim: adamw_bnb_8bit.
      • Use bfloat16 instead of float32.
      • Enable Flash Attention 2: use_flash_attention: true.
    • Throughput: Reduce batch size and increase gradient_accumulation_steps to maintain effective batch size.

  5. Capture Teacher Outputs for Offline Distillation

    main

    To create your own offline distillation dataset, use the distillkit.sample_logits_vllm module. This requires vLLM to be installed.

    Example command:

    python -m distillkit.sample_logits_vllm \
  --model meta-llama/Llama-3.1-70B \
  --dataset allenai/tulu-3-sft-mixture \
  --output ./llama3_70b_tulu_logits/ \
  --compression-config ./compression_config.yaml

  6. Install DistillKit

    main

    To install DistillKit, clone the repository and install it in editable mode using pip.

    To enable logit capture capabilities (required for creating your own offline datasets), install the [capture] extra.

    git clone https://github.com/arcee-ai/distillkit.git
cd distillkit
pip install -e .

# Optional: For logit capture
pip install -e ".[capture]"

  7. Training Tips for Distillation

    main

    When configuring your distillation training, consider these recommended starting points:

    • Cross-entropy weight: Start with ~0.5, then tune based on dataset quality.
    • Distillation temperature: A value of temperature: 2.0 is a recommended starting point.
    • Missing probability handling:
      • Use zero to focus exclusively on the teacher's most confident predictions.
      • Use uniform to match the teacher's uncertainty levels.

  8. Quick Start: Offline Distillation

    main

    Offline distillation uses pre-captured teacher outputs from a dataset. This is ideal when VRAM is limited or you want to reuse teacher signals for multiple students.

    To run offline distillation, create a config.yaml specifying the model, dataset (with prepacked: true), teacher configuration (using kind: dataset), and loss functions. Then, execute the training using the distillkit command.

    # config.yaml
project_name: my-distillation
model: Qwen/Qwen3-8B
output_path: ./output
sequence_length: 8192

dataset:
  train_dataset:
    repo_id: arcee-ai/Qwen3-235B-Logits-Packed-8192
    split: train
  prepacked: true

teacher:
  kind: dataset
  logprob_compressor:
    d: 151936
    delta_encoding: true
    error_diffusion: false
    exact_dtype: float32
    exact_k: 32
    k: 128
    polynomial_terms: [0, 1, 2]
    residual_bins: []
    term_dtype: float32

loss_functions:
  - function: cross_entropy
    weight: 0.5
  - function: kl
    weight: 0.5
    temperature: 1.0
    missing_probability_handling: zero
    sparse_chunk_length: 1024

training_args:
  num_train_epochs: 1
  per_device_train_batch_size: 1
  gradient_accumulation_steps: 8
  learning_rate: 2.0e-6
  bf16: true
  optim: adamw_torch
  gradient_checkpointing: true
    distillkit config.yaml

  9. How DistillationTrainer computes loss

    main

    The DistillationTrainer overrides compute_loss to implement the distillation logic:

    1. Label Preparation: It ensures labels exist in the inputs. If config.dataset.eos_label_token_ids is configured, it replaces specific token IDs in the labels with the model's eos_token_id.
    2. Student Forward Pass: It performs a forward pass on the student model. If any configured loss functions require hidden states, output_hidden_states=True is passed to the model.
    3. Logit Truncation: If the student's logits dimension does not match true_vocab_size, the logits are truncated to match true_vocab_size.
    4. Total Loss Calculation: It calls total_distillation_loss, which retrieves the teacher signal via signal_source.get_signal(...) and aggregates the weighted losses from all configured loss functions.

  10. Configure Teacher Models and Logit Datasets

    main

    The teacher configuration in a DistillationRunConfig determines how the student learns. It uses a discriminated union between two types:

    1. TeacherModelConfig: Used for Online Distillation. A live model is loaded from a path (typically a Hugging Face model) and used to generate logits during training.

      • kind: Must be "hf".
      • path: The model path.
      • kwargs: Additional arguments for loading the model.
      • top_k: (Optional) Top-k filtering.

    2. TeacherDatasetConfig: Used for Offline Distillation. Instead of a live model, you provide a dataset containing pre-captured logits.

      • kind: Must be "dataset".
      • legacy_logit_compression: Configuration for legacy compression formats.
      • logprob_compressor: Configuration for DistributionQuantizationConfig to handle compressed logits.
    from distillkit.configuration import TeacherModelConfig, TeacherDatasetConfig

# Online: Live teacher model
tr_online = TeacherModelConfig(kind="hf", path="meta-llama/Llama-2-7b-hf")

# Offline: Pre-captured logits dataset
tr_offline = TeacherDatasetConfig(kind="dataset")

  11. Understand the SignalSource abstraction

    main

    The SignalSource is an abstract base class (ABC) that defines how distillation data is retrieved from a teacher. There are two primary implementations:

    1. OfflineSignalSource: Used for high-throughput distillation where logits are pre-calculated and stored on disk. It uses a LogprobCompressor to handle sparse data. It cannot provide hidden states.
    2. OnlineSignalSource: Used when the teacher model is loaded in memory. It performs live forward passes. It can provide hidden states if requested.

    Both implement get_signal(batch, return_hidden_states) to return a TeacherSignal (either DenseSignal or SparseSignal).

  12. Configure Hidden State Mapping (HSD)

    main

    If layer_mapping is provided in the configuration, DistillKit performs Hidden State Distillation (HSD).

    Requirements & Behavior:

    • Compatibility: HSD is only supported with Online signal sources (teacher models), not offline logprob sources.
    • Mapping Types:
      • If layer_mapping: "all", it maps every student layer to its corresponding teacher layer.
      • Otherwise, it accepts a list of tuples [(student_layer_idx, teacher_layer_idx), ...].
    • Projection: If force_hidden_state_projection is enabled, it uses a HiddenStateMapping to handle differences between student and teacher hidden dimensions.