ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: align docs with actual code (40+ inconsistencies) 

- Remove nonexistent Muon class from architecture diagram
- Fix Checkpoint/TrainConfig/TrainContext field names (iteration -> consumed_samples, start_batch -> start_samples)
- Add missing fields: neftune_alpha, val_split, grad_norm, optimizer_step, tool_calls/tools
- Fix CLI param defaults: --log_interval 1, --metrics [loss,lr,grad_norm], --start_samples
- Add missing scheduler CLI params; remove nonexistent --num_workers from preprocess docs
- Fix inference SSE format, stats response keys, error codes to match actual server output
- Fix preprocessing docs: BOS once, shard_0000 layout, from_json->from_file, GRPO prompts_mask
- Fix dataflow detect_format/_normalize descriptions; correct callback order in training.md
This commit is contained in:
ViperEkura 2026-06-30 20:47:23 +08:00
parent aabb0d83e9
commit dc9faca3b1
6 changed files with 86 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

46
assets/docs/architecture.md
View File

@ -21,6 +21,7 @@ classDiagram
class BaseModelConfig { class BaseModelConfig {
+Optional[str] model_type +Optional[str] model_type
+float neftune_alpha
+from_file(config_path) Self +from_file(config_path) Self
+to_file(config_path) +to_file(config_path)
} }
@ -58,10 +59,12 @@ classDiagram
+Optional[int] dim_ffn +Optional[int] dim_ffn
+Optional[int] max_len +Optional[int] max_len
+Optional[float] rope_theta +Optional[float] rope_theta
+str attn_type
+Optional[int] n_heads +Optional[int] n_heads
+Optional[int] n_kv_heads +Optional[int] n_kv_heads
+Optional[bool] use_qk_norm +Optional[bool] use_qk_norm
+Optional[bool] use_gated_attention +Optional[bool] use_gated_attention
+str ffn_type
+Optional[dict] rope_scaling +Optional[dict] rope_scaling
+Optional[str] pooling_type +Optional[str] pooling_type
+Optional[bool] normalize_embeddings +Optional[bool] normalize_embeddings
@ -118,7 +121,7 @@ classDiagram
+float max_grad_norm +float max_grad_norm
+list gradient_checkpointing_modules +list gradient_checkpointing_modules
+int start_epoch +int start_epoch
+int start_batch +int start_samples
+str ckpt_dir +str ckpt_dir
+int ckpt_interval +int ckpt_interval
+str log_dir +str log_dir
@ -136,7 +139,9 @@ classDiagram
+str start_method +str start_method
+str device_type +str device_type
+Optional[Dataset] val_dataset +Optional[Dataset] val_dataset
+Optional[float] val_split
+int val_step +int val_step
+float neftune_alpha
+str parallel_mode +str parallel_mode
+dict executor_kwargs +dict executor_kwargs
+dict extra_kwargs +dict extra_kwargs
@ -215,12 +220,13 @@ classDiagram
class Checkpoint { class Checkpoint {
+dict state_dict +dict state_dict
+int epoch +int epoch
+int iteration +int consumed_samples
+dict extra +dict extra
+dict meta +dict meta
+dict config +dict config
+save(save_dir) +save(save_dir)
+load(save_dir, broadcast) Checkpoint +load(save_dir, broadcast) Checkpoint
+load_any(save_dir, broadcast) Optional[Checkpoint]
} }
} }
@ -350,7 +356,9 @@ classDiagram
class Embedding { class Embedding {
+Parameter weight +Parameter weight
+float neftune_noise_alpha
+forward(x) Tensor +forward(x) Tensor
+set_neftune_alpha(alpha)
} }
} }
@ -407,7 +415,9 @@ classDiagram
+Dict _entries +Dict _entries
+register(name) decorator +register(name) decorator
+create(name, *args, **kwargs) T +create(name, *args, **kwargs) T
+get_component_class(name) Type
+list_registered() list +list_registered() list
+is_registered(name) bool
} }
class MaskBuilderFactory { class MaskBuilderFactory {
@ -436,13 +446,15 @@ classDiagram
+dict model_config +dict model_config
+BaseExecutor executor +BaseExecutor executor
+int epoch +int epoch
+int iteration +int consumed_samples
+float loss +float loss
+float grad_norm
+DataLoader val_dataloader +DataLoader val_dataloader
+float val_loss +float val_loss
+int world_size +int world_size
+int rank +int rank
+dict kwargs +dict kwargs
+optimizer_step() int
} }
class TrainContextBuilder { class TrainContextBuilder {
@ -594,18 +606,6 @@ classDiagram
+create(name, **kwargs) TrainCallback +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 { namespace inference {
@ -810,7 +810,9 @@ classDiagram
class ChatMessage { class ChatMessage {
+str role +str role
+str content +Optional[str] content
+Optional[List[Dict]] tool_calls
+Optional[str] tool_call_id
} }
class ChatCompletionRequest { class ChatCompletionRequest {
@ -827,6 +829,8 @@ classDiagram
+Optional[float] frequency_penalty +Optional[float] frequency_penalty
+Optional[Dict[int, float]] logit_bias +Optional[Dict[int, float]] logit_bias
+Optional[str] user +Optional[str] user
+Optional[List[ToolDef]] tools
+Optional[Union[str, Dict]] tool_choice
} }
class AnthropicMessage { class AnthropicMessage {
@ -850,7 +854,7 @@ classDiagram
<<abstract>> <<abstract>>
+prepare(request, engine) Tuple[str, GenContext, List[str]] +prepare(request, engine) Tuple[str, GenContext, List[str]]
+format_stream_start(ctx) 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_stream_end(ctx, stop) List[str]
+format_response(ctx, content, stop) Dict +format_response(ctx, content, stop) Dict
} }
@ -858,7 +862,7 @@ classDiagram
class OpenAIResponseBuilder { class OpenAIResponseBuilder {
+prepare(request, engine) Tuple +prepare(request, engine) Tuple
+format_stream_start(ctx) 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_stream_end(ctx, stop) List[str]
+format_response(ctx, content, stop) Dict +format_response(ctx, content, stop) Dict
} }
@ -866,7 +870,7 @@ classDiagram
class AnthropicResponseBuilder { class AnthropicResponseBuilder {
+prepare(request, engine) Tuple +prepare(request, engine) Tuple
+format_stream_start(ctx) 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_stream_end(ctx, stop) List[str]
+format_response(ctx, content, stop) Dict +format_response(ctx, content, stop) Dict
} }
@ -1171,10 +1175,10 @@ classDiagram
| **astrai.serialization** | Checkpoint | Model serialization | | **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.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.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.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.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 | | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers |
## Design Patterns ## Design Patterns

8
assets/docs/dataflow.md
View File

@ -46,10 +46,10 @@ The output `meta.json` records the storage format, key names, dtype, total token
### Format Detection ### Format Detection
`detect_format(load_path)` inspects the directory: `detect_format(load_path)` inspects the path:
- If `*.h5` files exist → `"h5"` (HDF5 backend) - If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, unknown suffix raises `ValueError`
- If `*.bin` + `meta.json` files exist → `"bin"` (memory-mapped backend) - If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, or `*.bin` + `**/meta.json` → `"bin"`
### Store Backends ### Store Backends
@ -83,7 +83,7 @@ DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_typ
→ detect_format(load_path) → detect_format(load_path)
→ StoreFactory.create(storage_type) → StoreFactory.create(storage_type)
→ Store.load(load_path) → 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]]] → Store._data[Dict[str, List[Tensor]]] + _cum[Dict[str, List[int]]]
→ BaseDataset.__getitem__(idx) → BaseDataset.__getitem__(idx)
→ get_index(idx) → [begin, end) → get_index(idx) → [begin, end)

40
assets/docs/inference.md
View File

@ -23,7 +23,7 @@ RoPE is applied **before** KV cache write, not after — otherwise position enco
## KVCache System ## KVCache System
Six classes (plus two helpers) working together: Seven classes working together:
``` ```
KVCache (facade) 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", data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai",
"choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]} "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}]} "choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
data: {"id":"chatcmpl-...","object":"chat.completion.chunk",..., data: {"id":"chatcmpl-...","object":"chat.completion.chunk","created":...,"model":"astrai",
"choices":[{"index":0,"delta":{},"finish_reason":"stop"}], "choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
"usage":{"prompt_tokens":5,"completion_tokens":1,"total_tokens":6}}
data: {"prompt_tokens":5,"completion_tokens":1,"total_tokens":6}
data: [DONE] data: [DONE]
``` ```
@ -167,7 +168,7 @@ data: [DONE]
``` ```
event: message_start event: message_start
data: {"type":"message_start","message":{"id":"msg_...","model":"astrai","role":"assistant", data: {"type":"message_start","message":{"id":"msg_...","model":"astrai","role":"assistant",
"content":[],"stop_reason":null,...}} "content":[],"usage":{"input_tokens":0}}}
event: content_block_start event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} 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} data: {"type":"content_block_stop","index":0}
event: message_delta 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 event: message_stop
data: {"type":"message_stop"} data: {"type":"message_stop"}
@ -187,26 +188,20 @@ data: {"type":"message_stop"}
### Error Responses ### 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 | | Status | Meaning |
|--------|---------| |--------|---------|
| 200 | Success | | 200 | Success |
| 400 | Invalid request (bad JSON, missing fields, validation error) |
| 405 | Method not allowed |
| 422 | Unprocessable entity (Pydantic validation) | | 422 | Unprocessable entity (Pydantic validation) |
| 500 | Internal server error (model crash, OOM, scheduler failure) |
| 503 | Service unavailable (model not loaded, engine not ready) | | 503 | Service unavailable (model not loaded, engine not ready) |
Error response body: Error response body (503):
```json ```json
{ {
"error": { "detail": "Engine not initialized"
"message": "Invalid request: max_tokens must be > 0",
"type": "invalid_request_error",
"code": 400
}
} }
``` ```
@ -220,16 +215,13 @@ Response:
```json ```json
{ {
"active_requests": 3, "total_tasks": 128,
"waiting_requests": 2, "total_tokens": 10240,
"total_requests": 128, "active_tasks": 3,
"cache_usage": 0.45, "waiting_queue": 2
"tokens_generated": 10240
} }
``` ```
`cache_usage` is the fraction of KV cache pages currently in use (0.0–1.0).
## Engine API ## Engine API
```python ```python

