docs: restructure to docs/, add guides and developer docs

- Rename assets/ to docs/, split into guides/ and developer/ - Add get-started.md: installation + 5-step quickstart - Add guides/evaluation.md: 7 eval scripts with CLI args - Add guides/distributed.md: DDP/FSDP, gradient accumulation, NCCL - Add developer/internals.md: loss formulas, RoPE, KV cache math - Add developer/cuda_kernels.md: build system, benchmarks, file layout - Fix storage_format doc in preprocessing.md - Update cross-references in README.md, README-zh-CN.md, Dockerfile
2026-07-30 00:49:04 +08:00
parent bcaa2d1ae0
commit c8b1e40f71
17 changed files with 1124 additions and 25 deletions
@@ -0,0 +1,255 @@
+# Distributed Training
+
+AstrAI supports three parallel modes: **single GPU** (`none`), **Data Parallel** (`ddp`), and **Fully Sharded Data Parallel** (`fsdp`). This guide covers when to use each, how to launch multi-GPU training, and how gradient accumulation works.
+
+## Quick Start
+
+### Single GPU
+
+```bash
+python scripts/tools/train.py \
+    --train_type=sft \
+    --param_path ./params \
+    --data_root_path ./dataset \
+    --parallel_mode=none \
+    --nprocs=1 \
+    --batch_per_device=4 \
+    --grad_accum_steps=8
+```
+
+### Multi-GPU DDP (4 GPUs)
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export NCCL_P2P_DISABLE=1
+export NCCL_NET_GDR_LEVEL=0
+
+python scripts/tools/train.py \
+    --train_type=sft \
+    --param_path ./params \
+    --data_root_path ./dataset \
+    --parallel_mode=ddp \
+    --nprocs=4 \
+    --batch_per_device=4 \
+    --grad_accum_steps=8
+```
+
+### Multi-GPU FSDP (4 GPUs)
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export NCCL_P2P_DISABLE=1
+export NCCL_NET_GDR_LEVEL=0
+
+python scripts/tools/train.py \
+    --train_type=sft \
+    --param_path ./params \
+    --data_root_path ./dataset \
+    --parallel_mode=fsdp \
+    --nprocs=4 \
+    --batch_per_device=4 \
+    --grad_accum_steps=8
+```
+
+> `--parallel_mode` defaults to `fsdp`. You can omit it for FSDP.
+
+## Parallel Modes
+
+| Mode | `--parallel_mode` | Param Layout | Memory | When to Use |
+|------|-------------------|--------------|--------|-------------|
+| Single GPU | `none` | Full, replicated | Highest | Small models, DPO/GRPO, debugging |
+| DDP | `ddp` | Full, replicated | High | Most multi-GPU training |
+| FSDP | `fsdp` | Sharded (DTensor) | Lowest | Large models that don't fit in single GPU |
+
+### NoneExecutor
+
+No wrapping. The model runs as-is on a single device. Gradient accumulation still works via `AccumOptimizer`/`AccumScheduler` (they gate `step()` on the sync counter). Checkpoint saving is a plain `state_dict()` call.
+
+### DDPExecutor
+
+Wraps the model with `torch.nn.parallel.DistributedDataParallel`. Each rank has a full copy of the model; gradients are all-reduced across ranks. Uses `gradient_as_bucket_view=True` and `broadcast_buffers=False` by default (hardcoded in `train.py`).
+
+During gradient accumulation, non-sync micro-steps use `model.no_sync()` to skip gradient all-reduce. Only the final micro-step triggers the all-reduce.
+
+### FSDPExecutor (FSDP2 / `fully_shard`)
+
+Uses PyTorch's FSDP2 per-module API (`torch.distributed.fsdp.fully_shard`). Each model child (e.g., each `DecoderBlock`) is individually sharded — parameters become `DTensor`s distributed across ranks. No `FlatParameter`, original parameter names are preserved.
+
+Key differences from DDP:
+- **Lower memory**: parameters are sharded, not replicated.
+- **Custom grad norm**: FSDP gradients are `DTensor`s, so `clip_grad_norm` computes the local norm, then all-reduces to get the global norm.
+- **Collective checkpoint ops**: `unshard()` and `full_tensor()` are collective — all ranks must call them even though only rank-0 saves. The executor handles this via `dist.barrier()` in `checkpoint_context`.
+- **Root skipped**: `fully_shard` is applied to direct children only (not the root model) due to an `ABC + Generic[T]` MRO incompatibility.
+
+## Gradient Accumulation
+
+Gradient accumulation lets you simulate a larger effective batch size by accumulating gradients over multiple micro-batches before calling `optimizer.step()`.
+
+```
+Effective batch = nprocs × batch_per_device × grad_accum_steps
+```
+
+Example: 4 GPUs × batch 4 × accum 8 = effective batch 256.
+
+### How it works
+
+Three cooperating layers:
+
+1. **`GradientState`** — tracks the micro-step counter. Fires `sync_gradients=True` every `grad_accum_steps` micro-batches.
+2. **`executor._no_sync(model)`** — suppresses gradient synchronization on non-sync micro-steps:
+   - `none`: `nullcontext` (nothing to skip)
+   - `ddp`: `model.no_sync()` (skips all-reduce)
+   - `fsdp`: `set_requires_gradient_sync(False)` on each `FSDPModule`
+3. **`AccumOptimizer` / `AccumScheduler`** — gate `step()` and `zero_grad()` on `sync_gradients`, so the optimizer only fires on the last micro-step.
+
+The loss is divided by `grad_accum_steps` before `backward()`, so gradients sum to the correct mean.
+
+## Process Launching
+
+AstrAI auto-detects the launch method:
+
+| Detection | Strategy | Use Case |
+|-----------|----------|----------|
+| `torchelastic` / `torchrun` env vars | `TorchrunStrategy` | External orchestrator (torchrun, SLURM, K8s) |
+| `RANK` + `WORLD_SIZE` env vars | `TorchrunStrategy` | External launch |
+| Neither | `LocalStrategy` | `python scripts/tools/train.py` (in-process spawn) |
+
+### Local (default)
+
+When you run `python scripts/tools/train.py --nprocs=4`, AstrAI uses `torch.multiprocessing.start_processes` to spawn 4 child processes. The parent process manages signal forwarding (SIGTERM/SIGINT) and waits for all children to finish.
+
+### Torchrun
+
+For multi-node or SLURM environments:
+
+```bash
+torchrun --nproc_per_node=4 scripts/tools/train.py \
+    --train_type=sft \
+    --parallel_mode=ddp \
+    --param_path ./params \
+    --data_root_path ./dataset \
+    --batch_per_device=4
+```
+
+When launched via torchrun, AstrAI reads `RANK`, `WORLD_SIZE`, `LOCAL_RANK` from the environment and uses `TorchrunStrategy`. The `--nprocs` flag is ignored (the orchestrator controls process count).
+
+## NCCL Environment Variables
+
+For multi-GPU training, you **must** set these environment variables:
+
+```bash
+export NCCL_P2P_DISABLE=1
+export NCCL_NET_GDR_LEVEL=0
+```
+
+These are required on certain hardware configurations (see `AGENTS.md`). Without them, NCCL may hang or crash during collective operations. These are set in the training shell scripts (`train-seq.sh`, `train-sft.sh`, `train-dpo.sh`) but not in Python code — you must export them before launching.
+
+## Checkpoint Saving
+
+Checkpoints are saved by **rank-0 only**. The flow:
+
+1. `executor.checkpoint_context(model)` — wraps with `dist.barrier()` before and after (distributed only).
+2. `executor.unwrap_model(model)` — gathers the full state dict:
+   - `none`: `model.state_dict()`
+   - `ddp`: `model.module.state_dict()`
+   - `fsdp`: `unshard()` → `full_tensor()` → `reshard()` (collective on all ranks, result kept only on rank-0)
+3. Non-rank-0 ranks get `None` — the save is skipped.
+4. Rank-0 writes `meta.json`, `config.json`, `model.safetensors`, and optional `{key}.pt` (optimizer/scheduler state).
+
+> **FSDP note**: Even though only rank-0 saves, all ranks must participate in `unwrap_model` because `unshard()` and `full_tensor()` are collective operations. The barriers in `checkpoint_context` keep all ranks in lockstep.
+
+## Total Steps Calculation
+
+The scheduler's total step count accounts for data-parallel sharding:
+
+```
+samples_per_replica = ceil(dataset_len / nprocs)
+batches_per_replica  = ceil(samples_per_replica / batch_per_device)
+total_steps          = (batches_per_replica // grad_accum_steps) * n_epoch
+```
+
+This ensures the LR schedule is correctly scaled regardless of the number of GPUs.
+
+## Real Examples
+
+### Pretraining (seq, DDP, 4 GPUs)
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export NCCL_P2P_DISABLE=1
+export NCCL_NET_GDR_LEVEL=0
+
+python scripts/tools/train.py \
+    --train_type=seq \
+    --param_path ./params \
+    --data_root_path ./dataset/cached \
+    --parallel_mode=ddp \
+    --nprocs=4 \
+    --n_epoch=1 \
+    --max_lr=2e-4 \
+    --schedule_type=wsd \
+    --warmup_ratio=0.02 \
+    --window_size=2048 \
+    --batch_per_device=4 \
+    --grad_accum_steps=32 \
+    --ckpt_interval=2000
+# Effective batch = 4 × 4 × 32 = 512
+```
+
+### SFT (DDP, 4 GPUs)
+
+```bash
+python scripts/tools/train.py \
+    --train_type=sft \
+    --param_path ./AstrAI-V1-base \
+    --data_root_path ./dataset/cached_sft \
+    --parallel_mode=ddp \
+    --nprocs=4 \
+    --n_epoch=2 \
+    --max_lr=2e-5 \
+    --schedule_type=cosine \
+    --warmup_ratio=0.02 \
+    --min_rate=0.05 \
+    --window_size=2048 \
+    --batch_per_device=4 \
+    --grad_accum_steps=8
+# Effective batch = 4 × 4 × 8 = 128
+```
+
+### DPO (Single GPU)
+
+```bash
+python scripts/tools/train.py \
+    --train_type=dpo \
+    --param_path ./checkpoint/epoch_1_step_6000 \
+    --data_root_path ./alpaca_dpo.jsonl \
+    --parallel_mode=none \
+    --nprocs=1 \
+    --max_lr=5e-6 \
+    --schedule_type=cosine \
+    --warmup_ratio=0.1 \
+    --min_rate=0.1 \
+    --window_size=1024 \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --dpo_beta=0.1 \
+    --max_grad_norm=50
+```
+
+## CLI Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--nprocs` | 1 | Number of GPUs / processes |
+| `--parallel_mode` | `fsdp` | `none`, `ddp`, or `fsdp` |
+| `--start_method` | `spawn` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) |
+| `--backend` | `nccl` | Distributed backend (`nccl`, `gloo`) |
+| `--master_addr` | `localhost` | Master node address |
+| `--master_port` | `29500` | Master node port |
+| `--device_type` | `cuda` | Device type |
+
+> `--tp_size` is parsed but **not yet wired** — tensor parallelism is future work. `ColumnParallelLinear` / `RowParallelLinear` exist in `astrai/parallel/module.py` but are not used by the model.
+
+Full parameter reference: [CLI Reference](params.md). Training loop and strategies: [Training Guide](training.md).
+
+> Document Update Time: 2026-07-30
@@ -0,0 +1,252 @@
+# Evaluation
+
+AstrAI provides 7 evaluation scripts in `scripts/eval/` covering code generation, knowledge QA, perplexity, summarization, data quality, instruction following, and weight analysis.
+
+## Overview
+
+| Script | Metric | Model Invocation | External Dataset |
+|--------|--------|-------------------|-------------------|
+| `evaluate_humaneval.py` | Code-gen pass@1/10/100 | `InferenceEngine.generate` | HF `openai/openai_humaneval` (auto-download) |
+| `evaluate_mmlu.py` | MCQ accuracy (log-likelihood) | Direct `model()` forward | HF `cais/mmlu` (auto-download) |
+| `evaluate_ppl.py` | Perplexity / token loss | Direct `model()` forward | User JSONL |
+| `evaluate_rouge.py` | ROUGE-1/2/L | None (pure metric) | User JSONL |
+| `evaluate_ifd.py` | Instruction-Following Difficulty | Direct `model()` forward | User JSONL |
+| `evaluate_ifeval.py` | Instruction-following constraints | `InferenceEngine.generate` | HF `google/IFEval` (auto-download) |
+| `analyze_weights.py` | SVD effective rank / weight stats | None (loads safetensors) | Checkpoint dir |
+
+Two invocation patterns exist:
+- **Generation benchmarks** (HumanEval, IFEval): use `InferenceEngine` to generate responses, then score them.
+- **Scoring benchmarks** (MMLU, PPL, IFD): call `model()` directly under `torch.inference_mode()` for log-likelihood computation.
+
+Common defaults: `--param_path` defaults to `./params`; dtype defaults to `bfloat16` on CUDA, `float32` on CPU.
+
+---
+
+## HumanEval (Code Generation)
+
+Generates completions for 164 programming problems, executes them against hidden tests, and reports pass@k.
+
+```bash
+python scripts/eval/evaluate_humaneval.py \
+    --param_path ./params \
+    --num_samples 20 \
+    --batch_size 32 \
+    --max_tokens 512 \
+    --output results/humaneval.json
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--param_path` | `./params` | Model directory |
+| `--data_path` | `./humaneval/HumanEval.jsonl` | HumanEval JSONL (auto-downloaded if missing) |
+| `--output` | None | Save results JSON (also writes `_completions.json`) |
+| `--test_only` | None | Test an existing completions JSON (skip generation) |
+| `--generate_only` | False | Only generate, skip execution/testing |
+| `--num_samples` | 200 | Completions per problem (pass@k needs >= k) |
+| `--max_tokens` | 512 | Max generation length |
+| `--temperature` | 0.8 | Sampling temperature |
+| `--top_p` | 0.95 | Nucleus sampling threshold |
+| `--top_k` | 50 | Top-k sampling |
+| `--batch_size` | 32 | Generation batch size |
+| `--test_workers` | 8 | ProcessPoolExecutor workers for test execution |
+| `--test_timeout` | 3.0 | Per-subprocess timeout (seconds) |
+| `--problems` | None | Restrict to specific problem indices |
+
+**Output**: stdout prints `pass@1`, `pass@10`, `pass@100`. With `--output`, writes per-problem results + `_summary` aggregate and a `_completions.json` file.
+
+**Data**: Auto-downloads `openai/openai_humaneval` from HuggingFace on first run. Each problem has `task_id`, `entry_point`, `prompt`, `test`.
+
+---
+
+## MMLU (Knowledge QA)
+
+57-subject multiple-choice accuracy via log-likelihood comparison. Supports n-shot few-shot prompting and option permutation.
+
+```bash
+python scripts/eval/evaluate_mmlu.py \
+    --param_path ./params \
+    --n_shot 5 \
+    --subjects math_algebra history_us \
+    --output results/mmlu.json
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--param_path` | `./params` | Model directory |
+| `--data_dir` | `./mmlu_data` | MMLU data directory (per-subject CSVs) |
+| `--download` | False | Force re-download |
+| `--n_shot` | 5 | Few-shot examples (0 = zero-shot) |
+| `--subjects` | all 57 | Specific subjects to evaluate |
+| `--output` | None | Output JSON path |
+| `--split` | `test` | `test` or `val` |
+| `--device` | auto | Device (`cuda` / `cpu`) |
+| `--dtype` | auto | `bfloat16` on CUDA, `float32` on CPU |
+| `--seed` | 0 | Seed for option permutation (0 = enabled, -1 = disabled) |
+
+**How it works**: For each question, builds a prompt with n-shot examples, then scores each choice (A/B/C/D) by computing the summed log-likelihood of the choice token given the context. The choice with the highest log-prob is the prediction.
+
+**Output**: stdout prints per-subject accuracy and overall. With `--output`, writes per-subject `{accuracy, correct, total}` + `_overall` aggregate.
+
+**Data**: Auto-downloads `cais/mmlu` from HuggingFace. Stored as per-subject CSVs in `<data_dir>/<split>/` and `<data_dir>/dev/` (for few-shot).
+
+---
+
+## Perplexity (PPL)
+
+Token-level negative-log-likelihood and perplexity on arbitrary text data. Supports streaming mode (memory-efficient) and non-streaming mode (exact per-token stats).
+
+```bash
+python scripts/eval/evaluate_ppl.py \
+    --param_path ./params \
+    --input_path data.jsonl \
+    --output_dir ppl_results/ \
+    --batch_size 4 \
+    --max_length 2048
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--param_path` | required | Model directory |
+| `--input_path` | required | Input file, glob, or directory |
+| `--output_dir` | required | Output directory for `summary.json` + token JSONL |
+| `--text_key` | `text` | Key for the text field in input data |
+| `--batch_size` | 4 | Batch size |
+| `--max_length` | 2048 | Max sequence length (tokens) |
+| `--token_level` | False | Store per-token log_probs + token-type analysis |
+| `--max_samples` | None | Random subsample per file |
+| `--device` | auto | Device |
+| `--dtype` | auto | Torch dtype |
+
+**Input**: JSONL or JSON files. Each item must have a field named by `--text_key` (default `text`). If `--input_path` is a directory, recursively collects `*.jsonl` and `*.json`.
+
+**Output**: `summary.json` with per-file stats (tokens, mean/median loss, perplexity, p50/p90/p95/p99). With `--token_level`, also writes per-token JSONL with token IDs and log-probs.
+
+---
+
+## ROUGE
+
+ROUGE-1/2/L (precision, recall, F1) for summarization. Self-contained implementation with no external dependencies.
+
+```bash
+python scripts/eval/evaluate_rouge.py \
+    --data_path predictions.jsonl \
+    --output results/rouge.json
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--data_path` | required | JSONL with `reference`/`candidate` per line |
+| `--output` | None | Output JSON path |
+
+**Input**: JSONL, one object per line:
+```json
+{"reference": "Ground truth text", "candidate": "Model output text"}
+```
+
+**Output**: stdout prints `rouge-1`, `rouge-2`, `rouge-l` each as P/R/F1. With `--output`, writes JSON with `aggregate` and `per_item` scores.
+
+Can also be imported as a library:
+```python
+from scripts.eval.evaluate_rouge import compute_rouge
+scores = compute_rouge(reference, candidate)
+```
+
+---
+
+## IFD (Instruction-Following Difficulty)
+
+Data quality metric: `IFD = L_conditional / L_unconditional`. Measures how much harder it is to predict a response given its instruction vs. without it. Useful for filtering instruction-tuning data.
+
+```bash
+python scripts/eval/evaluate_ifd.py \
+    --param_path ./params \
+    --input_path sft_data.jsonl \
+    --output_dir ifd_results/ \
+    --format messages \
+    --batch_size 8
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--param_path` | required | Model directory |
+| `--input_path` | required | Input file, glob, or directory |
+| `--output_dir` | required | Output directory |
+| `--max_len` | 2048 | Max token length |
+| `--format` | `plain` | `plain` (instruction/response fields) or `messages` (chat format) |
+| `--instr_key` | `instruction` | Instruction field key (plain format) |
+| `--resp_key` | `response` | Response field key (plain format) |
+| `--batch_size` | 8 | Items per model-forward flush |
+| `--device` | auto | Device |
+| `--dtype` | auto | Torch dtype |
+| `--sentinel_text` | `\n` | Prefix for unconditional pass (`""` → bos/pad fallback) |
+| `--per_token` | False | Include per-token IFD breakdown |
+| `--max_samples` | None | Random subsample per file |
+
+**How it works**: Two forward passes per batch — (1) conditional: packed BFD sequence with context + response, (2) unconditional: response prefixed with a sentinel. IFD = mean_conditional_loss / mean_unconditional_loss. IFD > 1 means the instruction makes the response harder to predict (higher quality data).
+
+**Output**: Per-file `<label>_ifd.jsonl` with IFD scores per item. `summary.json` aggregates per-file stats.
+
+---
+
+## IFEval (Instruction Following)
+
+Google's IFEval benchmark: generates responses and verifies 27 types of constraints (keywords, format, length, case, punctuation, etc.).
+
+```bash
+python scripts/eval/evaluate_ifeval.py \
+    --param_path ./params \
+    --num_samples 1 \
+    --max_tokens 512 \
+    --output results/ifeval.json
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--param_path` | `./params` | Model directory |
+| `--data_path` | `./ifeval/input_data.jsonl` | IFEval JSONL (auto-downloaded if missing) |
+| `--output` | None | Output JSON path |
+| `--max_tokens` | 512 | Max generation tokens |
+| `--temperature` | 0.1 | Sampling temperature (low for instruction-following) |
+| `--top_p` | 0.95 | Top-p sampling |
+| `--top_k` | 50 | Top-k sampling |
+| `--num_samples` | 1 | Samples per problem (best-of-n scoring) |
+| `--batch_size` | 1 | Inference batch size |
+| `--limit` | None | Limit to first N problems (quick testing) |
+| `--dump_responses` | None | Path to dump raw responses as JSONL |
+
+**Output**: stdout prints overall accuracy + per-constraint-type accuracy table. With `--output`, writes per-problem results + `_summary`.
+
+**Data**: Auto-downloads `google/IFEval` from HuggingFace. Each problem has `key`, `prompt`, `instruction_id_list`, `kwargs`.
+
+---
+
+## Weight Analysis
+
+SVD-based effective rank and weight statistics for checkpoint diagnostics. Does not load the model graph or run any forward pass.
+
+```bash
+python scripts/eval/analyze_weights.py \
+    --ckpt_dir ./checkpoint/epoch_1_step_6000 \
+    --output results/weights.json
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--ckpt_dir` | required | Checkpoint dir with `model.safetensors` + `config.json` |
+| `--compare` | None | Additional checkpoint dirs to compare |
+| `--no_svd` | False | Skip SVD; show only weight stats (faster) |
+| `--output` | None | Save results as JSON |
+| `--device` | `cuda` | Device for SVD |
+
+**Output**: SVD effective rank by component (ER@90/95/99%, entropic rank, condition number), per-layer effective rank grid, and weight value statistics (mean/std/min/max). Provides a utilization verdict (HIGH >0.85 / MODERATE >0.5 / LOW).
+
+---
+
+## Tips
+
+- **Quick test**: Use `--limit` (IFEval) or `--problems` (HumanEval) to run on a small subset first.
+- **Auto-download**: HumanEval, MMLU, and IFEval auto-download their datasets on first run. The other scripts expect user-provided data.
+- **Output formats**: `--output` writes a single JSON for most scripts. PPL and IFD write an `--output_dir` containing `summary.json` plus per-file artifacts.
+- **CPU mode**: All scripts auto-detect CUDA. To force CPU, use `--device cpu --dtype float32`.
+
+> Document Update Time: 2026-07-30
@@ -0,0 +1,252 @@
+# Inference
+
+## Contents
+
+- [KV Cache](#kv-cache)
+- [KVCache System](#kvcache-system)
+- [Continuous Batching](#continuous-batching)
+- [Sampling](#sampling-strategy-pattern)
+- [Protocol Handlers](#protocol-handlers-strategy-pattern)
+- [Engine & GenerateResult](#engine--generateresult)
+- [HTTP API](#http-api) — endpoints, SSE, errors, stats
+- [Engine API](#engine-api)
+
+## KV Cache
+
+At decode time, only the last query token matters. All previous K/V are cached to avoid recomputation:
+
+$$
+o_n = \sum_j \text{softmax}\left(\frac{q_n k_j}{\sqrt{d_k}}\right) v_j
+$$
+
+RoPE is applied **before** KV cache write, not after — otherwise position encoding drift occurs.
+
+## KVCache System
+
+Seven classes working together, with two concrete cache implementations:
+
+### ContiguousCache (default)
+
+```
+ContiguousCache (simple contiguous per-slot cache)
+  ├── ContiguousCacheView  bundles k/v tensors + slot indices for attention layers
+```
+
+Created by default when no cache is passed to `InferenceScheduler`. Each task occupies a fixed slot of `[max_seq_len, num_key_value_heads, head_dim]`. Simple and efficient for small-to-medium batch sizes.
+
+### PageCache (paged with prefix sharing)
+
+```
+PageCache (paged KV cache with prefix sharing, alternative)
+  ├── PagePool               orchestrates page allocation + prefix matching
+  │     ├── Allocator         bitmask-based page allocator + ref-count + LRU
+  │     └── PrefixCache       hash-based prefix matching (page_hash via polynomial hash)
+  ├── TaskTable              maps task_id → page_table + cached token count
+  ├── Storage                k_cache / v_cache tensors (num_hidden_layers × n_pages × page_size × num_key_value_heads × head_dim)
+  └── PageCacheView          bundles Storage + page_table + total_len for attention layers
+```
+
+`isinstance(cache, KVCache)` checks dispatch to the correct view. Both implement the abstract `KVCache` interface used by `Executor` and `InferenceScheduler`.
+
+## Continuous Batching
+
+`InferenceScheduler` runs a daemon thread with a 4-phase loop:
+
+```
+1. Cleanup → Remove finished tasks, free KV cache slots/pages
+2. Refill  → Pop from waiting_queue, task_alloc resources, activate
+3. Prefill → Group by (prompt_len, start_pos), run full forward
+4. Decode  → Run single-token forward for each same-position group
+```
+
+## Sampling (Strategy Pattern)
+
+```
+BaseSamplingStrategy (ABC)
+  ├── TemperatureStrategy
+  ├── TopKStrategy
+  ├── TopPStrategy
+  └── SamplingPipeline
+```
+
+`SamplingPipeline` composes them: Temperature → Top-K → Top-P → softmax → multinomial.  
+`sample()` is a convenience shortcut for one-shot usage.
+
+## Protocol Handlers (Strategy Pattern)
+
+```python
+class ProtocolHandler:  # concrete orchestrator
+    def __init__(self, request, engine, builder): ...
+    async def handle(self):
+        prompt, ctx, stops = builder.prepare(request, engine)
+        agen = engine.generate_async(prompt, ...)
+        if stream: self._handle_stream(agen, ctx, stops)
+        else:      return await self._handle_non_stream(agen, ctx, stops)
+```
+
+`ResponseBuilder` (ABC): `prepare()`, `format_stream_start()`, `format_chunk()`, `format_stream_end()`, `format_response()`.
+
+`OpenAIResponseBuilder` → `/v1/chat/completions`, `AnthropicResponseBuilder` → `/v1/messages`.
+
+Adding a protocol = one builder file, no handler subclassing needed.
+
+## Engine & GenerateResult
+
+```
+InferenceEngine
+  ├── generate(prompt, stream, ...) → str | List[str] | Generator
+  ├── generate_with_request(req)    → same
+  ├── generate_async(prompt, ...)   → AsyncGenerator
+  ├── get_stats()                   → Dict
+  └── shutdown()
+```
+
+`GenerateResult` uses `Condition` for non-streaming (`wait_completion()`) and `Event` for streaming (`wait()`). Stream callback is `cb(token)`.
+
+## HTTP API
+
+```
+POST /v1/chat/completions   OpenAI
+POST /v1/messages            Anthropic
+GET  /health                 {"status":"ok","model_loaded":true}
+GET  /stats                  scheduler statistics
+```
+
+### OpenAI
+
+```bash
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+Response:
+```json
+{
+  "id": "chatcmpl-abc123",
+  "object": "chat.completion",
+  "created": 1717000000,
+  "model": "astrai",
+  "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"}],
+  "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}
+}
+```
+
+Streaming SSE: `object: "chat.completion.chunk"` — starts with role delta, then token chunks, ends with finish chunk + usage stats, then `data: [DONE]`.
+
+### Anthropic
+
+```bash
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{"model":"astrai","system":"You are helpful.","messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+Supports `stop_sequences` and streaming via `event: content_block_delta`.
+
+### GenerationRequest Parameters
+
+| Param | Type | Default | Description |
+|-------|------|---------|-------------|
+| `messages` | List[dict] | required | Chat messages (role, content) |
+| `top_k` | int | 50 | Top-k count |
+| `top_p` | float | 1.0 | Nucleus threshold |
+| `temperature` | float | 1.0 | Sampling temperature (> 0.0) |
+| `max_tokens` | Optional[int] | None | Max generation length |
+| `stream` | bool | False | Stream output |
+
+### SSE Streaming Format
+
+**OpenAI** (`/v1/chat/completions`, `stream=true`):
+
+```
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai",
+       "choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":0,"model":"astrai",
+       "choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai",
+       "choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
+
+data: {"prompt_tokens":5,"completion_tokens":1,"total_tokens":6}
+
+data: [DONE]
+```
+
+**Anthropic** (`/v1/messages`, `stream=true`):
+
+```
+event: message_start
+data: {"type":"message_start","message":{"id":"msg_...","model":"astrai","role":"assistant",
+       "content":[],"usage":{"input_tokens":0}}}
+
+event: content_block_start
+data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
+
+event: content_block_stop
+data: {"type":"content_block_stop","index":0}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{...}}
+
+event: message_stop
+data: {"type":"message_stop"}
+```
+
+### Error Responses
+
+The server returns standard HTTP status codes. Pydantic validation errors (e.g. missing required fields)
+are handled automatically by FastAPI with 422 status. The only application-level error is engine initialization:
+
+| Status | Meaning |
+|--------|---------|
+| 200 | Success |
+| 422 | Unprocessable entity (Pydantic validation) |
+| 503 | Service unavailable (model not loaded, engine not ready) |
+
+Error response body (503):
+
+```json
+{
+    "detail": "Engine not initialized"
+}
+```
+
+### Stats Endpoint
+
+```
+GET /stats
+```
+
+Response:
+
+```json
+{
+    "total_tasks": 128,
+    "total_tokens": 10240,
+    "active_tasks": 3,
+    "waiting_queue": 2
+}
+```
+
+## Engine API
+
+```python
+# Non-streaming
+engine.generate("Hello", stream=False)          # -> str
+engine.generate(["A", "B"], stream=False)       # -> List[str]
+
+# Streaming
+engine.generate("Hello", stream=True)           # -> Generator[str]
+engine.generate(["A", "B"], stream=True)        # -> Generator[Tuple[int, str]]
+
+# Async
+async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[str]
+    print(token)
+```
+
+> Document Update Time: 2026-07-09
@@ -0,0 +1,230 @@
+# CLI Parameter Reference
+
+## Contents
+
+- [Training Parameters](#training-parameters)
+- [Inference Server](#inference-server-serverpy)
+- [Generate](#generate-generatepy)
+- [Preprocess](#preprocess-preprocesspy)
+
+## Training Parameters
+
+### Basic Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`) | required |
+| `--data_root_path` | Dataset root directory | required |
+| `--param_path` | Model parameters or checkpoint path | required |
+| `--n_epoch` | Total training epochs | 1 |
+| `--batch_per_device` | Batch size per device | 1 |
+| `--grad_accum_steps` | Gradient accumulation steps between optimizer steps | 1 |
+
+### Learning Rate Scheduling
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
+| `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
+| `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | 1.0 |
+
+### Optimizer (MuonMix)
+
+Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
+| `--muon_momentum` | Muon momentum factor | 0.95 |
+| `--muon_nesterov` | Enable Nesterov momentum for Muon | True |
+| `--muon_ns_steps` | Newton-Schulz iteration steps for Muon | 5 |
+| `--muon_adjust_lr` | Muon LR adjustment strategy (`original`, `match_rms_adamw`) | `match_rms_adamw` |
+
+### Data Loading
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--window_size` | Max input sequence length | model config `max_position_embeddings` |
+| `--stride` | Stride for sliding window over sequences | None |
+| `--random_seed` | Random seed for reproducibility | 3407 |
+| `--num_workers` | DataLoader worker processes | 4 |
+| `--no_pin_memory` | Disable pin_memory (enabled by default) | (flag) |
+
+### Checkpoint & Resume
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--ckpt_interval` | Iterations between checkpoints | 5000 |
+| `--ckpt_dir` | Checkpoint save directory | checkpoint |
+| `--start_epoch` | Resume from epoch (0 = from scratch) | 0 |
+| `--start_samples` | Resume from sample count per rank | 0 |
+
+### Validation
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--val_split` | Ratio to split from training dataset for validation (e.g. 0.05) | None |
+| `--val_step` | Number of optimizer steps between validation runs | 1000 |
+
+### Logging
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--log_dir` | Directory for metric logs | checkpoint/logs |
+| `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr", "grad_norm"] |
+
+### Gradient Checkpointing
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--gradient_checkpointing` | Enable activation checkpointing for DecoderBlock modules | False |
+
+### Distributed Training
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--nprocs` | Number of GPUs / processes | 1 |
+| `--parallel_mode` | Parallel strategy (`none`, `ddp`, `fsdp`) | fsdp |
+| `--device_type` | Device type | cuda |
+| `--start_method` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) | spawn |
+| `--backend` | Distributed training backend | nccl |
+| `--master_addr` | Master node address | localhost |
+| `--master_port` | Master node port | 29500 |
+
+### Strategy-specific
+
+| Parameter | Description | Default | Used by |
+|-----------|-------------|---------|---------|
+| `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` |
+| `--group_size` | GRPO group size | 4 | `grpo` |
+| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
+| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
+| `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` |
+
+### Online Rollout
+
+These options apply to `online_grpo` and `online_dpo`. Online strategies require
+a `BaseRewardModel` factory in `TrainConfig`; `train.py` does not currently
+provide a command-line option for configuring one.
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--rollout_interval` | Optimizer steps between rollout refreshes | 512 |
+| `--rollout_temperature` | Rollout sampling temperature | 0.7 |
+| `--rollout_top_k` | Rollout top-k filtering (`0` disables) | 0 |
+| `--rollout_top_p` | Rollout nucleus sampling threshold | 0.9 |
+| `--rollout_max_tokens` | Maximum generated tokens per response | 1024 |
+
+### Scheduler
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine |
+| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default: 0.05 for cosine/SGDR, 0.0 for WSD) |
+| `--cycle_length` | SGDR first cycle length in steps | None (total_steps - warmup_steps) |
+| `--t_mult` | SGDR cycle length multiplier per restart | 2 |
+| `--stable_steps` | WSD stable plateau steps | None (80% of post-warmup steps) |
+| `--decay_steps` | WSD decay steps | None (total_steps - warmup_steps - stable_steps) |
+
+### Usage Example
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --parallel_mode=ddp \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --weight_decay=0.1 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
+```
+
+---
+
+## Inference Server (`server.py`)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--host` | str | `0.0.0.0` | Host address |
+| `--port` | int | `8000` | Port number |
+| `--param_path` | path | `project_root/params` | Path to model parameters |
+| `--device` | str | `cuda` | Device to load model on |
+| `--dtype` | str | `bfloat16` | Model weights dtype (`bfloat16`, `float16`, `float32`) |
+| `--max_batch_size` | int | `16` | Maximum batch size for continuous batching |
+| `--max_seq_len` | int | model config `max_position_embeddings` | Maximum sequence length (KV cache size + prompt truncation) |
+| `--reload` | flag | `False` | Enable auto-reload for development |
+
+Usage:
+```bash
+python scripts/tools/server.py --param_path ./params --device cuda --dtype bfloat16
+```
+
+See [Inference Guide](inference.md) for HTTP API documentation.
+
+# Preprocess
+
+```bash
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c config.json
+```
+
+See [Preprocessing Guide](preprocessing.md) for config file format and examples.
+
+## Generate (`generate.py`)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--param_path` | str | required | Path to the model directory |
+| `--input_json_file` | str | required | Path to the input JSONL file |
+| `--output_json_file` | str | required | Path to the output JSONL file |
+| `--question_key` | str | `question` | Key for the question in input JSON |
+| `--response_key` | str | `response` | Key for the response in output JSON |
+| `--temperature` | float | `0.60` | Sampling temperature |
+| `--top_k` | int | `30` | Top-k filtering |
+| `--top_p` | float | `0.95` | Nucleus sampling threshold |
+| `--batch_size` | int | `1` | Batch size for generation |
+| `--num_samples` | int | `1` | Responses per prompt |
+| `--max_tokens` | int | model config `max_position_embeddings` | Maximum tokens to generate |
+| `--cache_len` | int | `2048` | KV cache length |
+| `--frequency_penalty` | float | `0.0` | Frequency penalty |
+| `--rep_window` | int | `64` | Window size for frequency penalty |
+
+Usage:
+```bash
+python scripts/tools/generate.py \
+    --param_path ./params \
+    --input_json_file input.jsonl \
+    --output_json_file output.jsonl
+```
+
+## Preprocess (`preprocess.py`)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) |
+| `--output_dir`, `-o` | path | required | Output directory for processed data |
+| `--config`, `-c` | path | required | Preprocessing pipeline config (JSON) |
+| `--tokenizer_path` | str | `params` | Path to tokenizer directory |
+
+Usage:
+```bash
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c sft.json
+```
+
+See [Preprocessing Guide](preprocessing.md) for config file format and examples.
+
+---
+
+> Document Update Time: 2026-07-20
@@ -0,0 +1,365 @@
+# Preprocessing Pipeline
+
+Declarative JSON-driven data preprocessing. `MaskBuilderFactory` supports three registered builders: `"single"` (single-output via `input.sections`), `"multi"` (multi-output via `input.sources`), and `"sectioned"` (façade dispatching to `single` or `multi` based on config).
+
+## Contents
+
+- [Philosophy](#philosophy)
+- [Config Structure](#config-structure)
+- [Quick Start](#quick-start) — SFT Chat, SFT Instruction, Pretrain, DPO, GRPO examples
+- [Configuration Reference](#configuration-reference) — all fields
+- [Mask Algorithm](#mask-algorithm)
+- [Output Layout](#output-layout)
+- [CLI](#cli)
+- [Python API](#python-api)
+
+## Philosophy
+
+| Component | Responsibility |
+|-----------|---------------|
+| `tokenizer_config.json` (`chat_template`) | Formatting -- how roles become tokens |
+| `pipeline.json` (`mask`) | Masking -- which roles participate in training |
+
+A single config file captures the entire pipeline, reusable and version-controllable.
+
+## Config Structure
+
+```json
+{
+  "version":       1,
+  "input":         {},   // sections (single) or sources (multi)
+  "mask":          {},   // role -> "train" | "mask"
+  "mask_default":  "mask",
+  "preprocessing": {},
+  "output":        {}
+}
+```
+
+### Section Fields
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `field` | str | -- | JSONL key to read |
+| `action` | str | -- | `"train"` / `"mask"` / `"$role"` |
+| `template` | bool | `false` | Apply `chat_template` per message |
+| `add_special_tokens` | bool | `true` for first non-template section | Add special tokens during encode |
+
+### Source Fields (multi-output mode)
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `sections` | list[dict] | -- | Same as single-output section list |
+| `list_field` | bool | `false` | JSONL field holds a list; tokenise each element |
+| `mask_key` | str | `"{key}_mask"` | Explicit output key for loss mask |
+
+---
+
+## Quick Start
+
+### SFT Chat
+
+Input JSONL:
+
+```json
+{"messages": [{"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sections": [
+      {"field": "messages", "action": "$role", "template": true}
+    ]
+  },
+  "mask": {
+    "system": "mask",
+    "user": "mask",
+    "assistant": "train"
+  },
+  "mask_default": "mask",
+  "preprocessing": {
+    "max_seq_len": 2048
+  },
+  "output": {
+    "storage_format": "bin",
+    "dtype": {"loss_mask": "bool"}
+  }
+}
+```
+
+Output keys: `sequence` (int32), `loss_mask` (bool)
+
+### SFT Instruction
+
+Input JSONL:
+
+```json
+{"prompt": "Translate to French: Hello", "response": "Bonjour"}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sections": [
+      {"field": "prompt",   "action": "mask", "add_special_tokens": true},
+      {"field": "response", "action": "train"}
+    ]
+  },
+  "mask_default": "mask",
+  "preprocessing": {
+    "max_seq_len": 2048
+  }
+}
+```
+
+Output keys: `sequence`, `loss_mask`
+
+### Pretrain
+
+Input JSONL:
+
+```json
+{"text": "Artificial Intelligence is a field of computer science..."}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sections": [
+      {"field": "text", "action": "train"}
+    ]
+  },
+  "preprocessing": {
+    "max_seq_len": 8192,
+    "min_chars": 100
+  }
+}
+```
+
+Output keys: `sequence` (no `loss_mask` — all tokens trained)
+
+### DPO
+
+Input JSONL:
+
+```json
+{"chosen": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"}], "rejected": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "5"}]}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sources": {
+      "chosen": {
+        "sections": [
+          {"field": "chosen", "action": "$role", "template": true}
+        ]
+      },
+      "rejected": {
+        "sections": [
+          {"field": "rejected", "action": "$role", "template": true}
+        ]
+      }
+    }
+  },
+  "mask": {
+    "user": "mask",
+    "assistant": "train"
+  },
+  "mask_default": "mask"
+}
+```
+
+Output keys: `chosen`, `chosen_mask`, `rejected`, `rejected_mask`
+
+### GRPO
+
+Input JSONL:
+
+```json
+{"prompt": [{"role": "user", "content": "What is 2+2?"}], "responses": ["4", "Five", "Four"], "rewards": [1.0, 0.3, 0.8]}
+```
+
+Config:
+
+```json
+{
+  "input": {
+    "sources": {
+      "prompts": {
+        "sections": [
+          {"field": "prompt", "action": "mask", "template": true}
+        ]
+      },
+      "responses": {
+        "sections": [
+          {"field": "responses", "action": "train"}
+        ],
+        "list_field": true,
+        "mask_key": "masks"
+      },
+      "rewards": {
+        "sections": [
+          {"field": "rewards", "action": "value"}
+        ]
+      }
+    }
+  },
+  "mask": {
+    "user": "mask",
+    "assistant": "train"
+  },
+  "mask_default": "mask"
+}
+```
+
+Output keys: `prompts`, `prompts_mask`, `responses`, `masks`, `rewards` (float32)
+
+- `action: "value"` — extract raw values from JSONL without tokenisation
+- `list_field: true` — tokenise each list element independently, then concatenate
+- `mask_key: "masks"` — rename the auto-generated mask key (default: `responses_mask`)
+- `prompts_mask` is auto-generated (all masked) and unused by GRPOStrategy
+
+---
+
+## Configuration Reference
+
+### `input`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `sections` | list[dict] or null | `null` | Section specs for single-output mode |
+| `sources` | dict[str, dict] or null | `null` | Source specs for multi-output mode (DPO/GRPO) |
+
+When `sources` is set, `sections` is ignored.
+
+### `mask`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `mask` | dict | `{}` | `{role: "train" \| "mask"}` |
+| `mask_default` | str | `"mask"` | Default action for unlisted roles |
+
+### `preprocessing`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `max_seq_len` | int | `2048` | Truncate sequences to this length |
+| `min_chars` | int | `50` | Skip text-mode items shorter than this |
+| `max_chars` | int | `2000000` | Skip text-mode items longer than this |
+| `max_items` | int or null | `null` | Stop after N documents |
+| `batch_size` | int | `256` | Records per tokenization batch |
+| `packing_strategy` | str | `"simple"` | Packing strategy: `"simple"`, `"bfd"`, `"bfd_split"` |
+| `max_packed_len` | int | `8192` | Maximum length of a packed bin |
+| `truncation_mode` | str | `"keep_start"` | How to truncate sequences: `"keep_start"` or `"keep_end"` |
+
+### `output`
+
+| Field | Type | Default | Description |
+|-------|------|---------|-------------|
+| `domain_key` | str or null | `null` | JSONL key for domain grouping |
+| `storage_format` | str | `"bin"` | `"bin"` (mmap). Reading also supports `"jsonl"` for on-the-fly tokenization |
+| `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens |
+| `dtype` | dict[str, str] | `{}` | Per-key tensor dtype override (e.g. `{"loss_mask": "bool"}`) |
+| `position_ids_mode` | str | `"doc_reset"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` |
+
+---
+
+## Mask Algorithm
+
+### Template mode (`template: true`)
+
+1. Prepend BOS token (masked)
+2. For each message in the field's array:
+   1. Render through `chat_template` for that single message
+   2. Encode rendered text
+   3. Apply mask rule for the message's role
+
+### Non-template mode
+
+Encode the field value as text. Mask value is 1 (train) or 0 (mask) per the section's `action`.
+
+### Text config detection
+
+When no section uses `template` and all sections have `action: "train"`, the builder omits `loss_mask` from the output — all tokens are trained.
+
+---
+
+## Output Layout
+
+### Single-Shard (`bin`)
+
+```
+output/
+  __default__/
+    shard_0000/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+  wiki/
+    shard_0000/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+```
+
+### Multi-Shard (`bin`)
+
+When `max_tokens_per_shard` is exceeded:
+
+```
+output/
+  __default__/
+    shard_0000/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+    shard_0001/
+      meta.json
+      sequence.bin
+      loss_mask.bin
+```
+
+For `bin` format, `MmapStore` discovers all shards under the domain directory via `rglob("meta.json")`. For `h5` format, `H5Store` discovers `.h5`/`.hdf5` files via recursive glob.
+
+---
+
+## CLI
+
+```bash
+# SFT
+python scripts/tools/preprocess.py data/sft/*.jsonl -o output/sft/ -c configs/sft_chat.json
+
+# DPO
+python scripts/tools/preprocess.py data/dpo/*.jsonl -o output/dpo/ -c configs/dpo.json --tokenizer_path params
+
+# GRPO
+python scripts/tools/preprocess.py data/grpo/*.jsonl -o output/grpo/ -c configs/grpo.json
+```
+
+---
+
+## Python API
+
+```python
+from astrai.preprocessing.pipeline import Pipeline
+from astrai.config.preprocess_config import PipelineConfig
+
+config = PipelineConfig.from_file("sft.json")
+Pipeline(
+    config,
+    ["data_part1.jsonl", "data_part2.jsonl"],
+    output_dir="output/",
+    tokenizer_path="params",
+).run()
+```
+
+> Document Update Time: 2026-07-09
@@ -0,0 +1,235 @@
+# Training
+
+## Contents
+
+- [Autoregression](#autoregression)
+- [Causal Mask](#causal-mask)
+- [Rotary Position Embedding (RoPE)](#rotary-position-embedding-rope)
+- [Training Loop](#training-loop)
+- [Strategies](#strategies) — SEQ, SFT, DPO, GRPO, online rollout
+- [LR Schedulers](#lr-schedulers)
+- [Gradient Checkpointing](#gradient-checkpointing)
+- [Checkpoint](#checkpoint)
+- [TrainContextBuilder](#traincontextbuilder-builder-pattern)
+- [Training CLI](#training-cli)
+
+### Autoregression
+
+Given a token sequence, the model predicts the probability of the next token. Each generated token is appended to the input and fed back, repeating until an end-of-sequence token or max length.
+
+### Causal Mask
+
+```
+sequence : [[1, 2, 3, 4, 5, 6]]
+input_ids: [[1, 2, 3, 4, 5]]
+target_ids: [[2, 3, 4, 5, 6]]
+```
+
+Lower-triangular mask prevents attending to future positions:
+
+```
+[[0, -inf, -inf, -inf, -inf],
+ [0,    0, -inf, -inf, -inf],
+ [0,    0,    0, -inf, -inf],
+ [0,    0,    0,    0, -inf],
+ [0,    0,    0,    0,    0]]
+```
+
+### Rotary Position Embedding (RoPE)
+
+RoPE embeds position into Q/K vectors via complex rotation:
+
+$$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
+
+The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers.
+
+## Training Loop
+
+Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
+
+```
+on_train_begin
+  model.train()
+  on_epoch_begin
+    for batch in dataloader:
+      on_batch_begin
+      with executor.accumulate(model):
+        loss = strategy.compute_loss(batch)
+        context.loss = loss.item()
+        stand_loss = loss / executor.grad_accum_steps
+        executor.backward(stand_loss)
+        context.consumed_samples += (
+            context.config.batch_per_device * context.world_size
+        )
+        on_batch_end
+
+        if executor.sync_gradients:
+          on_optimizer_step
+          optimizer.step()
+          optimizer.zero_grad()
+          if scheduler:
+            scheduler.step()
+    on_epoch_end
+on_train_end
+```
+
+### Callback Lifecycle
+
+| Hook | Fires | Default callback |
+|------|-------|-----------------|
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
+| `on_batch_begin` | Every batch | — |
+| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricCallback`, `ProgressBarCallback` |
+| `on_batch_end` | Every batch | `CheckpointCallback` |
+| `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
+| `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
+| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricCallback`, `GradientCheckpointingCallback` |
+
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
+
+## Strategies
+
+### SEQ (Pre-training)
+
+Next-token cross-entropy with optional label smoothing:
+
+$$
+L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
+$$
+
+Keys: `input_ids`, `target_ids`. Optional: `label_smoothing`.
+
+### SFT (Supervised Fine-Tuning)
+
+Masked cross-entropy (`ignore_index=-100`) over response tokens:
+
+$$
+L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
+$$
+
+Keys: `input_ids`, `target_ids`, `loss_mask`, `position_ids`. Optional: `label_smoothing`.
+
+### DPO (Direct Preference Optimization)
+
+Frozen reference model, preference margin via log-ratio:
+
+$$
+L_{\text{DPO}} = -\mathbb{E}\left[\log\sigma\left(\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)} - \beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}\right)\right]
+$$
+
+Parameters: `beta=0.1`, `reduction="sum"`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
+
+### GRPO (Group Relative Policy Optimization)
+
+Token-level PPO with group-normalized advantages. Advantages are derived from
+scalar per-response rewards, group-normalized, and broadcast across all response
+tokens. Only response tokens contribute to the loss (prompt tokens are masked
+out):
+
+$$
+\text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
+$$
+
+$$
+L_{\text{GRPO}} = -\mathbb{E}_t\left[\min\left(\rho_t A,\; \text{clip}\left(\rho_t, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}_t\left[\frac{\pi_{\text{ref}}}{\pi_\theta} - \log\frac{\pi_{\text{ref}}}{\pi_\theta} - 1\right]
+$$
+
+where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the
+per-token importance sampling ratio against the behaviour policy
+(`old_model`, synced externally between data-generation rounds) and the
+expectations are over valid response tokens. The KL term regularises
+$\pi_\theta$ towards a frozen reference model (`ref_model`, typically
+the SFT checkpoint).
+
+Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. External sync of `old_model` weights via `sync_old_model()` between data-generation rounds.
+
+Keys: `prompts`, `responses`, `masks`, `rewards`.
+
+### Online Rollout
+
+`online_grpo` and `online_dpo` use the respective GRPO and DPO strategies with
+a `RolloutRunner`. The runner renders prompts through the tokenizer chat
+template, generates grouped responses through `InferenceScheduler`, then scores
+them with a `BaseRewardModel`. It refreshes cached rollouts every
+`rollout_interval` optimizer steps. `online_grpo` synchronizes `old_model` when
+a fresh rollout is produced.
+
+Online strategies require `TrainConfig.reward_model_fn`. `train.py` exposes the
+rollout sampling parameters but does not yet offer a CLI argument for the reward
+model factory.
+
+## LR Schedulers
+
+| Type | Class | Description |
+|------|-------|-------------|
+| Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
+| SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
+| WSD | `WSDScheduler` | Warmup-Stable-Decay with sqrt cooldown |
+
+Created by `SchedulerFactory.create(schedule_type, optimizer, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`, `"wsd"`. Omit to use no scheduler.
+
+## Gradient Checkpointing
+
+Trades compute for memory by recomputing activations during backward pass. Specify module types via `gradient_checkpointing_modules`:
+
+```python
+from astrai.model.components.decoder_block import DecoderBlock
+
+config = TrainConfig(..., gradient_checkpointing_modules=[DecoderBlock])
+```
+
+Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoint(use_reentrant=False)`, compatible with `torch.compile`. Uses `nn.Module.apply()` for traversal — works through DDP wrappers without manual unwrap. Empty list (default) means no-op.
+
+## Checkpoint
+
+```
+Checkpoint(state_dict, epoch, consumed_samples, extra, meta, config)
+  ├── save(save_dir)    rank-0 only: meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
+  └── load(save_dir, broadcast=False)    loads from local disk; set broadcast=True to broadcast metadata from rank-0
+```
+
+Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
+Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.
+
+## TrainContextBuilder (Builder Pattern)
+
+```python
+context = TrainContextBuilder(config).with_param_path(param_path, resume=True).build()
+# Returns TrainContext with model, strategy, optimizer, scheduler, dataloader, checkpoint
+```
+
+- Loads checkpoint weights before the model is wrapped
+- Creates executor via `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)`
+- Calls `executor.prepare(model_fn, optimizer_fn, scheduler_fn, before_wrap=...)`; the executor creates, wraps, then builds the optimizer and scheduler for the wrapped model
+- Creates `RDSampler` for shuffle+resume
+- Builds strategy via `StrategyFactory.create(train_type, model, device, **kwargs)`
+
+## Training CLI
+
+```bash
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --parallel_mode=ddp \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --weight_decay=0.1 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
+```
+
+Full parameter reference at [params.md](params.md).
+
+> Document Update Time: 2026-07-20