104 lines
4.0 KiB
Markdown
104 lines
4.0 KiB
Markdown
# 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)
|
|
|
|
```bash
|
|
# 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
|
|
|
|
```bash
|
|
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
|
|
|
|
```bash
|
|
python scripts/convert_model.py # torch <-> transformers format, LoRA merge
|
|
```
|
|
|
|
### Dependencies
|
|
|
|
```bash
|
|
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
|