docs: sync architecture/dataflow/training/params with code

- dataflow.md: update DatasetFactory.load signature, stream vs record access, Store._offsets
- architecture.md: add tokenizer to Pipeline, TokenizeTransform class, RecordDataset, Streamable/Recordable mixins, fix GRPOStrategy (old_model/sync_old_model)
- training.md: DPO reduction="sum", GRPO rho_t uses pi_old, gradient_clipping always registered
- params.md: --max_grad_norm default None

assets/docs/architecture.md

@@ -117,7 +117,7 @@ classDiagram
             +int n_epoch
             +int batch_per_device
             +int grad_accum_steps
-            +float max_grad_norm
+            +Optional[float] max_grad_norm
             +list gradient_checkpointing_modules
             +int start_epoch
             +int start_samples
@@ -166,6 +166,13 @@ classDiagram
             +__getitem__(index) Dict
         }
 
+        class RecordDataset {
+            +Optional[Callable] processor
+            +load(load_path, storage_type)
+            +__getitem__(index)
+            +__len__()
+        }
+
         class DPODataset {
             +__getitem__(index) Dict
         }
@@ -177,13 +184,26 @@ classDiagram
         class Store {
             +Dict[str, List[Tensor]] _data
             +Dict[str, List[int]] _cum
+            +Dict[str, List[int]] _offsets
             +int _length
+            +int _num_records
             +keys (property)
             +load(path)
-            +fetch(begin, end, keys)
             +__len__()
-            -_fetch_key(key, begin, end) Tensor
+            -_normalize(raw, offsets)
-            -_normalize(raw)
+        }
+
+        class Streamable {
+            <<mixin>>
+            +fetch(begin, end, keys)
+            -_fetch_stream_key(key, begin, end) Tensor
+        }
+
+        class Recordable {
+            <<mixin>>
+            +num_records (property)
+            +fetch_record(index, keys)
+            -_fetch_record_key(key, index) Tensor
         }
 
         class H5Store {
@@ -195,6 +215,13 @@ classDiagram
             +load(path)
         }
 
+        class JsonlStore {
+            +JsonlSource _source
+            +Callable _processor
+            +load(path, transform, processor)
+            +fetch_record(index, keys)
+        }
+
         class ResumableDistributedSampler {
             +int epoch
             +int iter
@@ -210,7 +237,7 @@ classDiagram
             +Dict _entries
             +register(name) decorator
             +create(train_type, window_size, stride) BaseDataset
-            +load(train_type, load_path, window_size, stride, storage_type) BaseDataset
+            +load(train_type, load_path, window_size, stride, storage_type, tokenizer_path, max_len, store) BaseDataset
         }
     }
 
@@ -378,6 +405,7 @@ classDiagram
             +List[str] paths
             +str output_dir
             +str tokenizer_path
+            +AutoTokenizer tokenizer
             +BaseMaskBuilder mask_builder
             +PackingStrategy _packer
             +PositionIdStrategy _position_id
@@ -385,6 +413,18 @@ classDiagram
             +transform(item) Optional[dict]
             +run()
             +_flush(domains, shard_idx)
+            +_inject_doc_reset_position_ids(keys, mode, seqs) Dict
+            +_inject_continuous_position_ids(tensors, mode, seqs) Dict
+            +_to_tensors(keys) Dict
+        }
+
+        class TokenizeTransform {
+            +PipelineConfig config
+            +AutoTokenizer tokenizer
+            +BaseMaskBuilder mask_builder
+            +PositionIdStrategy position_strategy
+            +from_config_file(path) TokenizeTransform
+            +apply(records) Dict[str, list]
         }
     }
 
@@ -495,13 +535,13 @@ classDiagram
         }
 
         class GRPOStrategy {
+            +nn.Module old_model
             +nn.Module ref_model
             +float clip_eps
             +float kl_coef
             +int group_size
-            +int sync_interval
             +compute_loss(batch) Tensor
-            +sync_ref_model()
+            +sync_old_model()
         }
 
         class BaseScheduler {
@@ -551,7 +591,7 @@ classDiagram
         }
 
         class GradientClippingCallback {
-            +float max_grad_norm
+            +Optional[float] max_grad_norm
             +on_optimizer_step(context)
         }
 
@@ -1064,11 +1104,18 @@ classDiagram
     TrainCallback <|-- MetricCallback
     BaseDataset <|-- SEQDataset
     BaseDataset <|-- SFTDataset
