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
@@ -4,6 +4,6 @@
 # Allow necessary files
 !astrai/
 !scripts/
-!assets/
+!docs/
 !pyproject.toml
 !README.md
@@ -24,7 +24,7 @@
 !/.dockerignore
 !/Dockerfile
 !/docker-compose.yml
-!/assets/**
+!/docs/**
 !/CONTRIBUTING.md
 !/LICENSE
 !/pyproject.toml
@@ -43,7 +43,7 @@ ENV PATH="/opt/venv/bin:$PATH"
 # Copy application code
 COPY astrai/ ./astrai/
 COPY scripts/ ./scripts/
-COPY assets/ ./assets/
+COPY docs/ ./docs/
 COPY pyproject.toml .
 COPY README.md .
@@ -1,6 +1,6 @@
 <div align="center">
-  <img src="assets/images/logo.png" width="auto" alt="Logo">
+  <img src="docs/images/logo.png" width="auto" alt="Logo">
  <p>
    <strong>A lightweight Transformer training & inference framework</strong>
  </p>
@@ -17,7 +17,7 @@
 <div align="center">
  <a href="#english">English</a> •
-  <a href="assets/docs/README-zh-CN.md">中文</a> •
+  <a href="docs/README-zh-CN.md">中文</a> •
  <a href="https://github.com/ViperEkura/AstrAI/issues">Issue Tracker</a> •
  <a href="https://github.com/ViperEkura/AstrAI/discussions">Discussions</a> •
  <a href="https://huggingface.co/ViperEkura">HuggingFace</a>
@@ -213,18 +213,23 @@ curl -X POST http://localhost:8000/v1/messages \
 curl http://localhost:8000/health
 ```
-See [Inference Guide](assets/docs/inference.md) for SSE streaming format, error codes, and stats endpoint.
+See [Inference Guide](docs/guides/inference.md) for SSE streaming format, error codes, and stats endpoint.
 ### Documentation
 | Document | Description |
 |----------|-------------|
-| [CLI Reference](./assets/docs/params.md) | Parameters for all CLI tools (train, server, generate, preprocess) |
+| [Get Started](./docs/get-started.md) | Installation and quickstart |
-| [Architecture](./assets/docs/architecture.md) | System architecture, class diagram & design patterns |
+| [CLI Reference](./docs/guides/params.md) | Parameters for all CLI tools (train, server, generate, preprocess) |
-| [Training](./assets/docs/training.md) | Training loop, strategies & formulas |
+| [Preprocessing](./docs/guides/preprocessing.md) | Declarative JSON-driven data preprocessing |
-| [Inference](./assets/docs/inference.md) | KVCache, continuous batching, sampling & HTTP API |
+| [Training](./docs/guides/training.md) | Training loop, strategies & formulas |
-| [Data Flow](./assets/docs/dataflow.md) | Data pipeline, storage backends & dataset architecture |
+| [Inference](./docs/guides/inference.md) | KVCache, continuous batching, sampling & HTTP API |
-| [Preprocessing](./assets/docs/preprocessing.md) | Declarative JSON-driven data preprocessing |
+| [Evaluation](./docs/guides/evaluation.md) | HumanEval, MMLU, PPL, ROUGE, IFD, IFEval |
 | [Distributed](./docs/guides/distributed.md) | Multi-GPU DDP / FSDP training |
 | [Architecture](./docs/developer/architecture.md) | System architecture, class diagram & design patterns |
 | [Data Flow](./docs/developer/dataflow.md) | Data pipeline, storage backends & dataset architecture |
 | [Internals](./docs/developer/internals.md) | Training internals: loss formulas, callback lifecycle, KV cache |
 | [CUDA Kernels](./docs/developer/cuda_kernels.md) | Custom CUDA attention kernels & benchmarks |
 ### Contributing
@@ -1,9 +1,9 @@
 <div align="center">
-  <img src="../images/logo.png" width="auto" alt="Logo">
+  <img src="./images/logo.png" width="auto" alt="Logo">
  <div>
-    <a href="../../README.md">English</a> • 
+    <a href="../README.md">English</a> • 
    <a href="#chinese">中文</a>
  </div>
@@ -23,7 +23,7 @@
 <br>
 <div align="center">
-  <a href="../../README.md">English</a> •
+  <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> •
@@ -219,18 +219,23 @@ curl -X POST http://localhost:8000/v1/messages \
 curl http://localhost:8000/health
 ```
-SSE 流式格式、错误码和统计端点详见[推理文档](./inference.md)。
+SSE 流式格式、错误码和统计端点详见[推理文档](guides/inference.md)。
 ### 文档
 | 文档 | 说明 |
 |------|------|
-| [CLI 参考](./params.md) | 所有 CLI 工具参数（训练、服务、生成、预处理） |
+| [快速上手](./get-started.md) | 安装与快速入门 |
-| [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
+| [CLI 参考](./guides/params.md) | 所有 CLI 工具参数（训练、服务、生成、预处理） |
-| [训练文档](./training.md) | 训练循环、策略与公式 |
+| [数据预处理](./guides/preprocessing.md) | 声明式 JSON 驱动数据预处理 |
-| [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
+| [训练文档](./guides/training.md) | 训练循环、策略与公式 |
-| [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 |
+| [推理文档](./guides/inference.md) | KVCache、连续批处理、采样与 HTTP API |
-| [数据预处理](./preprocessing.md) | 声明式 JSON 驱动数据预处理 |
+| [评估文档](./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 注意力内核与基准测试 |
 ### 贡献
@@ -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
@@ -1,6 +1,6 @@
 # Data Flow
-This document describes the data pipeline: from raw text to model input tensors. For creating preprocessing configs, see [Preprocessing Guide](preprocessing.md).
+This document describes the data pipeline: from raw text to model input tensors. For creating preprocessing configs, see [Preprocessing Guide](../guides/preprocessing.md).
 ## Contents
@@ -33,7 +33,7 @@ Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or
 ### Tokenization
-The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](preprocessing.md)), and produces flat token sequences:
+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
@@ -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
@@ -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
@@ -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
@@ -174,6 +174,14 @@ python scripts/tools/server.py --param_path ./params --device cuda --dtype bfloa
 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 |
@@ -266,7 +266,7 @@ When `sources` is set, `sections` is ignored.
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `domain_key` | str or null | `null` | JSONL key for domain grouping |
-| `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` |
+| `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"` |