perf: Muon step 改用 torch._foreach_* 批处理并移除 NS 迭代的冗余 bf16 转换

fix: decode后task_extend失败时提前中止，scheduler崩溃时通知waiting任务
fix: 修复 to_dict list 类型丢失与 OpenAI stop 参数失效
2026-05-23 19:50:12 +08:00 · 2026-05-20 19:23:13 +08:00 · 2026-05-19 21:07:07 +08:00 · 2026-05-19 19:12:38 +08:00 · 2026-05-19 17:50:24 +08:00 · 2026-05-19 17:34:01 +08:00
62 changed files with 2786 additions and 1365 deletions
@@ -2,7 +2,7 @@
 name: Bug report
 about: Create a report to help us improve
 title: "[BUG]"
-labels: enhancement
+labels: bug
 assignees: ''

 ---
@@ -16,9 +16,9 @@ Please delete options that are not relevant.
 Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce.

 ## Checklist:
- [ ] My code follows the style guidelines of this project (run `ruff format .` and `ruff check --fix .`)
+- [ ] My code follows the style guidelines of this project (run `ruff format .` and `ruff check . --select I`)
 - [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
+- [ ] Code is self-documenting (no unnecessary comments)
 - [ ] I have made corresponding changes to the documentation
 - [ ] My changes generate no new warnings
 - [ ] I have added tests that prove my fix is effective or that my feature works
@@ -1,68 +1,100 @@
 # Contributing to AstrAI

-Thank you for your interest in contributing to AstrAI! This document provides guidelines and steps for contributing.
+Thank you for your interest in contributing! This document provides step-by-step guidelines.

-## How to Contribute
+## Quick Start

-### Reporting Issues
-If you encounter a bug or have a feature request, please open an issue on GitHub. Include as much detail as possible:
- A clear description of the problem or request.
- Steps to reproduce (for bugs).
- Your environment (Python version, OS, etc.).
+```bash
+git clone https://github.com/your-username/AstrAI.git
+cd AstrAI
+pip install -e ".[dev]"     # install with dev dependencies (pytest, ruff)
+```

-### Submitting Changes
-1. **Fork** the repository.
-2. **Clone** your fork:
-   ```bash
-   git clone https://github.com/your-username/AstrAI.git
-   cd AstrAI
-   ```
-3. **Create a feature branch**:
-   ```bash
-   git checkout -b feature/your-feature-name
-   ```
-4. **Make your changes**. Follow the code style guidelines below.
-5. **Commit your changes** with a descriptive commit message:
-   ```bash
-   git commit -m "Add: brief description of the change"
-   ```
-6. **Push** to your fork:
-   ```bash
-   git push origin feature/your-feature-name
-   ```
-7. **Open a Pull Request** (PR) against the `main` branch of the upstream repository.
+## Before You Commit

-## Code Style
+Run the following checks **in order** — CI will reject if any fail.

-AstrAI uses [Ruff](https://docs.astral.sh/ruff/) for code formatting and linting. Please ensure your code is formatted before submitting.
+### 1. Format

- Run Ruff to format and lint (requires conda environment `nlp`):
-  ```bash
-  conda run -n nlp ruff format .
-  conda run -n nlp ruff check --fix .
-  ```
- The project uses **double quotes** for strings and **4‑space indentation** (as configured in `pyproject.toml`).
+```bash
+ruff format .
+```

-## Testing
+> **Note**: `ruff format` may rename parameters (e.g. `mask` → `attn_mask`).  
+> Always review the diff after formatting.

-If you add or modify functionality, please include appropriate tests.
+### 2. Import sorting

- Run the test suite with:
-  ```bash
-  conda run -n nlp python -u -m pytest
-  ```
- Ensure all tests pass before submitting your PR.
+```bash
+ruff check . --select I
+```
+
+If this fails, **manually fix** import ordering (ruff does not auto-fix in this project's CI):
+
+```bash
+ruff check . --select I --fix .
+ruff format .    # re-format after fix
+```
+
+### 3. Run tests
+
+```bash
+python -u -m pytest tests/ -v
+```
+
+> Failed tests may leave orphan tempdirs under `%TEMP%`. Clean them manually if needed.
+
+### 4. (Optional) Full pre-commit check
+
+If you have Git Bash available:
+
+```bash
+bash scripts/pre_commit.sh
+```
+
+This runs format check, import sort check, and tests in one go.
+
+## Commit Style
+
+```
+fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description (~50 chars)
+
+- bullet point body (each ~60 chars)
+```
+
+- **Type** must be one of: `fix`, `feat`, `chore`, `docs`, `refactor`, `perf`, `test`, `style`, `ci`, `build`, `revert`.
+- **Subject line** ends with no period.
+- **Body** uses bullet points starting with `-`.
+- No `(scope)` parentheses.
+
+## Common Issues
+
+| Problem | Cause | Fix |
+|---------|-------|-----|
+| `ruff check --select I` fails | Wrong import order | `ruff check . --select I --fix .` then `ruff format .` |
+| `ruff format` changed many files | Not formatted before commit | Review diff carefully before staging |
+| Pre-commit hook rejects | Tests or lint failed | Fix individually, do not `--no-verify` |
+| Tests fail with tempdir left | Test crash | Clean `%TEMP%` manually |
+
+## Submitting Changes
+
+1. Fork the repo.
+2. Create a feature branch: `git checkout -b feat/my-feature`
+3. Make changes following the steps above.
+4. Commit with the commit style above.
+5. Push: `git push origin feat/my-feature`
+6. Open a Pull Request against `main`.

 ## Code Review

-All submissions will be reviewed. We may request changes or discuss alternatives. Please be responsive to feedback.
+- All PRs are reviewed. We may request changes.
+- CI runs `ruff format --check .` then `ruff check . --select I` (no `--fix` in CI).
+- Ensure all tests pass.

 ## License

-By contributing, you agree that your contributions will be licensed under the same [GPL-3.0 License](LICENSE) that covers the project.
+By contributing, you agree that your contributions will be licensed under the [GPL-3.0 License](LICENSE).

 ---

-If you have any questions, feel free to ask in the [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions) or open an issue.
-
-Happy contributing!
+Questions? Ask in [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions) or open an issue.
@@ -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
@@ -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,16 +213,17 @@ 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

 | Document | Description |
 |----------|-------------|
 | [Parameter Guide](./assets/docs/params.md) | Training & inference parameters |
-| [Design Document](./assets/docs/design.md) | Framework architecture & module design |
-| [Data Flow](./assets/docs/dataflow.md) | Data processing pipeline details |
-| [Model Introduction](./assets/docs/introduction.md) | Model architecture & technical details |
+| [Architecture](./assets/docs/architecture.md) | System architecture, class diagram & design patterns |
+| [Training](./assets/docs/training.md) | Training loop, strategies & formulas |
+| [Inference](./assets/docs/inference.md) | KVCache, continuous batching, sampling & HTTP API |
+| [Data Flow](./assets/docs/dataflow.md) | Data pipeline, storage backends & dataset architecture |

 ### Contributing

@@ -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,16 +219,17 @@ 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) 上的视频演示。

 ### 文档

 | 文档 | 说明 |
 |------|------|
 | [参数说明](./params.md) | 训练与推理参数配置 |
-| [设计文档](./design.md) | 系统架构与模块设计 |
-| [数据流程](./dataflow.md) | 数据处理管道详解 |
-| [模型介绍](./introduction.md) | 模型架构与技术细节 |
+| [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
+| [训练文档](./training.md) | 训练循环、策略与公式 |
+| [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
+| [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 |

 ### 贡献

@@ -1,15 +1,22 @@
-## 1. Why I Created This Project
+# AstrAI Architecture

-There are many large language models on the market today, such as GPT, LLaMA, and others, with tens of billions or even hundreds of billions of parameters. But honestly, these models have extremely high hardware requirements, making them inaccessible for ordinary developers. I thought: **Can we create a model that is both useful and can run on ordinary computers?** This is also what most people currently hope for - a locally deployable AI project that achieves complete privatization while maintaining some level of intelligence.
-
-Thus, the AstrAI project was born - 1B parameters, Chinese-English bilingual, supporting dialogue, text generation, and the training code is open source!
-
-## 2. System Architecture
+## Class Diagram

 ```mermaid
 classDiagram
    namespace config {
-        class ModelConfig {
+        class BaseConfig {
+            +to_dict() Dict
+            +from_dict(d) Self
+        }
+
+        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
@@ -18,12 +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
-            +load(config_path) ModelConfig
-            +save(config_path)
+            +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
+            +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 {
@@ -33,16 +69,20 @@ 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
-            +int prefetch_factor
+            +Optional[int] prefetch_factor
            +bool pin_memory
            +int nprocs
            +str backend
@@ -50,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()
        }
@@ -61,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__()
@@ -122,11 +165,17 @@ classDiagram
            +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
        }
    }

@@ -135,6 +184,8 @@ classDiagram
            +dict state_dict
            +int epoch
            +int iteration
+            +dict extra
+            +dict meta
            +save(save_dir)
            +load(save_dir) Checkpoint
        }
@@ -142,31 +193,43 @@ 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) Tensor
+            +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 {
-            +GQA attention
+            +nn.Module attention  # GQA or MLA via AttnFactory
            +RMSNorm input_norm
-            +MLP mlp
+            +nn.Module mlp        # MLP or DeepSeekMoE via FFNFactory
            +RMSNorm post_attention_norm
            +forward(x, rotary_emb, attention_mask, paged_cache) Tensor
        }
@@ -175,8 +238,13 @@ classDiagram
            +int n_heads
            +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
-            +RMSNorm q_norm, k_norm
+            +Linear gate  # only if use_gated_attention
+            +RMSNorm q_norm, k_norm  # only if use_qk_norm
            +forward(x, rotary_emb, attn_mask, paged_cache) Tensor
        }

@@ -187,8 +255,12 @@ classDiagram
            +int kv_lora_rank
            +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
+            +Linear gate  # only if use_gated_attention
            +RMSNorm kv_norm
            +forward(x, rotary_emb, attn_mask, paged_cache) Tensor
        }
@@ -198,15 +270,36 @@ classDiagram
            +forward(x) Tensor
        }

+        class DeepSeekMoE {
+            +int dim
+            +int n_routed_experts
+            +int n_shared_experts
+            +int n_activated_experts
+            +str topk_method
+            +Linear router
+            +ModuleList shared_experts
+            +ModuleList routed_experts
+            +forward(x) Tensor
+        }
+
+        class AttnFactory {
+            +create(attn_type, **kwargs) nn.Module
+        }
+
+        class FFNFactory {
+            +create(ffn_type, dim, dim_ffn, **kwargs) nn.Module
+        }
+
        class RMSNorm {
            +Parameter weight
            +float norm_eps
+            +tuple normalized_shape
            +forward(x) Tensor
        }

        class Linear {
            +Parameter weight
-            +Parameter bias
+            +Optional[Parameter] bias  # only if bias=True
            +forward(x) Tensor
        }

@@ -264,7 +357,6 @@ classDiagram
            +TrainConfig train_config
            +List[TrainCallback] callbacks
            +train(checkpoint)
-            +_build_context(checkpoint) TrainContext
            +_get_default_callbacks() List[TrainCallback]
        }

@@ -275,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 {
@@ -289,7 +385,7 @@ classDiagram
        }

        class BaseStrategy {
-            +nn.Module model
+            +Union[Callable, nn.Module] model
            +str device
            +compute_loss(batch) Tensor
        }
@@ -297,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 {
@@ -325,6 +421,7 @@ classDiagram
            +str reduction
            +int sync_interval
            +compute_loss(batch) Tensor
+            +sync_ref_model()
        }

        class BaseScheduler {
@@ -352,6 +449,7 @@ classDiagram
        }

        class TrainCallback {
+            <<protocol>>
            +on_train_begin(context)
            +on_train_end(context)
            +on_epoch_begin(context)
@@ -368,17 +466,32 @@ classDiagram
            +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)
@@ -387,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 {
@@ -396,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 {
@@ -410,15 +539,21 @@ classDiagram
            +shutdown()
        }

-        class InferenceScheduler {
-            +nn.Module model
+        class Executor {
+            +AutoModel model
            +AutoTokenizer tokenizer
+            +KVCache page_cache
+            +execute_prefill(tasks, prompt_len, start_pos)
+            +execute_decode(tasks) List[int]
+        }
+
+        class InferenceScheduler {
            +KVCache _page_cache
-            +int max_batch_size
-            +int max_seq_len
-            +int max_prompt_len
-            +int page_size
+            +Executor _executor
            +TaskManager _task_mgr
+            +bool _running
+            +Thread _loop_thread
+            +int max_seq_len
            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
            +remove_task(task_id)
            +start()
@@ -428,8 +563,8 @@ classDiagram

        class Allocator {
            +int _free_mask
-            +int refs_count
-            +LRU _lru
+            +List[int] _refs
+            +OrderedDict _lru
            +alloc() int
            +free(idx, keep_cached)
            +inc_ref(idx)
@@ -456,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)
@@ -523,6 +655,19 @@ classDiagram
            ABORTED
        }

+        class TaskManager {
+            +AutoTokenizer tokenizer
+            +Deque waiting_queue
+            +List active_tasks
+            +add_task(prompt, **kwargs) str
+            +remove_task(task_id) List[Task]
+            +remove_finished_tasks(stop_ids) List[Task]
+            +pull_candidates(n) List[Task]
+            +activate(task)
+            +return_to_waiting(tasks)
+            +get_active_tasks() List[Task]
+        }
+
        class GenerationRequest {
            +List[Dict] messages
            +int top_k
@@ -553,7 +698,7 @@ classDiagram
        }

        class SamplingPipeline {
-            +List strategies
+            +List[BaseSamplingStrategy] strategies
            +apply(logits, filter_value) Tensor
            +sample(logits, filter_value) Tensor
        }
@@ -564,9 +709,9 @@ classDiagram
            +List[bool] _done
            +append(token, idx)
            +get_results() List[str]
-            +pop_all() List[str]
+            +pop_all() List[Tuple[int, str]]
            +wait(timeout) bool
-            +wait_completion()
+            +wait_completion(timeout)
        }

        class ChatMessage {
@@ -575,24 +720,97 @@ classDiagram
        }

        class ChatCompletionRequest {
+            +str model
            +List[ChatMessage] messages
+            +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 {
+            +str role
+            +Union[str, List[Dict]] content
+        }
+
+        class MessagesRequest {
+            +str model
+            +List[AnthropicMessage] messages
+            +Optional[str] system
            +float temperature
            +float top_p
            +int top_k
            +int max_tokens
            +bool stream
-            +Optional[str] stop
-            +Optional[int] n
+            +Optional[List[str]] stop_sequences
+        }
+
+        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]
+            +format_non_stream_response(ctx, content) Dict
+            +handle() Union[StreamingResponse, Dict]
+        }
+
+        class OpenAIHandler {
+            +build_prompt() str
+            +create_response_id() str
+        }
+
+        class AnthropicHandler {
+            +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
+        }
+
+        class StreamContext {
+            +str resp_id
+            +int created
+            +str model
+            +int prompt_tokens
+            +int completion_tokens
+            +str accumulated
+            +Optional[str] stop_matched
+            +str last_yield_trimmed
+        }
+
+        class app {
+            <<singleton>>
+            +FastAPI app
        }
    }

    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 {
@@ -610,170 +828,172 @@ classDiagram
        }
    }

-    %% Relationships
-    TrainConfig --> BaseDataset : uses
-    TrainConfig ..> BaseStrategy : selects
-    StrategyFactory ..> BaseStrategy : creates
+    %% Relationships — UML notation: <|-- generalization, *-- composition, o-- aggregation, --> association, ..> dependency
+
+    %% --- Generalization (inheritance) ---
    BaseStrategy <|-- SEQStrategy
    BaseStrategy <|-- SFTStrategy
    BaseStrategy <|-- DPOStrategy
    BaseStrategy <|-- GRPOStrategy
-    DPOStrategy --> Transformer : uses
-    GRPOStrategy --> Transformer : uses
-    Trainer --> TrainConfig : uses
-    Trainer --> TrainContextBuilder : uses
-    Trainer --> TrainCallback : manages
-    TrainContextBuilder --> TrainContext : creates
-    TrainContextBuilder --> StrategyFactory : uses
-    Checkpoint ..> Checkpoint : serializes
-    TrainContext --> Checkpoint : manages
-    TrainContext --> BaseStrategy : uses
-    TrainContext --> BaseScheduler : uses
-    SchedulerFactory ..> BaseScheduler : creates
    BaseScheduler <|-- CosineScheduler
    BaseScheduler <|-- SGDRScheduler
-    CallbackFactory ..> TrainCallback : creates
    TrainCallback <|-- GradientClippingCallback
+    TrainCallback <|-- GradientCheckpointingCallback
    TrainCallback <|-- CheckpointCallback
    TrainCallback <|-- ProgressBarCallback
    TrainCallback <|-- MetricLoggerCallback
-    PagePool --> Allocator : composes
-    PagePool --> PrefixCache : composes
-    KVCache --> PagePool : composes
-    KVCache --> Storage : composes
-    KVCache --> TaskTable : composes
-    KvcacheView --> Storage : wraps
-    InferenceEngine --> InferenceScheduler : uses
-    InferenceEngine --> GenerationRequest : uses
-    InferenceEngine --> GenerateResult : creates
-    InferenceScheduler --> Task : manages
-    InferenceScheduler --> TaskStatus : uses
-    InferenceScheduler --> KVCache : uses
-    InferenceScheduler --> Transformer : uses
-    Task --> TaskStatus : uses
-    InferenceEngine --> Transformer : uses
-    BaseSamplingStrategy <|-- TemperatureStrategy
-    BaseSamplingStrategy <|-- TopKStrategy
-    BaseSamplingStrategy <|-- TopPStrategy
-    SamplingPipeline --> BaseSamplingStrategy : composes
    BaseDataset <|-- SEQDataset
    BaseDataset <|-- SFTDataset
    BaseDataset <|-- DPODataset
    BaseDataset <|-- GRPODataset
-    DatasetFactory ..> BaseDataset : creates
    BaseStorage <|-- H5Storage
    BaseStorage <|-- JSONStorage
-    BaseDataset --> BaseStorage : uses
-    MultiSegmentFetcher --> BaseSegmentFetcher : uses
-    AutoModel <|-- Transformer
-    AutoModel --> ModelConfig : contains
-    Transformer --> DecoderBlock : uses
-    Transformer --> RotaryEmbedding : uses
-    Transformer --> Embedding : uses
-    DecoderBlock --> GQA : uses
-    DecoderBlock --> MLP : uses
-    DecoderBlock --> RMSNorm : uses
-    TrainContextBuilder --> ResumableDistributedSampler : creates
-    ResumableDistributedSampler --> BaseDataset : samples
+    BaseSamplingStrategy <|-- TemperatureStrategy
+    BaseSamplingStrategy <|-- TopKStrategy
+    BaseSamplingStrategy <|-- TopPStrategy
+    BaseSamplingStrategy <|-- SamplingPipeline
    ParallelModel <|-- RowParallelLinear
    ParallelModel <|-- ColumnParallelLinear
-    AutoTokenizer --> ChatTemplate : uses
+    AutoModel <|-- AutoRegressiveLM
+    AutoModel <|-- EmbeddingEncoder
+    BaseConfig <|-- BaseModelConfig
+    BaseConfig <|-- TrainConfig
+    BaseModelConfig <|-- AutoRegressiveLMConfig
+    BaseModelConfig <|-- EncoderConfig
    BaseFactory <|-- AutoModel
+    BaseFactory <|-- AttnFactory
+    BaseFactory <|-- FFNFactory
    BaseFactory <|-- DatasetFactory
    BaseFactory <|-- StrategyFactory
    BaseFactory <|-- SchedulerFactory
    BaseFactory <|-- CallbackFactory
+    BaseFactory <|-- StorageFactory
+    BaseFactory <|-- ConfigFactory
+    TrainCallback <|-- ValidationCallback
+    ProtocolHandler <|-- OpenAIHandler
+    ProtocolHandler <|-- AnthropicHandler
+
+    %% --- Composition (strong ownership, part destroyed with whole) ---
+    KVCache *-- PagePool
+    KVCache *-- Storage
+    KVCache *-- TaskTable
+    PagePool *-- Allocator
+    PagePool *-- PrefixCache
+    InferenceEngine *-- InferenceScheduler
+    InferenceScheduler *-- KVCache
+    InferenceScheduler *-- Executor
+    InferenceScheduler *-- TaskManager
+    AutoRegressiveLM *-- DecoderBlock
+    AutoRegressiveLM *-- RotaryEmbedding
+    AutoRegressiveLM *-- Embedding
+    EmbeddingEncoder *-- DecoderBlock
+    EmbeddingEncoder *-- RotaryEmbedding
+    EmbeddingEncoder *-- Embedding
+    DecoderBlock *-- RMSNorm
+    ChatCompletionRequest *-- ChatMessage
+    MessagesRequest *-- AnthropicMessage
+    AutoTokenizer *-- ChatTemplate
+    BaseFactory *-- Registry
+
+    %% --- Aggregation (weak ownership) ---
+    AutoModel o-- BaseModelConfig
+    Trainer o-- TrainCallback
+    TrainContext o-- BaseStrategy
+    TrainContext o-- BaseScheduler
+    TrainContext o-- Checkpoint
+    KvcacheView o-- Storage
+    SamplingPipeline o-- BaseSamplingStrategy
+    BaseDataset o-- BaseStorage
+
+    %% --- Dependency (uses temporarily) ---
+    TrainConfig ..> BaseStrategy : selects
+    StrategyFactory ..> BaseStrategy : creates
+    SchedulerFactory ..> BaseScheduler : creates
+    DatasetFactory ..> BaseDataset : creates
+    CallbackFactory ..> TrainCallback : creates
+    AttnFactory ..> GQA : creates
+    AttnFactory ..> MLA : creates
+    FFNFactory ..> MLP : creates
+    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
+    Checkpoint ..> Checkpoint : serializes
+    CheckpointCallback ..> Checkpoint : creates
+    KVCache ..> KvcacheView : binds
+    InferenceEngine ..> GenerationRequest : uses
+    InferenceEngine ..> GenerateResult : creates
+    OpenAIHandler ..> ChatCompletionRequest : receives
+    AnthropicHandler ..> MessagesRequest : receives
+    ProtocolHandler ..> StopChecker : creates
+    ProtocolHandler ..> StreamContext : creates
+
+    %% --- Association (general usage) ---
+    Trainer --> TrainConfig
+    DPOStrategy --> AutoModel
+    GRPOStrategy --> AutoModel
+    InferenceScheduler --> Task
+    InferenceScheduler --> TaskStatus
+    Task --> TaskStatus
+    InferenceEngine --> AutoModel
+    Executor --> AutoModel
+    Executor --> AutoTokenizer
+    TaskManager --> AutoTokenizer
+    MultiSegmentFetcher --> BaseSegmentFetcher
+    ResumableDistributedSampler --> BaseDataset
+
 ```

-### Module Overview
+
+## Module Overview

 | Module | Components | Description |
 |--------|------------|-------------|
-| **astrai.config** | ModelConfig, TrainConfig | Configuration management |
-| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, BaseStorage, H5Storage, JSONStorage, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory, save_h5, load_h5 | Dataset loading and management |
-| **astrai.serialization** | Checkpoint | Model serialization and checkpoint management |
-| **astrai.model** | AutoModel, Transformer, DecoderBlock, GQA, MLA, MLP, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
+| **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, 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, StrategyFactory, BaseScheduler, SchedulerFactory, TrainCallback, CallbackFactory | Training workflow management |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, KVCache, KvcacheView, Allocator, PrefixCache, PagePool, Storage, TaskTable, Task, TaskStatus, GenerationRequest, BaseSamplingStrategy, TemperatureStrategy, TopKStrategy, TopPStrategy, SamplingPipeline, ChatMessage, ChatCompletionRequest | Inference service with continuous batching and paged KV cache |
-| **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank, get_world_size, get_current_device, ParallelModel, ColumnParallelLinear, RowParallelLinear | Distributed parallel |
-| **astrai.factory** | Registry, BaseFactory | Generic component registration |
+| **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 |

-### Design Patterns
+## Design Patterns

 | Pattern | Classes | Purpose |
 |---------|---------|---------|
-| **Strategy** | `BaseStrategy`, `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy`, `StrategyFactory` | Flexible training strategy switching, supports SEQ/SFT/DPO/GRPO |
-| **Builder** | `TrainContextBuilder` | Chain-building training context, step-by-step initialization of components |
-| **Factory** | `StrategyFactory`, `SchedulerFactory`, `DatasetFactory`, `CallbackFactory`, `BaseFactory` | Decorator registration mechanism, dynamically create training strategies, schedulers, datasets, and callbacks |
-| **Observer** | `TrainCallback`, `CallbackFactory` | Callback mechanism for training process monitoring (checkpoint, early stopping, metrics) |
-| **Context** | `TrainContext` | Training process state container with model, optimizer, scheduler and checkpoint |
-| **Registry** | `BaseFactory`, `Registry` | Generic component registration with category and priority support |
-| **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with O(1) alloc/free via bitmask + LRU eviction |
-| **Strategy (Sampling)** | `BaseSamplingStrategy`, `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations with temperature, top-k, top-p |
-| **Producer-Consumer** | `InferenceScheduler`, `Task`, `waiting_queue`, `active_tasks` | Continuous batching with dynamic task queue management |
-| **Event-Driven** | `threading.Event`, `_task_event` | Non-blocking wait mechanism for task scheduling using Python's `threading` module |
-| **AutoModel Registry** | `AutoModel`, `Transformer` | Model type registration and dynamic loading via decorator pattern |
-| **Generator Pattern** | `GenerateResult`, `GenerationRequest` | Event-based result notification for streaming/non-streaming generation |
+| **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 |
+| **Template Method** | `ProtocolHandler`, `OpenAIHandler`, `AnthropicHandler` | HTTP API handler with format hooks |
+| **Builder** | `TrainContextBuilder` | Chain-building training context |
+| **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
+| **Context** | `TrainContext` | Unified training state bag |
+| **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`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |

-### Core Relationships
+## Core Relationships

-1. **Configuration → Training**: `TrainConfig` holds model, dataset, optimizer_fn, scheduler_fn and other training configuration references
-2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` to compute loss
-3. **Strategy Selection**: `StrategyFactory` creates corresponding strategy instance based on `train_type`
-4. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `Transformer`, uses `KVCache` (backed by `Allocator` + `PrefixCache` + `PagePool` + `Storage`) for paged KV cache management and `SamplingPipeline` for efficient continuous batching with streaming/non-streaming
-5. **Distributed Support**: `spawn_parallel_fn` and `setup_parallel` provide multi-process training capability for `Trainer`
-6. **Dataset Loading**: `DatasetFactory` creates datasets (SEQDataset, SFTDataset, DPODataset, GRPODataset), supports HDF5 loading via `BaseSegmentFetcher` and `MultiSegmentFetcher`
-7. **Checkpoint Management**: `Checkpoint` handles model state serialization/deserialization with safetensors
-8. **Scheduler Support**: `SchedulerFactory` creates learning rate schedulers (CosineScheduler, SGDRScheduler)
-9. **AutoModel Loading**: `AutoModel.from_pretrained()` dynamically loads model based on `config.json` model_type, uses `Registry` pattern for model type registration
+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` → `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

-## 3. Training Process
-
-The common training process for large language models (LLM) typically includes three stages: **Pre-training (SEQ)**, **Supervised Fine-Tuning (SFT)**, and **Reinforcement Learning from Human Feedback (DPO/GRPO)**. This system is designed to support seamless end-to-end flow, achieving efficient switching and state management of different training stages through modular strategies.
-
-### Core Formulas
-
-**Pre-training (SEQ):**
-
-$$
-L_{\text{PT}} = - \sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
-$$
-
-**SFT:**
-
-$$
-L_{\text{SFT}} = - \sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
-$$
-
-**DPO:**
-
-$$
-L_{\text{DPO}} = -\mathbb{E}_{(x, y_w, y_l) \sim D} \left[ \log \sigma\left( \beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)} \right) \right]
-$$
-
-**GRPO:**
-
-GRPO (Group Relative Policy Optimization) computes advantages from multiple responses to the same prompt, then optimizes using a PPO-style clipped objective:
-
-$$
-\text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
-$$
-
-Where $r_i$ is the reward for the $i$-th response, $\mu$ and $\sigma$ are the mean and standard deviation of group rewards.
-
-$$
-L_{\text{GRPO}} = -\mathbb{E} \left[ \min\left( \frac{\pi_\theta(a|s)}{\pi_{\text{ref}}(a|s)} \cdot A, \text{clip}\left(\frac{\pi_\theta(a|s)}{\pi_{\text{ref}}(a|s)}, 1-\epsilon, 1+\epsilon\right) \cdot A \right) \right] + \lambda \cdot D_{KL}
-$$
-
-The KL divergence term uses mean squared error approximation:
-
-$$
-L_{KL} = \lambda \cdot \mathbb{E} \left[ (\log \pi_\theta - \log \pi_{\text{ref}})^2 \right]
-$$
-
-The final loss is the sum of both: $L = L_{\text{policy}} + L_{KL}$
-
-Through the above three-stage progressive training, the model completes its evolution from a general language foundation to a specialized, highly-aligned dialogue intelligence.
-
-> Document Update Time: 2026-05-14
+> Document Update Time: 2026-05-17
@@ -1,237 +1,57 @@
-# AstrAI Data Flow Documentation
+# Data Flow

-This document describes the data flow of the AstrAI project (a training and inference framework for autoregressive Transformer language models). It covers the complete flow from raw data to model training and inference.
+This document describes the data pipeline: from raw text to model input tensors.

 ## Overview

-AstrAI adopts a modular design with the following main components:
- **Dataset Module** (`astrai/dataset/`): Dataset, sampler, serialization tools
- **Model Module** (`astrai/model/`): AutoModel, Transformer model and its submodules
- **Training Module** (`astrai/trainer/`): Trainer, training context, strategies, schedulers, callbacks, metric utilities
- **Inference Module** (`astrai/inference/`): Inference engine with continuous batching, streaming generation
- **Config Module** (`astrai/config/`): ModelConfig, TrainConfig
- **Factory Module** (`astrai/factory/`): Registry, BaseFactory for component registration
- **Parallel Module** (`astrai/parallel/`): Distributed training support
- **Serialization** (`astrai/serialization.py`): Checkpoint management with safetensors
-
-## Data Flow Diagram
-
-```mermaid
-flowchart LR
-    subgraph A[Data Preparation]
-        direction TB
-        A1[Raw Text] --> A2[AutoTokenizer]
-        A2 --> A3[Tokenized .h5 files]
-        A3 --> A4[BaseDataset]
-        A4 --> A5[ResumableDistributedSampler]
-        A5 --> A6[DataLoader]
-    end
-
-    subgraph B[Training]
-        direction TB
-        B1[DataLoader] --> B2[BaseStrategy]
-        B2 --> B3[Transformer Forward]
-        B3 --> B4[Loss + Backward]
-        B4 --> B5[Gradient Accumulation]
-        B5 -->|every accum_steps| B6[Optimizer Step]
-        B6 --> B7[LR Scheduler]
-        B7 -->|next batch| B2
-        B6 --> B8[CheckpointCallback]
-    end
-
-    subgraph C[Inference]
-        direction TB
-        C1[Checkpoint] --> C2[AutoModel]
-        C1 --> C3[AutoTokenizer]
-        C2 --> C4[InferenceEngine]
-        C3 --> C4
-        C4 --> C5[InferenceScheduler]
-        C5 --> C6[Transformer Forward]
-        C6 --> C7[sample]
-        C7 --> C8{End?}
-        C8 -->|No| C6
-        C8 -->|Yes| C9[Generated Text]
-    end
-
-    A --> B
-    B --> C
+```
+Raw Text → AutoTokenizer → Token IDs → .h5/.json → Dataset → Sampler → DataLoader → Training/Inference
 ```

-## Detailed Module Descriptions
+## Data Preparation

-### 1. Data Serialization (`astrai/dataset/storage.py` & `astrai/serialization.py`)
+Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or JSON (`.json`/`.jsonl`) files with keyed tensor groups.

- **`save_h5`**: Saves tensors by groups as HDF5 files (`.h5`), each key maps to a list of tensors
- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory
- **`Checkpoint`**: Encapsulates model state dict + epoch + iteration; uses safetensors
-
-### 2. Dataset Module
-
-#### 2.1 Dataset (`dataset.py`)
- **`BaseDataset`**: Abstract base class for windowed sequence sampling
- **`BaseSegmentFetcher` / `MultiSegmentFetcher`**: Fetch tensor segments by index range
- **`DatasetFactory`**: Creates dataset instances by `train_type` (`seq`, `sft`, `dpo`, `grpo`)
- Data keys: `"sequence"` (SEQ), `"loss_mask"` (SFT), `"chosen_mask"/"rejected_mask"` (DPO), `"masks"` (GRPO)
-
-#### 2.2 Sampler (`sampler.py`)
- **`ResumableDistributedSampler`**: Tracks `epoch` and `iter` for breakpoint resume; supports shuffle and drop_last
-
-### 3. Model Module
-
-#### 3.1 Transformer / AutoModel
- **`AutoModel`**: Base class with `from_pretrained()` / `save_pretrained()`
- **`Transformer`**: Decoder-only architecture, registered via `@AutoModel.register('transformer')`
- Embedding → N×DecoderBlock → RMSNorm → Linear lm_head
- RoPE position encoding, optional weight tying
-
-#### 3.2 Submodules (`module.py`)
- **`DecoderBlock`**: GQA attention + residual + MLP + RMSNorm
- **`GQA`**: Grouped Query Attention (also `MLA` for multi-latent attention)
- **`MLP`**: `SiLU(gate(x)) * up(x)` → down projection
- **`RotaryEmbedding`**: RoPE complex cache (freqs_cis)
- **`RMSNorm`**: Layer normalization
-
-### 4. Training Module
-
-#### 4.1 Training Context (`train_context.py`)
- **`TrainContext`**: Dataclass holding model, optimizer, dataloader, strategy, scheduler, checkpoint state
- **`TrainContextBuilder`**: Builder pattern — takes checkpoint for resume, builds all components
-
-#### 4.2 Trainer (`trainer.py`)
-
-The training loop is nested: **epoch** → **batch** (with step phase interspersed):
+Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:

 ```
-on_train_begin
-  on_epoch_begin
-    for each accumulation window of batches:        ← step phase
-      on_step_begin
-        for each batch in window:                    ← batch phase
-          on_batch_begin → strategy(batch) → loss → backward → on_batch_end
-          iteration += 1
-      on_step_end
-      optimizer.step() → zero_grad
-
-    on_epoch_end
-on_train_end
+StorageFactory.create("h5")   → H5Storage
+StorageFactory.create("json") → JSONStorage
 ```

-Key points:
- `on_step_*` fires every `accumulation_steps` batches, wrapping optimizer step AFTER the hook
- `on_batch_*` fires every batch, wrapping loss computation
- `GradientClippingCallback` fires on `on_step_end`
- LR scheduler steps inline (no `SchedulerCallback` class)
+Both support shared memory via `.share_memory_()`.

-#### 4.3 Strategy (`strategy.py`)
- **`SEQStrategy`**: Next-token prediction, cross-entropy with label smoothing
- **`SFTStrategy`**: Supervised fine-tuning with loss masking
- **`DPOStrategy`**: Direct Preference Optimization with reference model
- **`GRPOStrategy`**: Group Relative Policy Optimization with clipped ratio
+## Data Keys by Training Type

-#### 4.4 Scheduler (`schedule.py`)
- **`CosineScheduler`**: Cosine decay + linear warmup
- **`SGDRScheduler`**: Cosine annealing with warm restarts
- Created by `SchedulerFactory` and bound to optimizer
+| Type | Storage Keys |
+|------|-------------|
+| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) |
+| `sft` | `sequence`, `loss_mask` |
+| `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` |
+| `grpo` | `prompts`, `responses`, `masks`, `rewards` |

-#### 4.5 Callbacks
- **`CheckpointCallback`**: Saves safetensors at `ckpt_interval` iterations
- **`ProgressBarCallback`**: tqdm progress display
- **`MetricLoggerCallback`**: Writes JSONL metrics to `{ckpt_dir}/logs/`
- **`GradientClippingCallback`**: `clip_grad_norm_` on `on_step_end`
-
-### 5. Inference Module
-
-#### 5.1 Inference Engine (`engine.py`)
- **`InferenceEngine`**: Facade over scheduler; provides `generate()`, `generate_with_request()`, `generate_async()`
- Accepts `prompt: str | List[str]`, returns generator (stream) or string (non-stream)
-
-#### 5.2 Scheduler 4-Phase Loop (`scheduler.py`)
-
-Background thread runs continuously:
+## Dataset Architecture

 ```
-1. Cleanup → Remove finished tasks, free KV cache pages
-2. Refill  → Pop from waiting_queue, alloc pages, add to active
-3. Prefill → Group active tasks by prompt_len, run full forward pass
-4. Decode  → Pick largest same-position group, run single-token forward
+DatasetFactory.load(train_type, path, window_size, stride)
+  → StorageFactory.create(detect_format(path))
+    → MultiSegmentFetcher(BaseSegmentFetcher per key)
+      → BaseDataset.__getitem__(idx)
+        → sliding window [begin, end) via get_index(idx)
 ```

- **`Task`**: Tracks prompt_ids, output_ids, status (PENDING/RUNNING/FINISHED/ABORTED)
- **`KVCache`**: Facade over `Allocator` + `PrefixCache` + `PagePool` + `Storage` for paged KV cache
- **`KvcacheView`**: Batch view bundling cache + page table for attention layers
- **`sample()`**: Temperature → top-k → top-p → multinomial
+`window_size` = max input length, `stride` = step between consecutive samples.

-#### 5.3 Server (`server.py`)
- FastAPI with OpenAI `/v1/chat/completions` and Anthropic `/v1/messages` endpoints
- Streaming via SSE, health check at `/health`, stats at `/stats`
+## Sampler

-### 6. Tokenizer Module
+`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:

- **`AutoTokenizer`**: Wraps HuggingFace tokenizers (BBPE); `encode`/`decode`/`apply_chat_template`
- **`ChatTemplate`**: Jinja2-based template rendering for multi-turn chat
+- Tracks `start_epoch` / `start_iter` for resume
+- Shuffle via `torch.Generator(seed + epoch)`
+- Per-replica index slicing for DDP

-### 7. Factory & Parallel
+## DataLoader

- **`Registry` / `BaseFactory`**: Decorator-based component registration
- **`spawn_parallel_fn`**: Multi-process DDP launcher with NCCL backend
- **`ParallelModel` / `ColumnParallelLinear` / `RowParallelLinear`**: Tensor model parallelism
+Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.

-## Training Data Flow — Detailed Steps
-
-1. **Data Preparation**
-   - Raw text → token IDs via `AutoTokenizer.encode()`
-   - Save as `.h5` files (groups of tensor lists per data key)
-
-2. **Dataset Loading**
-   - `BaseDataset.load()` calls `load_h5()`, builds `MultiSegmentFetcher`
-   - Sliding window of `window_size` with `stride` determines sample boundaries
-
-3. **Sampling & Batching**
-   - `ResumableDistributedSampler` produces shuffled index sequences
-   - `DataLoader` fetches `[batch_size, window_size]` tensors via `__getitem__`
-
-4. **Strategy Forward**
-   - Strategy receives batch, calls `Transformer.forward()` for logits
-   - Computes task-specific loss (cross-entropy, DPO, GRPO)
-
-5. **Backward & Accumulation**
-   - `loss = raw_loss / accumulation_steps`
-   - `loss.backward()` accumulates gradients
-   - Every `accumulation_steps` batches: `optimizer.step()` → `zero_grad()`
-   - Every batch: `scheduler.step()` updates learning rate
-
-6. **Checkpoint**
-   - `CheckpointCallback` saves `model.state_dict()` + metadata to safetensors at `ckpt_interval` iterations
-   - Does NOT save optimizer/scheduler state (resume resets those)
-
-## Inference Data Flow — Detailed Steps
-
-1. **Model Loading**
-   - `AutoModel.from_pretrained(path)` loads weights from safetensors
-   - `torch.inference_mode()` wraps generation
-
-2. **Prompt Construction**
-   - Messages → `apply_chat_template(messages, tokenize=False)` → prompt string
-   - `tokenizer.encode(prompt)` → token IDs (truncated to `max_prompt_len`)
-
-3. **Continuous Batching Loop**
-   - **Cleanup**: Finished tasks → `stream_callback(STOP)`, free KV pages
-   - **Refill**: Pop from waiting queue, `PagePool.task_alloc()` for prompt pages
-   - **Prefill**: Group by prompt length, run full forward with `start_pos=0`
-   - **Decode**: Pick position group with most tasks, single-token forward:
-     - Model forward → `logits` → `sample()` → next token ID
-     - Append to `output_ids`, update `output_tokens`
-      - `PagePool.task_alloc()` allocates pages as needed
-      - `stream_callback(token)` for streaming clients
-
-4. **Output**
-   - `tokenizer.decode(output_ids)` → text
-   - Return to caller (streaming: token-by-token; non-streaming: complete string)
-
-## Checkpoint & Serialization
-
- **Training Checkpoint**: safetensors weights + epoch/iteration metadata. Optimizer/scheduler state is NOT persisted.
- **Inference Loading**: `AutoModel.from_pretrained()` loads from the same safetensors format.
- **Dataset Serialization**: HDF5 with shared memory support for large-scale pre-training data.
-
-> Document Update Time: 2026-05-14
+> Document Update Time: 2026-05-17
@@ -0,0 +1,140 @@
+# Inference
+
+## KV Cache
+
+At decode time, only the last query token matters. All previous K/V are cached to avoid recomputation:
+
+$$
+o_n = \sum_j \text{softmax}\left(\frac{q_n k_j}{\sqrt{d_k}}\right) v_j
+$$
+
+RoPE is applied **before** KV cache write, not after — otherwise position encoding drift occurs.
+
+## KVCache System
+
+Six classes working together:
+
+```
+KVCache (facade)
+  ├── Allocator        bitmask-based page allocator + ref-count + LRU eviction
+  ├── PrefixCache      hash-based prefix matching (page_hash via rolling hash)
+  ├── PagePool         orchestrates Allocator + PrefixCache
+  ├── TaskTable        maps task_id → page_table + cached token count
+  ├── Storage          k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim)
+  └── KvcacheView      bundles Storage + page_table + total_len for attention layers
+```
+
+`KVCache.bind(page_table, total_len)` returns a `KvcacheView` used by attention layers via `write()` / `gather()`.
+
+## Continuous Batching
+
+`InferenceScheduler` runs a daemon thread with a 4-phase loop:
+
+```
+1. Cleanup → Remove finished tasks, free KV pages
+2. Refill  → Pop from waiting_queue, task_alloc pages, activate
+3. Prefill → Group by (prompt_len, start_pos), run full forward
+4. Decode  → Pick largest same-position group, single-token forward
+```
+
+## Sampling (Strategy Pattern)
+
+```
+BaseSamplingStrategy → TemperatureStrategy → TopKStrategy → TopPStrategy
+```
+
+`SamplingPipeline` composes them: Temperature → Top-K → Top-P → softmax → multinomial.  
+`sample()` is a convenience shortcut for one-shot usage.
+
+## Protocol Handlers (Template Method)
+
+```python
+class ProtocolHandler(ABC):
+    def handle(self):
+        ctx = StreamContext(...)
+        agen = engine.generate_async(prompt, ...)
+        if stream: self._handle_stream(agen, ctx)
+        else:      self._handle_non_stream(agen, ctx)
+```
+
+Subclass hooks: `build_prompt()`, `create_response_id()`, `format_stream_start/token/end()`, `format_non_stream_response()`.
+
+`OpenAIHandler` → `/v1/chat/completions`, `AnthropicHandler` → `/v1/messages`.
+
+## Engine & GenerateResult
+
+```
+InferenceEngine
+  ├── generate(prompt, stream, ...) → str | List[str] | Generator
+  ├── generate_with_request(req)    → same
+  └── generate_async(prompt, ...)   → AsyncGenerator
+```
+
+`GenerateResult` uses `Condition` for non-streaming (`wait_completion()`) and `Event` for streaming (`wait()`). Stream callback is `cb(token)`.
+
+## HTTP API
+
+```
+POST /v1/chat/completions   OpenAI
+POST /v1/messages            Anthropic
+GET  /health                 {"status":"ok","model_loaded":true}
+GET  /stats                  scheduler statistics
+```
+
+### OpenAI
+
+```bash
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+Response:
+```json
+{
+  "id": "chatcmpl-abc123",
+  "object": "chat.completion",
+  "choices": [{"message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"}],
+  "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}
+}
+```
+
+Streaming SSE: `data: {"choices":[{"delta":{"role":"assistant"}}]}` → token chunks → `data: [DONE]`
+
+### Anthropic
+
+```bash
+curl -X POST http://localhost:8000/v1/messages \
+  -H "Content-Type: application/json" \
+  -d '{"model":"astrai","system":"You are helpful.","messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
+```
+
+Supports `stop_sequences` and streaming via `event: content_block_delta`.
+
+### GenerationRequest Parameters
+
+| Param | Type | Default | Description |
+|-------|------|---------|-------------|
+| `messages` | List[dict] | required | Chat messages (role, content) |
+| `temperature` | float | 1.0 | Sampling temperature (0.0–2.0) |
+| `top_p` | float | 1.0 | Nucleus threshold |
+| `top_k` | int | 50 | Top-k count |
+| `max_tokens` | int | None | Max generation length |
+| `stream` | bool | False | Stream output |
+
+## Engine API
+
+```python
+# Non-streaming
+engine.generate("Hello", stream=False)          # -> str
+engine.generate(["A", "B"], stream=False)       # -> List[str]
+
+# Streaming
+engine.generate("Hello", stream=True)           # -> Generator[str]
+engine.generate(["A", "B"], stream=True)        # -> Generator[Tuple[int, str]]
+
+# Async
+await engine.generate_async("Hello", ...)       # -> AsyncGenerator[str]
+```
+
+> Document Update Time: 2026-05-17
@@ -1,334 +0,0 @@
-## Model Introduction
-
-### 1. Model Architecture
-
-This model uses the Transformer architecture with GQA mechanism (q_head=24, kv_head=4), which saves KV cache memory compared to traditional MHA. The model is built by stacking multiple layers of Transformer blocks, with 1.0 billion parameters. Transformer is an autoregressive model that calculates the relationship between all previous tokens to obtain the probability distribution of the next token.
-
-The model now uses the **AutoModel** base class for flexible loading and saving:
-
-```python
-from astrai.model import AutoModel
-
-# Load model from checkpoint
-model = AutoModel.from_pretrained("path/to/model")
-
-# Save model to new directory
-model.save_pretrained("path/to/save")
-```
-
-The Transformer model is registered via `@AutoModel.register('transformer')` decorator, allowing easy extension for new model types.
-
-```mermaid
-flowchart TB
-    subgraph Layers["Transformer Layers"]
-        direction TB
-        A[Input Embedding] --> B[Transformer Block\nLayer 1]
-        B --> C[Transformer Block\nLayer ...]
-        C -->         D[Transformer Block\nLayer ...]
-        D --> E[RMSNorm]
-        E --> F[Linear]
-        F --> G[SoftMax]
-    end
-
-    subgraph TransformerBlock["Transformer Block"]
-        direction TB
-        H[x] --> I[RMSNorm]
-        I --> J[Linear → Q/K/V]
-        J --> K[Q]
-        J --> L[K]
-        J --> M[V]
-        K --> N[RoPE]
-        L --> O[RoPE]
-        N --> P["Q @ K^T / sqrt(d)"]
-        O --> P
-        P --> Q[Masked SoftMax]
-        Q --> R[S @ V]
-        M --> R
-        R --> S[Linear]
-        S --> T[+]
-        H --> T
-        T --> U[RMSNorm]
-        U --> V["Linear (gate)"]
-        U --> W["Linear (up)"]
-        V --> X[SiLU]
-        X --> Y[×]
-        W --> Y
-        Y --> Z["Linear (down)"]
-        Z --> AA[+]
-        T --> AA
-        AA --> BB[x']
-    end
-
-    classDef main fill:#e6f3ff,stroke:#0066cc;
-    classDef block fill:#fff2e6,stroke:#cc6600;
-    class Layers main;
-    class TransformerBlock block;
-```
-
-What is an autoregressive model? After splitting a sentence into tokens, the model predicts the probability distribution of the next token. This means the model calculates the probability of the next possible token and its corresponding probability based on the given context (the sequence of tokens that have already appeared).
-
-#### 1. Autoregression
-
-In autoregressive modeling, when a sentence is tokenized into a sequence of tokens, the model learns to predict what comes next. Given a sequence of tokens as input, the model calculates a probability distribution over all possible next tokens. This distribution tells us how likely each potential next token is, given the current context.
-
-For instance, if the input sequence contains tokens representing a question, the model might predict that certain response tokens have higher probabilities than others. The sampling process then selects one token from this distribution—controlled by parameters like top_k, top_p, and temperature—to serve as the next token in the sequence.
-
-Once a token is selected, it is appended to the input sequence, and the model repeats this process. The updated sequence is then fed back into the model to predict the next token. This iterative process continues until either a special end-of-sequence token is generated, or the maximum sequence length is reached. These control tokens are essential because without them, the model would continue generating tokens indefinitely, eventually exhausting available memory.
-
-#### 2. Causal Mask
-
-Transformers use attention mechanism. The input shape is generally [bsz, seq_len], and the output is [bsz, seq_len, n_dim]. To predict the next token, the model's input and output must be offset by one position. The target predicted by the model must be offset by one position, and during training we also use the offset-by-one method:
-
-```
-sequence : [[1, 2, 3, 4, 5, 6]]
-input_ids: [[1, 2, 3, 4, 5]]
-target_ids: [[2, 3, 4, 5, 6]]
-```
-
-The attention score calculation formula is:
-
-$$ s_{ij} = softmax(\frac{q_i^Tk_j}{\sqrt{d_k}}) $$
-$$ s_{ij} := s_{ij} + mask_{ij} $$
-
-Here, the attention score represents the degree to which the model attends to the similarity between two tokens.
-
-For decoder-only structure models, to prevent the model from "stealing" information from future positions, a mask needs to be added during attention calculation. We need to apply a mask before attention score calculation. This mask is typically a lower triangular matrix, and for a sequence of length n, its shape is [n, n]. Below is an example of how to create such a causal mask matrix for a sequence of length 5:
-
-```
-[[0, -inf, -inf, -inf, -inf],
- [0,    0, -inf, -inf, -inf],
- [0,    0,    0, -inf, -inf],
- [0,    0,    0,    0, -inf],
- [0,    0,    0,    0,    0]]
-```
-
-In this matrix, 0 represents positions that can be attended to, while -inf represents positions that should be masked (i.e., should not be attended to). Because this matrix ensures that after the softmax, the parts of the attention scores where $j > i$ change from `inf` to 0, meaning the model cannot see future information.
-
-#### 3. Rotary Position Embedding
-
-Rotary Position Embedding (RoPE) is a position encoding method designed to solve the problem of lacking direct modeling of sequence position information in Transformer models. Unlike traditional position encodings (such as sine and cosine function position encodings), RoPE embeds position information directly into the Query (Q) and Key (K) vectors, allowing the model to more naturally handle relative position relationships in sequences.
-
-$$ q_i = R_i W_q x_i $$
-$$ k_j = R_j W_k x_j $$
-$$ q_i^T k_j = (R_i W_q x_i)^T( R_j W_k x_j) = x_i^T W_q^T R_{i-j} W_k x_j $$
-
-The $R_{i-j}$ controls the attenuation of attention for different tokens at different relative distances. When the absolute value of $i - j$ is larger, the degree of attenuation is stronger. This approach allows the model to learn relative position relationships, enabling the model to scale and adapt to longer sequences.
-
-## KV Cache Implementation
-
-According to the attention calculation formula:
-
-$$
-\begin{align*}
-o_i &= \sum_j s_{ij} v_{j} \newline
-s_{ij} &= \text{softmax}\left( \frac{q_{i} k_{j}}{\sqrt{d_k}} \right)
-\end{align*}
-$$
-
-Since the model is an autoregressive model, we only need to calculate for the last part of the sequence, meaning the index $i$ is fixed as the last element of the sequence, and we compute $o_{n}$:
-
-$$
-\begin{align*}
-o_n &= \sum_j s_{j}v_{j} \newline
-s_j &= \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}} \right)
-\end{align*}
-$$
-
-If we expand the expression:
-
-$$
-o_n = \sum_j \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}}\right)v_{j}
-$$
-
-In the above expression, only k and v have length indices, while $q$ does not. Therefore, during the calculation process, the input of $q$ is fixed as the last token from the previous input, while $k$ and $v$ need to be cached for parts of different lengths. Also, when caching, note that position encoding calculation should be performed before KV cache computation, otherwise there will be position encoding calculation errors.
-
-### 4. AutoModel Loading
-
-The project now uses the **AutoModel** base class for flexible model loading and saving:
-
-```python
-from astrai.model import AutoModel
-
-# Load model from checkpoint
-model = AutoModel.from_pretrained("path/to/model")
-
-# Save model to new directory
-model.save_pretrained("path/to/save")
-```
-
-The Transformer model is registered via `@AutoModel.register('transformer')` decorator, allowing easy extension for new model types. The `from_pretrained` method automatically loads the `config.json` to determine the model type and uses safetensors format for weights.
-
-### 5. Continuous Batching Inference
-
-The inference engine supports **continuous batching** for efficient batch processing:
-
-```python
-from astrai.inference import InferenceEngine, GenerationRequest
-
-# Create inference engine with continuous batching
-engine = InferenceEngine(
-    model=model,
-    tokenizer=tokenizer,
-)
-
-# Use GenerationRequest 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,
-    stream=True,
-)
-
-# Generate with streaming
-for token in engine.generate_with_request(request):
-    print(token, end="", flush=True)
-```
-
-The continuous batching feature allows dynamic batch composition where new requests can join at any time and completed requests are released immediately.
-
-## HTTP API Usage
-
-The inference server provides HTTP endpoints for remote inference. Start the server first:
-
-```bash
-python -m scripts.tools.server --port 8000
-```
-
-### OpenAI-Compatible Endpoint
-
-The server provides an OpenAI-compatible chat completion endpoint at `/v1/chat/completions`:
-
-```bash
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [
-      {"role": "system", "content": "You are a helpful assistant."},
-      {"role": "user", "content": "Hello, how are you?"}
-    ],
-    "temperature": 0.8,
-    "max_tokens": 2048,
-    "stream": false
-  }'
-```
-
-**Request Parameters:**
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `messages` | List[dict] | Required | Chat messages with role and content |
-| `temperature` | float | 1.0 | Sampling temperature (0.0-2.0) |
-| `top_p` | float | 1.0 | Nucleus sampling threshold |
-| `top_k` | int | 50 | Top-k sampling parameter |
-| `max_tokens` | int | 1024 | Maximum tokens to generate |
-| `stream` | bool | false | Enable streaming response |
-
-**Response (non-streaming):**
-```json
-{
-  "id": "chatcmpl-1234567890",
-  "object": "chat.completion",
-  "created": 1234567890,
-  "model": "astrai",
-  "choices": [
-    {
-      "index": 0,
-      "message": {"role": "assistant", "content": "Hello! I'm doing well..."},
-      "finish_reason": "stop"
-    }
-  ],
-  "usage": {
-    "prompt_tokens": 20,
-    "completion_tokens": 15,
-    "total_tokens": 35
-  }
-}
-```
-
-### Streaming Response
-
-Enable streaming for real-time token-by-token output:
-
-```bash
-curl -X POST http://localhost:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": "Write a story"}],
-    "stream": true,
-    "max_tokens": 500
-  }'
-```
-
-The server uses Server-Sent Events (SSE) with content type `text/event-stream`.
-
-### Anthropic-Compatible Endpoint
-
-The server also provides an Anthropic-compatible endpoint at `/v1/messages`:
-
-```bash
-curl -X POST http://localhost:8000/v1/messages \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "You are a helpful assistant.",
-    "messages": [{"role": "user", "content": "Hello, how are you?"}],
-    "max_tokens": 2048
-  }'
-```
-
-Response:
-```json
-{
-  "id": "msg_abc123...",
-  "type": "message",
-  "role": "assistant",
-  "model": "astrai",
-  "content": [{"type": "text", "text": "Hello! I am doing well..."}],
-  "stop_reason": "end_turn",
-  "stop_sequence": null,
-  "usage": {"input_tokens": 20, "output_tokens": 15}
-}
-```
-
-Streaming:
-```bash
-curl -X POST http://localhost:8000/v1/messages \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "system": "You are a helpful assistant.",
-    "messages": [{"role": "user", "content": "Write a short poem"}],
-    "max_tokens": 500,
-    "stream": true
-  }'
-```
-
-Supports `stop_sequences` for early termination:
-```bash
-curl -X POST http://localhost:8000/v1/messages \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "astrai",
-    "messages": [{"role": "user", "content": "Write a story"}],
-    "max_tokens": 500,
-    "stop_sequences": ["The end", "THE END"]
-  }'
-```
-
-### Health Check
-
-Monitor server and model status:
-
-```bash
-curl http://localhost:8000/health
-# {"status": "ok", "model_loaded": true}
-
-curl http://localhost:8000/stats
-# {"total_tasks": 10, "total_tokens": 5000, "active_tasks": 1, "waiting_queue": 0}
-```
-
-> Document Update Time: 2026-05-14
@@ -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 | `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 (unlimited) |
-| `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-14
+> Document Update Time: 2026-05-17
@@ -0,0 +1,225 @@
+# Training
+
+## Model Architecture
+
+The model uses a decoder-only Transformer with **GQA** (Grouped Query Attention) and optional **MLA** (Multi-head Latent Attention). 1.0 billion parameters, Chinese–English bilingual.
+
+```mermaid
+flowchart TB
+    subgraph Layers["Transformer Layers"]
+        direction TB
+        A[Input Embedding] --> B[Transformer Block\nLayer 1]
+        B --> C[Transformer Block\nLayer ...]
+        C --> D[Transformer Block\nLayer ...]
+        D --> E[RMSNorm]
+        E --> F[Linear]
+        F --> G[SoftMax]
+    end
+
+    subgraph TransformerBlock["Transformer Block"]
+        direction TB
+        H[x] --> I[RMSNorm]
+        I --> J[Linear → Q/K/V]
+        J --> K[Q]; J --> L[K]; J --> M[V]
+        K --> N[RoPE]; L --> O[RoPE]
+        N --> P["Q @ K^T / sqrt(d)"]; O --> P
+        P --> Q[Masked SoftMax]; Q --> R[S @ V]; M --> R
+        R --> S[Linear]; S --> T[+]; H --> T
+        T --> U[RMSNorm]
+        U --> V["Linear (gate)"]; U --> W["Linear (up)"]
+        V --> X[SiLU]; X --> Y[×]; W --> Y
+        Y --> Z["Linear (down)"]; Z --> AA[+]; T --> AA
+        AA --> BB[x']
+    end
+```
+
+### Autoregression
+
+Given a token sequence, the model predicts the probability of the next token. Each generated token is appended to the input and fed back, repeating until an end-of-sequence token or max length.
+
+### Causal Mask
+
+```
+sequence : [[1, 2, 3, 4, 5, 6]]
+input_ids: [[1, 2, 3, 4, 5]]
+target_ids: [[2, 3, 4, 5, 6]]
+```
+
+Lower-triangular mask prevents attending to future positions:
+
+```
+[[0, -inf, -inf, -inf, -inf],
+ [0,    0, -inf, -inf, -inf],
+ [0,    0,    0, -inf, -inf],
+ [0,    0,    0,    0, -inf],
+ [0,    0,    0,    0,    0]]
+```
+
+### Rotary Position Embedding (RoPE)
+
+RoPE embeds position into Q/K vectors via complex rotation:
+
+$$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
+
+The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers.
+
+## Training Loop
+
+Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
+
+```
+on_train_begin
+  on_epoch_begin
+    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
+```
+
+### Callback Lifecycle
+
+| Hook | Fires | Default callback |
+|------|-------|-----------------|
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_step_begin` | Every accumulation window | `GradientClippingCallback` |
+| `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
+| `on_step_end` | Every accumulation window | `ValidationCallback` |
+| `on_train_end` | Training ends | `CheckpointCallback`, `MetricLoggerCallback` (final save) |
+
+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
+
+### SEQ (Pre-training)
+
+Next-token cross-entropy with optional label smoothing:
+
+$$
+L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
+$$
+
+Keys: `input_ids`, `target_ids`
+
+### SFT (Supervised Fine-Tuning)
+
+Masked cross-entropy (`ignore_index=-100`) over response tokens:
+
+$$
+L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
+$$
+
+Keys: `input_ids`, `target_ids`, `loss_mask`
+
+### DPO (Direct Preference Optimization)
+
+Frozen reference model, preference margin via log-ratio:
+
+$$
+L_{\text{DPO}} = -\mathbb{E}\left[\log\sigma\left(\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)} - \beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}\right)\right]
+$$
+
+Parameters: `beta=0.1`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
+
+### GRPO (Group Relative Policy Optimization)
+
+On-policy PPO with group-normalized advantages:
+
+$$
+\text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
+$$
+
+$$
+L_{\text{GRPO}} = -\mathbb{E}\left[\min\left(\frac{\pi_\theta}{\pi_{\text{ref}}}A,\; \text{clip}\left(\frac{\pi_\theta}{\pi_{\text{ref}}}, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}\left[(\log\pi_\theta - \log\pi_{\text{ref}})^2\right]
+$$
+
+Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`, `sync_interval=200`.
+
+Keys: `prompts`, `responses`, `masks`, `rewards`.
+
+## LR Schedulers
+
+| Type | Class | Description |
+|------|-------|-------------|
+| Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
+| SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
+
+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, 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 persisted by default via `Checkpoint.extra`.  
+Training config (`TrainConfig.to_dict()`) saved into `meta.json` during training via `CheckpointCallback`.
+
+## TrainContextBuilder (Builder Pattern)
+
+```python
+context = (
+    TrainContextBuilder(config)
+        .with_checkpoint(checkpoint)
+        .build()
+)
+# Returns TrainContext with model, strategy, optimizer, scheduler, dataloader, checkpoint
+```
+
+- Loads checkpoint weights if provided
+- Wraps model with `parallel_wrapper` if `nprocs > 1`
+- Creates `ResumableDistributedSampler` for shuffle+resume
+- Builds strategy via `StrategyFactory.create(train_type, ...)`
+
+## Training CLI
+
+```bash
+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-17
@@ -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",
@@ -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",
 ]
@@ -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
@@ -1,42 +1,90 @@
 import json
-from dataclasses import asdict, dataclass
-from typing import Optional, Self
+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 ModelConfig:
-    # basic config
+class BaseModelConfig(BaseConfig):
+    """Base config with ``model_type`` dispatch and file I/O."""
+
    model_type: Optional[str] = None
+
+    @classmethod
+    def from_file(cls, config_path: str) -> Self:
+        with open(config_path, "r") as f:
+            raw: Dict[str, Any] = json.load(f)
+        return cls.from_dict(raw)
+
+    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)
+
+
+@dataclass
+@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

