minimind/CLAUDE.md

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

MiniMind is an ultra-small language model (~26M-104M params) designed for training from scratch on consumer GPUs. It demonstrates the complete LLM training pipeline at minimal cost. The codebase is a Chinese-language educational project with English code.

Common Commands

Training (all run from trainer/ directory)

# Stage 1: Pretraining
cd trainer && torchrun --nproc_per_node=1 train_pretrain.py

# Stage 2: Supervised Fine-Tuning
cd trainer && torchrun --nproc_per_node=1 train_full_sft.py

# Optional: LoRA / DPO / PPO / GRPO / Agent RL / Distillation
cd trainer && torchrun --nproc_per_node=1 train_lora.py
cd trainer && torchrun --nproc_per_node=1 train_dpo.py
cd trainer && torchrun --nproc_per_node=1 train_ppo.py
cd trainer && torchrun --nproc_per_node=1 train_grpo.py
cd trainer && torchrun --nproc_per_node=1 train_agent.py
cd trainer && torchrun --nproc_per_node=1 train_distillation.py

Multi-GPU: change --nproc_per_node=N. Resume training with --from_resume 1.

Inference

python eval_llm.py                                          # CLI inference
python eval_llm.py --from_transformers                      # Load transformers-format model
python eval_llm.py --lora_path <path>                       # With LoRA adapter
python scripts/serve_openai_api.py                          # OpenAI-compatible API server

Model Conversion

python scripts/convert_model.py                             # torch <-> transformers format, LoRA merge

Dependencies

pip install -r requirements.txt
# PyTorch must be installed separately (commented out in requirements.txt)

Architecture

Training Pipeline (sequential)

Pretrain → Full SFT → optional (LoRA / DPO / PPO / GRPO / Agent RL / Distillation)

Model (model/model_minimind.py)

• MiniMindConfig: HuggingFace PretrainedConfig with hidden_size=768, 8 layers, vocab_size=6400
• MiniMindForCausalLM: Main model class (inherits PreTrainedModel + GenerationMixin)
• Components: RMSNorm, RoPE + YaRN extrapolation, GQA (8 query heads, 4 KV heads), SwiGLU FFN, MoE with top-K routing and aux loss
• Weight tying between input embedding and LM head (tie_word_embeddings)
• Custom generate() with top-k/top-p sampling, repetition penalty, KV-cache
• Model sizes controlled by hidden_size: 512 (minimind-3, 4 layers), 640 (minimind-2, 6 layers), 768 (default, 8 layers)

LoRA (model/model_lora.py)

Pure PyTorch LoRA implementation (rank-16 A*B matrices). Functions: apply_lora(), load_lora(), save_lora(), merge_lora().

Datasets (dataset/lm_dataset.py)

• PretrainDataset: plain text JSONL
• SFTDataset: multi-turn chat conversations
• DPODataset: chosen/rejected pairs
• RLAIFDataset: prompts for RL training
• AgentRLDataset: multi-turn with tool calls

All data is JSONL format in dataset/.

Shared Utilities (trainer/trainer_utils.py)

• init_model(): model + tokenizer initialization, handles DDP unwrapping and torch.compile
• lm_checkpoint(): atomic checkpoint save/load with os.replace(), handles multi-GPU world_size changes on resume
• init_distributed_mode(): DDP setup
• setup_seed(): deterministic seeding
• LMForRewardModel: wrapper for reward models (uses InternLM2-1.8B-Reward)

Logging

SwanLab (default) or WandB (--use_wandb). Checkpoints store wandb_id for resume.

Tokenizer (model/)

Custom BPE tokenizer with 6,400 vocab, 36 special tokens. Chat template supports tools, multi-turn, and reasoning_content. train_tokenizer.py is reference-only (pre-trained tokenizer is used).

Key Conventions

• No test framework exists; scripts/eval_toolcall.py provides tool call evaluation
• All training scripts use argparse with defaults matching the minimind-768 config
• Checkpoints are saved to ../checkpoints/ relative to trainer/
• Code comments are in Chinese