-    BaseDataset <|-- DPODataset
+    BaseDataset <|-- RecordDataset
-    BaseDataset <|-- GRPODataset
+    RecordDataset <|-- DPODataset
+    RecordDataset <|-- GRPODataset
     Store <|-- H5Store
     Store <|-- MmapStore
     Store <|-- JsonlStore
+    H5Store --|> Streamable
+    H5Store --|> Recordable
+    MmapStore --|> Streamable
+    MmapStore --|> Recordable
+    JsonlStore --|> Streamable
+    JsonlStore --|> Recordable
     BaseSamplingStrategy <|-- TemperatureStrategy
     BaseSamplingStrategy <|-- TopKStrategy
     BaseSamplingStrategy <|-- TopPStrategy
@@ -1143,6 +1190,9 @@ classDiagram
     BaseDataset o-- Store
     Pipeline o-- PipelineConfig
     Pipeline o-- BaseMaskBuilder
+    Pipeline o-- AutoTokenizer
+    TokenizeTransform o-- AutoTokenizer
+    TokenizeTransform o-- BaseMaskBuilder
 
     %% --- Dependency (uses temporarily) ---
     TrainConfig ..> BaseStrategy : selects
@@ -1186,7 +1236,7 @@ classDiagram
     %% --- Association (general usage) ---
     Trainer --> TrainConfig
     DPOStrategy --> AutoModel
-    GRPOStrategy --> AutoModel
+    GRPOStrategy --> AutoModel : policy/old/ref
     InferenceScheduler --> Task
     InferenceScheduler --> TaskStatus
     Task --> TaskStatus
@@ -1203,8 +1253,8 @@ classDiagram
 | Module | Components | Description |
 |--------|------------|-------------|
 | **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig, PipelineConfig, InputConfig, ProcessingConfig, OutputConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
-| **astrai.preprocessing** | BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, Pipeline, filter_by_length, PackingStrategy, PackingStrategyFactory, PositionIdStrategy, PositionIdStrategyFactory, StoreWriter, StoreWriterFactory | Declarative JSON-driven data preprocessing |
+| **astrai.preprocessing** | BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, SingleOutputMaskBuilder, MultiOutputMaskBuilder, Pipeline, TokenizeTransform, filter_by_length, PackingStrategy, PackingStrategyFactory, plan_bfd, PositionIdStrategy, PositionIdStrategyFactory, StoreWriter, StoreWriterFactory, core (shared helpers) | Declarative JSON-driven data preprocessing |
-| **astrai.dataset** | BaseDataset–GRPODataset, Store–JsonlStore/MmapStore/H5Store, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
+| **astrai.dataset** | BaseDataset–RecordDataset–DPO/GRPODataset, SEQDataset, SFTDataset, Store, Streamable, Recordable, H5Store, MmapStore, JsonlStore, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
 | **astrai.model** | AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
@@ -1246,4 +1296,4 @@ classDiagram
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-19

assets/docs/dataflow.md