-    # GQA
+    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

-    def load(self, config_path: str) -> Self:
-        config = {}
-        with open(config_path, "r") as f:
-            config.update(json.load(f))
+    kv_lora_rank: Optional[int] = None
+    qk_nope_head_dim: Optional[int] = None
+    qk_rope_head_dim: Optional[int] = None

-        for key, value in config.items():
-            if hasattr(self, key):
-                setattr(self, key, value)
+    ffn_type: str = "mlp"
+    n_routed_experts: Optional[int] = None
+    n_shared_experts: Optional[int] = None
+    n_activated_experts: Optional[int] = None
+    topk_method: Optional[str] = None

-        return self

-    def save(self, config_path: str):
-        config_dict = {k: v for k, v in asdict(self).items() if v is not None}
-        with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
+@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
@@ -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.")
@@ -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",
@@ -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)

@@ -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())
@@ -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
@@ -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(
@@ -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,
    )
@@ -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

@@ -148,11 +156,15 @@ class InferenceScheduler:
                            t.output_ids.append(ntok)
                            t.output_tokens += 1
                            pos = t.input_tokens + t.output_tokens
-                            self._page_cache.task_extend(t.task_id, pos)
+                            extend_ok = self._page_cache.task_extend(t.task_id, pos)
                            if t.stream_callback:
                                t.stream_callback(
                                    self._task_mgr.tokenizer.decode([ntok])
                                )
