ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-07-30 00:49:04 +08:00
parent bcaa2d1ae0
commit c8b1e40f71
17 changed files with 1124 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/README-zh-CN.md
+265
View File
@@ -0,0 +1,265 @@
<div align="center">
<img src="./images/logo.png" width="auto" alt="Logo">
<div>
<a href="../README.md">English</a> •
<a href="#chinese">中文</a>
</div>
<p>
<strong>轻量级 Transformer 训练与推理框架</strong>
</p>
</div>
<div align="center">
<img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
<img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
<img src="https://img.shields.io/github/v/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
<img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
<img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
</div>
<br>
<div align="center">
<a href="../README.md">English</a> •
<a href="#chinese">中文</a> •
<a href="https://github.com/ViperEkura/AstrAI/issues">问题追踪</a> •
<a href="https://github.com/ViperEkura/AstrAI/discussions">讨论区</a> •
<a href="https://huggingface.co/ViperEkura">HuggingFace</a>
</div>
<br>
## 📖 目录
- [特性](#特性)
- [快速上手](#快速上手)
- [演示](#演示)
- [文档](#文档)
- [贡献](#贡献)
- [社区](#社区)
- [许可证](#许可证)
---
<a id="chinese"></a>
## 中文
### 特性
- 🚀 **高性能**: 训练与推理双向优化，高效并行。
- 🔧 **灵活**: 支持 seq/sft/dpo/grpo 多种训练方式，可定制模型架构。
- 💡 **易用**: 简洁的 API 与丰富的示例、演示。
- 📦 **轻量**: 依赖少，部署简单。
- 🔬 **研究友好**: 模块化设计，便于实验新想法。
- 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
- 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
### 快速上手
端到端演示，只需 5 步：
**1. 安装**
```bash
git clone https://github.com/ViperEkura/AstrAI.git
cd AstrAI
pip install -e . # 纯 PyTorch（不含 CUDA 内核）
# CSRC_KERNELS=true pip install -e . --no-build-isolation # 可选：融合 CUDA 内核加速
# pip install -e ".[dev]" # 可选：开发依赖（pytest, ruff）
```
**2. 下载模型**
```bash
python scripts/demo/download.py # 下载 1B 检查点到 params/
```
**3. 预处理数据**
创建 `pretrain.json`（`seq` 策略的预处理配置）：
```json
{
"version": 1,
"input": {"sections": [{"field": "text", "action": "train"}]},
"preprocessing": {"max_seq_len": 2048},
"output": {"storage_format": "bin"}
}
```
```bash
python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
```
**4. 训练**
```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 &
```
**5. 启动服务并调用**
```bash
# 终端 1：启动服务
python scripts/tools/server.py --param_path ./params --device cuda
# 终端 2：发起请求
curl http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"messages":[{"role":"user","content":"你好"}],"max_tokens":512}'
```
### 演示
查看 `scripts/demo/` 文件夹中的演示：
```bash
# 下载模型权重（运行演示前必需）
python scripts/demo/download.py # model → params/
# 交互式流式聊天（多轮对话，保持历史记录）
python scripts/demo/stream_chat.py
# 在 >> 后输入消息，输入 !exit 退出
# 批量生成（5 条硬编码提示词，非流式）
python scripts/demo/generate_batch.py
# 单条提示词自回归流式生成
python scripts/demo/generate_ar.py
```
所有生成演示默认使用 `temperature=0.8`、`top_p=0.95`、`top_k=50`、`max_tokens=2048`，需要 `params/` 目录包含模型权重（请先运行 `download.py`）。
观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
---
更多选项请参考[文档](#文档)。
#### 文本生成
从 JSONL 文件批量生成：
```bash
python scripts/tools/generate.py \
--param_path ./params \
--input_json_file input.jsonl \
--output_json_file output.jsonl
```
#### Docker
使用 Docker 构建和运行（推荐用于 GPU 环境）：
```bash
# 构建镜像
docker build -t astrai:latest .
# 启用 GPU 运行
docker run --gpus all -it astrai:latest
# 运行推理服务
docker run --gpus all -p 8000:8000 astrai:latest \
python -m scripts.tools.server --port 8000 --device cuda
# 挂载数据卷
docker run --gpus all -v /path/to/data:/data -it astrai:latest
# Docker Compose（GPU，默认）
docker compose up -d
# Docker Compose（仅 CPU）
docker compose --profile cpu up -d
```
> **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。
#### HTTP API 示例
除[快速上手](#快速上手)流程外，更多请求示例：
```bash
# OpenAI 兼容流式
curl -X POST http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"messages":[{"role":"user","content":"讲个故事"}],"stream":true,"max_tokens":500}'
# Anthropic 兼容
curl -X POST http://localhost:8000/v1/messages \
-H "Content-Type: application/json" \
-d '{"model":"astrai","system":"你是一个乐于助人的助手。","messages":[{"role":"user","content":"你好"}],"max_tokens":512}'
# Anthropic 兼容流式并设置停止序列
curl -X POST http://localhost:8000/v1/messages \
-H "Content-Type: application/json" \
-d '{"model":"astrai","messages":[{"role":"user","content":"写个故事"}],"max_tokens":500,"stream":true,"stop_sequences":["结束"]}'
# 健康检查
curl http://localhost:8000/health
```
SSE 流式格式、错误码和统计端点详见[推理文档](guides/inference.md)。
### 文档
| 文档 | 说明 |
|------|------|
| [快速上手](./get-started.md) | 安装与快速入门 |
| [CLI 参考](./guides/params.md) | 所有 CLI 工具参数（训练、服务、生成、预处理） |
| [数据预处理](./guides/preprocessing.md) | 声明式 JSON 驱动数据预处理 |
| [训练文档](./guides/training.md) | 训练循环、策略与公式 |
| [推理文档](./guides/inference.md) | KVCache、连续批处理、采样与 HTTP API |
| [评估文档](./guides/evaluation.md) | HumanEval、MMLU、PPL、ROUGE、IFD、IFEval |
| [分布式训练](./guides/distributed.md) | 多卡 DDP / FSDP 训练 |
| [架构文档](./developer/architecture.md) | 系统架构、类图与设计模式 |
| [数据流程](./developer/dataflow.md) | 数据管道、存储后端与数据集架构 |
| [内部实现](./developer/internals.md) | 训练原理：损失公式、回调生命周期、KV Cache |
| [CUDA 内核](./developer/cuda_kernels.md) | 自定义 CUDA 注意力内核与基准测试 |
### 贡献
我们欢迎贡献！请参阅[贡献指南](../../CONTRIBUTING.md)了解详情。
1. Fork 本仓库。
2. 创建功能分支。
3. 提交更改。
4. 发起 Pull Request。
重大更改请先开 issue 讨论。
### 社区
- **GitHub Issues**: [问题追踪](https://github.com/ViperEkura/AstrAI/issues)
- **Discussions**: [GitHub 讨论区](https://github.com/ViperEkura/AstrAI/discussions)
- **HuggingFace**: [模型中心](https://huggingface.co/ViperEkura)
### 许可证
本项目采用 [GPL-3.0 许可证](../../LICENSE)。
---
<div align="center">
<em>专为高性能与易用性设计的轻量级 Transformer 框架。</em>
</div>
docs/developer/architecture.md
+1476
View File
File diff suppressed because it is too large Load Diff
docs/developer/cuda_kernels.md
+131
View File
@@ -0,0 +1,131 @@
# CUDA Kernels
AstrAI includes optional custom CUDA attention kernels for decode and prefill. These are **not built by default** and are **not yet wired into the model or inference path** — they are standalone kernels with benchmarks and tests.
## Overview
| Kernel | File | Description |
|--------|------|-------------|
| `attn_decode` | `attn_decode.cu` | Basic GQA decode attention |
| `attn_prefill` | `attn_prefill.cu` | Basic GQA prefill attention |
| `attn_paged_decode` | `attn_paged_decode.cu` | Paged KV cache decode attention |
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
| Variant | File | Optimization |
|---------|------|--------------|
| Split-KV MMA decode | `attn_decode_split_kv_mma.cuh` | Split KV across waraps + MMA (sm_80+) |
| Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across waraps + MMA (sm_80+) |
| Paged split-KV MMA decode | `attn_paged_decode_split_kv_mma.cuh` | Paged cache + split-KV + MMA |
## Build System
### Auto-detection
Kernels are built when **both** of these conditions are met:
1. `nvcc` is available on `PATH`
2. `torch.cuda.is_available()` returns `True`
Unless `CSRC_KERNELS=false` is set explicitly.
### Manual build
```bash
# During install
CSRC_KERNELS=true pip install -e . --no-build-isolation
# Rebuild after editing .cu/.cuh files
CSRC_KERNELS=true python setup.py build_ext --inplace
# Output: astrai/extension/*.so
```
### Architecture flags
`csrc/build.py` auto-detects the GPU compute capability and generates the appropriate `nvcc` gencode flag:
- **sm_80+** (Ampere and later): enables tensor-core MMA path (`mma.sync.m16n8k16.bf16`)
- **Below sm_80**: adds `-DASTRAI_NO_MMA` to disable the MMA path at compile time
### Build configuration
```
NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
--ptxas-options=-O3,-v --extra-device-vectorization --threads=8
```
The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 3). Each entry maps a kernel name to its source files and build flags.
## Python Wrappers
`astrai/extension/ops.py` provides Python wrappers for each compiled kernel. When the `.so` is not available, wrappers **fall back to `torch.nn.functional.scaled_dot_product_attention`** (SDPA).
Interface:
```
causal_offset: -1 = non-causal; >=0 = absolute position of first Q token
mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool)
scale: 0.0 = auto (1/sqrt(head_dim)); >0 = explicit
layout: "bhld" (default) or "blhd"
```
> **Note**: Wrappers are not yet called from `model/transformer.py` or `inference/`. The model uses PyTorch's built attention. Integration is future work.
## Standalone Testing
Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment. Example:
```bash
nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
--ptxas-options=-O3,-v --extra-device-vectorization \
csrc/tests/attn_decode_test.cu -o /tmp/test && /tmp/test
```
Test files:
- `attn_decode_test.cu` — basic decode kernel
- `attn_paged_decode_test.cu` — paged decode kernel
- `attn_prefill_test.cu` — prefill kernel
## Benchmarks
Hardware: NVIDIA L20 (sm_89, 46 GB), CUDA 12.8, driver 570.86.
Reproduce:
```bash
nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
--ptxas-options=-O3,-v --extra-device-vectorization \
csrc/tests/attn_<name>_test.cu -o /tmp/test && /tmp/test
```
## Known Optimization Targets
- **Decode D=256**: spill eliminated (BC=16 + STAGES=2), but still 248 regs — further tiling could help.
- **Prefill single-batch**: bandwidth low (52 GB/s at q=kv=2048) — likely compute-bound but near L20 bf16 ceiling (~94 TFLOP/s).
- **Decode single-batch**: bandwidth low (309 GB/s at kv=512) — L20 HBM ~864 GB/s theoretical; small kv underutilizes SMs despite split-KV.
## File Layout
```
csrc/
├── build.py # Build system: REGISTRY, _arch_flags, nvcc flags
├── kernels/
│ ├── attn_common.h # Shared attention utilities
│ ├── attn_decode.cu # Basic decode kernel (registered)
│ ├── attn_prefill.cu # Basic prefill kernel (registered)
│ ├── attn_paged_decode.cu # Paged decode kernel (registered)
│ ├── attn_decode_split_kv.cuh # Split-KV variant
│ ├── attn_decode_split_kv_mma.cuh # Split-KV + MMA variant
│ ├── attn_prefill_split_q.cuh # Split-Q variant
│ ├── attn_prefill_split_q_mma.cuh # Split-Q + MMA variant
│ ├── attn_paged_decode_split_kv.cuh # Paged + split-KV variant
│ ├── attn_paged_decode_split_kv_mma.cuh # Paged + split-KV + MMA variant
│ ├── attn_dispatchers.cuh # Kernel dispatch macros
│ ├── attn_entry_utils.cuh # Entry point helpers
│ ├── attn_mma_utils.cuh # MMA utilities
│ └── attn_warp_utils.cuh # Warp-level utilities
└── tests/
├── test_utils.cuh # Shared test utilities
├── attn_decode_test.cu # Decode kernel test
├── attn_paged_decode_test.cu # Paged decode test
└── attn_prefill_test.cu # Prefill kernel test
```
> Document Update Time: 2026-07-30
docs/developer/dataflow.md
+132
View File
@@ -0,0 +1,132 @@
# Data Flow
This document describes the data pipeline: from raw text to model input tensors. For creating preprocessing configs, see [Preprocessing Guide](../guides/preprocessing.md).
## Contents
- [Overview](#overview)
- [Data Preparation](#data-preparation) — tokenization, format detection, backends
- [Data Keys by Training Type](#data-keys-by-training-type)
- [Dataset Architecture](#dataset-architecture)
- [Sampler](#sampler)
- [DataLoader](#dataloader)
## Overview
```
JSONL Lines → Pipeline (mask builder) → Tokenized Tensors
↓
.h5 or .bin storage
↓
Store.load()
↓
Store.fetch(begin, end, keys)
↓
BaseDataset.__getitem__(idx)
↓
Sampler → DataLoader → Training / Inference
```
## Data Preparation
Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
### Tokenization
The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](../guides/preprocessing.md)), and produces flat token sequences:
```python
# Per JSONL line: messages → chat template → token IDs + loss mask
tokens = tokenizer.encode(rendered_text) # List[int]
loss_mask = [0, 0, 0, 1, 1, 1, 1, 1, 1] # 0=masked, 1=train
# Stored as flat tensors, packed with other lines by packing strategy
```
The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
### Format Detection
`detect_format(load_path)` inspects the path:
- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, `.jsonl` → `"jsonl"`, unknown suffix raises `ValueError`
- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, `*.bin` + `**/meta.json` → `"bin"`, or `*.jsonl` + `dataset_config.json` → `"jsonl"`
### Store Backends
Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
```
StoreFactory.create("h5") → H5Store
StoreFactory.create("bin") → MmapStore
StoreFactory.create("jsonl") → JsonlStore
```
All three inherit `Store` (base, owns `_data`/`_cum`/`_offsets`/`_normalize`) plus the `Streamable` and `Recordable` mixins, so every backend supports both `fetch(begin, end, keys)` (stream) and `fetch_record(index, keys)` (record) APIs.
**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage. `segments_are_records=True` — each `data_i` dataset is one record.
**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. `segments_are_records=False` — bin segments are contiguous streams; record access is driven by `_offsets` (written when `save_bin(..., record_keys=...)` was used at preprocessing time).
**JsonlStore**: On-the-fly tokenization of raw JSONL files at load time. Requires a `dataset_config.json` alongside the `.jsonl` files following the same `PipelineConfig` schema with an additional `tokenizer_path` field. Two modes: eager (default, applies `TokenizeTransform` to all records at load) and lazy (`processor=fn` given, defers tokenisation to `fetch_record` — used by DPO/GRPO).
All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record-mode indexing). Nested keys (GRPO `responses`/`masks` as `List[List[Tensor]]`) are stored as-is and excluded from both bookkeepings — they are only accessed record-by-record.
## Data Keys by Training Type
| Type | Storage Keys | Access Mode |
|------|-------------|-------------|
| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | stream (`fetch`) |
| `sft` | `sequence`, `loss_mask`, `position_ids` | stream (`fetch`) |
| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | record (`fetch_record`) |
| `grpo` | `prompts`, `responses`, `masks`, `rewards` | record (`fetch_record`) |
## Dataset Architecture
```
DatasetFactory.load(
train_type, load_path=None, window_size=0, stride=None,
storage_type=None, tokenizer_path=None,
max_len=2048, store=None
)
→ BaseDataset.load(load_path, storage_type=None)
→ detect_format(load_path)
→ StoreFactory.create(storage_type)
→ Store.load(load_path)
→ _normalize(raw) # base Store, shared by both backends
→ Store._data[Dict[str, List[Tensor]]]
+ _cum[Dict[str, List[int]]] (stream mode)
+ _offsets[Dict[str, List[int]]] (record mode)
Stream datasets (SEQ/SFT):
BaseDataset.__getitem__(idx)
→ get_index(idx) → [begin, end)
→ Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
Record datasets (DPO/GRPO via RecordDataset):
RecordDataset.__getitem__(idx)
→ Store.fetch_record(idx, keys) → Tensor / Dict[str, Tensor]
```
Class hierarchy: `BaseDataset` ← `SEQDataset` / `SFTDataset` (stream); `BaseDataset` ← `RecordDataset` ← `DPODataset` / `GRPODataset` (record).
`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). Only meaningful for stream datasets — record datasets ignore both. `storage_type` defaults to `None` (auto-detect via `detect_format`).
`tokenizer_path` triggers lazy on-the-fly tokenisation for record datasets on raw JSONL (DPO builds a `dpo_processor`; SEQ/SFT/pre-tokenised backends ignore it). `store` (pre-built `Store`) bypasses `load_path`/`storage_type`/`tokenizer_path` entirely — the caller controls Store construction.
`Store.fetch(begin, end, keys)` (stream mode, on `Streamable`): accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present (bin layout with per-record offsets), otherwise indexes `_data[key]` directly (H5/JSONL where each segment is one record).
## Sampler
`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
- Tracks `start_epoch` / `start_iter` for resume
- Shuffle via `torch.Generator(seed + epoch)`
- Per-replica index slicing for DDP
## DataLoader
Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
> Document Update Time: 2026-07-19
docs/developer/internals.md
+208
View File
@@ -0,0 +1,208 @@
# Internals
Mathematical foundations and internal algorithms for AstrAI's training, inference, and preprocessing pipelines. For practical usage guides, see [Training](../guides/training.md), [Inference](../guides/inference.md), and [Preprocessing](../guides/preprocessing.md).
## Contents
- [Autoregression & Causal Masking](#autoregression--causal-masking)
- [Rotary Position Embedding (RoPE)](#rotary-position-embedding-rope)
- [Training Loss Formulas](#training-loss-formulas)
- [Training Loop Internals](#training-loop-internals)
- [Callback Lifecycle](#callback-lifecycle)
- [KV Cache Mathematics](#kv-cache-mathematics)
- [Mask Algorithm Internals](#mask-algorithm-internals)
- [Gradient Accumulation Mechanics](#gradient-accumulation-mechanics)
## Autoregression & Causal Masking
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.
```
sequence : [[1, 2, 3, 4, 5, 6]]
input_ids: [[1, 2, 3, 4, 5]]
target_ids: [[2, 3, 4, 5, 6]]
```
A lower-triangular causal 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]]
```
This ensures position $i$ can only attend to positions $\leq i$, which is essential for autoregressive generation.
## 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. The key property is that the dot product $q_i^T k_j$ depends only on the relative position $i - j$, not the absolute positions.
**Critical for inference**: RoPE is applied **before** KV cache write, not after. If applied after caching, position encoding drift occurs because cached K/V would have stale rotation factors.
## Training Loss Formulas
### 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) $$
### SFT (Supervised Fine-Tuning)
Masked cross-entropy (`ignore_index=-100`) over response tokens only:
$$ L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta) $$
Prompt tokens are masked out via `loss_mask`; only response tokens contribute to the loss.
### 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"`.
### GRPO (Group Relative Policy Optimization)
Token-level PPO with group-normalized advantages:
$$ \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. Advantages are derived from scalar per-response rewards, group-normalized, and broadcast across all response tokens. Only response tokens contribute to the loss.
Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`.
## Training Loop Internals
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
```
The loss is divided by `grad_accum_steps` before `backward()`, so accumulated gradients sum to the correct mean.
## 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`).
## KV Cache Mathematics
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 $$
The cache stores $k_j$ and $v_j$ for all previous positions. At each decode step, only $q_n$ (the current query) is computed fresh, and attention is computed against the cached K/V.
**RoPE ordering**: RoPE is applied to Q/K **before** writing to the KV cache. This is essential because:
1. The cached K values already contain the rotation for their original positions.
2. The new Q is rotated for its current position.
3. The dot product $q_n^T k_j$ then correctly depends on $n - j$ (relative position).
If RoPE were applied after caching, the rotation factors would be inconsistent between cached and new tokens.
### Cache Implementations
- **ContiguousCache**: Each task gets a fixed slot of `[max_seq_len, num_key_value_heads, head_dim]`. Simple, efficient for small-to-medium batch sizes.
- **PageCache**: Paged KV cache with prefix sharing. Uses `PagePool` (allocator + LRU + prefix matching) and `Storage` (page tensors). Enables sharing of common prompt prefixes across requests.
## Mask Algorithm Internals
### 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.
### Position ID strategies
| Mode | Behavior |
|------|----------|
| `none` | No position IDs generated |
| `doc_reset` | Reset position to 0 at each document boundary in packed sequences |
| `continuous` | Continuous position IDs across packed documents |
Default is `doc_reset`, which ensures each document in a packed bin starts from position 0, preventing position encoding drift between unrelated documents.
## Gradient Accumulation Mechanics
Three cooperating layers enable gradient accumulation:
1. **`GradientState`** — tracks the micro-step counter. Fires `sync_gradients=True` every `grad_accum_steps` micro-batches. The counter is incremented at the **start** of `accumulate()`, before the forward pass.
2. **`executor._no_sync(model)`** — suppresses gradient synchronization on non-sync micro-steps:
- `NoneExecutor`: `nullcontext` (nothing to skip)
- `DDPExecutor`: `model.no_sync()` (PyTorch's built-in — skips all-reduce of gradient buckets)
- `FSDPExecutor`: `set_requires_gradient_sync(False, recurse=True)` on each `FSDPModule` (FSDP2's native mechanism)
3. **`AccumOptimizer` / `AccumScheduler`** — wrap the real optimizer/scheduler. `step()` and `zero_grad()` are gated on `sync_gradients` — they only forward to the inner optimizer when the sync flag is True.
The loss is divided by `grad_accum_steps` before `backward()`, so gradients sum to the correct mean across micro-steps. `consumed_samples` increments by `batch_per_device * world_size` every micro-batch.
### Effective batch size
$$ \text{Effective batch} = \text{nprocs} \times \text{batch\_per\_device} \times \text{grad\_accum\_steps} $$
### Total optimizer steps
```
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 accounts for data-parallel sharding — each rank processes `1/nprocs` of the dataset.
> Document Update Time: 2026-07-30
docs/get-started.md
+235
View File
@@ -0,0 +1,235 @@
# Getting Started
This guide walks you through installing AstrAI, downloading a model, running inference, preprocessing data, and launching your first training job.
## Prerequisites
- **Python 3.12+**
- **PyTorch 2.11+** (CUDA 12.8 recommended for GPU support)
- NVIDIA GPU with CUDA (optional but recommended; CPU works for inference)
## 1. Install
```bash
git clone https://github.com/ViperEkura/AstrAI.git
cd AstrAI
# Basic install (pure PyTorch, no custom CUDA kernels)
pip install -e .
# With CUDA kernels (optional, for fused attention)
# CSRC_KERNELS=true pip install -e . --no-build-isolation
# With dev dependencies (pytest, ruff)
# pip install -e ".[dev]"
```
> **CUDA kernels** are opt-in. They are not built by default and are not yet wired into the model or inference path. You can skip them for normal usage.
## 2. Download Model Weights
AstrAI uses HuggingFace-style model directories. Download the default 1B instruction-tuned model:
```bash
python scripts/demo/download.py
# → Downloads to params/
```
To use a different model:
```bash
python scripts/demo/download.py --repo-id <HF_REPO_ID> --local-dir ./my_model
```
The model directory contains:
- `config.json` — model architecture configuration
- `model.safetensors` — model weights
- `tokenizer.json` + `tokenizer_config.json` — tokenizer files (including chat template)
## 3. Run Inference
### Interactive Chat (Simplest)
```bash
python scripts/demo/stream_chat.py
# Type your message after >>, type !exit to quit
```
This starts a multi-turn interactive chat session with streaming output.
### Start an HTTP Server
```bash
# Terminal 1: start server
python scripts/tools/server.py --param_path ./params --device cuda
# Terminal 2: query (OpenAI-compatible API)
curl -X POST http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
```
The server also supports the Anthropic API at `/v1/messages`. See [Inference Guide](guides/inference.md) for full API documentation.
### Batch Generation from a File
Create an input JSONL file (one JSON object per line):
```json
{"question": "What is machine learning?"}
{"question": "Explain gradient descent."}
```
```bash
python scripts/tools/generate.py \
--param_path ./params \
--input_json_file input.jsonl \
--output_json_file output.jsonl
```
## 4. Preprocess Data
AstrAI uses a declarative JSON config to define the preprocessing pipeline. Create a config file for your training type:
### Pretraining (seq)
Input JSONL:
```json
{"text": "Artificial intelligence is..."}
```
Config (`pretrain.json`):
```json
{
"input": {
"sections": [{"field": "text", "action": "train"}]
},
"preprocessing": {"max_seq_len": 2048},
"output": {"storage_format": "bin"}
}
```
### SFT (Supervised Fine-Tuning)
Input JSONL:
```json
{"messages": [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]}
```
Config (`sft.json`):
```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"}}
}
```
### Run Preprocessing
```bash
python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
```
See [Preprocessing Guide](guides/preprocessing.md) for DPO/GRPO configs and all options.
## 5. Train
### Single GPU
```bash
python scripts/tools/train.py \
--train_type=seq \
--data_root_path=/path/to/dataset \
--param_path=./params \
--batch_per_device=4 \
--grad_accum_steps=8 \
--max_lr=1e-4 \
--window_size=2048 \
--ckpt_dir=./checkpoint \
--nprocs=1 \
--parallel_mode=none
```
### Multi-GPU (DDP)
```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 \
--data_root_path=/path/to/dataset \
--param_path=./params \
--parallel_mode=ddp \
--nprocs=4 \
--batch_per_device=4 \
--grad_accum_steps=8 \
--max_lr=1e-4 \
--window_size=2048 \
--ckpt_dir=./checkpoint
```
### Training Types
| `--train_type` | Description | Data Keys |
|----------------|-------------|-----------|
| `seq` | Pre-training (next-token prediction) | `sequence` |
| `sft` | Supervised fine-tuning (masked loss) | `sequence`, `loss_mask` |
| `dpo` | Direct Preference Optimization | `chosen`, `rejected`, `*_mask` |
| `grpo` | Group Relative Policy Optimization | `prompts`, `responses`, `masks`, `rewards` |
See [Training Guide](guides/training.md) for loss formulas and strategies. See [Distributed Guide](guides/distributed.md) for DDP/FSDP details.
## 6. Evaluate
```bash
# HumanEval (code generation, auto-downloads dataset)
python scripts/eval/evaluate_humaneval.py --param_path ./params --num_samples 20
# MMLU (knowledge, auto-downloads dataset)
python scripts/eval/evaluate_mmlu.py --param_path ./params --n_shot 5
# Perplexity on custom data
python scripts/eval/evaluate_ppl.py --param_path ./params --input_path data.jsonl --output_dir ppl_results/
```
See [Evaluation Guide](guides/evaluation.md) for all benchmarks.
## 7. Docker
```bash
# Build
docker build -t astrai:latest .
# Run inference server with GPU
docker run --gpus all -p 8000:8000 astrai:latest \
python -m scripts.tools.server --port 8000 --device cuda
# Docker Compose (GPU)
docker compose up -d
```
## Next Steps
| Topic | Document |
|-------|----------|
| CLI parameters (train, server, generate, preprocess) | [CLI Reference](guides/params.md) |
| Preprocessing pipeline details | [Preprocessing Guide](guides/preprocessing.md) |
| Training loop, strategies, schedulers | [Training Guide](guides/training.md) |
| KV cache, continuous batching, HTTP API | [Inference Guide](guides/inference.md) |
| Evaluation benchmarks | [Evaluation Guide](guides/evaluation.md) |
| Multi-GPU DDP / FSDP | [Distributed Guide](guides/distributed.md) |
| System architecture | [Architecture](developer/architecture.md) |
| Data pipeline internals | [Data Flow](developer/dataflow.md) |
> Document Update Time: 2026-07-30
docs/guides/distributed.md
+255
View File
@@ -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
docs/guides/evaluation.md
+252
View File
@@ -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
docs/guides/inference.md
+252
View File
@@ -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
docs/guides/params.md
+230
View File
@@ -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
docs/guides/preprocessing.md
+365
View File
@@ -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
docs/guides/training.md
+235
View File
@@ -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
docs/images/logo.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 281 KiB