@@ -61,41 +61,59 @@ StoreFactory.create("bin")   → MmapStore
 StoreFactory.create("jsonl") → JsonlStore
 ```
 
-**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage.
+All three inherit `Store` (base, owns `_data`/`_cum`/`_offsets`/`_normalize`) plus the `Streamable` and `Recordable` mixins, so every backend supports both `fetch(begin, end, keys)` (stream) and `fetch_record(index, keys)` (record) APIs.
 
-**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
+**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage. `segments_are_records=True` — each `data_i` dataset is one record.
 
-**JsonlStore**: On-the-fly tokenization of raw JSONL files at load time. Requires a `dataset_config.json` alongside the `.jsonl` files following the same `PipelineConfig` schema with an additional `tokenizer_path` field.
+**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. `segments_are_records=False` — bin segments are contiguous streams; record access is driven by `_offsets` (written when `save_bin(..., record_keys=...)` was used at preprocessing time).
 
-All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based indexing).
+**JsonlStore**: On-the-fly tokenization of raw JSONL files at load time. Requires a `dataset_config.json` alongside the `.jsonl` files following the same `PipelineConfig` schema with an additional `tokenizer_path` field. Two modes: eager (default, applies `TokenizeTransform` to all records at load) and lazy (`processor=fn` given, defers tokenisation to `fetch_record` — used by DPO/GRPO).
+
+All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record-mode indexing). Nested keys (GRPO `responses`/`masks` as `List[List[Tensor]]`) are stored as-is and excluded from both bookkeepings — they are only accessed record-by-record.
 
 ## Data Keys by Training Type
 
-| Type | Storage Keys |
+| Type | Storage Keys | Access Mode |
-|------|-------------|
+|------|-------------|-------------|
-| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) |
+| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | stream (`fetch`) |
-| `sft` | `sequence`, `loss_mask`, `position_ids` |
+| `sft` | `sequence`, `loss_mask`, `position_ids` | stream (`fetch`) |
-| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` |
+| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | record (`fetch_record`) |
-| `grpo` | `prompts`, `responses`, `masks`, `rewards` |
+| `grpo` | `prompts`, `responses`, `masks`, `rewards` | record (`fetch_record`) |
 
 ## Dataset Architecture
 
 ```
-DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_type=None)
+DatasetFactory.load(train_type, load_path, window_size, stride=None,
+                    storage_type=None, tokenizer_path=None,
+                    max_len=2048, store=None)
   → BaseDataset.load(load_path, storage_type=None)
     → detect_format(load_path)
     → StoreFactory.create(storage_type)
     → Store.load(load_path)
       → _normalize(raw)  # base Store, shared by both backends
-        → Store._data[Dict[str, List[Tensor]]] + _cum[Dict[str, List[int]]]
+        → Store._data[Dict[str, List[Tensor]]]
-          → BaseDataset.__getitem__(idx)
+          + _cum[Dict[str, List[int]]]   (stream mode)
+          + _offsets[Dict[str, List[int]]]  (record mode)
+
+Stream datasets (SEQ/SFT):
+  BaseDataset.__getitem__(idx)
     → get_index(idx) → [begin, end)
     → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
+
+Record datasets (DPO/GRPO via RecordDataset):
+  RecordDataset.__getitem__(idx)
+    → Store.fetch_record(idx, keys) → Tensor / Dict[str, Tensor]
 ```
 
-`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). `storage_type` defaults to `None` (auto-detect via `detect_format`).
+Class hierarchy: `BaseDataset` ← `SEQDataset` / `SFTDataset` (stream); `BaseDataset` ← `RecordDataset` ← `DPODataset` / `GRPODataset` (record).
 
-`Store.fetch(begin, end, keys)` accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
+`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). Only meaningful for stream datasets — record datasets ignore both. `storage_type` defaults to `None` (auto-detect via `detect_format`).
+
+`tokenizer_path` triggers lazy on-the-fly tokenisation for record datasets on raw JSONL (DPO builds a `dpo_processor`; SEQ/SFT/pre-tokenised backends ignore it). `store` (pre-built `Store`) bypasses `load_path`/`storage_type`/`tokenizer_path` entirely — the caller controls Store construction.
+
+`Store.fetch(begin, end, keys)` (stream mode, on `Streamable`): accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
+
+`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present (bin layout with per-record offsets), otherwise indexes `_data[key]` directly (H5/JSONL where each segment is one record).
 
 ## Sampler
 
@@ -109,4 +127,4 @@ DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_typ
 
 Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-19

assets/docs/params.md

@@ -26,7 +26,7 @@
 |-----------|-------------|---------|
 | `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
-| `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
+| `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | None |
 
 ### Optimizer (MuonMix)
 
@@ -201,4 +201,4 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples.
 
 ---
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-19

assets/docs/training.md

@@ -86,7 +86,7 @@ on_train_end
 | `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`.
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
 
 ## Strategies
 
@@ -118,7 +118,7 @@ $$
 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="mean"`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
+Parameters: `beta=0.1`, `reduction="sum"`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
 
 ### GRPO (Group Relative Policy Optimization)
 
@@ -135,10 +135,14 @@ $$
 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{ref}}(a_t|s_t)$ is the
+where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the
-per-token probability ratio and the expectations are over valid response tokens.
+per-token importance sampling ratio against the behaviour policy
+(`old_model`, synced externally between data-generation rounds) and the
+expectations are over valid response tokens. The KL term regularises
+$\pi_\theta$ towards a frozen reference model (`ref_model`, typically
+the SFT checkpoint).
 
-Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`, `sync_interval=200`.
+Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. External sync of `old_model` weights via `sync_old_model()` between data-generation rounds.
 
 Keys: `prompts`, `responses`, `masks`, `rewards`.
 
@@ -218,4 +222,4 @@ nohup python scripts/tools/train.py \
 
 Full parameter reference at [params.md](params.md).
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-19