+                            if not extend_ok:
+                                t.status = TaskStatus.ABORTED
+                                if t.stream_callback:
+                                    t.stream_callback(STOP)

                        for t in valid:
                            if t.is_finished(stop_ids):
@@ -165,6 +177,9 @@ class InferenceScheduler:
                if task.stream_callback:
                    task.stream_callback(STOP)
                self._page_cache.task_free(task.task_id)
+            for task in self._task_mgr.get_waiting_tasks():
+                if task.stream_callback:
+                    task.stream_callback(STOP)
            self._task_mgr.clear_queues()
            raise

@@ -193,6 +193,10 @@ class TaskManager:
        with self._lock:
            return list(self.active_tasks)

+    def get_waiting_tasks(self) -> List[Task]:
+        with self._lock:
+            return list(self.waiting_queue)
+
    def clear_queues(self) -> None:
        with self._lock:
            self.waiting_queue.clear()
@@ -1,12 +1,11 @@
 from astrai.model.automodel import AutoModel
-from astrai.model.module import (
-    GQA,
-    MLP,
-    DecoderBlock,
-    Linear,
-    RMSNorm,
-)
-from astrai.model.transformer import Transformer
+from astrai.model.components.attention import GQA
+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.encoder import EmbeddingEncoder
+from astrai.model.transformer import AutoRegressiveLM

 __all__ = [
    # Modules
@@ -16,6 +15,7 @@ __all__ = [
    "GQA",
    "DecoderBlock",
    # Models
-    "Transformer",
+    "AutoRegressiveLM",
+    "EmbeddingEncoder",
    "AutoModel",
 ]
@@ -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"))
@@ -0,0 +1,25 @@
+from astrai.model.components.attention import GQA, MLA, repeat_kv
+from astrai.model.components.decoder_block import DecoderBlock
+from astrai.model.components.embedding import Embedding
+from astrai.model.components.linear import Linear
+from astrai.model.components.mlp import MLP
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import (
+    RotaryEmbedding,
+    apply_rotary_emb,
+    get_rotary_emb,
+)
+
+__all__ = [
+    "Linear",
+    "RMSNorm",
+    "MLP",
+    "Embedding",
+    "GQA",
+    "MLA",
+    "DecoderBlock",
+    "RotaryEmbedding",
+    "apply_rotary_emb",
+    "get_rotary_emb",
+    "repeat_kv",
+]
@@ -5,11 +5,14 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor

+from astrai.factory import BaseFactory
 from astrai.inference.core.cache import KvcacheView
+from astrai.model.components.linear import Linear
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import apply_rotary_emb


 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
-    """Repeat KV heads n_rep times for GQA."""
    bs, slen, n_heads, head_dim = x.shape
    if n_rep == 1:
        return x
@@ -20,88 +23,13 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
    )


