diff --git a/assets/docs/architecture.md b/assets/docs/architecture.md index 30caff4..4ba3211 100644 --- a/assets/docs/architecture.md +++ b/assets/docs/architecture.md @@ -125,7 +125,6 @@ classDiagram +str ckpt_dir +int ckpt_interval +str log_dir - +int log_interval +List[str] metrics +Optional[LoRAConfig] lora +int random_seed @@ -559,7 +558,7 @@ classDiagram } class GradientCheckpointingCallback { - +tuple modules + +Optional[List[type]] modules +on_train_begin(context) +on_train_end(context) } @@ -573,31 +572,29 @@ classDiagram +on_batch_end(context) +on_train_end(context) +on_error(context) - +save_extra(context) dict$ + +save_extra(context) dict } class ProgressBarCallback { +int num_epoch +int log_interval +IO file + +tqdm progress_bar +on_epoch_begin(context) - +on_batch_end(context) + +on_optimizer_step(context) +on_epoch_end(context) } - class MetricLoggerCallback { + class MetricCallback { +Path log_dir +int save_interval - +int log_interval +List[str] metrics - +on_batch_end(context) + +int val_step + +on_optimizer_step(context) + +on_epoch_end(context) +on_train_end(context) +on_error(context) - } - - class ValidationCallback { -_run_validation(context) - +on_optimizer_step(context) } class CallbackFactory { @@ -684,20 +681,44 @@ classDiagram } class KVCache { - -PagePool _pool - -Storage _storage - -TaskTable _table - +int page_size + <> +task_alloc(task_id, prompt_ids) bool +task_free(task_id) +task_extend(task_id, pos) bool +task_cached(task_id) int +task_record_hashes(task_id, prompt_ids, start_logical_page) - +make_table_tensor(task_ids, device) Tensor - +bind(page_table, total_len) KvcacheView + +bind_tasks(task_ids, total_len, device) CacheView } - class KvcacheView { + class PageCache { + +int page_size + -PagePool _pool + -Storage _storage + -TaskTable _table + +task_alloc(task_id, prompt_ids) bool + +task_free(task_id) + +task_extend(task_id, pos) bool + +task_cached(task_id) int + +task_record_hashes(task_id, prompt_ids, start_logical_page) + +bind_tasks(task_ids, total_len, device) PageCacheView + } + + class ContiguousCache { + +int max_seq_len + +Tensor k, v + +task_alloc(task_id, prompt_ids) bool + +task_free(task_id) + +task_extend(task_id, pos) bool + +bind_tasks(task_ids, total_len, device) ContiguousCacheView + } + + class CacheView { + <> + +write(layer_id, k, v) + +gather(layer_id) Tuple[Tensor, Tensor] + } + + class PageCacheView { -Storage _storage +Tensor _page_table +int _total_len @@ -705,6 +726,14 @@ classDiagram +gather(layer_id) Tuple[Tensor, Tensor] } + class ContiguousCacheView { + -ContiguousCache _cache + +Tensor _batch_indices + +int _total_len + +write(layer_id, k, v) + +gather(layer_id) Tuple[Tensor, Tensor] + } + class TaskTable { +set(task_id, page_table, cached) +get(task_id) List[int] @@ -1035,14 +1064,14 @@ classDiagram TrainCallback <|-- GradientCheckpointingCallback TrainCallback <|-- CheckpointCallback TrainCallback <|-- ProgressBarCallback - TrainCallback <|-- MetricLoggerCallback - TrainCallback <|-- ValidationCallback + TrainCallback <|-- MetricCallback BaseDataset <|-- SEQDataset BaseDataset <|-- SFTDataset BaseDataset <|-- DPODataset BaseDataset <|-- GRPODataset Store <|-- H5Store Store <|-- MmapStore + Store <|-- JsonlStore BaseSamplingStrategy <|-- TemperatureStrategy BaseSamplingStrategy <|-- TopKStrategy BaseSamplingStrategy <|-- TopPStrategy @@ -1075,11 +1104,15 @@ classDiagram ResponseBuilder <|-- OpenAIResponseBuilder ResponseBuilder <|-- AnthropicResponseBuilder BaseMaskBuilder <|-- SectionedMaskBuilder + KVCache <|-- PageCache + KVCache <|-- ContiguousCache + CacheView <|-- PageCacheView + CacheView <|-- ContiguousCacheView %% --- Composition (strong ownership, part destroyed with whole) --- - KVCache *-- PagePool - KVCache *-- Storage - KVCache *-- TaskTable + PageCache *-- PagePool + PageCache *-- Storage + PageCache *-- TaskTable InferenceEngine *-- InferenceScheduler InferenceScheduler *-- KVCache InferenceScheduler *-- Executor @@ -1107,7 +1140,8 @@ classDiagram TrainContext o-- BaseScheduler TrainContext o-- Checkpoint TrainContext o-- BaseExecutor - KvcacheView o-- Storage + PageCacheView o-- Storage + ContiguousCacheView o-- ContiguousCache SamplingPipeline o-- BaseSamplingStrategy BaseDataset o-- Store Pipeline o-- PipelineConfig @@ -1129,6 +1163,7 @@ classDiagram DecoderBlock ..> FFNFactory : uses StoreFactory ..> H5Store : creates StoreFactory ..> MmapStore : creates + StoreFactory ..> JsonlStore : creates ConfigFactory ..> AutoRegressiveLMConfig : creates ConfigFactory ..> EncoderConfig : creates ExecutorFactory ..> NoneExecutor : creates @@ -1142,7 +1177,8 @@ classDiagram TrainContextBuilder ..> ResumableDistributedSampler : creates Checkpoint ..> Checkpoint : serializes CheckpointCallback ..> Checkpoint : creates - KVCache ..> KvcacheView : binds + PageCache ..> PageCacheView : binds + ContiguousCache ..> ContiguousCacheView : binds InferenceEngine ..> GenerationRequest : uses InferenceEngine ..> GenerateResult : creates OpenAIResponseBuilder ..> ChatCompletionRequest : receives @@ -1171,12 +1207,12 @@ classDiagram |--------|------------|-------------| | **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–MmapStore, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management | +| **astrai.dataset** | BaseDataset–GRPODataset, Store–JsonlStore/MmapStore/H5Store, 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 | -| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–ValidationCallback, CallbackFactory | Training workflow | -| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–KvcacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, ChatMessage–MessagesRequest, app | Inference service | +| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory | Training workflow | +| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–ContiguousCache/PageCache, CacheView–ContiguousCacheView/PageCacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, ChatMessage–MessagesRequest, app | Inference service | | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation | | **astrai.factory** | BaseFactory | Component registration | | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers | @@ -1195,7 +1231,7 @@ classDiagram | **Context** | `TrainContext` | Unified training state bag | | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction | | **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution | -| **Storage** | `Store`, `H5Store`, `MmapStore` | Format-agnostic data access with multi-segment support | +| **Storage** | `Store`, `H5Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support | | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching | | **AutoModel Registry** | `AutoModel`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading | @@ -1207,10 +1243,10 @@ classDiagram 4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor` 5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `KVCache` + `SamplingPipeline` 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP -7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore) loads data with explicit `_length` and multi-segment `_data` +7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore/JsonlStore) loads data with explicit `_length` and multi-segment `_data` 8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only), extra state saved as `{key}.pt` 9. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`/`WSDScheduler` 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-05-30 +> Document Update Time: 2026-07-05 diff --git a/assets/docs/dataflow.md b/assets/docs/dataflow.md index 310018b..0b5a7f1 100644 --- a/assets/docs/dataflow.md +++ b/assets/docs/dataflow.md @@ -48,23 +48,26 @@ The output `meta.json` records the storage format, key names, dtype, total token `detect_format(load_path)` inspects the path: -- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, unknown suffix raises `ValueError` -- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, or `*.bin` + `**/meta.json` → `"bin"` +- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, `.jsonl` → `"jsonl"`, unknown suffix raises `ValueError` +- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, `*.bin` + `**/meta.json` → `"bin"`, or `*.jsonl` + `dataset_config.json` → `"jsonl"` ### Store Backends Storage format is auto-detected by `detect_format()`; backends are dispatched via registry: ``` -StoreFactory.create("h5") → H5Store -StoreFactory.create("bin") → MmapStore +StoreFactory.create("h5") → H5Store +StoreFactory.create("bin") → MmapStore +StoreFactory.create("jsonl") → JsonlStore ``` **H5Store**: Reads HDF5 files, supports `share_memory_()` for multi-process DataLoader workers (copies tensors to shared memory). **MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. -Both 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. + +All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based indexing). ## Data Keys by Training Type @@ -106,4 +109,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-06-19 +> Document Update Time: 2026-07-05 diff --git a/assets/docs/inference.md b/assets/docs/inference.md index 71be84f..b8f49a8 100644 --- a/assets/docs/inference.md +++ b/assets/docs/inference.md @@ -23,27 +23,38 @@ RoPE is applied **before** KV cache write, not after — otherwise position enco ## KVCache System -Seven classes working together: +Seven classes working together, with two concrete cache implementations: + +### ContiguousCache (default) ``` -KVCache (facade) - ├── PagePool orchestrates page allocation + prefix matching - │ ├── Allocator bitmask-based page allocator + ref-count + LRU eviction (inside PagePool) - │ └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) (inside PagePool) - ├── TaskTable maps task_id → page_table + cached token count - ├── Storage k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim) - └── KvcacheView bundles Storage + page_table + total_len for attention layers (returned by bind()) +ContiguousCache (simple contiguous per-slot cache) + ├── ContiguousCacheView bundles k/v tensors + slot indices for attention layers ``` -`KVCache.bind(page_table, total_len)` returns a `KvcacheView` used by attention layers via `write()` / `gather()`. +Created by default when no cache is passed to `InferenceScheduler`. Each task occupies a fixed slot of `[max_seq_len, n_kv_heads, head_dim]`. Simple and efficient for small-to-medium batch sizes. + +### PageCache (paged with prefix sharing) + +``` +PageCache (paged KV cache with prefix sharing, alternative) + ├── PagePool orchestrates page allocation + prefix matching + │ ├── Allocator bitmask-based page allocator + ref-count + LRU + │ └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) + ├── TaskTable maps task_id → page_table + cached token count + ├── Storage k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim) + └── PageCacheView bundles Storage + page_table + total_len for attention layers +``` + +`isinstance(cache, KVCache)` checks dispatch to the correct view. Both implement the abstract `KVCache` interface used by `Executor` and `InferenceScheduler`. ## Continuous Batching `InferenceScheduler` runs a daemon thread with a 4-phase loop: ``` -1. Cleanup → Remove finished tasks, free KV pages -2. Refill → Pop from waiting_queue, task_alloc pages, activate +1. Cleanup → Remove finished tasks, free KV cache slots/pages +2. Refill → Pop from waiting_queue, task_alloc resources, activate 3. Prefill → Group by (prompt_len, start_pos), run full forward 4. Decode → Pick largest same-position group, single-token forward ``` @@ -238,4 +249,4 @@ async for token in engine.generate_async("Hello", ...): # -> AsyncGenerator[s print(token) ``` -> Document Update Time: 2026-06-19 +> Document Update Time: 2026-07-05 diff --git a/assets/docs/params.md b/assets/docs/params.md index 2654272..25494bc 100644 --- a/assets/docs/params.md +++ b/assets/docs/params.md @@ -67,7 +67,6 @@ | Parameter | Description | Default | |-----------|-------------|---------| | `--log_dir` | Directory for metric logs | checkpoint/logs | -| `--log_interval` | Number of optimizer steps between metric logs | 1 | | `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr", "grad_norm"] | ### Gradient Checkpointing @@ -200,4 +199,4 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples. --- -> Document Update Time: 2026-06-19 \ No newline at end of file +> Document Update Time: 2026-07-05 \ No newline at end of file diff --git a/assets/docs/preprocessing.md b/assets/docs/preprocessing.md index 15afd5c..bd8e786 100644 --- a/assets/docs/preprocessing.md +++ b/assets/docs/preprocessing.md @@ -361,4 +361,4 @@ Pipeline( ).run() ``` -> Document Update Time: 2026-06-03 +> Document Update Time: 2026-07-05 diff --git a/assets/docs/training.md b/assets/docs/training.md index dd160b3..d3afd9a 100644 --- a/assets/docs/training.md +++ b/assets/docs/training.md @@ -80,13 +80,13 @@ on_train_end | `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`, `MetricLoggerCallback`, `ValidationCallback` | -| `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` | -| `on_epoch_end` | End of each epoch | `ProgressBarCallback` | -| `on_error` | On exception during training | `CheckpointCallback`, `MetricLoggerCallback` | -| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricLoggerCallback`, `GradientCheckpointingCallback` | +| `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), `validation` (periodic validation on val_dataset), `metric_logger` (JSONL, 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`. ## Strategies @@ -108,7 +108,7 @@ $$ L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta) $$ -Keys: `input_ids`, `target_ids`, `loss_mask`. Optional: `label_smoothing`. +Keys: `input_ids`, `target_ids`, `loss_mask`, `position_ids`. Optional: `label_smoothing`. ### DPO (Direct Preference Optimization) @@ -214,4 +214,4 @@ nohup python scripts/tools/train.py \ Full parameter reference at [params.md](params.md). -> Document Update Time: 2026-05-30 +> Document Update Time: 2026-07-05