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

- Rename assets/ to docs/, split into guides/ and developer/ - Add get-started.md: installation + 5-step quickstart - Add guides/evaluation.md: 7 eval scripts with CLI args - Add guides/distributed.md: DDP/FSDP, gradient accumulation, NCCL - Add developer/internals.md: loss formulas, RoPE, KV cache math - Add developer/cuda_kernels.md: build system, benchmarks, file layout - Fix storage_format doc in preprocessing.md - Update cross-references in README.md, README-zh-CN.md, Dockerfile
2026-07-30 00:49:04 +08:00
parent bcaa2d1ae0
commit c8b1e40f71
17 changed files with 1124 additions and 25 deletions
@@ -0,0 +1,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