Retnet-Distillation

KO | EN

gitlite — search

#javascript #react #machine-learning #deep-learning #linux #docker #cli #awesome-list #vue #security #pytorch #mcp #ios #open-source

Retnet-Distillation

(★ 11)

No description available.

Retnet-Distillation Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Second Place Winner of LocalLLM Hackathon
https://www.reddit.com/r/LocalLLM/comments/1qseno9/mod_post_announcing_the_winners_of_the_rlocalllm/

# Post Mortem Thoughts
Having worked with Retnet and other non quadratic designs, I can say that retention is terrible to work with. It is extremely unstable and the torchscale upstream is incomplete and orphaned by Microsoft.

DeltaNet is an all around superior non quadratic architecture. The KL distillation pipeline in this project is solid and versatile. Most of the code is for dealing with instabilities from Retnet.

# RetNet Distillation Pipeline

A hackathon project focused on making next-generation recurrent architectures (RetNet) accessible and trainable on consumer hardware. While Transformers dominate the landscape, their O(N²) complexity limits context scaling. RetNet offers what the authors call the "impossible triangle": **O(1) inference, O(N) training, and competitive performance**.

## Features

### 1. High-Performance Distillation Engine

The core of the project is a modular distillation system that supports three modes:
(VRAM is with out optimizer or additional allocations)
| Mode | Description | VRAM | Best For |
|------|-------------|------|----------|
| **Direct** | Teacher + student on same GPU | ~14GB | Fast iteration, single GPU |
| **Cached** | Pre-computed teacher logits from disk | ~9GB | Production training, reproducibility |
| **Network** | Teacher on separate vLLM server | ~9GB | Multi-node, large teachers |

- **Direct Mode**: Loads Llama-3.2-1B (teacher) and RetNet-500M (student) simultaneously. At 1K sequence length: **0.1s per optimizer step**. At 4K: **0.3s per step**.
- **Cached Mode**: Pre-compute teacher logits to disk, then train without teacher overhead.
- **Network Mode**: Offloads the teacher to a vLLM-compatible server with a custom `/v1/topk` endpoint for **100x faster** logit extraction.

### 2. Advanced Training Stability

Chasing down bugs in the original Titans implementation led to building a comprehensive system for detecting training pathologies:

- **Saddle Point Escape**: Automated detection when the model gets stuck in local minima, with interventions (LR spikes) to escape.
- **Muon Optimizer**: Integrated for superior RetNet performance vs AdamW. Uses Muon for 2D+ parameters, AdamW for 1D.
- **Diversity Regularization**: Custom loss components to ensure the student learns the distribution, not just the mode.

### 3. Production-Ready Infrastructure

- **Pre-tokenized Data Pipeline**: Custom `PretokenizedShardDataset` handles massive datasets with minimal RAM usage.
- **Fragmented Memory Fixes**: Custom PyTorch CUDA allocator configurations to prevent OOM during long runs.
- **WandB Integration**: Full telemetry for loss, gradient norms, evaluations, saddle detection, and memory usage.
- **Teacher Finetuning Pipeline**: **Critical for convergence.** Microsoft research shows 4.5× faster convergence when finetuning the teacher first. Without it, distillation on pure logits struggles to converge.

### 4. TorchScale Patch

RetNet is still experimental in TorchScale. This repo includes a patched version of TorchScale with necessary fixes for stable training.

## Training Metrics

The pipeline includes comprehensive WandB integration for real-time monitoring:

![Training Metrics](https://raw.githubusercontent.com/bigwolfeman/Retnet-Distillation/master/docs/fin-wandb.png)

*Example training run showing loss curves, gradient norms, learning rate schedule, and performance metrics.*

## Quick Start

See **[QUICKSTART.md](QUICKSTART.md)** for the complete walkthrough.

**TL;DR:**

```bash
# 1. Install
git clone <repo-url> && cd retnet-distillation
pip install -e .

# 2. Download data (~5GB)
cd data && python download_curriculum_datasets.py --category instruction_chat

# 3. Tokenize
python preprocess_to_parquet.py \
    --input data/distillation \
    --output data/distillation_preprocessed \
    --recursive

# 4. (Strongly Recommended) Finetune teacher
python scripts/finetune_teacher.py --config configs/teacher_ft.yaml

# 5. Train!
python -m src.distillation.scripts.train --config configs/train_direct.yaml
```

## Architecture

The current model is a **500M parameter RetNet** trained on instruction-tuning data. By distilling from a finetuned Llama-3.2-1B-Instruct, we bypass the trillions of tokens usually required for pre-training and jump straight to a usable, instruction-following recurrent model.

**Model Variants:**
- `350M`: Lighter, faster training
- `500M`: Default, better quality (configured in `configs/train_direct.yaml`)

## Training Modes

### Direct Mode (Recommended for Single GPU)

```bash
python -m src.distillation.scripts.train --config configs/train_direct.yaml
```

- Loads both models on same GPU
- Fastest feedback loop, no network overhead
- Requires ~14GB VRAM (teacher 4GB + student 2GB + activations)

### Cached Mode (Best for Production)

```bash
# 1. Cache teacher logits
python scripts/cache_teacher_logits.py \
    --data-path data/distillation_preprocessed \
    --output-dir data/teacher_cache

# 2. Train from cache
python -m src.distillation.scripts.train --config configs/train_cached.yaml
```

- Pre-compute teacher logits once, reuse forever
- Reproducible training (exact same logits each epoch)
- Only ~9GB VRAM during training

### Network Mode (Multi-Node / Large Teachers)

```bash
# Terminal 1: Start vLLM server (on GPU machine)
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.2-1B-Instruct \
    --port 8080

# Terminal 2: Train (can be on different machine)
python -m src.distillation.scripts.train --config configs/train_network.yaml
```

- Decouples teacher from training process
- Scale to larger teachers (70B+) on separate hardware
- Custom `/v1/topk` endpoint for 100x faster logit extraction

See **[README_FAST_ENDPOINT.md](README_FAST_ENDPOINT.md)** for the custom vLLM endpoint setup.

## Configuration

All training configs are in `configs/`:

| Config | Mode | Use Case |
|--------|------|----------|
| `train_direct.yaml` | Direct | Default, single GPU training |
| `train_cached.yaml` | Cached | Production, reproducible |
| `train_network.yaml` | Network | Multi-node, large teachers |
| `teacher_ft.yaml` | - | Teacher finetuning |

Key parameters in `train_direct.yaml`:

```yaml
model_variant: "500M"           # Student size
max_seq_length: 1024            # Sequence length
teacher_model: "meta-llama/Llama-3.2-1B-Instruct"
teacher_topk: 512               # Top-k logits from teacher
distill_alpha: 0.5              # CE vs KL balance
```

## Project Structure

```
├── configs/                 # Training configurations
├── data/                    # Data scripts and storage
│   ├── download_curriculum_datasets.py  # Main data downloader
│   ├── preprocess_to_parquet.py         # Tokenization script
│   └── QUICKSTART.md                    # Data pipeline guide
├── scripts/                 # Utility scripts
│   ├── finetune_teacher.py  # Teacher finetuning
│   ├── cache_teacher_logits.py
│   └── evaluate_checkpoint.py
├── src/
│   ├── distillation/        # Core distillation code
│   │   ├── scripts/train.py # Main training script
│   │   └── dataset.py       # Data loaders
│   └── models/              # Model architectures
├── torchscale/              # Patched TorchScale (RetNet)
└── tests/                   # Test suite
```

## Legacy Code Notice

To preserve lessons from the original Titans/RheaNet exploration, this repository still contains legacy modules (for example under `src/models/hrm/` and the `tests/titans/` suites). These files are **not** part of the current distillation pipeline but remain for reference. When extending the project, focus on the RetNet distillation paths documented above and treat the Titans/HRM code as archival.

## Requirements

- **Python**: 3.10 or 3.11
- **GPU**: NVIDIA with ≥16GB VRAM (24GB+ recommended)
- **CUDA**: 11.8+
- **HuggingFace Token**: Required for Llama models

## License

MIT License. See LICENSE file.

The bundled TorchScale library is also under MIT License (Copyright Microsoft Corporation).

## Acknowledgments

- Microsoft TorchScale team for the RetNet implementation
- The knowledge distillation research community (MiniLLM, NeMo-Aligner)
- The Titans paper authors for the original inspiration

---

**Happy Distilling!** 🚀

Retnet-Distillation

Retnet-Distillation Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Retnet-Distillation

Retnet-Distillation Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README