# Model Training Guide Learn how to train MiniMind language models from scratch using pure PyTorch. ## ๐Ÿ“Š Training Overview MiniMind implements a complete training pipeline: ``` Tokenizer Training โ†“ Pretraining (Learn knowledge) โ†“ SFT (Learn conversation) โ†“ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ†“ โ†“ โ†“ โ†“ LoRA DPO/RLHF RLAIF (PPO/GRPO/SPO) Distillation (Domain adapt) (Preference) (Reinforcement Learn) (Reasoning) ``` ## ๐Ÿ’ฐ Training Costs (Single NVIDIA 3090) | Model | Dataset | Duration | Cost (RMB) | Quality | |-------|---------|----------|-----------|---------| | MiniMind2-Small | pretrain_hq + sft_mini_512 | 2.1h | โ‰ˆ3 | ๐Ÿ˜Š๐Ÿ˜Š | | MiniMind2-Small | Full dataset | 38h | โ‰ˆ50 | ๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š | | MiniMind2 | pretrain_hq + sft_mini_512 | 3.3h | โ‰ˆ5 | ๐Ÿ˜Š๐Ÿ˜Š | | MiniMind2 | Full dataset | 122h | โ‰ˆ160 | ๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š๐Ÿ˜Š | !!! success "Ultra-Fast Training" **Just 2.1 hours + $3 = Functional ChatBot!** Use `pretrain_hq.jsonl` + `sft_mini_512.jsonl` for fastest reproduction ## ๐Ÿ“‹ Data Preparation ### 1. Download Datasets Download from [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset) or [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset): ```bash mkdir -p dataset cd dataset # Download required files ``` ### 2. Dataset Directory Structure ``` ./dataset/ โ”œโ”€โ”€ pretrain_hq.jsonl โœจ (1.6GB, required for pretraining) โ”œโ”€โ”€ sft_mini_512.jsonl โœจ (1.2GB, fastest SFT) โ”œโ”€โ”€ sft_512.jsonl (7.5GB, standard SFT) โ”œโ”€โ”€ sft_1024.jsonl (5.6GB, longer SFT) โ”œโ”€โ”€ sft_2048.jsonl (9GB, very long SFT) โ”œโ”€โ”€ dpo.jsonl (909MB, DPO training) โ”œโ”€โ”€ r1_mix_1024.jsonl (340MB, reasoning distillation) โ”œโ”€โ”€ rlaif-mini.jsonl (1MB, RLAIF algorithms) โ”œโ”€โ”€ lora_identity.jsonl (22.8KB, identity LoRA) โ””โ”€โ”€ lora_medical.jsonl (34MB, medical domain LoRA) ``` ### 3. Data Formats **Pretraining Data** (`pretrain_hq.jsonl`): ```json {"text": "How to overcome procrastination? Overcoming procrastination is not easy, but these suggestions may help..."} ``` **SFT Data** (`sft_*.jsonl`): ```json { "conversations": [ {"role": "user", "content": "Hello!"}, {"role": "assistant", "content": "Hello! How can I help?"}, {"role": "user", "content": "Tell me a joke."}, {"role": "assistant", "content": "Why did the scarecrow win an award? Because he was outstanding in his field!"} ] } ``` **DPO Data** (`dpo.jsonl`): ```json { "chosen": [ {"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "2+2 equals 4."} ], "rejected": [ {"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "2+2 equals 5."} ] } ``` **LoRA Domain Data** (`lora_*.jsonl`): ```json { "conversations": [ {"role": "user", "content": "What's the treatment for cervical spondylosis?"}, {"role": "assistant", "content": "Cervical spondylosis treatment typically includes..."} ] } ``` ## ๐ŸŽฏ Complete Training Pipeline All training scripts are in the `./trainer` directory. ```bash cd trainer ``` ### Stage 1: Pretraining **Purpose**: Learn foundational knowledge (word continuation) ```bash # Single GPU python train_pretrain.py # Multi-GPU (DDP) torchrun --nproc_per_node 2 train_pretrain.py # Multi-GPU (DeepSpeed) deepspeed --master_port 29500 --num_gpus=2 train_pretrain.py ``` **Key Parameters**: - `max_seq_len`: 512 (adjust based on GPU memory) - `learning_rate`: 1e-4 - `epochs`: Adjust based on dataset size **Output**: `./out/pretrain_*.pth` **Training Duration**: - MiniMind2-Small (26M): ~1.1h - MiniMind2 (104M): ~3.9h !!! tip "Pretraining Tips" - Start with `pretrain_hq.jsonl` for best results - Quality > Quantity for pretraining data - Monitor loss curve to detect overfitting ### Stage 2: Supervised Fine-Tuning (SFT) **Purpose**: Teach conversation patterns and chat templates ```bash # Single GPU python train_full_sft.py # Multi-GPU torchrun --nproc_per_node 2 train_full_sft.py ``` **Configuration**: - Load pretrained model from Stage 1 - Use SFT dataset (`sft_mini_512.jsonl` or `sft_512.jsonl`) - Adjust `max_seq_len` to match training data **Output**: `./out/full_sft_*.pth` **Training Duration**: - With sft_mini_512: 1-3 hours - With full sft_512: 20-25 hours !!! warning "SFT Data Selection" - `sft_mini_512.jsonl`: Fastest, ~1.2GB, 512 tokens max - `sft_512.jsonl`: Standard, ~7.5GB, 512 tokens max - `sft_1024.jsonl`: Longer, ~5.6GB, 1024 tokens max - `sft_2048.jsonl`: Extended, ~9GB, 2048 tokens max ### Stage 3: LoRA Fine-Tuning (Optional) **Purpose**: Parameter-efficient domain adaptation **Use Cases**: - Medical Q&A knowledge - Personal identity/self-awareness - Proprietary domain knowledge ```bash # Edit train_lora.py to set correct dataset and base model python train_lora.py # Multi-GPU torchrun --nproc_per_node 2 train_lora.py ``` **Output**: `./out/lora/lora_*.pth` **Example 1: Medical Domain** Prepare `dataset/lora_medical.jsonl`: ```json { "conversations": [ {"role": "user", "content": "What's the correct pillow height for cervical spondylosis?"}, {"role": "assistant", "content": "For cervical spondylosis, pillow height should be..."} ] } ``` Train: ```bash # Modify train_lora.py: lora_name = 'medical' python train_lora.py ``` **Example 2: Identity/Self-Awareness** Prepare `dataset/lora_identity.jsonl`: ```json { "conversations": [ {"role": "user", "content": "Who are you?"}, {"role": "assistant", "content": "I am MiniMind..."} ] } ``` ### Stage 4: Direct Preference Optimization (DPO) **Purpose**: Align model responses with human preferences DPO eliminates the need for separate reward models by directly optimizing preference pairs. ```bash python train_dpo.py # Multi-GPU torchrun --nproc_per_node 2 train_dpo.py ``` **Output**: `./out/rlhf_*.pth` **Key Features**: - Off-policy training (reuse data across epochs) - No separate reward model needed - Better sample efficiency than PPO - Stable training convergence **Training Duration**: ~1-3 hours ### Stage 5: Reinforcement Learning from AI Feedback (RLAIF) RLAIF is an advanced training approach using AI-generated rewards instead of human annotations. MiniMind implements three modern algorithms: #### 5.1 PPO (Proximal Policy Optimization) Classical on-policy RL algorithm with proven stability. ```bash python train_ppo.py # Multi-GPU torchrun --nproc_per_node 2 train_ppo.py ``` **Algorithm**: $$\mathcal{L}_{PPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \text{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)\right] + \beta \cdot \mathbb{E}[\text{KL}]$$ **Characteristics**: - Stable but slower reward improvement - Requires both Actor and Critic networks - High memory usage (1.5-2ร— single network) - Good for exploration **Output**: `./out/ppo_actor_*.pth` **Training Duration**: ~1-3 hours #### 5.2 GRPO (Group Relative Policy Optimization) Modern algorithm used in DeepSeek-R1, with faster convergence. ```bash python train_grpo.py # Multi-GPU torchrun --nproc_per_node 2 train_grpo.py ``` **Algorithm**: $$\mathcal{L}_{GRPO} = -\mathbb{E}\left[r_t \cdot A_t - \beta \cdot \text{KL}_t\right]$$ Where advantage is computed as: $$A_t = \frac{R - \mu_{group}}{\sigma_{group}}$$ **Characteristics**: - Single-network design (memory efficient) - Faster reward improvement - Group normalization removes bias - Better convergence stability **Output**: `./out/grpo_*.pth` **Training Duration**: ~1-3 hours #### 5.3 SPO (Single-stream Policy Optimization) Newest algorithm (2025) addressing GRPO's degenerate group problem. ```bash python train_spo.py # Multi-GPU torchrun --nproc_per_node 2 train_spo.py ``` **Algorithm**: $$\mathcal{L}_{SPO} = -\mathbb{E}\left[\log \pi_\theta(a_t|s) \cdot A_t - \beta \cdot \text{KL}_t\right]$$ With adaptive baseline: $B_t^{adaptive}$ **Characteristics**: - No group dependency (1 input โ†’ 1 training sample) - Adaptive value tracking - Better handling of difficult examples - Experimental on small models **Output**: `./out/spo_*.pth` **Training Duration**: ~1-3 hours #### RLAIF Dataset Preparation All RLAIF algorithms use `rlaif-mini.jsonl` (1MB, 10k examples): ```bash # Download dataset # Format: Same as SFT, but assistant content is "ๆ— " (none) { "conversations": [ {"role": "user", "content": "Explain photosynthesis briefly."}, {"role": "assistant", "content": "ๆ— "} ] } ``` The model generates completions during training, which are scored by a **Reward Model** (e.g., InternLM2-1.8B-Reward). **Reward Model Setup**: ```bash # Download reward model to parent directory cd ../ git clone https://huggingface.co/internlm/internlm2-1_8b-reward # Directory structure should be: # project/ # โ”œโ”€โ”€ minimind/ # โ””โ”€โ”€ internlm2-1_8b-reward/ ``` #### RLAIF vs DPO Comparison | Aspect | DPO | RLAIF (PPO/GRPO/SPO) | |--------|-----|---------------------| | Training Type | Off-policy | On-policy | | Data Freshness | Static pairs | Dynamic (generated) | | Reward Source | Implicit | Explicit model | | Convergence | Fast | Slower | | Memory Usage | Lower | Higher | | Best For | Preference refinement | Capability improvement | ### Stage 6: Reasoning Model Distillation **Purpose**: Distill DeepSeek-R1-style reasoning into MiniMind ```bash python train_distill_reason.py # Multi-GPU torchrun --nproc_per_node 2 train_distill_reason.py ``` **Data Format** (`r1_mix_1024.jsonl`): ```json { "conversations": [ { "role": "user", "content": "Solve: 5 + 3 = ?" }, { "role": "assistant", "content": "\nI need to add 5 and 3.\n5 + 3 = 8\n\n\n5 + 3 = 8\n" } ] } ``` **Output**: `./out/reason_*.pth` **Training Features**: - Enforces `` and `` tags - Penalty loss for format violations - Mixed data (reasoning + multi-turn + English) ## ๐Ÿ”ง Multi-GPU Training ### DDP (Distributed Data Parallel) Best for single-machine multi-GPU: ```bash torchrun --nproc_per_node N train_xxx.py # N = number of GPUs ``` ### DeepSpeed For advanced optimization: ```bash deepspeed --master_port 29500 --num_gpus=N train_xxx.py ``` ### Wandb Monitoring Track training progress: ```bash # Login first wandb login # Enable wandb logging torchrun --nproc_per_node N train_xxx.py --use_wandb # Or SwanLab (China-friendly alternative) python train_xxx.py --use_wandb # Automatically uses SwanLab if available ``` ## ๐Ÿงช Model Testing ### Evaluate Pretrain Model ```bash python eval_model.py --model_mode 0 ``` ### Evaluate Chat Model ```bash python eval_model.py --model_mode 1 ``` ### Evaluate with LoRA ```bash python eval_model.py --lora_name 'lora_medical' --model_mode 1 ``` ### Evaluate Reasoning Model ```bash python eval_model.py --model_mode 3 ``` ### Evaluate RLAIF Models ```bash # PPO model python eval_model.py --model_mode 4 # GRPO model python eval_model.py --model_mode 4 ``` ### RoPE Length Extrapolation Test with extended context: ```bash python eval_model.py --model_mode 1 --inference_rope_scaling True ``` ## ๐Ÿ“ Model Architecture ### MiniMind Structure **Decoder-Only Transformer** (similar to Llama3): ``` Input Tokens โ†“ Token Embedding (6400 vocab) โ†“ Rotary Embeddings (RoPE) [with YaRN for length extrapolation] โ†“ [Transformer Blocks] ร—N โ”œโ”€ Attention (Multi-Head) โ”œโ”€ RMSNorm โ”œโ”€ SwiGLU FFN [or MoE for MoE variant] โ””โ”€ Residual Connections โ†“ RMSNorm โ†“ LM Head (โ†’ 6400 vocab logits) โ†“ Output Probabilities ``` ### Model Configurations | Config | MiniMind2-Small | MiniMind2 | MiniMind2-MoE | |--------|-----------------|----------|---------------| | Parameters | 26M | 104M | 145M | | Hidden Dim | 512 | 768 | 640 | | Layers | 8 | 16 | 8 | | KV Heads | 2 | 2 | 2 | | Q Heads | 8 | 8 | 8 | | Vocab Size | 6,400 | 6,400 | 6,400 | | Context Length | 2,048 | 2,048 | 2,048 | ### Modify Architecture Edit `./model/LMConfig.py`: ```python class LMConfig: hidden_size: int = 768 num_layers: int = 16 num_heads: int = 8 num_kv_heads: int = 2 # ... other configs ``` ## ๐Ÿ” Training Tips & Best Practices ### Data Quality > Quantity - High-quality pretraining data accelerates convergence - `pretrain_hq.jsonl` is carefully curated for quality - Consider data deduplication and cleaning ### Learning Rate Scheduling ```python # Recommended schedules - Linear warmup then decay - Initial: 1e-4 to 5e-4 - Warmup steps: 10% of total - Final: 10% of initial LR ``` ### Batch Size & Sequence Length ```python # Balance between GPU memory and convergence - Pretraining: max_seq_len=512, batch_size=32 - SFT: max_seq_len=512, batch_size=16 - LoRA: max_seq_len=512, batch_size=16 ``` ### Memory Optimization ```bash # Reduce batch size if OOM python train_xxx.py --batch_size 8 # Or use gradient accumulation python train_xxx.py --gradient_accumulation_steps 4 ``` ### Checkpoint Management - Saves every 100 steps by default - Each new save overwrites the old one - Automatic backup before training ## ๐Ÿšจ Common Issues & Solutions ### Issue: CUDA Out of Memory ```bash # Solution 1: Reduce batch size python train_xxx.py --batch_size 4 # Solution 2: Use gradient accumulation python train_xxx.py --batch_size 16 --gradient_accumulation_steps 2 # Solution 3: Use smaller model # Edit trainer script to use MiniMind2-Small instead ``` ### Issue: Training Not Converging ```python # Possible causes: 1. Learning rate too high/low 2. Data quality issues 3. Model capacity mismatch # Solutions: - Reduce learning rate: --learning_rate 1e-5 - Check data format and quality - Try smaller model first ``` ### Issue: Multi-GPU Sync Errors ```bash # Ensure: 1. All GPUs visible: nvidia-smi 2. Same CUDA version across all GPUs 3. Network connectivity for distributed training # Debug: torchrun --nproc_per_node 2 train_xxx.py --debug ``` ### Issue: Different Results Than Expected ```python # Check: 1. Random seed set (reproducibility) 2. Correct model checkpoint loaded 3. Correct dataset being used 4. Same hyperparameters as reference ``` ## ๐Ÿ“ˆ Training Progression Typical training curves: ``` Pretraining Loss: โ†˜โ†˜โ†˜ (steep decline, then plateau) SFT Loss: โ†˜ (steady decline) PPO Reward: โ†— (rising, may plateau) GRPO Reward: โ†—โ†— (faster rise, more stable) ``` ## ๐ŸŽ“ Advanced Topics ### Custom Datasets Create your own dataset: ```python # Format: JSONL with conversations list # Each line is one training example # Ensure consistent quality and format ``` ### Model Quantization (Post-training) ```bash # 4-bit quantization for inference # Use tools like: # - llama.cpp (gguf format) # - bitsandbytes (dynamic quantization) # - AutoGPTQ (static quantization) ``` ### Model Merging ```python # Merge base model + LoRA weights # Use tools like: peft, llama.cpp ``` ## ๐Ÿ“š References - [Scaling Laws](https://arxiv.org/pdf/2001.08361.pdf) - [RoPE Position Embeddings](https://arxiv.org/abs/2104.09864) - [YaRN Length Extrapolation](https://arxiv.org/abs/2309.00071) - [PPO Algorithm](https://arxiv.org/abs/1707.06347) - [GRPO (DeepSeek)](https://arxiv.org/pdf/2402.03300) - [SPO Algorithm](https://arxiv.org/abs/2509.13232) - [DPO](https://arxiv.org/abs/2305.18290) --- **Next**: Deploy your trained model or explore [advanced inference options](quickstart.md#third-party-inference-frameworks)