-def get_rotary_emb(
-    dim: int,
-    max_len: int,
-    base: float = 10000,
-    device: Optional[torch.device] = None,
-) -> Tensor:
-    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
-    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
-    freqs = torch.outer(t, theta).float()
-    cos = torch.cos(freqs)
-    sin = torch.sin(freqs)
-    return torch.complex(cos, sin)
-
-
-def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
-    dtype = x.dtype
-    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
-    x_complex = torch.view_as_complex(x_)
-    freqs_cis = freqs_cis.unsqueeze(2)
-    x_rotated = x_complex * freqs_cis
-    x_out = torch.view_as_real(x_rotated).flatten(-2)
-    return x_out.to(dtype)
-
-
-class RotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, max_len: int, base: int = 10000):
-        super().__init__()
-        self.dim = dim
-        self.max_len = max_len
-        self.base = base
-        self._set_rotary_buffer(self.max_len)
-
-    def _set_rotary_buffer(self, max_len: int):
-        rotary_emb = get_rotary_emb(self.dim, max_len, self.base)
-        freqs_cis = torch.view_as_real(rotary_emb)
-        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
-
-    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
-        if position_ids is None:
-            position_ids = (
-                torch.arange(x.size(1), device=x.device)
-                .unsqueeze(0)
-                .expand(x.size(0), -1)
-            )
-        position_freq_cis = self.freqs_cis[position_ids].float()
-        return torch.view_as_complex(position_freq_cis)
-
-
-class Linear(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
-        super().__init__()
-        self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
-        self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
-
-    def forward(self, x: Tensor) -> Tensor:
-        return F.linear(x, self.weight, self.bias)
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim, norm_eps):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.normalized_shape = (dim,)
-        self.norm_eps = norm_eps
-
-    def forward(self, x: Tensor) -> Tensor:
-        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)
-
-
-class MLP(nn.Module):
-    def __init__(self, dim: int, dim_feed_forward: int):
-        super().__init__()
-        self.up = Linear(dim, dim_feed_forward)
-        self.gate = Linear(dim, dim_feed_forward)
-        self.down = Linear(dim_feed_forward, dim)
-
-    def forward(self, x: Tensor) -> Tensor:
-        gated = self.up(x) * F.silu(self.gate(x))
-        out = self.down(gated)
-        return out
+class AttnFactory(BaseFactory[nn.Module]):
+    @classmethod
+    def create(cls, attn_type: str, **kwargs) -> nn.Module:
+        return super().create(attn_type, **kwargs)


