MiniLLM

MiniLLM is a lightweight LLM training and inference project.

Folder Structure

minillm/
├─dataset/
│  ├─__init__.py
│  └─lm_dataset.py
├─infer/
│  ├─__init__.py
│  ├─chat_openai_api.py
│  ├─convert_model.py
│  ├─serve_openai_api.py
│  └─web_demo.py
├─model/
│  ├─config.py
│  ├─model_lora.py
│  ├─model_minillm.py
│  └─triton_flash_attn.py
├─test/
│  ├─__init__.py
│  └─eval_llm.py
├─tokenizer/
│  ├─tokenizer.json
│  └─tokenizer_config.json
└─trainer/
   ├─__init__.py
   ├─trainer_utils.py
   ├─train_pretrain.py
   ├─train_full_sft.py
   ├─train_dpo.py
   ├─train_reason.py
   ├─train_distillation.py
   ├─train_lora.py
   ├─train_ppo.py
   ├─train_grpo.py
   ├─train_spo.py
   └─train_tokenizer.py

Quick Start

# 1) create virtual env
uv venv --prompt minillm --python 3.12
source .venv/bin/activate

# 2) install dependencies
uv sync

Training Entrypoints

Run from repository root using module mode (python -m ...).

Pretraining

Script: minillm/trainer/train_pretrain.py

torchrun --standalone --nproc_per_node=1 \
  minillm/trainer/train_pretrain.py \
  --model_config examples/configs/common/model_config.small.json \
  --train_config examples/configs/pretrain/train_config.json \
  --script_config examples/configs/pretrain/script_config.json

Test the trained Model

python minillm/test/eval_llm.py \
  --load_from out/pretrain-small \
  --tokenizer_path minillm/tokenizer \
  --data_path datasets/sft_t2t_mini.jsonl \
  --sample_index 0 \
  --max_samples 5 \
  --max_new_tokens 128 \
  --temperature 0.8 \
  --top_p 0.9

Full SFT

Script: minillm/trainer/train_full_sft.py

torchrun --standalone --nproc_per_node=1 \
  minillm/trainer/train_full_sft.py \
  --model_config examples/configs/common/model_config.small.json \
  --train_config examples/configs/full_sft/train_config.json \
  --script_config examples/configs/full_sft/script_config.json

DPO

DPO即相对偏好优化，通过最大化chosen和rejected样本的相对概率差异来训练模型。其损失计算公式为

$$ \mathcal{L}_{\mathrm{DPO}}(\theta)=-\log \sigma\Big(\beta \big[\log \pi_\theta(y^+|x) - \log \pi_\theta(y^-|x)-\big(\log \pi_{\mathrm{ref}}(y^+|x) - \log \pi_{\mathrm{ref}}(y^-|x)\big)\big]\Big) $$

Script: minillm/trainer/train_dpo.py

torchrun --standalone --nproc_per_node=1 \
  minillm/trainer/train_dpo.py \
  --model_config examples/configs/common/model_config.small.json \
  --train_config examples/configs/dpo/train_config.json \
  --script_config examples/configs/dpo/script_config.json

Reason Distillation

Script: minillm/trainer/train_reason.py

torchrun --standalone --nproc_per_node=1 \
  minillm/trainer/train_dpo.py \
  --model_config examples/configs/common/model_config.small.json \
  --train_config examples/configs/dpo/train_config.json \
  --script_config examples/configs/dpo/script_config.json

Knowledge Distillation

Script: minillm/trainer/train_distillation.py

python -m minillm.trainer.train_distillation \
  --data_path ./dataset/sft_mini_512.jsonl \
  --save_dir ./out \
  --epochs 6 \
  --batch_size 32

LoRA Fine-tuning

Script: minillm/trainer/train_lora.py

python -m minillm.trainer.train_lora \
  --data_path ./dataset/lora_identity.jsonl \
  --save_dir ./out/lora \
  --epochs 50 \
  --batch_size 32

RLHF-style Trainers

• PPO: minillm/trainer/train_ppo.py
• GRPO: minillm/trainer/train_grpo.py
• SPO: minillm/trainer/train_spo.py

Example:

python -m minillm.trainer.train_grpo \
  --data_path ./dataset/rlaif-mini.jsonl \
  --save_dir ./out \
  --epochs 1 \
  --batch_size 2

Inference

Local Chat / Eval

Script: minillm/test/eval_llm.py

python -m minillm.test.eval_llm \
  --load_from model \
  --save_dir out \
  --weight full_sft

OpenAI-Compatible API Server

Script: minillm/infer/serve_openai_api.py

python -m minillm.infer.serve_openai_api \
  --load_from ../model \
  --save_dir out \
  --weight full_sft \
  --hidden_size 512 \
  --num_hidden_layers 8

Streamlit Web Demo

Script: minillm/infer/web_demo.py

streamlit run minillm/infer/web_demo.py

Model Conversion

Script: minillm/infer/convert_model.py

python -m minillm.infer.convert_model

Triton Flash Attention

Triton implementation has been migrated to minillm/model/triton_flash_attn.py and integrated into minillm/model/model_minillm.py.

The model will use Triton flash-attention only when all conditions are met:

• CUDA is available
• inference/eval mode (not training)
• Q/K/V are CUDA tensors with same dtype
• dtype is float16 or bfloat16
• head dimension is one of 16, 32, 64, 128
• attention mask is all-ones (or None)

If conditions are not met, it automatically falls back to SDPA.

Acknowledgements

• MiniMind
• llm-tap-rl
• trl
• simple_grpo

License

MIT License. See LICENSE.