diff --git a/assets/docs/architecture.md b/assets/docs/architecture.md index d1e2655..10b326b 100644 --- a/assets/docs/architecture.md +++ b/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) + -_normalize(raw, offsets) + } + + class Streamable { + <> + +fetch(begin, end, keys) + -_fetch_stream_key(key, begin, end) Tensor + } + + class Recordable { + <> + +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 <|-- GRPODataset + BaseDataset <|-- RecordDataset + 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.dataset** | BaseDataset–GRPODataset, Store–JsonlStore/MmapStore/H5Store, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management | +| **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–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 diff --git a/assets/docs/dataflow.md b/assets/docs/dataflow.md index db7f3e6..fa5bca4 100644 --- a/assets/docs/dataflow.md +++ b/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 | -|------|-------------| -| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | -| `sft` | `sequence`, `loss_mask`, `position_ids` | -| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | -| `grpo` | `prompts`, `responses`, `masks`, `rewards` | +| Type | Storage Keys | Access Mode | +|------|-------------|-------------| +| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | stream (`fetch`) | +| `sft` | `sequence`, `loss_mask`, `position_ids` | stream (`fetch`) | +| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | record (`fetch_record`) | +| `grpo` | `prompts`, `responses`, `masks`, `rewards` | record (`fetch_record`) | ## Dataset Architecture ``` -DatasetFactory.load(train_type, load_path, 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]]] - → BaseDataset.__getitem__(idx) - → get_index(idx) → [begin, end) - → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor] + → Store._data[Dict[str, List[Tensor]]] + + _cum[Dict[str, List[int]]] (stream mode) + + _offsets[Dict[str, List[int]]] (record mode) + +Stream datasets (SEQ/SFT): + BaseDataset.__getitem__(idx) + → get_index(idx) → [begin, end) + → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor] + +Record datasets (DPO/GRPO via RecordDataset): + RecordDataset.__getitem__(idx) + → Store.fetch_record(idx, keys) → Tensor / Dict[str, Tensor] ``` -`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 diff --git a/assets/docs/params.md b/assets/docs/params.md index 4d5edf6..53c547e 100644 --- a/assets/docs/params.md +++ b/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 \ No newline at end of file +> Document Update Time: 2026-07-19 \ No newline at end of file diff --git a/assets/docs/training.md b/assets/docs/training.md index 9b138f3..c81b704 100644 --- a/assets/docs/training.md +++ b/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 -per-token probability ratio and the expectations are over valid response tokens. +where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the +per-token importance sampling ratio against the behaviour policy +(`old_model`, synced externally between data-generation rounds) and the +expectations are over valid response tokens. The KL term regularises +$\pi_\theta$ towards a frozen reference model (`ref_model`, typically +the SFT checkpoint). -Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`, `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