+@AttnFactory.register("gqa")
 class GQA(nn.Module):
    def __init__(
        self,
@@ -152,7 +80,6 @@ class GQA(nn.Module):
    ) -> Tensor:
        is_causal = attn_mask is None

-        # (bsz, seq_len, dim) -> (bsz, seq_len, n_heads, head_dim)
        q = self._split_heads(self.q_proj(x), self.n_heads)
        k = self._split_heads(self.k_proj(x), self.n_kv_heads)
        v = self._split_heads(self.v_proj(x), self.n_kv_heads)
@@ -167,7 +94,6 @@ class GQA(nn.Module):

        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)

-        # (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
        sdqa_out = (
            F.scaled_dot_product_attention(q, k, v, attn_mask, is_causal=is_causal)
@@ -183,6 +109,7 @@ class GQA(nn.Module):
        return out


+@AttnFactory.register("mla")
 class MLA(nn.Module):
    def __init__(
        self,
@@ -193,6 +120,7 @@ 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,
    ):
@@ -206,16 +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)

-        # fused KV: (k_nope, k_rope, v)
        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)
@@ -248,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)
@@ -256,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)
@@ -274,57 +210,3 @@ class MLA(nn.Module):

        out = self.o_proj(attn_out)
        return out
-
-
-class DecoderBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        n_heads: int,
-        dim_ffn: int,
-        n_kv_heads: int,
-        norm_eps: int,
-        use_qk_norm: bool,
-        use_gated_attention: bool,
-        layer_id: int,
-    ):
-        super().__init__()
-        self.attention = GQA(
-            dim,
-            n_heads,
-            n_kv_heads,
-            use_qk_norm,
-            norm_eps,
-            use_gated_attention,
-            layer_id,
-        )
-        self.input_norm = RMSNorm(dim, norm_eps)
-        self.mlp = MLP(dim, dim_ffn)
-        self.post_attention_norm = RMSNorm(dim, norm_eps)
-
-    def forward(
-        self,
-        x: Tensor,
-        rotary_emb: Tensor,
-        attention_mask: Optional[Tensor] = None,
-        paged_cache: Optional[KvcacheView] = None,
-    ) -> Tensor:
-        attn_output = self.attention(
-            self.input_norm(x),
-            rotary_emb,
-            attention_mask,
-            paged_cache,
-        )
-        x = attn_output + x
-        x = self.mlp(self.post_attention_norm(x)) + x
-
-        return x
-
-
-class Embedding(nn.Module):
-    def __init__(self, vocab_size: int, embedding_dim: int):
-        super().__init__()
-        self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
-
-    def forward(self, x: Tensor) -> Tensor:
-        return F.embedding(x, self.weight)
@@ -0,0 +1,59 @@
+from typing import Optional
+
+import torch.nn as nn
+from torch import Tensor
+
+from astrai.inference.core.cache import KvcacheView
+from astrai.model.components.attention import AttnFactory
+from astrai.model.components.mlp import FFNFactory
+from astrai.model.components.norm import RMSNorm
+
+
+class DecoderBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        dim_ffn: int,
+        n_kv_heads: int,
+        norm_eps: float,
+        use_qk_norm: bool,
+        use_gated_attention: bool,
+        layer_id: int,
+        attn_type: str = "gqa",
+        ffn_type: str = "mlp",
+        **kwargs,
+    ):
+        super().__init__()
+        self.attention = AttnFactory.create(
+            attn_type,
+            dim=dim,
+            n_heads=n_heads,
+            n_kv_heads=n_kv_heads,
+            use_qk_norm=use_qk_norm,
+            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, **kwargs)
+
+    def forward(
+        self,
+        x: Tensor,
+        rotary_emb: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        paged_cache: Optional[KvcacheView] = None,
+    ) -> Tensor:
+        attn_output = self.attention(
+            self.input_norm(x),
+            rotary_emb,
+            attention_mask,
+            paged_cache,
+        )
+        x = attn_output + x
+        x = self.mlp(self.post_attention_norm(x)) + x
+
+        return x
@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class Embedding(nn.Module):
+    def __init__(self, vocab_size: int, embedding_dim: int):
+        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)
@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class Linear(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
+        super().__init__()
+        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)
@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from astrai.factory import BaseFactory
+from astrai.model.components.linear import Linear
+
+
+class FFNFactory(BaseFactory[nn.Module]):
+    @classmethod
+    def create(cls, ffn_type: str, dim: int, dim_ffn: int, **kwargs) -> nn.Module:
+        return super().create(ffn_type, dim, dim_ffn, **kwargs)
+
+
+@FFNFactory.register("mlp")
+class MLP(nn.Module):
+    def __init__(self, dim: int, dim_ffn: int):
+        super().__init__()
+        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))
+        out = self.down(gated)
+        return out
+
+
+@FFNFactory.register("moe")
+class DeepSeekMoE(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_ffn: int,
+        n_routed_experts: int,
+        n_shared_experts: int = 1,
+        n_activated_experts: int = 2,
+        topk_method: str = "greedy",
+    ):
+        super().__init__()
+        self.dim = dim
+        self.n_routed_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.n_activated_experts = n_activated_experts
+        self.topk_method = topk_method
+
+        self.router = Linear(dim, n_routed_experts, bias=False)
+
+        self.shared_experts = nn.ModuleList(
+            [MLP(dim, dim_ffn) for _ in range(n_shared_experts)]
+        )
+        self.routed_experts = nn.ModuleList(
+            [MLP(dim, dim_ffn) for _ in range(n_routed_experts)]
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        bsz, seq_len, dim = x.shape
+        x_flat = x.view(-1, dim)
+
+        shared_out = self._shared_forward(x_flat)
+        routed_out = self._routed_forward(x_flat)
+
+        out = (shared_out + routed_out).view(bsz, seq_len, dim)
+        return out
+
+    def _shared_forward(self, x: Tensor) -> Tensor:
+        if self.n_shared_experts == 0:
+            return torch.zeros_like(x)
+        return sum(e(x) for e in self.shared_experts) / self.n_shared_experts
+
+    def _routed_forward(self, x: Tensor) -> Tensor:
+        N, D = x.shape
+        K = self.n_activated_experts
+
+        router_logits = self.router(x)
+        router_probs = torch.softmax(router_logits.float(), dim=-1).to(x.dtype)
+
+        topk_weights, topk_indices = torch.topk(router_probs, K, dim=-1)
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+        output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
+        for expert_idx in range(self.n_routed_experts):
+            expert_mask = topk_indices == expert_idx
+            token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
+            if token_idx.numel() == 0:
+                continue
+            expert_input = x[token_idx]
+            expert_output = self.routed_experts[expert_idx](expert_input)
+            weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
+            output.index_add_(0, token_idx, expert_output * weights)
+
+        return output
@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, norm_eps):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.normalized_shape = (dim,)
+        self.norm_eps = norm_eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)
@@ -0,0 +1,53 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+def get_rotary_emb(
+    dim: int,
+    max_len: int,
+    base: float = 10000,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
+    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
+    freqs = torch.outer(t, theta).float()
+    cos = torch.cos(freqs)
+    sin = torch.sin(freqs)
+    return torch.complex(cos, sin)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
+    dtype = x.dtype
+    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
+    x_complex = torch.view_as_complex(x_)
+    freqs_cis = freqs_cis.unsqueeze(2)
+    x_rotated = x_complex * freqs_cis
+    x_out = torch.view_as_real(x_rotated).flatten(-2)
+    return x_out.to(dtype)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, max_len: int, base: float = 10000):
+        super().__init__()
+        self.dim = dim
+        self.max_len = max_len
+        self.base = base
+        self._set_rotary_buffer(self.max_len)
+
+    def _set_rotary_buffer(self, max_len: int):
+        rotary_emb = get_rotary_emb(self.dim, max_len, self.base)
+        freqs_cis = torch.view_as_real(rotary_emb)
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
+        if position_ids is None:
+            position_ids = (
+                torch.arange(x.size(1), device=x.device)
+                .unsqueeze(0)
+                .expand(x.size(0), -1)
+            )
+        position_freq_cis = self.freqs_cis[position_ids].float()
+        return torch.view_as_complex(position_freq_cis)
@@ -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
@@ -4,16 +4,14 @@ 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.module import (
-    DecoderBlock,
-    Embedding,
-    Linear,
-    RMSNorm,
-    RotaryEmbedding,
-)
+from astrai.model.components.decoder_block import DecoderBlock
+from astrai.model.components.embedding import Embedding
+from astrai.model.components.linear import Linear
+from astrai.model.components.norm import RMSNorm
+from astrai.model.components.rope import RotaryEmbedding


 def process_attention_mask(
@@ -48,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(
@@ -71,6 +73,15 @@ class Transformer(AutoModel):
                    config.use_qk_norm,
                    config.use_gated_attention,
                    layer_id,
+                    attn_type=config.attn_type,
+                    ffn_type=config.ffn_type,
+                    n_routed_experts=config.n_routed_experts,
+                    n_shared_experts=config.n_shared_experts,
+                    n_activated_experts=config.n_activated_experts,
+                    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)
            ]
@@ -79,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"
@@ -95,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]
@@ -111,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]
@@ -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,
    )
@@ -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,
        )
@@ -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):
@@ -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",
@@ -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)

@@ -0,0 +1,143 @@
+import torch
+from torch.optim import Optimizer
+
+
+def _zeropower_via_newtonschulz(G: torch.Tensor, steps: int = 5):
+    assert G.ndim == 2
+    X = G
+    scale = max(1, G.size(0) / G.size(1)) ** 0.5
+    X = X / (X.norm() + 1e-7) * scale
+    if steps == 0:
+        return X
+    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
+
+
+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:
+            params_2d, params_1d = [], []
+            grads_2d, grads_1d = [], []
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                if p.grad.is_sparse:
+                    raise RuntimeError("Muon does not support sparse gradients")
+                if p.ndim >= 2:
+                    params_2d.append(p)
+                    grads_2d.append(p.grad)
+                else:
+                    params_1d.append(p)
+                    grads_1d.append(p.grad)
+
+            if params_2d:
+                self._muon_update_foreach(params_2d, grads_2d, group)
+            if params_1d:
+                self._adamw_update_foreach(params_1d, grads_1d, group)
+
+        return loss
+
+    def _muon_update_foreach(self, params_2d, grads_2d, group):
+        lr = group["lr"]
+        momentum = group["momentum"]
+        wd = group["weight_decay"]
+        nesterov = group["nesterov"]
+        ns_steps = group["ns_steps"]
+
+        if wd != 0:
+            torch._foreach_mul_(params_2d, 1 - lr * wd)
+
+        if nesterov:
+            grads_2d = torch._foreach_add(grads_2d, params_2d, alpha=wd)
+
+        bufs = []
+        for p, grad in zip(params_2d, grads_2d):
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(grad)
+            bufs.append(state["momentum_buffer"])
+
+        torch._foreach_lerp_(bufs, grads_2d, 1 - momentum)
+
+        for p, buf in zip(params_2d, bufs):
+            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_foreach(self, params_1d, grads_1d, group):
+        lr = group["adamw_lr"]
+        betas = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        wd = group["adamw_wd"]
+
+        steps: list[int] = []
+        exp_avgs, exp_avg_sqs = [], []
+        has_state = []
+        for p in params_1d:
+            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)
+                has_state.append(False)
+            else:
+                has_state.append(True)
+            state["step"] += 1
+            steps.append(state["step"])
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+
+        beta1, beta2 = betas
+
+        torch._foreach_lerp_(exp_avgs, grads_1d, 1 - beta1)
+        grads_sq = torch._foreach_mul(grads_1d, grads_1d)
+        torch._foreach_lerp_(exp_avg_sqs, grads_sq, 1 - beta2)
+
+        bias_correction1 = [1 - beta1**s for s in steps]
+        bias_correction2 = [1 - beta2**s for s in steps]
+
+        if wd != 0:
+            torch._foreach_mul_(params_1d, 1 - lr * wd)
+
+        exp_avg_corrected = torch._foreach_div(exp_avgs, bias_correction1)
+        denom = torch._foreach_div(exp_avg_sqs, bias_correction2)
+        denom = torch._foreach_sqrt(denom)
+        torch._foreach_add_(denom, eps)
+        torch._foreach_addcdiv_(params_1d, exp_avg_corrected, denom, value=-lr)
@@ -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)
@@ -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,
@@ -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,
+        )
@@ -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"]
@@ -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__":
@@ -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"
@@ -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(
@@ -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,
    )