19
assets/docs/params.md
View File

@ -53,7 +53,7 @@
| `--ckpt_interval` | Iterations between checkpoints | 5000 | | `--ckpt_interval` | Iterations between checkpoints | 5000 |
| `--ckpt_dir` | Checkpoint save directory | checkpoint | | `--ckpt_dir` | Checkpoint save directory | checkpoint |
| `--start_epoch` | Resume from epoch (0 = from scratch) | 0 | | `--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 ### Validation
@ -67,8 +67,8 @@
| Parameter | Description | Default | | Parameter | Description | Default |
|-----------|-------------|---------| |-----------|-------------|---------|
| `--log_dir` | Directory for metric logs | checkpoint/logs | | `--log_dir` | Directory for metric logs | checkpoint/logs |
| `--log_interval` | Number of batch iterations between metric logs | 100 | | `--log_interval` | Number of optimizer steps between metric logs | 1 |
| `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr"] | | `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr", "grad_norm"] |
### Gradient Checkpointing ### Gradient Checkpointing
@ -100,6 +100,17 @@
| `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` | | `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` |
| `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` | | `--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 ### Usage Example
```bash ```bash
@ -178,7 +189,7 @@ python scripts/tools/generate.py \
| `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) | | `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) |
| `--output_dir`, `-o` | path | required | Output directory for processed data | | `--output_dir`, `-o` | path | required | Output directory for processed data |
| `--config`, `-c` | path | required | Preprocessing pipeline config (JSON) | | `--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: Usage:
```bash ```bash

35
assets/docs/preprocessing.md
View File

@ -26,8 +26,9 @@ A single config file captures the entire pipeline, reusable and version-controll
```json ```json
{ {
"version": 1,
"input": {}, // sections (single) or sources (multi) "input": {}, // sections (single) or sources (multi)
"mask": {}, // role → "train" | "mask" "mask": {}, // role -> "train" | "mask"
"mask_default": "mask", "mask_default": "mask",
"preprocessing": {}, "preprocessing": {},
"output": {} "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 - `action: "value"` — extract raw values from JSONL without tokenisation
- `list_field: true` — tokenise each list element independently, then concatenate - `list_field: true` — tokenise each list element independently, then concatenate
- `mask_key: "masks"` — rename the auto-generated mask key (default: `responses_mask`) - `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`) ### Template mode (`template: true`)
For each message in the field's array:
1. Prepend BOS token (masked) 1. Prepend BOS token (masked)
2. Render through `chat_template` for that single message 2. For each message in the field's array:
3. Encode rendered text 1. Render through `chat_template` for that single message
4. Apply mask rule for the message's role 2. Encode rendered text
3. Apply mask rule for the message's role
### Non-template mode ### 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 ### 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/ output/
__default__/ __default__/
meta.json shard_0000/
sequence.bin meta.json
loss_mask.bin sequence.bin
loss_mask.bin
wiki/ wiki/
meta.json shard_0000/
sequence.bin meta.json
loss_mask.bin sequence.bin
loss_mask.bin
``` ```
### Multi-Shard (`bin`) ### Multi-Shard (`bin`)
@ -324,7 +327,7 @@ output/
loss_mask.bin 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.preprocessing.pipeline import Pipeline
from astrai.config.preprocess_config import PipelineConfig from astrai.config.preprocess_config import PipelineConfig
config = PipelineConfig.from_json("sft.json") config = PipelineConfig.from_file("sft.json")
Pipeline( Pipeline(
config, config,
["data_part1.jsonl", "data_part2.jsonl"], ["data_part1.jsonl", "data_part2.jsonl"],

12
assets/docs/training.md
View File

@ -58,7 +58,9 @@ on_train_begin
context.loss = loss.item() context.loss = loss.item()
stand_loss = loss / executor.grad_accum_steps stand_loss = loss / executor.grad_accum_steps
executor.backward(stand_loss) executor.backward(stand_loss)
context.iteration += 1 context.consumed_samples += (
context.config.batch_per_device * context.world_size
)
on_batch_end on_batch_end
if executor.sync_gradients: if executor.sync_gradients:
@ -78,13 +80,13 @@ on_train_end
| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` | | `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
| `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` | | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
| `on_batch_begin` | Every batch | — | | `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_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
| `on_epoch_end` | End of each epoch | `ProgressBarCallback` | | `on_epoch_end` | End of each epoch | `ProgressBarCallback` |
| `on_error` | On exception during training | `CheckpointCallback`, `MetricLoggerCallback` | | `on_error` | On exception during training | `CheckpointCallback`, `MetricLoggerCallback` |
| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricLoggerCallback`, `GradientCheckpointingCallback` | | `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 ## Strategies
@ -158,8 +160,8 @@ Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoi
## Checkpoint ## Checkpoint
``` ```
Checkpoint(state_dict, epoch, iteration, extra, meta, config) Checkpoint(state_dict, epoch, consumed_samples, 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) ├── 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 └── load(save_dir, broadcast=False) loads from local disk; set broadcast=True to broadcast metadata from rank-0
``` ```