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:
parent
aabb0d83e9
commit
dc9faca3b1
|
|
@ -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
|
|||
<<abstract>>
|
||||
+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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue