fix: 修复 to_dict list 类型丢失与 OpenAI stop 参数失效

- to_dict() 增加 list 类型序列化支持，metrics 等字段不再丢失
- OpenAIHandler 补充 get_stop_sequences/on_token，读取 request.stop 并检测停止序列
- 文档类图补充缺失字段、修正关系分类、ChatCompletionRequest 字段增加 Optional

Dockerfile

@@ -1,7 +1,7 @@
 # AstrAI Dockerfile - Multi-stage Build (Optimized)
 
 # Build stage - use base image with minimal build tools
-FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS builder
+FROM ubuntu:24.04 AS builder
 
 WORKDIR /app
 
@@ -18,7 +18,7 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
 RUN python3.12 -m venv --copies /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
 
-# Copy source code and install dependencies
+# Copy source code and install (deps read from pyproject.toml)
 COPY astrai/ ./astrai/
 COPY pyproject.toml .
 RUN pip install --no-cache-dir --upgrade pip \
@@ -26,13 +26,14 @@ RUN pip install --no-cache-dir --upgrade pip \
     --extra-index-url https://download.pytorch.org/whl/cu126
 
 # Production stage
-FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production
+FROM ubuntu:24.04 AS production
 
 WORKDIR /app
 
-# Install Python 3.12 runtime
+# Install Python 3.12 runtime and healthcheck dependency
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     python3.12 \
+    curl \
     && rm -rf /var/lib/apt/lists/*
 
 # Copy virtual environment from builder

README.md

@@ -78,15 +78,27 @@ Or download manually from [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) i
 #### Train a Model
 
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
-    --train_type seq \
-    --data_root_path /path/to/dataset \
-    --param_path /path/to/model \
-    --batch_size 4 \
-    --accumulation_steps 8 \
-    --max_lr 3e-4 \
-    --warmup_steps 1000 \
-    --n_epoch 1
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
 ```
 
 Full reference at [Parameter Guide](assets/docs/params.md).
@@ -201,7 +213,7 @@ python scripts/demo/generate_batch.py
 python scripts/demo/generate_ar.py
 ```
 
-Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd).
+Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).
 
 ### Documentation
 

assets/docs/README-zh-CN.md

@@ -84,15 +84,27 @@ python scripts/demo/download.py
 #### 训练模型
 
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
-    --train_type seq \
-    --data_root_path /path/to/dataset \
-    --param_path /path/to/model \
-    --batch_size 4 \
-    --accumulation_steps 8 \
-    --max_lr 3e-4 \
-    --warmup_steps 1000 \
-    --n_epoch 1
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
 ```
 
 完整参数列表见[参数说明](./params.md)。
@@ -207,7 +219,7 @@ python scripts/demo/generate_batch.py
 python scripts/demo/generate_ar.py
 ```
 
-观看 [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd) 上的视频演示。
+观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
 
 ### 文档
 

assets/docs/architecture.md

@@ -5,13 +5,18 @@
 ```mermaid
 classDiagram
     namespace config {
-        class BaseModelConfig {
-            +Optional[str] model_type
-            +load(config_path) Self
-            +save(config_path)
+        class BaseConfig {
+            +to_dict() Dict
+            +from_dict(d) Self
         }
 
-        class ModelConfig {
+        class BaseModelConfig {
+            +Optional[str] model_type
+            +from_file(config_path) Self
+            +to_file(config_path)
+        }
+
+        class AutoRegressiveLMConfig {
             +int vocab_size
             +int dim
             +int n_layers
@@ -20,18 +25,41 @@ classDiagram
             +bool tie_weight
             +int max_len
             +float rope_theta
+            +str attn_type
             +int n_heads
             +int n_kv_heads
             +bool use_qk_norm
             +bool use_gated_attention
-            +str attn_type
+            +Optional[int] kv_lora_rank
+            +Optional[int] qk_nope_head_dim
+            +Optional[int] qk_rope_head_dim
             +str ffn_type
             +int n_routed_experts
             +int n_shared_experts
             +int n_activated_experts
-            +str moe_topk_method
-            +load(config_path) ModelConfig
-            +save(config_path)
+            +Optional[str] topk_method
+        }
+
+        class EncoderConfig {
+            +int vocab_size
+            +int dim
+            +int n_layers
+            +float norm_eps
+            +int dim_ffn
+            +int max_len
+            +float rope_theta
+            +int n_heads
+            +int n_kv_heads
+            +bool use_qk_norm
+            +bool use_gated_attention
+            +Optional[str] pooling_type
+            +Optional[bool] normalize_embeddings
+        }
+
+        class ConfigFactory {
+            +Registry _registry
+            +register(name) decorator
+            +load(raw) BaseConfig
         }
 
         class TrainConfig {
@@ -41,13 +69,17 @@ classDiagram
             +Callable optimizer_fn
             +Callable scheduler_fn
             +int n_epoch
-            +int batch_size
-            +int accumulation_steps
+            +int batch_per_device
+            +int grad_accum_steps
             +float max_grad_norm
+            +list gradient_checkpointing_modules
             +int start_epoch
             +int start_batch
             +str ckpt_dir
             +int ckpt_interval
+            +str log_dir
+            +int log_interval
+            +List[str] metrics
             +int random_seed
             +int num_workers
             +Optional[int] prefetch_factor
@@ -58,7 +90,10 @@ classDiagram
             +str master_port
             +Callable parallel_wrapper
             +Callable state_dict_fn
+            +str start_method
             +str device_type
+            +Optional[Dataset] val_dataset
+            +int val_step
             +dict extra_kwargs
             +validate()
         }
@@ -69,7 +104,7 @@ classDiagram
         class BaseDataset {
             +int window_size
             +int stride
-            +BaseStorage storage
+            +Optional[BaseStorage] storage
             +load(load_path, storage_type, tokenizer)
             +__getitem__(index)
             +__len__()
@@ -126,15 +161,21 @@ classDiagram
         }
 
         class ResumableDistributedSampler {
-            +int start_epoch
-            +int start_iter
+            +int epoch
+            +int iter
+        }
+
+        class StorageFactory {
+            +Registry _registry
+            +register(name) decorator
+            +create(storage_type) BaseStorage
         }
 
         class DatasetFactory {
             +Registry _registry
             +register(name) decorator
             +create(train_type, window_size, stride) BaseDataset
-            +load(train_type, load_path, window_size, stride) BaseDataset
+            +load(train_type, load_path, window_size, stride, storage_type, tokenizer) BaseDataset
         }
     }
 
@@ -144,6 +185,7 @@ classDiagram
             +int epoch
             +int iteration
             +dict extra
+            +dict meta
             +save(save_dir)
             +load(save_dir) Checkpoint
         }
@@ -151,27 +193,39 @@ classDiagram
 
     namespace model {
         class AutoModel {
-            +ModelConfig config
+            +BaseModelConfig config
             +Registry _registry
             +register(model_type) decorator
             +get_component_class(model_type) Type
-            +from_pretrained(path, disable_random_init) nn.Module
+            +from_pretrained(path, disable_random_init, strict) nn.Module
             +save_pretrained(save_directory)
             +to(*args, **kwargs) Self
         }
 
-        class Transformer {
-            +ModelConfig config
+        class AutoRegressiveLM {
+            +AutoRegressiveLMConfig config
             +RotaryEmbedding rotary_embedding
             +Embedding embed_tokens
             +ModuleList layers
             +RMSNorm norm
             +Linear lm_head
-            +forward(input_ids, input_mask, paged_cache, position_ids) Dict
+            +forward(input_ids, input_mask, paged_cache, position_ids) Dict[str, Tensor]
             +load_state_dict(state_dict)
             +state_dict()
         }
 
+        class EmbeddingEncoder {
+            +EncoderConfig config
+            +RotaryEmbedding rotary_embedding
+            +Embedding embed_tokens
+            +ModuleList layers
+            +RMSNorm norm
+            +str pooling_type
+            +bool normalize_embeddings
+            +forward(input_ids, input_mask, position_ids) Tensor
+            +load_state_dict(state_dict)
+        }
+
         class DecoderBlock {
             +nn.Module attention  # GQA or MLA via AttnFactory
             +RMSNorm input_norm
@@ -185,6 +239,7 @@ classDiagram
             +int n_kv_heads
             +int head_dim
             +int n_rep
+            +int layer_id
             +bool use_qk_norm
             +bool use_gated_attention
             +Linear q_proj, k_proj, v_proj, o_proj
@@ -201,6 +256,7 @@ classDiagram
             +int qk_nope_head_dim
             +int qk_rope_head_dim
             +int n_rep
+            +int layer_id
             +bool use_gated_attention
             +Linear q_proj, kv_a_proj, kv_b_proj
             +Linear o_proj
@@ -215,6 +271,7 @@ classDiagram
         }
 
         class DeepSeekMoE {
+            +int dim
             +int n_routed_experts
             +int n_shared_experts
             +int n_activated_experts
@@ -236,6 +293,7 @@ classDiagram
         class RMSNorm {
             +Parameter weight
             +float norm_eps
+            +tuple normalized_shape
             +forward(x) Tensor
         }
 
@@ -299,7 +357,6 @@ classDiagram
             +TrainConfig train_config
             +List[TrainCallback] callbacks
             +train(checkpoint)
-            +_build_context(checkpoint) TrainContext
             +_get_default_callbacks() List[TrainCallback]
         }
 
@@ -310,11 +367,15 @@ classDiagram
             +Optimizer optimizer
             +LRScheduler scheduler
             +Checkpoint checkpoint
+            +TrainConfig config
             +int epoch
             +int iteration
             +float loss
+            +DataLoader val_dataloader
+            +float val_loss
             +int world_size
             +int rank
+            +dict kwargs
         }
 
         class TrainContextBuilder {
@@ -324,7 +385,7 @@ classDiagram
         }
 
         class BaseStrategy {
-            +nn.Module model
+            +Union[Callable, nn.Module] model
             +str device
             +compute_loss(batch) Tensor
         }
@@ -332,7 +393,7 @@ classDiagram
         class StrategyFactory {
             +Registry _registry
             +register(name) decorator
-            +create(model, train_type, device, **kwargs) BaseStrategy
+            +create(train_type, model, device, **kwargs) BaseStrategy
         }
 
         class SEQStrategy {
@@ -360,6 +421,7 @@ classDiagram
             +str reduction
             +int sync_interval
             +compute_loss(batch) Tensor
+            +sync_ref_model()
         }
 
         class BaseScheduler {
@@ -387,6 +449,7 @@ classDiagram
         }
 
         class TrainCallback {
+            <<protocol>>
             +on_train_begin(context)
             +on_train_end(context)
             +on_epoch_begin(context)
@@ -400,20 +463,35 @@ classDiagram
 
         class GradientClippingCallback {
             +float max_grad_norm
-            +on_step_end(context)
+            +on_step_begin(context)
+        }
+
+        class GradientCheckpointingCallback {
+            +tuple modules
+            +on_train_begin(context)
+            +on_train_end(context)
         }
 
         class CheckpointCallback {
             +str save_dir
             +int interval
+            +bool weight_only
+            +Callable state_dict_fn
+            +Callable save_extra_fn
+            +Callable load_extra_fn
             +_save_checkpoint(context)
+            +on_train_begin(context)
             +on_batch_end(context)
             +on_train_end(context)
             +on_error(context)
+            +save_extra(context)$
+            +load_extra(extra, context)$
         }
 
         class ProgressBarCallback {
             +int num_epoch
+            +int log_interval
+            +IO file
             +on_epoch_begin(context)
             +on_batch_end(context)
             +on_epoch_end(context)
@@ -422,8 +500,16 @@ classDiagram
         class MetricLoggerCallback {
             +str log_dir
             +int save_interval
+            +int log_interval
+            +List[str] metrics
             +on_batch_end(context)
             +on_train_end(context)
+            +on_error(context)
+        }
+
+        class ValidationCallback {
+            +_run_validation(context)
+            +on_step_end(context)
         }
 
         class CallbackFactory {
@@ -431,6 +517,14 @@ classDiagram
             +register(name) decorator
             +create(name, **kwargs) TrainCallback
         }
+
+        class Muon {
+            +float lr
+            +float momentum
+            +float weight_decay
+            +int ns_steps
+            +step(closure) Optional[float]
+        }
     }
 
     namespace inference {
@@ -459,10 +553,7 @@ classDiagram
             +TaskManager _task_mgr
             +bool _running
             +Thread _loop_thread
-            +int max_batch_size
             +int max_seq_len
-            +int max_prompt_len
-            +int page_size
             +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
             +remove_task(task_id)
             +start()
@@ -500,10 +591,7 @@ classDiagram
         }
 
         class Storage {
-            +int n_layers
             +int page_size
-            +int head_dim
-            +int n_kv_heads
             +Tensor k_cache
             +Tensor v_cache
             +write(layer_id, page_table, start_pos, k, v)
@@ -610,7 +698,7 @@ classDiagram
         }
 
         class SamplingPipeline {
-            +List strategies
+            +List[BaseSamplingStrategy] strategies
             +apply(logits, filter_value) Tensor
             +sample(logits, filter_value) Tensor
         }
@@ -632,14 +720,19 @@ classDiagram
         }
 
         class ChatCompletionRequest {
+            +str model
             +List[ChatMessage] messages
-            +float temperature
-            +float top_p
-            +int top_k
-            +int max_tokens
-            +bool stream
-            +Optional[str] stop
+            +Optional[float] temperature
+            +Optional[float] top_p
+            +Optional[int] top_k
+            +Optional[int] max_tokens
+            +Optional[bool] stream
+            +Optional[Union[str, List[str]]] stop
             +Optional[int] n
+            +Optional[float] presence_penalty
+            +Optional[float] frequency_penalty
+            +Optional[Dict[int, float]] logit_bias
+            +Optional[str] user
         }
 
         class AnthropicMessage {
@@ -648,6 +741,7 @@ classDiagram
         }
 
         class MessagesRequest {
+            +str model
             +List[AnthropicMessage] messages
             +Optional[str] system
             +float temperature
@@ -660,8 +754,13 @@ classDiagram
 
         class ProtocolHandler {
             <<abstract>>
+            +request
+            +engine
             +build_prompt() str
             +create_response_id() str
+            +get_stop_sequences() List[str]
+            +create_stop_checker() StopChecker
+            +on_token(ctx, token, stop_checker) Optional[str]
             +format_stream_start(ctx) List[str]
             +format_stream_token(ctx, token) str
             +format_stream_end(ctx) List[str]
@@ -675,13 +774,13 @@ classDiagram
         }
 
         class AnthropicHandler {
-            +List[str] stop_sequences
             +build_prompt() str
             +create_response_id() str
             +on_token(ctx, token, stop_checker) Optional[str]
         }
 
         class StopChecker {
+            +has_sequences (property) bool
             +check(text) Optional[str]
             +trim(text, matched) str
         }
@@ -694,6 +793,7 @@ classDiagram
             +int completion_tokens
             +str accumulated
             +Optional[str] stop_matched
+            +str last_yield_trimmed
         }
 
         class app {
@@ -704,11 +804,13 @@ classDiagram
 
     namespace parallel {
         class Functions {
-            +spawn_parallel_fn(fn, nprocs)
+            <<module>>
+            +spawn_parallel_fn(func, world_size, backend, master_addr, master_port, device_type, start_method, **kwargs)
             +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type)
             +get_current_device() str
             +get_world_size() int
             +get_rank() int
+            +only_on_rank(rank, sync) decorator
         }
 
         class ParallelModel {
@@ -736,6 +838,7 @@ classDiagram
     BaseScheduler <|-- CosineScheduler
     BaseScheduler <|-- SGDRScheduler
     TrainCallback <|-- GradientClippingCallback
+    TrainCallback <|-- GradientCheckpointingCallback
     TrainCallback <|-- CheckpointCallback
     TrainCallback <|-- ProgressBarCallback
     TrainCallback <|-- MetricLoggerCallback
@@ -748,10 +851,15 @@ classDiagram
     BaseSamplingStrategy <|-- TemperatureStrategy
     BaseSamplingStrategy <|-- TopKStrategy
     BaseSamplingStrategy <|-- TopPStrategy
+    BaseSamplingStrategy <|-- SamplingPipeline
     ParallelModel <|-- RowParallelLinear
     ParallelModel <|-- ColumnParallelLinear
-    AutoModel <|-- Transformer
-    BaseModelConfig <|-- ModelConfig
+    AutoModel <|-- AutoRegressiveLM
+    AutoModel <|-- EmbeddingEncoder
+    BaseConfig <|-- BaseModelConfig
+    BaseConfig <|-- TrainConfig
+    BaseModelConfig <|-- AutoRegressiveLMConfig
+    BaseModelConfig <|-- EncoderConfig
     BaseFactory <|-- AutoModel
     BaseFactory <|-- AttnFactory
     BaseFactory <|-- FFNFactory
@@ -759,6 +867,9 @@ classDiagram
     BaseFactory <|-- StrategyFactory
     BaseFactory <|-- SchedulerFactory
     BaseFactory <|-- CallbackFactory
+    BaseFactory <|-- StorageFactory
+    BaseFactory <|-- ConfigFactory
+    TrainCallback <|-- ValidationCallback
     ProtocolHandler <|-- OpenAIHandler
     ProtocolHandler <|-- AnthropicHandler
 
@@ -766,31 +877,33 @@ classDiagram
     KVCache *-- PagePool
     KVCache *-- Storage
     KVCache *-- TaskTable
-    KVCache *-- Allocator
-    KVCache *-- PrefixCache
+    PagePool *-- Allocator
+    PagePool *-- PrefixCache
     InferenceEngine *-- InferenceScheduler
     InferenceScheduler *-- KVCache
     InferenceScheduler *-- Executor
     InferenceScheduler *-- TaskManager
-    SamplingPipeline *-- BaseSamplingStrategy
-    TrainContextBuilder *-- TrainContext
-    Transformer *-- DecoderBlock
-    Transformer *-- RotaryEmbedding
-    Transformer *-- Embedding
+    AutoRegressiveLM *-- DecoderBlock
+    AutoRegressiveLM *-- RotaryEmbedding
+    AutoRegressiveLM *-- Embedding
+    EmbeddingEncoder *-- DecoderBlock
+    EmbeddingEncoder *-- RotaryEmbedding
+    EmbeddingEncoder *-- Embedding
     DecoderBlock *-- RMSNorm
-    BaseDataset *-- BaseStorage
     ChatCompletionRequest *-- ChatMessage
     MessagesRequest *-- AnthropicMessage
+    AutoTokenizer *-- ChatTemplate
+    BaseFactory *-- Registry
 
     %% --- Aggregation (weak ownership) ---
-    AutoModel o-- ModelConfig
+    AutoModel o-- BaseModelConfig
     Trainer o-- TrainCallback
     TrainContext o-- BaseStrategy
     TrainContext o-- BaseScheduler
     TrainContext o-- Checkpoint
-    AutoTokenizer o-- ChatTemplate
     KvcacheView o-- Storage
-    BaseFactory o-- Registry
+    SamplingPipeline o-- BaseSamplingStrategy
+    BaseDataset o-- BaseStorage
 
     %% --- Dependency (uses temporarily) ---
     TrainConfig ..> BaseStrategy : selects
@@ -804,7 +917,12 @@ classDiagram
     FFNFactory ..> DeepSeekMoE : creates
     DecoderBlock ..> AttnFactory : uses
     DecoderBlock ..> FFNFactory : uses
+    StorageFactory ..> H5Storage : creates
+    StorageFactory ..> JSONStorage : creates
+    ConfigFactory ..> AutoRegressiveLMConfig : creates
+    ConfigFactory ..> EncoderConfig : creates
     Trainer ..> TrainContextBuilder : uses
+    TrainContextBuilder ..> TrainContext : creates
     Trainer ..> Functions : spawns
     TrainContextBuilder ..> StrategyFactory : uses
     TrainContextBuilder ..> ResumableDistributedSampler : creates
@@ -820,13 +938,13 @@ classDiagram
 
     %% --- Association (general usage) ---
     Trainer --> TrainConfig
-    DPOStrategy --> Transformer
-    GRPOStrategy --> Transformer
+    DPOStrategy --> AutoModel
+    GRPOStrategy --> AutoModel
     InferenceScheduler --> Task
     InferenceScheduler --> TaskStatus
     Task --> TaskStatus
-    InferenceEngine --> Transformer
-    Executor --> Transformer
+    InferenceEngine --> AutoModel
+    Executor --> AutoModel
     Executor --> AutoTokenizer
     TaskManager --> AutoTokenizer
     MultiSegmentFetcher --> BaseSegmentFetcher
@@ -839,12 +957,12 @@ classDiagram
 
 | Module | Components | Description |
 |--------|------------|-------------|
-| **astrai.config** | ModelConfig, TrainConfig | Configuration management |
-| **astrai.dataset** | BaseDataset–GRPODataset, BaseStorage–JSONStorage, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
+| **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
+| **astrai.dataset** | BaseDataset–GRPODataset, BaseStorage–JSONStorage, StorageFactory, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
-| **astrai.model** | AutoModel, Transformer, 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.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–SGDRScheduler, SchedulerFactory, TrainCallback–MetricLoggerCallback, CallbackFactory | Training workflow |
+| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–SGDRScheduler, SchedulerFactory, TrainCallback(Protocol)–ValidationCallback, CallbackFactory, Muon | Training workflow |
 | **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–KvcacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler–AnthropicHandler, ChatMessage–MessagesRequest, app | Inference service |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel |
 | **astrai.factory** | Registry, BaseFactory[T] | Component registration |
@@ -853,7 +971,7 @@ classDiagram
 
 | Pattern | Classes | Purpose |
 |---------|---------|---------|
-| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory` | Decorator-based component creation |
+| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StorageFactory`, `ConfigFactory` | Decorator-based component creation |
 | **Registry** | `BaseFactory`, `Registry` | Component registration with category/priority |
 | **Strategy** | `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy` | Training strategy switching |
 | **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations |
@@ -864,18 +982,18 @@ classDiagram
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
 | **Storage** | `BaseStorage`, `H5Storage`, `JSONStorage` | Format-agnostic data access |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
-| **AutoModel Registry** | `AutoModel`, `Transformer` | Model-type dynamic loading |
+| **AutoModel Registry** | `AutoModel`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |
 
 ## Core Relationships
 
 1. **Config → Training**: `TrainConfig` holds model, dataset, optimizer_fn, scheduler_fn
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
-4. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `Transformer`, backed by `KVCache` + `SamplingPipeline`
+4. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `KVCache` + `SamplingPipeline`
 5. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 6. **Dataset Loading**: `DatasetFactory` creates datasets, `BaseStorage` (H5Storage/JSONStorage) loads via `BaseSegmentFetcher` + `MultiSegmentFetcher`
 7. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only)
 8. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`
 9. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 
-> Document Update Time: 2026-05-15
+> Document Update Time: 2026-05-17

assets/docs/dataflow.md

@@ -15,8 +15,8 @@ Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or
 Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
 
 ```
-create_storage("h5")   → H5Storage
-create_storage("json") → JSONStorage
+StorageFactory.create("h5")   → H5Storage
+StorageFactory.create("json") → JSONStorage
 ```
 
 Both support shared memory via `.share_memory_()`.
@@ -34,7 +34,7 @@ Both support shared memory via `.share_memory_()`.
 
 ```
 DatasetFactory.load(train_type, path, window_size, stride)
-  → create_storage(detect_format(path))
+  → StorageFactory.create(detect_format(path))
     → MultiSegmentFetcher(BaseSegmentFetcher per key)
       → BaseDataset.__getitem__(idx)
         → sliding window [begin, end) via get_index(idx)
@@ -54,4 +54,4 @@ DatasetFactory.load(train_type, path, window_size, stride)
 
 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-05-15
+> Document Update Time: 2026-05-17

assets/docs/inference.md

@@ -137,4 +137,4 @@ engine.generate(["A", "B"], stream=True)        # -> Generator[Tuple[int, str]]
 await engine.generate_async("Hello", ...)       # -> AsyncGenerator[str]
 ```
 
-> Document Update Time: 2026-05-15
+> Document Update Time: 2026-05-17

assets/docs/params.md

@@ -10,14 +10,14 @@
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
 | `--n_epoch` | Total training epochs | 1 |
-| `--batch_size` | Batch size | 1 |
-| `--accumulation_steps` | Gradient accumulation steps between optimizer steps | 1 |
+| `--batch_per_device` | Batch size per device | 1 |
+| `--grad_accum_steps` | Gradient accumulation steps between optimizer steps | 1 |
 
 ### Learning Rate Scheduling
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--warmup_steps` | Warmup steps | 1000 |
+| `--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 |
 
@@ -60,7 +60,7 @@
 | Parameter | Description | Default | Used by |
 |-----------|-------------|---------|---------|
 | `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
-| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.1 (CLI) / 0.0 (strategy default) | `seq`, `sft` |
+| `--label_smoothing` | Label smoothing for cross-entropy loss | 0.05 | `seq`, `sft` |
 | `--group_size` | GRPO group size | 4 | `grpo` |
 | `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
 | `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
@@ -69,90 +69,29 @@
 ### Usage Example
 
 ```bash
-python scripts/tools/train.py \
-  --train_type seq \
-  --data_root_path /path/to/dataset \
-  --param_path /path/to/model \
-  --n_epoch 3 \
-  --batch_size 4 \
-  --accumulation_steps 8 \
-  --max_lr 3e-4 \
-  --warmup_steps 2000 \
-  --max_grad_norm 1.0 \
-  --ckpt_interval 5000 \
-  --ckpt_dir ./checkpoints \
-  --num_workers 4 \
-  --nprocs 1 \
-  --device_type cuda
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
 ```
 
 ---
 
-## Generation Parameters
-
-### GenerationRequest Parameters
-
-| Parameter | Description | Default Value |
-|-----------|-------------|---------------|
-| `messages` | List of message dictionaries (role, content) | required |
-| `temperature` | Sampling temperature (higher = more random) | 1.0 |
-| `top_p` | Nucleus sampling threshold | 1.0 |
-| `top_k` | Top-k sampling count | 50 |
-| `max_tokens` | Maximum generation length | None (defaults to max_seq_len - prompt_len) |
-| `stream` | Whether to stream output | False |
-
-### Usage Example
-
-```python
-import torch
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
-from astrai.inference import InferenceEngine, GenerationRequest
-
-# Load model using AutoModel
-model = AutoModel.from_pretrained("your_model_dir")
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("your_model_dir")
-
-# Create engine with separate model and tokenizer
-engine = InferenceEngine(
-    model=model,
-    tokenizer=tokenizer,
-)
-
-# Build request with messages format
-request = GenerationRequest(
-    messages=[
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": "Hello"},
-    ],
-    temperature=0.8,
-    top_p=0.95,
-    top_k=50,
-    max_tokens=None,
-)
-
-# Generate (streaming)
-for token in engine.generate_with_request(request):
-    print(token, end="", flush=True)
-
-# Or use simple generate interface
-result = engine.generate(
-    prompt="Hello",
-    stream=False,
-    max_tokens=1024,
-    temperature=0.8,
-    top_p=0.95,
-    top_k=50,
-)
-```
-
-### Generation Modes
-
-| Mode | Description |
-|------|-------------|
-| `stream=True` | Streaming output, yields token by token |
-| `stream=False` | Non-streaming output, returns complete result |
-
-> Document Update Time: 2026-05-15
+> Document Update Time: 2026-05-17

assets/docs/training.md

@@ -65,24 +65,24 @@ The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per posi
 
 ## Training Loop
 
-Nested loop: **epoch** → **step** (accumulation window) → **batch**.
+Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
 
 ```
 on_train_begin
   on_epoch_begin
-    for steps in batched(dataloader, accumulation_steps):
-      on_step_begin
-        step_batch_nums = len(steps)
-        for batch in steps:
-          on_batch_begin
-          loss = strategy(batch)
-          (loss / step_batch_nums).backward()
-          iteration += 1
-          on_batch_end
-      on_step_end
-      optimizer.step()
-      optimizer.zero_grad()
-      scheduler.step()
+    for batch in dataloader:
+      on_batch_begin
+      loss = strategy(batch)
+      (loss / grad_accum_steps).backward()
+      iteration += 1
+      on_batch_end
+
+      if iteration % grad_accum_steps == 0:
+        on_step_begin
+        optimizer.step()
+        optimizer.zero_grad()
+        on_step_end
+        scheduler.step()
     on_epoch_end
 on_train_end
 ```
@@ -91,11 +91,13 @@ on_train_end
 
 | Hook | Fires | Default callback |
 |------|-------|-----------------|
-| `on_step_end` | Every accumulation window | `GradientClippingCallback` |
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_step_begin` | Every accumulation window | `GradientClippingCallback` |
 | `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
-| `on_train_end` | Training ends | `CheckpointCallback` (final save) |
+| `on_step_end` | Every accumulation window | `ValidationCallback` |
+| `on_train_end` | Training ends | `CheckpointCallback`, `MetricLoggerCallback` (final save) |
 
-Default callbacks: `progress_bar` (tqdm), `checkpoint` (safetensors, rank-0), `metric_logger` (JSONL, rank-0), `gradient_clipping`.
+Default callbacks: `gradient_checkpointing` (activation checkpointing, optional), `progress_bar` (tqdm), `checkpoint` (safetensors, rank-0), `metric_logger` (JSONL, rank-0), `gradient_clipping`, `validation` (periodic validation on val_dataset).
 
 ## Strategies
 
@@ -154,15 +156,27 @@ Keys: `prompts`, `responses`, `masks`, `rewards`.
 
 Created by `SchedulerFactory.create(optimizer, schedule_type, **kwargs)`.
 
+## Gradient Checkpointing
+
+Trades compute for memory by recomputing activations during backward pass. Specify module types via `gradient_checkpointing_modules`:
+
+```python
+from astrai.model.components.decoder_block import DecoderBlock
+config = TrainConfig(..., gradient_checkpointing_modules=[DecoderBlock])
+```
+
+Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoint(use_reentrant=False)`, compatible with `torch.compile`. Uses `nn.Module.apply()` for traversal — works through DDP wrappers without manual unwrap. Empty list (default) means no-op.
+
 ## Checkpoint
 
 ```
-Checkpoint(state_dict, epoch, iteration, extra)
-  ├── save(save_dir)    rank-0 only: meta.json + state_dict.safetensors + optional extra.pt
+Checkpoint(state_dict, epoch, iteration, extra, meta)
+  ├── save(save_dir)    rank-0 only: meta.json (includes training config) + state_dict.safetensors + optional extra.pt
   └── load(save_dir)    broadcasts metadata from rank-0
 ```
 
-Optimizer/scheduler state NOT persisted by default; `Checkpoint.extra` can store arbitrary data.
+Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
+Training config (`TrainConfig.to_dict()`) saved into `meta.json` during training via `CheckpointCallback`.
 
 ## TrainContextBuilder (Builder Pattern)
 
@@ -183,17 +197,29 @@ context = (
 ## Training CLI
 
 ```bash
-python scripts/tools/train.py \
-    --train_type seq \
-    --data_root_path /path/to/data \
-    --param_path /path/to/model \
-    --batch_size 4 \
-    --accumulation_steps 8 \
-    --max_lr 3e-4 \
-    --warmup_steps 1000 \
-    --n_epoch 1
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+nohup python scripts/tools/train.py \
+    --nprocs=4 \
+    --train_type=seq \
+    --data_root_path=/path/to/dataset \
+    --param_path=/path/to/model \
+    --batch_per_device=4 \
+    --grad_accum_steps=8 \
+    --warmup_ratio=0.05 \
+    --max_lr=1e-4 \
+    --max_grad_norm=1.0 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
+    --window_size=2048 \
+    --ckpt_interval=10000 \
+    --ckpt_dir=./checkpoint \
+    --random_seed=3407 \
+    --label_smoothing=0.05 \
+    > out.log 2> err.log &
 ```
 
 Full parameter reference at [params.md](params.md).
 
-> Document Update Time: 2026-05-15
+> Document Update Time: 2026-05-17

astrai/__init__.py

@@ -1,8 +1,9 @@
-__version__ = "1.3.5"
+__version__ = "1.3.6"
 __author__ = "ViperEkura"
 
 from astrai.config import (
-    ModelConfig,
+    AutoRegressiveLMConfig,
+    EncoderConfig,
     TrainConfig,
 )
 from astrai.dataset import DatasetFactory
@@ -11,13 +12,14 @@ from astrai.inference import (
     GenerationRequest,
     InferenceEngine,
 )
-from astrai.model import AutoModel, Transformer
+from astrai.model import AutoModel, AutoRegressiveLM
 from astrai.tokenize import AutoTokenizer
 from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer
 
 __all__ = [
-    "Transformer",
-    "ModelConfig",
+    "AutoRegressiveLM",
+    "AutoRegressiveLMConfig",
+    "EncoderConfig",
     "TrainConfig",
     "DatasetFactory",
     "AutoTokenizer",

astrai/config/__init__.py

@@ -1,8 +1,16 @@
-from astrai.config.model_config import ModelConfig
+from astrai.config.model_config import (
+    AutoRegressiveLMConfig,
+    BaseModelConfig,
+    ConfigFactory,
+    EncoderConfig,
+)
 from astrai.config.train_config import TrainConfig
 
 __all__ = [
     # Model configuration
-    "ModelConfig",
+    "BaseModelConfig",
+    "AutoRegressiveLMConfig",
+    "EncoderConfig",
+    "ConfigFactory",
     "TrainConfig",
 ]

astrai/config/base.py

@@ -0,0 +1,77 @@
+import json
+from dataclasses import MISSING, dataclass, fields
+from typing import Any, Dict, Optional, Self, get_type_hints
+
+
+@dataclass
+class BaseConfig:
+    def to_dict(self) -> Dict[str, Any]:
+        d = {}
+        for fld in fields(self):
+            v = getattr(self, fld.name)
+            if isinstance(v, (str, int, float, bool)):
+                d[fld.name] = v
+            elif v is None:
+                d[fld.name] = None
+            elif isinstance(v, (dict, list)):
+                try:
+                    json.dumps(v)
+                    d[fld.name] = v
+                except (TypeError, ValueError):
+                    pass
+        return d
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Self:
+        hints = get_type_hints(cls)
+        inst = cls.__new__(cls)
+        for fld in fields(cls):
+            if fld.name in d:
+                v = d[fld.name]
+                target = cls._unwrap_optional(hints.get(fld.name))
+                if target is not None:
+                    try:
+                        v = cls._coerce(v, target)
+                    except (TypeError, ValueError):
+                        pass
+                object.__setattr__(inst, fld.name, v)
+            elif fld.default is not MISSING:
+                object.__setattr__(inst, fld.name, fld.default)
+            elif fld.default_factory is not MISSING:
+                object.__setattr__(inst, fld.name, fld.default_factory())
+            else:
+                object.__setattr__(inst, fld.name, None)
+        return inst
+
+    @staticmethod
+    def _unwrap_optional(tp) -> Optional[type]:
+        if tp is None:
+            return None
+        origin = getattr(tp, "__origin__", None)
+        if origin is not None:
+            args = getattr(tp, "__args__", ())
+            non_none = [a for a in args if a is not type(None)]
+            return non_none[0] if non_none else None
+        return tp
+
+    @staticmethod
+    def _coerce(value: Any, target_type: type) -> Any:
+        if target_type is bool and isinstance(value, bool):
+            return value
+        if (
+            target_type is int
+            and isinstance(value, (int, float))
+            and not isinstance(value, bool)
+        ):
+            return int(value)
+        if (
+            target_type is float
+            and isinstance(value, (int, float))
+            and not isinstance(value, bool)
+        ):
+            return float(value)
+        if target_type is str and isinstance(value, str):
+            return value
+        if isinstance(value, target_type):
+            return value
+        raise TypeError

astrai/config/model_config.py

@@ -1,114 +1,90 @@
 import json
-import sys
-from dataclasses import dataclass, fields
-from typing import Any, Dict, Optional, Self, get_type_hints
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Self
+
+from astrai.config.base import BaseConfig
+from astrai.factory import BaseFactory
+
+
+class ConfigFactory(BaseFactory[BaseConfig]):
+    """Factory that dispatches config classes by ``model_type``."""
+
+    @classmethod
+    def load(cls, raw: Dict[str, Any]) -> BaseConfig:
+        model_type = raw.get("model_type") or "autoregressive_lm"
+        config_cls = cls.get_component_class(model_type)
+        return config_cls.from_dict(raw)
 
 
 @dataclass
-class BaseModelConfig:
-    """Field-aware JSON load/save for dataclass configs.
-
-    Subclass with additional fields. The base ``model_type`` field
-    enables ``AutoModel`` to pick the correct subclass.
-    """
+class BaseModelConfig(BaseConfig):
+    """Base config with ``model_type`` dispatch and file I/O."""
 
     model_type: Optional[str] = None
 
-    def load(self, config_path: str) -> Self:
-        raw: Dict[str, Any] = {}
+    @classmethod
+    def from_file(cls, config_path: str) -> Self:
         with open(config_path, "r") as f:
-            raw.update(json.load(f))
+            raw: Dict[str, Any] = json.load(f)
+        return cls.from_dict(raw)
 
-        hints = get_type_hints(type(self))
-        valid = {fld.name for fld in fields(self)}
-        for key, value in raw.items():
-            if key not in valid:
-                sys.stderr.write(f"WARNING: unknown config key '{key}'\n")
-                continue
-
-            target_type = self._unwrap_optional(hints.get(key))
-            if target_type is None:
-                continue
-
-            try:
-                value = self._coerce(value, target_type)
-            except (TypeError, ValueError):
-                sys.stderr.write(
-                    f"WARNING: cannot coerce '{key}' = {value!r} to {target_type}\n"
-                )
-                continue
-
-            setattr(self, key, value)
-
-        return self
-
-    def save(self, config_path: str):
-        config_dict: Dict[str, Any] = {}
-        for fld in fields(self):
-            v = getattr(self, fld.name)
-            if v is not None:
-                config_dict[fld.name] = v
+    def to_file(self, config_path: str):
+        d = self.to_dict()
+        config_dict = {k: v for k, v in d.items() if v is not None}
         with open(config_path, "w") as f:
             json.dump(config_dict, f, indent=4)
 
-    @staticmethod
-    def _unwrap_optional(tp: type) -> Optional[type]:
-        if tp is None:
-            return None
-        origin = getattr(tp, "__origin__", None)
-        if origin is not None:
-            args = getattr(tp, "__args__", ())
-            non_none = [a for a in args if a is not type(None)]
-            return non_none[0] if non_none else None
-        return tp
-
-    @staticmethod
-    def _coerce(value: Any, target_type: type) -> Any:
-        if target_type is bool and isinstance(value, bool):
-            return value
-        if (
-            target_type is int
-            and isinstance(value, (int, float))
-            and not isinstance(value, bool)
-        ):
-            return int(value)
-        if (
-            target_type is float
-            and isinstance(value, (int, float))
-            and not isinstance(value, bool)
-        ):
-            return float(value)
-        if target_type is str and isinstance(value, str):
-            return value
-        if isinstance(value, target_type):
-            return value
-        raise TypeError
-
 
 @dataclass
-class ModelConfig(BaseModelConfig):
+@ConfigFactory.register("autoregressive_lm")
+class AutoRegressiveLMConfig(BaseModelConfig):
+    """Configuration for autoregressive language model."""
+
     vocab_size: Optional[int] = None
     dim: Optional[int] = None
-
     n_layers: Optional[int] = None
     norm_eps: Optional[float] = None
     dim_ffn: Optional[int] = None
     tie_weight: Optional[bool] = None
 
-    # RoPE
     max_len: Optional[int] = None
     rope_theta: Optional[float] = None
 
-    # attention
     attn_type: str = "gqa"
     n_heads: Optional[int] = None
     n_kv_heads: Optional[int] = None
     use_qk_norm: Optional[bool] = None
     use_gated_attention: Optional[bool] = None
 
-    # MoE
+    kv_lora_rank: Optional[int] = None
+    qk_nope_head_dim: Optional[int] = None
+    qk_rope_head_dim: Optional[int] = None
+
     ffn_type: str = "mlp"
     n_routed_experts: Optional[int] = None
     n_shared_experts: Optional[int] = None
     n_activated_experts: Optional[int] = None
-    moe_topk_method: Optional[str] = None
+    topk_method: Optional[str] = None
+
+
+@dataclass
+@ConfigFactory.register("embedding")
+class EncoderConfig(BaseModelConfig):
+    """Configuration for embedding encoder model."""
+
+    vocab_size: Optional[int] = None
+    dim: Optional[int] = None
+    n_layers: Optional[int] = None
+    norm_eps: Optional[float] = None
+    dim_ffn: Optional[int] = None
+
+    max_len: Optional[int] = None
+    rope_theta: Optional[float] = None
+
+    n_heads: Optional[int] = None
+    n_kv_heads: Optional[int] = None
+    use_qk_norm: Optional[bool] = None
+    use_gated_attention: Optional[bool] = None
+
+    pooling_type: Optional[str] = None
+    normalize_embeddings: Optional[bool] = None

astrai/config/train_config.py

@@ -1,32 +1,48 @@
-from dataclasses import dataclass, field
-from typing import Callable, Optional
+from dataclasses import dataclass, field, fields
+from typing import Callable, List, Optional
 
 import torch.nn as nn
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import Dataset
 
+from astrai.config.base import BaseConfig
+
+
+def required(**kw):
+    return {"required": True, **kw}
+
 
 @dataclass
-class TrainConfig:
+class TrainConfig(BaseConfig):
     # basic setting
-    model: nn.Module = field(default=None, metadata={"help": "Model for training."})
-    strategy: str = field(default=None, metadata={"help": "Training strategy."})
-    dataset: Dataset = field(default=None, metadata={"help": "Dataset for training."})
+    model: nn.Module = field(
+        default=None, metadata=required(help="Model for training.")
+    )
+    strategy: str = field(default=None, metadata=required(help="Training strategy."))
+    dataset: Dataset = field(
+        default=None, metadata=required(help="Dataset for training.")
+    )
     optimizer_fn: Callable[[nn.Module], Optimizer] = field(
-        default=None, metadata={"help": "Optimizer factory for training."}
+        default=None, metadata=required(help="Optimizer factory for training.")
     )
     scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
-        default=None, metadata={"help": "Scheduler factory for training."}
+        default=None, metadata=required(help="Scheduler factory for training.")
     )
     n_epoch: int = field(default=1, metadata={"help": "Number of epochs for training."})
-    batch_size: int = field(default=4, metadata={"help": "Batch size for training."})
-    accumulation_steps: int = field(
+    batch_per_device: int = field(
+        default=4, metadata={"help": "Batch size per device."}
+    )
+    grad_accum_steps: int = field(
         default=1, metadata={"help": "Number of iterations between steps."}
     )
     max_grad_norm: float = field(
         default=1.0, metadata={"help": "Maximum gradient norm."}
     )
+    gradient_checkpointing_modules: list = field(
+        default_factory=list,
+        metadata={"help": "Module types to enable activation checkpointing for."},
+    )
 
     # checkpoint setting
     start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
@@ -40,6 +56,19 @@ class TrainConfig:
         default=5000, metadata={"help": "Number of iterations between checkpoints."}
     )
 
+    # metric setting
+    log_dir: str = field(
+        default="./checkpoint/logs", metadata={"help": "Directory for metric logs."}
+    )
+    log_interval: int = field(
+        default=100,
+        metadata={"help": "Number of batch iterations between metric logs."},
+    )
+    metrics: List[str] = field(
+        default_factory=lambda: ["loss", "lr"],
+        metadata={"help": "Metrics to record during training."},
+    )
+
     # dataloader setting
     random_seed: int = field(default=3407, metadata={"help": "Random seed."})
     num_workers: int = field(
@@ -72,11 +101,23 @@ class TrainConfig:
     state_dict_fn: Optional[Callable] = field(
         default=None, metadata={"help": "Parallel function for state  dict saving."}
     )
+    start_method: str = field(
+        default="spawn",
+        metadata={"help": "Multiprocessing start method (spawn/fork/forkserver)."},
+    )
 
     # others
     device_type: str = field(
         default="cuda", metadata={"help": "Device type for distributed training."}
     )
+    val_dataset: Optional[Dataset] = field(
+        default=None, metadata={"help": "Dataset for validation."}
+    )
+    val_step: int = field(
+        default=1000,
+        metadata={"help": "Number of optimizer steps between validation runs."},
+    )
+
     extra_kwargs: dict = field(
         default_factory=dict, metadata={"help": "Other arguments."}
     )
@@ -85,14 +126,6 @@ class TrainConfig:
         self.validate()
 
     def validate(self):
-        required_fields = [
-            "model",
-            "strategy",
-            "dataset",
-            "optimizer_fn",
-            "scheduler_fn",
-        ]
-
-        for field_name in required_fields:
-            if getattr(self, field_name) is None:
-                raise ValueError(f"{field_name} is required.")
+        for fld in fields(self):
+            if fld.metadata.get("required") and getattr(self, fld.name) is None:
+                raise ValueError(f"TrainConfig.{fld.name} is required but got None.")

astrai/dataset/__init__.py

@@ -9,8 +9,7 @@ from astrai.dataset.storage import (
     H5Storage,
     JSONStorage,
     MultiSegmentFetcher,
-    available_storage_types,
-    create_storage,
+    StorageFactory,
     detect_format,
     load_h5,
     load_json,
@@ -26,9 +25,8 @@ __all__ = [
     "BaseStorage",
     "H5Storage",
     "JSONStorage",
-    "create_storage",
+    "StorageFactory",
     "detect_format",
-    "available_storage_types",
     "save_h5",
     "load_h5",
     "save_json",

astrai/dataset/dataset.py

@@ -9,7 +9,7 @@ from torch.utils.data import Dataset
 
 from astrai.dataset.storage import (
     BaseStorage,
-    create_storage,
+    StorageFactory,
     detect_format,
 )
 from astrai.factory import BaseFactory
@@ -28,6 +28,26 @@ class BaseDataset(Dataset, ABC):
         self.stride = stride
         self.storage: Optional[BaseStorage] = None
 
+    @property
+    def required_keys(self) -> List[str]:
+        """Return required storage keys for this dataset type.
+
+        Subclasses should override to specify expected keys.
+        """
+        return []
+
+    def _validate_keys(self):
+        if not self.required_keys:
+            return
+        actual_keys = set(self.storage.keys)
+        missing = [k for k in self.required_keys if k not in actual_keys]
+        if missing:
+            raise KeyError(
+                f"Dataset {type(self).__name__} requires keys {self.required_keys}, "
+                f"but storage at {self._load_path} only has {sorted(actual_keys)}. "
+                f"Missing: {missing}"
+            )
+
     def load(self, load_path: str, storage_type: Optional[str] = None, tokenizer=None):
         """Load dataset from the given path.
 
@@ -39,11 +59,16 @@ class BaseDataset(Dataset, ABC):
                           or None for auto-detection
             tokenizer: Callable str -> List[int], used to tokenize raw text
                        in JSON files. Ignored for HDF5.
+
+        Raises:
+            KeyError: If the loaded storage is missing required keys.
         """
         if storage_type is None:
             storage_type = detect_format(load_path)
-        self.storage = create_storage(storage_type)
+        self.storage = StorageFactory.create(storage_type)
+        self._load_path = load_path
         self.storage.load(load_path, tokenizer=tokenizer)
+        self._validate_keys()
 
     def load_json(self, load_path: str, tokenizer=None):
         """Load dataset from JSON files explicitly.
@@ -186,6 +211,10 @@ class SEQDataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["sequence"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
         return self.storage.fetch(begin_idx, end_idx, "sequence")
 
@@ -205,6 +234,10 @@ class SFTDataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["sequence", "loss_mask"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
         return self.storage.fetch(begin_idx, end_idx, key)
 
@@ -229,6 +262,10 @@ class DPODataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
         return self.storage.fetch(begin_idx, end_idx, key)
 
@@ -259,6 +296,10 @@ class GRPODataset(BaseDataset):
     def __init__(self, window_size: int, stride: int):
         super().__init__(window_size, stride)
 
+    @property
+    def required_keys(self) -> List[str]:
+        return ["prompts", "responses", "masks", "rewards"]
+
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
         return self.storage.fetch(begin_idx, end_idx, key)
 

astrai/dataset/storage.py

@@ -15,6 +15,8 @@ import h5py
 import torch
 from torch import Tensor
 
+from astrai.factory import BaseFactory
+
 
 def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
     os.makedirs(file_path, exist_ok=True)
@@ -258,6 +260,24 @@ class BaseStorage(ABC):
         return self._fetcher.multi_keys
 
 
+class StorageFactory(BaseFactory["BaseStorage"]):
+    """Factory for creating storage backends by type name.
+
+    Example:
+        @StorageFactory.register("custom")
+        class CustomStorage(BaseStorage):
+            ...
+
+        storage = StorageFactory.create("custom")
+    """
+
+    @classmethod
+    def _validate_component(cls, storage_cls: type) -> None:
+        if not issubclass(storage_cls, BaseStorage):
+            raise TypeError(f"{storage_cls.__name__} must inherit from BaseStorage")
+
+
+@StorageFactory.register("h5")
 class H5Storage(BaseStorage):
     """HDF5-based storage backend (pre-tokenized data)."""
 
@@ -266,6 +286,7 @@ class H5Storage(BaseStorage):
         self._fetcher = MultiSegmentFetcher(segments)
 
 
+@StorageFactory.register("json")
 class JSONStorage(BaseStorage):
     """JSON-based storage backend.
 
@@ -278,35 +299,3 @@ class JSONStorage(BaseStorage):
     def load(self, load_path: str, tokenizer=None) -> None:
         segments = load_json(load_path, tokenizer=tokenizer)
         self._fetcher = MultiSegmentFetcher(segments)
-
-
-_STORAGE_REGISTRY: Dict[str, type] = {
-    "h5": H5Storage,
-    "json": JSONStorage,
-}
-
-
-def create_storage(storage_type: str) -> BaseStorage:
-    """Create a storage instance by type name.
-
-    Args:
-        storage_type: Storage type name ("h5", "json")
-
-    Returns:
-        Storage instance
-
-    Raises:
-        ValueError: If the storage type is unknown
-    """
-    storage_cls = _STORAGE_REGISTRY.get(storage_type)
-    if storage_cls is None:
-        raise ValueError(
-            f"Unknown storage type: '{storage_type}'. "
-            f"Available: {sorted(_STORAGE_REGISTRY.keys())}"
-        )
-    return storage_cls()
-
-
-def available_storage_types() -> List[str]:
-    """Return list of registered storage type names."""
-    return sorted(_STORAGE_REGISTRY.keys())

astrai/factory.py

@@ -1,5 +1,6 @@
 """Base factory class for extensible component registration."""
 
+import inspect
 from abc import ABC
 from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar
 
@@ -122,6 +123,10 @@ class BaseFactory(ABC, Generic[T]):
     def create(cls, name: str, *args, **kwargs) -> T:
         """Create a component instance by name.
 
+        Filters kwargs to match the component's __init__ signature,
+        so components don't need to declare **kwargs just to absorb
+        parameters meant for other components.
+
         Args:
             name: Registered name of the component
             *args: Positional arguments passed to component constructor
@@ -139,6 +144,17 @@ class BaseFactory(ABC, Generic[T]):
                 f"Supported types: {sorted(cls._registry.list_names())}"
             )
         component_cls = cls._registry.get(name)
+        sig = inspect.signature(component_cls.__init__)
+        has_var_kwargs = any(
+            p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
+        )
+        if not has_var_kwargs:
+            valid = {
+                p.name
+                for p in sig.parameters.values()
+                if p.name != "self" and p.kind != inspect.Parameter.VAR_KEYWORD
+            }
+            kwargs = {k: v for k, v in kwargs.items() if k in valid}
         return component_cls(*args, **kwargs)
 
     @classmethod

astrai/inference/api/protocol.py

@@ -226,6 +226,17 @@ class OpenAIHandler(ProtocolHandler):
     def create_response_id(self) -> str:
         return f"chatcmpl-{uuid.uuid4().hex[:12]}"
 
+    def get_stop_sequences(self) -> List[str]:
+        stop = self.request.stop
+        if stop is None:
+            return []
+        return [stop] if isinstance(stop, str) else stop
+
+    def on_token(
+        self, ctx: StreamContext, token: str, stop_checker: StopChecker
+    ) -> Optional[str]:
+        return stop_checker.check(ctx.accumulated)
+
     def format_stream_start(self, ctx: StreamContext) -> List[str]:
         return [
             _sse_event(

astrai/inference/api/server.py

@@ -12,7 +12,7 @@ from typing import Any, Dict, List, Optional, Union
 
 import torch
 import uvicorn
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field
 
 from astrai.inference.api.protocol import AnthropicHandler, OpenAIHandler
@@ -67,6 +67,24 @@ class MessagesRequest(BaseModel):
     stop_sequences: Optional[List[str]] = None
 
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    config = app.state.server_config
+    if not config.get("_test", False):
+        try:
+            app.state.engine = _create_engine(**config)
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            raise
+    yield
+    if app.state.engine:
+        app.state.engine.shutdown()
+        logger.info("Inference engine shutdown complete")
+
+
+app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
+
+
 def _create_engine(
     param_path: Optional[Path] = None,
     device: str = "cuda",
@@ -92,54 +110,36 @@ def _create_engine(
     return engine
 
 
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    config = app.state.server_config
-    if not config.get("_test", False):
-        try:
-            app.state.engine = _create_engine(**config)
-        except Exception as e:
-            logger.error(f"Failed to load model: {e}")
-            raise
-    yield
-    if app.state.engine:
-        app.state.engine.shutdown()
-        logger.info("Inference engine shutdown complete")
-
-
-app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
-
-
-def _get_engine(request: Request) -> InferenceEngine:
-    engine = request.app.state.engine
+def _get_engine() -> InferenceEngine:
+    engine = app.state.engine
     if engine is None:
         raise HTTPException(status_code=503, detail="Engine not initialized")
     return engine
 
 
 @app.get("/health")
-async def health(request: Request):
+async def health():
     return {
         "status": "ok",
-        "model_loaded": request.app.state.engine is not None,
+        "model_loaded": app.state.engine is not None,
     }
 
 
 @app.get("/stats")
-async def get_stats(request: Request):
-    return _get_engine(request).get_stats()
+async def get_stats():
+    return _get_engine().get_stats()
 
 
 @app.post("/v1/chat/completions")
-async def chat_completion(request: ChatCompletionRequest, req: Request):
-    engine = _get_engine(req)
+async def chat_completion(request: ChatCompletionRequest):
+    engine = _get_engine()
     handler = OpenAIHandler(request, engine)
     return await handler.handle()
 
 
 @app.post("/v1/messages")
-async def create_message(request: MessagesRequest, req: Request):
-    engine = _get_engine(req)
+async def create_message(request: MessagesRequest):
+    engine = _get_engine()
     handler = AnthropicHandler(request, engine)
     return await handler.handle()
 
@@ -163,4 +163,5 @@ def run_server(
         app,
         host=host,
         port=port,
+        reload=reload,
     )

astrai/inference/core/scheduler.py

@@ -22,14 +22,22 @@ class InferenceScheduler:
         tokenizer: AutoTokenizer,
         max_batch_size: int = 16,
         max_seq_len: Optional[int] = None,
-        max_prompt_len: int = 512,
+        max_prompt_len: int = 2048,
         page_size: int = 64,
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
     ):
         config = model.config
 
-        self.max_seq_len = max_seq_len or config.max_len
+        if max_seq_len is not None:
+            self.max_seq_len = max_seq_len
+        elif config.max_len is not None:
+            self.max_seq_len = config.max_len
+        else:
+            raise ValueError(
+                "max_seq_len must be provided either as argument "
+                "or in model config (config.max_len)"
+            )
         self.device = device or next(model.parameters()).device
         self.dtype = dtype or next(model.parameters()).dtype
 

astrai/model/__init__.py

@@ -4,7 +4,8 @@ from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.linear import Linear
 from astrai.model.components.mlp import MLP
 from astrai.model.components.norm import RMSNorm
-from astrai.model.transformer import Transformer
+from astrai.model.encoder import EmbeddingEncoder
+from astrai.model.transformer import AutoRegressiveLM
 
 __all__ = [
     # Modules
@@ -14,6 +15,7 @@ __all__ = [
     "GQA",
     "DecoderBlock",
     # Models
-    "Transformer",
+    "AutoRegressiveLM",
+    "EmbeddingEncoder",
     "AutoModel",
 ]

astrai/model/automodel.py

@@ -2,6 +2,7 @@
 AutoModel base class for model loading and saving.
 """
 
+import json
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Self, Union
@@ -9,7 +10,7 @@ from typing import Self, Union
 import safetensors.torch as st
 import torch.nn as nn
 
-from astrai.config import ModelConfig
+from astrai.config.model_config import BaseModelConfig, ConfigFactory
 from astrai.factory import BaseFactory
 
 
@@ -45,7 +46,7 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
     Provides model loading/saving, registration, and generation.
     """
 
-    def __init__(self, config: ModelConfig):
+    def __init__(self, config: BaseModelConfig):
         super().__init__()
         self.config = config
 
@@ -60,14 +61,15 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
         model_path = Path(path)
 
         # Load config
-        config = ModelConfig()
         config_path = model_path / "config.json"
         if config_path.exists():
-            config.load(str(config_path))
+            with open(config_path, "r") as f:
+                raw = json.load(f)
+            config = ConfigFactory.load(raw)
+            model_type = config.model_type or "autoregressive_lm"
         else:
             raise FileNotFoundError(f"Config file not found: {config_path}")
 
-        model_type = config.model_type or "transformer"
         actual_cls = AutoModel.get_component_class(model_type)
 
         with _disable_random_init(enable=disable_random_init):
@@ -89,7 +91,7 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
         save_path.mkdir(parents=True, exist_ok=True)
 
         # Save config
-        self.config.save(str(save_path / "config.json"))
+        self.config.to_file(str(save_path / "config.json"))
 
         # Save weights
         st.save_file(self.state_dict(), str(save_path / "model.safetensors"))

astrai/model/components/attention.py

@@ -40,7 +40,6 @@ class GQA(nn.Module):
         norm_eps: float,
         use_gated_attention: bool,
         layer_id: int,
-        **kwargs,
     ):
         super().__init__()
         assert dim % n_heads == 0
@@ -121,9 +120,9 @@ class MLA(nn.Module):
         qk_nope_head_dim: int,
         qk_rope_head_dim: int,
         norm_eps: float,
+        use_qk_norm: bool,
         use_gated_attention: bool,
         layer_id: int,
-        **kwargs,
     ):
         super().__init__()
         self.dim = dim
@@ -135,15 +134,20 @@ class MLA(nn.Module):
         self.head_dim = qk_nope_head_dim + qk_rope_head_dim
         self.layer_id = layer_id
         self.n_rep = n_heads // n_kv_heads
+        self.use_qk_norm = use_qk_norm
         self.use_gated_attention = use_gated_attention
 
         self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
+
+        if self.use_qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, norm_eps)
         self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
         self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
 
         self.kv_b_proj = Linear(
             kv_lora_rank,
-            n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
+            n_kv_heads * (2 * self.head_dim),
         )
 
         self.o_proj = Linear(dim, dim, bias=False)
@@ -176,7 +180,7 @@ class MLA(nn.Module):
 
         q_nope, q_rope = (
             q[..., : self.qk_nope_head_dim],
-            q[..., self.qk_rope_head_dim :],
+            q[..., self.qk_nope_head_dim :],
         )
         q_rope = apply_rotary_emb(q_rope, rotary_emb)
         k_rope = apply_rotary_emb(k_rope, rotary_emb)
@@ -184,6 +188,10 @@ class MLA(nn.Module):
         q = torch.cat([q_nope, q_rope], dim=-1)
         k = torch.cat([k_nope, k_rope], dim=-1)
 
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
         if paged_cache is not None:
             paged_cache.write(self.layer_id, k, v)
             k, v = paged_cache.gather(self.layer_id)

astrai/model/components/decoder_block.py

@@ -16,13 +16,13 @@ class DecoderBlock(nn.Module):
         n_heads: int,
         dim_ffn: int,
         n_kv_heads: int,
-        norm_eps: int,
+        norm_eps: float,
         use_qk_norm: bool,
         use_gated_attention: bool,
         layer_id: int,
         attn_type: str = "gqa",
         ffn_type: str = "mlp",
-        **moe_kwargs,
+        **kwargs,
     ):
         super().__init__()
         self.attention = AttnFactory.create(
@@ -34,10 +34,11 @@ class DecoderBlock(nn.Module):
             norm_eps=norm_eps,
             use_gated_attention=use_gated_attention,
             layer_id=layer_id,
+            **kwargs,
         )
         self.input_norm = RMSNorm(dim, norm_eps)
         self.post_attention_norm = RMSNorm(dim, norm_eps)
-        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **moe_kwargs)
+        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **kwargs)
 
     def forward(
         self,

astrai/model/components/embedding.py

@@ -9,5 +9,8 @@ class Embedding(nn.Module):
         super().__init__()
         self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
 
+    def reset_parameters(self):
+        nn.init.normal_(self.weight, mean=0.0, std=0.02)
+
     def forward(self, x: Tensor) -> Tensor:
         return F.embedding(x, self.weight)

astrai/model/components/linear.py

@@ -10,5 +10,12 @@ class Linear(nn.Module):
         self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
         self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
 
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / (fan_in**0.5)
+            nn.init.uniform_(self.bias, -bound, bound)
+
     def forward(self, x: Tensor) -> Tensor:
         return F.linear(x, self.weight, self.bias)

astrai/model/components/mlp.py

@@ -15,11 +15,11 @@ class FFNFactory(BaseFactory[nn.Module]):
 
 @FFNFactory.register("mlp")
 class MLP(nn.Module):
-    def __init__(self, dim: int, dim_feed_forward: int, **kwargs):
+    def __init__(self, dim: int, dim_ffn: int):
         super().__init__()
-        self.up = Linear(dim, dim_feed_forward)
-        self.gate = Linear(dim, dim_feed_forward)
-        self.down = Linear(dim_feed_forward, dim)
+        self.up = Linear(dim, dim_ffn)
+        self.gate = Linear(dim, dim_ffn)
+        self.down = Linear(dim_ffn, dim)
 
     def forward(self, x: Tensor) -> Tensor:
         gated = self.up(x) * F.silu(self.gate(x))
@@ -32,12 +32,11 @@ class DeepSeekMoE(nn.Module):
     def __init__(
         self,
         dim: int,
-        dim_feed_forward: int,
+        dim_ffn: int,
         n_routed_experts: int,
         n_shared_experts: int = 1,
         n_activated_experts: int = 2,
         topk_method: str = "greedy",
-        **kwargs,
     ):
         super().__init__()
         self.dim = dim
@@ -49,10 +48,10 @@ class DeepSeekMoE(nn.Module):
         self.router = Linear(dim, n_routed_experts, bias=False)
 
         self.shared_experts = nn.ModuleList(
-            [MLP(dim, dim_feed_forward) for _ in range(n_shared_experts)]
+            [MLP(dim, dim_ffn) for _ in range(n_shared_experts)]
         )
         self.routed_experts = nn.ModuleList(
-            [MLP(dim, dim_feed_forward) for _ in range(n_routed_experts)]
+            [MLP(dim, dim_ffn) for _ in range(n_routed_experts)]
         )
 
     def forward(self, x: Tensor) -> Tensor:

astrai/model/components/rope.py

@@ -30,7 +30,7 @@ def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
 
 
 class RotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, max_len: int, base: int = 10000):
+    def __init__(self, dim: int, max_len: int, base: float = 10000):
         super().__init__()
         self.dim = dim
         self.max_len = max_len

astrai/model/encoder.py

@@ -0,0 +1,100 @@
+from typing import Any, Mapping, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from astrai.config.model_config import EncoderConfig
+from astrai.model.automodel import AutoModel
+from astrai.model.components.decoder_block import DecoderBlock
+from astrai.model.components.embedding import Embedding
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import RotaryEmbedding
+from astrai.model.transformer import process_attention_mask
+
+
+@AutoModel.register("embedding")
+class EmbeddingEncoder(AutoModel):
+    def __init__(self, config: EncoderConfig):
+        super().__init__(config)
+        self.config = config
+        rope_dim = config.dim // config.n_heads
+        rope_base = config.rope_theta if config.rope_theta is not None else 10000
+        self.rotary_embedding = RotaryEmbedding(rope_dim, config.max_len, rope_base)
+        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+
+        self.layers = nn.ModuleList(
+            [
+                DecoderBlock(
+                    config.dim,
+                    config.n_heads,
+                    config.dim_ffn,
+                    config.n_kv_heads,
+                    config.norm_eps,
+                    config.use_qk_norm,
+                    config.use_gated_attention,
+                    layer_id,
+                )
+                for layer_id in range(config.n_layers)
+            ]
+        )
+
+        self.norm = RMSNorm(config.dim, config.norm_eps)
+
+        self.pooling_type = config.pooling_type or "mean"
+        self.normalize_embeddings = config.normalize_embeddings or False
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if hasattr(module, "reset_parameters"):
+            module.reset_parameters()
+
+    def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
+        state_dict = dict(state_dict)
+        state_dict.pop("lm_head.weight", None)
+        return super().load_state_dict(state_dict, strict=strict, assign=assign)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        input_mask: Optional[Tensor] = None,
+        position_ids: Optional[Tensor] = None,
+    ) -> Tensor:
+        assert input_ids.ndim == 2
+        B, S = input_ids.shape
+
+        x = self.embed_tokens(input_ids)
+
+        if position_ids is None:
+            position_ids = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1)
+
+        rotary_emb = self.rotary_embedding(x, position_ids)
+        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=False)
+
+        for layer in self.layers:
+            x = layer(x, rotary_emb, attn_mask, paged_cache=None)
+
+        hidden_states = self.norm(x)
+
+        if self.pooling_type == "cls":
+            pooled = hidden_states[:, 0]
+        elif self.pooling_type == "last":
+            if input_mask is not None:
+                lengths = input_mask.sum(dim=1) - 1
+                pooled = hidden_states[torch.arange(B, device=x.device), lengths]
+            else:
+                pooled = hidden_states[:, -1]
+        else:
+            if input_mask is not None:
+                mask = input_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
+                pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(
+                    min=1.0
+                )
+            else:
+                pooled = hidden_states.mean(dim=1)
+
+        if self.normalize_embeddings:
+            pooled = torch.nn.functional.normalize(pooled, p=2, dim=-1)
+
+        return pooled

astrai/model/transformer.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
-from astrai.config.model_config import ModelConfig
+from astrai.config.model_config import AutoRegressiveLMConfig
 from astrai.inference.core.cache import KvcacheView
 from astrai.model.automodel import AutoModel
 from astrai.model.components.decoder_block import DecoderBlock
@@ -46,16 +46,20 @@ def process_attention_mask(
     ).masked_fill_(attend.unsqueeze(1), 0.0)
 
 
-@AutoModel.register("transformer")
-class Transformer(AutoModel):
-    """Transformer language model with paged KV cache."""
+@AutoModel.register("autoregressive_lm")
+class AutoRegressiveLM(AutoModel):
+    """Autoregressive language model with paged KV cache."""
 
-    def __init__(self, config: ModelConfig):
+    def __init__(self, config: AutoRegressiveLMConfig):
         super().__init__(config)
         self.config = config
-        self.rotary_embedding = RotaryEmbedding(
-            config.dim // config.n_heads, config.max_len
+        rope_dim = (
+            config.qk_rope_head_dim
+            if config.attn_type == "mla"
+            else config.dim // config.n_heads
         )
+        rope_base = config.rope_theta if config.rope_theta is not None else 10000
+        self.rotary_embedding = RotaryEmbedding(rope_dim, config.max_len, rope_base)
         self.embed_tokens = Embedding(config.vocab_size, config.dim)
 
         self.layers = nn.ModuleList(
@@ -74,7 +78,10 @@ class Transformer(AutoModel):
                     n_routed_experts=config.n_routed_experts,
                     n_shared_experts=config.n_shared_experts,
                     n_activated_experts=config.n_activated_experts,
-                    topk_method=config.moe_topk_method,
+                    topk_method=config.topk_method,
+                    kv_lora_rank=config.kv_lora_rank,
+                    qk_nope_head_dim=config.qk_nope_head_dim,
+                    qk_rope_head_dim=config.qk_rope_head_dim,
                 )
                 for layer_id in range(config.n_layers)
             ]
@@ -83,15 +90,14 @@ class Transformer(AutoModel):
         self.norm = RMSNorm(config.dim, config.norm_eps)
         self.lm_head = Linear(config.dim, config.vocab_size)
 
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
             self.lm_head.weight = self.embed_tokens.weight
 
-        self._init_weights()
+        self.apply(self._init_weights)
 
-    def _init_weights(self):
-        for param in self.parameters():
-            if param.dim() > 1:
-                nn.init.normal_(param, mean=0.0, std=0.006)
+    def _init_weights(self, module):
+        if hasattr(module, "reset_parameters"):
+            module.reset_parameters()
 
     def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
         lm_head_key = "lm_head.weight"
@@ -99,7 +105,7 @@ class Transformer(AutoModel):
 
         state_dict = dict(state_dict)
 
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
             # same tensor for embed and lm_head
             if embed_key in state_dict:
                 state_dict[lm_head_key] = state_dict[embed_key]
@@ -115,7 +121,7 @@ class Transformer(AutoModel):
             destination=destination, prefix=prefix, keep_vars=keep_vars
         )
 
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
             lm_head_key = prefix + "lm_head.weight"
             if lm_head_key in state_dict:
                 del state_dict[lm_head_key]

astrai/parallel/setup.py

@@ -123,6 +123,7 @@ def spawn_parallel_fn(
     master_addr: str = "localhost",
     master_port: str = "29500",
     device_type: str = "cuda",
+    start_method: str = "spawn",
     **kwargs,
 ):
     # clear environment variables
@@ -156,6 +157,10 @@ def spawn_parallel_fn(
         kwargs,
     )
 
-    mp.spawn(
-        wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
+    mp.start_processes(
+        wrapper_spawn_func,
+        args=wrapper_spawn_func_args,
+        nprocs=world_size,
+        start_method=start_method,
+        join=True,
     )

astrai/serialization.py

@@ -1,4 +1,5 @@
 import json
+import time
 from pathlib import Path
 from typing import Any, Dict, Optional
 
@@ -16,11 +17,13 @@ class Checkpoint:
         epoch: int = 0,
         iteration: int = 0,
         extra: Optional[Dict[str, Any]] = None,
+        meta: Optional[Dict[str, Any]] = None,
     ):
         self.state_dict = state_dict
         self.epoch = epoch
         self.iteration = iteration
         self.extra = extra or {}
+        self.meta = meta or {}
 
     def save(
         self,
@@ -35,13 +38,16 @@ class Checkpoint:
             meta = {
                 "epoch": self.epoch,
                 "iteration": self.iteration,
+                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             }
+            meta.update(self.meta)
             with open(save_path / "meta.json", "w") as f:
                 json.dump(meta, f, indent=2)
 
             st.save_file(self.state_dict, save_path / "state_dict.safetensors")
             if self.extra:
-                torch.save(self.extra, save_path / "extra.pt")
+                for key, value in self.extra.items():
+                    torch.save(value, save_path / f"{key}.pt")
 
     @classmethod
     def load(
@@ -64,14 +70,14 @@ class Checkpoint:
 
         state_dict = st.load_file(save_path / "state_dict.safetensors")
 
-        extra = None
-        extra_path = save_path / "extra.pt"
-        if extra_path.exists():
-            extra = torch.load(extra_path, map_location="cpu", weights_only=False)
+        extra = {}
+        for f in save_path.iterdir():
+            if f.suffix == ".pt" and f.stem not in ("meta",):
+                extra[f.stem] = torch.load(f, map_location="cpu", weights_only=False)
 
         return cls(
             state_dict=state_dict,
             epoch=meta["epoch"],
             iteration=meta["iteration"],
-            extra=extra,
+            extra=extra or None,
         )

astrai/tokenize/tokenizer.py

@@ -51,9 +51,26 @@ class AutoTokenizer:
                 self.set_chat_template(config["chat_template"])
 
     @classmethod
-    def from_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoTokenizer":
-        """Load tokenizer from pretrained directory."""
+    def from_pretrained(cls, path: Union[str, Path]) -> "AutoTokenizer":
+        """Load tokenizer from pretrained directory.
+
+        Raises:
+            FileNotFoundError: If tokenizer.json is missing.
+            RuntimeError: If tokenizer failed to initialize.
+        """
+        path = Path(path)
+        tokenizer_file = path / "tokenizer.json"
+        if not tokenizer_file.exists():
+            raise FileNotFoundError(
+                f"Tokenizer file not found: {tokenizer_file}. "
+                "A valid tokenizer.json is required."
+            )
         instance = cls(path)
+        if instance._tokenizer is None:
+            raise RuntimeError(
+                f"Failed to load tokenizer from {path}. "
+                "The tokenizer.json may be corrupted or incompatible."
+            )
         return instance
 
     def save_pretrained(self, save_path: str):

astrai/trainer/__init__.py

@@ -1,3 +1,4 @@
+from astrai.trainer.optim import Muon
 from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 from astrai.trainer.train_callback import (
@@ -9,6 +10,8 @@ from astrai.trainer.trainer import Trainer
 __all__ = [
     # Main trainer
     "Trainer",
+    # Optimizer
+    "Muon",
     # Strategy factory
     "StrategyFactory",
     "BaseStrategy",

astrai/trainer/metric_util.py

@@ -47,6 +47,10 @@ def ctx_get_lr(ctx):
     return ctx.optimizer.param_groups[-1]["lr"]
 
 
+def ctx_get_val_loss(ctx):
+    return ctx.val_loss
+
+
 def ctx_get_grad_norm(ctx):
     return grad_norm(ctx.model)
 

astrai/trainer/optim.py

@@ -0,0 +1,113 @@
+import torch
+from torch.optim import Optimizer
+
+
+def _zeropower_via_newtonschulz(G: torch.Tensor, steps: int = 5):
+    assert G.ndim == 2
+    X = G.bfloat16()
+    scale = max(1, G.size(0) / G.size(1)) ** 0.5
+    X = X / (X.norm() + 1e-7) * scale
+    if steps == 0:
+        return X.type_as(G)
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    for _ in range(steps):
+        A = X @ X.T
+        B = A @ X
+        X = a * X + b * B + c * (A @ B)
+    return X.type_as(G)
+
+
+class Muon(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr: float = 2e-3,
+        momentum: float = 0.95,
+        weight_decay: float = 0.0,
+        nesterov: bool = True,
+        ns_steps: int = 5,
+        adamw_lr: float = None,
+        adamw_betas: tuple = (0.9, 0.95),
+        adamw_eps: float = 1e-8,
+        adamw_wd: float = 0.0,
+    ):
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            ns_steps=ns_steps,
+            adamw_lr=adamw_lr if adamw_lr is not None else lr * 0.1,
+            adamw_betas=adamw_betas,
+            adamw_eps=adamw_eps,
+            adamw_wd=adamw_wd,
+        )
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError("Muon does not support sparse gradients")
+                if p.ndim >= 2:
+                    self._muon_update(p, grad, group)
+                else:
+                    self._adamw_update(p, grad, group)
+        return loss
+
+    def _muon_update(self, p, grad, group):
+        lr = group["lr"]
+        momentum = group["momentum"]
+        wd = group["weight_decay"]
+        nesterov = group["nesterov"]
+        ns_steps = group["ns_steps"]
+        state = self.state[p]
+
+        p.mul_(1 - lr * wd)
+
+        if nesterov:
+            grad = grad.add(p, alpha=wd)
+
+        if "momentum_buffer" not in state:
+            state["momentum_buffer"] = torch.zeros_like(grad)
+        buf = state["momentum_buffer"]
+        buf.lerp_(grad, 1 - momentum)
+
+        update = _zeropower_via_newtonschulz(buf, steps=ns_steps)
+        scale = max(1, p.size(0) / p.size(1)) ** 0.5
+        p.add_(update, alpha=-lr * scale)
+
+    def _adamw_update(self, p, grad, group):
+        lr = group["adamw_lr"]
+        betas = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        wd = group["adamw_wd"]
+        state = self.state[p]
+
+        if not state:
+            state["step"] = 0
+            state["exp_avg"] = torch.zeros_like(p)
+            state["exp_avg_sq"] = torch.zeros_like(p)
+
+        state["step"] += 1
+        exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+        beta1, beta2 = betas
+
+        exp_avg.lerp_(grad, 1 - beta1)
+        exp_avg_sq.lerp_(grad.square(), 1 - beta2)
+
+        step = state["step"]
+        bias1 = 1 - beta1**step
+        bias2 = 1 - beta2**step
+
+        p.mul_(1 - lr * wd)
+        denom = exp_avg_sq.sqrt().div_(bias2**0.5).add_(eps)
+        p.addcdiv_(exp_avg / bias1, denom, value=-lr)

astrai/trainer/train_callback.py

@@ -1,15 +1,21 @@
 import json
+import logging
 import os
+import sys
 import time
 from pathlib import Path
-from typing import Callable, List, Optional, Protocol, runtime_checkable
+from typing import IO, Callable, List, Optional, Protocol, runtime_checkable
 
+import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch.nn.utils import clip_grad_norm_
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
 from tqdm import tqdm
 
 from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
+from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
     ctx_get_grad_max,
@@ -20,9 +26,12 @@ from astrai.trainer.metric_util import (
     ctx_get_grad_std,
     ctx_get_loss,
     ctx_get_lr,
+    ctx_get_val_loss,
 )
 from astrai.trainer.train_context import TrainContext
 
+logger = logging.getLogger(__name__)
+
 
 @runtime_checkable
 class TrainCallback(Protocol):
@@ -79,17 +88,53 @@ class GradientClippingCallback(TrainCallback):
     def __init__(self, max_grad_norm: float):
         self.max_grad_norm = max_grad_norm
 
-    def on_step_end(self, context: TrainContext):
-        _ = context
+    def on_step_begin(self, context: TrainContext):
         clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
 
 
+@CallbackFactory.register("gradient_checkpointing")
+class GradientCheckpointingCallback(TrainCallback):
+    """
+    Activation checkpointing callback — trades compute for memory
+    by recomputing specified module activations during the backward pass.
+
+    Args:
+        modules: Module types to apply checkpointing to.
+    """
+
+    def __init__(self, modules: Optional[List[type]] = None):
+        self.modules = tuple(modules) if modules else ()
+
+    def _enable(self, module: nn.Module):
+        if self.modules and isinstance(module, self.modules):
+            fn = module.forward
+            module._original_forward = fn
+            module.forward = lambda *a, **kw: torch_checkpoint(
+                fn, *a, use_reentrant=False, **kw
+            )
+
+    @staticmethod
+    def _disable(module: nn.Module):
+        if hasattr(module, "_original_forward"):
+            module.forward = module._original_forward
+            del module._original_forward
+
+    def on_train_begin(self, context: TrainContext):
+        context.model.apply(self._enable)
+        logger.info("Gradient checkpointing enabled")
+
+    def on_train_end(self, context: TrainContext):
+        context.model.apply(self._disable)
+
+
 @CallbackFactory.register("checkpoint")
 class CheckpointCallback(TrainCallback):
     """
     Checkpoint callback for trainer.
     """
 
+    extra_keys = ("optimizer", "scheduler")
+
     def __init__(
         self,
         save_dir: str,
@@ -97,12 +142,14 @@ class CheckpointCallback(TrainCallback):
         weight_only: bool = False,
         state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
         save_extra_fn: Optional[Callable[["TrainContext"], dict]] = None,
+        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
     ):
         self.save_dir = save_dir
         self.interval = interval
         self.weight_only = weight_only
         self.state_dict_fn = state_dict_fn
-        self.save_extra_fn = save_extra_fn
+        self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
+        self.load_extra_fn = load_extra_fn or CheckpointCallback.load_extra
         self.last_ckpt_iter = 0
 
     @only_on_rank(0)
@@ -116,17 +163,22 @@ class CheckpointCallback(TrainCallback):
             else context.model.state_dict()
         )
 
-        extra = self.save_extra_fn(context) if self.save_extra_fn else None
+        extra = self.save_extra_fn(context)
         context.checkpoint = Checkpoint(
             state_dict=state_dict,
             epoch=context.epoch,
             iteration=context.iteration,
             extra=extra,
+            meta=context.config.to_dict(),
         )
 
         context.checkpoint.save(save_path)
         self.last_ckpt_iter = context.iteration
 
+    def on_train_begin(self, context: TrainContext):
+        if context.checkpoint and context.checkpoint.extra:
+            self.load_extra_fn(context.checkpoint.extra, context)
+
     def on_batch_end(self, context: TrainContext):
         if context.iteration - self.last_ckpt_iter >= self.interval:
             self._save_checkpoint(context)
@@ -138,6 +190,21 @@ class CheckpointCallback(TrainCallback):
     def on_error(self, context: TrainContext):
         self._save_checkpoint(context)
 
+    @staticmethod
+    def save_extra(context: TrainContext) -> dict:
+        extra = {}
+        for name in CheckpointCallback.extra_keys:
+            obj = getattr(context, name, None)
+            if obj:
+                extra[name] = obj.state_dict()
+        return extra
+
+    @staticmethod
+    def load_extra(extra: dict, context: TrainContext):
+        for name in CheckpointCallback.extra_keys:
+            if name in extra:
+                getattr(context, name).load_state_dict(extra[name])
+
 
 @CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):
@@ -145,8 +212,12 @@ class ProgressBarCallback(TrainCallback):
     Progress bar callback for trainer.
     """
 
-    def __init__(self, num_epoch: int):
+    def __init__(
+        self, num_epoch: int, log_interval: int = 100, file: IO[str] = sys.stdout
+    ):
         self.num_epoch = num_epoch
+        self.log_interval = log_interval
+        self.file = file
         self.progress_bar: tqdm = None
 
     @only_on_rank(0)
@@ -155,16 +226,18 @@ class ProgressBarCallback(TrainCallback):
             context.dataloader,
             desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
             dynamic_ncols=True,
+            file=self.file,
         )
 
     @only_on_rank(0)
     def on_batch_end(self, context: TrainContext):
-        self.progress_bar.set_postfix(
-            {
-                "loss": f"{context.loss:.4f}",
-                "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
-            }
-        )
+        postfix = {
+            "loss": f"{context.loss:.4f}",
+            "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
+        }
+        if context.val_loss > 0:
+            postfix["val_loss"] = f"{context.val_loss:.4f}"
+        self.progress_bar.set_postfix(postfix)
         self.progress_bar.update(1)
 
     @only_on_rank(0)
@@ -196,6 +269,7 @@ class MetricLoggerCallback(TrainCallback):
         self._metric_funcs = {
             "loss": ctx_get_loss,
             "lr": ctx_get_lr,
+            "val_loss": ctx_get_val_loss,
             "grad_norm": ctx_get_grad_norm,
             "grad_std": ctx_get_grad_std,
             "grad_max": ctx_get_grad_max,
@@ -206,7 +280,7 @@ class MetricLoggerCallback(TrainCallback):
 
     def _get_log_data(self, context: TrainContext):
         return {
-            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             "epoch": context.epoch,
             "iter": context.iteration,
             **{m: self._metric_funcs[m](context) for m in self.metrics},
@@ -239,3 +313,43 @@ class MetricLoggerCallback(TrainCallback):
 
     def on_error(self, context):
         self._save_log(context.epoch, context.iteration)
+
+
+@CallbackFactory.register("validation")
+class ValidationCallback(TrainCallback):
+    def _run_validation(self, context: TrainContext):
+        context.model.eval()
+
+        total_loss = 0.0
+        num_batches = 0
+
+        with torch.no_grad():
+            for batch in context.val_dataloader:
+                loss = context.strategy(batch)
+                total_loss += loss.item()
+                num_batches += 1
+
+        avg_loss = total_loss / max(num_batches, 1)
+
+        if context.world_size > 1 and dist.is_initialized():
+            loss_tensor = torch.tensor([avg_loss], device=get_current_device())
+            dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
+            avg_loss = loss_tensor.item()
+
+        context.val_loss = avg_loss
+        context.model.train()
+
+        step_count = context.iteration // context.config.grad_accum_steps
+        logger.info(
+            f"Epoch {context.epoch + 1}, Step {step_count}, Val Loss: {avg_loss:.4f}"
+        )
+
+    def on_step_end(self, context: TrainContext):
+        if context.val_dataloader is None:
+            return
+        cfg = context.config
+        if cfg.val_step <= 0:
+            return
+        step_count = context.iteration // cfg.grad_accum_steps
+        if step_count % cfg.val_step == 0:
+            self._run_validation(context)

astrai/trainer/train_context.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Callable, Optional, Self
+from typing import Optional, Self
 
 import torch.nn as nn
 from torch.optim import Optimizer
@@ -21,10 +21,13 @@ class TrainContext:
     optimizer: Optimizer = field(default=None)
     scheduler: LRScheduler = field(default=None)
     checkpoint: Checkpoint = field(default=None)
+    config: TrainConfig = field(default=None)
 
     epoch: int = field(default=0)
     iteration: int = field(default=0)
     loss: float = field(default=0.0)
+    val_dataloader: DataLoader = field(default=None)
+    val_loss: float = field(default=0.0)
 
     world_size: int = field(default=1)
     rank: int = field(default=0)
@@ -35,11 +38,9 @@ class TrainContextBuilder:
     def __init__(
         self,
         config: TrainConfig,
-        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
     ):
         self.config = config
         self._checkpoint: Optional[Checkpoint] = None
-        self._load_extra_fn = load_extra_fn
 
     def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
         self._checkpoint = checkpoint
@@ -50,6 +51,7 @@ class TrainContextBuilder:
             model=self.config.model,
             world_size=get_world_size(),
             rank=get_rank(),
+            config=self.config,
         )
 
         device = get_current_device()
@@ -71,11 +73,8 @@ class TrainContextBuilder:
         context.optimizer = self.config.optimizer_fn(context.model)
         context.scheduler = self.config.scheduler_fn(context.optimizer)
 
-        if self._checkpoint and self._checkpoint.extra and self._load_extra_fn:
-            self._load_extra_fn(self._checkpoint.extra, context)
-
         cfg = self.config
-        sampler_offset = context.iteration * cfg.batch_size
+        sampler_offset = context.iteration * cfg.batch_per_device
         sampler = ResumableDistributedSampler(
             data_source=cfg.dataset,
             start_epoch=context.epoch,
@@ -84,13 +83,30 @@ class TrainContextBuilder:
         )
         context.dataloader = DataLoader(
             cfg.dataset,
-            batch_size=cfg.batch_size,
+            batch_size=cfg.batch_per_device,
             sampler=sampler,
             num_workers=cfg.num_workers,
             pin_memory=cfg.pin_memory,
             prefetch_factor=cfg.prefetch_factor,
         )
 
+        if cfg.val_dataset is not None:
+            val_sampler = ResumableDistributedSampler(
+                data_source=cfg.val_dataset,
+                start_epoch=0,
+                start_iter=0,
+                seed=cfg.random_seed,
+                shuffle=False,
+            )
+            context.val_dataloader = DataLoader(
+                cfg.val_dataset,
+                batch_size=cfg.batch_per_device,
+                sampler=val_sampler,
+                num_workers=cfg.num_workers,
+                pin_memory=cfg.pin_memory,
+                prefetch_factor=cfg.prefetch_factor,
+            )
+
         context.strategy = StrategyFactory.create(
             model=context.model,
             train_type=self.config.strategy,

astrai/trainer/trainer.py

@@ -1,5 +1,4 @@
 import logging
-from itertools import batched
 from typing import List, Optional
 
 from astrai.config import TrainConfig
@@ -26,17 +25,29 @@ class Trainer:
 
     def _get_default_callbacks(self) -> List[TrainCallback]:
         cfg = self.train_config
-        return [
+        callbacks = [
+            CallbackFactory.create(
+                "gradient_checkpointing",
+                modules=cfg.gradient_checkpointing_modules,
+            ),
+            CallbackFactory.create(
+                "checkpoint",
+                cfg.ckpt_dir,
+                cfg.ckpt_interval,
+                state_dict_fn=cfg.state_dict_fn,
+            ),
+            CallbackFactory.create(
+                "metric_logger",
+                log_dir=cfg.log_dir,
+                save_interval=cfg.ckpt_interval,
+                log_interval=cfg.log_interval,
+                metrics=cfg.metrics,
+            ),
             CallbackFactory.create("progress_bar", cfg.n_epoch),
-            CallbackFactory.create("checkpoint", cfg.ckpt_dir, cfg.ckpt_interval),
-            CallbackFactory.create("metric_logger", cfg.ckpt_dir, cfg.ckpt_interval),
             CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
+            CallbackFactory.create("validation"),
         ]
-
-    def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
-        return (
-            TrainContextBuilder(self.train_config).with_checkpoint(checkpoint).build()
-        )
+        return callbacks
 
     def _call_callbacks(self, method_name: str, context: TrainContext):
         for callback in self.callbacks:
@@ -44,50 +55,36 @@ class Trainer:
             if method:
                 method(context)
 
-    def train(self, checkpoint: Optional[Checkpoint] = None):
-        config = self.train_config
-        spawn_parallel_fn(
-            self._train_impl,
-            backend=config.backend,
-            world_size=config.nprocs,
-            master_addr=config.master_addr,
-            master_port=config.master_port,
-            device_type=config.device_type,
-            checkpoint=checkpoint,
-        )
-
-    def _train_impl(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
-        context = self._build_context(checkpoint)
+    def _trainer_loop(self, checkpoint: Optional[Checkpoint] = None):
+        cfg = self.train_config
+        context = TrainContextBuilder(cfg).with_checkpoint(checkpoint).build()
         self._call_callbacks("on_train_begin", context)
 
         try:
             context.model.train()
-            accumulation_steps = max(self.train_config.accumulation_steps, 1)
+            grad_accum_steps = cfg.grad_accum_steps
 
-            for epoch in range(context.epoch, self.train_config.n_epoch):
+            for epoch in range(context.epoch, cfg.n_epoch):
                 context.epoch = epoch
                 self._call_callbacks("on_epoch_begin", context)
 
-                for steps in batched(context.dataloader, accumulation_steps):
-                    self._call_callbacks("on_step_begin", context)
+                for batch in context.dataloader:
+                    self._call_callbacks("on_batch_begin", context)
+                    loss = context.strategy(batch)
+                    context.loss = loss.item()
+                    stand_loss = loss / grad_accum_steps
+                    stand_loss.backward()
+                    context.iteration += 1
+                    self._call_callbacks("on_batch_end", context)
 
-                    step_batch_nums = len(steps)
-                    for batch in steps:
-                        self._call_callbacks("on_batch_begin", context)
-                        loss = context.strategy(batch)
-                        context.loss = loss.item()
-                        context.iteration += 1
+                    if context.iteration % grad_accum_steps == 0:
+                        self._call_callbacks("on_step_begin", context)
+                        context.optimizer.step()
+                        context.optimizer.zero_grad()
+                        self._call_callbacks("on_step_end", context)
 
-                        stand_loss = loss / step_batch_nums
-                        stand_loss.backward()
-                        self._call_callbacks("on_batch_end", context)
-
-                    self._call_callbacks("on_step_end", context)
-                    context.optimizer.step()
-                    context.optimizer.zero_grad()
-
-                    if context.scheduler:
-                        context.scheduler.step()
+                        if context.scheduler:
+                            context.scheduler.step()
 
                 self._call_callbacks("on_epoch_end", context)
 
@@ -97,3 +94,16 @@ class Trainer:
             raise
         finally:
             self._call_callbacks("on_train_end", context)
+
+    def train(self, checkpoint: Optional[Checkpoint] = None):
+        cfg = self.train_config
+        spawn_parallel_fn(
+            self._trainer_loop,
+            backend=cfg.backend,
+            world_size=cfg.nprocs,
+            master_addr=cfg.master_addr,
+            master_port=cfg.master_port,
+            device_type=cfg.device_type,
+            start_method=cfg.start_method,
+            checkpoint=checkpoint,
+        )

docker-compose.yml

@@ -1,12 +1,13 @@
 services:
   server:
-    build: .
-    image: astrai:latest
+    build:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
       - ./params:/app/params:ro
-      - ./checkpoints:/app/checkpoints
     command: python -m scripts.tools.server --port 8000 --device cuda
     deploy:
       resources:
@@ -25,13 +26,14 @@ services:
 
   server-cpu:
     profiles: [cpu]
-    build: .
-    image: astrai:latest
+    build:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
       - ./params:/app/params:ro
-      - ./checkpoints:/app/checkpoints
     command: python -m scripts.tools.server --port 8000 --device cpu
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8000/health"]

scripts/demo/generate_ar.py

@@ -11,7 +11,6 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
 
 
 def generate_text():
-    # Load model from pretrained
     model = AutoModel.from_pretrained(PARAMETER_ROOT)
     tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
     model.to(device="cuda", dtype=torch.bfloat16)
@@ -22,16 +21,15 @@ def generate_text():
         model=model,
         tokenizer=tokenizer,
     )
-    response = engine.generate(
+    for token in engine.generate(
         prompt=query,
-        stream=False,
+        stream=True,
         max_tokens=2048,
         temperature=0.8,
         top_p=0.95,
         top_k=50,
-    )
-
-    print(response)
+    ):
+        print(token, end="", flush=True)
 
 
 if __name__ == "__main__":

scripts/docker.sh

@@ -16,6 +16,7 @@ NC='\033[0m' # No Color
 IMAGE_NAME="astrai"
 IMAGE_TAG="latest"
 REGISTRY=""
+CONTAINER_ID=""
 
 # Print colored messages
 print_info() {
@@ -175,6 +176,10 @@ main() {
                 PORT="$2"
                 shift 2
                 ;;
+            --container)
+                CONTAINER_ID="$2"
+                shift 2
+                ;;
             --gpu)
                 GPU=true
                 shift
@@ -197,6 +202,7 @@ main() {
                 echo "  --dockerfile FILE  Dockerfile path (default: Dockerfile)"
                 echo "  --context PATH     Build context (default: .)"
                 echo "  --port PORT        Port for run (default: 8000)"
+                echo "  --container ID     Container ID for logs"
                 echo "  --gpu              Enable GPU support"
                 echo "  --help             Show this help message"
                 echo ""
@@ -205,6 +211,7 @@ main() {
                 echo "  $0 build --tag v1.0.0"
                 echo "  $0 run --port 8080"
                 echo "  $0 run --gpu"
+                echo "  $0 logs --container abc123"
                 echo "  $0 push --registry ghcr.io/username"
                 exit 0
                 ;;
@@ -237,7 +244,7 @@ main() {
             show_info
             ;;
         logs)
-            show_logs "$2"
+            show_logs "$CONTAINER_ID"
             ;;
         "")
             print_error "No command specified. Use --help for usage"

scripts/tools/benchmark.py

@@ -1,13 +1,13 @@
-"""Benchmark Transformer with KVCache"""
+"""Benchmark AutoRegressiveLM with KVCache"""
 
 from dataclasses import dataclass
 from typing import Any, Dict
 
 import torch
 
-from astrai.config import ModelConfig
+from astrai.config import AutoRegressiveLMConfig
 from astrai.inference import KVCache
-from astrai.model.transformer import Transformer
+from astrai.model.transformer import AutoRegressiveLM
 
 
 @dataclass
@@ -21,7 +21,7 @@ class BenchmarkResult:
 class GenerationBenchmark:
     def __init__(
         self,
-        config: ModelConfig,
+        config: AutoRegressiveLMConfig,
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
         page_size: int = 128,
@@ -29,7 +29,7 @@ class GenerationBenchmark:
         self.config = config
         self.device = device
         self.dtype = dtype
-        self.model = Transformer(config).to(device=device, dtype=dtype)
+        self.model = AutoRegressiveLM(config).to(device=device, dtype=dtype)
         self.model.eval()
         head_dim = config.dim // config.n_heads
         n_pages = (config.max_len * 4 + page_size - 1) // page_size
@@ -216,7 +216,7 @@ def print_benchmark_result(result: BenchmarkResult):
 
 
 if __name__ == "__main__":
-    config = ModelConfig(
+    config = AutoRegressiveLMConfig(
         vocab_size=10000,
         dim=1536,
         n_heads=24,
@@ -230,7 +230,7 @@ if __name__ == "__main__":
     benchmark = GenerationBenchmark(config)
 
     print("=" * 80)
-    print("Running Transformer Generation Benchmark (KVCache)")
+    print("Running AutoRegressiveLM Generation Benchmark (KVCache)")
     print("=" * 80)
 
     prefill_result = benchmark.run_prefill_benchmark(

scripts/tools/train.py

@@ -8,16 +8,16 @@ import torch.nn as nn
 import torch.optim as optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from astrai.config import ModelConfig, TrainConfig
+from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory
-from astrai.model import Transformer
+from astrai.model import AutoRegressiveLM
 from astrai.parallel import get_rank
 from astrai.trainer import SchedulerFactory, Trainer
 
 
 def parse_args() -> argparse.Namespace:
 
-    parser = argparse.ArgumentParser(description="Train the Transformer model.")
+    parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.")
 
     parser.add_argument(
         "--train_type",
@@ -42,18 +42,20 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--n_epoch", type=int, default=1, help="Number of epochs to train."
     )
-    parser.add_argument("--batch_size", type=int, default=1, help="Batch size per GPU.")
     parser.add_argument(
-        "--accumulation_steps",
+        "--batch_per_device", type=int, default=1, help="Batch size per GPU."
+    )
+    parser.add_argument(
+        "--grad_accum_steps",
         type=int,
         default=1,
         help="Number of iterations between each optimizer step.",
     )
     parser.add_argument(
-        "--warmup_steps",
-        type=int,
-        default=1000,
-        help="Number of warmup steps for LR scheduler.",
+        "--warmup_ratio",
+        type=float,
+        default=0.05,
+        help="Fraction of total steps used for LR warmup.",
     )
     parser.add_argument(
         "--max_lr", type=float, default=3e-4, help="Max learning rate for training."
@@ -68,13 +70,13 @@ def parse_args() -> argparse.Namespace:
         "--adamw_beta1",
         type=float,
         default=0.9,
-        help="Beta values for AdamW optimizer.",
+        help="Beta1 for AdamW optimizer.",
     )
     parser.add_argument(
         "--adamw_beta2",
         type=float,
         default=0.95,
-        help="Beta values for AdamW optimizer.",
+        help="Beta2 for AdamW optimizer.",
     )
     parser.add_argument(
         "--adamw_weight_decay",
@@ -114,7 +116,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--label_smoothing",
         type=float,
-        default=0.1,
+        default=0.05,
         help="cross_entropy function label smoothing parameter",
     )
 
@@ -147,6 +149,13 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--device_type", type=str, default="cuda", help="Device type to use."
     )
+    parser.add_argument(
+        "--start_method",
+        type=str,
+        default="spawn",
+        choices=["spawn", "fork", "forkserver"],
+        help="Multiprocessing start method.",
+    )
 
     args = parser.parse_args()
 
@@ -178,7 +187,26 @@ def create_scheduler(
 
 
 def prepare_checkpoint(model: nn.Module) -> dict:
-    return model.module.state_dict()
+    if isinstance(model, DDP):
+        return model.module.state_dict()
+    return model.state_dict()
+
+
+def compute_total_steps(
+    dataset_len: int,
+    n_epoch: int,
+    batch_per_device: int,
+    nprocs: int,
+    grad_accum_steps: int,
+) -> int:
+
+    def ceil_div(a: int, b: int) -> int:
+        return (a + b - 1) // b
+
+    samples_per_replica = ceil_div(dataset_len, nprocs)
+    batches_per_replica = ceil_div(samples_per_replica, batch_per_device)
+    total_steps = (batches_per_replica // grad_accum_steps) * n_epoch
+    return total_steps
 
 
 def train(
@@ -187,11 +215,11 @@ def train(
     data_root_path: str,
     max_lr: float,
     n_epoch: int,
-    batch_size: int,
+    batch_per_device: int,
     start_epoch: int,
     start_batch: int,
-    accumulation_steps: int,
-    warmup_steps: int,
+    grad_accum_steps: int,
+    warmup_ratio: float,
     ckpt_interval: int,
     ckpt_dir: str,
     dpo_beta: float,
@@ -211,21 +239,20 @@ def train(
     stride: int,
     nprocs: int,
     device_type: str,
+    start_method: str,
 ):
     assert train_type in ["seq", "sft", "dpo", "grpo"]
     assert os.path.exists(param_path)
 
     # Load config
-    config = ModelConfig()
     config_path = os.path.join(param_path, "config.json")
-    if os.path.exists(config_path):
-        config.load(config_path)
+    config = AutoRegressiveLMConfig.from_file(config_path)
 
     if window_size is None:
         window_size = config.max_len
 
-    # Create bare Transformer (for training, no tokenizer needed)
-    model = Transformer(config)
+    # Create bare AutoRegressiveLM (for training, no tokenizer needed)
+    model = AutoRegressiveLM(config)
 
     # Load weights if available
     weights_path = os.path.join(param_path, "model.safetensors")
@@ -236,7 +263,7 @@ def train(
     model = model.to(dtype=torch.bfloat16)
 
     strategy_kwargs = {
-        "dpo_beta": dpo_beta,
+        "beta": dpo_beta,
         "label_smoothing": label_smoothing,
         "clip_eps": grpo_clip_eps,
         "kl_coef": grpo_kl_coef,
@@ -260,13 +287,17 @@ def train(
         },
     )
 
-    total_steps = len(dataset) * n_epoch // (batch_size * nprocs)
+    total_steps = compute_total_steps(
+        len(dataset), n_epoch, batch_per_device, nprocs, grad_accum_steps
+    )
+    warmup_steps = int(warmup_ratio * total_steps)
+
     scheduler_fn = partial(
         create_scheduler,
         **{
             "schedule_type": "cosine",
-            "warmup_steps": warmup_steps,
-            "lr_decay_steps": total_steps - warmup_steps,
+            "warmup_steps": min(warmup_steps, total_steps),
+            "lr_decay_steps": total_steps - min(warmup_steps, total_steps),
         },
     )
 
@@ -278,11 +309,11 @@ def train(
         scheduler_fn=scheduler_fn,
         ckpt_dir=ckpt_dir,
         n_epoch=n_epoch,
-        batch_size=batch_size,
+        batch_per_device=batch_per_device,
         start_epoch=start_epoch,
         start_batch=start_batch,
         ckpt_interval=ckpt_interval,
-        accumulation_steps=accumulation_steps,
+        grad_accum_steps=grad_accum_steps,
         max_grad_norm=max_grad_norm,
         random_seed=random_seed,
         num_workers=num_workers,
@@ -291,6 +322,7 @@ def train(
         parallel_wrapper=ddp_wrap,
         state_dict_fn=prepare_checkpoint,
         device_type=device_type,
+        start_method=start_method,
         extra_kwargs=strategy_kwargs,
     )
 

tests/conftest.py

@@ -8,8 +8,8 @@ import torch
 from tokenizers import Tokenizer, models, pre_tokenizers, trainers
 from torch.utils.data import Dataset
 
-from astrai.config.model_config import ModelConfig
-from astrai.model.transformer import Transformer
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.transformer import AutoRegressiveLM
 from astrai.tokenize import AutoTokenizer
 
 
@@ -104,19 +104,19 @@ def test_tokenizer():
 
 @pytest.fixture(scope="session")
 def test_model():
-    """Session-scoped small Transformer model, created once."""
-    config = ModelConfig(
+    """Session-scoped small AutoRegressiveLM model, created once."""
+    config = AutoRegressiveLMConfig(
         vocab_size=1000,
-        dim=16,
-        n_heads=4,
-        n_kv_heads=2,
-        dim_ffn=32,
-        max_len=1024,
-        n_layers=4,
+        dim=8,
+        n_heads=2,
+        n_kv_heads=1,
+        dim_ffn=16,
+        max_len=64,
+        n_layers=2,
         norm_eps=1e-5,
     )
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = Transformer(config).to(device=device)
+    model = AutoRegressiveLM(config).to(device=device)
 
     return {
         "model": model,
@@ -137,12 +137,12 @@ def base_test_env(test_model, test_tokenizer):
         json.dump(
             {
                 "vocab_size": 1000,
-                "dim": 16,
-                "n_heads": 4,
-                "n_kv_heads": 2,
-                "dim_ffn": 32,
-                "max_len": 1024,
-                "n_layers": 4,
+                "dim": 8,
+                "n_heads": 2,
+                "n_kv_heads": 1,
+                "dim_ffn": 16,
+                "max_len": 64,
+                "n_layers": 2,
                 "norm_eps": 1e-5,
             },
             f,

tests/data/test_checkpoint.py

@@ -35,6 +35,33 @@ def test_single_process():
         assert loaded_checkpoint.iteration == 30
 
 
+def test_checkpoint_with_extra():
+    """Verify extra keys are saved as individual .pt files and loaded back."""
+    model = torch.nn.Linear(10, 5)
+    optimizer = AdamW(model.parameters(), lr=1e-3)
+    optimizer.step()
+
+    extra = {
+        "optimizer": optimizer.state_dict(),
+        "scheduler": {"last_epoch": 5},
+    }
+    checkpoint = Checkpoint(
+        state_dict=model.state_dict(), epoch=1, iteration=10, extra=extra
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        checkpoint.save(tmpdir)
+
+        import os
+
+        assert os.path.exists(os.path.join(tmpdir, "optimizer.pt"))
+        assert os.path.exists(os.path.join(tmpdir, "scheduler.pt"))
+
+        loaded = Checkpoint.load(tmpdir)
+        assert loaded.extra["scheduler"]["last_epoch"] == 5
+        assert "state" in loaded.extra["optimizer"]
+
+
 def simple_training():
     model = torch.nn.Linear(10, 5)
     optimizer = AdamW(model.parameters(), lr=1e-3)

tests/data/test_dataset.py

@@ -10,7 +10,7 @@ from astrai.dataset.storage import (
     BaseSegmentFetcher,
     H5Storage,
     MultiSegmentFetcher,
-    create_storage,
+    StorageFactory,
     detect_format,
     load_json,
     save_h5,
@@ -368,9 +368,9 @@ def test_detect_format_unsupported_file(base_test_env):
 
 
 def test_create_storage_invalid_type():
-    """create_storage raises ValueError for unknown type"""
-    with pytest.raises(ValueError, match="Unknown storage type"):
-        create_storage("parquet")
+    """StorageFactory.create raises ValueError for unknown type"""
+    with pytest.raises(ValueError, match="Unknown component"):
+        StorageFactory.create("parquet")
 
 
 def test_json_pretokenized_without_tokenizer(base_test_env):

tests/inference/test_server.py

@@ -157,5 +157,60 @@ def test_messages_with_system(client, loaded_model):
     assert data["type"] == "message"
 
 
+def test_chat_completions_stop_sequence(client, loaded_model):
+    """POST /v1/chat/completions with stop parameter truncates at stop sequence."""
+
+    async def async_gen():
+        yield "Hello"
+        yield "X"
+        yield "world"
+
+    app.state.engine = loaded_model
+    loaded_model.generate_async.return_value = async_gen()
+    response = client.post(
+        "/v1/chat/completions",
+        json={
+            "messages": [{"role": "user", "content": "Hello"}],
+            "max_tokens": 100,
+            "stream": False,
+            "stop": ["X"],
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    content = data["choices"][0]["message"]["content"]
+    assert "X" in content
+    assert "world" not in content
+
+
+def test_chat_completions_stop_sequence_stream(client, loaded_model):
+    """POST /v1/chat/completions with stop parameter truncates SSE stream."""
+
+    async def async_gen():
+        yield "Hello"
+        yield "X"
+        yield "world"
+
+    app.state.engine = loaded_model
+    loaded_model.generate_async.return_value = async_gen()
+    response = client.post(
+        "/v1/chat/completions",
+        json={
+            "messages": [{"role": "user", "content": "Hello"}],
+            "max_tokens": 100,
+            "stream": True,
+            "stop": ["X"],
+        },
+        headers={"Accept": "text/event-stream"},
+    )
+    assert response.status_code == 200
+    content = response.content.decode("utf-8")
+    assert "Hello" in content
+    assert "world" not in content
+    assert any(
+        "finish_reason" in line for line in content.split("\n") if "stop" in line
+    )
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/module/test_encoder.py

@@ -0,0 +1,166 @@
+import torch
+
+from astrai.config.model_config import EncoderConfig
+from astrai.model.encoder import EmbeddingEncoder
+
+TINY_CONFIG = dict(
+    vocab_size=128,
+    dim=8,
+    n_heads=2,
+    n_kv_heads=1,
+    dim_ffn=16,
+    max_len=64,
+    n_layers=2,
+    norm_eps=1e-5,
+)
+
+
+def test_encoder_forward_mean():
+    config = EncoderConfig(**TINY_CONFIG)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_forward_cls():
+    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "cls"})
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_forward_last():
+    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "last"})
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_forward_with_padding():
+    config = EncoderConfig(**TINY_CONFIG)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+    input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
+    input_mask[:, 4:] = False
+
+    with torch.no_grad():
+        output = model(input_ids, input_mask=input_mask)
+
+    assert output.shape == (batch_size, config.dim)
+    assert not torch.isnan(output).any()
+
+
+def test_encoder_normalize():
+    config = EncoderConfig(
+        **{**TINY_CONFIG, "pooling_type": "mean", "normalize_embeddings": True}
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    norms = output.norm(p=2, dim=-1)
+    assert torch.allclose(norms, torch.ones_like(norms), atol=1e-4)
+
+
+def test_encoder_register():
+    from astrai.model.automodel import AutoModel
+
+    assert AutoModel.is_registered("embedding")
+    cls = AutoModel.get_component_class("embedding")
+    assert cls is EmbeddingEncoder
+
+
+def test_encoder_from_transformer_checkpoint():
+    config = EncoderConfig(**TINY_CONFIG)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = EmbeddingEncoder(config).to(device=device)
+
+    state_dict = model.state_dict()
+    state_dict["lm_head.weight"] = torch.randn(
+        config.vocab_size, config.dim, device=device
+    )
+
+    new_model = EmbeddingEncoder(config).to(device=device)
+    new_model.load_state_dict(state_dict, strict=True)
+
+    for key in model.state_dict():
+        assert torch.equal(new_model.state_dict()[key], model.state_dict()[key])
+
+
+def test_encoder_save_load():
+    import json
+    import os
+    import tempfile
+
+    import safetensors.torch as st
+
+    test_dir = tempfile.mkdtemp(prefix="encoder_test_")
+    config_path = os.path.join(test_dir, "config.json")
+    weights_path = os.path.join(test_dir, "model.safetensors")
+
+    try:
+        config_data = {**TINY_CONFIG, "pooling_type": "mean"}
+        with open(config_path, "w") as f:
+            json.dump(config_data, f)
+
+        config = EncoderConfig.from_file(config_path)
+        original = EmbeddingEncoder(config)
+        st.save_file(original.state_dict(), weights_path)
+
+        loaded = EmbeddingEncoder(config)
+        loaded.load_state_dict(st.load_file(weights_path))
+
+        for key in original.state_dict():
+            assert torch.equal(original.state_dict()[key], loaded.state_dict()[key])
+    finally:
+        if os.path.exists(test_dir):
+            for f in os.listdir(test_dir):
+                os.remove(os.path.join(test_dir, f))
+            os.rmdir(test_dir)

tests/module/test_forward_configs.py

@@ -0,0 +1,108 @@
+import pytest
+import torch
+
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.transformer import AutoRegressiveLM
+
+TINY_CONFIG = dict(
+    vocab_size=128,
+    dim=8,
+    n_heads=2,
+    n_kv_heads=1,
+    dim_ffn=16,
+    max_len=64,
+    n_layers=2,
+    norm_eps=1e-5,
+)
+
+
+CONFIGS = [
+    pytest.param(
+        {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp"},
+        id="gqa_mlp",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "mla",
+            "ffn_type": "mlp",
+            "kv_lora_rank": 4,
+            "qk_nope_head_dim": 2,
+            "qk_rope_head_dim": 2,
+        },
+        id="mla_mlp",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+        },
+        id="gqa_moe",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "mlp",
+            "rope_theta": 100000.0,
+        },
+        id="gqa_rope_theta",
+    ),
+    pytest.param(
+        {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp", "use_qk_norm": True},
+        id="gqa_qk_norm",
+    ),
+    pytest.param(
+        {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp", "tie_weight": True},
+        id="gqa_tie_weight",
+    ),
+]
+
+
+@pytest.mark.parametrize("config_kwargs", CONFIGS)
+def test_model_forward(config_kwargs):
+    config = AutoRegressiveLMConfig(**config_kwargs)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoRegressiveLM(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+
+    with torch.no_grad():
+        output = model(input_ids)
+
+    assert "logits" in output
+    assert "hidden_states" in output
+    assert output["logits"].shape == (batch_size, seq_len, config.vocab_size)
+    assert output["hidden_states"].shape == (batch_size, seq_len, config.dim)
+    assert not torch.isnan(output["logits"]).any()
+    assert not torch.isnan(output["hidden_states"]).any()
+
+
+@pytest.mark.parametrize("config_kwargs", CONFIGS)
+def test_model_forward_with_padding(config_kwargs):
+    config = AutoRegressiveLMConfig(**config_kwargs)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoRegressiveLM(config).to(device=device)
+    model.eval()
+
+    batch_size, seq_len = 2, 8
+    input_ids = torch.randint(
+        0, config.vocab_size, (batch_size, seq_len), device=device
+    )
+    input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
+    input_mask[:, 4:] = False
+
+    with torch.no_grad():
+        output = model(input_ids, input_mask=input_mask)
+
+    assert output["logits"].shape == (batch_size, seq_len, config.vocab_size)
+    assert not torch.isnan(output["logits"]).any()

tests/module/test_tie_weight.py

@@ -6,8 +6,8 @@ import pytest
 import safetensors.torch as st
 import torch
 
-from astrai.config.model_config import ModelConfig
-from astrai.model.transformer import Transformer
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.transformer import AutoRegressiveLM
 
 
 @pytest.fixture
@@ -17,10 +17,10 @@ def transformer_test_env():
 
     config = {
         "vocab_size": 1000,
-        "dim": 128,
-        "n_heads": 4,
-        "n_kv_heads": 2,
-        "dim_ffn": 256,
+        "dim": 8,
+        "n_heads": 2,
+        "n_kv_heads": 1,
+        "dim_ffn": 16,
         "max_len": 64,
         "n_layers": 2,
         "norm_eps": 1e-5,
@@ -50,8 +50,8 @@ def test_tie_weight_init(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    config = ModelConfig().load(config_path)
-    model = Transformer(config)
+    config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(config)
 
     assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
     assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
@@ -68,8 +68,8 @@ def test_tie_weight_init(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    config = ModelConfig().load(config_path)
-    model = Transformer(config)
+    config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(config)
 
     assert not torch.equal(model.lm_head.weight, model.embed_tokens.weight)
     assert model.lm_head.weight.data_ptr() != model.embed_tokens.weight.data_ptr()
@@ -94,13 +94,13 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    config = ModelConfig().load(config_path)
-    original_model = Transformer(config)
+    config = AutoRegressiveLMConfig.from_file(config_path)
+    original_model = AutoRegressiveLM(config)
 
     st.save_file(original_model.state_dict(), model_path)
 
-    loaded_config = ModelConfig().load(config_path)
-    model = Transformer(loaded_config)
+    loaded_config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(loaded_config)
     model.load_state_dict(st.load_file(model_path))
 
     assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
@@ -112,8 +112,8 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 
-    loaded_config = ModelConfig().load(config_path)
-    model = Transformer(loaded_config)
+    loaded_config = AutoRegressiveLMConfig.from_file(config_path)
+    model = AutoRegressiveLM(loaded_config)
     model.load_state_dict(st.load_file(model_path))
 
     assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)

tests/trainer/conftest.py

@@ -31,8 +31,8 @@ def create_train_config(
     device: str,
     strategy: str = "seq",
     n_epoch: int = 1,
-    batch_size: int = 2,
-    accumulation_steps: int = 1,
+    batch_per_device: int = 2,
+    grad_accum_steps: int = 1,
     max_grad_norm: float = 1.0,
     ckpt_interval: int = 5,
     random_seed: int = 42,
@@ -47,8 +47,8 @@ def create_train_config(
         device: Device type ("cuda" or "cpu")
         strategy: Training strategy type (default: "seq")
         n_epoch: Number of epochs (default: 1)
-        batch_size: Batch size (default: 2)
-        accumulation_steps: Gradient accumulation steps (default: 1)
+        batch_per_device: Batch size per device (default: 2)
+        grad_accum_steps: Gradient accumulation steps (default: 1)
         max_grad_norm: Maximum gradient norm for clipping (default: 1.0)
         ckpt_interval: Checkpoint save interval in iterations (default: 5)
         random_seed: Random seed for reproducibility (default: 42)
@@ -74,9 +74,9 @@ def create_train_config(
         scheduler_fn=scheduler_fn,
         ckpt_dir=test_dir,
         n_epoch=n_epoch,
-        batch_size=batch_size,
+        batch_per_device=batch_per_device,
         ckpt_interval=ckpt_interval,
-        accumulation_steps=accumulation_steps,
+        grad_accum_steps=grad_accum_steps,
         max_grad_norm=max_grad_norm,
         random_seed=random_seed,
         device_type=device,

tests/trainer/test_callbacks.py

@@ -1,11 +1,130 @@
 import torch
 
 from astrai.config.train_config import TrainConfig
+from astrai.model.components.decoder_block import DecoderBlock
 from astrai.trainer.schedule import SchedulerFactory
-from astrai.trainer.train_callback import TrainCallback
+from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
 from astrai.trainer.trainer import Trainer
 
 
+def test_gradient_checkpointing_enable_disable(test_model):
+    """Enable wraps forward, _disable restores it."""
+    model = test_model["model"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    originals = [layer.forward for layer in model.layers]
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    for layer in model.layers:
+        assert hasattr(layer, "_original_forward")
+        assert layer.forward is not originals[0]
+
+    for layer in model.layers:
+        callback._disable(layer)
+
+    for layer in model.layers:
+        assert not hasattr(layer, "_original_forward")
+
+
+def test_gradient_checkpointing_empty_modules_noop(test_model):
+    """modules=None should leave forwards untouched."""
+    model = test_model["model"]
+    callback = GradientCheckpointingCallback()
+
+    originals = [layer.forward for layer in model.layers]
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    for layer, orig in zip(model.layers, originals):
+        assert layer.forward is orig
+
+
+def test_gradient_checkpointing_forward_unchanged(test_model):
+    """Forward output unchanged after patching (no_grad)."""
+    model = test_model["model"]
+    device = test_model["device"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    input_ids = torch.randint(0, 1000, (2, 32)).to(device)
+
+    with torch.no_grad():
+        ref = model(input_ids)["logits"].clone()
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    with torch.no_grad():
+        out = model(input_ids)["logits"]
+
+    assert torch.equal(ref, out)
+
+
+def test_gradient_checkpointing_backward(test_model):
+    """backward passes gradients through checkpointed layers."""
+    model = test_model["model"]
+    device = test_model["device"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    input_ids = torch.randint(0, 1000, (2, 32)).to(device)
+    target_ids = torch.randint(0, 1000, (2, 32)).to(device)
+
+    logits = model(input_ids)["logits"]
+    loss = torch.nn.functional.cross_entropy(
+        logits.flatten(0, 1).float(), target_ids.flatten()
+    )
+    loss.backward()
+
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            assert param.grad is not None, f"{name} gradient is None"
+
+    for layer in model.layers:
+        callback._disable(layer)
+
+    model.zero_grad()
+    for name, p in model.named_parameters():
+        assert p.grad is None or p.grad.sum().item() == 0, f"{name} grad not zeroed"
+
+
+def test_gradient_checkpointing_trainer_integration(base_test_env, random_dataset):
+    """Gradient checkpointing runs end-to-end via Trainer."""
+
+    def optimizer_fn(model):
+        return torch.optim.AdamW(model.parameters())
+
+    def scheduler_fn(optim):
+        return SchedulerFactory.create(
+            optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
+        )
+
+    train_config = TrainConfig(
+        model=base_test_env["model"],
+        strategy="seq",
+        dataset=random_dataset,
+        optimizer_fn=optimizer_fn,
+        scheduler_fn=scheduler_fn,
+        ckpt_dir=base_test_env["test_dir"],
+        n_epoch=1,
+        batch_per_device=2,
+        ckpt_interval=3,
+        grad_accum_steps=1,
+        max_grad_norm=1.0,
+        random_seed=42,
+        device_type=base_test_env["device"],
+        gradient_checkpointing_modules=[DecoderBlock],
+    )
+
+    trainer = Trainer(train_config)
+    trainer.train()
+    # no crash = callback correctly enabled/disabled
+
+
 def test_callback_integration(base_test_env, random_dataset):
     """Test that all callbacks are properly integrated"""
 
@@ -25,9 +144,9 @@ def test_callback_integration(base_test_env, random_dataset):
         scheduler_fn=scheduler_fn,
         ckpt_dir=base_test_env["test_dir"],
         n_epoch=1,
-        batch_size=2,
+        batch_per_device=2,
         ckpt_interval=3,
-        accumulation_steps=1,
+        grad_accum_steps=1,
         max_grad_norm=1.0,
         random_seed=42,
         device_type=base_test_env["device"],

tests/trainer/test_early_stopping.py

@@ -28,9 +28,9 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
         dataset=early_stopping_dataset,
         ckpt_dir=base_test_env["test_dir"],
         n_epoch=2,
-        batch_size=2,
+        batch_per_device=2,
         ckpt_interval=1,
-        accumulation_steps=2,
+        grad_accum_steps=2,
         random_seed=np.random.randint(1e4),
         device_type=base_test_env["device"],
     )

tests/trainer/test_trainer.py

@@ -7,45 +7,45 @@ def test_different_batch_sizes(base_test_env, random_dataset, train_config_facto
     """Test training with different batch sizes"""
     batch_sizes = [1, 2, 4, 8]
 
-    for batch_size in batch_sizes:
+    for batch_per_device in batch_sizes:
         train_config = train_config_factory(
             model=base_test_env["model"],
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
-            batch_size=batch_size,
+            batch_per_device=batch_per_device,
         )
 
-        assert train_config.batch_size == batch_size
+        assert train_config.batch_per_device == batch_per_device
 
 
 def test_gradient_accumulation(base_test_env, random_dataset, train_config_factory):
     """Test training with different gradient accumulation steps"""
-    accumulation_steps_list = [1, 2, 4]
+    grad_accum_steps_list = [1, 2, 4]
 
-    for accumulation_steps in accumulation_steps_list:
+    for grad_accum_steps in grad_accum_steps_list:
         train_config = train_config_factory(
             model=base_test_env["model"],
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
-            batch_size=2,
-            accumulation_steps=accumulation_steps,
+            batch_per_device=2,
+            grad_accum_steps=grad_accum_steps,
         )
 
         trainer = Trainer(train_config)
         trainer.train()
 
-        assert train_config.accumulation_steps == accumulation_steps
+        assert train_config.grad_accum_steps == grad_accum_steps
 
 
 def test_memory_efficient_training(base_test_env, random_dataset, train_config_factory):
     """Test training with memory-efficient configurations"""
     # Test with smaller batch sizes and gradient checkpointing
     small_batch_configs = [
-        {"batch_size": 1, "accumulation_steps": 8},
-        {"batch_size": 2, "accumulation_steps": 4},
-        {"batch_size": 4, "accumulation_steps": 2},
+        {"batch_per_device": 1, "grad_accum_steps": 8},
+        {"batch_per_device": 2, "grad_accum_steps": 4},
+        {"batch_per_device": 4, "grad_accum_steps": 2},
     ]
 
     for config in small_batch_configs:
@@ -54,8 +54,9 @@ def test_memory_efficient_training(base_test_env, random_dataset, train_config_f
             dataset=random_dataset,
             test_dir=base_test_env["test_dir"],
             device=base_test_env["device"],
-            batch_size=config["batch_size"],
-            accumulation_steps=config["accumulation_steps"],
+            batch_per_device=config["batch_per_device"],
+            grad_accum_steps=config["grad_accum_steps"],
         )
 
-        assert train_config.accumulation_steps == config["accumulation_steps"]
+        assert train_config.grad_accum_steps == config["grad_accum_steps"]
+        assert train_config.batch_per_device == config["batch_per_device"]