moshilite

Lightweight compression toolkit for the Moshi speech-language model.

moshilite provides a reproducible, stage-by-stage pipeline for compressing Moshi through structured pruning, knowledge distillation, and quantization — with optional LoRA fine-tuning.

---

Project Structure

moshilite/
├── README.md                       # This file
├── .gitignore
├── pyproject.toml                  # pip install -e .
│
├── notebooks/                      # One per major stage, numbered
│   ├── 00_setup_and_encoding.ipynb
│   ├── 01_layer_analysis.ipynb
│   ├── 02_structured_pruning.ipynb
│   ├── 03_depth_transformer_eval.ipynb
│   ├── 04a_teacher_precomputation.ipynb
│   ├── 04b_student_kd_training.ipynb
│   ├── 06_quantization.ipynb
│   └── 07_lora_finetuning.ipynb
│
├── scripts/                        # Batch / long-running jobs
│   ├── encode_tokens.py            # Phase 0 – Mimi pre-encoding
│   ├── precompute_teacher.py       # Stage 4a – teacher target generation
│   ├── export_model.py             # Stage 6 – model export + manifest update
│   └── run_sqa_eval.py             # SQA benchmark pipeline
│
├── src/moshilite/                  # Installable Python package
│   ├── data/                       # Dataset loading, Mimi encoding, dataloader
│   ├── analysis/                   # Layer importance, BI scores
│   ├── pruning/                    # Structured pruning
│   ├── distillation/               # KD losses, student training loop
│   ├── eval/                       # Evaluation pipeline, SQA benchmarks
│   └── utils/                      # Checkpointing, W&B logging, experiment helpers
│
├── configs/                        # Hyperparameter YAML configs per stage
│   ├── stage1_analysis.yaml
│   ├── stage2_pruning.yaml
│   ├── stage4b_training.yaml
│   └── stage6_quantization.yaml
│
└── tests/                          # Unit + integration tests
    ├── test_token_encoding.py
    ├── test_kd_losses.py
    ├── test_checkpoint_resume.py
    └── test_experiment_helpers.py

---

Setup

1. Clone and install

git clone https://github.com/AZMA1N/moshilite.git
cd moshilite
pip install -e ".[dev]"

2. (Optional) install extras

# Jupyter notebooks
pip install -e ".[notebooks]"

# Weights & Biases logging
pip install -e ".[wandb]"

# LoRA fine-tuning
pip install -e ".[lora]"

# Everything at once
pip install -e ".[all]"

3. Run tests

pytest

---

Stage-by-Stage Guide

Stage 0 – Mimi Pre-encoding

Pre-encode your audio dataset with Mimi to avoid redundant inference during training.

python scripts/encode_tokens.py \
  --audio_dir /path/to/audio \
  --output_dir data/tokens \
  --model_id kyutai/moshika-pytorch-bf16

Stage 1 – Layer Analysis

Open notebooks/01_layer_analysis.ipynb and follow the cells to compute per-layer importance scores (gradient-based) and BI scores for the depth transformer.

Config: configs/stage1_analysis.yaml

Stage 2 – Structured Pruning

Remove the N least-important depth-transformer layers identified in Stage 1.

Config: configs/stage2_pruning.yaml

Stage 3 – Depth-Transformer Evaluation

Evaluate the pruned model on your validation set before distillation.

Stage 4a – Teacher Pre-computation

Cache teacher hidden-state targets once to speed up KD training.

python scripts/precompute_teacher.py \
  --config configs/stage4b_training.yaml \
  --output_dir data/teacher_targets

Stage 4b – Knowledge Distillation

Train the student model with soft-label and feature-level KD.

Config: configs/stage4b_training.yaml

Stage 6 – Quantization & Export

Apply post-training quantization and export the final model.

python scripts/export_model.py \
  --checkpoint experiments/latest/checkpoint_best.pt \
  --output_dir exports/

Config: configs/stage6_quantization.yaml

Stage 7 – LoRA Fine-tuning (Optional)

Fine-tune the quantized student with LoRA on a domain-specific dataset.

---

SQA Evaluation

python scripts/run_sqa_eval.py \
  --model_path exports/moshilite_final \
  --dataset /path/to/sqa_dataset \
  --output_dir results/sqa

---

Key Concepts

• Mimi – Kyutai's neural audio codec used as Moshi's tokenizer.
• BI Score – Block Importance score: measures how much a layer contributes to output quality; used to rank layers for pruning.
• KD – Knowledge Distillation: the student learns from the teacher's soft targets and intermediate hidden states.
• SQA – Speech Question Answering benchmark used to assess end-to-end model quality.

---

License

MIT