strong-distill

KO | EN

gitlite — search

#javascript #python #golang #go #machine-learning #nodejs #java #windows #awesome-list #vue #security #mcp #ios #open-source

(★ 9)

No description available.

strong-distill Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Strong Teacher Not Needed? On Distillation in LLM Pretraining

[arxiv] [Taiming Lu](https://taiminglu.com/), [Zhuang Liu](https://liuzhuang13.github.io)  
Princeton University

---

Knowledge distillation generally assumes a **strong-to-weak** relationship where stronger teachers yield better students. In this work we examine this assumption in **large language model pretraining**. By varying architecture sizes and training token budgets, we create strong-to-weak, same-level, and weak-to-strong teacher–student relationships, and study distillation's effectiveness under each with different mixes between language-modeling loss and distillation loss. Three findings emerge:

1. With proper loss mixing, **weak-to-strong and same-level distillation improves over standard pretraining** — even small and undertrained teachers benefit large students.
2. Making the teacher stronger can lead to **saturated or even reversed gains**.
3. Distillation improves **generalization** (out-of-distribution, downstream) more readily than in-domain fitting.

Our results challenge the common belief that distillation pretraining requires a strong teacher.

## Prerequisites

- **Hardware** — a Google Cloud TPU VM pod. Paper runs use v6e-64 and v4-128; smaller topologies work after adjusting `per_device_batch_size` in the generator templates.
- **Environment** — clone this repository into `~/maxtext` on the coordinator host and install dependencies into a virtualenv at `~/maxtext_env`:
  ```bash
  git clone <REPO_URL> ~/maxtext
  cd ~/maxtext
  python3.10 -m venv ~/maxtext_env
  source ~/maxtext_env/bin/activate
  pip install -r requirements.txt
  ```
- **Data** — pretokenize a FineWeb-Edu split into `array_record` shards with [`quickstart/stream_fineweb_to_tokens.py`](quickstart/stream_fineweb_to_tokens.py).
- **Runtime env vars** — export before running any script:
  ```bash
  export BUCKET_NAME=<your-gcs-bucket>
  export TPU_PREFIX=<your-tpu-prefix>
  export WANDB_API_KEY=<your-wandb-key>   # optional; omit to disable wandb
  ```

## Pipeline

### 1. Pretrain the teacher

Generate a per-configuration job script, then launch it:

```bash
python scripts/pretrain/generate_scripts.py \
    --archs 1b --tokens 50B --seeds 42 \
    --output-dir scripts/pretrain
bash scripts/pretrain/llama1b-finewebedu-vanilla-s42-50b.sh
```

Full checkpoints are written to `gs://$BUCKET_NAME/ckpts/pretrain/<run_name>/checkpoints/<step>/items`.

### 2. Convert the teacher to a param-only checkpoint

Strip optimizer state and downcast to bfloat16 so the teacher can be loaded as a frozen distribution during student training. The `run_name` below must match the naming convention expected by the distillation scripts: `llama{arch}-vanilla-{tokens}-s{seed}`.

```bash
python -m MaxText.generate_param_only_checkpoint MaxText/configs/base.yml \
    model_name=llama3.1-1b \
    load_full_state_path=gs://$BUCKET_NAME/ckpts/pretrain/llama3.1-1b-finewebedu-vanilla-s42-50b/checkpoints/24999/items \
    base_output_directory=$HOME/ckpts/pretrain_param_only \
    run_name=llama1b-vanilla-50B-s42
```

### 3. Distill into the student

Pick the experiment to reproduce. Each generator defaults to the paper's sweep; narrow with flags.

```bash
# Experiment 1: vary teacher size at fixed teacher tokens (50B), sweep alpha
python scripts/exp1/generate_scripts.py --output-dir scripts/exp1
bash scripts/exp1/exp1_llama1b-A1BT50BS42-a05-s43.sh

# Experiment 2: vary teacher training compute at fixed teacher size, sweep alpha
python scripts/exp2/generate_scripts.py --teacher-archs 1b --output-dir scripts/exp2
bash scripts/exp2/exp2_llama1b-A1BT80BS42-a05-s43.sh
```

Full sweep definitions and env-var documentation live in [`scripts/README.md`](scripts/README.md).

## Implementation

The distillation loss, teacher-distribution filtering (top-$k$ / top-$p$ with an optional "OTHER" bucket), temperature scaling, and hard-label variants are implemented on top of [MaxText](https://github.com/google/maxtext) in `MaxText/train.py` and `MaxText/max_utils.py`.

- **[Pretrain](scripts/pretrain/)** — Vanilla CE pretraining for teacher checkpoints
- **[Experiment 1](scripts/exp1/)** — Distillation ablation varying teacher size
- **[Experiment 2](scripts/exp2/)** — Distillation ablation varying teacher training compute
- **[Quickstart](quickstart/)** — Single-job reproductions and FineWeb-Edu data preparation

## Acknowledgement

This repository is built on top of [MaxText](https://github.com/google/maxtext), Google's high-performance JAX/TPU training framework for large language models. We gratefully acknowledge the generous support of the [Google TPU Research Cloud (TRC)](https://sites.research.google/trc/about/), which provided the computational resources used to train every model in this study.

## License

Apache 2.0, inherited from MaxText. See the [LICENSE](LICENSE) file for details.

## Citation

```bibtex
@misc{lu2026strongKD,
  title={Strong Teacher Not Needed? On Distillation in LLM Pretraining},
  author={Lu, Taiming and Liu, Zhuang},
  year={2026},
}
```

strong-distill

strong-distill Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

strong-distill

strong-distill Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README