diff --git a/assets/docs/architecture.md b/assets/docs/architecture.md index 2590e63..30caff4 100644 --- a/assets/docs/architecture.md +++ b/assets/docs/architecture.md @@ -21,6 +21,7 @@ classDiagram class BaseModelConfig { +Optional[str] model_type + +float neftune_alpha +from_file(config_path) Self +to_file(config_path) } @@ -58,10 +59,12 @@ classDiagram +Optional[int] dim_ffn +Optional[int] max_len +Optional[float] rope_theta + +str attn_type +Optional[int] n_heads +Optional[int] n_kv_heads +Optional[bool] use_qk_norm +Optional[bool] use_gated_attention + +str ffn_type +Optional[dict] rope_scaling +Optional[str] pooling_type +Optional[bool] normalize_embeddings @@ -118,7 +121,7 @@ classDiagram +float max_grad_norm +list gradient_checkpointing_modules +int start_epoch - +int start_batch + +int start_samples +str ckpt_dir +int ckpt_interval +str log_dir @@ -136,7 +139,9 @@ classDiagram +str start_method +str device_type +Optional[Dataset] val_dataset + +Optional[float] val_split +int val_step + +float neftune_alpha +str parallel_mode +dict executor_kwargs +dict extra_kwargs @@ -215,12 +220,13 @@ classDiagram class Checkpoint { +dict state_dict +int epoch - +int iteration + +int consumed_samples +dict extra +dict meta +dict config +save(save_dir) +load(save_dir, broadcast) Checkpoint + +load_any(save_dir, broadcast) Optional[Checkpoint] } } @@ -350,7 +356,9 @@ classDiagram class Embedding { +Parameter weight + +float neftune_noise_alpha +forward(x) Tensor + +set_neftune_alpha(alpha) } } @@ -407,7 +415,9 @@ classDiagram +Dict _entries +register(name) decorator +create(name, *args, **kwargs) T + +get_component_class(name) Type +list_registered() list + +is_registered(name) bool } class MaskBuilderFactory { @@ -436,13 +446,15 @@ classDiagram +dict model_config +BaseExecutor executor +int epoch - +int iteration + +int consumed_samples +float loss + +float grad_norm +DataLoader val_dataloader +float val_loss +int world_size +int rank +dict kwargs + +optimizer_step() int } class TrainContextBuilder { @@ -594,18 +606,6 @@ classDiagram +create(name, **kwargs) TrainCallback } - class Muon { - +float lr - +float momentum - +float weight_decay - +bool nesterov - +int ns_steps - +Optional[float] adamw_lr - +tuple adamw_betas - +float adamw_eps - +float adamw_wd - +step(closure) Optional[float] - } } namespace inference { @@ -810,7 +810,9 @@ classDiagram class ChatMessage { +str role - +str content + +Optional[str] content + +Optional[List[Dict]] tool_calls + +Optional[str] tool_call_id } class ChatCompletionRequest { @@ -827,6 +829,8 @@ classDiagram +Optional[float] frequency_penalty +Optional[Dict[int, float]] logit_bias +Optional[str] user + +Optional[List[ToolDef]] tools + +Optional[Union[str, Dict]] tool_choice } class AnthropicMessage { @@ -850,7 +854,7 @@ classDiagram <> +prepare(request, engine) Tuple[str, GenContext, List[str]] +format_stream_start(ctx) List[str] - +format_chunk(token) str + +format_chunk(token) List[str] +format_stream_end(ctx, stop) List[str] +format_response(ctx, content, stop) Dict } @@ -858,7 +862,7 @@ classDiagram class OpenAIResponseBuilder { +prepare(request, engine) Tuple +format_stream_start(ctx) List[str] - +format_chunk(token) str + +format_chunk(token) List[str] +format_stream_end(ctx, stop) List[str] +format_response(ctx, content, stop) Dict } @@ -866,7 +870,7 @@ classDiagram class AnthropicResponseBuilder { +prepare(request, engine) Tuple +format_stream_start(ctx) List[str] - +format_chunk(token) str + +format_chunk(token) List[str] +format_stream_end(ctx, stop) List[str] +format_response(ctx, content, stop) Dict } @@ -1171,10 +1175,10 @@ classDiagram | **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, Muon | Training workflow | +| **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.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** | Registry, BaseFactory[T] | Component registration | +| **astrai.factory** | BaseFactory | Component registration | | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers | ## Design Patterns diff --git a/assets/docs/dataflow.md b/assets/docs/dataflow.md index 08cef80..310018b 100644 --- a/assets/docs/dataflow.md +++ b/assets/docs/dataflow.md @@ -46,10 +46,10 @@ The output `meta.json` records the storage format, key names, dtype, total token ### Format Detection -`detect_format(load_path)` inspects the directory: +`detect_format(load_path)` inspects the path: -- If `*.h5` files exist → `"h5"` (HDF5 backend) -- If `*.bin` + `meta.json` files exist → `"bin"` (memory-mapped backend) +- 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"` ### Store Backends @@ -83,7 +83,7 @@ DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_typ → detect_format(load_path) → StoreFactory.create(storage_type) → Store.load(load_path) - → H5Store._normalize() / MmapStore._normalize() + → _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) diff --git a/assets/docs/inference.md b/assets/docs/inference.md index 764058b..71be84f 100644 --- a/assets/docs/inference.md +++ b/assets/docs/inference.md @@ -23,7 +23,7 @@ RoPE is applied **before** KV cache write, not after — otherwise position enco ## KVCache System -Six classes (plus two helpers) working together: +Seven classes working together: ``` KVCache (facade) @@ -152,12 +152,13 @@ Supports `stop_sequences` and streaming via `event: content_block_delta`. data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai", "choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]} -data: {"id":"chatcmpl-...","object":"chat.completion.chunk",..., +data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":0,"model":"astrai", "choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]} -data: {"id":"chatcmpl-...","object":"chat.completion.chunk",..., - "choices":[{"index":0,"delta":{},"finish_reason":"stop"}], - "usage":{"prompt_tokens":5,"completion_tokens":1,"total_tokens":6}} +data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai", + "choices":[{"index":0,"delta":{},"finish_reason":"stop"}]} + +data: {"prompt_tokens":5,"completion_tokens":1,"total_tokens":6} data: [DONE] ``` @@ -167,7 +168,7 @@ data: [DONE] ``` event: message_start data: {"type":"message_start","message":{"id":"msg_...","model":"astrai","role":"assistant", - "content":[],"stop_reason":null,...}} + "content":[],"usage":{"input_tokens":0}}} event: content_block_start data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} @@ -179,7 +180,7 @@ event: content_block_stop data: {"type":"content_block_stop","index":0} event: message_delta -data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{...}} +data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{...}} event: message_stop data: {"type":"message_stop"} @@ -187,26 +188,20 @@ data: {"type":"message_stop"} ### Error Responses -All endpoints use standard HTTP status codes: +The server returns standard HTTP status codes. Pydantic validation errors (e.g. missing required fields) +are handled automatically by FastAPI with 422 status. The only application-level error is engine initialization: | Status | Meaning | |--------|---------| | 200 | Success | -| 400 | Invalid request (bad JSON, missing fields, validation error) | -| 405 | Method not allowed | | 422 | Unprocessable entity (Pydantic validation) | -| 500 | Internal server error (model crash, OOM, scheduler failure) | | 503 | Service unavailable (model not loaded, engine not ready) | -Error response body: +Error response body (503): ```json { - "error": { - "message": "Invalid request: max_tokens must be > 0", - "type": "invalid_request_error", - "code": 400 - } + "detail": "Engine not initialized" } ``` @@ -220,16 +215,13 @@ Response: ```json { - "active_requests": 3, - "waiting_requests": 2, - "total_requests": 128, - "cache_usage": 0.45, - "tokens_generated": 10240 + "total_tasks": 128, + "total_tokens": 10240, + "active_tasks": 3, + "waiting_queue": 2 } ``` -`cache_usage` is the fraction of KV cache pages currently in use (0.0–1.0). - ## Engine API ```python diff --git a/assets/docs/params.md b/assets/docs/params.md index ee91ee5..2654272 100644 --- a/assets/docs/params.md +++ b/assets/docs/params.md @@ -53,7 +53,7 @@ | `--ckpt_interval` | Iterations between checkpoints | 5000 | | `--ckpt_dir` | Checkpoint save directory | checkpoint | | `--start_epoch` | Resume from epoch (0 = from scratch) | 0 | -| `--start_batch` | Resume from batch iteration | 0 | +| `--start_samples` | Resume from sample count per rank | 0 | ### Validation @@ -67,8 +67,8 @@ | Parameter | Description | Default | |-----------|-------------|---------| | `--log_dir` | Directory for metric logs | checkpoint/logs | -| `--log_interval` | Number of batch iterations between metric logs | 100 | -| `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr"] | +| `--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 @@ -100,6 +100,17 @@ | `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` | | `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` | +### Scheduler + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine | +| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default) | +| `--cycle_length` | SGDR first cycle length in steps | None (total_steps - warmup_steps) | +| `--t_mult` | SGDR cycle length multiplier per restart | 2 | +| `--stable_steps` | WSD stable plateau steps | None (required for wsd) | +| `--decay_steps` | WSD decay steps | None (total_steps - warmup_steps - stable_steps) | + ### Usage Example ```bash @@ -178,7 +189,7 @@ python scripts/tools/generate.py \ | `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) | | `--output_dir`, `-o` | path | required | Output directory for processed data | | `--config`, `-c` | path | required | Preprocessing pipeline config (JSON) | -| `--num_workers` | int | `4` | Number of parallel workers | +| `--tokenizer_path` | str | `params` | Path to tokenizer directory | Usage: ```bash diff --git a/assets/docs/preprocessing.md b/assets/docs/preprocessing.md index c049288..6246ea4 100644 --- a/assets/docs/preprocessing.md +++ b/assets/docs/preprocessing.md @@ -26,8 +26,9 @@ A single config file captures the entire pipeline, reusable and version-controll ```json { + "version": 1, "input": {}, // sections (single) or sources (multi) - "mask": {}, // role → "train" | "mask" + "mask": {}, // role -> "train" | "mask" "mask_default": "mask", "preprocessing": {}, "output": {} @@ -220,11 +221,12 @@ Config: } ``` -Output keys: `prompts`, `responses`, `masks`, `rewards` (float32) +Output keys: `prompts`, `prompts_mask`, `responses`, `masks`, `rewards` (float32) - `action: "value"` — extract raw values from JSONL without tokenisation - `list_field: true` — tokenise each list element independently, then concatenate - `mask_key: "masks"` — rename the auto-generated mask key (default: `responses_mask`) +- `prompts_mask` is auto-generated (all masked) and unused by GRPOStrategy --- @@ -274,12 +276,11 @@ When `sources` is set, `sections` is ignored. ### Template mode (`template: true`) -For each message in the field's array: - 1. Prepend BOS token (masked) -2. Render through `chat_template` for that single message -3. Encode rendered text -4. Apply mask rule for the message's role +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 @@ -287,7 +288,7 @@ Encode the field value as text. Mask value is 1 (train) or 0 (mask) per the sect ### Text config detection -When no section uses `template` and all sections have `action: "train"`, the builder skips mask generation entirely — all tokens are trained. +When no section uses `template` and all sections have `action: "train"`, the builder omits `loss_mask` from the output — all tokens are trained. --- @@ -298,13 +299,15 @@ When no section uses `template` and all sections have `action: "train"`, the bui ``` output/ __default__/ - meta.json - sequence.bin - loss_mask.bin + shard_0000/ + meta.json + sequence.bin + loss_mask.bin wiki/ - meta.json - sequence.bin - loss_mask.bin + shard_0000/ + meta.json + sequence.bin + loss_mask.bin ``` ### Multi-Shard (`bin`) @@ -324,7 +327,7 @@ output/ loss_mask.bin ``` -`MmapStore` discovers all shards under the domain directory via `rglob("meta.json")`. +For `bin` format, `MmapStore` discovers all shards under the domain directory via `rglob("meta.json")`. For `h5` format, `H5Store` discovers `.h5`/`.hdf5` files via recursive glob. --- @@ -349,7 +352,7 @@ python scripts/tools/preprocess.py data/grpo/*.jsonl -o output/grpo/ -c configs/ from astrai.preprocessing.pipeline import Pipeline from astrai.config.preprocess_config import PipelineConfig -config = PipelineConfig.from_json("sft.json") +config = PipelineConfig.from_file("sft.json") Pipeline( config, ["data_part1.jsonl", "data_part2.jsonl"], diff --git a/assets/docs/training.md b/assets/docs/training.md index 443aaec..dd160b3 100644 --- a/assets/docs/training.md +++ b/assets/docs/training.md @@ -58,7 +58,9 @@ on_train_begin context.loss = loss.item() stand_loss = loss / executor.grad_accum_steps executor.backward(stand_loss) - context.iteration += 1 + context.consumed_samples += ( + context.config.batch_per_device * context.world_size + ) on_batch_end if executor.sync_gradients: @@ -78,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`, `ValidationCallback` | +| `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` | -Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric_logger` (JSONL, rank-0), `progress_bar` (tqdm), `gradient_clipping`, `validation` (periodic validation on val_dataset). +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`. ## Strategies @@ -158,8 +160,8 @@ Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoi ## Checkpoint ``` -Checkpoint(state_dict, epoch, iteration, extra, meta, config) - ├── save(save_dir) rank-0 only: meta.json (epoch/iteration/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt) +Checkpoint(state_dict, epoch, consumed_samples, extra, meta, config) + ├── save(save_dir) rank-0 only: meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt) └── load(save_dir, broadcast=False) loads from local disk; set broadcast=True to broadcast metadata from rank-0 ```