# MiniMind Admin & Developer Guide

This guide is the architectural and operational reference for MiniMind — for developers, ML engineers, and architects who want to understand how the project is built, extend it, or reproduce its results. Where the User Guide explains *how to use* MiniMind, this guide explains *how it works*: the model architecture, the training pipeline internals, the reinforcement-learning algorithms, the data abstractions, the full configuration surface, and the security and quality characteristics of the codebase.

MiniMind is an educational, from-scratch LLM project. Every core algorithm — the transformer, the tokenizer, LoRA, DPO, PPO, GRPO, CISPO, distillation, and the agentic RL loop — is implemented in native PyTorch without relying on high-level framework abstractions (`trl`, `peft`, etc., are present as utilities but the learning algorithms are hand-written). This makes the repository unusually readable and an excellent reference implementation, at the cost of being a research/learning codebase rather than a production system.

---

## Table of Contents

1. [Executive Summary](#1-executive-summary)
2. [Repository Structure](#2-repository-structure)
3. [Architecture Overview](#3-architecture-overview)
4. [Model Architecture Deep Dive](#4-model-architecture-deep-dive)
5. [The Tokenizer and Chat Template](#5-the-tokenizer-and-chat-template)
6. [Data Layer](#6-data-layer)
7. [Training Pipeline Internals](#7-training-pipeline-internals)
8. [Reinforcement Learning Subsystem](#8-reinforcement-learning-subsystem)
9. [Inference and Serving](#9-inference-and-serving)
10. [Configuration Reference](#10-configuration-reference)
11. [Checkpointing and Resume](#11-checkpointing-and-resume)
12. [Security Analysis](#12-security-analysis)
13. [Code Quality and Patterns](#13-code-quality-and-patterns)
14. [Testing and Reproducibility](#14-testing-and-reproducibility)
15. [Potential Improvements](#15-potential-improvements)
16. [Quick Reference](#16-quick-reference)

---

## 1. Executive Summary

MiniMind is a single-machine, script-driven training and inference codebase for small (26M–198M parameter) decoder-only language models. It has no service layer, no database, and no orchestration — operations are performed by running Python scripts directly. State lives entirely on disk as model weights (`./out/`), resumable checkpoints (`./checkpoints/`), and JSONL datasets (`./dataset/`).

**Architectural highlights:**

- **Model** (`model/model_minimind.py`): a Qwen3-aligned decoder-only transformer with RMSNorm pre-normalization, SwiGLU feed-forward, Rotary Position Embeddings (RoPE) with optional YaRN extrapolation, grouped-query attention (GQA) with per-head QK normalization, and an optional top-1 Mixture-of-Experts feed-forward. It subclasses Hugging Face `PreTrainedModel` and `GenerationMixin` for ecosystem compatibility, but implements its own `generate()`.
- **Training** (`trainer/`): eleven scripts, each a self-contained ~9-section scaffold covering pretraining, SFT, LoRA, DPO, PPO, GRPO, CISPO, agentic RL, distillation, and tokenizer training. All support DDP multi-GPU, mixed precision, gradient accumulation, checkpoint resume, and WandB/SwanLab logging.
- **Data** (`dataset/lm_dataset.py`): five `torch.utils.data.Dataset` subclasses that turn JSONL into model-ready tensors, plus chat-template preprocessing that handles system-prompt augmentation, empty-think-tag stripping, and tool-call expansion.
- **RL decoupling** (`trainer/rollout_engine.py`): a pluggable rollout abstraction with a native PyTorch engine and an SGLang HTTP engine, separating policy optimization from trajectory generation ("train/inference separation").
- **Serving** (`scripts/serve_openai_api.py`): a FastAPI OpenAI-compatible endpoint supporting streaming, `reasoning_content`, `tool_calls`, and `open_thinking`.

**Scale:** ~49 files, ~207K characters of source. Complexity tier: small. The single largest source artifact is the bundled `model/tokenizer.json` (the 6400-token BPE vocabulary), not hand-written code.

---

## 2. Repository Structure

```text
minimind/
├── eval_llm.py              # CLI inference / interactive chat entry point
├── requirements.txt         # Python deps (torch/torchvision intentionally commented out)
├── README.md / README_en.md # Primary documentation (the canonical project spec)
│
├── model/
│   ├── model_minimind.py    # MiniMindConfig + full transformer (Dense + MoE) + generate()
│   ├── model_lora.py        # From-scratch LoRA: apply/load/save/merge
│   ├── tokenizer.json       # 6400-token BPE vocabulary (ByteLevel)
│   └── tokenizer_config.json# Special tokens + chat_template (Jinja2)
│
├── dataset/
│   └── lm_dataset.py        # PretrainDataset, SFTDataset, DPODataset, RLAIFDataset, AgentRLDataset
│
├── trainer/
│   ├── train_pretrain.py    # Stage 1: next-token pretraining
│   ├── train_full_sft.py    # Stage 2: supervised fine-tuning
│   ├── train_lora.py        # PEFT LoRA fine-tuning
│   ├── train_dpo.py         # RLHF via Direct Preference Optimization
│   ├── train_ppo.py         # RLAIF via PPO (Actor + Critic)
│   ├── train_grpo.py        # RLAIF via GRPO / CISPO
│   ├── train_agent.py       # Agentic RL (multi-turn tool use)
│   ├── train_distillation.py# White-box knowledge distillation
│   ├── train_tokenizer.py   # BPE tokenizer training (reference)
│   ├── rollout_engine.py    # Pluggable rollout: TorchRolloutEngine, SGLangRolloutEngine
│   └── trainer_utils.py     # Shared helpers: init_model, lm_checkpoint, schedulers, reward model
│
└── scripts/
    ├── serve_openai_api.py  # FastAPI OpenAI-compatible server (:8998)
    ├── chat_api.py          # OpenAI-SDK test client
    ├── web_demo.py          # Streamlit chat UI
    ├── eval_toolcall.py     # Tool-calling test harness
    └── convert_model.py     # torch <-> Transformers conversion; LoRA merge
```

Folder triage during analysis ranked `model/`, `trainer/`, `dataset/`, and the repository root as HIGH importance, with `scripts/` and `images/` as supporting material.

---

## 3. Architecture Overview

MiniMind has no runtime architecture in the service sense — it is a set of scripts that read data and weights from disk and write weights back. The conceptual data flow across the full lifecycle is:

```text
                         ┌─────────────────────────────────────────────┐
   raw JSONL data        │                model/                       │
   (./dataset/)          │  MiniMindConfig → MiniMindForCausalLM        │
        │                │  (Dense or MoE) + tokenizer (6400 BPE)       │
        ▼                └─────────────────────────────────────────────┘
   dataset/lm_dataset.py            ▲                     ▲
   (Dataset → tensors)              │ build/init          │ load/save weights
        │                           │                     │
        ▼                           │                     ▼
   ┌──────────────────────────────────────────────────────────────────┐
   │                          trainer/                                 │
   │  pretrain → full_sft → {lora | dpo | ppo | grpo | agent | dist}   │
   │     │            │              │ (RL stages use rollout_engine)  │
   │     ▼            ▼              ▼                                  │
   │  ./out/*.pth   (+ ./checkpoints/*_resume.pth for resume)          │
   └──────────────────────────────────────────────────────────────────┘
        │
        ▼
   ┌──────────────────────────────────────────────────────────────────┐
   │  Inference:  eval_llm.py (CLI)                                    │
   │              scripts/serve_openai_api.py (FastAPI :8998)          │
   │              scripts/web_demo.py (Streamlit)                      │
   │  External:   ollama / vLLM / SGLang / llama.cpp / MNN             │
   └──────────────────────────────────────────────────────────────────┘
```

**Key design decisions:**

- **Hugging Face compatibility by subclassing.** `MiniMindForCausalLM` extends `PreTrainedModel` and `GenerationMixin`, and `MiniMindConfig` extends `PretrainedConfig` with `model_type = "minimind"`. This means trained models export cleanly to the Transformers format and downstream to llama.cpp, vLLM, ollama, etc. — while the internals remain hand-written and readable.
- **Two weight formats.** Native training produces single-file `*.pth` state dicts (loaded by `MiniMindForCausalLM`); these convert to standard Transformers directories via `scripts/convert_model.py`. Inference paths branch on whether `--load_from` contains the substring `model` (native torch) or points at a Transformers directory.
- **Configuration via argparse, not files.** There are no YAML/JSON config files. Every hyperparameter is an `argparse` argument with a default. AST extraction found exactly one environment-variable read (`RANK`, for distributed init) and 246 CLI arguments across the scripts.

---

## 4. Model Architecture Deep Dive

The entire model is defined in `model/model_minimind.py`. It is a Pre-Norm, decoder-only transformer aligned with the Qwen3 / Qwen3-MoE ecosystem.

### 4.1 Configuration (`MiniMindConfig`)

`MiniMindConfig(PretrainedConfig)` carries all structural hyperparameters. Defaults for the mainline `minimind-3` (64M) model:

| Field | Default | Meaning |
|-------|---------|---------|
| `hidden_size` | `768` | Model dimension (`d_model`). |
| `num_hidden_layers` | `8` | Transformer block count. |
| `num_attention_heads` | `8` | Query heads. |
| `num_key_value_heads` | `4` | Key/value heads (GQA; `n_rep = 8/4 = 2`). |
| `head_dim` | `hidden_size // num_attention_heads` | Per-head dimension. |
| `vocab_size` | `6400` | Tokenizer vocabulary size. |
| `intermediate_size` | `ceil(hidden_size * π / 64) * 64` | SwiGLU inner dimension. |
| `hidden_act` | `silu` | Activation (SwiGLU uses SiLU). |
| `max_position_embeddings` | `32768` | Maximum positions. |
| `rope_theta` | `1e6` | RoPE base frequency. |
| `rms_norm_eps` | `1e-6` | RMSNorm epsilon. |
| `tie_word_embeddings` | `True` | Share embedding and LM-head weights. |
| `flash_attn` | `True` | Use `scaled_dot_product_attention` when available. |
| `dropout` | `0.0` | Dropout probability. |
| `inference_rope_scaling` | `False` | Enable YaRN extrapolation (sets a `rope_scaling` dict). |

MoE-specific fields (ignored when `use_moe=False`):

| Field | Default | Meaning |
|-------|---------|---------|
| `num_experts` | `4` | Routed experts. |
| `num_experts_per_tok` | `1` | Top-k routing (top-1 by default). |
| `moe_intermediate_size` | `intermediate_size` | Per-expert inner dimension. |
| `norm_topk_prob` | `True` | Normalize the top-k routing weights. |
| `router_aux_loss_coef` | `5e-4` | Load-balancing auxiliary-loss coefficient. |

### 4.2 Component layers

- **`RMSNorm`** — Root-mean-square layer normalization, computed in float32 then cast back to the input dtype for stability.
- **`precompute_freqs_cis` / `apply_rotary_pos_emb`** — RoPE rotary position embeddings. When `rope_scaling` is set, a YaRN ramp rescales the per-dimension frequencies to extrapolate beyond the original training length (`original_max_position_embeddings = 2048`, `factor = 16`).
- **`Attention`** — Grouped-query attention. Q/K/V projections are bias-free; `repeat_kv` expands KV heads to match query heads. Notably, MiniMind applies **per-head QK RMSNorm** (`q_norm` / `k_norm` on `head_dim`) before applying RoPE — a Qwen3-style stabilization. It uses flash attention (`F.scaled_dot_product_attention`) for the prefill path and falls back to a manual masked-softmax path for incremental decode or when an attention mask is present. A KV cache (`past_key_value`) supports autoregressive generation.
- **`FeedForward`** — SwiGLU: `down_proj(silu(gate_proj(x)) * up_proj(x))`.
- **`MOEFeedForward`** — Top-k expert routing. A linear `gate` produces softmax scores; `torch.topk` selects experts; outputs are accumulated with `index_add_`. During training it computes an auxiliary load-balancing loss (`aux_loss`) and includes an unused-expert gradient trick so all experts receive gradient flow.
- **`MiniMindBlock`** — Pre-Norm residual block: `x += attn(input_layernorm(x)); x += mlp(post_attention_layernorm(x))`. The MLP is `FeedForward` or `MOEFeedForward` depending on `use_moe`.
- **`MiniMindModel`** — Embedding + N blocks + final RMSNorm. RoPE buffers are registered non-persistently and recomputed if lost during meta-device initialization (a `transformers>=5.x` compatibility guard). Aggregates per-layer MoE `aux_loss`.
- **`MiniMindForCausalLM`** — Ties the LM head to the token embeddings when `tie_word_embeddings` is set. `forward()` returns a `MoeCausalLMOutputWithPast` with `loss`, `aux_loss`, `logits`, and `past_key_values`. Cross-entropy loss ignores index `-100` (the label-mask sentinel).

### 4.3 Custom `generate()`

`MiniMindForCausalLM.generate()` is a hand-written autoregressive loop (decorated `@torch.inference_mode()`), not the inherited Transformers implementation. It supports temperature scaling, repetition penalty, top-k and top-p (nucleus) filtering, KV caching, streaming via a `streamer`, multiple return sequences, and per-sequence EOS finishing. This is intentional — it keeps the generation logic transparent for learners while remaining `GenerationMixin`-compatible.

### 4.4 Published model variants

| Model | Params | `d_model` | `n_layers` | `kv_heads` | Notes |
|-------|--------|-----------|------------|------------|-------|
| minimind-3 | 64M | 768 | 8 | 4 | Dense, current mainline |
| minimind-3-moe | 198M-A64M | 768 | 8 | 4 | 4 experts / top-1 |
| minimind2-small | 26M | 512 | 8 | 2 | Legacy |
| minimind2-moe | 145M | 640 | 8 | 2 | Legacy |
| minimind2 | 104M | 768 | 16 | 2 | Legacy |

All variants share `vocab_size = 6400`, `max_position_embeddings = 32768`, and `rope_theta = 1e6`.

> [!NOTE]
> Per the project's design notes (echoing the MobileLLM study), for small models depth tends to matter more than width, but width below `d_model = 512` degrades sharply. The `768 × 8` mainline configuration is a deliberate balance of training speed, stability, and final quality.

---

## 5. The Tokenizer and Chat Template

MiniMind ships a custom 6400-token BPE tokenizer (`model/tokenizer.json`), built with a ByteLevel pre-tokenizer. The small vocabulary is a deliberate trade-off: it compresses the embedding and output-layer parameter share, which matters disproportionately for a tiny model.

**Special tokens** include chat-structure markers (`<|im_start|>`, `<|im_end|>`, `<|endoftext|>`), tool-call markers (`<tool_call>`, `</tool_call>`, `<tool_response>`, `</tool_response>`), reasoning markers (`<think>`, `</think>`), multimodal placeholders reserved for sibling projects, and a set of `<|buffer_N|>` tokens reserved for future expansion (36 special-token slots total).

**Chat template.** The `chat_template` (a Jinja2 string in `tokenizer_config.json`) is the heart of MiniMind's conversational behavior. It:

- Renders `system` / `user` / `assistant` / `tool` roles into `<|im_start|>...<|im_end|>` blocks.
- Emits a `# Tools` system block listing function signatures inside `<tools></tools>` when `tools` are provided.
- Expands `assistant.tool_calls` into `<tool_call>{...}</tool_call>` segments and `tool` messages into `<tool_response>...</tool_response>`.
- Implements the adaptive-thinking switch: when `add_generation_prompt` is set, it injects either an open `<think>\n` (if `open_thinking` is true) or an empty `<think>\n\n</think>\n\n` (default).

`train_tokenizer.py` documents how this vocabulary is produced (BPE over 10,000 sample lines) but is reference-only — retraining the tokenizer breaks weight, data, and ecosystem compatibility.

---

## 6. Data Layer

`dataset/lm_dataset.py` defines five dataset classes, each mapping a JSONL file to model-ready tensors. A shared preprocessing pair handles chat augmentation:

- `pre_processing_chat()` — with probability `add_system_ratio` (0.2), prepends a randomly chosen system prompt (English or Chinese) unless one is already present; tool-use samples are passed through untouched.
- `post_processing_chat()` — with probability `1 - empty_think_ratio` (0.8), strips empty `<think>\n\n</think>\n\n` blocks so the model sees a mix of "think" and "direct answer" patterns.

| Class | Source format | Output |
|-------|---------------|--------|
| `PretrainDataset` | `{"text": ...}` | `(input_ids, labels)`; BOS/EOS wrapped, padded; pad positions labeled `-100`. |
| `SFTDataset` | `{"conversations": [...]}` | `(input_ids, labels)`; labels masked to **assistant spans only** via `generate_labels()` (BOS-marker scanning). Applies the chat template with `tools`. |
| `DPODataset` | `{"chosen": [...], "rejected": [...]}` | dict of `x/y/mask` for chosen and rejected, with assistant-span loss masks. |
| `RLAIFDataset` | `{"conversations": [...]}` | `{prompt, answer:""}`; builds the prompt with a probabilistic `open_thinking` (`thinking_ratio`); the answer is filled by rollout. |
| `AgentRLDataset` | `{"conversations": [...], "gt": ...}` | `{messages, tools, gt}`; preserves ground-truth verification targets for trajectory rewards. |

The label-masking logic (`generate_labels`, `generate_loss_mask`) is the most subtle part of the data layer: it scans for the assistant-turn BOS marker (`<|im_start|>assistant\n`) and unmasks tokens until the EOS marker, so loss is computed only on assistant outputs — standard SFT practice, implemented from scratch here.

> [!IMPORTANT]
> `SFTDataset` uses an explicit Hugging Face `datasets` `Features` schema that includes `reasoning_content`, `tools`, and `tool_calls` fields. If your custom conversation data omits these keys, ensure they are at least present as empty/null so the schema validates.

---

## 7. Training Pipeline Internals

### 7.1 The common training scaffold

Every `trainer/train_*.py` script follows a near-identical nine-section structure, which is worth internalizing once:

1. **Distributed init** — `init_distributed_mode()` reads `RANK`/`LOCAL_RANK`; returns `0` for single-process runs.
2. **Config & checkpoint check** — build `MiniMindConfig`; if `--from_resume 1`, load the resume checkpoint.
3. **Mixed precision** — `torch.cuda.amp.autocast` with `bfloat16` (default) or `float16`; `nullcontext` on CPU.
4. **Experiment tracking** — if `--use_wandb`, `import swanlab as wandb` and init the run.
5. **Model / data / optimizer** — `init_model()` loads the base weights; an `AdamW` optimizer is created.
6. **Resume restore** — load model/optimizer/scaler/scheduler state and `start_epoch`/`start_step`.
7. **Compile & DDP** — optional `torch.compile`; wrap in `DistributedDataParallel` under distributed.
8. **Train loop** — per-epoch, with a `SkipBatchSampler` to skip already-completed steps on resume.
9. **Cleanup** — `dist.barrier()` and `destroy_process_group()`.

**Shared mechanics:**

- **Learning-rate schedule** — `get_lr()` implements a cosine schedule: `lr * (0.1 + 0.45 * (1 + cos(π * step / total)))`. The RL scripts instead use `CosineAnnealingLR`.
- **Gradient accumulation** — controlled by `--accumulation_steps`; the optimizer steps every N micro-batches, with a final flush for the trailing partial group.
- **Checkpoint saving** — at `--save_interval`, weights are saved as half-precision CPU tensors to `./out/<prefix>_<dim>[_moe].pth`, plus a full resume checkpoint via `lm_checkpoint()`.
- **Windows workaround** — every script does `import datasets` early (with a `# noqa` comment referencing issue #771) to avoid a pyarrow/torch DLL conflict on Windows.

### 7.2 Stage-specific loss functions

- **Pretrain / SFT** (`train_pretrain.py`, `train_full_sft.py`) — standard next-token cross-entropy (`res.loss`) plus the MoE `aux_loss`. SFT differs only in dataset (assistant-masked labels) and defaults.
- **LoRA** (`train_lora.py`) — same forward/loss as SFT, but only LoRA parameters are trainable; non-LoRA params are frozen, and `torch.compile` is force-disabled (incompatible with the monkey-patched LoRA forward).
- **Distillation** (`train_distillation.py`) — `loss = alpha * CE + (1 - alpha) * (T² · KL(teacher || student))`, with the teacher frozen in eval mode and its logits truncated to the student vocabulary.
- **DPO** (`train_dpo.py`) — concatenates chosen and rejected into one batch, computes per-token log-probs for policy and reference models, and applies `-logsigmoid(beta * ((π_chosen − π_rejected) − (ref_chosen − ref_rejected)))`.

---

## 8. Reinforcement Learning Subsystem

The RLAIF and agentic stages share a unified Policy-Optimization view: every algorithm optimizes a policy term `f(r_t)`, an advantage term `g(A_t)`, and a KL-regularization term `h(KL_t)`. The differences are entirely in how those three terms are instantiated.

| Algorithm | Policy term `f(r_t)` | Advantage `g(A_t)` | Regularizer | Trained networks |
|-----------|----------------------|--------------------|-------------|------------------|
| DPO | `log r_w − log r_l` | implicit (preference) | implicit in `β` | 1 (+1 frozen ref) |
| PPO | `min(r, clip(r))` | `R − V(s)` (Critic + GAE) | `β · E[KL]` | 2 (Actor + Critic) |
| GRPO | `min(r, clip(r))` | `(R − μ_group) / σ_group` | `β · KL_t` | 1 (+ frozen ref) |
| CISPO | `clip(r, 0, ε_max) · A_t · log π` | `(R − μ_group) / σ_group` | `β · KL_t` | 1 (+ frozen ref) |

### 8.1 Rollout engine abstraction

`trainer/rollout_engine.py` decouples policy optimization from trajectory generation. The `RolloutEngine` ABC defines `rollout()` and `update_policy()`; two implementations exist:

- **`TorchRolloutEngine`** — generates with the in-process model's `generate()`. Default.
- **`SGLangRolloutEngine`** — POSTs token IDs to an external SGLang server's `/generate` endpoint and syncs weights via `/update_weights_from_disk` (writing a shared checkpoint to disk). Adds `flush_cache()` and `health()` helpers. Selected with `--rollout_engine sglang`.

`compute_per_token_logps()` is the shared utility computing gathered log-probabilities for completion tokens. The factory `create_rollout_engine()` raises `ValueError: 不支持的引擎类型` (unsupported engine type) for unknown engine names, and `SGLangRolloutEngine.update_policy()` raises `RuntimeError: SGLang update_policy failed` on a weight-sync failure — the two error conditions surfaced by AST extraction.

### 8.2 GRPO / CISPO (`train_grpo.py`)

For each prompt, `num_generations` (default 6) completions are sampled; rewards are normalized within the group (`(r − μ)/σ`) to form advantages — eliminating the Critic network. The per-token loss switches on `--loss_type`: `cispo` (default) uses a clamped-ratio weighting that avoids hard gradient truncation, while `grpo` uses the standard min-of-clipped-ratio objective. A token-level KL penalty (`exp(kl) − kl − 1`) regularizes toward the reference model.

### 8.3 PPO (`train_ppo.py`)

A full Actor-Critic implementation. `CriticModel` subclasses `MiniMindForCausalLM`, replacing the LM head with a scalar `value_head`. The script computes Generalized Advantage Estimation (GAE) with `--gamma` and `--lam`, clips both the policy ratio (`--clip_epsilon`) and the value function (`--cliprange_value`), supports multiple PPO update iterations per rollout (`--ppo_update_iters`) with KL-based early stopping (`--early_stop_kl`), and carefully keeps the DDP forward/backward closed even when early-stopping to avoid deadlocks.

### 8.4 Agentic RL (`train_agent.py`)

The most sophisticated stage. It runs **multi-turn** rollouts: at each turn the model may emit a `<tool_call>`, which is parsed and executed against a set of six mock tools (`calculate_math`, `unit_converter`, `get_current_weather`, `get_current_time`, `get_exchange_rate`, `translate_text`), with the observation appended to the context for the next turn. The trajectory reward combines answer correctness against ground truth (`gt`), tool-call validity, format/tag closure, a reward-model score, and an unfinished-trajectory penalty:

```text
R(τ) = R_answer + R_tool + R_format + R_rm − R_unfinished
```

Tool execution is sandboxed with a `signal.alarm(1)` timeout and a restricted `eval` namespace (`{"__builtins__": {}, "math": math}`) for the math tool. Rewards are clipped to `[-3.0, 3.0]`.

### 8.5 Reward model

`LMForRewardModel` (in `trainer_utils.py`) wraps an external reward model — the project default is `InternLM2-1.8B-Reward`, placed as a sibling directory. `get_score()` formats the conversation, queries the reward model, and clamps the score to `[-3.0, 3.0]`. A continuous (not binary) reward is used deliberately to combat reward sparsity on a ~0.1B model.

---

## 9. Inference and Serving

### 9.1 CLI (`eval_llm.py`)

`init_model()` branches on `--load_from`: if the path contains `model`, it instantiates `MiniMindForCausalLM` and loads native `.pth` weights from `./out/`, optionally applying a LoRA adapter; otherwise it loads a Transformers model with `trust_remote_code=True`. The model is cast to half precision and `.eval()`. The script offers automatic (built-in prompts) or manual interactive modes and streams output via `TextStreamer`.

### 9.2 OpenAI-compatible API (`scripts/serve_openai_api.py`)

A FastAPI app exposing `POST /v1/chat/completions` on port `8998`. Notable internals:

- **`ChatRequest`** — a Pydantic model accepting `messages`, `temperature`, `top_p`, `max_tokens`, `tools`, `open_thinking`, and `chat_template_kwargs`. `get_open_thinking()` reconciles the `open_thinking` flag with `chat_template_kwargs` (also accepting an `enable_thinking` alias).
- **Streaming** — a `CustomStreamer` pushes tokens onto a `Queue` consumed by a generator that emits SSE chunks. It tracks the `</think>` boundary to split `reasoning_content` from `content` in the stream.
- **`parse_response()`** — for non-streaming responses, extracts `<think>...</think>` into `reasoning_content` and `<tool_call>...</tool_call>` into a structured `tool_calls` array (OpenAI function-call shape).

| Endpoint | Method | Auth | Purpose |
|----------|--------|------|---------|
| `/v1/chat/completions` | POST | None | Chat completion (streaming or single response). |

> [!WARNING]
> The API server binds `0.0.0.0:8998` with no authentication, no rate limiting, and no TLS. It is a development/demo server. Do not expose it directly to the public internet — see [Security Analysis](#12-security-analysis).

### 9.3 Format conversion (`scripts/convert_model.py`)

Converts between native torch `.pth` and Transformers directory formats (`torch2transformers` and back) and merges LoRA weights into a base model (`convert_merge_base_lora`, built on `model_lora.merge_lora`). Published mainline models are distributed in Transformers format.

---

## 10. Configuration Reference

All configuration is via `argparse`. The table below consolidates the arguments shared by the training scripts, followed by stage-specific parameters. There are no configuration files to manage.

### 10.1 Common training arguments

| Argument | Type | Typical default | Description |
|----------|------|-----------------|-------------|
| `--save_dir` | str | `../out` | Output weight directory. |
| `--save_weight` / `--lora_name` | str | per stage | Output weight prefix. |
| `--epochs` | int | per stage | Training epochs. |
| `--batch_size` | int | per stage | Batch size. |
| `--learning_rate` | float | per stage | Initial LR. |
| `--device` | str | `cuda:0`/`cpu` | Training device. |
| `--dtype` | str | `bfloat16` | `bfloat16` or `float16`. |
| `--num_workers` | int | `8` | DataLoader workers. |
| `--accumulation_steps` | int | `1` (8 for pretrain) | Gradient accumulation. |
| `--grad_clip` | float | `1.0` | Gradient-norm clip. |
| `--log_interval` | int | per stage | Steps between log lines. |
| `--save_interval` | int | per stage | Steps between saves. |
| `--hidden_size` | int | `768` | Model dimension. |
| `--num_hidden_layers` | int | `8` | Layer count. |
| `--max_seq_len` | int | per stage | Max sequence length (tokens). |
| `--use_moe` | int (0/1) | `0` | Train MoE model. |
| `--from_weight` | str | per stage | Base weight (`none` = scratch). |
| `--from_resume` | int (0/1) | `0` | Resume from checkpoint. |
| `--use_wandb` | flag | off | Enable WandB/SwanLab. |
| `--wandb_project` | str | per stage | Project name. |
| `--use_compile` | int (0/1) | `0` | `torch.compile` (disabled for LoRA). |

### 10.2 Stage-specific defaults

| Stage | `epochs` | `batch_size` | `learning_rate` | `max_seq_len` | `from_weight` | Distinctive args |
|-------|----------|--------------|-----------------|---------------|---------------|------------------|
| Pretrain | 2 | 32 | 5e-4 | 340 | `none` | `accumulation_steps 8` |
| Full SFT | 2 | 16 | 1e-5 | 768 | `pretrain` | — |
| LoRA | 10 | 32 | 1e-4 | 340 | `full_sft` | `--lora_name`, rank 16 |
| DPO | 1 | 4 | 4e-8 | 1024 | `full_sft` | `--beta 0.15` |
| Distillation | 6 | 32 | 5e-6 | 340 | `full_sft` | `--alpha 0.5`, `--temperature 1.5`, teacher/student size + MoE flags |
| GRPO/CISPO | 1 | 2 | 3e-7 | 768 | `full_sft` | `--num_generations 6`, `--loss_type cispo`, `--beta 0.1`, `--epsilon 0.2`, `--epsilon_high 5.0`, `--thinking_ratio 0.9`, `--max_gen_len 1024` |
| PPO | 1 | 2 | 3e-7 | 768 | `full_sft` | `--critic_learning_rate 5e-7`, `--clip_epsilon 0.2`, `--vf_coef 0.5`, `--kl_coef 0.02`, `--gamma 1.0`, `--lam 0.95`, `--cliprange_value 0.2`, `--ppo_update_iters 2`, `--early_stop_kl 0.25`, `--mini_batch_size 2` |
| Agentic RL | 1 | 2 | 3e-7 | 1024 | `full_sft`* | `--num_generations 4`, `--loss_type cispo`, `--max_gen_len 768`, `--max_total_len 2500`, `--thinking_ratio 0.1`, `--rollout_engine torch` |

\* Agentic RL uses `--from_weight full_sft` to initialize the policy and reference models.

**RL rollout / reward arguments** (PPO, GRPO, Agent): `--reward_model_path` (default `../../internlm2-1_8b-reward`), `--rollout_engine {torch,sglang}`, `--sglang_base_url`, `--sglang_model_path`, `--sglang_shared_path`, `--debug_mode`, `--debug_interval`.

### 10.3 Environment variables

| Variable | Used by | Purpose |
|----------|---------|---------|
| `RANK`, `LOCAL_RANK`, `WORLD_SIZE` | `init_distributed_mode()` | Standard `torchrun` distributed coordination. |
| `TOKENIZERS_PARALLELISM` | `dataset/lm_dataset.py` | Set to `false` to silence tokenizer fork warnings. |
| `HF_ENDPOINT` | benchmarking (optional) | Hugging Face mirror endpoint. |

---

## 11. Checkpointing and Resume

`lm_checkpoint()` in `trainer_utils.py` is the central checkpoint manager. It writes two artifacts atomically (write-to-`.tmp` then `os.replace`):

- **`./out/<prefix>_<dim>[_moe].pth`** — half-precision CPU weights, the deployable model.
- **`./checkpoints/<prefix>_<dim>[_moe]_resume.pth`** — a full resume bundle: model, optimizer, epoch, step, world size, WandB run id, and any extra state passed via `**kwargs` (scaler, scheduler, critic model/optimizer/scheduler for PPO).

**Cross-GPU-count resume:** if the saved `world_size` differs from the current one, the stored `step` is rescaled (`step * saved_ws // current_ws`) so training resumes at the equivalent point. The `SkipBatchSampler` then skips the already-completed batches within the resumed epoch. WandB/SwanLab runs resume under the saved run id.

> [!TIP]
> Because checkpoints save half-precision weights and the optimizer state separately, you can safely deploy the `./out/*.pth` file while a long run continues — but do not delete `./checkpoints/*_resume.pth` if you intend to resume.

---

## 12. Security Analysis

MiniMind is a research/training codebase, not a hardened service. The following are the security-relevant characteristics an operator should understand before exposing any part of it.

### 12.1 Network exposure

- **`serve_openai_api.py` binds `0.0.0.0:8998` with no authentication, rate limiting, or TLS.** Anyone who can reach the port can run inference on your GPU. Run it only on trusted networks (localhost/LAN) or behind an authenticating reverse proxy. The same applies to externally launched vLLM/SGLang servers.

### 12.2 Dynamic code execution in the agent tool sandbox

- The agentic RL math tool evaluates expressions with Python `eval`. It is **partially sandboxed**: the namespace is restricted to `{"__builtins__": {}, "math": math}` and wrapped in a 1-second `signal.alarm` timeout. This blocks most builtin-based escapes, but `eval`-on-untrusted-input is inherently risky. The inputs here are model-generated tool arguments during training (not external user input), which limits exposure, but treat any expansion of this mechanism to user-facing contexts with caution.

> [!CAUTION]
> If you adapt the agent tool-execution loop for a user-facing product, replace the `eval`-based math tool with a safe expression evaluator. Do not pass untrusted user input into `eval`, even with a restricted namespace.

### 12.3 Remote code trust

- Inference paths load Transformers models with `trust_remote_code=True`. This executes arbitrary model-repo code on load — only load model directories you trust.
- `signal.alarm`/`SIGALRM` is Unix-only; the tool-timeout sandbox does not function on Windows.

### 12.4 Data and weight integrity

- There is no authentication or signing on dataset or weight files. Training reads whatever JSONL/`.pth` files are present on disk. Validate the provenance of downloaded datasets and checkpoints.
- No secrets are required or stored by the core codebase. WandB/SwanLab use their own credential mechanisms; the OpenAI test client (`chat_api.py`) uses a placeholder key (`sk-123`) against a local server.

### 12.5 Overall assessment

For its intended use — local/single-machine training and experimentation — MiniMind's security posture is appropriate. The two items that require active attention before any production or shared-network use are (1) the unauthenticated API server and (2) the `eval`-based tool sandbox if extended to user input.

---

## 13. Code Quality and Patterns

**Strengths:**

- **Readability over abstraction.** Algorithms are implemented inline and from scratch, making the codebase a strong reference for learning. The unified PO-loss framing across DPO/PPO/GRPO/CISPO is reflected directly in the code structure.
- **Consistent scaffolding.** The nine-section training-script template means once you understand one trainer, you understand them all. Argument names and defaults are consistent across stages.
- **Ecosystem compatibility.** Subclassing `PreTrainedModel`/`GenerationMixin`/`PretrainedConfig` yields clean export to the broader inference ecosystem despite the from-scratch internals.
- **Robust resume.** Atomic checkpoint writes, cross-GPU-count step rescaling, and batch skipping make long runs resilient to interruption.

**Trade-offs and observations:**

- **Significant code duplication across trainers.** Each `train_*.py` repeats the full nine-section scaffold, the `rep_penalty`/`calculate_rewards` helpers (duplicated in `train_ppo.py`, `train_grpo.py`, and `train_agent.py`), and the optimizer/DDP setup. This is a deliberate pedagogical choice (each script is self-contained and readable end-to-end) but increases maintenance surface — a bug fixed in one reward function must be fixed in three.
- **Configuration is entirely argparse.** There is no config-file layer, so reproducing a specific run means recording a long command line. A YAML/config option would aid reproducibility.
- **Bilingual codebase.** Argument help text and comments are primarily Chinese; error messages mix Chinese and English (e.g. `不支持的引擎类型`). Non-Chinese-reading maintainers should keep a translation handy.
- **Mixed-precision and AMP** use the older `torch.cuda.amp.autocast`/`GradScaler` API, which newer PyTorch versions deprecate in favor of `torch.amp`.

---

## 14. Testing and Reproducibility

- **No automated test suite.** There is no `tests/` directory or CI; correctness is validated empirically through training-curve inspection (`images/*_loss.jpg`) and the eval scripts (`eval_llm.py`, `eval_toolcall.py`).
- **Determinism.** `setup_seed()` seeds Python, NumPy, and PyTorch (CPU + CUDA), and sets `cudnn.deterministic = True` / `benchmark = False`. Each epoch reseeds as `42 + epoch`, and distributed ranks offset by rank — giving reproducible-as-possible runs.
- **Benchmarking.** The project reports results on C-Eval, CMMLU, ARC-Easy, PIQA, OpenBookQA, HellaSwag, and Social-IQa via the external `lm-evaluation-harness`. Scores for these tiny models hover near the random-choice floor, as expected; benchmarks are framed as "for reference/amusement," not competitive claims.
- **Reproduction targets.** The Zero recipe (`pretrain_t2t_mini` + `sft_t2t_mini`, 1 epoch each) is the canonical fast-reproduction path: ~2.3 hours on a single RTX 3090.

---

## 15. Potential Improvements

Derived from the codebase as-is; these are opportunities, not defects.

- [ ] **Factor out shared RL helpers.** `rep_penalty()` and the reward-shaping logic are duplicated across `train_ppo.py`, `train_grpo.py`, and `train_agent.py`. Extracting them into `trainer_utils.py` would remove the three-way maintenance hazard while keeping per-stage scripts thin.
- [ ] **Add an optional config-file layer.** A thin YAML loader feeding `argparse` defaults would make runs reproducible by file rather than by long command lines.
- [ ] **Authentication for the API server.** Add an API-key check (and optional TLS guidance) to `serve_openai_api.py`, or document a reverse-proxy pattern, before any shared deployment.
- [ ] **Harden the agent tool sandbox.** Replace `eval` in the math tool with a dedicated safe-expression evaluator, and provide a cross-platform timeout (the current `signal.alarm` is Unix-only).
- [ ] **Migrate AMP APIs.** Move from `torch.cuda.amp` to `torch.amp` to track current PyTorch.
- [ ] **Smoke/unit tests at the data seams.** The label-masking logic (`generate_labels`, `generate_loss_mask`) and the chat-template expansion are subtle and high-impact; a small unit-test suite over a few fixtures would catch regressions cheaply.
- [ ] **Combined reasoning + tool-call data.** The project notes the model cannot reliably emit reasoning during tool calls because training data lacks joint samples; adding such data is a clear capability gap to close.

---

## 16. Quick Reference

**Train from scratch (Zero recipe):**

```bash
cd trainer
python train_pretrain.py          # out/pretrain_768.pth
python train_full_sft.py          # out/full_sft_768.pth
```

**Optional refinement stages:**

```bash
python train_lora.py  --lora_name lora_medical   # out/lora_medical_768.pth
python train_dpo.py                               # out/dpo_768.pth
python train_grpo.py                              # out/grpo_768.pth   (CISPO by default)
python train_ppo.py                               # out/ppo_actor_768.pth
python train_agent.py                             # out/agent_768.pth
```

**Multi-GPU:** prefix any script with `torchrun --nproc_per_node N`.
**Resume:** add `--from_resume 1`. **MoE:** add `--use_moe 1` everywhere. **Logging:** add `--use_wandb`.

**Inference & serving:**

```bash
python eval_llm.py --weight full_sft                       # CLI chat
python eval_toolcall.py --weight agent                     # tool-call test
cd scripts && python serve_openai_api.py                   # API server :8998
cd scripts && streamlit run web_demo.py                    # web UI
cd scripts && python convert_model.py                      # format conversion / LoRA merge
```

**Key paths:** weights → `./out/`; resume checkpoints → `./checkpoints/`; datasets → `./dataset/`; model code → `model/model_minimind.py`; tokenizer → `model/tokenizer.json`.

**Weight prefixes:** `pretrain`, `full_sft`, `dpo`, `grpo`, `ppo_actor`, `agent`, `full_dist`, `lora_*` (append `_moe` for MoE).

---

_Generated June 2026 from jingyaogong/minimind at 3f1a7cc. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._