From c8b1e40f71678779b0495421a9a27f855538e3a8 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 30 Jul 2026 00:49:04 +0800 Subject: [PATCH] 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 --- .dockerignore | 2 +- .gitignore | 2 +- Dockerfile | 2 +- README.md | 23 +- {assets/docs => docs}/README-zh-CN.md | 25 +- .../docs => docs/developer}/architecture.md | 0 docs/developer/cuda_kernels.md | 131 +++++++++ {assets/docs => docs/developer}/dataflow.md | 4 +- docs/developer/internals.md | 208 ++++++++++++++ docs/get-started.md | 235 ++++++++++++++++ docs/guides/distributed.md | 255 ++++++++++++++++++ docs/guides/evaluation.md | 252 +++++++++++++++++ {assets/docs => docs/guides}/inference.md | 0 {assets/docs => docs/guides}/params.md | 8 + {assets/docs => docs/guides}/preprocessing.md | 2 +- {assets/docs => docs/guides}/training.md | 0 {assets => docs}/images/logo.png | Bin 17 files changed, 1124 insertions(+), 25 deletions(-) rename {assets/docs => docs}/README-zh-CN.md (85%) rename {assets/docs => docs/developer}/architecture.md (100%) create mode 100644 docs/developer/cuda_kernels.md rename {assets/docs => docs/developer}/dataflow.md (98%) create mode 100644 docs/developer/internals.md create mode 100644 docs/get-started.md create mode 100644 docs/guides/distributed.md create mode 100644 docs/guides/evaluation.md rename {assets/docs => docs/guides}/inference.md (100%) rename {assets/docs => docs/guides}/params.md (97%) rename {assets/docs => docs/guides}/preprocessing.md (98%) rename {assets/docs => docs/guides}/training.md (100%) rename {assets => docs}/images/logo.png (100%) diff --git a/.dockerignore b/.dockerignore index a5bbc7b..9926b40 100644 --- a/.dockerignore +++ b/.dockerignore @@ -4,6 +4,6 @@ # Allow necessary files !astrai/ !scripts/ -!assets/ +!docs/ !pyproject.toml !README.md diff --git a/.gitignore b/.gitignore index 024aee7..60e8739 100644 --- a/.gitignore +++ b/.gitignore @@ -24,7 +24,7 @@ !/.dockerignore !/Dockerfile !/docker-compose.yml -!/assets/** +!/docs/** !/CONTRIBUTING.md !/LICENSE !/pyproject.toml diff --git a/Dockerfile b/Dockerfile index ceaeb29..c07428c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 . diff --git a/README.md b/README.md index 48fc831..e391e6e 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@
- Logo + Logo

A lightweight Transformer training & inference framework

@@ -17,7 +17,7 @@
English • - 中文 • + 中文 • Issue Tracker • Discussions • HuggingFace @@ -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) | -| [Architecture](./assets/docs/architecture.md) | System architecture, class diagram & design patterns | -| [Training](./assets/docs/training.md) | Training loop, strategies & formulas | -| [Inference](./assets/docs/inference.md) | KVCache, continuous batching, sampling & HTTP API | -| [Data Flow](./assets/docs/dataflow.md) | Data pipeline, storage backends & dataset architecture | -| [Preprocessing](./assets/docs/preprocessing.md) | Declarative JSON-driven data preprocessing | +| [Get Started](./docs/get-started.md) | Installation and quickstart | +| [CLI Reference](./docs/guides/params.md) | Parameters for all CLI tools (train, server, generate, preprocess) | +| [Preprocessing](./docs/guides/preprocessing.md) | Declarative JSON-driven data preprocessing | +| [Training](./docs/guides/training.md) | Training loop, strategies & formulas | +| [Inference](./docs/guides/inference.md) | KVCache, continuous batching, sampling & HTTP API | +| [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 diff --git a/assets/docs/README-zh-CN.md b/docs/README-zh-CN.md similarity index 85% rename from assets/docs/README-zh-CN.md rename to docs/README-zh-CN.md index 1070d1d..df638dd 100644 --- a/assets/docs/README-zh-CN.md +++ b/docs/README-zh-CN.md @@ -1,9 +1,9 @@
- Logo + Logo
- English • + English • 中文
@@ -23,7 +23,7 @@
- English • + English • 中文 • 问题追踪 • 讨论区 • @@ -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 工具参数（训练、服务、生成、预处理） | -| [架构文档](./architecture.md) | 系统架构、类图与设计模式 | -| [训练文档](./training.md) | 训练循环、策略与公式 | -| [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API | -| [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 | -| [数据预处理](./preprocessing.md) | 声明式 JSON 驱动数据预处理 | +| [快速上手](./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 注意力内核与基准测试 | ### 贡献 diff --git a/assets/docs/architecture.md b/docs/developer/architecture.md similarity index 100% rename from assets/docs/architecture.md rename to docs/developer/architecture.md diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md new file mode 100644 index 0000000..347f4f3 --- /dev/null +++ b/docs/developer/cuda_kernels.md @@ -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__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 diff --git a/assets/docs/dataflow.md b/docs/developer/dataflow.md similarity index 98% rename from assets/docs/dataflow.md rename to docs/developer/dataflow.md index 80e35a0..5fff96f 100644 --- a/assets/docs/dataflow.md +++ b/docs/developer/dataflow.md @@ -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 diff --git a/docs/developer/internals.md b/docs/developer/internals.md new file mode 100644 index 0000000..a2ec20d --- /dev/null +++ b/docs/developer/internals.md @@ -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 diff --git a/docs/get-started.md b/docs/get-started.md new file mode 100644 index 0000000..d1ac064 --- /dev/null +++ b/docs/get-started.md @@ -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 --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 diff --git a/docs/guides/distributed.md b/docs/guides/distributed.md new file mode 100644 index 0000000..8462fdc --- /dev/null +++ b/docs/guides/distributed.md @@ -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 diff --git a/docs/guides/evaluation.md b/docs/guides/evaluation.md new file mode 100644 index 0000000..80481f5 --- /dev/null +++ b/docs/guides/evaluation.md @@ -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 `//` and `/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 `