@@ -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,
@@ -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)
@@ -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):
@@ -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"])
@@ -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)
@@ -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()
@@ -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)
@@ -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,
@@ -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"],
@@ -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"],
    )
@@ -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"]
Author	SHA1	Message	Date
ViperEkura	0594ce1017	perf: Muon step 改用 torch._foreach_* 批处理并移除 NS 迭代的冗余 bf16 转换	2026-05-23 19:50:12 +08:00
ViperEkura	ff509ff39f	fix: decode后task_extend失败时提前中止，scheduler崩溃时通知waiting任务	2026-05-20 19:23:13 +08:00
ViperEkura	785d65436c	fix: 修复 to_dict list 类型丢失与 OpenAI stop 参数失效 - to_dict() 增加 list 类型序列化支持，metrics 等字段不再丢失 - OpenAIHandler 补充 get_stop_sequences/on_token，读取 request.stop 并检测停止序列 - 文档类图补充缺失字段、修正关系分类、ChatCompletionRequest 字段增加 Optional	2026-05-19 21:07:07 +08:00
ViperEkura	64be81b7b3	feat: ProgressBarCallback 支持日志行输出到 stdout - serialization 和 metric_logger 的 timestamp 统一使用 ISO 8601 格式 - ProgressBarCallback 新增 log_interval/file 参数，默认输出到 sys.stdout	2026-05-19 19:12:38 +08:00
ViperEkura	45479b5731	feat: metric 参数通过 TrainConfig 传递 - TrainConfig 新增 log_dir/log_interval/metrics 配置字段 - metric_logger 调用改用 **kwargs 传递，BaseFactory.create 自动过滤	2026-05-19 17:50:24 +08:00
ViperEkura	e0a3337c22	docs: 更新视频链接	2026-05-19 17:34:01 +08:00
ViperEkura	812238060b	fix: docker-compose UID/GID 添加默认值，修复 docker.sh logs 命令	2026-05-18 14:24:00 +08:00
ViperEkura	14b0d56197	fix: 修复无法创建子进程的问题 - mp.start_processes daemon=False	2026-05-18 09:40:32 +08:00
ViperEkura	6c8533f1d2	docs: 修正文档中类名/字段名与代码不一致之处 - ModelConfig → AutoRegressiveLMConfig, Transformer → AutoRegressiveLM - 新增缺失类: EncoderConfig, EmbeddingEncoder, ConfigFactory, StorageFactory, ValidationCallback - TrainConfig/TrainContext/ChatCompletionRequest 补充缺失字段 - dataflow.md 中 create_storage → StorageFactory.create - 示例 --train_type=pt → seq 与代码一致	2026-05-17 21:02:21 +08:00
ViperEkura	2c2697390d	feat: 新增 GradientCheckpointingCallback - TrainConfig.gradient_checkpointing_modules 指定模块类型 - apply 递归遍历，兼容 DDP，不硬编码模型结构 - modules=None 时静默跳过，零开销	2026-05-17 18:21:05 +08:00
ViperEkura	7621f05d3f	docs: AdamW beta 默认值改为 (0.9, 0.95) - 与 Muon 优化器的 AdamW 子优化器保持一致 - 同步更新 train.py/training.md/params.md/README	2026-05-17 17:08:31 +08:00
ViperEkura	10ebd7211f	feat: 新增 Muon 优化器 - 2D 参数用 Newton-Schulz 正交化 + Nesterov 动量更新 - 1D 参数用 AdamW 更新 - 支持 lr/momentum/weight_decay/ns_steps 配置	2026-05-17 16:44:03 +08:00
ViperEkura	42a391f0fb	feat: 训练中新增验证循环 - TrainConfig 添加 val_dataset/val_step 字段 - TrainContext 添加 val_dataloader/val_loss 字段 - 新增 ValidationCallback 按 step 触发验证 + 训练结束时验证 - ProgressBar/MetricLogger 支持 val_loss 展示与记录	2026-05-17 16:12:42 +08:00
ViperEkura	97c7ac0f4f	refactor: Transformer更名为AutoRegressiveLM并新增EmbeddingEncoder - AutoRegressiveLM 注册名改为 autoregressive_lm - 新增 EmbeddingEncoder 支持 mean/cls/last pooling - ModelConfig 增加 pooling_type / normalize_embeddings 字段 - 导入、注释、测试全部同步更新	2026-05-17 15:29:20 +08:00
ViperEkura	8f1b32f2b6	fix: 移除多余 request 参数并增强 tokenizer 健壮性 - 路由和 _get_engine 不再需要 request 参数，直接引用模块级 app - from_pretrained 增加文件完整性校验，缺 tokenizer.json 则抛 FileNotFoundError - 移除 from_pretrained 中未使用的 **kwargs	2026-05-17 12:52:18 +08:00
ViperEkura	c241a5dcef	refactor: 优化并行训练配置与启动管理 - 配置新增 start_method 支持 spawn/fork/forkserver 选择 - 启动方式 mp.spawn 改为 mp.start_processes，支持 daemon=True - validate() 改为基于 metadata 的反射式校验，不再硬编码字段列表 - CLI 新增 --start_method 参数	2026-05-17 12:33:10 +08:00
ViperEkura	44dab27fdc	feat: 数据集加载时校验必填字段 - BaseDataset.required_keys 属性声明所需存储 key - load() 时自动校验，缺失立即抛 KeyError - SEQ/SFT/DPO/GRPO 各自声明 required_keys	2026-05-17 11:50:38 +08:00
ViperEkura	a44fd22a99	fix: 修复训练与模型参数传递问题 - state_dict_fn 传入 CheckpointCallback，修复多卡 DDP 下 key 前缀丢失 - MLA 增加 use_qk_norm 支持，消除参数静默丢失 - moe_topk_method 统一命名为 topk_method - checkpoint 回调移至最前	2026-05-17 11:20:13 +08:00
ViperEkura	8a11a7d444	fix: 修复训练脚本两处参数传递问题 - prepare_checkpoint 增加 DDP 判断，单卡时不访问 .module - dpo_beta 改为 beta，对齐 DPOStrategy 参数名	2026-05-17 11:04:40 +08:00
ViperEkura	1d54491809	refactor: 改用递归子模块 init 替代统一 normal_(0.006) - Embedding.reset_parameters: normal_(std=0.02) - Linear.reset_parameters: kaiming_uniform_ + uniform_ bias - Transformer._init_weights 通过 apply 递归调用子模块 reset_parameters - 移除全局 normal_(0.006) 覆盖，各模块使用更合适的分布	2026-05-17 10:44:18 +08:00
ViperEkura	ad9f4d9cf6	refactor: generate_ar 改用流式输出并去除冗余注释	2026-05-17 10:23:42 +08:00
ViperEkura	e1638a7ade	fix: 修正AdamW超参数默认值与文档示例 - 交换adamw_beta1/adamw_beta2默认值：beta1=0.95, beta2=0.99 - label_smoothing默认值改为0.05 - 文档示例统一更新：train_type=pt, weight_decay=0.01 - 移除文档中过时的strategy default标注	2026-05-16 22:46:17 +08:00
ViperEkura	f91bfee33e	refactor: Config序列化统一BaseConfig基类 - 新增astrai/config/base.py，提供to_dict/from_dict基类 - 统一命名：load/save → from_file/to_file - Checkpoint.meta合并训练配置到meta.json - sys.stderr.warn → warnings.warn - from_file改为classmethod	2026-05-16 22:06:39 +08:00
ViperEkura	d7a7f570ed	refactor: 训练循环改为两重迭代并统一参数命名 - 训练循环从三重(epoch→batched→batch)改为二重(epoch→batch) - batch_size → batch_per_device, accumulation_steps → grad_accum_steps - scheduler 移入 step block 对齐 optimizer 更新步 - GradientClippingCallback 改用 on_step_begin 避免零梯度裁剪 - 移除 _train_impl 误导性的 -> Checkpoint 标注 - total_steps 修除为向下取整并精简为一行 - warmup_steps 改为 warmup_ratio (默认0.05)	2026-05-16 21:27:35 +08:00
ViperEkura	7dea929788	refactor: checkpoint 按 HF 方式存独立 .pt 文件，callback 接管恢复 - Checkpoint.save/load: extra 逐 key 写为 {key}.pt 而非单个 extra.pt - meta.json 新增 timestamp - CheckpointCallback: save_extra/load_extra 静态方法 + extra_keys 类属性 - on_train_begin 接管 optimizer/scheduler 恢复，TrainContextBuilder 不再传 load_extra_fn	2026-05-16 18:29:04 +08:00
ViperEkura	026d1fc33d	fix: total_steps 改用 ceiling 匹配实际步数原公式全用 floor 少算 optimizer step，改用逐层 ceiling （ceil_div via (a+b-1)//b）对齐 DDP sampler padding + DataLoader drop_last=False 尾批 + batched 尾组截断。	2026-05-16 17:53:18 +08:00
ViperEkura	7242eedbf4	fix: 学习率调度按 optimizer step 计数并防止 warmup 越界 - total_steps 除以 accumulation_steps，匹配 optimizer.step() 频率 - warmup_steps 用 min 截断，避免 lr_decay_steps 为负	2026-05-16 17:07:36 +08:00
ViperEkura	04c0dc7a47	refactor: Storage 改用工厂模式，server reload 接入 uvicorn - 新增 StorageFactory(BaseFactory[BaseStorage]) 替代手写 dict 注册 - H5Storage / JSONStorage 通过 @StorageFactory.register 注册 - dataset.py 使用 StorageFactory.create() 替代 create_storage() - 删除 create_storage / available_storage_types 死函数 - server.py reload 参数正式传入 uvicorn.run()	2026-05-16 17:00:26 +08:00
ViperEkura	48a53121ba	refactor: 工厂 kwargs 过滤及组件参数清理 - BaseFactory.create() 按 __init__ 签名过滤多余 kwargs - 移除 GQA/MLA/MLP/DeepSeekMoE 中多余的 **kwargs - MLP/DeepSeekMoE 参数名统一为 dim_ffn - scheduler max_seq_len 增加 None 显式判断 - 默认 max_prompt_len 提升至 2048	2026-05-16 16:47:41 +08:00
ViperEkura	0ba8c70ce1	fix: 修复 MLA 多个 bug 并缩小测试模型参数 - MLA kv_b_proj 输出维度和 q_rope 切分偏移修复 - 打通 MLA 配置从 ModelConfig 到 DecoderBlock 的传递路径 - rope_theta 配置不再被忽略，MLA 使用 qk_rope_head_dim - tie_weight 使用 is True 避免 None 隐式生效 - norm_eps/rope base 类型标注修正 - 测试模型参数缩小 (dim=8, head_dim=4) - 新增 6 种架构配置 × 2 场景的前向传播测试	2026-05-16 14:57:43 +08:00
ViperEkura	3d12a03909	docs : 拆分文档并补充类图缺失类和关系线 - 将 design.md 拆分为 architecture.md / inference.md / training.md - 精简 dataflow.md 为纯数据管道 - 删除 design.md 和 introduction.md - 更新 README.md 和 README-zh-CN.md 链接 - 补充 ChatMessage / AnthropicMessage 等 6 条孤立类关系线 - 补充 BaseModelConfig 和 TaskManager 两个缺失类	2026-05-15 23:38:26 +08:00
ViperEkura	c169659611	docs: 修正 assets/docs/ 类图、数据流、参数文档及贡献指南 - design.md: 新增 ProtocolHandler/OpenAIHandler/AnthropicHandler 等缺失类 - design.md: 新增 Template Method、Storage 设计模式 - dataflow.md: 修正 GQA/MLA 为独立条目，补充 JSON 存储后端 - params.md: 标注 label_smoothing CLI 默认与 strategy 默认差异 - introduction.md: 修正 max_tokens 默认值 1024→2048 - CONTRIBUTING.md: 重写（纯 Python 无 conda、补充 CI 步骤与常见问题） - .github/PULL_REQUEST_TEMPLATE.md: 修正 lint 命令，去除多余注释要求 - .github/ISSUE_TEMPLATE/bug_report.md: 修正 label（enhancement→bug）	2026-05-15 22:54:41 +08:00
ViperEkura	e12f1a7ee5	feat: BaseModelConfig + DeepSeekMoE + 工厂模式替代 if/else - BaseModelConfig: fields() 精确字段匹配 + 类型矫正 + 未知key警告 - DeepSeekMoE: 共享专家 + 路由专家 + top-K 门控 - AttnFactory/FFNFactory: 装饰器注册，DecoderBlock 零分支 - config 用 attn_type/ffn_type 驱动组件选择	2026-05-15 20:34:52 +08:00
ViperEkura	ef25efffa2	refactor: 拆分 module.py 为 components 子包 - rope/linear/norm/embedding/mlp/attention/decoder_block 各自独立文件 - 依赖单向无循环 - 公开接口不变，外部无需修改	2026-05-15 20:08:36 +08:00