ci: 添加 Docker 镜像自动构建工作流

.dockerignore

@@ -0,0 +1,9 @@
+# Ignore everything
+*
+
+# Allow necessary files
+!astrai/
+!scripts/
+!assets/
+!pyproject.toml
+!README.md

.gitattributes

@@ -0,0 +1,19 @@
+# Auto detect text files
+* text=auto
+
+# Files that MUST use LF (Unix/Linux execution)
+*.sh text eol=lf
+*.py text eol=lf
+*.md text eol=lf
+*.yml text eol=lf
+
+Dockerfile text eol=lf
+.dockerignore text eol=lf
+
+.gitignore text eol=lf
+.gitattributes text eol=lf
+
+# Windows scripts - use CRLF
+*.bat text eol=crlf
+*.cmd text eol=crlf
+*.ps1 text eol=crlf

.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,27 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: "[BUG]"
+labels: enhancement
+assignees: ''
+
+---
+
+## Description
+A clear and concise description of what the bug is.
+## Steps to Reproduce
+1.  ...
+2.  ...
+3.  ...
+## Expected Behavior
+What you expected to happen.
+## Actual Behavior
+What actually happened.
+## Environment
+- Python version:
+- AstrAI version (or commit hash):
+- Operating System:
+- GPU (if applicable):
+- CUDA/cuDNN version (if applicable):
+## Additional Context
+Add any other context, screenshots, or logs here.

.github/ISSUE_TEMPLATE/custom.md

@@ -0,0 +1,10 @@
+---
+name: Custom issue template
+about: Describe this issue template's purpose here.
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+

.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,19 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: "[FEAT]"
+labels: ''
+assignees: ''
+
+---
+
+## Description
+A clear and concise description of the feature you'd like to see.
+## Problem Statement
+What problem does this feature solve? Why is it needed?
+## Proposed Solution
+Describe the solution you'd like. Include any design ideas, API changes, or implementation details.
+## Alternatives Considered
+Describe any alternative solutions or features you've considered.
+## Additional Context
+Add any other context, screenshots, or references here.

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,26 @@
+## Description
+Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context.
+
+Fixes # (issue number)
+
+## Type of Change
+Please delete options that are not relevant.
+
+- [ ] Bug fix (non-breaking change which fixes an issue)
+- [ ] New feature (non-breaking change which adds functionality)
+- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
+- [ ] Documentation update
+- [ ] Other (please describe):
+
+## How Has This Been Tested?
+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 .`)
+- [ ] I have performed a self-review of my own code
+- [ ] I have commented my code, particularly in hard-to-understand areas
+- [ ] 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
+- [ ] New and existing unit tests pass locally with my changes
+- [ ] Any dependent changes have been merged and published in downstream modules

.github/workflows/docker.yml

@@ -0,0 +1,50 @@
+name: Build and Push Docker Image
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Extract metadata
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ghcr.io/${{ github.repository }}
+          tags: |
+            type=ref,event=tag
+            type=raw,value=latest
+
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          platforms: linux/amd64
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

.github/workflows/lint.yml

@@ -0,0 +1,31 @@
+name: Lint
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.12
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install .[dev]
+
+      - name: Check formatting with ruff
+        run: |
+          ruff format --check .
+
+      - name: Check import sorting
+        run: |
+          ruff check . --select I

.github/workflows/spell-check.yml

@@ -1,17 +0,0 @@
-name: Spell Check
-on: [push, pull_request]
-
-permissions:
-  contents: read
-
-jobs:
-  spellcheck:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - name: Check spelling in specific files
-        uses: codespell-project/actions-codespell@v2
-        with:
-          check_filenames: true
-          only_warn: false
-          path: "**/*.{md, py}"

.github/workflows/tests.yml

@@ -0,0 +1,31 @@
+name: Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install .[dev]
+
+      - name: Run tests with pytest
+        run: |
+          python -m pytest tests/ -v

.gitignore

@@ -6,7 +6,17 @@
 
 # Allow specific file types and root files
 !*.py
-!*.md
+!*.sh
-!*.png
+
-!LICENSE
+# Allow GitHub files
-!pyproject.toml
+!/.github/**
+
+# Allow root files
+!/.gitattributes
+!/.dockerignore
+!/Dockerfile
+!/assets/**
+!/CONTRIBUTING.md
+!/LICENSE
+!/pyproject.toml
+!/README.md

CONTRIBUTING.md

@@ -0,0 +1,68 @@
+# Contributing to AstrAI
+
+Thank you for your interest in contributing to AstrAI! This document provides guidelines and steps for contributing.
+
+## How to Contribute
+
+### 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.).
+
+### 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.
+
+## Code Style
+
+AstrAI uses [Ruff](https://docs.astral.sh/ruff/) for code formatting and linting. Please ensure your code is formatted before submitting.
+
+- Run Ruff to format and lint:
+  ```bash
+  ruff format .
+  ruff check --fix .
+  ```
+- The project uses **double quotes** for strings and **4‑space indentation** (as configured in `pyproject.toml`).
+
+## Testing
+
+If you add or modify functionality, please include appropriate tests.
+
+- Run the test suite with:
+  ```bash
+  pytest
+  ```
+- Ensure all tests pass before submitting your PR.
+
+## Code Review
+
+All submissions will be reviewed. We may request changes or discuss alternatives. Please be responsive to feedback.
+
+## License
+
+By contributing, you agree that your contributions will be licensed under the same [GPL-3.0 License](LICENSE) that covers the project.
+
+---
+
+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!

Dockerfile

@@ -0,0 +1,54 @@
+# 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
+
+WORKDIR /app
+
+# Install Python 3.12 and minimal build dependencies
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    python3.12 \
+    python3.12-dev \
+    python3.12-venv \
+    gcc \
+    g++ \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create isolated virtual environment
+RUN python3.12 -m venv --copies /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# Copy source code and install dependencies
+COPY astrai/ ./astrai/
+COPY pyproject.toml .
+RUN pip install --no-cache-dir --upgrade pip \
+    && pip install --no-cache-dir . \
+    --extra-index-url https://download.pytorch.org/whl/cu126
+
+# Production stage
+FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production
+
+WORKDIR /app
+
+# Install Python 3.12 runtime
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    python3.12 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy virtual environment from builder
+COPY --from=builder /opt/venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# Copy application code
+COPY astrai/ ./astrai/
+COPY scripts/ ./scripts/
+COPY assets/ ./assets/
+COPY pyproject.toml .
+COPY README.md .
+
+# Create non-root user
+RUN useradd -m astrai && chown -R astrai:astrai /app
+USER astrai
+
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1

README.md

@@ -1,286 +1,197 @@
-![image-20250306182014120](/assets/images/project_logo_clipped.png)
+<div align="center">
-
-<div style="display: flex; flex-direction: column; align-items: center; justify-content: center; text-align: center; font-size: 16px; font-weight: bold; margin-top: 50px;">
   
-  <div>
+  <img src="assets/images/logo.png" width="auto" alt="Logo">
-    <a href="#english" style="text-decoration: none; margin: 0 10px; color: blue;">English</a> | 
+  <p>
-    <a href="#chinese" style="text-decoration: none; margin: 0 10px; color: blue;">中文</a>
+    <strong>A lightweight Transformer training & inference framework</strong>
-  </div>
+  </p>
-
-  <h1 style="margin: 20px 0 0 0; font-size: 2.5em; font-weight: bold;">KHAOSZ </h1>
 </div>
 
-<h2 id="english">English Version</h2>
+<div align="center">
+  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
+  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
+  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?color=76bad9" alt="release">
+  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.stargazers_count&label=stars&suffix=%20stars&color=76bad9" alt="stars">
+  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.forks_count&label=forks&suffix=%20forks&color=76bad9" alt="forks">
+</div>
+<br>
 
-A training and inference framework for autoregressive Transformer language models.
+<div align="center">
+  <a href="#english">English</a> •
+  <a href="assets/docs/README-zh-CN.md">中文</a> •
+  <a href="https://github.com/ViperEkura/AstrAI/issues">Issue Tracker</a> •
+  <a href="https://github.com/ViperEkura/AstrAI/discussions">Discussions</a> •
+  <a href="https://huggingface.co/ViperEk/">HuggingFace</a>
+</div>
 
-**Model Download Options (choose one):**
+<br>
 
-1. Visit [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) and check **Files and versions**
+## 📖 Table of Contents
-2. Run `scripts/download.py` to download model parameters
 
-**Demo Video:** [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd)
+<details open>
+<summary><b>English</b></summary>
 
-For training data sources, please refer to the **Model Card** section on the HuggingFace download page.
+- [Features](#features)
+- [Quick Start](#quick-start)
+- [Documentation](#documentation)
+- [Contributing](#contributing)
+- [Community](#community)
+- [License](#license)
 
-**License:** The code follows the GPL-3.0 license. Please provide attribution when using it.
+</details>
 
-- **📊 Device Selection:** Uses CUDA for training by default
+---
-- **🌐 Performance Optimization:** Enable `dtype=torch.bfloat16` to accelerate training and reduce memory usage. Ensure your hardware supports this feature
-- **🤖 Language Support:** The model supports training in Chinese and English. Since the BBPE tokenizer hasn't been trained on multilingual text, OOV (Out-of-Vocabulary) issues are minimal for Chinese and English, but may exist for other languages
 
+<a id="english"></a>
+## English
 
-### 📌 Training Guide
+### Features
 
-To train this Transformer model, follow these steps:
+- 🚀 **High Performance**: Optimized for both training and inference with efficient parallelization.
+- 🔧 **Flexible**: Support for seq/sft/dpo/grpo training, customizable model architectures.
+- 💡 **Easy to Use**: Simple API with comprehensive examples and demos.
+- 📦 **Lightweight**: Minimal dependencies, easy to deploy.
+- 🔬 **Research‑Friendly**: Modular design, easy to experiment with new ideas.
+- 🤗 **HuggingFace Integration**: Compatible with HuggingFace models and datasets.
 
-**(1). Prepare the Dataset:**
+### Quick Start
 
-Place the dataset in the specified root directory. This system uses the BBPE tokenizer for tokenization and requires training with pre-tokenized segments (stored as *.h5 format files).
+#### Installation
-
-**(2). Install Dependencies:**
 
 ```bash
+git clone https://github.com/ViperEkura/AstrAI.git
+cd AstrAI
 pip install -e .
 ```
 
-**(3). Run the Training Script:**
+For development dependencies:
 
 ```bash
-python train.py \
+pip install -e ".[dev]"
---train_type=train_type[seq, sft, dpo] \
---data_root_path=/path/to/dataset \
---param_path=/path/to/param_path \
---n_epoch=5 \
---batch_size=8 \
---max_lr=2e-4 \
---checkpoint_interval=10000 \
---checkpoint_dir=checkpoints 
 ```
 
-**Parameter Explanation:**
+#### Train a Model
-- `--train_type`: Training type (seq, sft, dpo)
-- `--data_root_path`: Dataset root directory
-- `--param_path`: Path to model training parameters
-- `--n_epoch`: Total number of training epochs
-- `--batch_size`: Batch size
-- `--accumulation_steps`: Number of batches per training step
-- `--warmup_steps`: Warmup steps
-- `--max_lr`: Maximum learning rate (using warmup + cosine decay)
-- `--checkpoint_interval`: Checkpoint saving interval
-- `--checkpoint_dir`: Checkpoint saving directory
-- `--resume_dir`: Resume training from specified path
-
-
-
-### 👉 Usage Guide
-
-**(1). Chat with the Model:**
-
-Open `chat.py` or use the streaming/non-streaming interfaces:
-
-**Streaming Output:**
-```python
-import torch
-from khaosz import Khaosz
-
-model_dir = "your_model_parameter_dir"
-model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-history = []
-
-while True:
-    query = input(">> ")
-    if query == "!exit":
-        break
-    
-    response_size = 0
-    for response, history in model.stream_generate(
-        query=query, 
-        history=history,
-        temperature=0.85,
-        top_p=0.95,
-        top_k=50
-    ):
-        print(response[response_size:], end="")
-        response_size = len(response)       
-```
-
-**Non-streaming Output:**
-```python
-import torch
-from khaosz import Khaosz
-
-model_dir = "your_model_parameter_dir"
-model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-history = []
-
-while True:
-    query = input(">> ")
-    if query == "!exit":
-        break
-    
-    response = model.generate(
-        query=query, 
-        history=history,
-        temperature=0.85,
-        top_p=0.95,
-        top_k=50
-    )
-    print(response)
-```
-
-**(2). Retrieval-Augmented Generation (RAG):**
-
-```python
-import torch
-from khaosz import Khaosz
-
-model_dir = "your_model_parameter_dir"
-model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-
-retrieved_content = model.retrieve_generate(
-    query=query,
-    retrieve_top_k=5,
-    temperature=0.6,
-    top_k=30,
-    top_p=0.95
-)
-print(retrieved_content)
-```
-
-<h2 id="chinese">中文版本</h2>
-这是一个支持基于自回归模式的 Transfomer 语言模型训练以及推理框架
-
-**模型下载选项（任选其一）：**
-
-1. 访问 [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) 查看 **Files and versions**
-2. 运行 `scripts/download.py` 下载模型参数
-
-**演示视频：** [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd)
-
-训练数据来源请参见 HuggingFace 下载页面中的 **Model Card** 部分。
-
-**许可证：** 代码遵循 GPL-3.0 协议，使用时请注明出处。
-
-- **📊 设备选择：** 默认使用 CUDA 进行训练
-- **🌐 性能优化：** 启用 `dtype=torch.bfloat16` 以加速训练并减少内存占用，请确保硬件支持该特性
-- **🤖 语言支持：** 模型支持中文和英文训练。由于 BBPE 分词器未使用多语言文本训练，因此中英文的 OOV（未登录词）问题较少，其他语言可能存在 OOV 问题
-
-
-### 📌 训练指南
-
-要训练该 Transformer 模型，请按照以下步骤操作：
-
-**(1). 准备数据集：**
-
-将数据集放置在指定的根目录下， 本系统采用 BBPE 分词器进行分词，并且要求使用已经经过分词的 token 分段训练（分段存储为 *.h5 格式）
-
-**(2). 安装依赖：**
 
 ```bash
-pip install -e .
+python scripts/tools/train.py \
+  --train_type=seq \
+  --data_root_path=/path/to/dataset \
+  --param_path=/path/to/param_path
 ```
 
-**(3). 运行训练脚本：**
+#### Generate Text
 
 ```bash
-python train.py \
+python scripts/tools/generate.py --param_path=/path/to/param_path
---train_type=train_type[seq, sft, dpo] \
---data_root_path=/path/to/dataset \
---param_path=/path/to/param_path \
---n_epoch=5 \
---batch_size=8 \
---max_lr=2e-4 \
---checkpoint_interval=10000 \
---checkpoint_dir=checkpoints 
 ```
 
-**参数说明：**
+#### Docker
-- `--train_type`: 训练类型（seq, sft, dpo）
-- `--data_root_path`: 数据集根目录
-- `--param_path`: 模型训练参数路径
-- `--n_epoch`: 总训练轮数
-- `--batch_size`: 批量大小
-- `--accumulation_steps`: 每个训练步骤的 batch 数量
-- `--warmup_steps`: 预热步数（warmup steps）
-- `--max_lr`: 最大学习率（使用预热 + 余弦衰减）
-- `--checkpoint_interval`: 检查点保存间隔
-- `--checkpoint_dir`: 检查点保存目录
-- `--resume_dir`: 从指定路径恢复训练
 
+Build and run with Docker (recommended for GPU environments):
 
+```bash
+# Build image
+docker build -t astrai:latest .
 
-### 👉 使用指南
+# Run with GPU support
+docker run --gpus all -it astrai:latest
 
-**(1). 与模型对话：**
+# Run with specific GPUs
+docker run --gpus '"device=0,1"' -it astrai:latest
 
-打开 `chat.py` 或使用流式/非流式接口：
+# Run inference server
+docker run --gpus all -p 8000:8000 astrai:latest \
+  python -m scripts.tools.server --port 8000 --device cuda
 
-**流式输出：**
+# Run with volume mount for data
-```python
+docker run --gpus all -v /path/to/data:/data -it astrai:latest
-import torch
-from khaosz import Khaosz
-
-model_dir = "your_model_parameter_dir"
-model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-history = []
-
-while True:
-    query = input(">> ")
-    if query == "!exit":
-        break
-    
-    response_size = 0
-    for response, history in model.stream_generate(
-        query=query, 
-        history=history,
-        temperature=0.85,
-        top_p=0.95,
-        top_k=50
-    ):
-        print(response[response_size:], end="")
-        response_size = len(response)       
 ```
 
-**非流式输出：**
+> **Note**: `--gpus all` is required for CUDA support. Without it, `torch.cuda.is_available()` will return `False`.
-```python
-import torch
-from khaosz import Khaosz
 
-model_dir = "your_model_parameter_dir"
+#### Start HTTP Server
-model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-history = []
 
-while True:
+Start the inference server with OpenAI-compatible HTTP API:
-    query = input(">> ")
+
-    if query == "!exit":
+```bash
-        break
+python -m scripts.tools.server --port 8000 --device cuda
-    
-    response = model.generate(
-        query=query, 
-        history=history,
-        temperature=0.85,
-        top_p=0.95,
-        top_k=50
-    )
-    print(response)
 ```
 
-**(2). 基于检索的生成（RAG）：**
+Make requests:
 
-```python
+```bash
-import torch
+# Chat API (OpenAI compatible)
-from khaosz import Khaosz
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "Hello"}],
+    "max_tokens": 512
+  }'
 
-model_dir = "your_model_parameter_dir"
+# Streaming response
-model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "Tell a story"}],
+    "stream": true,
+    "max_tokens": 500
+  }'
 
-retrieved_content = model.retrieve_generate(
+# Health check
-    query=query,
+curl http://localhost:8000/health
-    retrieve_top_k=5,
+```
-    temperature=0.6,
+
-    top_k=30,
+#### Demo
-    top_p=0.95
+
-)
+Check out the demos in the `scripts/demo/` folder:
-print(retrieved_content)
+
-```
+```bash
+# Download pre‑processed data (required before running demos)
+python scripts/demo/download.py
+
+# Interactive streaming chat
+python scripts/demo/stream_chat.py
+
+# Batch generation
+python scripts/demo/generate_batch.py
+
+# Auto‑regressive generation
+python scripts/demo/generate_ar.py
+```
+
+Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd).
+
+### 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 |
+
+### Contributing
+
+We welcome contributions! Please see our [Contributing Guidelines](CONTRIBUTING.md) for details.
+
+1. Fork the repository.
+2. Create a feature branch.
+3. Commit your changes.
+4. Open a Pull Request.
+
+For major changes, please open an issue first to discuss what you would like to change.
+
+### Community
+
+- **GitHub Issues**: [Issue Tracker](https://github.com/ViperEkura/AstrAI/issues)
+- **Discussions**: [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions)
+- **HuggingFace**: [Model Hub](https://huggingface.co/ViperEk)
+
+### License
+
+This project is licensed under the [GPL-3.0 License](LICENSE).
+
+---
+
+<div align="center">
+  <em>A lightweight Transformer framework designed for both high performance and ease of use.</em>
+</div>

assets/docs/README-zh-CN.md

@@ -0,0 +1,198 @@
+<div align="center">
+  
+  <img src="../images/logo.png" width="auto" alt="Logo">
+  
+  <div>
+    <a href="../../README.md">English</a> • 
+    <a href="#chinese">中文</a>
+  </div>
+  
+  <p>
+    <strong>轻量级 Transformer 训练与推理框架</strong>
+  </p>
+</div>
+
+<div align="center">
+  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
+  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
+  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?color=76bad9" alt="release">
+  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.stargazers_count&label=stars&suffix=%20stars&color=76bad9" alt="stars">
+  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.forks_count&label=forks&suffix=%20forks&color=76bad9" alt="forks">
+</div>
+
+<br>
+
+<div align="center">
+  <a href="../../README.md">English</a> •
+  <a href="#chinese">中文</a> •
+  <a href="https://github.com/ViperEkura/AstrAI/issues">问题追踪</a> •
+  <a href="https://github.com/ViperEkura/AstrAI/discussions">讨论区</a> •
+  <a href="https://huggingface.co/ViperEk">HuggingFace</a>
+</div>
+<br>
+
+## 📖 目录
+
+- [特性](#特性)
+- [快速开始](#快速开始)
+- [文档](#文档)
+- [贡献](#贡献)
+- [社区](#社区)
+- [许可证](#许可证)
+
+---
+
+<a id="chinese"></a>
+## 中文
+
+### 特性
+
+- 🚀 **高性能**: 训练与推理双向优化，高效并行。
+- 🔧 **灵活**: 支持 seq/sft/dpo/grpo 多种训练方式，可定制模型架构。
+- 💡 **易用**: 简洁的 API 与丰富的示例、演示。
+- 📦 **轻量**: 依赖少，部署简单。
+- 🔬 **研究友好**: 模块化设计，便于实验新想法。
+- 🤗 **HuggingFace 集成**: 兼容 HuggingFace 模型与数据集。
+
+### 快速开始
+
+#### 安装
+
+```bash
+git clone https://github.com/ViperEkura/AstrAI.git
+cd AstrAI
+pip install -e .
+```
+
+安装开发依赖：
+
+```bash
+pip install -e ".[dev]"
+```
+
+#### 训练模型
+
+```bash
+python scripts/tools/train.py \
+  --train_type=seq \
+  --data_root_path=/path/to/dataset \
+  --param_path=/path/to/param_path
+```
+
+#### 文本生成
+
+```bash
+python scripts/tools/generate.py --param_path=/path/to/param_path
+```
+
+#### Docker
+
+使用 Docker 构建和运行（推荐用于 GPU 环境）：
+
+```bash
+# 构建镜像
+docker build -t astrai:latest .
+
+# 启用 GPU 运行
+docker run --gpus all -it astrai:latest
+
+# 指定特定 GPU
+docker run --gpus '"device=0,1"' -it astrai:latest
+
+# 运行推理服务
+docker run --gpus all -p 8000:8000 astrai:latest \
+  python -m scripts.tools.server --port 8000 --device cuda
+
+# 挂载数据卷
+docker run --gpus all -v /path/to/data:/data -it astrai:latest
+```
+
+> **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。
+
+#### 启动 HTTP 服务
+
+启动推理服务器，支持 OpenAI 兼容的 HTTP API：
+
+```bash
+python -m scripts.tools.server --port 8000 --device cuda
+```
+
+发起请求：
+
+```bash
+# Chat API（OpenAI 兼容）
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "你好"}],
+    "max_tokens": 512
+  }'
+
+# 流式响应
+curl -X POST http://localhost:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": "讲个故事"}],
+    "stream": true,
+    "max_tokens": 500
+  }'
+
+# 健康检查
+curl http://localhost:8000/health
+```
+
+#### 演示
+
+查看 `scripts/demo/` 文件夹中的演示：
+
+```bash
+# 下载预处理数据（运行演示前必需）
+python scripts/demo/download.py
+
+# 交互式流式聊天
+python scripts/demo/stream_chat.py
+
+# 批量生成
+python scripts/demo/generate_batch.py
+
+# 自回归生成
+python scripts/demo/generate_ar.py
+```
+
+观看 [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd) 上的视频演示。
+
+### 文档
+
+| 文档 | 说明 |
+|------|------|
+| [参数说明](./params.md) | 训练与推理参数配置 |
+| [设计文档](./design.md) | 系统架构与模块设计 |
+| [数据流程](./dataflow.md) | 数据处理管道详解 |
+| [模型介绍](./introduction.md) | 模型架构与技术细节 |
+
+### 贡献
+
+我们欢迎贡献！请参阅[贡献指南](../../CONTRIBUTING.md)了解详情。
+
+1. Fork 本仓库。
+2. 创建功能分支。
+3. 提交更改。
+4. 发起 Pull Request。
+
+重大更改请先开 issue 讨论。
+
+### 社区
+
+- **GitHub Issues**: [问题追踪](https://github.com/ViperEkura/AstrAI/issues)
+- **Discussions**: [GitHub 讨论区](https://github.com/ViperEkura/AstrAI/discussions)
+- **HuggingFace**: [模型中心](https://huggingface.co/ViperEk)
+
+### 许可证
+
+本项目采用 [GPL-3.0 许可证](../../LICENSE)。
+
+---
+
+<div align="center">
+  <em>专为高性能与易用性设计的轻量级 Transformer 框架。</em>
+</div>

assets/docs/architecture.md

@@ -1,220 +0,0 @@
-## 1. 为什么我要做这个项目？
-
-现在市面上有很多大模型，比如GPT、LLaMA这些，动不动就是几十亿甚至上千亿参数。但说实话，这些模型对硬件要求太高了，普通开发者根本玩不起。我就想：**能不能做一个既好用又能在普通电脑上跑起来的模型呢？** 这其实也是目前大部分人的期望， 能有一个可以本地部署的ai小型项目，实现完全私有化并且有一定的智能能力。
-
-于是就有了这个KHAOSZ项目，1B参数，中英双语，支持对话、文本生成、RAG检索，而且训练代码都是开源的！
-
-## 2. 系统架构
-
-系统分为以下板块
-
-```mermaid
-graph LR
-    %% 样式定义
-    classDef config fill:#e1f5fe,stroke:#01579b;
-    classDef trainer fill:#f3e5f5,stroke:#4a148c;
-    classDef data fill:#e8f5e8,stroke:#1b5e20;
-    classDef model fill:#fff3e0,stroke:#e65100;
-    classDef inference fill:#fce4ec,stroke:#880e4f;
-    classDef parallel fill:#e0f2f1,stroke:#004d40;
-    
-    %% 配置模块
-    subgraph Config["Config(配置模块)"]
-        C1[model_config.py]
-        C2[train_config.py]
-        C3[scheduler_config.py]
-    end
-    class Config config;
-
-    %% 训练器模块
-    subgraph Trainer["Trainer(训练器模块)"]
-        T1[trainer.py]
-        T2[train_content.py]
-        T3[schedule.py]
-        T4[strategy.py]
-        T5[train_callback.py]
-    end
-    class Trainer trainer;
-
-    %% 数据模块
-    subgraph Data["Data(数据模块)"]
-        D1[dataset.py]
-        D2[sampler.py]
-        D3[mmap.py]
-        D4[tokenizer.py]
-        D5[checkpoint.py]
-    end
-    class Data data;
-
-    %% 模型模块
-    subgraph Model["Model(模型模块)"]
-        M1[transformer.py]
-        M2[module.py]
-    end
-    class Model model;
-
-    %% 推理模块
-    subgraph Inference["Inference(推理模块)"]
-        I1[generator.py]
-        I2[core.py]
-    end
-    class Inference inference;
-
-    %% 并行模块
-    subgraph Parallel["Parallel(并行模块)"]
-        P1[setup.py]
-        P2[module.py]
-    end
-    class Parallel parallel;
-
-    %% 配置依赖
-    C2 -.-> T1
-    C1 -.-> M1
-    C3 -.-> T3
-
-    %% 训练器内部依赖
-    T1 --> T5
-    T1 --> T2
-    T2 --> T3
-    T2 --> T4
-
-    %% 数据流
-    D1 --> D2
-    D1 --> D3
-    D1 --> D4
-    D1 --> D5
-
-    %% 模型依赖
-    M1 --> M2
-
-    %% 推理依赖
-    I1 --> I2
-
-    %% 跨模块依赖
-    T2 -.-> M1
-    I1 -.-> M1
-    T2 -.-> D1
-    T1 -.-> P1
-```
-
-
-### 1. 配置管理（/config/）
-- **模型配置**：定义模型结构参数（如层数、头数、维度等），通过 `ModelConfig` 统一管理。
-- **训练配置**：设置训练参数（如批次大小、训练阶段 PT/SFT/DPO、优化器等），由 `TrainConfig` 加载。
-- **调度配置**：控制学习率策略（如余弦退火）和训练进度。
-
-### 2. 硬件与并行（/parallel/）
-- **分布式初始化**：通过 `setup_parallel` 函数，根据配置初始化多卡/多机训练环境。
-
-### 3. 数据处理（/data/）
-- **高效加载**：使用内存映射（mmap）技术加载超大语料，避免内存溢出，实现零拷贝读取。
-
-### 4. 模型与训练（/model/, /trainer/）
-- **统一模型架构**：基于 Transformer，支持灵活配置不同规模（如7B、13B）。
-- **策略化训练器**：`Trainer` 根据训练阶段（PT/SFT/DPO）自动切换训练策略，复用同一训练循环。
-- **训练上下文管理**：统一管理模型、优化器、调度器和指标，支持多阶段无缝衔接。
-
-### 5. 推理服务（/inference/, /utils/）
-- **统一生成接口**：提供同步、批量、流式生成方法，适配所有训练阶段。
-- **KV缓存优化**：在自回归生成中缓存 Key/Value，昇腾XPU下利用高速片上内存加速。
-- **RAG支持**：结合检索器和嵌入模型，从外部知识库注入相关信息，提升回答质量。
-- **智能文本分割**：
-  - **结构优先分割**：按标题、段落等切分；
-  - **语义分割**：基于句子嵌入相似度，确保片段语义完整，提升微调效果。
-
-
-## 3. 训练流程
-
-常见大语言模型（Large Language Model, LLM）的训练流程通常包含三个阶段：**预训练（Pre-training, PT）**、**监督微调（Supervised Fine-Tuning, SFT）** 以及 **基于人类反馈的强化学习（Reinforcement Learning from Human Feedback, RLHF）**。本系统设计支持全流程无缝衔接，通过模块化策略实现不同训练阶段的高效切换与状态管理，确保模型能力从通用语言理解逐步对齐至符合人类偏好的对话与指令执行。
-
-### **2.1 预训练阶段**
-
-预训练阶段旨在构建模型的基础语言能力与通用知识表示。该阶段在大规模、无标注的语料库（通常涵盖数百GB至数TB的文本数据）上进行自监督学习。模型架构基于标准的Transformer Decoder，通过掩码语言建模（如因果语言建模）目标进行训练，使模型能够学习词汇、语法、语义及蕴含于文本中的世界知识。
-
-**核心公式：因果语言建模（Causal Language Modeling）**
-
-$$
-L_{\text{PT}} = - \sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
-$$
-
-**符号说明：**
-
-- $T$：序列长度
-- $x_t$：序列中第 $ t $ 个词元（token）
-- $x_{<t}$：位置 $ t $ 之前的所有词元
-- $\theta$：模型参数
-- $P(x_t \mid x_{<t}; \theta)$：模型在给定上文条件下预测下一个词元的概率
-
-
-
-本阶段的核心在于利用分布式的并行计算资源，实现模型参数的稳定优化。训练器模块中的`PTStrategy`策略，专门负责管理预训练特有的数据采样、长序列分段与梯度累积逻辑。同时，硬件适配模块会根据运行环境（如华为昇腾NPU集群或标准GPU集群）自动选择最优的并行通信后端（如HCCL或NCCL），并进行计算图优化，以最大化硬件利用率和训练吞吐量。
-
-另外系统通过数据模块中的高效内存映射加载器（`MmapFileHandler`），实现海量数据的零拷贝读取，以克服传统IO瓶颈。
-
-
-
-### **2.2 监督微调阶段**
-
-预训练模型虽具备强大的语言生成能力，但尚未对齐至遵循人类指令、进行安全有益对话的行为模式。监督微调阶段旨在弥合这一差距。该阶段使用由人工精心编写的、高质量的“指令-响应”配对数据集。
-
-
-
-**核心公式：序列到序列条件语言建模**
-
-设完整序列 $S = [s_1, s_2, \ldots, s_{P+L}]$，其中：
-
-- 前 $P$ 个token是prompt 以及对应控制token： $X = [s_1, \ldots, s_P]$
-- 后 $L$ 个token是response以及对应控制token： $Y = [s_{P+1}, \ldots, s_{P+L}]$
-
-损失函数为：
-
-$$
-L_{\text{SFT}} = - \sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
-$$
-
-训练器模块将动态切换到`SFTStrategy`策略。此策略的核心是引入序列级的监督学习目标，例如预测给定指令下完整、正确的响应序列。训练上下文管理器（`TrainContext`）负责平滑地从PT阶段检查点加载模型状态，并初始化新的优化器和学习率调度器。本阶段不仅优化模型参数，更重要的是引导模型学习“对话”这一特定任务范式，使其输出风格、内容与格式均符合人类期望。
-
-
-
-### **2.3 基于人类反馈的强化学习阶段**
-
-为生成更具帮助性、无害性且符合人类偏好的高质量输出，系统进一步集成强化学习阶段。传统的RLHF流程包括**奖励模型训练**与**策略模型微调**两个核心步骤。系统支持以直接偏好优化(Direct Preference Optimization，DPO)算法为代表的策略微调，并针对稳定性与收敛性进行了多项工程优化。
-
-
-
-#### **2.3.1 传统 RLHF（奖励模型训练）**
-
-$$
-L_{\text{RM}} = -\mathbb{E}_{(x, y_w, y_l) \sim D} \left[ \log \sigma\left( r_\phi(x, y_w) - r_\phi(x, y_l) \right) \right]
-$$
-
-**符号说明：**
-
-- $r_\phi(x, y)$：参数为 $phi$ 的奖励模型给出的标量分数
-- $y_w, y_l $：同一提示 $ x $ 下的优选和劣选回答
-- $\sigma $：sigmoid 函数
-- $\mathcal{D} $：人类偏好数据集
-
-
-#### **2.3.2 DPO 直接偏好优化**（推荐）
-
-$$
-L_{\text{DPO}} = -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]
-$$
-
-**符号说明：**
-
-- $\pi_\theta(y \mid x) $：当前策略模型生成回答的概率
-- $\pi_{\text{ref}}(y \mid x) $：参考模型生成回答的概率
-- $\beta $：温度参数（通常设为 0.1-0.5）
-- 注意：隐式学习奖励函数 $r(x, y) = \beta \log \frac{\pi_\theta(y \mid x)}{\pi_{\text{ref}}(y \mid x)} $
-
-
-在本阶段，训练器模块启用`RLHFStrategy`策略（或类似的`DPOStrategy`直接偏好优化策略）。该策略管理一个复杂的训练循环，其中包含策略模型（待优化的LLM）、参考模型（通常为SFT后的模型快照）和奖励模型。系统流程如下：
-
-1. **偏好数据收集与奖励建模**：首先，通过收集人类标注员对同一提示词下多个模型生成结果的排序偏好数据，训练一个独立的奖励模型（Reward Model, RM）。该模型学习为生成文本输出一个标量奖励分数，以量化其符合人类偏好的程度。
-2. **策略优化**：随后，使用奖励模型作为优化信号，通过强化学习算法对SFT模型（作为策略）进行微调。策略优化的目标是最大化从奖励模型获得的期望累计奖励，同时通过KL散度惩罚项约束策略模型与参考模型的输出分布不过度偏离，以防止模式崩溃并保持生成多样性。训练上下文管理器在此阶段同时维护策略模型、参考模型和奖励模型（或价值函数模型）的状态，并协调复杂的多阶段梯度计算。
-
-
-
-通过上述三阶段的递进式训练，模型完成了从通用语言基座到专业化、高对齐度对话智能体的进化。系统通过统一的`Trainer`接口和策略模式设计，使得各阶段训练在代码层面高度复用，在流程层面清晰解耦，为大规模语言模型的研发与迭代提供了高效、灵活且可扩展的工程基础。

assets/docs/dataflow.md

@@ -0,0 +1,269 @@
+# AstrAI Data Flow Documentation
+
+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.
+
+## 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
+- **Inference Module** (`astrai/inference/`): Inference engine with continuous batching, streaming generation
+- **Config Module** (`astrai/config/`): Model, training, scheduler, and other configurations
+- **Factory Module** (`astrai/factory/`): Registry, BaseFactory for component registration
+- **Parallel Module** (`astrai/parallel/`): Distributed training support
+- **Serialization Module** (`astrai/serialization/`): HDF5 data loading, checkpoint management
+
+The data flow can generally be divided into two main lines: **Training Data Flow** and **Inference Data Flow**.
+
+## Data Flow Diagram
+
+```mermaid
+flowchart LR
+    subgraph A[Data Preparation]
+        direction TB
+        A1[Raw Text] --> A2[AutoTokenizer]
+        A2 --> A3[Serialize to .h5 files]
+        A3 --> A4[BaseDataset]
+        A4 --> A5[ResumableDistributedSampler]
+        A5 --> A6[PyTorch DataLoader]
+    end
+
+    subgraph B[Training]
+        direction TB
+        B1[Batch Data] --> B2[TrainContextBuilder]
+        B2 --> B3[TrainContext]
+        B3 --> B4[BaseStrategy]
+        B4 --> B5[Transformer]
+        B5 --> B6[Compute Loss]
+        B6 --> B7[Backward]
+        B7 --> B8[Optimizer]
+        B8 --> B9[LRScheduler]
+        B9 --> B10[CheckpointCallback]
+    end
+
+    subgraph C[Inference]
+        direction TB
+        C1[Checkpoint] --> C2[AutoModel]
+        C2 --> C3[Transformer + Tokenizer]
+        C3 --> C4[GenerationRequest + apply_chat_template]
+        C4 --> C5[InferenceEngine]
+        C5 --> C6[InferenceScheduler]
+        C6 --> C7[apply_sampling_strategies]
+        C7 --> C8[Transformer Forward]
+        C8 --> C9[KV Cache + Prefix Cache]
+        C9 --> C10{End Condition?}
+        C10 -->|No| C8
+        C10 -->|Yes| C11[Output Text]
+    end
+
+    A --> B
+    B --> C
+```
+
+## Detailed Module Descriptions
+
+### 1. Dataset Module
+
+#### 1.1 Serialization (`serialization.py`)
+- **`save_h5`**: Saves multiple tensors by groups as HDF5 files (`.h5`), each key corresponds to a list of tensors
+- **`load_h5`**: Loads `.h5` files, returns `Dict[str, List[Tensor]]`, supports shared memory (`share_memory=True`)
+- **`Checkpoint` class**: Encapsulates model state dict, training epoch, iteration count; supports safetensors format for saving and loading
+
+#### 1.2 Dataset (`dataset.py`)
+- **`BaseDataset`**: Abstract base class, defines common logic for window sampling, stride, etc.
+- **`BaseSegmentFetcher`** and **`MultiSegmentFetcher`**: Efficiently fetch data from specified index ranges in multiple segments
+- **`DatasetFactory`**: Factory pattern, supports dynamic registration of dataset types (`seq`, `sft`, `dpo`, `grpo`)
+- After dataset loading, multiple data keys (such as `"sequence"`, `"mask"`) are managed through `MultiSegmentFetcher`
+
+#### 1.3 Sampler (`sampler.py`)
+- **`ResumableDistributedSampler`**: Resumable sampler supporting distributed training
+- Records current epoch and iteration position, enabling training resume from breakpoints
+- Supports shuffle and drop_last options
+
+### 2. Model Module
+
+#### 2.1 Transformer / AutoModel (`transformer.py`, `automodel.py`)
+- **`AutoModel`**: Base class for autoregressive language models with `from_pretrained()` and `save_pretrained()` methods
+- **`Transformer`**: Core autoregressive decoder architecture (registered via `@AutoModel.register('transformer')`)
+- Contains embedding layer, multi-layer `DecoderBlock`, RMSNorm, and linear output head
+- Supports weight tying (`tie_weight=True`) to reduce parameter count
+- Uses Rotary Position Embedding (RoPE) to inject position information
+- Supports loading from safetensors format with automatic model type detection from `config.json`
+
+#### 2.2 Submodules (`module.py`)
+- **`RotaryEmbedding`**: Generates RoPE cos/sin cache
+- **`DecoderBlock`**: Contains multi-head attention (supports GQA and MLA), feedforward network (FFN), residual connections
+- **`GQA`**: Grouped Query Attention implementation
+- **`MLA`**: Multi-Latent Attention implementation (like Qwen2-VL)
+- **`MLP`**: Feed-forward network with SiLU activation and gated mechanism
+- **`RMSNorm`**: Layer normalization variant
+- **`Linear`**, **`Embedding`**: Custom linear layer and embedding layer, supporting parallelism wrappers
+
+### 3. Training Module
+
+#### 3.1 Training Context (`train_context.py`)
+- **`TrainContext`**: Data class encapsulating all components needed for training (model, optimizer, data loader, strategy, etc.)
+- **`TrainContextBuilder`**: Builder pattern, progressively assembles training context, supports resume from checkpoint
+
+#### 3.2 Trainer (`trainer.py`)
+- **`Trainer`**: Main training loop, manages callbacks (progress bar, checkpoint, metric logging, gradient clipping, scheduler)
+- Supports distributed training (launches multi-process via `spawn_parallel_fn`)
+- Training steps include:
+  1. `on_train_begin` → 2. `on_epoch_begin` → 3. `on_batch_begin` → 4. Forward/loss calculation → 5. `on_batch_end` → 6. Gradient accumulation → 7. `on_step_begin` → 8. Optimizer update → 9. `on_step_end` → 10. `on_epoch_end`
+
+#### 3.3 Strategy (`strategy.py`)
+- **`BaseStrategy`**: Defines training strategy interface
+- **`SEQStrategy`**: Standard next-token prediction training
+- **`SFTStrategy`**: Supervised Fine-tuning with loss masking
+- **`DPOStrategy`**: Direct Preference Optimization
+- **`GRPOStrategy`**: Group Relative Policy Optimization
+- Strategy receives batch data, executes model forward pass, loss calculation, returns loss tensor
+- Created dynamically by `StrategyFactory` according to configuration
+
+#### 3.4 Scheduler (`schedule.py`)
+- **`BaseScheduler`**: Abstract base class defining learning rate scheduling interface
+- **`CosineScheduler`**: Cosine decay scheduler with warmup
+- **`SGDRScheduler`**: Stochastic Gradient Descent with Warm Restarts
+- **`SchedulerFactory`**: Factory pattern, supports registration of various schedulers
+- Scheduler is automatically created according to configuration and bound to optimizer
+
+#### 3.5 Callbacks (`train_callback.py`)
+- **`TrainCallback`**: Protocol interface for trainer callbacks
+- **`CheckpointCallback`**: Saves model checkpoints at configurable intervals
+- **`ProgressBarCallback`**: Displays training progress
+- **`MetricLoggerCallback`**: Logs training metrics to JSON files
+- **`GradientClippingCallback`**: Clips gradient norms
+- **`SchedulerCallback`**: Steps learning rate scheduler
+
+### 4. Factory Module
+
+#### 4.1 Registry and BaseFactory (`factory.py`)
+- **`Registry`**: Flexible registry for component classes with category and priority support
+- **`BaseFactory`**: Generic factory class for component registration and creation
+- Supports decorator-based registration pattern for extensible components
+- Provides methods for registration, retrieval, and listing with filtering
+
+### 5. Parallel Module
+
+#### 5.1 Setup (`setup.py`)
+- **`spawn_parallel_fn`**: Spawns multiple processes for distributed training using PyTorch multiprocessing
+- **`setup_parallel`**: Context manager for initializing distributed process group (NCCL/CCL backend)
+- **`only_on_rank`**: Decorator to execute functions only on specific ranks
+- **`get_rank`**: Returns current process rank in distributed group
+- **`get_world_size`**: Returns total number of processes in distributed group
+- **`get_current_device`**: Returns current device from environment
+
+#### 5.2 Parallel Layers (`module.py`)
+- **`ParallelModel`**: Base class for parallel models with process group
+- **`ColumnParallelLinear`**: Column-parallel linear layer with input splitting and output gathering
+- **`RowParallelLinear`**: Row-parallel linear layer with output reduction
+
+### 6. Inference Module
+
+#### 6.1 Inference Engine (`engine.py`)
+- **`InferenceEngine`**: Unified inference interface, supports streaming and non-streaming generation
+- **`InferenceScheduler`**: Continuous batching scheduler with dynamic batch composition
+- **`GenerationRequest`**: Encapsulates generation parameters (top_k, top_p, temperature, max_len, messages, etc.)
+- **`messages` format**: List of message dictionaries with `role` (system/user/assistant) and `content`
+- **`apply_chat_template`** (from `tokenizer.py`): Converts messages into prompt string using ChatML format
+- Provides streaming (`stream=True`) and non-streaming (`stream=False`) generation interfaces
+- Supports continuous batching with `max_batch_size` and `max_seq_len` parameters
+- Uses separate model and tokenizer initialization for flexibility
+
+#### 6.2 Scheduler (`scheduler.py`)
+- **`Task`**: Individual generation task with state management (PENDING, RUNNING, FINISHED, ABORTED)
+- **`TaskStatus`**: Task state enumeration
+- **`apply_sampling_strategies`**: Applies temperature, top-k, top-p sampling to logits
+- **`PrefixCacheManager`**: Radix tree-based prefix cache with LRU eviction for efficient KV cache reuse
+- **`RadixNode`**: Tree node structure for prefix caching
+- Continuous batching: new requests can join at any time, completed requests are released immediately
+
+#### 6.3 Server (`server.py`)
+- FastAPI-based HTTP inference server
+- OpenAI-compatible `/v1/chat/completions` endpoint
+- Health check and statistics endpoints
+- Supports both streaming and non-streaming responses
+
+### 7. Tokenizer Module
+
+#### 7.1 Tokenizer (`tokenizer.py`)
+- Implemented based on HuggingFace tokenizers library (Byte-Level BPE)
+- **`AutoTokenizer`**: Auto-loading tokenizer class
+- Supports special tokens: `<｜begin▁of▁sentence｜>`, `<｜end▁of▁sentence｜>`, `<｜▁pad▁｜>`, `<｜im▁start｜>`, `<｜im▁end｜>`
+- Provides `encode`/`decode` methods for mutual conversion between text and token IDs
+- Uses `AutoTokenizer` for loading pre-trained tokenizers
+
+#### 7.2 Chat Template (`chat_template.py`)
+- **`ChatTemplate`**: Jinja2-based chat template with rendering support
+- Handles multi-role message formatting (system, user, assistant)
+- Supports dynamic prompts and generation prompts
+
+## Training Data Flow - Detailed Steps
+
+1. **Data Preparation**
+   - Raw text is converted to token ID sequences through AutoTokenizer
+   - Token ID sequences (possibly with masks, labels, etc.) are saved by groups as `.h5` files
+   - Files can contain multiple segments, each segment corresponds to a tensor
+
+2. **Dataset Loading**
+   - `BaseDataset`'s `load` method calls `load_h5`, obtaining `segments` dictionary
+   - Create `MultiSegmentFetcher` to manage data for multiple keys
+   - Calculate total sample count, and determine start/end indices for each sample based on window size and stride
+
+3. **Sampling and Batch Loading**
+   - `ResumableDistributedSampler` generates index sequence based on current epoch and iteration position
+   - PyTorch `DataLoader` uses sampler to get indices, calls dataset's `__getitem__` to get actual data
+   - Batch data shape is `[batch_size, window_size]` (or varies according to specific dataset type)
+
+4. **Strategy Forward and Loss Calculation**
+   - Batch data is passed to strategy (such as `SEQStrategy`)
+   - Strategy internally calls `Transformer` model, obtaining logits
+   - Calculate cross-entropy loss (or DPO loss, etc.) according to task type
+   - Return loss tensor
+
+5. **Backpropagation and Optimization**
+   - Loss is normalized by dividing by accumulation steps, then `loss.backward()` is executed
+   - After accumulating `accumulation_steps` batches, optimizer `step()` and `zero_grad()` are executed
+   - Learning rate scheduler updates learning rate after each step
+
+6. **Checkpoint Saving**
+   - `CheckpointCallback` saves checkpoints at set intervals
+   - Checkpoints contain model state dict, current epoch, iteration, and other metadata
+   - Saved in safetensors format, ensuring safety and efficiency
+
+## Inference Data Flow - Detailed Steps
+
+1. **Model Loading**
+   - Load `Transformer` model from checkpoint via `AutoModel.from_pretrained()`
+   - Set model to evaluation mode (`model.eval()`), enable inference mode (`torch.inference_mode`)
+
+2. **Prompt Construction and Encoding**
+   - User messages (list of dict with role and content) are converted to ChatML format string through `apply_chat_template` method in tokenizer
+   - Tokenizer encodes prompt string to token ID sequence `input_ids`
+   - For batch generation, use `pad_sequence` for padding
+
+3. **Autoregressive Generation Loop**
+   - Initialize KV cache (optional) and prefix cache
+   - Loop until generating `max_len` tokens or encountering stop token:
+     - Input current `input_ids` (or cached new token) to model, obtain `logits`
+     - Apply `apply_sampling_strategies` (temperature, top-k, top-p) to `logits`
+     - Sample next token ID from the processed distribution
+     - Append new token to `input_ids`, while updating KV cache
+     - For streaming generation, yield each token to caller immediately
+
+4. **Decoding and Output**
+   - Decode generated token ID sequence to text through tokenizer
+   - Remove special tokens, return plain text response
+
+## Checkpoint and Serialization
+
+- **Training Checkpoint**: Saves model parameters, optimizer state, scheduler state, current epoch and iteration
+- **Model Parameters**: Supports safetensors format, automatically handles special logic like weight tying during loading
+- **Dataset Serialization**: HDF5 format supports efficient random access and shared memory, suitable for large-scale pre-training data
+
+## Summary
+
+The data flow design of AstrAI reflects the characteristics of modularity, extensibility, and resumability. The training data flow supports large-scale distributed training through chunk loading, resumable sampling, gradient accumulation, and other mechanisms; the inference data flow achieves efficient text generation using KV cache, prefix caching, and sampling strategies. Clear interfaces between modules facilitate customization and extension.
+
+> Document Update Time: 2026-04-09

assets/docs/design.md

@@ -0,0 +1,694 @@
+## 1. Why I Created This Project
+
+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
+
+```mermaid
+classDiagram
+    namespace config {
+        class ModelConfig {
+            +int vocab_size
+            +int dim
+            +int n_layers
+            +float norm_eps
+            +int dim_ffn
+            +bool tie_weight
+            +int max_len
+            +float rope_theta
+            +int n_heads
+            +int n_kv_heads
+            +bool use_qk_norm
+            +bool use_gated_attention
+            +load(config_path) ModelConfig
+            +save(config_path)
+        }
+
+        class TrainConfig {
+            +nn.Module model
+            +str strategy
+            +Dataset dataset
+            +Callable optimizer_fn
+            +Callable scheduler_fn
+            +int n_epoch
+            +int batch_size
+            +int accumulation_steps
+            +float max_grad_norm
+            +int start_epoch
+            +int start_batch
+            +str ckpt_dir
+            +int ckpt_interval
+            +int random_seed
+            +int num_workers
+            +int prefetch_factor
+            +bool pin_memory
+            +int nprocs
+            +str backend
+            +str master_addr
+            +str master_port
+            +Callable parallel_wrapper
+            +Callable state_dict_fn
+            +List[int] device_ids
+            +str device_type
+            +dict extra_kwargs
+            +validate()
+        }
+
+    }
+
+    namespace dataset {
+        class BaseDataset {
+            +int window_size
+            +int stride
+            +MultiSegmentFetcher fetcher
+            +load(load_path)
+            +__getitem__(index)
+            +__len__()
+        }
+
+        class SEQDataset {
+            +__getitem__(index) Dict
+        }
+
+        class SFTDataset {
+            +__getitem__(index) Dict
+        }
+
+        class DPODataset {
+            +__getitem__(index) Dict
+        }
+
+        class GRPODataset {
+            +__getitem__(index) Dict
+        }
+
+        class BaseSegmentFetcher {
+            +List~Tensor~ segments
+            +List~int~ cum_lengths
+            +int total_length
+            +fetch_data(begin_idx, end_idx) Tensor
+        }
+
+        class MultiSegmentFetcher {
+            +Dict multi_fetchers
+            +List multi_keys
+            +key_fetch(begin_idx, end_idx, keys) Dict
+            +fetch_data(begin_idx, end_idx) Dict
+        }
+
+        class ResumableDistributedSampler {
+            +int start_epoch
+            +int start_iter
+        }
+
+        class DatasetFactory {
+            +Registry _registry
+            +register(name) decorator
+            +create(train_type, window_size, stride) BaseDataset
+            +load(train_type, load_path, window_size, stride) BaseDataset
+        }
+
+        class Checkpoint {
+            +dict state_dict
+            +int epoch
+            +int iteration
+            +save(save_dir)
+            +load(save_dir) Checkpoint
+        }
+    }
+
+    namespace model {
+        class AutoModel {
+            +ModelConfig config
+            +Dict _registry
+            +register(model_type) decorator
+            +get_model_class(model_type) Type
+            +from_pretrained(path, disable_random_init) nn.Module
+            +save_pretrained(save_directory)
+            +to(*args, **kwargs) Self
+        }
+
+        class Transformer {
+            +ModelConfig config
+            +RotaryEmbedding rotary_embedding
+            +Embedding embed_tokens
+            +ModuleList layers
+            +RMSNorm norm
+            +Linear lm_head
+            +forward(input_ids, input_mask, persistent_key_values, start_pos) Dict
+            +load_state_dict(state_dict)
+            +state_dict()
+        }
+
+        class DecoderBlock {
+            +GQA attention
+            +RMSNorm input_norm
+            +MLP mlp
+            +RMSNorm post_attention_norm
+            +forward(x, rotary_emb, attention_mask, kv_cache, start_pos) Tensor
+        }
+
+        class GQA {
+            +int n_heads
+            +int n_kv_heads
+            +int head_dim
+            +Linear q_proj, k_proj, v_proj, o_proj
+            +RMSNorm q_norm, k_norm
+            +forward(x, rotary_emb, mask, kv_cache, start_pos) Tensor
+        }
+
+        class MLA {
+            +int n_heads
+            +int n_kv_heads
+            +int head_dim
+            +Linear q_a_proj, q_b_proj, q_c_proj
+            +Linear kv_a_proj, kv_b_proj, kv_c_proj
+            +Linear o_proj
+            +RMSNorm q_norm, k_norm
+            +forward(x, rotary_emb, mask, kv_cache, start_pos) Tensor
+        }
+
+        class MLP {
+            +Linear up, gate, down
+            +forward(x) Tensor
+        }
+
+        class RMSNorm {
+            +Parameter weight
+            +float norm_eps
+            +forward(x) Tensor
+        }
+
+        class Linear {
+            +Parameter weight
+            +Parameter bias
+            +forward(x) Tensor
+        }
+
+        class RotaryEmbedding {
+            +int dim
+            +int max_len
+            +float base
+            +forward(x, start_pos) Tuple~Tensor, Tensor~
+        }
+
+        class Embedding {
+            +Parameter weight
+            +forward(x) Tensor
+        }
+    }
+
+    namespace tokenize {
+        class AutoTokenizer {
+            +List~str~ stop_ids
+            +int bos_id
+            +int eos_id
+            +int pad_id
+            +vocab_size int
+            +encode(tokens, out_ids, add_special_tokens) List~int~
+            +decode(tokens, skip_special_tokens) str
+            +apply_chat_template(messages, tokenize) Union~str, List[int]~
+            +set_chat_template(template)
+            +load(path)
+            +from_pretrained(path) AutoTokenizer
+            +save_pretrained(save_path)
+        }
+
+        class ChatTemplate {
+            +String template_str
+            +render(messages, add_generation_prompt) str
+            +from_string(template) ChatTemplate
+        }
+    }
+
+    namespace factory {
+        class Registry {
+            +Dict _entries
+            +register(name, component_cls, category, priority)
+            +get(name) Type
+            +list_names() List~str~
+        }
+
+        class BaseFactory {
+            +Registry _registry
+            +register(name, category, priority) decorator
+            +create(name, *args, **kwargs) T
+            +list_registered() list
+        }
+    }
+
+    namespace trainer {
+        class Trainer {
+            +TrainConfig train_config
+            +List~TrainCallback~ callbacks
+            +train(checkpoint)
+            +_build_context(checkpoint) TrainContext
+            +_get_default_callbacks() List~TrainCallback~
+        }
+
+        class TrainContext {
+            +nn.Module model
+            +BaseStrategy strategy
+            +DataLoader dataloader
+            +Optimizer optimizer
+            +LRScheduler scheduler
+            +Checkpoint checkpoint
+            +int epoch
+            +int iteration
+            +float loss
+            +int world_size
+            +int rank
+        }
+
+        class TrainContextBuilder {
+            +TrainConfig config
+            +with_checkpoint(checkpoint) TrainContextBuilder
+            +with_dataloader() TrainContextBuilder
+            +with_strategy() TrainContextBuilder
+            +build() TrainContext
+        }
+
+        class BaseStrategy {
+            +nn.Module model
+            +str device
+            +compute_loss(batch) Tensor
+        }
+
+        class StrategyFactory {
+            +Registry _registry
+            +register(name) decorator
+            +create(model, train_type, device, **kwargs) BaseStrategy
+        }
+
+        class SEQStrategy {
+            +float label_smoothing
+            +compute_loss(batch) Tensor
+        }
+
+        class SFTStrategy {
+            +float label_smoothing
+            +compute_loss(batch) Tensor
+        }
+
+        class DPOStrategy {
+            +nn.Module ref_model
+            +float beta
+            +str reduction
+            +compute_loss(batch) Tensor
+        }
+
+        class GRPOStrategy {
+            +nn.Module ref_model
+            +float clip_eps
+            +float kl_coef
+            +int group_size
+            +compute_loss(batch) Tensor
+        }
+
+        class BaseScheduler {
+            +get_lr() List~float~
+            +step()
+        }
+
+        class SchedulerFactory {
+            +Registry _registry
+            +register(name) decorator
+            +create(optimizer, schedule_type, **kwargs) BaseScheduler
+        }
+
+        class CosineScheduler {
+            +int warmup_steps
+            +int lr_decay_steps
+            +float min_rate
+        }
+
+        class SGDRScheduler {
+            +int warmup_steps
+            +int cycle_length
+            +float min_rate
+            +int t_mult
+        }
+
+        class TrainCallback {
+            +on_train_begin(context)
+            +on_train_end(context)
+            +on_epoch_begin(context)
+            +on_epoch_end(context)
+            +on_step_begin(context)
+            +on_step_end(context)
+            +on_batch_begin(context)
+            +on_batch_end(context)
+            +on_error(context)
+        }
+
+        class GradientClippingCallback {
+            +float max_grad_norm
+            +on_step_begin(context)
+        }
+
+        class SchedulerCallback {
+            +on_train_begin(context)
+            +on_batch_end(context)
+        }
+
+        class CheckpointCallback {
+            +str save_dir
+            +int interval
+            +_save_checkpoint(context)
+            +on_batch_end(context)
+            +on_train_end(context)
+            +on_error(context)
+        }
+
+        class ProgressBarCallback {
+            +int num_epoch
+            +on_epoch_begin(context)
+            +on_batch_end(context)
+            +on_epoch_end(context)
+        }
+
+        class MetricLoggerCallback {
+            +str log_dir
+            +int save_interval
+            +on_batch_end(context)
+            +on_train_end(context)
+        }
+
+        class CallbackFactory {
+            +Registry _registry
+            +register(name) decorator
+            +create(name, **kwargs) TrainCallback
+        }
+    }
+
+    namespace inference {
+        class InferenceEngine {
+            +nn.Module model
+            +AutoTokenizer tokenizer
+            +InferenceScheduler scheduler
+            +int max_batch_size
+            +Optional int max_seq_len
+            +int max_prefix_len
+            +int cache_capacity
+            +Tensor kv_cache
+            +Tensor seq_mask
+            +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
+            +generate_with_request(request) Union[Generator, str, List[str]]
+            +get_stats() Dict
+            +shutdown()
+        }
+
+        class InferenceScheduler {
+            +nn.Module model
+            +AutoTokenizer tokenizer
+            +ModelConfig config
+            +Tuple kv_cache
+            +Tensor seq_mask
+            +PrefixCacheManager prefix_cache
+            +List waiting_queue
+            +List active_tasks
+            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
+            +remove_task(task_id)
+            +start()
+            +stop()
+            +get_stats() Dict
+        }
+
+        class PrefixCacheManager {
+            +RadixNode root
+            +int max_capacity
+            +List lru
+            +insert(token_ids, slot)
+            +find_longest_prefix(token_ids) Tuple[int, int]
+            +release(token_ids)
+        }
+
+        class RadixNode {
+            +Dict children
+            +int hash
+            +int slot
+            +int ref_count
+            +float last_access
+            +List token_sequence
+        }
+
+        class Task {
+            +str task_id
+            +List prompt_ids
+            +int max_tokens
+            +float temperature
+            +float top_p
+            +int top_k
+            +TaskStatus status
+            +List output_ids
+            +int input_tokens
+            +int output_tokens
+            +int slot
+            +Callable stream_callback
+            +is_finished(stop_ids) bool
+        }
+
+        class TaskStatus {
+            +str PENDING
+            +str RUNNING
+            +str FINISHED
+            +str ABORTED
+        }
+
+        class Server {
+            +start()
+            +predict(request)
+        }
+
+        class GenerationRequest {
+            +int top_k
+            +float top_p
+            +float temperature
+            +int max_len
+            +List~Dict~ messages
+            +stream bool
+        }
+
+        class _Result {
+            +List~str~ tokens
+            +List~str~ results
+            +List~bool~ done_flags
+            +append(token, idx)
+            +get_results() List~str~
+        }
+
+        class ChatMessage {
+            +str role
+            +str content
+        }
+
+        class ChatCompletionRequest {
+            +List~ChatMessage~ messages
+            +float temperature
+            +float top_p
+            +int top_k
+            +int max_tokens
+            +bool stream
+            +Optional~str~ system_prompt
+        }
+
+        class CompletionResponse {
+            +str id
+            +str object
+            +int created
+            +str model
+            +List~Dict~ choices
+        }
+    }
+
+    namespace parallel {
+        class ParallelSetup {
+            +spawn_parallel_fn(fn, nprocs)
+            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type, device_ids)
+        }
+
+        class ParallelModel {
+            +dist.ProcessGroup process_group
+            +int rank
+            +int world_size
+        }
+
+        class ColumnParallelLinear {
+            +forward(x) Tensor
+        }
+
+        class RowParallelLinear {
+            +forward(x) Tensor
+        }
+    }
+
+    %% Relationships
+    TrainConfig --> ModelConfig : uses
+    TrainConfig --> BaseDataset : uses
+    TrainConfig --> StrategyFactory : selects
+    StrategyFactory ..> BaseStrategy : creates
+    BaseStrategy <|-- SEQStrategy
+    BaseStrategy <|-- SFTStrategy
+    BaseStrategy <|-- DPOStrategy
+    BaseStrategy <|-- GRPOStrategy
+    DPOStrategy --> Transformer : uses
+    GRPOStrategy --> Transformer : uses
+    Trainer --> TrainConfig : configures
+    Trainer --> TrainContextBuilder : builds
+    Trainer --> TrainCallback : manages
+    TrainContextBuilder --> TrainContext : creates
+    TrainContext --> Checkpoint : manages
+    TrainContext --> BaseStrategy : uses
+    TrainContext --> BaseScheduler : uses
+    AutoModel --> ModelConfig : contains
+    SchedulerFactory ..> BaseScheduler : creates
+    BaseScheduler <|-- CosineScheduler
+    BaseScheduler <|-- SGDRScheduler
+    CallbackFactory ..> TrainCallback : creates
+    TrainCallback <|-- GradientClippingCallback
+    TrainCallback <|-- SchedulerCallback
+    TrainCallback <|-- CheckpointCallback
+    TrainCallback <|-- ProgressBarCallback
+    TrainCallback <|-- MetricLoggerCallback
+    InferenceEngine --> InferenceScheduler : uses
+    InferenceScheduler --> Task : manages
+    InferenceScheduler --> TaskStatus : uses
+    InferenceScheduler --> Transformer : uses
+    InferenceEngine --> Transformer : uses
+    InferenceEngine --> GenerationRequest : uses
+    Server --> InferenceEngine : uses
+    Server --> ChatMessage : uses
+    Server --> ChatCompletionRequest : uses
+    Server --> CompletionResponse : uses
+    ParallelSetup --> Trainer : enables
+    BaseDataset <|-- SEQDataset
+    BaseDataset <|-- SFTDataset
+    BaseDataset <|-- DPODataset
+    BaseDataset <|-- GRPODataset
+    DatasetFactory ..> BaseDataset : creates
+    BaseSegmentFetcher --> MultiSegmentFetcher : used by
+    MultiSegmentFetcher --> BaseDataset : used by
+    AutoModel <|-- Transformer
+    AutoModel --> ModelConfig : contains
+    Transformer --> DecoderBlock : uses
+    Transformer --> RotaryEmbedding : uses
+    Transformer --> Embedding : uses
+    DecoderBlock --> GQA : uses
+    DecoderBlock --> MLA : uses
+    DecoderBlock --> MLP : uses
+    DecoderBlock --> RMSNorm : uses
+    TrainContextBuilder --> ResumableDistributedSampler : creates
+    ResumableDistributedSampler --> BaseDataset : samples
+    ParallelModel <|-- RowParallelLinear
+    ParallelModel <|-- ColumnParallelLinear
+    AutoTokenizer --> ChatTemplate : uses
+    InferenceScheduler --> PrefixCacheManager : uses
+    InferenceScheduler --> RadixNode : uses
+    Checkpoint ..> Checkpoint : saves/loads
+    TrainConfig --> DatasetFactory : selects
+    TrainConfig --> SchedulerFactory : selects
+    TrainConfig --> CallbackFactory : selects
+    AutoModel ..> AutoTokenizer : loads with
+    BaseFactory <|-- DatasetFactory
+    BaseFactory <|-- StrategyFactory
+    BaseFactory <|-- SchedulerFactory
+    BaseFactory <|-- CallbackFactory
+```
+
+### Module Overview
+
+| Module | Components | Description |
+|--------|------------|-------------|
+| **astrai.config** | ModelConfig, TrainConfig | Configuration management |
+| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory, Checkpoint | Dataset loading and management |
+| **astrai.model** | AutoModel, Transformer, DecoderBlock, GQA, MLA, MLP, 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, Task, TaskStatus, Server, GenerationRequest, PrefixCacheManager, ChatMessage, ChatCompletionRequest, CompletionResponse | Inference service with continuous batching |
+| **astrai.parallel** | ParallelSetup, ColumnParallelLinear, RowParallelLinear | Distributed parallel |
+| **astrai.factory** | Registry, BaseFactory | Generic component registration |
+
+### 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) |
+| **Singleton** | `TrainContext` | Training process global state management |
+| **Registry** | `BaseFactory`, `Registry` | Generic component registration with category and priority support |
+| **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** | `_Result`, `GenerationRequest` | Event-based result notification for streaming/non-streaming generation |
+
+### Core Relationships
+
+1. **Configuration → Training**: `TrainConfig` contains `ModelConfig`, holds model, dataset, optimizer and other 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**: `Server` → `InferenceEngine` → `InferenceScheduler` → `Transformer`, supports continuous batching with streaming/non-streaming
+5. **Distributed Support**: `ParallelSetup` provides 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
+
+## 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}
+$$
+
+In this implementation, an off-policy approach is used ($\pi_\theta = \pi_{\text{ref}}$), and the policy loss simplifies to:
+
+$$
+L_{\text{policy}} = -\mathbb{E}[A]
+$$
+
+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-04-09

assets/docs/introduction.md

@@ -1,50 +1,82 @@
-## 模型介绍
+## 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 32 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.
 
-### 1. 模型搭建
+The model now uses the **AutoModel** base class for flexible loading and saving:
 
-本模型采用Transformer架构， 使用GQA（q_head=24, kv_head=4） 机制，相较于传统的MHA可以节省KV cache 的显存占用（但是目前没有做KV cache），通过堆叠24层Transformer实现模型的搭建， 参数量为1.0b。Transformer 是自回归模型， 是通过计算前面所有的token的关系得到下一个token的概率分布
+```python
+from astrai.model import AutoModel
 
-![structure](../images/structure.png)
+# Load model from checkpoint
+model = AutoModel.from_pretrained("path/to/model")
 
-什么是自回归模型呢， 在把句子拆分成token之后, 模型会预测下一个token的概率分布。这意味着模型会根据给定的上下文（即已经出现的tokens序列），计算出下一个可能的token及其对应的概率。
+# Save model to new directory
-
+model.save_pretrained("path/to/save")
-
-
-#### 1. 自回归
-
-假设我们有一个句子被拆分成如下tokens列表：
-
-```
-["你好", "，" "今天", "天气"]
 ```
 
-接下来，模型会基于这个序列预测下一个可能出现的token。这通常以概率分布的形式给出，比如：
+The Transformer model is registered via `@AutoModel.register('transformer')` decorator, allowing easy extension for new model types.
 
-```
+```mermaid
--> {"token": "不错", "probability": 0.4}
+flowchart TB
--> {"token": "晴朗", "probability": 0.2}
+    subgraph Layers["Transformer Layers"]
--> ......
+        direction TB
+        A[Input Embedding] --> B[Transformer Block\nLayer 1]
+        B --> C[Transformer Block\nLayer ...]
+        C --> D[Transformer Block\nLayer 32]
+        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]
+        V --> W[SiLU]
+        V --> X[×]
+        W --> X
+        X --> Y[Linear]
+        Y --> Z[+]
+        T --> Z
+        Z --> AA[x']
+    end
+
+    classDef main fill:#e6f3ff,stroke:#0066cc;
+    classDef block fill:#fff2e6,stroke:#cc6600;
+    class Layers main;
+    class TransformerBlock block;
 ```
 
-这里，“不错”和“晴朗”是两个可能跟随在“天气”之后的tokens，并且给出了每个token成为下一个token的可能性大小。
+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).
 
-之后，我们通过采样（通过top_k, top_p, temperature参数调整采样后的结果）得到下一个token并且将下一个token加入序列作为输入
+#### 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.
-["你好", "，" "今天", "天气", "不错"]
-```
 
-之后都是在重复这个流程， 直到遇到控制流程结束的token（<|end_of_seqence|>）模型停止处理（一般模型都会设置控制token， 不然模型会一直输出到显存爆炸）。 
+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:
-
-#### 2. 因果掩码
-
-transformer 中采用注意力机制，输入的形状一般为[bsz, seq_len]， 输出为[bsz, seq_len，n_dim]， 为了实现预测下一个token， 模型的输入和输出必须错开来一个位置。模型预测的target必须错开一个位置， 在训练的时候我们也采用错开一个位置的方法
 
 ```
 sequence : [[1, 2, 3, 4, 5, 6]]
@@ -52,18 +84,14 @@ 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.
 
-其中注意力得分代表了模型对两个token之间相似程度的关注程度
+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:
-
-对于decoder only结构的模型， 为了防止模型从未来的位置偷到信息， 在注意力的计算过程中需要增加掩码，我们需要在注意力得分计算之前应用一个掩码。这个掩码通常是一个下三角矩阵，对于长度为n的序列，它的形状是[n, n]。下面以一个长度为5的序列为例，展示如何创建这样的因果掩码矩阵：
 
 ```
 [[0, -inf, -inf, -inf, -inf],
@@ -73,17 +101,199 @@ $$ s_{ij} := s_{ij} + mask_{ij} $$
  [0,    0,    0,    0,    0]]
 ```
 
-在这个矩阵中，0表示可以注意到的位置，而-inf表示应该被掩盖（即不应注意到）的位置。因为这个句子保证了注意力得分中 $j > i$  的部分通过softmax 之后由`inf` 变成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.
-#### 3. 旋转位置编码
-
-旋转位置编码（Rotary Position Embedding, RoPE）是一种为了解决Transformer模型中缺乏对序列位置信息直接建模的问题而设计的位置编码方法。与传统的位置编码（如正弦和余弦函数的位置编码）不同，RoPE通过将位置信息直接嵌入到查询（Query, Q）和键（Key, K）向量中来实现，使得模型能够更自然地处理序列中的相对位置关系。
-
 
 $$ 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 $$
 
-其中的 $R_{i-j}$ 控制了模型的不同token 在不同相对距离上注意力的衰减，在 $i - 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,
+    max_batch_size=8,
+    max_seq_len=4096,
+)
+
+# 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_len=1024,
+    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 | 0.8 | Sampling temperature (0.0-2.0) |
+| `top_p` | float | 0.95 | Nucleus sampling threshold |
+| `top_k` | int | 50 | Top-k sampling parameter |
+| `max_tokens` | int | 2048 | Maximum tokens to generate |
+| `stream` | bool | false | Enable streaming response |
+| `system_prompt` | str | None | System prompt override |
+
+**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"
+    }
+  ]
+}
+```
+
+### 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`.
+
+### Simple Generation Endpoint
+
+For basic text generation without chat format:
+
+```bash
+curl -X POST "http://localhost:8000/generate?query=Hello&max_len=1000" \
+  -H "Content-Type: application/json"
+```
+
+Or with conversation history:
+
+```bash
+curl -X POST "http://localhost:8000/generate" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "query": "What is AI?",
+    "history": [["Hello", "Hi there!"], ["How are you?", "I'm doing well"]],
+    "temperature": 0.8,
+    "max_len": 2048
+  }'
+```
+
+### Health Check
+
+Monitor server and model status:
+
+```bash
+curl http://localhost:8000/health
+# {"status": "ok", "model_loaded": true, "engine_ready": true}
+
+curl http://localhost:8000/stats
+# {"requests_total": 10, "tokens_generated": 5000, ...}
+```
+
+> Document Update Time: 2026-04-09

assets/docs/kvcache.md

@@ -1,27 +0,0 @@
-## kv_cache 实现
-
-根据注意力的计算公式
-
-$$
-\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*}
-$$
-
-由于模型是自回归模型, 我们只用求序列最后一个部分，也就是说 $ i $ 的下标是确定的, 是序列最后一个元素, 我们求的是 $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*}
-$$
-
-如果我们把式子展开
-
-$$
-o_n = \sum_j \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}}\right)v_{j}
-$$
-
-以上表达式只有k和v存在长度下标, 而 $q$ 没有， 所以计算过程中 $q$ 的输入是确定的上次输入的最后一个token, 而 $k,  v$ 是需要对不同长度的部分进行缓存的，同时缓存的时候应该注意位置编码的计算应该在kvcache的计算之前进行，否则会存在位置编码的计算错误

assets/docs/params.md

@@ -0,0 +1,141 @@
+# Parameter Documentation
+
+## Training Parameters
+
+### Basic Parameters
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `--train_type` | Training type (seq, sft, dpo, grpo) | required |
+| `--model_type` | Model type for AutoModel loading (e.g., transformer) | transformer |
+| `--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 | 4 |
+| `--accumulation_steps` | Gradient accumulation steps | 1 |
+
+### Learning Rate Scheduling
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `--warmup_steps` | Warmup steps | 1000 |
+| `--max_lr` | Maximum learning rate (warmup + cosine decay) | 3e-4 |
+| `--max_grad_norm` | Maximum gradient norm | 1.0 |
+
+### Checkpoint
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `--ckpt_interval` | Checkpoint save interval (iterations) | 5000 |
+| `--ckpt_dir` | Checkpoint save directory | checkpoint |
+| `--resume_dir` | Resume training from specified path | - |
+
+### Optimizer Parameters
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `--adamw_beta1` | AdamW beta1 | 0.9 |
+| `--adamw_beta2` | AdamW beta2 | 0.95 |
+| `--adamw_weight_decay` | AdamW weight decay | 0.01 |
+
+### Data Loading
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `--random_seed` | Random seed | 3407 |
+| `--num_workers` | DataLoader workers | 0 |
+| `--prefetch_factor` | Prefetch factor for dataloader | None |
+| `--pin_memory` | Enable pin_memory | False |
+| `--no_pin_memory` | Disable pin_memory | - |
+
+### Distributed Training
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `--nprocs` | Number of GPUs | 1 |
+| `--device_type` | Device type (cuda/cpu) | cuda |
+
+### Other Parameters
+
+| Parameter | Description | Default Value |
+|-----------|-------------|---------------|
+| `--window_size` | Maximum input sequence length | model config max_len |
+| `--stride` | Input sequence stride | - |
+| `--dpo_beta` | DPO beta value | 0.1 |
+| `--grpo_clip_eps` | GRPO clip epsilon | 0.2 |
+| `--grpo_kl_coef` | GRPO KL coefficient | 0.01 |
+| `--grpo_group_size` | GRPO group size | 4 |
+| `--label_smoothing` | Label smoothing parameter | 0.1 |
+| `--start_epoch` | Starting epoch | 0 |
+| `--start_batch` | Starting batch | 0 |
+
+---
+
+## 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_len` | Maximum generation length | 1024 |
+| `stream` | Whether to stream output | False |
+
+### Usage Example
+
+```python
+import torch
+from astrai.model import AutoModel
+from astrai.tokenize import Tokenizer
+from astrai.inference import InferenceEngine, GenerationRequest
+
+# Load model using AutoModel
+model = AutoModel.from_pretrained("your_model_dir")
+
+# Load tokenizer
+tokenizer = Tokenizer("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_len=1024,
+)
+
+# 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-04-09

astrai/__init__.py

@@ -0,0 +1,32 @@
+__version__ = "1.3.3"
+__author__ = "ViperEkura"
+
+from astrai.config import (
+    ModelConfig,
+    TrainConfig,
+)
+from astrai.dataset import DatasetFactory
+from astrai.factory import BaseFactory
+from astrai.inference import (
+    GenerationRequest,
+    InferenceEngine,
+)
+from astrai.model import AutoModel, Transformer
+from astrai.tokenize import AutoTokenizer
+from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer
+
+__all__ = [
+    "Transformer",
+    "ModelConfig",
+    "TrainConfig",
+    "DatasetFactory",
+    "AutoTokenizer",
+    "GenerationRequest",
+    "InferenceEngine",
+    "Trainer",
+    "CallbackFactory",
+    "StrategyFactory",
+    "SchedulerFactory",
+    "BaseFactory",
+    "AutoModel",
+]

astrai/config/__init__.py

@@ -0,0 +1,8 @@
+from astrai.config.model_config import ModelConfig
+from astrai.config.train_config import TrainConfig
+
+__all__ = [
+    # Model configuration
+    "ModelConfig",
+    "TrainConfig",
+]

astrai/config/model_config.py

@@ -1,5 +1,4 @@
 import json
-
 from dataclasses import asdict, dataclass
 from typing import Optional, Self
 
@@ -7,6 +6,7 @@ from typing import Optional, Self
 @dataclass
 class ModelConfig:
     # basic config
+    model_type: Optional[str] = None
     vocab_size: Optional[int] = None
     dim: Optional[int] = None
 
@@ -14,30 +14,29 @@ class ModelConfig:
     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
     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:
+        with open(config_path, "r") as f:
-            config.update(json.load(f)) 
+            config.update(json.load(f))
-            
+
         for key, value in config.items():
             if hasattr(self, key):
                 setattr(self, key, value)
-    
+
         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:
+        with open(config_path, "w") as f:
             json.dump(config_dict, f, indent=4)

astrai/config/train_config.py

@@ -0,0 +1,101 @@
+from dataclasses import dataclass, field
+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
+
+
+@dataclass
+class TrainConfig:
+    # 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."})
+    optimizer_fn: Callable[[nn.Module], Optimizer] = field(
+        default=None, metadata={"help": "Optimizer factory for training."}
+    )
+    scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
+        default=None, metadata={"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(
+        default=1, metadata={"help": "Number of iterations between steps."}
+    )
+    max_grad_norm: float = field(
+        default=1.0, metadata={"help": "Maximum gradient norm."}
+    )
+
+    # checkpoint setting
+    start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
+    start_batch: int = field(
+        default=0, metadata={"help": "Start batch iteration for training."}
+    )
+    ckpt_dir: str = field(
+        default="./checkpoint", metadata={"help": "Checkpoint directory."}
+    )
+    ckpt_interval: int = field(
+        default=5000, metadata={"help": "Number of iterations between checkpoints."}
+    )
+
+    # dataloader setting
+    random_seed: int = field(default=3407, metadata={"help": "Random seed."})
+    num_workers: int = field(
+        default=0, metadata={"help": "Number of workers for dataloader."}
+    )
+    prefetch_factor: Optional[int] = field(
+        default=None, metadata={"help": "Prefetch factor for dataloader."}
+    )
+    pin_memory: bool = field(
+        default=False, metadata={"help": "Pin memory for dataloader."}
+    )
+
+    # distributed training
+    nprocs: int = field(
+        default=1, metadata={"help": "Number of processes for distributed training."}
+    )
+    backend: str = field(
+        default="nccl", metadata={"help": "Distributed training backend."}
+    )
+    master_addr: str = field(
+        default="localhost",
+        metadata={"help": "Master address for distributed training."},
+    )
+    master_port: str = field(
+        default="29500", metadata={"help": "Master port for distributed training."}
+    )
+    parallel_wrapper: Optional[Callable] = field(
+        default=None, metadata={"help": "Parallel function for training."}
+    )
+    state_dict_fn: Optional[Callable] = field(
+        default=None, metadata={"help": "Parallel function for state  dict saving."}
+    )
+
+    # others
+    device_ids: Optional[List[int]] = field(
+        default=None, metadata={"help": "Device ids for distributed training."}
+    )
+    device_type: str = field(
+        default="cuda", metadata={"help": "Device type for distributed training."}
+    )
+    extra_kwargs: dict = field(
+        default_factory=dict, metadata={"help": "Other arguments."}
+    )
+
+    def __post_init__(self):
+        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.")

astrai/dataset/__init__.py

@@ -0,0 +1,19 @@
+from astrai.dataset.dataset import (
+    BaseDataset,
+    BaseSegmentFetcher,
+    DatasetFactory,
+    MultiSegmentFetcher,
+)
+from astrai.dataset.sampler import ResumableDistributedSampler
+
+__all__ = [
+    # Base classes
+    "BaseDataset",
+    # Factory
+    "DatasetFactory",
+    # Fetchers
+    "BaseSegmentFetcher",
+    "MultiSegmentFetcher",
+    # Sampler
+    "ResumableDistributedSampler",
+]

astrai/dataset/dataset.py

@@ -0,0 +1,338 @@
+"""Dataset implementations with factory pattern for training."""
+
+import bisect
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Union
+
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+
+from astrai.factory import BaseFactory
+from astrai.serialization import load_h5
+
+
+class BaseSegmentFetcher:
+    """Fetches data segments across multiple tensor segments.
+
+    Maintains cumulative lengths for efficient range queries across
+    multiple discontinuous segments.
+    """
+
+    def __init__(self, segments: List[Tensor]):
+        self.segments = segments
+        self.cum_lengths = []
+
+        total = 0
+        for seg in segments:
+            total += torch.numel(seg)
+            self.cum_lengths.append(total)
+
+        self.total_length = total
+
+    def __len__(self) -> int:
+        return self.total_length
+
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        """Fetch data in the range [begin_idx, end_idx).
+
+        Args:
+            begin_idx: Starting index (inclusive)
+            end_idx: Ending index (exclusive)
+
+        Returns:
+            Concatenated tensor of data in the specified range
+        """
+        if not (
+            0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length
+        ):
+            raise ValueError("begin_idx or end_idx out of bounds")
+        if begin_idx >= end_idx:
+            return torch.tensor([], dtype=torch.long)
+
+        # Find segment boundaries for the range
+        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
+        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
+
+        result_segments = []
+
+        for i in range(seg_start_idx, seg_end_idx + 1):
+            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
+            start = max(begin_idx - prev_cum, 0)
+            end = min(end_idx - prev_cum, len(self.segments[i]))
+            data = self.segments[i][start:end]
+            result_segments.append(data)
+
+        return torch.cat(result_segments, dim=0)
+
+
+class MultiSegmentFetcher:
+    """Manages multiple segment fetchers for different data keys.
+
+    Each key corresponds to a different type of data (e.g., "sequence", "mask").
+    """
+
+    def __init__(self, multi_segments: Dict):
+        self.multi_keys = list(multi_segments.keys())
+        self.multi_fetchers = {
+            key: BaseSegmentFetcher(segments)
+            for key, segments in multi_segments.items()
+        }
+
+    def __len__(self) -> int:
+        """Returns the minimum length across all fetchers."""
+        len_list = [len(seg) for seg in self.multi_fetchers.values()]
+        return min(len_list)
+
+    def key_fetch(
+        self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]
+    ) -> Dict:
+        """Fetch data for specific keys.
+
+        Args:
+            begin_idx: Starting index
+            end_idx: Ending index
+            keys: Single key or list of keys to fetch
+
+        Returns:
+            Dictionary of tensors if multiple keys, single tensor if one key
+        """
+        fetch_dict = {}
+        keys = [keys] if isinstance(keys, str) else keys
+
+        for key in keys:
+            fetcher = self.multi_fetchers[key]
+            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
+            fetch_dict[key] = fetch_tensor
+
+        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
+
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
+        """Fetch all keys."""
+        return self.key_fetch(begin_idx, end_idx, self.multi_keys)
+
+
+class BaseDataset(Dataset, ABC):
+    """Abstract base class for all dataset types.
+
+    Implements common functionality for window-based data fetching.
+    """
+
+    def __init__(self, window_size: int, stride: int):
+        super().__init__()
+        self.segments = {}
+        self.window_size = window_size
+        self.stride = stride
+        self.total_samples = None
+        self.fetcher: Optional[MultiSegmentFetcher] = None
+
+    def load(self, load_path: str):
+        """Load dataset from HDF5 file.
+
+        Args:
+            load_path: Path to the HDF5 data file
+        """
+        self.segments = load_h5(load_path)
+        self.fetcher = MultiSegmentFetcher(self.segments)
+        self.total_samples = len(self.fetcher)
+
+    def get_index(self, index: int) -> tuple:
+        """Calculate begin and end indices for a sample.
+
+        Args:
+            index: Sample index
+
+        Returns:
+            Tuple of (begin_idx, end_idx)
+        """
+        assert self.total_samples > self.window_size
+
+        begin_idx = min(index * self.stride, self.total_samples - 1 - self.window_size)
+        end_idx = min(begin_idx + self.window_size, self.total_samples - 1)
+
+        return begin_idx, end_idx
+
+    @abstractmethod
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        """Get a single sample by index.
+
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError
+
+    def __len__(self) -> int:
+        assert self.total_samples is not None
+        if self.total_samples <= self.window_size:
+            return 0
+        return (self.total_samples - 1 - self.window_size) // self.stride + 1
+
+
+class DatasetFactory(BaseFactory["BaseDataset"]):
+    """Factory class for creating dataset instances.
+
+    Supports decorator-based registration for extensible dataset types.
+    All default dataset types (seq, sft, dpo, grpo) are registered automatically
+    when their classes are defined with the decorator.
+
+    Example usage:
+        @DatasetFactory.register("custom")
+        class CustomDataset(BaseDataset):
+            ...
+
+        dataset = DatasetFactory.create("custom", window_size, stride)
+    """
+
+    @classmethod
+    def _validate_component(cls, dataset_cls: type) -> None:
+        """Validate that the dataset class inherits from BaseDataset."""
+        if not issubclass(dataset_cls, BaseDataset):
+            raise TypeError(f"{dataset_cls.__name__} must inherit from BaseDataset")
+
+    @classmethod
+    def create(cls, train_type: str, window_size: int, stride: int) -> "BaseDataset":
+        """Create a dataset instance.
+
+        Args:
+            train_type: Type of training ("seq", "sft", "dpo", "grpo")
+            window_size: Window size for data sampling
+            stride: Stride between consecutive samples
+
+        Returns:
+            Dataset instance
+        """
+        return super().create(train_type, window_size, stride)
+
+    @classmethod
+    def load(
+        cls,
+        train_type: str,
+        load_path: str,
+        window_size: int,
+        stride: Optional[int] = None,
+    ) -> "BaseDataset":
+        """Create and load a dataset in one step.
+
+        Args:
+            train_type: Type of training dataset
+            load_path: Path to the data file
+            window_size: Window size for data sampling
+            stride: Stride between consecutive samples (default: same as window_size)
+
+        Returns:
+            Loaded dataset instance
+        """
+        if stride is None:
+            stride = window_size
+
+        dataset = cls.create(train_type, window_size, stride)
+        dataset.load(load_path)
+
+        return dataset
+
+    @classmethod
+    def available_types(cls) -> list:
+        """Return list of registered dataset type names."""
+        return cls.list_registered()
+
+
+# ============== Dataset Classes ==============
+# All dataset classes are registered at class definition time using the decorator
+
+
+@DatasetFactory.register("seq")
+class SEQDataset(BaseDataset):
+    """Dataset for sequential next-token prediction training."""
+
+    def __init__(self, window_size: int, stride: int):
+        super().__init__(window_size, stride)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
+
+    def __getitem__(self, index):
+        begin_idx, end_idx = self.get_index(index)
+
+        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
+
+        return {"input_ids": x, "target_ids": y}
+
+
+@DatasetFactory.register("sft")
+class SFTDataset(BaseDataset):
+    """Dataset for supervised fine-tuning with loss masking."""
+
+    def __init__(self, window_size: int, stride: int):
+        super().__init__(window_size, stride)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+
+    def __getitem__(self, index):
+        begin_idx, end_idx = self.get_index(index)
+
+        x = self._fetch_data(begin_idx, end_idx, "sequence").to(dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(
+            dtype=torch.long
+        )
+        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask").to(
+            dtype=torch.bool
+        )
+
+        return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask}
+
+
+@DatasetFactory.register("dpo")
+class DPODataset(BaseDataset):
+    """Dataset for Direct Preference Optimization training."""
+
+    def __init__(self, window_size: int, stride: int):
+        super().__init__(window_size, stride)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+
+    def __getitem__(self, index: int):
+        begin_idx, end_idx = self.get_index(index)
+
+        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
+        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
+        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(
+            dtype=torch.bool
+        )
+        rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(
+            dtype=torch.bool
+        )
+
+        return {
+            "chosen": chosen,
+            "rejected": rejected,
+            "chosen_mask": chosen_mask,
+            "rejected_mask": rejected_mask,
+        }
+
+
+@DatasetFactory.register("grpo")
+class GRPODataset(BaseDataset):
+    """Dataset for Group Relative Policy Optimization training."""
+
+    def __init__(self, window_size: int, stride: int):
+        super().__init__(window_size, stride)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        begin_idx, end_idx = self.get_index(index)
+
+        prompts = self._fetch_data(begin_idx, end_idx, "prompts")
+        responses = self._fetch_data(begin_idx, end_idx, "responses")
+        masks = self._fetch_data(begin_idx, end_idx, "masks")
+        rewards = self._fetch_data(begin_idx, end_idx, "rewards")
+
+        return {
+            "prompts": prompts,
+            "responses": responses,
+            "masks": masks,
+            "rewards": rewards,
+        }

astrai/dataset/sampler.py

@@ -1,51 +1,51 @@
+from typing import Optional
+
 import torch
 import torch.distributed as dist
-
 from torch.utils.data import Dataset, Sampler
-from typing import Optional
 
 
 class ResumableDistributedSampler(Sampler[int]):
     def __init__(
-        self, 
+        self,
         data_source: Dataset,
-        start_epoch: int=0, 
+        start_epoch: int = 0,
-        start_iter: int=0, 
+        start_iter: int = 0,
-        seed: int=42,
+        seed: int = 42,
-        drop_last: bool=False,
+        drop_last: bool = False,
-        shuffle: bool=True,
+        shuffle: bool = True,
-        process_group: Optional[dist.ProcessGroup]=None,
+        process_group: Optional[dist.ProcessGroup] = None,
     ):
         self.epoch = start_epoch
         self.iter = start_iter
         self.seed = seed
         self.num_samples = len(data_source)
-        
+
         if process_group is not None:
             # input process group
             self.rank = dist.get_rank(process_group)
             self.num_replicas = dist.get_world_size(process_group)
-            
+
         elif dist.is_available() and dist.is_initialized():
             # use default process group
             process_group = dist.group.WORLD
             self.rank = dist.get_rank()
             self.num_replicas = dist.get_world_size()
-            
+
         else:
             # single process
             self.rank = 0
             self.num_replicas = 1
-        
+
         self.drop_last = drop_last
         self.shuffle = shuffle
-        
+
-        offset = 0 if drop_last else  self.num_replicas - 1
+        offset = 0 if drop_last else self.num_replicas - 1
         self.num_samples_per_replica = (self.num_samples + offset) // self.num_replicas
         self.total_size = self.num_samples_per_replica * self.num_replicas
-        
+
         self._indices = None
-    
+
     def _get_indices(self):
         if self.shuffle:
             generator = torch.Generator()
@@ -53,26 +53,26 @@ class ResumableDistributedSampler(Sampler[int]):
             indices = torch.randperm(self.num_samples, generator=generator).tolist()
         else:
             indices = torch.arange(self.num_samples).tolist()
-        
+
         if not self.drop_last and self.num_samples < self.total_size:
             padding_size = self.total_size - len(indices)
             indices += indices[:padding_size]
-        
+
-        local_indices = indices[self.rank:self.total_size:self.num_replicas]
+        local_indices = indices[self.rank : self.total_size : self.num_replicas]
-        
+
         self.iter = self.iter % self.num_samples_per_replica
-        self._indices = local_indices[self.iter:]
+        self._indices = local_indices[self.iter :]
-    
+
     def __iter__(self):
         if self._indices is None:
             self._get_indices()
-        
+
         for i in self._indices:
             self.iter += 1
             yield i
-        
+
         self.epoch += 1
         self._indices = None
-    
+
     def __len__(self):
-        return self.num_samples_per_replica
+        return self.num_samples_per_replica

astrai/factory.py

@@ -0,0 +1,190 @@
+"""Base factory class for extensible component registration."""
+
+from abc import ABC
+from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+
+T = TypeVar("T")
+
+
+class Registry:
+    """Flexible registry for component classes with category and priority support.
+
+    This registry stores component classes with optional metadata (category, priority).
+    It provides methods for registration, retrieval, and listing with filtering.
+    """
+
+    def __init__(self):
+        self._entries = {}  # name -> (component_cls, category, priority)
+
+    def register(
+        self,
+        name: str,
+        component_cls: Type,
+        category: Optional[str] = None,
+        priority: int = 0,
+    ) -> None:
+        """Register a component class with optional category and priority."""
+        if name in self._entries:
+            raise ValueError(f"Component '{name}' is already registered")
+        self._entries[name] = (component_cls, category, priority)
+
+    def get(self, name: str) -> Type:
+        """Get component class by name."""
+        if name not in self._entries:
+            raise KeyError(f"Component '{name}' not found in registry")
+        return self._entries[name][0]
+
+    def get_with_metadata(self, name: str) -> Tuple[Type, Optional[str], int]:
+        """Get component class with its metadata."""
+        entry = self._entries.get(name)
+        if entry is None:
+            raise KeyError(f"Component '{name}' not found in registry")
+        return entry
+
+    def contains(self, name: str) -> bool:
+        """Check if a name is registered."""
+        return name in self._entries
+
+    def list_names(self) -> List[str]:
+        """Return list of registered component names."""
+        return sorted(self._entries.keys())
+
+    def list_by_category(self, category: str) -> List[str]:
+        """Return names of components belonging to a specific category."""
+        return sorted(
+            name for name, (_, cat, _) in self._entries.items() if cat == category
+        )
+
+    def list_by_priority(self, reverse: bool = False) -> List[str]:
+        """Return names sorted by priority (default ascending)."""
+        return sorted(
+            self._entries.keys(),
+            key=lambda name: self._entries[name][2],
+            reverse=reverse,
+        )
+
+    def entries(self) -> Dict[str, Tuple[Type, Optional[str], int]]:
+        """Return raw entries dictionary."""
+        return self._entries.copy()
+
+
+class BaseFactory(ABC, Generic[T]):
+    """Generic factory class for component registration and creation.
+
+    This base class provides a decorator-based registration pattern
+    for creating extensible component factories.
+
+    Example usage:
+        class MyFactory(BaseFactory[MyBaseClass]):
+            pass
+
+        @MyFactory.register("custom")
+        class CustomComponent(MyBaseClass):
+            ...
+
+        component = MyFactory.create("custom", *args, **kwargs)
+    """
+
+    _registry: Registry
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls._registry = Registry()
+
+    @classmethod
+    def register(
+        cls, name: str, category: Optional[str] = None, priority: int = 0
+    ) -> Callable[[Type[T]], Type[T]]:
+        """Decorator to register a component class with optional category and priority.
+
+        Args:
+            name: Registration name for the component
+            category: Optional category for grouping components
+            priority: Priority for ordering (default 0)
+
+        Returns:
+            Decorator function that registers the component class
+
+        Raises:
+            TypeError: If the decorated class doesn't inherit from the base type
+        """
+
+        def decorator(component_cls: Type[T]) -> Type[T]:
+            cls._validate_component(component_cls)
+            cls._registry.register(
+                name, component_cls, category=category, priority=priority
+            )
+            return component_cls
+
+        return decorator
+
+    @classmethod
+    def create(cls, name: str, *args, **kwargs) -> T:
+        """Create a component instance by name.
+
+        Args:
+            name: Registered name of the component
+            *args: Positional arguments passed to component constructor
+            **kwargs: Keyword arguments passed to component constructor
+
+        Returns:
+            Component instance
+
+        Raises:
+            ValueError: If the component name is not registered
+        """
+        if not cls._registry.contains(name):
+            raise ValueError(
+                f"Unknown component: '{name}'. "
+                f"Supported types: {sorted(cls._registry.list_names())}"
+            )
+        component_cls = cls._registry.get(name)
+        return component_cls(*args, **kwargs)
+
+    @classmethod
+    def _validate_component(cls, component_cls: Type[T]) -> None:
+        """Validate that the component class is valid for this factory.
+
+        Override this method in subclasses to add custom validation.
+
+        Args:
+            component_cls: Component class to validate
+
+        Raises:
+            TypeError: If the component class is invalid
+        """
+        pass
+
+    @classmethod
+    def list_registered(cls) -> list:
+        """List all registered component names.
+
+        Returns:
+            List of registered component names
+        """
+        return cls._registry.list_names()
+
+    @classmethod
+    def is_registered(cls, name: str) -> bool:
+        """Check if a component name is registered.
+
+        Args:
+            name: Component name to check
+
+        Returns:
+            True if registered, False otherwise
+        """
+        return cls._registry.contains(name)
+
+    @classmethod
+    def list_by_category(cls, category: str) -> List[str]:
+        """List registered component names in a category."""
+        return cls._registry.list_by_category(category)
+
+    @classmethod
+    def list_by_priority(cls, reverse: bool = False) -> List[str]:
+        """List registered component names sorted by priority."""
+        return cls._registry.list_by_priority(reverse)
+
+
+__all__ = ["Registry", "BaseFactory"]

astrai/inference/__init__.py

@@ -0,0 +1,25 @@
+"""Inference module for continuous batching."""
+
+from astrai.inference.engine import (
+    GenerationRequest,
+    InferenceEngine,
+)
+from astrai.inference.scheduler import (
+    InferenceScheduler,
+    Task,
+    TaskStatus,
+    apply_sampling_strategies,
+)
+
+__all__ = [
+    # Engine
+    "InferenceEngine",
+    # Scheduler
+    "InferenceScheduler",
+    "Task",
+    "TaskStatus",
+    # Request
+    "GenerationRequest",
+    # Sampling
+    "apply_sampling_strategies",
+]

astrai/inference/engine.py

@@ -0,0 +1,292 @@
+"""Unified inference engine."""
+
+import gc
+import logging
+import threading
+from typing import Any, Dict, Generator, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from astrai.inference.scheduler import InferenceScheduler
+from astrai.tokenize import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class GenerationRequest:
+    """Request parameters for text generation."""
+
+    def __init__(
+        self,
+        messages: List[Dict[str, str]],
+        top_k: int = 50,
+        top_p: float = 1.0,
+        temperature: float = 1.0,
+        max_len: int = 1024,
+        stream: bool = False,
+    ):
+        self.messages = messages
+        self.top_k = top_k
+        self.top_p = top_p
+        self.temperature = temperature
+        self.max_len = max_len
+        self.stream = stream
+
+        self._validate()
+
+    def _validate(self):
+        """Validate request parameters."""
+        if not (isinstance(self.top_k, int) and self.top_k >= 0):
+            raise ValueError("top_k must be a non-negative integer")
+        if not (0.0 <= self.top_p <= 1.0):
+            raise ValueError("top_p must be a float between 0.0 and 1.0")
+        if not (isinstance(self.temperature, (int, float)) and self.temperature >= 0):
+            raise ValueError("temperature must be a non-negative number")
+
+
+class _Result:
+    """Unified result holder for streaming/non-streaming modes."""
+
+    def __init__(self, count: int = 1, stream: bool = False):
+        self._stream = stream
+        self._lock = threading.Lock()
+        self._event = threading.Event()
+        self.tokens: List[str] = []
+        self.results: List[str] = [""] * count if count > 1 else [""]
+        self.done_flags: List[bool] = [False] * count
+        self._completed_count = 0
+
+    def append(self, token: str, idx: int = 0):
+        with self._lock:
+            if self._stream:
+                self.tokens.append(token)
+            else:
+                if token == "[DONE]":
+                    if not self.done_flags[idx]:
+                        self.done_flags[idx] = True
+                        self._completed_count += 1
+                        if self._completed_count == len(self.results):
+                            self._event.set()
+                else:
+                    self.results[idx] += token
+        self._event.set()
+
+    def pop_all(self) -> List[str]:
+        with self._lock:
+            tokens = self.tokens.copy()
+            self.tokens.clear()
+            if not tokens:
+                self._event.clear()
+            return tokens
+
+    def wait(self, timeout: float = None) -> bool:
+        return self._event.wait(timeout=timeout)
+
+    def get_results(self) -> List[str]:
+        with self._lock:
+            return self.results.copy()
+
+
+class InferenceEngine:
+    """Unified inference engine for continuous batching."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        tokenizer: AutoTokenizer,
+        max_batch_size: int = 1,
+        max_seq_len: Optional[int] = None,
+        max_prefix_len: int = 512,
+        cache_capacity: int = 1000,
+    ):
+        """
+        Initialize inference engine with separate model and tokenizer.
+
+        Args:
+            model: The language model for inference (nn.Module, e.g., Transformer)
+            tokenizer: The tokenizer for encoding/decoding text
+            config: Model configuration
+            max_batch_size: Maximum batch size for continuous batching
+            max_seq_len: Maximum sequence length (defaults to config.max_len)
+            max_prefix_len: Maximum prefix length for cache (default: 512)
+            cache_capacity: Maximum number of cached prefixes (default: 1000)
+        """
+        self.model = model
+        self.tokenizer = tokenizer
+
+        # Get device and dtype from model parameters
+        try:
+            first_param = next(model.parameters())
+            device = first_param.device
+            dtype = first_param.dtype
+        except StopIteration:
+            # Model has no parameters, use default device/dtype
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            dtype = torch.float32
+
+        self.scheduler = InferenceScheduler(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            max_batch_size=max_batch_size,
+            max_seq_len=max_seq_len,
+            max_prefix_len=max_prefix_len,
+            cache_capacity=cache_capacity,
+            device=device,
+            dtype=dtype,
+        )
+
+        self.kv_cache = self.scheduler.kv_cache
+        self.seq_mask = self.scheduler.seq_mask
+
+        self.scheduler.start()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Handle exceptions on exit."""
+        self.shutdown()
+        return False
+
+    def generate(
+        self,
+        prompt: Union[str, List[str]],
+        stream: bool = False,
+        max_tokens: int = 1024,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+        abort_on_exception: bool = True,
+    ) -> Union[Generator[str, None, None], str, List[str]]:
+        """Unified generation interface.
+
+        Args:
+            abort_on_exception: If True, abort the generation when consumer
+                stops iterating (GeneratorExit/StopIteration). Default: True.
+        """
+        is_batch = isinstance(prompt, list)
+        prompts = prompt if is_batch else [prompt]
+
+        if stream:
+            return self._generate_streaming(
+                prompts,
+                is_batch,
+                max_tokens,
+                temperature,
+                top_p,
+                top_k,
+                abort_on_exception,
+            )
+        else:
+            return self._generate_non_streaming(
+                prompts, is_batch, max_tokens, temperature, top_p, top_k
+            )
+
+    def generate_with_request(
+        self, request: GenerationRequest
+    ) -> Union[Generator[str, None, None], str, List[str]]:
+        """Generate with GenerationRequest object."""
+        # Use tokenizer's chat template with messages
+        prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
+
+        return self.generate(
+            prompt=prompt,
+            stream=request.stream,
+            max_tokens=request.max_len,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            top_k=request.top_k,
+        )
+
+    def _generate_streaming(
+        self,
+        prompts: List[str],
+        is_batch: bool,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+        abort_on_exception: bool = True,
+    ) -> Union[Generator[str, None, None], List[Generator[str, None, None]]]:
+        """Generate with streaming output.
+
+        Args:
+            abort_on_exception: If True, abort the task when generator is
+                stopped early by consumer (GeneratorExit/StopIteration).
+        """
+        if is_batch:
+            raise NotImplementedError("Batch streaming is not implemented yet")
+
+        result = _Result(stream=True)
+
+        task_id = self.scheduler.add_task(
+            prompt=prompts[0],
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            stream_callback=result.append,
+        )
+
+        def gen():
+            try:
+                while True:
+                    tokens = result.pop_all()
+                    for token in tokens:
+                        if token == "[DONE]":
+                            return
+                        yield token
+                    result.wait(timeout=0.05)
+            except Exception:
+                # Consumer stopped iterating - abort the task
+                if abort_on_exception:
+                    self.scheduler.remove_task(task_id)
+                raise
+
+        gen.task_id = task_id
+        return gen()
+
+    def _generate_non_streaming(
+        self,
+        prompts: List[str],
+        is_batch: bool,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+    ) -> Union[str, List[str]]:
+        """Generate without streaming."""
+        result = _Result(count=len(prompts))
+
+        for i, p in enumerate(prompts):
+            # Create closure to capture current index value using factory function
+            def make_callback(idx):
+                def callback(token):
+                    result.append(idx, token)
+
+                return callback
+
+            self.scheduler.add_task(
+                prompt=p,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                stream_callback=make_callback(i),
+            )
+
+        result.wait()
+        results = result.get_results()
+        return results if is_batch else results[0]
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get engine statistics."""
+        return self.scheduler.get_stats()
+
+    def shutdown(self) -> None:
+        """Shutdown the engine and release all resources."""
+        self.scheduler.stop()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()

astrai/inference/scheduler.py

@@ -0,0 +1,637 @@
+"""Inference scheduler for continuous batching."""
+
+import threading
+import time
+import uuid
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from astrai.model.automodel import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+
+class RadixNode:
+    """Radix tree node for prefix cache."""
+
+    def __init__(self):
+        self.children: Dict[int, "RadixNode"] = {}  # token_id -> child node
+        self.hash: Optional[int] = None  # 64-bit hash of the prefix
+        self.slot: int = -1  # KV Cache slot, valid only for leaf nodes
+        self.ref_count: int = 0  # number of tasks referencing this prefix
+        self.last_access: float = 0.0  # timestamp for LRU
+        self.token_sequence: list = []  # full token sequence from root to this node
+
+
+class PrefixCacheManager:
+    """Prefix cache manager using Radix tree with LRU eviction."""
+
+    def __init__(self, max_capacity: int = 1000, base: int = 131, mod: int = 10**9 + 7):
+        self.root = RadixNode()
+        self.base = base
+        self.mod = mod
+        self.max_capacity = max_capacity
+        self.lru: List[Tuple[float, RadixNode]] = []  # (timestamp, node) for LRU
+
+    def insert(self, token_ids: Tuple[int, ...], slot: int) -> None:
+        """Insert a prefix, increase ref_count if already exists, otherwise create new node."""
+        node = self.root
+        path = []
+        h = 0
+        for i, token_id in enumerate(token_ids):
+            if token_id not in node.children:
+                node.children[token_id] = RadixNode()
+            node = node.children[token_id]
+            h = (h * self.base + token_id) % self.mod
+            node.hash = h
+            path.append(token_id)
+            node.token_sequence = list(
+                path
+            )  # store full sequence for exact verification
+
+        # Leaf node: set slot and increase ref_count
+        if node.slot == -1:
+            node.slot = slot
+        node.ref_count += 1
+        node.last_access = time.time()
+        self._update_lru(node)
+        self._evict_if_needed()
+
+    def find_longest_prefix(self, token_ids: List[int]) -> Optional[Tuple[int, int]]:
+        """Find longest matching prefix, return (prefix_len, slot).
+
+        During traversal, compute hash per token and compare with node hash.
+        If hash matches, perform full token sequence verification to avoid
+        hash collision errors.
+        """
+        node = self.root
+        best_len = 0
+        best_slot = -1
+        h = 0
+
+        for i, token_id in enumerate(token_ids):
+            if token_id not in node.children:
+                break
+            node = node.children[token_id]
+            h = (h * self.base + token_id) % self.mod
+            if node.hash == h:  # hash matches
+                # Exact verification: compare full token sequence
+                if node.token_sequence == token_ids[: i + 1]:
+                    best_len = i + 1
+                    best_slot = node.slot
+                    node.last_access = time.time()
+                    self._update_lru(node)
+
+        if best_len > 0:
+            return (best_len, best_slot)
+        return None
+
+    def release(self, token_ids: Tuple[int, ...]) -> None:
+        """Release reference to a prefix, decrease ref_count. If zero, mark as evictable."""
+        node = self.root
+        for token_id in token_ids:
+            if token_id not in node.children:
+                return
+            node = node.children[token_id]
+        if node.ref_count > 0:
+            node.ref_count -= 1
+            if node.ref_count == 0:
+                node.slot = -1  # slot can be reused
+
+    def _update_lru(self, node: RadixNode) -> None:
+        """Update LRU list, move node to most recently used position."""
+        self.lru = [(ts, n) for (ts, n) in self.lru if n is not node]
+        self.lru.append((node.last_access, node))
+
+    def _evict_if_needed(self) -> None:
+        """If cache entries exceed capacity, evict least recently used leaf nodes (ref_count must be 0)."""
+        if len(self.lru) <= self.max_capacity:
+            return
+        # Sort by timestamp
+        self.lru.sort(key=lambda x: x[0])
+        for ts, node in self.lru:
+            if node.ref_count == 0:
+                # Remove leaf node from tree (need to recursively delete empty branches)
+                self._remove_node(node)
+                self.lru.remove((ts, node))
+                if len(self.lru) <= self.max_capacity:
+                    break
+
+    def _remove_node(
+        self,
+        node: RadixNode,
+        parent: Optional[RadixNode] = None,
+        child_key: Optional[int] = None,
+    ) -> None:
+        """Remove node from tree, including empty parent nodes."""
+        # First, recursively remove all children
+        for child_key, child_node in list(node.children.items()):
+            self._remove_node(child_node, node, child_key)
+
+        # Clear the node's leaf properties
+        node.slot = -1
+        node.hash = None
+        node.token_sequence = []
+        node.children.clear()
+
+        # If this node has no children and has a parent, remove the reference from parent
+        if parent is not None and child_key is not None and len(node.children) == 0:
+            if child_key in parent.children:
+                del parent.children[child_key]
+
+
+class TaskStatus:
+    """Task state for continuous batching."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    FINISHED = "finished"
+    ABORTED = "aborted"
+
+
+class Task:
+    """Individual task for continuous batching."""
+
+    def __init__(
+        self,
+        task_id: str,
+        prompt_ids: List[int],
+        max_tokens: int = 1024,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ):
+        self.task_id = task_id
+        self.prompt_ids = prompt_ids
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
+
+        self.status = TaskStatus.PENDING
+        self.output_ids: List[int] = []
+        self.input_tokens: int = 0
+        self.output_tokens: int = 0
+        self.slot: int = -1
+        self.prefix_len: int = 0  # prefix cache matched length
+        self.arrival_time = time.time()
+        self.finish_time: Optional[float] = None
+
+        self.stream_callback = stream_callback
+
+    def is_finished(self, stop_ids: List[int]) -> bool:
+        """Check if task is finished."""
+        return (
+            bool(self.output_ids and self.output_ids[-1] in stop_ids)
+            or self.output_tokens >= self.max_tokens
+        )
+
+
+def apply_sampling_strategies(
+    logits: Tensor,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    filter_value: float = -float("inf"),
+) -> Tensor:
+    """Apply sampling strategies to the logits tensor."""
+    # Clone logits to avoid inplace updates on inference tensor
+    logits = logits.clone()
+
+    if temperature != 1.0:
+        logits = logits / temperature
+
+    if top_k > 0:
+        top_k = min(top_k, logits.size(-1))
+        indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+
+        sorted_indices_to_remove = cumulative_probs > top_p
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
+        indices_to_remove.scatter_(
+            dim=1, index=sorted_indices, src=sorted_indices_to_remove
+        )
+
+        logits[indices_to_remove] = filter_value
+
+    return logits
+
+
+class InferenceScheduler:
+    """Inference scheduler with continuous batching support."""
+
+    def __init__(
+        self,
+        model: AutoModel,
+        tokenizer: AutoTokenizer,
+        max_batch_size: int = 16,
+        max_seq_len: Optional[int] = None,
+        max_prefix_len: int = 512,
+        cache_capacity: int = 1000,
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        config = model.config
+
+        self.model = model
+        self.tokenizer = tokenizer
+        self.max_batch_size = max_batch_size
+        self.max_seq_len = max_seq_len or config.max_len
+        self.max_prefix_len = max_prefix_len
+        self.device = device or next(model.parameters()).device
+        self.dtype = dtype or next(model.parameters()).dtype
+
+        # Initialize prefix cache
+        self.prefix_cache = PrefixCacheManager(max_capacity=cache_capacity)
+
+        num_kv_heads = config.n_kv_heads
+        head_dim = config.dim // config.n_heads
+        n_layers = config.n_layers
+
+        k_cache = torch.empty(
+            (
+                max_batch_size,
+                self.max_seq_len,
+                n_layers,
+                num_kv_heads,
+                head_dim,
+            ),
+            device=self.device,
+            dtype=self.dtype,
+        )
+        v_cache = torch.empty(
+            (
+                max_batch_size,
+                self.max_seq_len,
+                n_layers,
+                num_kv_heads,
+                head_dim,
+            ),
+            device=self.device,
+            dtype=self.dtype,
+        )
+        self.kv_cache = (k_cache, v_cache)
+        self.seq_mask = torch.ones(
+            (max_batch_size, self.max_seq_len), device=self.device, dtype=torch.bool
+        )
+
+        self.waiting_queue: List[Task] = []
+        self.active_tasks: List[Task] = []
+
+        self._running = False
+        self._task_event = threading.Event()
+        self._lock = threading.Lock()
+
+        self._total_tasks = 0
+        self._total_tokens = 0
+
+    def add_task(
+        self,
+        prompt: str,
+        max_tokens: int = 1024,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+        stream_callback: Optional[Callable[[str], None]] = None,
+    ) -> str:
+        """Add a new task to the waiting queue."""
+        task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
+        prompt_ids = self.tokenizer.encode(prompt)
+
+        # Truncate if exceeds max_prefix_len
+        if len(prompt_ids) > self.max_prefix_len:
+            prompt_ids = prompt_ids[: self.max_prefix_len]
+
+        task = Task(
+            task_id=task_id,
+            prompt_ids=prompt_ids,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            stream_callback=stream_callback,
+        )
+
+        # Find longest matching prefix from cache
+        match = self.prefix_cache.find_longest_prefix(prompt_ids)
+        if match:
+            prefix_len, slot = match
+            task.prefix_len = prefix_len
+            task.slot = slot
+        else:
+            task.prefix_len = 0
+            task.slot = -1
+
+        with self._lock:
+            self.waiting_queue.append(task)
+            self._total_tasks += 1
+
+        self._task_event.set()
+        return task_id
+
+    def remove_task(self, task_id: str) -> None:
+        """Remove a task from the scheduler."""
+        with self._lock:
+            self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
+            self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
+
+    def _remove_finished_tasks(self) -> None:
+        """Remove finished tasks from active batch."""
+        finished = []
+        for task in self.active_tasks:
+            if task.is_finished(self.tokenizer.stop_ids):
+                task.status = TaskStatus.FINISHED
+                task.finish_time = time.time()
+                finished.append(task)
+                self._total_tokens += task.output_tokens
+
+        for task in finished:
+            slot = task.slot
+            if slot >= 0 and slot < len(self.active_tasks):
+                self.seq_mask[slot, :] = False
+
+            # Release prefix cache reference
+            if task.prefix_len > 0:
+                self.prefix_cache.release(tuple(task.prompt_ids[: task.prefix_len]))
+
+            task.slot = -1
+
+        self.active_tasks = [
+            t for t in self.active_tasks if t.status != TaskStatus.FINISHED
+        ]
+
+    def _refill_active_batch(self) -> None:
+        """Refill active batch with waiting tasks."""
+        available_slots = self.max_batch_size - len(self.active_tasks)
+        if available_slots <= 0:
+            return
+
+        with self._lock:
+            to_add = [
+                self.waiting_queue.pop(0)
+                for _ in range(min(available_slots, len(self.waiting_queue)))
+            ]
+            for task in to_add:
+                task.slot = self._allocate_slot()
+                task.status = TaskStatus.RUNNING
+                self.active_tasks.append(task)
+
+    def _allocate_slot(self) -> int:
+        """Allocate an available slot for a task."""
+        for i in range(self.max_batch_size):
+            if not any(t.slot == i for t in self.active_tasks):
+                return i
+        return -1
+
+    def _execute_prefill(self, tasks: List[Task]) -> None:
+        """Execute Prefill phase with incremental prefill support."""
+        if not tasks:
+            return
+
+        # Group tasks by prefix cache status
+        fully_cached, partial, full = [], [], []
+        for task in tasks:
+            total_len, prefix_len = len(task.prompt_ids), task.prefix_len
+            if prefix_len == total_len:
+                fully_cached.append(task)
+            elif prefix_len > 0:
+                partial.append(task)
+            else:
+                full.append(task)
+
+        # Handle fully cached tasks
+        for t in fully_cached:
+            t.input_tokens, t.output_tokens = len(t.prompt_ids), 0
+            if t.slot >= 0:
+                self.seq_mask[t.slot, : t.input_tokens] = True
+
+        if full:
+            self._execute_full_prefill(full)
+        if partial:
+            self._execute_partial_prefill(partial)
+
+    def _execute_full_prefill(self, tasks: List[Task]) -> None:
+        """Execute full prefill for tasks without prefix cache."""
+        if not tasks:
+            return
+
+        tasks = sorted(tasks, key=lambda t: t.slot)
+
+        prompt_lens = [len(task.prompt_ids) for task in tasks]
+        max_len = max(prompt_lens)
+
+        input_ids = torch.zeros(
+            len(tasks), max_len, dtype=torch.long, device=self.device
+        )
+        for i, task in enumerate(tasks):
+            if len(task.prompt_ids) > 0:
+                input_ids[i, : len(task.prompt_ids)] = torch.tensor(
+                    task.prompt_ids, device=self.device
+                )
+
+        if self.tokenizer.pad_id is not None:
+            input_mask = torch.ne(input_ids, self.tokenizer.pad_id)
+        else:
+            input_mask = torch.ones(
+                input_ids.shape, dtype=torch.bool, device=self.device
+            )
+
+        with torch.inference_mode():
+            self.model(
+                input_ids,
+                input_mask=input_mask,
+                start_pos=0,
+                persistent_key_values=self.kv_cache,
+            )
+
+        for i, task in enumerate(tasks):
+            task.input_tokens = prompt_lens[i]
+            task.output_tokens = 0
+            # Insert new prefix into cache
+            self.prefix_cache.insert(tuple(task.prompt_ids), task.slot)
+
+        for task in tasks:
+            if task.slot >= 0:
+                self.seq_mask[task.slot, : task.input_tokens] = True
+
+    def _execute_partial_prefill(self, tasks: List[Task]) -> None:
+        """Execute incremental prefill for tasks with partial prefix cache match."""
+        for task in tasks:
+            total_len = len(task.prompt_ids)
+            prefix_len = task.prefix_len
+
+            if prefix_len >= total_len:
+                task.input_tokens = total_len
+                task.output_tokens = 0
+                continue
+
+            # Get new tokens that need prefill
+            new_ids = task.prompt_ids[prefix_len:]
+            new_len = len(new_ids)
+
+            if new_len == 0:
+                task.input_tokens = total_len
+                task.output_tokens = 0
+                continue
+
+            # Build input for incremental prefill
+            input_ids = torch.tensor([new_ids], dtype=torch.long, device=self.device)
+
+            # Input mask should cover from position 0 to prefix_len + new_len
+            # The prefix part uses cached KV, new part needs computation
+            input_mask = torch.ones(
+                (1, prefix_len + new_len), dtype=torch.bool, device=self.device
+            )
+
+            with torch.inference_mode():
+                self.model(
+                    input_ids,
+                    input_mask=input_mask,
+                    start_pos=prefix_len,
+                    persistent_key_values=self.kv_cache,
+                )
+
+            task.input_tokens = total_len
+            task.output_tokens = 0
+
+            # Insert full prefix into cache (ref_count already increased in add_task)
+            self.prefix_cache.insert(tuple(task.prompt_ids), task.slot)
+
+            if task.slot >= 0:
+                self.seq_mask[task.slot, : task.input_tokens] = True
+
+    def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
+        """Execute Decode phase."""
+        if not tasks:
+            return
+
+        tasks = sorted(tasks, key=lambda t: t.slot)
+
+        input_ids = torch.zeros(len(tasks), dtype=torch.long, device=self.device)
+        for i, task in enumerate(tasks):
+            if task.output_ids:
+                input_ids[i] = task.output_ids[-1]
+            else:
+                input_ids[i] = task.prompt_ids[-1]
+
+        input_tensor = input_ids.unsqueeze(1)
+        active_mask = torch.ones((len(tasks), 1), dtype=torch.bool, device=self.device)
+
+        with torch.inference_mode():
+            outputs = self.model(
+                input_tensor,
+                input_mask=active_mask,
+                persistent_key_values=self.kv_cache,
+                start_pos=start_pos,
+            )
+            logits = outputs["logits"][:, -1, :]
+
+        next_token_ids = []
+        for i, task in enumerate(tasks):
+            logit = logits[i : i + 1]
+            logit = apply_sampling_strategies(
+                logit,
+                task.temperature,
+                task.top_k,
+                task.top_p,
+            )
+            probs = torch.softmax(logit, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            next_token_ids.append(next_token.item())
+
+        for task, next_token in zip(tasks, next_token_ids):
+            task.output_ids.append(next_token)
+            task.output_tokens += 1
+
+            pos = task.input_tokens + task.output_tokens
+            if task.slot >= 0 and pos < self.max_seq_len:
+                self.seq_mask[task.slot, pos] = True
+
+            if task.stream_callback:
+                token_str = self.tokenizer.decode([next_token])
+                task.stream_callback(token_str)
+
+        for task in tasks:
+            if task.output_tokens >= task.max_tokens or (
+                task.output_ids and task.output_ids[-1] in self.tokenizer.stop_ids
+            ):
+                if task.stream_callback:
+                    task.stream_callback("[DONE]")
+
+    def _run_generation_loop(self) -> None:
+        """Main generation loop."""
+        while self._running:
+            self._remove_finished_tasks()
+            self._refill_active_batch()
+
+            if not self.active_tasks:
+                self._task_event.wait(timeout=0.01)
+                self._task_event.clear()
+                continue
+
+            new_tasks = [t for t in self.active_tasks if t.output_tokens == 0]
+            decode_tasks = [t for t in self.active_tasks if t.output_tokens > 0]
+
+            if decode_tasks:
+                start_pos = max(t.input_tokens + t.output_tokens for t in decode_tasks)
+            else:
+                start_pos = 0
+
+            if new_tasks:
+                self._execute_prefill(new_tasks)
+                decode_tasks = new_tasks
+                start_pos = max(t.input_tokens for t in decode_tasks)
+
+            if decode_tasks:
+                self._execute_decode(decode_tasks, start_pos)
+
+            if not self.active_tasks and not self.waiting_queue:
+                self._task_event.wait(timeout=0.05)
+                self._task_event.clear()
+
+    def start(self) -> None:
+        """Start the generation loop."""
+        if not self._running:
+            self._running = True
+            self._loop_thread = threading.Thread(target=self._run_generation_loop)
+            self._loop_thread.daemon = True
+            self._loop_thread.start()
+
+    def stop(self) -> None:
+        """Stop the generation loop."""
+        self._running = False
+        if hasattr(self, "_loop_thread"):
+            self._loop_thread.join(timeout=1.0)
+
+        # Clear KV cache to free GPU memory
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            if k_cache is not None:
+                k_cache.detach()
+            if v_cache is not None:
+                v_cache.detach()
+
+        # Clear seq mask
+        self.seq_mask.detach()
+
+        # Clear task lists
+        self.waiting_queue.clear()
+        self.active_tasks.clear()
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get scheduler statistics."""
+        return {
+            "total_tasks": self._total_tasks,
+            "total_tokens": self._total_tokens,
+            "active_tasks": len(self.active_tasks),
+            "waiting_queue": len(self.waiting_queue),
+        }

astrai/inference/server.py

@@ -0,0 +1,321 @@
+"""
+Inference Server with Continuous Batching Support
+
+FastAPI server for inference with continuous batching.
+Provides OpenAI-compatible chat completion endpoints.
+"""
+
+import json
+import logging
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+
+from astrai.inference.engine import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+# Global model parameter and engine (loaded once)
+_engine: Optional[InferenceEngine] = None
+_model_param: Optional[Any] = None
+_project_root = Path(__file__).parent.parent.parent
+
+# Server configuration (set before running server)
+_server_config: Dict[str, Any] = {
+    "device": "cuda",
+    "dtype": torch.bfloat16,
+    "param_path": None,
+    "max_batch_size": 16,
+}
+
+
+def configure_server(
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    param_path: Optional[Path] = None,
+    max_batch_size: int = 16,
+):
+    """Configure server settings before starting.
+
+    Args:
+        device: Device to load model on (e.g., "cuda", "cpu", "cuda:0")
+        dtype: Data type for model weights (e.g., torch.bfloat16, torch.float16)
+        param_path: Path to model parameters directory
+        max_batch_size: Maximum batch size for continuous batching
+    """
+    _server_config["device"] = device
+    _server_config["dtype"] = dtype
+    _server_config["param_path"] = param_path
+    _server_config["max_batch_size"] = max_batch_size
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan context manager for startup and shutdown events."""
+    global _model_param, _engine
+    # Startup: Load model with configured settings
+    try:
+        load_model(
+            param_path=_server_config["param_path"],
+            device=_server_config["device"],
+            dtype=_server_config["dtype"],
+            max_batch_size=_server_config["max_batch_size"],
+        )
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+    yield
+    # Shutdown: Cleanup engine
+    if _engine:
+        _engine.shutdown()
+        logger.info("Inference engine shutdown complete")
+
+
+app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
+
+
+def load_model(
+    param_path: Optional[Path] = None,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    max_batch_size: int = 16,
+):
+    """Load model parameters and initialize inference engine."""
+    global _model_param, _engine
+    if param_path is None:
+        param_path = _project_root / "params"
+    if not param_path.exists():
+        raise FileNotFoundError(f"Parameter directory not found: {param_path}")
+
+    # Load tokenizer separately
+    tokenizer = AutoTokenizer.from_pretrained(param_path)
+    _model_param = AutoModel.from_pretrained(param_path)
+    _model_param.to(device=device, dtype=dtype)
+    logger.info(f"Model loaded on {device} with dtype {dtype}")
+
+    # Initialize inference engine with separate model and tokenizer
+    _engine = InferenceEngine(
+        model=_model_param,
+        tokenizer=tokenizer,
+        max_batch_size=max_batch_size,
+    )
+    logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
+
+
+# Pydantic models for API request/response
+class ChatMessage(BaseModel):
+    role: str  # "user", "assistant", "system"
+    content: str
+
+
+class ChatCompletionRequest(BaseModel):
+    messages: List[ChatMessage]
+    temperature: float = Field(0.8, ge=0.0, le=2.0)
+    top_p: float = Field(0.95, ge=0.0, le=1.0)
+    top_k: int = Field(50, ge=0)
+    max_tokens: int = Field(2048, ge=1)
+    stream: bool = False
+    system_prompt: Optional[str] = None
+
+
+class CompletionResponse(BaseModel):
+    id: str = "chatcmpl-default"
+    object: str = "chat.completion"
+    created: int = 0
+    model: str = "astrai"
+    choices: List[Dict[str, Any]]
+
+
+@app.get("/health")
+async def health():
+    return {
+        "status": "ok",
+        "model_loaded": _model_param is not None,
+        "engine_ready": _engine is not None,
+    }
+
+
+@app.get("/stats")
+async def get_stats():
+    """Get inference engine statistics."""
+    if _engine is None:
+        raise HTTPException(status_code=503, detail="Engine not initialized")
+    return _engine.get_stats()
+
+
+@app.post("/v1/chat/completions", response_model=CompletionResponse)
+async def chat_completion(request: ChatCompletionRequest):
+    """OpenAI-compatible chat completion endpoint.
+
+    Supports both streaming and non-streaming modes with continuous batching.
+    """
+    if _engine is None:
+        raise HTTPException(status_code=503, detail="Engine not initialized")
+
+    # Convert messages to prompt using engine's tokenizer
+    # Extract system prompt if present, then apply chat template
+    # Apply chat template directly with messages
+    prompt = _engine.tokenizer.apply_chat_template(
+        [{"role": m.role, "content": m.content} for m in request.messages],
+        tokenize=False,
+    )
+
+    if request.stream:
+        # Streaming response (use synchronous generator)
+        generator = _engine.generate(
+            prompt=prompt,
+            stream=True,
+            max_tokens=request.max_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            top_k=request.top_k,
+        )
+
+        def generate_stream():
+            for token in generator:
+                if token == "[DONE]":
+                    break
+                yield f"data: {json.dumps({'choices': [{'delta': {'content': token}}]})}\n\n"
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(
+            generate_stream(),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
+        )
+    else:
+        # Non-streaming response
+        result = _engine.generate(
+            prompt=prompt,
+            stream=False,
+            max_tokens=request.max_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            top_k=request.top_k,
+        )
+
+        # Build OpenAI-style response
+        import time
+
+        resp = CompletionResponse(
+            id=f"chatcmpl-{int(time.time())}",
+            created=int(time.time()),
+            choices=[
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": result},
+                    "finish_reason": "stop",
+                }
+            ],
+        )
+        return resp
+
+
+@app.post("/generate")
+async def generate(
+    query: str,
+    history: Optional[List[List[str]]] = None,
+    temperature: float = 0.8,
+    top_p: float = 0.95,
+    top_k: int = 50,
+    max_len: int = 2048,
+    stream: bool = False,
+):
+    """Simple generation endpoint.
+
+    Args:
+        query: Input query string
+        history: Conversation history as list of [user, assistant] pairs
+        temperature: Sampling temperature
+        top_p: Top-p sampling parameter
+        top_k: Top-k sampling parameter
+        max_len: Maximum tokens to generate
+        stream: Enable streaming output
+
+    Returns:
+        dict: Generation result with response field
+    """
+    if _engine is None:
+        raise HTTPException(status_code=503, detail="Engine not initialized")
+
+    # Build messages for chat template
+    messages = []
+    if history:
+        # Convert history format: List[List[str]] -> List[Dict]
+        for h in history:
+            if len(h) >= 2:
+                messages.append({"role": "user", "content": h[0]})
+                messages.append({"role": "assistant", "content": h[1]})
+    messages.append({"role": "user", "content": query})
+
+    # Use tokenizer's chat template
+    prompt = _engine.tokenizer.apply_chat_template(messages, tokenize=False)
+
+    if stream:
+        # Synchronous streaming
+        result = _engine.generate(
+            prompt=prompt,
+            stream=True,
+            max_tokens=max_len,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+        )
+
+        def stream_generator():
+            for token in result:
+                yield token + "\n"
+
+        return StreamingResponse(stream_generator(), media_type="text/plain")
+    else:
+        result = _engine.generate(
+            prompt=prompt,
+            stream=False,
+            max_tokens=max_len,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+        )
+        return {"response": result}
+
+
+def run_server(
+    host: str = "0.0.0.0",
+    port: int = 8000,
+    reload: bool = False,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    param_path: Optional[Path] = None,
+    max_batch_size: int = 16,
+):
+    """Run the FastAPI server with uvicorn.
+
+    Args:
+        host: Server host address
+        port: Server port number
+        reload: Enable auto-reload for development
+        device: Device to load model on (e.g., "cuda", "cpu", "cuda:0")
+        dtype: Data type for model weights (e.g., torch.bfloat16, torch.float16)
+        param_path: Path to model parameters directory
+        max_batch_size: Maximum batch size for continuous batching
+    """
+    configure_server(
+        device=device,
+        dtype=dtype,
+        param_path=param_path,
+        max_batch_size=max_batch_size,
+    )
+    uvicorn.run(
+        "astrai.inference.server:app",
+        host=host,
+        port=port,
+        reload=reload,
+    )

astrai/model/__init__.py

@@ -0,0 +1,21 @@
+from astrai.model.automodel import AutoModel
+from astrai.model.module import (
+    GQA,
+    MLP,
+    DecoderBlock,
+    Linear,
+    RMSNorm,
+)
+from astrai.model.transformer import Transformer
+
+__all__ = [
+    # Modules
+    "Linear",
+    "RMSNorm",
+    "MLP",
+    "GQA",
+    "DecoderBlock",
+    # Models
+    "Transformer",
+    "AutoModel",
+]

astrai/model/automodel.py

@@ -0,0 +1,134 @@
+"""
+AutoModel base class for model loading and saving.
+"""
+
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Dict, Self, Type, Union
+
+import safetensors.torch as st
+import torch.nn as nn
+
+from astrai.config import ModelConfig
+
+
+@contextmanager
+def _disable_random_init(enable: bool = True):
+    init_functions = [
+        "xavier_normal_",
+        "xavier_uniform_",
+        "kaiming_normal_",
+        "kaiming_uniform_",
+        "zeros_",
+        "ones_",
+        "constant_",
+        "normal_",
+        "uniform_",
+    ]
+    original_funcs = {}
+    for name in init_functions:
+        if enable and hasattr(nn.init, name):
+            original_funcs[name] = getattr(nn.init, name)
+            setattr(nn.init, name, lambda *args, **kwargs: None)
+    try:
+        yield
+    finally:
+        if enable:
+            for name, orig_func in original_funcs.items():
+                setattr(nn.init, name, orig_func)
+
+
+class AutoModel(nn.Module):
+    """
+    Autoregressive language model base class.
+    Provides model loading/saving and generation capabilities.
+    """
+
+    # Model registry - stored as class attribute
+    _registry: Dict[str, Type["AutoModel"]] = {}
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+
+    @classmethod
+    def register(cls, model_type: str):
+        """
+        Class method decorator to register model type.
+
+        Usage:
+            @AutoModel.register('transformer')
+            class Transformer(AutoModel):
+                ...
+        """
+
+        def decorator(sub_cls: Type["AutoModel"]) -> Type["AutoModel"]:
+            cls._registry[model_type.lower()] = sub_cls
+            return sub_cls
+
+        return decorator
+
+    @classmethod
+    def get_model_class(cls, model_type: str) -> Type["AutoModel"]:
+        """Get model class by model_type string."""
+        model_type = model_type.lower()
+        if model_type not in cls._registry:
+            available = list(cls._registry.keys())
+            raise ValueError(
+                f"Unknown model_type: {model_type}. Available: {available}"
+            )
+        return cls._registry[model_type]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        path: Union[str, Path],
+        disable_random_init: bool = True,
+    ) -> 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))
+        else:
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        # If called from base class, use model_type to determine actual model class
+        if cls is AutoModel:
+            model_type = config.model_type or "transformer"
+            actual_cls = cls.get_model_class(model_type)
+        else:
+            raise ValueError(
+                f"Cannot call from_pretrained() on subclass {cls.__name__}"
+            )
+
+        with _disable_random_init(enable=disable_random_init):
+            model = actual_cls(config)
+
+        # Load weights
+        weights_path = model_path / "model.safetensors"
+        if weights_path.exists():
+            state_dict = st.load_file(str(weights_path))
+            model.load_state_dict(state_dict, strict=False)
+
+        return model
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, Path],
+    ) -> None:
+        save_path = Path(save_directory)
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        # Save config
+        self.config.save(str(save_path / "config.json"))
+
+        # Save weights
+        st.save_file(self.state_dict(), str(save_path / "model.safetensors"))
+
+    def to(self, *args, **kwargs) -> Self:
+        """Move model to device/dtype."""
+        return super().to(*args, **kwargs)

astrai/model/module.py

@@ -1,13 +1,13 @@
+from typing import Optional, Tuple
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 from torch import Tensor
-from typing import Optional, Tuple
 
 
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
-    """ 
+    """
     Repeat k times along the dimension for attention heads.
     Args:
         x (Tensor): The input tensor.
@@ -15,7 +15,7 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
     Returns:
         Tensor: The repeated tensor.
     """
-    
+
     bs, slen, n_heads, head_dim = x.shape
     if n_rep == 1:
         return x
@@ -25,27 +25,31 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
         .reshape(bs, slen, n_heads * n_rep, head_dim)
     )
 
+
 def get_rotary_emb(
-        dim: int, 
+    dim: int,
-        max_len: int, 
+    max_len: int,
-        base: float = 10000, 
+    base: float = 10000,
-    ) -> Tuple[Tensor, Tensor]:
+    device: Optional[torch.device] = None,
-    """ 
+) -> Tuple[Tensor, Tensor]:
+    """
     Get the rotary embedding for the given dimension and maximum length.
     Args:
         dim (int): The dimension of the input.
         max_len (int): The maximum length of the input.
         base (float, optional): The base for the frequency. Defaults to 10000.
+        device (optional): The device to create tensors on. Defaults to None.
     Returns:
         Tensor: The rotary embedding tensor.
     """
 
-    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64) / dim)
+    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
-    t = torch.arange(0, max_len, dtype=torch.float64)
+    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
     freqs = torch.outer(t, theta)
 
     return torch.cos(freqs).float(), torch.sin(freqs).float()
 
+
 def apply_rotary_emb(x: torch.Tensor, rotary_emb: Tuple[Tensor, Tensor]) -> Tensor:
     """
     Apply rotary embedding to the input tensor using cos/sin form.
@@ -55,49 +59,49 @@ def apply_rotary_emb(x: torch.Tensor, rotary_emb: Tuple[Tensor, Tensor]) -> Tens
     Returns:
         Tensor: The output tensor (rotated, same shape as input).
     """
-    
+
     dtype = x.dtype
     cos, sin = rotary_emb
-    
+
     cos = cos.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, dim//2]
     sin = sin.unsqueeze(0).unsqueeze(2)  # [1, seq_len, 1, dim//2]
-    
+
     x_real = x[..., 0::2]  # [batch, seq_len, dim//2]
     x_imag = x[..., 1::2]  # [batch, seq_len, dim//2]
-    
+
     x_real_rot = x_real * cos - x_imag * sin
     x_imag_rot = x_real * sin + x_imag * cos
-    
+
     x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1)  # [batch, seq_len, dim//2, 2]
-    x_out = x_out.view(*x_out.shape[:-2], -1)              # [batch, seq_len, dim]
+    x_out = x_out.view(*x_out.shape[:-2], -1)  # [batch, seq_len, dim]
-    
+
     return x_out.to(dtype)
 
 
 class RotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, max_len: int, base: int=10000):
+    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.max_len_cached = None
-        self._set_rotary_buffer(self.max_len)
+        self._set_rotary_buffer(self.max_len, None)
-    
+
-    def _set_rotary_buffer(self, max_len: int):
+    def _set_rotary_buffer(self, max_len: int, device: Optional[torch.device] = None):
-        cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base)
+        cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base, device)
         self.register_buffer("cos_cached", cos_cached, persistent=False)
         self.register_buffer("sin_cached", sin_cached, persistent=False)
         self.max_len_cached = max_len
-    
+
-    def forward(self, x: Tensor, start_pos: int=0) -> Tuple[Tensor, Tensor]:
+    def forward(self, x: Tensor, start_pos: int = 0) -> Tuple[Tensor, Tensor]:
         seq_len = x.size(1)
-        
+
         if self.max_len_cached < seq_len + start_pos:
-            self._set_rotary_buffer(seq_len)
+            self._set_rotary_buffer(self.max_len_cached * 2, x.device)
-        
+
         cos = self.cos_cached[start_pos : start_pos + seq_len]
         sin = self.sin_cached[start_pos : start_pos + seq_len]
-        
+
         return (cos, sin)
 
 
@@ -115,52 +119,41 @@ class RMSNorm(nn.Module):
     def __init__(self, dim, norm_eps):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(dim))
-        self.normalized_shape = (dim, )
+        self.normalized_shape = (dim,)
         self.norm_eps = norm_eps
 
     def forward(self, x: Tensor) -> Tensor:
-        rms =  F.rms_norm(x.float(), self.normalized_shape, self.weight, self.norm_eps)
+        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)
-        return rms.to(x.dtype)
+
-    
+
-    
 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 Attention(nn.Module):
-    
-    def forward(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None, is_causal: bool= False):
-        # (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)
-        # (bsz, n_heads, seq_len, head_dim) - > (bsz, seq_len, n_heads*head_dim)
-        sdqa_out = F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal).permute(0, 2, 1, 3).contiguous().flatten(2)
-        
-        return sdqa_out
-
 
 class GQA(nn.Module):
     def __init__(
-        self, 
+        self,
-        dim: int, 
+        dim: int,
-        n_heads: int, 
+        n_heads: int,
-        n_kv_heads: int, 
+        n_kv_heads: int,
         use_qk_norm: bool,
         norm_eps: float,
         use_gated_attention: bool,
-        layer_id: int
+        layer_id: int,
     ):
         super().__init__()
         assert dim % n_heads == 0
         assert n_heads % n_kv_heads == 0
-        
+
         self.head_dim = dim // n_heads
         self.layer_id = layer_id
         self.dim = dim
@@ -169,18 +162,16 @@ class GQA(nn.Module):
         self.n_rep = n_heads // n_kv_heads
         self.use_qk_norm = use_qk_norm
         self.use_gated_attention = use_gated_attention
-        
-        self.attention = Attention()
 
         self.q_proj = Linear(dim, n_heads * self.head_dim)
         self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
         self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
         self.o_proj = Linear(dim, dim)
-        
+
         if self.use_qk_norm:
             self.q_norm = RMSNorm(self.head_dim, norm_eps)
             self.k_norm = RMSNorm(self.head_dim, norm_eps)
-        
+
         if self.use_gated_attention:
             self.gate = Linear(dim, dim)
 
@@ -188,62 +179,176 @@ class GQA(nn.Module):
         batch_size, seq_len, _ = x.shape
         x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
         return x
-    
+
     def forward(
         self,
-        x: Tensor, 
+        x: Tensor,
-        rotary_emb: Tuple[Tensor, Tensor], 
+        rotary_emb: Tuple[Tensor, Tensor],
         mask: Tensor = None,
         kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
-        start_pos: int = 0
+        start_pos: int = 0,
     ) -> Tensor:
         bsz, seq_len, _ = x.size()
+        is_causal = mask is None
+
         # x(bsz, seq_len, n_heads * head_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)
         q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
-        
+
         if self.use_qk_norm:
             q, k = self.q_norm(q), self.k_norm(k)
-        
+
         if kv_cache is not None:
             k_cache, v_cache = kv_cache
-            
+
             # copy to cache
-            k_cache[:bsz, start_pos:start_pos + seq_len, self.layer_id] = k
+            k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
-            v_cache[:bsz, start_pos:start_pos + seq_len, self.layer_id] = v
+            v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
-            
+
             # get cache
-            k = k_cache[:bsz, :start_pos + seq_len, self.layer_id]
+            k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
-            v = v_cache[:bsz, :start_pos + seq_len, self.layer_id]
+            v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
-        
+
         k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
-        sdqa_out = self.attention(q, k, v, mask, is_causal=(mask == None))
+
-        
+        # (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)
+        # (bsz, n_heads, seq_len, head_dim) - > (bsz, seq_len, n_heads*head_dim)
+        sdqa_out = (
+            F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+            .flatten(2)
+        )
+
         if self.use_gated_attention:
             sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
-        
+
         out = self.o_proj(sdqa_out)
 
         return out
 
 
+class MLA(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        norm_eps: float,
+        use_gated_attention: bool,
+        layer_id: int,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        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_gated_attention = use_gated_attention
+
+        self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
+        self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
+        self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
+
+        # 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),
+        )
+
+        self.o_proj = Linear(dim, dim, bias=False)
+
+        if use_gated_attention:
+            self.gate = Linear(dim, dim, bias=False)
+
+    def forward(
+        self,
+        x: Tensor,
+        rotary_emb: Tuple[Tensor, Tensor],
+        mask: Tensor = None,
+        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        start_pos: int = 0,
+    ) -> Tensor:
+        bsz, seq_len, _ = x.size()
+        is_causal = mask is None
+
+        q = self.q_proj(x)
+        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
+
+        kv_compressed = self.kv_a_proj(x)
+        kv_compressed = self.kv_norm(kv_compressed)
+
+        kv = self.kv_b_proj(kv_compressed)
+        kv = kv.view(bsz, seq_len, self.n_kv_heads, -1)
+
+        k_nope, k_rope, v = torch.split(
+            kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
+        )
+
+        q_nope, q_rope = (
+            q[..., : self.qk_nope_head_dim],
+            q[..., self.qk_rope_head_dim :],
+        )
+        q_rope = apply_rotary_emb(q_rope, rotary_emb)
+        k_rope = apply_rotary_emb(k_rope, rotary_emb)
+
+        q = torch.cat([q_nope, q_rope], dim=-1)
+        k = torch.cat([k_nope, k_rope], dim=-1)
+
+        if kv_cache is not None:
+            k_cache, v_cache = kv_cache
+            k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
+            v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
+            k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
+            v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
+
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+
+        attn_out = F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
+        attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
+
+        if self.use_gated_attention:
+            attn_out = attn_out * F.sigmoid(self.gate(x))
+
+        out = self.o_proj(attn_out)
+
+        return out
+
+
 class DecoderBlock(nn.Module):
     def __init__(
-        self, 
+        self,
-        dim: int, 
+        dim: int,
-        n_heads: int, 
+        n_heads: int,
-        dim_ffn: int, 
+        dim_ffn: int,
-        n_kv_heads: int, 
+        n_kv_heads: int,
-        norm_eps: int, 
+        norm_eps: int,
-        use_qk_norm: bool, 
+        use_qk_norm: bool,
-        use_gated_attention: bool, 
+        use_gated_attention: bool,
-        layer_id: int
+        layer_id: int,
     ):
         super().__init__()
-        self.attention = GQA(dim, n_heads, n_kv_heads, 
+        self.attention = GQA(
-                             use_qk_norm, norm_eps, use_gated_attention, layer_id)
+            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)
@@ -251,24 +356,20 @@ class DecoderBlock(nn.Module):
     def forward(
         self,
         x: Tensor,
-        rotary_emb: Tuple[Tensor, Tensor], 
+        rotary_emb: Tuple[Tensor, Tensor],
         attention_mask: Optional[Tensor] = None,
         kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
-        start_pos: int = 0
+        start_pos: int = 0,
     ) -> Tensor:
         # attention
         attn_output = self.attention(
-            self.input_norm(x), 
+            self.input_norm(x), rotary_emb, attention_mask, kv_cache, start_pos
-            rotary_emb, 
-            attention_mask, 
-            kv_cache, 
-            start_pos
         )
         x = attn_output + x
-        
+
         # feed forward
         x = self.mlp(self.post_attention_norm(x)) + x
-        
+
         return x
 
 
@@ -276,6 +377,6 @@ 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)
+        return F.embedding(x, self.weight)

astrai/model/transformer.py

@@ -1,18 +1,26 @@
+from typing import Any, Mapping, Optional, Tuple
+
 import torch
 import torch.nn as nn
-
 from torch import Tensor
-from typing import Any, Mapping, Optional, Tuple
+
-from khaosz.config.model_config import ModelConfig
+from astrai.config.model_config import ModelConfig
-from khaosz.model.module import Embedding, DecoderBlock, Linear, RMSNorm, RotaryEmbedding
+from astrai.model.automodel import AutoModel
+from astrai.model.module import (
+    DecoderBlock,
+    Embedding,
+    Linear,
+    RMSNorm,
+    RotaryEmbedding,
+)
 
 
 def process_attention_mask(
-        seq_mask: Tensor, 
+    seq_mask: Tensor,
-        input_tensor: Tensor,
+    input_tensor: Tensor,
-        start_pos: int = 0,
+    start_pos: int = 0,
-        is_causal: bool = False,
+    is_causal: bool = False,
-    ) -> Tensor:
+) -> Tensor:
     """
     Create attention mask for GQA
     Args:
@@ -26,109 +34,132 @@ def process_attention_mask(
     device = input_tensor.device
     dtype = input_tensor.dtype
     seq_len = input_tensor.size(1)
-    
+
     if seq_mask is None:
         if start_pos != 0:
             # for single prompt chat
             seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
         else:
             return None
-    
+
     if seq_mask.dim() > 2:
         # shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos)
         # if ndim > 2, it's 4D tensor
         return seq_mask
-    
+
     batch_size = seq_mask.size(0)
-    seq_mask = seq_mask[:, :start_pos + seq_len].to(device=device, dtype=torch.bool)
+    seq_mask = seq_mask[:, : start_pos + seq_len].to(device=device, dtype=torch.bool)
     # (bsz, start_pos + seq_len)
-    expanded_mask = seq_mask.unsqueeze(1).expand(batch_size, seq_len, start_pos + seq_len)
+    expanded_mask = seq_mask.unsqueeze(1).expand(
+        batch_size, seq_len, start_pos + seq_len
+    )
     # (bsz, seq_len, start_pos + seq_len)
-    
+
     if is_causal:
         expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)
-    
+
     attention_mask = torch.zeros_like(expanded_mask, dtype=dtype, device=device)
-    attention_mask = attention_mask.masked_fill_(~expanded_mask, -torch.finfo(dtype).max / 2).unsqueeze(1)
+    attention_mask = attention_mask.masked_fill_(
+        ~expanded_mask, -torch.finfo(dtype).max / 2
+    ).unsqueeze(1)
     # (bsz, 1, seq_len, seq_len + start_pos)
-    
+
     return attention_mask
 
 
-class Transformer(nn.Module):
+@AutoModel.register("transformer")
+class Transformer(AutoModel):
+    """
+    Transformer language model.
+    """
+
     def __init__(self, config: ModelConfig):
-        super().__init__()
+        super().__init__(config)
         self.config = config
-        self.rotary_embeding = RotaryEmbedding(config.dim // config.n_heads, config.max_len)
+        self.rotary_embedding = RotaryEmbedding(
+            config.dim // config.n_heads, config.max_len
+        )
         self.embed_tokens = Embedding(config.vocab_size, config.dim)
-        
+
-        self.layers = nn.ModuleList([
+        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)
+                DecoderBlock(
-            for layer_id in range(config.n_layers)
+                    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.lm_head = Linear(config.dim, config.vocab_size)
-        
+
-        if self.config.tie_weight == True:
+        if self.config.tie_weight:
             self.lm_head.weight = self.embed_tokens.weight
 
-        self._init_parameters()
+        self._init_weights()
-    
-    def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
-        lm_head_key = 'lm_head.weight'
-        embed_key = 'embed_tokens.weight'
 
-        if self.config.tie_weight == True:
+    def _init_weights(self):
-            # same tensor
-            state_dict[lm_head_key] = state_dict[embed_key]
-        else:
-            if lm_head_key not in state_dict and embed_key in state_dict:
-                # use clone to avoid sharing the same tensor
-                state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
-        
-        return super().load_state_dict(state_dict, strict, assign)
-    
-    def state_dict(self, destination=None, prefix='', keep_vars=False):
-        state_dict = super().state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
-        
-        if self.config.tie_weight == True:
-            lm_head_key = prefix + 'lm_head.weight'
-            if lm_head_key in state_dict:
-                del state_dict[lm_head_key]
-        
-        return state_dict
-    
-    def _init_parameters(self):
         for param in self.parameters():
             if param.dim() > 1:
                 nn.init.normal_(param, mean=0.0, std=0.006)
-    
+
+    def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
+        lm_head_key = "lm_head.weight"
+        embed_key = "embed_tokens.weight"
+
+        # Make a copy to avoid modifying the original state_dict
+        state_dict = dict(state_dict)
+
+        if self.config.tie_weight:
+            # same tensor
+            if embed_key in state_dict:
+                state_dict[lm_head_key] = state_dict[embed_key]
+        else:
+            # If lm_head.weight exists in checkpoint, use it directly
+            # If not, copy from embed_tokens.weight
+            if lm_head_key not in state_dict and embed_key in state_dict:
+                # use clone to avoid sharing the same tensor
+                state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
+
+        return super().load_state_dict(state_dict, strict, assign)
+
+    def state_dict(self, destination=None, prefix="", keep_vars=False):
+        state_dict = super().state_dict(
+            destination=destination, prefix=prefix, keep_vars=keep_vars
+        )
+
+        if self.config.tie_weight:
+            lm_head_key = prefix + "lm_head.weight"
+            if lm_head_key in state_dict:
+                del state_dict[lm_head_key]
+
+        return state_dict
+
     def forward(
-        self, 
+        self,
-        input_ids: Tensor, 
+        input_ids: Tensor,
-        input_mask: Optional[Tensor]=None,
+        input_mask: Optional[Tensor] = None,
-        persistent_key_values: Optional[Tuple[Tensor, Tensor]]=None,
+        persistent_key_values: Optional[Tuple[Tensor, Tensor]] = None,
-        start_pos: int = 0
+        start_pos: int = 0,
     ) -> Tensor:
         assert input_ids.ndim == 2
-        
+
         x = self.embed_tokens(input_ids)
-        rotary_emb = self.rotary_embeding(x, start_pos)
+        rotary_emb = self.rotary_embedding(x, start_pos)
-        
+
-        attn_mask = process_attention_mask(
+        attn_mask = process_attention_mask(input_mask, x, start_pos, is_causal=True)
-            input_mask, x, start_pos, is_causal=True
+
-        )
-        
         for layer in self.layers:
             x = layer(x, rotary_emb, attn_mask, persistent_key_values, start_pos)
-        
+
-        hidden_states = self.norm(x)        
+        hidden_states = self.norm(x)
         logits = self.lm_head(hidden_states)
-        
+
-        return {
+        return {"logits": logits, "hidden_states": hidden_states}
-            "logits": logits,
-            "hidden_states": hidden_states
-        }
-        

astrai/parallel/__init__.py

@@ -0,0 +1,20 @@
+from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
+from astrai.parallel.setup import (
+    get_current_device,
+    get_rank,
+    get_world_size,
+    only_on_rank,
+    setup_parallel,
+    spawn_parallel_fn,
+)
+
+__all__ = [
+    "get_world_size",
+    "get_rank",
+    "get_current_device",
+    "only_on_rank",
+    "setup_parallel",
+    "spawn_parallel_fn",
+    "RowParallelLinear",
+    "ColumnParallelLinear",
+]

astrai/parallel/module.py

@@ -1,10 +1,10 @@
+from typing import Dict
+
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
-import torch.distributed as dist
-
 from torch import Tensor
-from typing import Dict
 
 
 class ParallelModel(nn.Module):
@@ -17,91 +17,99 @@ class ParallelModel(nn.Module):
 
 class RowParallelLinear(ParallelModel):
     def __init__(
-        self, 
+        self,
         process_group: dist.ProcessGroup,
-        in_features: int, 
+        in_features: int,
-        out_features: int, 
+        out_features: int,
         bias: bool = True,
-        reduce_results: bool = True
+        reduce_results: bool = True,
     ):
         super().__init__(process_group)
-        
+
         self.in_features = in_features
         self.out_features = out_features
         self.in_features_per_rank = in_features // self.world_size
         self.reduce_results = reduce_results
-        
+
         if in_features % self.world_size != 0:
-            raise ValueError(f"in_features must be divisible by world_size. Got {in_features} and {self.world_size}")
+            raise ValueError(
-        
+                f"in_features must be divisible by world_size. Got {in_features} and {self.world_size}"
+            )
+
         self.weight = nn.Parameter(torch.empty(out_features, self.in_features_per_rank))
         self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None
-        
+
     def forward(self, input: Tensor) -> Tensor:
         output = F.linear(input, self.weight)
-        
+
         if self.reduce_results:
             dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.process_group)
-        
+
         if self.bias is not None:
             output += self.bias
-        
+
         return output
-    
+
-    def load_state_dict(self, state_dict: Dict[str, Tensor]):        
+    def load_state_dict(self, state_dict: Dict[str, Tensor]):
-        full_weight = state_dict.get('weight')
+        full_weight = state_dict.get("weight")
-        full_bias = state_dict.get('bias')
+        full_bias = state_dict.get("bias")
-        
+
         start_idx = self.rank * self.in_features_per_rank
         end_idx = start_idx + self.in_features_per_rank
         weight_slice = full_weight[:, start_idx:end_idx]
         self.weight.data.copy_(weight_slice)
-        
+
         if self.bias is not None:
             self.bias.data.copy_(full_bias)
 
 
 class ColumnParallelLinear(ParallelModel):
     def __init__(
-        self, 
+        self,
         process_group: dist.ProcessGroup,
-        in_features: int, 
+        in_features: int,
-        out_features: int, 
+        out_features: int,
         bias: bool = True,
-        gather_results: bool = True
+        gather_results: bool = True,
     ):
         super().__init__(process_group)
-        
+
         self.in_features = in_features
         self.out_features = out_features
         self.out_features_per_rank = out_features // self.world_size
         self.gather_results = gather_results
 
         if out_features % self.world_size != 0:
-            raise ValueError(f"out_features must be divisible by world_size. Got {out_features} and {self.world_size}")
+            raise ValueError(
+                f"out_features must be divisible by world_size. Got {out_features} and {self.world_size}"
+            )
+
+        self.weight = nn.Parameter(
+            torch.empty(self.out_features_per_rank, self.in_features)
+        )
+        self.bias = (
+            nn.Parameter(torch.zeros(self.out_features_per_rank)) if bias else None
+        )
 
-        self.weight = nn.Parameter(torch.empty(self.out_features_per_rank, self.in_features))
-        self.bias = nn.Parameter(torch.zeros(self.out_features_per_rank)) if bias else None
-        
     def forward(self, input: Tensor) -> Tensor:
         output = F.linear(input, self.weight, self.bias)
-        
+
         if self.gather_results:
             output_list = [torch.empty_like(output) for _ in range(self.world_size)]
             dist.all_gather(output_list, output, group=self.process_group)
             output = torch.cat(output_list, dim=-1)
-        
+
         return output
-    
+
-    def load_state_dict(self, state_dict: Dict[str, Tensor]):        
+    def load_state_dict(self, state_dict: Dict[str, Tensor]):
-        full_weight = state_dict.get('weight')
+        full_weight = state_dict.get("weight")
-        full_bias = state_dict.get('bias')
+        full_bias = state_dict.get("bias")
-        
+
         start_idx = self.rank * self.out_features_per_rank
         end_idx = start_idx + self.out_features_per_rank
         weight_slice = full_weight[start_idx:end_idx, :]
         self.weight.data.copy_(weight_slice)
-        
+
         if self.bias is not None:
             bias_slice = full_bias[start_idx:end_idx]
-            self.bias.data.copy_(bias_slice)
+            self.bias.data.copy_(bias_slice)

astrai/parallel/setup.py

@@ -1,151 +1,170 @@
 import os
+from contextlib import contextmanager
+from functools import wraps
+from typing import Callable, List, Optional
+
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from functools import wraps
-from contextlib import contextmanager
-from typing import Callable, List, Optional
-
 
 def get_current_device():
     return os.environ["LOCAL_DEVICE"]
 
+
 def get_world_size() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_world_size()
     else:
         return 1
 
+
 def get_rank() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_rank()
     else:
         return 0
 
+
 @contextmanager
 def setup_parallel(
-    rank: int, 
+    rank: int,
     world_size: int,
     backend: str = "nccl",
     master_addr: str = "localhost",
     master_port: str = "29500",
     device_type: str = "cuda",
-    device_ids: Optional[List[int]] = None
+    device_ids: Optional[List[int]] = None,
 ):
 
     if dist.is_available() and dist.is_initialized():
         yield dist.group.WORLD
-        return 
+        return
 
     if world_size <= 1:
         yield None
         return
-    
+
     if device_ids is None:
         device_ids = [i for i in range(world_size)]
-    
+
     rank = device_ids[rank % len(device_ids)]
     device_id = torch.device(device_type, device_ids[rank])
-    
+
-    os.environ['MASTER_ADDR'] = master_addr
+    os.environ["MASTER_ADDR"] = master_addr
-    os.environ['MASTER_PORT'] = master_port
+    os.environ["MASTER_PORT"] = master_port
-    
+
-    os.environ['LOCAL_RANK'] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
-    os.environ['WORLD_SIZE'] = str(world_size)
+    os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["LOCAL_DEVICE"] = str(device_id)
-    
+
     dist.init_process_group(
-        rank=rank,
+        rank=rank, world_size=world_size, backend=backend, device_id=device_id
-        world_size=world_size,
-        backend=backend,
-        device_id=device_id
     )
-    
+
     try:
         if backend == "nccl" and torch.cuda.is_available():
             torch.cuda.set_device(device_id)
-        elif backend == "ccl" and hasattr(torch, 'xpu') and torch.xpu.is_available():
+        elif backend == "ccl" and hasattr(torch, "xpu") and torch.xpu.is_available():
             torch.xpu.set_device(device_id)
-        
+
         yield dist.group.WORLD
     finally:
         if dist.is_initialized():
             dist.destroy_process_group()
 
+
 def only_on_rank(rank, sync=False):
     """
     decorator to run a function only on a specific rank.
     """
-    
+
     def decorator(func):
         @wraps(func)
         def wrapper(*args, **kwargs):
+            ret_args = None
             if get_rank() == rank:
-                return func(*args, **kwargs)
+                ret_args = func(*args, **kwargs)
-            if sync:
+
+            if sync and dist.is_available() and dist.is_initialized():
                 dist.barrier()
-        
+
+            return ret_args
+
         return wrapper
-    
+
     return decorator
 
+
 def wrapper_spawn_func(
-    rank: int, 
+    rank: int,
-    world_size: int, 
+    world_size: int,
-    backend: str, 
+    backend: str,
-    master_addr: str, 
+    master_addr: str,
     master_port: str,
     device_type: str,
-    device_ids: List[int], 
+    device_ids: List[int],
-    func: Callable, 
+    func: Callable,
-    kwargs: dict
+    kwargs: dict,
 ):
     try:
         with setup_parallel(
-            rank=rank, 
+            rank=rank,
-            world_size=world_size, 
+            world_size=world_size,
-            backend=backend, 
+            backend=backend,
-            master_addr=master_addr, 
+            master_addr=master_addr,
             master_port=master_port,
             device_type=device_type,
-            device_ids=device_ids
+            device_ids=device_ids,
         ):
             func(**kwargs)
-        
+
     except Exception as e:
         print(f"Error in rank {rank}: {e}")
         raise
 
+
 def spawn_parallel_fn(
-    func: Callable, 
+    func: Callable,
-    world_size: int, 
+    world_size: int,
     backend: str = "nccl",
     master_addr: str = "localhost",
     master_port: str = "29500",
     device_type: str = "cuda",
     device_ids: Optional[List[int]] = None,
-    **kwargs
+    **kwargs,
 ):
     # clear environment variables
-    for key in ['MASTER_ADDR', 'MASTER_PORT', 'RANK', 'WORLD_SIZE', 'LOCAL_RANK', 'LOCAL_DEVICE']:
+    for key in [
+        "MASTER_ADDR",
+        "MASTER_PORT",
+        "RANK",
+        "WORLD_SIZE",
+        "LOCAL_RANK",
+        "LOCAL_DEVICE",
+    ]:
         if key in os.environ:
             del os.environ[key]
-    
+
     if world_size == 1:
         device_ids = device_ids or [0]
-        deice_id = torch.device(device_type, device_ids[0])
+        device_id = torch.device(device_type, device_ids[0])
-        os.environ["LOCAL_DEVICE"] = str(deice_id)
+        os.environ["LOCAL_DEVICE"] = str(device_id)
-        
+
         func(**kwargs)
         return
 
-    wrapper_spawn_func_args = (world_size, backend, master_addr, master_port, 
+    wrapper_spawn_func_args = (
-                               device_type, device_ids, func, kwargs)
+        world_size,
+        backend,
+        master_addr,
+        master_port,
+        device_type,
+        device_ids,
+        func,
+        kwargs,
+    )
 
     mp.spawn(
-        wrapper_spawn_func, 
+        wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
-        nprocs=world_size, 
+    )
-        args=wrapper_spawn_func_args, 
-        join=True
-    )

astrai/serialization.py

@@ -0,0 +1,106 @@
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, List
+
+import h5py
+import safetensors.torch as st
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+from astrai.parallel.setup import get_rank
+
+
+def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
+    os.makedirs(file_path, exist_ok=True)
+    full_file_path = os.path.join(file_path, f"{file_name}.h5")
+    with h5py.File(full_file_path, "w") as f:
+        for key, tensors in tensor_group.items():
+            grp = f.create_group(key)
+            for idx, tensor in enumerate(tensors):
+                arr = tensor.cpu().numpy()
+                grp.create_dataset(f"data_{idx}", data=arr)
+
+
+def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
+    tensor_group: Dict[str, List[Tensor]] = {}
+
+    root_path = Path(file_path)
+    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
+
+    for h5_file in h5_files:
+        with h5py.File(h5_file, "r") as f:
+            for key in f.keys():
+                grp = f[key]
+                dsets = []
+                for dset_name in grp.keys():
+                    dset = grp[dset_name]
+                    tensor = torch.from_numpy(dset[:])
+                    if share_memory:
+                        tensor = tensor.share_memory_()
+                    dsets.append(tensor)
+
+                if tensor_group.get(key) is None:
+                    tensor_group[key] = []
+                tensor_group[key].extend(dsets)
+
+    return tensor_group
+
+
+class Checkpoint:
+    def __init__(
+        self,
+        state_dict: Dict[str, Any],
+        epoch: int = 0,
+        iteration: int = 0,
+    ):
+        self.state_dict = state_dict
+        self.epoch = epoch
+        self.iteration = iteration
+
+    def save(
+        self,
+        save_dir: str,
+    ) -> None:
+
+        save_path = Path(save_dir)
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        rank = get_rank()
+        if rank == 0:
+            meta = {
+                "epoch": self.epoch,
+                "iteration": self.iteration,
+            }
+            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")
+
+    @classmethod
+    def load(
+        cls,
+        save_dir: str,
+    ) -> "Checkpoint":
+
+        rank = get_rank()
+        save_path = Path(save_dir)
+
+        meta = {}
+        if rank == 0:
+            with open(Path(save_dir) / "meta.json", "r") as f:
+                meta = json.load(f)
+
+        if dist.is_initialized():
+            meta_list = [meta]
+            dist.broadcast_object_list(meta_list, src=0)
+            meta = meta_list[0]
+
+        state_dict = st.load_file(save_path / "state_dict.safetensors")
+
+        return cls(
+            state_dict=state_dict,
+            epoch=meta["epoch"],
+            iteration=meta["iteration"],
+        )

astrai/tokenize/__init__.py

@@ -0,0 +1,8 @@
+from astrai.tokenize.chat_template import ChatTemplate, MessageType
+from astrai.tokenize.tokenizer import AutoTokenizer
+
+__all__ = [
+    "AutoTokenizer",
+    "ChatTemplate",
+    "MessageType",
+]

astrai/tokenize/chat_template.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+from jinja2 import Template
+
+# Message type for chat messages
+type MessageType = Dict[str, Any]
+
+
+@dataclass
+class ChatTemplate:
+    """A chat template with Jinja2 rendering support.
+
+    Attributes:
+        name: Unique identifier for the template.
+        template_str: Jinja2 template string.
+        description: Optional description.
+        default_variables: Optional dictionary of default variable values
+            that will be passed to the template if not overridden during rendering.
+        special_tokens: Optional dictionary mapping token names to their string values.
+            These tokens are automatically added to the template variables.
+    """
+
+    name: str
+    template_str: str
+    description: str = ""
+    default_variables: Dict[str, Any] = None
+    special_tokens: Dict[str, str] = None
+
+    def __post_init__(self):
+        if self.default_variables is None:
+            self.default_variables = {}
+        if self.special_tokens is None:
+            self.special_tokens = {}
+
+    @classmethod
+    def from_string(
+        cls,
+        template_str: str,
+        description: str = "",
+        default_variables: Optional[Dict[str, Any]] = None,
+        special_tokens: Optional[Dict[str, str]] = None,
+    ) -> "ChatTemplate":
+        """Create a ChatTemplate instance directly from a template string."""
+        return cls(
+            name="",  # empty name for ad‑hoc templates
+            template_str=template_str,
+            description=description,
+            default_variables=default_variables,
+            special_tokens=special_tokens,
+        )
+
+    def render(
+        self,
+        messages: List[MessageType],
+        system_prompt: Optional[str] = None,
+        **extra_variables: Any,
+    ) -> str:
+        """Render the template with given messages and variables.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'.
+            system_prompt: Optional system prompt string.
+            **extra_variables: Additional variables to pass to the template.
+                These override default_variables and special_tokens.
+
+        Returns:
+            Rendered prompt string.
+        """
+        # Merge default variables, special tokens, and extra variables
+        variables = {**self.default_variables, **self.special_tokens, **extra_variables}
+        variables["messages"] = messages
+        if system_prompt is not None:
+            variables["system_prompt"] = system_prompt
+
+        jinja_template = Template(self.template_str)
+        return jinja_template.render(**variables)

astrai/tokenize/tokenizer.py

@@ -0,0 +1,242 @@
+"""
+Tokenizer module with implementation and auto-loading support.
+"""
+
+import json
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+from tokenizers import Tokenizer
+
+from astrai.tokenize.chat_template import ChatTemplate
+
+
+class AutoTokenizer:
+    """Base tokenizer class with automatic loading support"""
+
+    TOKENIZER_CLASSES = {}  # Registry for auto-loading
+
+    def __init__(
+        self,
+        path: Optional[Union[str, Path]] = None,
+        special_token_map: Optional[Dict[str, str]] = None,
+        chat_template: Optional[str] = None,
+    ):
+        self._tokenizer: Tokenizer = None
+        self._chat_template: Optional[ChatTemplate] = None
+        self._special_token_map: Optional[Dict] = special_token_map or {}
+
+        if chat_template:
+            self.set_chat_template(chat_template)
+
+        if path:
+            self.load(path)
+
+    def load(self, path: Union[str, Path]):
+        """Load tokenizer from directory."""
+        path = Path(path)
+        tokenizer_file = path / "tokenizer.json"
+        config_file = path / "tokenizer_config.json"
+        self._tokenizer = Tokenizer.from_file(str(tokenizer_file))
+
+        if config_file.exists():
+            with open(config_file, "r", encoding="utf-8") as f:
+                config = json.load(f)
+
+            if "special_tokens" in config:
+                self._special_token_map.update(config["special_tokens"])
+
+            # Load chat template from config
+            if "chat_template" in config:
+                self.set_chat_template(config["chat_template"])
+
+    @classmethod
+    def from_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoTokenizer":
+        """Load tokenizer from pretrained directory."""
+        instance = cls(path)
+        return instance
+
+    def save_pretrained(self, save_path: str):
+        """
+        Save tokenizer to pretrained directory.
+
+        Args:
+            save_path: Path to save the tokenizer
+        """
+
+        save_path = Path(save_path)
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        # Save tokenizer
+        self._tokenizer.save(str(save_path / "tokenizer.json"))
+
+        # Save tokenizer config
+        config = {}
+        if self._special_token_map is not None:
+            config["special_tokens"] = self._special_token_map
+        if self._chat_template is not None:
+            config["chat_template"] = self._chat_template.template_str
+
+        with open(save_path / "tokenizer_config.json", "w", encoding="utf-8") as f:
+            json.dump(config, f, ensure_ascii=False, indent=2)
+
+    @classmethod
+    def register_tokenizer(cls, name: str, tokenizer_class: type):
+        """
+        Register a new tokenizer class.
+
+        Args:
+            name: Name to register the tokenizer class under
+            tokenizer_class: The tokenizer class to register
+        """
+        cls.TOKENIZER_CLASSES[name] = tokenizer_class
+
+    def encode(
+        self,
+        tokens: Union[str, List[str]],
+        out_ids: bool = True,
+        is_pretokenized: bool = False,
+        add_special_tokens: bool = True,
+    ) -> List:
+        """Encode text to tokens or token IDs."""
+        if self._tokenizer is None:
+            raise RuntimeError(
+                "Tokenizer not initialized. Load or create a tokenizer first."
+            )
+
+        if isinstance(tokens, str):
+            encoded = self._tokenizer.encode(
+                tokens,
+                is_pretokenized=is_pretokenized,
+                add_special_tokens=add_special_tokens,
+            )
+            return encoded.ids if out_ids else encoded.tokens
+        else:
+            encoded_list = self._tokenizer.encode_batch(
+                tokens,
+                is_pretokenized=is_pretokenized,
+                add_special_tokens=add_special_tokens,
+            )
+            return [
+                encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
+            ]
+
+    def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
+        """Decode token IDs to text."""
+        if self._tokenizer is None:
+            raise RuntimeError(
+                "Tokenizer not initialized. Load or create a tokenizer first."
+            )
+
+        return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
+
+    def __len__(self) -> int:
+        if self._tokenizer is None:
+            return 0
+        return self._tokenizer.get_vocab_size()
+
+    def __getattr__(self, key: str):
+        """
+        Dynamically intercept special token attribute access.
+        Supports three forms:
+          - tokenizer.bos_token   → returns string
+          - tokenizer.bos_token_id → returns corresponding integer ID
+          - tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens
+        """
+        # Handle stop_ids - return IDs for all special tokens
+        if key == "stop_ids":
+            stop_ids = []
+
+            if self._tokenizer is None:
+                return stop_ids
+
+            for val in self._special_token_map.values():
+                token_id = self._tokenizer.token_to_id(val)
+                if token_id is not None:
+                    stop_ids.append(token_id)
+
+            return stop_ids
+
+        # Handle _id suffix (e.g., bos_token_id -> bos_token)
+        if key.endswith("_id"):
+            base_attr = key[:-3]  # Remove "_id"
+            token_str = self._special_token_map.get(base_attr)
+            if token_str is None:
+                return None
+            if self._tokenizer is None:
+                raise RuntimeError("Tokenizer not loaded, cannot convert token to id.")
+            return self._tokenizer.token_to_id(token_str)
+
+        # Handle regular string attributes
+        if key in self._special_token_map:
+            return self._special_token_map.get(key)
+
+        # Other attributes trigger default AttributeError
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{key}'")
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self)
+
+    def set_chat_template(self, template: Union[str, ChatTemplate]):
+        """
+        Set the chat template for the tokenizer.
+
+        Args:
+            template: Either a template name (str) registered in the global registry,
+                      or a ChatTemplate instance, or a Jinja2 template string.
+
+        Raises:
+            KeyError: If template name is not registered.
+        """
+        if isinstance(template, str):
+            self._chat_template = ChatTemplate.from_string(template)
+        elif isinstance(template, ChatTemplate):
+            self._chat_template = template
+        else:
+            raise ValueError("Invalid template type, must be str or ChatTemplate.")
+
+    def apply_chat_template(
+        self,
+        messages: List[Dict[str, str]],
+        system_prompt: Optional[str] = None,
+        tokenize: bool = True,
+        add_generation_prompt: bool = True,
+        **kwargs,
+    ) -> Union[str, List[int]]:
+        """
+        Apply the chat template to messages and optionally tokenize the result.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'.
+            system_prompt: Optional system prompt string (auto-converted to first message).
+            tokenize: Whether to return token IDs (True) or raw string (False).
+            add_generation_prompt: Whether to add the generation prompt (default: True).
+            **kwargs: Additional variables to pass to the template.
+
+        Returns:
+            Either the rendered string or list of token IDs.
+
+        Raises:
+            RuntimeError: If chat template is not set.
+        """
+        if self._chat_template is None:
+            raise RuntimeError(
+                "Chat template not set. Use set_chat_template() to set a template first."
+            )
+
+        # Auto-convert system_prompt to first message if provided
+        if system_prompt:
+            messages = [{"role": "system", "content": system_prompt}] + list(messages)
+
+        # Render the template
+        rendered = self._chat_template.render(
+            messages=messages,
+            add_generation_prompt=add_generation_prompt,
+            **kwargs,
+        )
+
+        if tokenize:
+            return self.encode(rendered)
+
+        return rendered

astrai/trainer/__init__.py

@@ -0,0 +1,21 @@
+from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
+from astrai.trainer.train_callback import (
+    CallbackFactory,
+    TrainCallback,
+)
+from astrai.trainer.trainer import Trainer
+
+__all__ = [
+    # Main trainer
+    "Trainer",
+    # Strategy factory
+    "StrategyFactory",
+    "BaseStrategy",
+    # Scheduler factory
+    "SchedulerFactory",
+    "BaseScheduler",
+    # Callback factory
+    "TrainCallback",
+    "CallbackFactory",
+]

astrai/trainer/metric_util.py

@@ -1,8 +1,10 @@
-import torch.nn as nn
 from typing import Dict
 
+import torch.nn as nn
+
+
 def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
-    """ Compute gradient norm for each parameter in the model. """
+    """Compute gradient norm for each parameter in the model."""
     norms = {}
     for name, param in model.named_parameters():
         norms[name] = 0.0
@@ -11,8 +13,9 @@ def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
             norms[name] = norm
     return norms
 
+
 def grad_std(model: nn.Module) -> Dict[str, float]:
-    """ Compute standard deviation of gradients for each parameter. """
+    """Compute standard deviation of gradients for each parameter."""
     stds = {}
     for name, param in model.named_parameters():
         stds[name] = 0.0
@@ -21,41 +24,45 @@ def grad_std(model: nn.Module) -> Dict[str, float]:
             stds[name] = std
     return stds
 
+
 def grad_max(model: nn.Module) -> Dict[str, float]:
-    """ Find the maximum absolute gradient value for each parameter. """
+    """Find the maximum absolute gradient value for each parameter."""
     max_vals = {}
     for name, param in model.named_parameters():
-        max_vals[name] = -float('inf')
+        max_vals[name] = -float("inf")
         if param.grad:
             max_val = param.grad.data.max().item()
             max_vals[name] = max_val
-            
+
     return max_vals
 
+
 def grad_min(model: nn.Module) -> Dict[str, float]:
-    """ Find the minimum absolute gradient value for each parameter. """
+    """Find the minimum absolute gradient value for each parameter."""
     min_vals = {}
     for name, param in model.named_parameters():
-        min_vals[name] = float('inf')
+        min_vals[name] = float("inf")
         if param.grad:
             min_val = param.grad.data.min().item()
             min_vals[name] = min_val
-            
+
     return min_vals
 
+
 def grad_mean(model: nn.Module) -> Dict[str, float]:
-    """ Compute mean of gradients for each parameter. """
+    """Compute mean of gradients for each parameter."""
     means = {}
     for name, param in model.named_parameters():
         means[name] = 0.0
         if param.grad:
             mean = param.grad.data.mean().item()
             means[name] = mean
-            
+
     return means
 
+
 def grad_nan_num(model: nn.Module) -> Dict[str, int]:
-    """ Count the number of NaNs in gradients for each parameter. """
+    """Count the number of NaNs in gradients for each parameter."""
     nan_nums = {}
     for name, param in model.named_parameters():
         nan_nums[name] = 0
@@ -64,26 +71,34 @@ def grad_nan_num(model: nn.Module) -> Dict[str, int]:
             nan_nums[name] = nan_num
     return nan_nums
 
+
 def ctx_get_loss(ctx):
     return ctx.loss
 
+
 def ctx_get_lr(ctx):
-    return ctx.optimizer.param_groups[-1]['lr']
+    return ctx.optimizer.param_groups[-1]["lr"]
+
 
 def ctx_get_grad_norm(ctx):
     return grad_norm(ctx.model)
 
+
 def ctx_get_grad_std(ctx):
     return grad_std(ctx.model)
 
+
 def ctx_get_grad_max(ctx):
     return grad_max(ctx.model)
 
+
 def ctx_get_grad_min(ctx):
     return grad_min(ctx.model)
 
+
 def ctx_get_grad_mean(ctx):
     return grad_mean(ctx.model)
 
+
 def ctx_get_grad_nan_num(ctx):
-    return grad_nan_num(ctx.model)
+    return grad_nan_num(ctx.model)

astrai/trainer/schedule.py

@@ -0,0 +1,194 @@
+"""Learning rate scheduler implementations with factory pattern."""
+
+import math
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Type
+
+from torch.optim.lr_scheduler import LRScheduler
+
+from astrai.factory import BaseFactory
+
+
+class BaseScheduler(LRScheduler, ABC):
+    """Base scheduler class for all other schedulers."""
+
+    def __init__(self, optimizer, last_epoch: int = -1):
+        super().__init__(optimizer, last_epoch)
+
+    @abstractmethod
+    def get_lr(self) -> List[float]:
+        """Calculate the current learning rate."""
+        raise NotImplementedError
+
+    def state_dict(self) -> Dict[str, Any]:
+        return super().state_dict()
+
+    def load_state_dict(self, state_dict: Dict[str, Any]):
+        super().load_state_dict(state_dict)
+
+
+class SchedulerFactory(BaseFactory["BaseScheduler"]):
+    """Factory class for creating learning rate schedulers.
+
+    Supports decorator-based registration for extensible scheduler types.
+    Also supports creation from ScheduleConfig objects.
+
+    Example usage:
+        @SchedulerFactory.register("custom")
+        class CustomScheduler(BaseScheduler):
+            ...
+
+        scheduler = SchedulerFactory.create("custom", optimizer, **kwargs)
+    """
+
+    @classmethod
+    def _validate_component(cls, scheduler_cls: Type[BaseScheduler]) -> None:
+        """Validate that the scheduler class inherits from BaseScheduler."""
+        if not issubclass(scheduler_cls, BaseScheduler):
+            raise TypeError(f"{scheduler_cls.__name__} must inherit from BaseScheduler")
+
+    @classmethod
+    def create(
+        cls, optimizer, schedule_type: str = "none", **kwargs
+    ) -> "BaseScheduler":
+        """Create a scheduler instance by type name.
+
+        Args:
+            optimizer: PyTorch optimizer
+            schedule_type: Type of scheduler ("cosine", "sgdr")
+            **kwargs: Arguments passed to the scheduler constructor
+
+        Returns:
+            Scheduler instance
+        """
+        return super().create(schedule_type, optimizer, **kwargs)
+
+    @classmethod
+    def available_types(cls) -> list:
+        """Return list of registered scheduler type names."""
+        return cls.list_registered()
+
+
+# ----------- Scheduler implementations -----------
+
+
+@SchedulerFactory.register("cosine")
+class CosineScheduler(BaseScheduler):
+    """Cosine decay scheduler with warmup, implemented as PyTorch LRScheduler."""
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_steps: int,
+        lr_decay_steps: int,
+        min_rate: float = 0.05,
+        last_epoch: int = -1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.lr_decay_steps = lr_decay_steps
+        self.min_rate = min_rate
+        self.total_steps = warmup_steps + lr_decay_steps
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> List[float]:
+        # warmup
+        if self.last_epoch < self.warmup_steps:
+            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
+            return [base_lr * warmup_factor for base_lr in self.base_lrs]
+
+        # cosine decay
+        decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
+        decay_progress = min(decay_progress, 1.0)
+        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
+        decay_factor = max(self.min_rate, cosine_decay)
+        return [base_lr * decay_factor for base_lr in self.base_lrs]
+
+    def state_dict(self):
+        state = super().state_dict()
+        state.update(
+            {
+                "warmup_steps": self.warmup_steps,
+                "lr_decay_steps": self.lr_decay_steps,
+                "min_rate": self.min_rate,
+                "total_steps": self.total_steps,
+            }
+        )
+        return state
+
+    def load_state_dict(self, state_dict):
+        self.warmup_steps = state_dict.pop("warmup_steps")
+        self.lr_decay_steps = state_dict.pop("lr_decay_steps")
+        self.min_rate = state_dict.pop("min_rate")
+        self.total_steps = state_dict.pop("total_steps")
+        super().load_state_dict(state_dict)
+
+
+@SchedulerFactory.register("sgdr")
+class SGDRScheduler(BaseScheduler):
+    """SGDR (Stochastic Gradient Descent with Warm Restarts) scheduler."""
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_steps: int,
+        cycle_length: int,
+        min_rate: float = 0.05,
+        t_mult: int = 2,
+        last_epoch: int = -1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.cycle_length = cycle_length
+        self.min_rate = min_rate
+        self.t_mult = t_mult
+
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        # warmup
+        if self.last_epoch < self.warmup_steps:
+            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
+            return [base_lr * warmup_factor for base_lr in self.base_lrs]
+
+        # SGDR
+        steps_since_warmup = self.last_epoch - self.warmup_steps
+
+        # 1. Calculate current cycle and position within cycle
+        current_cycle_length = self.cycle_length
+        total_cycles_length = 0
+        cycle_num = 0
+
+        while total_cycles_length + current_cycle_length <= steps_since_warmup:
+            total_cycles_length += current_cycle_length
+            current_cycle_length *= self.t_mult
+            cycle_num += 1
+
+        steps_in_cycle = steps_since_warmup - total_cycles_length
+
+        # 2. Cosine annealing within the current cycle
+        cosine_factor = 0.5 * (
+            1 + math.cos(math.pi * steps_in_cycle / current_cycle_length)
+        )
+        learning_rate_factor = self.min_rate + (1 - self.min_rate) * cosine_factor
+
+        return [base_lr * learning_rate_factor for base_lr in self.base_lrs]
+
+    def state_dict(self):
+        """Returns the state of the scheduler as a dict."""
+        state = super().state_dict()
+        state.update(
+            {
+                "warmup_steps": self.warmup_steps,
+                "cycle_length": self.cycle_length,
+                "min_rate": self.min_rate,
+                "t_mult": self.t_mult,
+            }
+        )
+        return state
+
+    def load_state_dict(self, state_dict):
+        """Loads the scheduler's state."""
+        self.warmup_steps = state_dict.pop("warmup_steps")
+        self.cycle_length = state_dict.pop("cycle_length")
+        self.min_rate = state_dict.pop("min_rate")
+        self.t_mult = state_dict.pop("t_mult")
+        super().load_state_dict(state_dict)

astrai/trainer/strategy.py

@@ -0,0 +1,330 @@
+"""Training strategy implementations with factory pattern."""
+
+import copy
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from astrai.factory import BaseFactory
+
+
+def unwrap_model(model: nn.Module) -> nn.Module:
+    """Unwrap DDP wrapper if present to get the original model."""
+    if isinstance(model, DDP):
+        return model.module
+    return model
+
+
+def create_ref_model(model: nn.Module) -> nn.Module:
+    """Create a reference model for DPO/GRPO training.
+
+    Handles DDP-wrapped models safely by unwrapping first,
+    then creating a deep copy with frozen gradients.
+    """
+    original_model = unwrap_model(model)
+    ref_model = copy.deepcopy(original_model)
+    ref_model.requires_grad_(False)
+    ref_model.eval()
+    return ref_model
+
+
+def move_to_device(batch: Dict[str, Tensor], device: str) -> Any:
+    """Move batch tensors to specified device with non-blocking transfer."""
+    return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
+
+
+def get_logprobs(
+    model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
+    input_ids: Tensor,
+    mask: Tensor,
+    reduction: str,
+):
+    """Compute token-wise log probabilities from model outputs.
+
+    Args:
+        model: The language model
+        input_ids: Input token IDs of shape [batch_size, seq_len]
+        mask: Attention mask of shape [batch_size, seq_len]
+        reduction: How to reduce over sequence dimension ("mean", "sum", "none")
+
+    Returns:
+        Log probabilities with reduction applied over sequence dimension
+    """
+    allowed_reductions = ["mean", "sum", "none"]
+    if reduction not in allowed_reductions:
+        raise ValueError(
+            f"reduction must be one of {allowed_reductions}, got '{reduction}'"
+        )
+
+    shifted_input_ids = input_ids[:, 1:]
+    shifted_mask = mask[:, 1:]
+
+    logits = model(input_ids[:, :-1], mask[:, :-1])["logits"]
+    log_probs = torch.log_softmax(logits.float(), dim=-1)
+
+    token_logprobs = torch.gather(
+        log_probs, dim=-1, index=shifted_input_ids.unsqueeze(-1)
+    ).squeeze(-1)
+
+    if reduction == "mean":
+        return (token_logprobs * shifted_mask).sum(dim=-1) / shifted_mask.sum(
+            dim=-1
+        ).clamp(min=1.0)
+    elif reduction == "sum":
+        return (token_logprobs * shifted_mask).sum(dim=-1)
+    else:
+        return token_logprobs * shifted_mask
+
+
+class BaseStrategy(ABC):
+    """Abstract base class for training strategies."""
+
+    def __init__(
+        self, model: Union[Callable[..., Dict[str, Tensor]]], device: str, **kwargs
+    ):
+        self.model = model
+        self.device = device
+        self.extra_kwargs = kwargs
+
+    @abstractmethod
+    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        """Compute loss for the given batch.
+
+        Args:
+            batch: Dictionary containing batch tensors
+
+        Returns:
+            Computed loss tensor
+        """
+        raise NotImplementedError
+
+    def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
+        """Allow calling strategy directly as a callable."""
+        return self.compute_loss(batch)
+
+
+class StrategyFactory(BaseFactory["BaseStrategy"]):
+    """Factory class for creating training strategy instances.
+
+    Supports decorator-based registration for extensible strategy types.
+    All default strategies (seq, sft, dpo, grpo) are automatically registered.
+
+    Example usage:
+        @StrategyFactory.register("custom")
+        class CustomStrategy(BaseStrategy):
+            ...
+
+        strategy = StrategyFactory.create("custom", model, device)
+    """
+
+    @classmethod
+    def _validate_component(cls, strategy_cls: type) -> None:
+        """Validate that the strategy class inherits from BaseStrategy."""
+        if not issubclass(strategy_cls, BaseStrategy):
+            raise TypeError(f"{strategy_cls.__name__} must inherit from BaseStrategy")
+
+    @classmethod
+    def create(cls, train_type: str, model, device: str, **kwargs) -> "BaseStrategy":
+        """Create a strategy instance based on training type.
+
+        Args:
+            train_type: Type of training ("seq", "sft", "dpo", "grpo")
+            model: Model instance for the strategy
+            device: Device to run the strategy on
+            **kwargs: Additional arguments passed to strategy constructor
+
+        Returns:
+            Strategy instance
+        """
+        return super().create(train_type, model, device, **kwargs)
+
+    @classmethod
+    def available_strategies(cls) -> list:
+        """Return list of registered strategy names."""
+        return cls.list_registered()
+
+
+# ============== Strategy Classes ==============
+# All strategies are registered at class definition time using the decorator
+
+
+@StrategyFactory.register("seq")
+class SEQStrategy(BaseStrategy):
+    """Standard next-token prediction training strategy.
+
+    Computes cross-entropy loss for next token prediction.
+    """
+
+    def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
+        super().__init__(model, device, **kwargs)
+        self.label_smoothing = label_smoothing
+
+    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        batch = move_to_device(batch, self.device)
+        input_ids, target_ids = batch["input_ids"], batch["target_ids"]
+        logits = self.model(input_ids=input_ids)["logits"]
+
+        loss = F.cross_entropy(
+            input=logits.flatten(0, 1).float(),
+            target=target_ids.flatten(),
+            label_smoothing=self.label_smoothing,
+        )
+
+        return loss
+
+
+@StrategyFactory.register("sft")
+class SFTStrategy(BaseStrategy):
+    """Supervised Fine-tuning strategy with loss masking.
+
+    Applies cross-entropy loss only to tokens where loss_mask is True.
+    """
+
+    def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
+        super().__init__(model, device, **kwargs)
+        self.label_smoothing = label_smoothing
+
+    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        batch = move_to_device(batch, self.device)
+        input_ids, target_ids, loss_mask = (
+            batch["input_ids"],
+            batch["target_ids"],
+            batch["loss_mask"],
+        )
+
+        ignore_index = -100
+        logits = self.model(input_ids=input_ids)["logits"]
+        target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
+
+        loss = F.cross_entropy(
+            input=logits.flatten(0, 1).float(),
+            target=target_ids.flatten(),
+            ignore_index=ignore_index,
+            label_smoothing=self.label_smoothing,
+        )
+
+        return loss
+
+
+@StrategyFactory.register("dpo")
+class DPOStrategy(BaseStrategy):
+    """Direct Preference Optimization strategy.
+
+    Implements the DPO loss from the paper "Direct Preference Optimization".
+    Uses a reference model to compute KL divergence penalty.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        device: str,
+        beta: float = 0.1,
+        reduction: str = "mean",
+        **kwargs,
+    ):
+        super().__init__(model, device, **kwargs)
+        self.ref_model = create_ref_model(model)
+        self.beta = beta
+        self.reduction = reduction
+
+    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        batch = move_to_device(batch, self.device)
+        chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
+        chosen_mask, rejected_mask = batch["chosen_mask"], batch["rejected_mask"]
+
+        concat_ids = torch.cat([chosen_ids, rejected_ids], dim=0)
+        concat_mask = torch.cat([chosen_mask, rejected_mask], dim=0)
+
+        log_pi = get_logprobs(self.model, concat_ids, concat_mask, self.reduction)
+
+        with torch.no_grad():
+            log_ref = get_logprobs(
+                self.ref_model, concat_ids, concat_mask, self.reduction
+            )
+
+        log_pi_chosen = log_pi[: chosen_ids.shape[0]]
+        log_pi_rejected = log_pi[chosen_ids.shape[0] :]
+        log_ref_chosen = log_ref[: chosen_ids.shape[0]]
+        log_ref_rejected = log_ref[chosen_ids.shape[0] :]
+
+        pi_log_ratio = log_pi_chosen - log_pi_rejected
+        ref_log_ratio = log_ref_chosen - log_ref_rejected
+
+        ratio_diff = pi_log_ratio - ref_log_ratio
+        dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
+
+        return dpo_loss
+
+
+@StrategyFactory.register("grpo")
+class GRPOStrategy(BaseStrategy):
+    """Group Relative Policy Optimization strategy.
+
+    Implements GRPO with clipping and KL penalty.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        device: str,
+        clip_eps: float = 0.2,
+        kl_coef: float = 0.01,
+        group_size: int = 4,
+        reduction: str = "mean",
+        **kwargs,
+    ):
+        super().__init__(model, device, **kwargs)
+        self.ref_model = create_ref_model(model)
+        self.clip_eps = clip_eps
+        self.kl_coef = kl_coef
+        self.group_size = group_size
+        self.reduction = reduction
+
+    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        batch = move_to_device(batch, self.device)
+        prompts = batch["prompts"]
+        responses = batch["responses"]
+        masks = batch["masks"]
+        rewards = batch["rewards"]
+
+        batch_size, group_size, response_len = responses.shape
+        responses_flat = responses.view(-1, response_len)
+        masks_flat = masks.view(-1, response_len)
+        prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
+
+        # Shape: (batch_size * group_size, seq_len + response_len)
+        full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
+        full_masks = torch.cat([torch.ones_like(prompt_expanded), masks_flat], dim=-1)
+
+        log_probs_policy = get_logprobs(
+            self.model, full_sequences, full_masks, self.reduction
+        )
+        log_probs_policy = log_probs_policy.view(batch_size, group_size)
+
+        with torch.no_grad():
+            log_probs_ref = get_logprobs(
+                self.ref_model, full_sequences, full_masks, self.reduction
+            )
+            log_probs_ref = log_probs_ref.view(batch_size, group_size)
+
+        # Compute advantages from rewards with normalization
+        eps = torch.finfo(log_probs_policy.dtype).eps
+        mean = rewards.mean(dim=-1, keepdim=True)
+        std = rewards.std(dim=-1, keepdim=True)
+        advantages = (rewards - mean) / (std + eps)
+
+        # PPO-style clipped surrogate objective
+        ratio = torch.exp(0)  # Off-policy: policy_model = old_model
+        surr1 = ratio * advantages
+        surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
+
+        policy_loss = -torch.min(surr1, surr2).mean()
+        kl_penalty = self.kl_coef * (log_probs_policy - log_probs_ref).square().mean()
+        total_loss = policy_loss + kl_penalty
+
+        return total_loss

astrai/trainer/train_callback.py

@@ -1,161 +1,191 @@
-import os
 import json
+import os
 import time
-import torch.nn as nn
-
 from pathlib import Path
-from tqdm import tqdm
+from typing import Callable, List, Optional, Protocol, runtime_checkable
-from torch.nn.utils import clip_grad_norm_
-from torch.optim.lr_scheduler import LRScheduler
-from typing import Callable, List, Optional, Protocol
 
-from khaosz.parallel import only_on_rank
+import torch.nn as nn
-from khaosz.trainer.metric_util import (
+from torch.nn.utils import clip_grad_norm_
+from tqdm import tqdm
+
+from astrai.factory import BaseFactory
+from astrai.parallel import only_on_rank
+from astrai.serialization import Checkpoint
+from astrai.trainer.metric_util import (
+    ctx_get_grad_max,
+    ctx_get_grad_mean,
+    ctx_get_grad_min,
+    ctx_get_grad_nan_num,
+    ctx_get_grad_norm,
+    ctx_get_grad_std,
     ctx_get_loss,
     ctx_get_lr,
-    ctx_get_grad_max,
-    ctx_get_grad_min,
-    ctx_get_grad_norm,
-    ctx_get_grad_mean,
-    ctx_get_grad_std,
-    ctx_get_grad_nan_num
 )
-from khaosz.data.checkpoint import Checkpoint
+from astrai.trainer.train_context import TrainContext
-from khaosz.trainer.train_context import TrainContext
 
 
+@runtime_checkable
 class TrainCallback(Protocol):
-    """ 
+    """
     Callback interface for trainer.
     """
 
     def on_train_begin(self, context: TrainContext):
-        """ Called at the beginning of training. """
+        """Called at the beginning of training."""
-    
+
     def on_train_end(self, context: TrainContext):
-        """ Called at the end of training. """
+        """Called at the end of training."""
 
     def on_epoch_begin(self, context: TrainContext):
-        """ Called at the beginning of each epoch. """
+        """Called at the beginning of each epoch."""
 
     def on_epoch_end(self, context: TrainContext):
-        """  Called at the end of each epoch. """
+        """Called at the end of each epoch."""
-    
+
     def on_step_begin(self, context: TrainContext):
-        """ Called at the beginning of each step. """
+        """Called at the beginning of each step."""
 
     def on_step_end(self, context: TrainContext):
-        """ Called at the end of each step."""
+        """Called at the end of each step."""
 
     def on_batch_begin(self, context: TrainContext):
-        """ Called at the beginning of each batch. """
+        """Called at the beginning of each batch."""
 
     def on_batch_end(self, context: TrainContext):
-        """ Called at the end of each batch. """
+        """Called at the end of each batch."""
-    
+
     def on_error(self, context: TrainContext):
-        """ Called when an error occurs during training. """
+        """Called when an error occurs during training."""
 
 
+class CallbackFactory(BaseFactory[TrainCallback]):
+    """Factory for registering and creating training callbacks.
+
+    Example:
+        @CallbackFactory.register("my_callback")
+        class MyCallback(TrainCallback):
+            ...
+
+        callback = CallbackFactory.create("my_callback", **kwargs)
+    """
+
+    @classmethod
+    def _validate_component(cls, callback_cls: type) -> None:
+        """Validate that the callback class inherits from TrainCallback."""
+        if not issubclass(callback_cls, TrainCallback):
+            raise TypeError(f"{callback_cls.__name__} must inherit from TrainCallback")
+
+
+@CallbackFactory.register("gradient_clipping")
 class GradientClippingCallback(TrainCallback):
-    """ 
+    """
     Gradient clipping callback for trainer.
     """
+
     def __init__(self, max_grad_norm: float):
         self.max_grad_norm = max_grad_norm
-    
+
     def on_step_begin(self, context: TrainContext):
         _ = context
         clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
 
 
+@CallbackFactory.register("scheduler")
 class SchedulerCallback(TrainCallback):
     """
     Scheduler callback for trainer.
     """
+
     def __init__(self):
-        self.scheduler: LRScheduler = None
+        pass
-    
+
     def on_train_begin(self, context: TrainContext):
         for group in context.optimizer.param_groups:
             if "initial_lr" not in group:
-                group["initial_lr"] = group["lr"] 
+                group["initial_lr"] = group["lr"]
-        
+
-        self.scheduler = context.scheduler
-    
     def on_batch_end(self, context: TrainContext):
-        _ = context
+        if context.scheduler:
-        if self.scheduler:
+            context.scheduler.step()
-            self.scheduler.step()
 
 
+@CallbackFactory.register("checkpoint")
 class CheckpointCallback(TrainCallback):
-    """ 
+    """
     Checkpoint callback for trainer.
     """
+
     def __init__(
-        self, 
+        self,
-        save_dir: str,  
+        save_dir: str,
         interval: int,
         weight_only: bool = False,
-        state_dict_fn: Optional[Callable[[nn.Module], dict]] = None
+        state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
     ):
         self.save_dir = save_dir
         self.interval = interval
         self.weight_only = weight_only
         self.state_dict_fn = state_dict_fn
         self.last_ckpt_iter = 0
-    
+
     @only_on_rank(0)
     def _save_checkpoint(self, context: TrainContext):
-        save_path = os.path.join(self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}")
+        save_path = os.path.join(
-        state_dict = self.state_dict_fn(context.model) if self.state_dict_fn else context.model.state_dict()
+            self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
-        
+        )
+        state_dict = (
+            self.state_dict_fn(context.model)
+            if self.state_dict_fn
+            else context.model.state_dict()
+        )
+
         context.checkpoint = Checkpoint(
-            state_dict=state_dict,
+            state_dict=state_dict, epoch=context.epoch, iteration=context.iteration
-            epoch=context.epoch,
-            iteration=context.iteration
         )
 
         context.checkpoint.save(save_path)
         self.last_ckpt_iter = context.iteration
-    
+
     def on_batch_end(self, context: TrainContext):
         if context.iteration - self.last_ckpt_iter >= self.interval:
             self._save_checkpoint(context)
-    
+
     def on_train_end(self, context: TrainContext):
         if context.iteration != self.last_ckpt_iter:
             self._save_checkpoint(context)
-    
+
     def on_error(self, context: TrainContext):
         self._save_checkpoint(context)
 
 
+@CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):
-    """ 
+    """
     Progress bar callback for trainer.
     """
+
     def __init__(self, num_epoch: int):
         self.num_epoch = num_epoch
         self.progress_bar: tqdm = None
-    
+
     @only_on_rank(0)
     def on_epoch_begin(self, context: TrainContext):
         self.progress_bar = tqdm(
-            context.dataloader, 
+            context.dataloader,
-            desc=f"Epoch {context.epoch+1}/{self.num_epoch}", 
+            desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
-            dynamic_ncols=True
+            dynamic_ncols=True,
         )
-    
+
     @only_on_rank(0)
     def on_batch_end(self, context: TrainContext):
-        self.progress_bar.set_postfix({
+        self.progress_bar.set_postfix(
-            "loss": f"{context.loss:.4f}",
+            {
-            "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}"
+                "loss": f"{context.loss:.4f}",
-        })
+                "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
+            }
+        )
         self.progress_bar.update(1)
-    
+
     @only_on_rank(0)
     def on_epoch_end(self, context: TrainContext):
         _ = context
@@ -163,70 +193,68 @@ class ProgressBarCallback(TrainCallback):
             self.progress_bar.close()
 
 
+@CallbackFactory.register("metric_logger")
 class MetricLoggerCallback(TrainCallback):
     def __init__(
-        self, 
+        self,
-        log_dir:str,
+        log_dir: str,
-        save_interval:int,
+        save_interval: int,
-        log_interval:int=10,
+        log_interval: int = 10,
-        metrics:List[str]=None
+        metrics: List[str] = None,
     ):
-        self.step_num = 0
+        self.last_log_iter = 0
-        self.last_save_step = 0
         self.save_interval = save_interval
         self.log_interval = log_interval
-        self.metrics = metrics or ['loss', 'lr']
+        self.metrics = metrics or ["loss", "lr"]
-        
+
         self.log_dir = Path(log_dir) if log_dir else Path.cwd() / "logs"
         self.log_dir.mkdir(parents=True, exist_ok=True)
-        
+
         self.log_cache = []
-        
+
         self._metric_funcs = {
-            'loss': ctx_get_loss,
+            "loss": ctx_get_loss,
-            'lr': ctx_get_lr,
+            "lr": ctx_get_lr,
-            'grad_norm': ctx_get_grad_norm,
+            "grad_norm": ctx_get_grad_norm,
-            'grad_std': ctx_get_grad_std,
+            "grad_std": ctx_get_grad_std,
-            'grad_max': ctx_get_grad_max,
+            "grad_max": ctx_get_grad_max,
-            'grad_min': ctx_get_grad_min,
+            "grad_min": ctx_get_grad_min,
-            'grad_mean': ctx_get_grad_mean,
+            "grad_mean": ctx_get_grad_mean,
-            'grad_nan_num': ctx_get_grad_nan_num
+            "grad_nan_num": ctx_get_grad_nan_num,
         }
 
     def _get_log_data(self, context: TrainContext):
         return {
-            "timestamp": time.strftime('%Y-%m-%d %H:%M:%S'),
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
             "epoch": context.epoch,
             "iter": context.iteration,
-            **{m: self._metric_funcs[m](context) for m in self.metrics}
+            **{m: self._metric_funcs[m](context) for m in self.metrics},
         }
-    
+
     @only_on_rank(0)
     def _add_log(self, log_data):
         self.log_cache.append(log_data)
-    
+
     @only_on_rank(0)
     def _save_log(self, epoch, iter):
         log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
-        
+
-        with open(log_file, 'w') as f:
+        with open(log_file, "w") as f:
             for log in self.log_cache:
-                f.write(json.dumps(log) + '\n')
+                f.write(json.dumps(log) + "\n")
-    
+
     def on_batch_end(self, context):
-        if self.step_num % self.log_interval == 0:
+        if context.iteration % self.log_interval == 0:
             log_data = self._get_log_data(context)
             self._add_log(log_data)
-        
+
-        if self.step_num - self.last_save_step >= self.save_interval:
+        if context.iteration - self.last_log_iter >= self.save_interval:
             self._save_log(context.epoch, context.iteration)
-            self.last_save_step = self.step_num
+            self.last_log_iter = context.iteration
-        
+
-        self.step_num += 1
-    
     def on_train_end(self, context):
-        self._save_log(context.epoch, context.iteration)
+        if context.iteration != self.last_log_iter:
-    
+            self._save_log(context.epoch, context.iteration)
+
     def on_error(self, context):
         self._save_log(context.epoch, context.iteration)
-    

astrai/trainer/train_context.py

@@ -1,16 +1,16 @@
+from dataclasses import dataclass, field
+from typing import Optional, Self
+
 import torch.nn as nn
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 
-from khaosz.data import ResumableDistributedSampler
+from astrai.config.train_config import TrainConfig
-from khaosz.data.checkpoint import Checkpoint
+from astrai.dataset import ResumableDistributedSampler
-from khaosz.trainer.strategy import StrategyFactory, BaseStrategy
+from astrai.parallel.setup import get_current_device, get_rank, get_world_size
-from khaosz.config.train_config import TrainConfig
+from astrai.serialization import Checkpoint
-from khaosz.parallel.setup import get_current_device, get_world_size, get_rank
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
-
-from dataclasses import dataclass, field
-from typing import Optional, Self
 
 
 @dataclass
@@ -21,11 +21,11 @@ class TrainContext:
     optimizer: Optimizer = field(default=None)
     scheduler: LRScheduler = field(default=None)
     checkpoint: Checkpoint = field(default=None)
-    
+
     epoch: int = field(default=0)
     iteration: int = field(default=0)
     loss: float = field(default=0.0)
-    
+
     world_size: int = field(default=1)
     rank: int = field(default=0)
     kwargs: dict = field(default_factory=dict)
@@ -39,17 +39,17 @@ class TrainContextBuilder:
             world_size=get_world_size(),
             rank=get_rank(),
         )
-        
+
         device = get_current_device()
         self._context.model = self._context.model.to(device=device)
-        
+
         if self.config.nprocs > 1:
             fn = self.config.parallel_wrapper
             self._context.model = fn(self._context.model)
-        
+
         self._context.optimizer = self.config.optimizer_fn(self._context.model)
         self._context.scheduler = self.config.scheduler_fn(self._context.optimizer)
-    
+
     def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
         if checkpoint is None:
             checkpoint = Checkpoint(
@@ -60,10 +60,10 @@ class TrainContextBuilder:
             self._context.epoch = max(checkpoint.epoch, self.config.start_epoch)
             self._context.iteration = max(checkpoint.iteration, self.config.start_batch)
             self._context.model.load_state_dict(checkpoint.state_dict)
-        
+
         self._context.checkpoint = checkpoint
         return self
-    
+
     def with_dataloader(self) -> Self:
         # fix: change batch level iteration to sample level offset
         config = self.config
@@ -72,28 +72,28 @@ class TrainContextBuilder:
             data_source=config.dataset,
             start_epoch=self._context.epoch,
             start_iter=sampler_offset,
-            seed=config.random_seed
+            seed=config.random_seed,
         )
-        
+
         dataloader = DataLoader(
-            config.dataset, 
+            config.dataset,
-            batch_size=config.batch_size, 
+            batch_size=config.batch_size,
             sampler=resumeable_sampler,
             num_workers=config.num_workers,
             pin_memory=config.pin_memory,
-            prefetch_factor=config.prefetch_factor
+            prefetch_factor=config.prefetch_factor,
         )
         self._context.dataloader = dataloader
         return self
-    
+
     def with_strategy(self) -> Self:
-        self._context.strategy = StrategyFactory.load(
+        self._context.strategy = StrategyFactory.create(
-            model=self.config.model,
+            model=self._context.model,
             train_type=self.config.strategy,
             device=get_current_device(),
-            **self.config.extra_kwargs
+            **self.config.extra_kwargs,
         )
         return self
-    
+
     def build(self) -> TrainContext:
-        return self._context
+        return self._context

astrai/trainer/trainer.py

@@ -0,0 +1,102 @@
+import logging
+from typing import List, Optional
+
+from astrai.config import TrainConfig
+from astrai.parallel.setup import spawn_parallel_fn
+from astrai.serialization import Checkpoint
+from astrai.trainer.train_callback import (
+    CallbackFactory,
+    TrainCallback,
+)
+from astrai.trainer.train_context import TrainContext, TrainContextBuilder
+
+logger = logging.getLogger(__name__)
+
+
+class Trainer:
+    def __init__(
+        self, train_config: TrainConfig, callbacks: Optional[List[TrainCallback]] = None
+    ):
+        self.train_config = train_config
+        default_callbacks = self._get_default_callbacks()
+        self.callbacks = (
+            default_callbacks + callbacks if callbacks else default_callbacks
+        )
+
+    def _get_default_callbacks(self) -> List[TrainCallback]:
+        cfg = self.train_config
+        return [
+            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("scheduler"),
+        ]
+
+    def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
+        return (
+            TrainContextBuilder(self.train_config)
+            .with_checkpoint(checkpoint)
+            .with_dataloader()
+            .with_strategy()
+            .build()
+        )
+
+    def _call_callbacks(self, method_name: str, context: TrainContext):
+        for callback in self.callbacks:
+            method = getattr(callback, method_name, None)
+            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,
+            device_ids=config.device_ids,
+            checkpoint=checkpoint,
+        )
+
+    def _train_impl(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
+        context = self._build_context(checkpoint)
+        self._call_callbacks("on_train_begin", context)
+
+        try:
+            context.model.train()
+            # 1.epoch
+            for epoch in range(context.epoch, self.train_config.n_epoch):
+                context.epoch = epoch
+                self._call_callbacks("on_epoch_begin", context)
+
+                for batch in context.dataloader:
+                    if context.iteration % self.train_config.accumulation_steps == 0:
+                        # 2. step
+                        self._call_callbacks("on_step_begin", context)
+                        context.optimizer.step()
+                        context.optimizer.zero_grad()
+                        self._call_callbacks("on_step_end", context)
+
+                    # 3. batch
+                    self._call_callbacks("on_batch_begin", context)
+                    loss = context.strategy(batch)
+                    context.loss = loss.item()
+                    context.iteration += 1
+
+                    # to make the loss normalized by accumulation steps
+                    stand_loss = loss / self.train_config.accumulation_steps
+                    stand_loss.backward()
+
+                    self._call_callbacks("on_batch_end", context)
+
+                self._call_callbacks("on_epoch_end", context)
+
+        except Exception as e:
+            logger.error(f"Training failed: {str(e)}", exc_info=True)
+            self._call_callbacks("on_error", context)
+            raise
+        finally:
+            self._call_callbacks("on_train_end", context)

demo/download.py

@@ -1,14 +0,0 @@
-import os
-from huggingface_hub import snapshot_download
-
-
-PROJECT_ROOT = os.path.dirname(
-    os.path.dirname(os.path.abspath(__file__)))
-
-
-if __name__ == "__main__":
-    snapshot_download(
-        repo_id="ViperEk/KHAOSZ",
-        local_dir=os.path.join(PROJECT_ROOT, "params"),  
-        force_download=True
-    )

demo/generate_ar.py

@@ -1,27 +0,0 @@
-import os
-import torch
-from khaosz import Khaosz
-
-
-PROJECT_ROOT = os.path.dirname(
-    os.path.dirname(os.path.abspath(__file__)))
-
-def generate_text():
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-
-    query = input(">> ")
-    
-    response = model.text_generate(
-        query=query, 
-        temperature=0.8,
-        top_p=0.95,
-        top_k=50
-    )
-    
-    print(response)
-
-    
-
-if __name__ == "__main__":
-    generate_text()

demo/generate_batch.py

@@ -1,25 +0,0 @@
-import os
-import torch
-from khaosz import Khaosz
-
-
-PROJECT_ROOT = os.path.dirname(
-    os.path.dirname(os.path.abspath(__file__)))
-
-def batch_generate():
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-    inputs = ["你好", "请问什么是人工智能", "今天天气如何", "我感到焦虑， 请问我应该怎么办", "请问什么是显卡"]
-    
-    responses = model.batch_generate(
-        queries=inputs,
-        temperature=0.8,
-        top_p=0.95,
-        top_k=50
-    )
-    
-    for q, r in zip(inputs, responses):
-        print((q, r))
-
-if __name__ == "__main__":
-    batch_generate()

demo/generate_retrieve.py

@@ -1,42 +0,0 @@
-import os
-import torch
-from khaosz import Khaosz, SemanticTextSplitter, Retriever
-
-
-PROJECT_ROOT = os.path.dirname(
-    os.path.dirname(os.path.abspath(__file__)))
-
-if __name__ == "__main__":
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    context_path = os.path.join(PROJECT_ROOT, "README.md")
-    
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-    spliter = SemanticTextSplitter(model.encode)
-    retriever = Retriever()
-    text = open(context_path, "r", encoding="utf-8").read()
-    
-    res = spliter.split(text, threshold=0.8, window_size=1)
-    # print(("\n" + "+"*100 + "\n").join(res))
-
-    res_embs = model.encode(res)
-    for sentence, emb in zip(res, res_embs):
-        retriever.add_vector(sentence, emb)
-
-    retrive_top_k = 5
-    query = "作者设计了一个怎样的模型"
-    emb_query = model.encode(query)
-    retrieved = retriever.retrieve(emb_query, retrive_top_k)
-    
-    retrive_response = model.retrieve_generate(
-        retrieved=retrieved,
-        query=query,
-        temperature=0.8,
-        top_p=0.95,
-        top_k=50
-    )
-
-    print("retrieve content:")
-    print("\n".join([f"{idx + 1}. " + text for idx, (text, _) in enumerate(retrieved)]))
-
-    print("\n\nretrive generate:")
-    print(retrive_response)

demo/stream_chat.py

@@ -1,32 +0,0 @@
-import os
-import torch
-from khaosz import Khaosz
-
-
-PROJECT_ROOT = os.path.dirname(
-    os.path.dirname(os.path.abspath(__file__)))
-
-def chat():
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-
-    history = []
-    while True:
-        query = input(">> ")
-        if query == "!exit":
-            break
-        
-        response_size = 0
-        for response, history in model.stream_generate(
-            query=query, 
-            history=history,
-            temperature=0.8,
-            top_p=0.95,
-            top_k=50
-        ):
-            print(response[response_size:], end="", flush=True)
-            response_size = len(response)
-    
-
-if __name__ == "__main__":
-    chat()

khaosz/__init__.py

@@ -1,59 +0,0 @@
-__version__ = "1.3.2"
-__author__ = "ViperEkura"
-
-from khaosz.api import Khaosz
-from khaosz.config import (
-    ModelConfig,
-    TrainConfig,
-)
-from khaosz.model.transformer import Transformer
-from khaosz.utils.retriever import Retriever
-from khaosz.utils.splitter import (
-    SemanticTextSplitter, 
-    PriorityTextSplitter
-)
-from khaosz.data import (
-    DatasetLoader,
-    BpeTokenizer
-)
-from khaosz.inference.generator import (
-    TextGenerator,
-    ChatGenerator, 
-    StreamGenerator, 
-    BatchGenerator, 
-    RetrievalGenerator, 
-    EmbeddingEncoder
-)
-
-from khaosz.trainer import (
-    Trainer,
-    StrategyFactory,
-    SchedulerFactory
-)
-
-__all__ = [
-    "Khaosz",
-    
-    "Transformer",
-    
-    "Retriever",
-    "SemanticTextSplitter",
-    "PriorityTextSplitter",
-    
-    "ModelConfig",
-    "TrainConfig",
-    
-    "DatasetLoader",
-    "BpeTokenizer",
-    
-    "TextGenerator",
-    "ChatGenerator",
-    "StreamGenerator",
-    "BatchGenerator",
-    "RetrievalGenerator",
-    "EmbeddingEncoder",
-    
-    "Trainer",
-    "StrategyFactory",
-    "SchedulerFactory"
-]

khaosz/api.py

@@ -1,113 +0,0 @@
-from torch import Tensor
-from typing import List, Tuple, Generator, Union
-
-from khaosz.inference.generator import (
-    TextGenerator,
-    ChatGenerator, 
-    StreamGenerator, 
-    BatchGenerator, 
-    RetrievalGenerator, 
-    EmbeddingEncoder
-)
-from khaosz.config.param_config import ModelParameter
-
-
-class Khaosz:
-    def __init__(self, model_dir: str):
-        self.parameter = ModelParameter()
-        self.parameter.load(model_dir)
-    
-    def to(self, *args, **kwargs):
-        self.parameter.to(*args, **kwargs)
-        return self
-
-    def generate(
-            self, 
-            query: str, 
-            history: List[Tuple[str, str]]=None,
-            temperature: float=0.8,
-            top_k: int=50,
-            top_p: float=0.95,
-        ) -> str:
-            generator = ChatGenerator(self.parameter)
-            return generator.generate(
-                query, 
-                history=history,
-                temperature=temperature, 
-                top_k=top_k, 
-                top_p=top_p,
-            )
-    
-    def batch_generate(
-            self, 
-            queries: List[str],
-            histories: List[Tuple[str, str]]=None,
-            temperature: float=0.8,
-            top_k: int=50,
-            top_p: float=0.95,
-        ) -> List[str]:
-            generator = BatchGenerator(self.parameter)
-            return generator.generate(
-                queries, 
-                histories=histories,
-                temperature=temperature, 
-                top_k=top_k, 
-                top_p=top_p,
-            )
-        
-    
-    def stream_generate(
-            self, 
-            query: str, 
-            history: List[Tuple[str, str]]=None,
-            temperature: float=0.8,
-            top_k: int=50,
-            top_p: float=0.95,
-        ) -> Generator[Tuple[str, List[Tuple[str, str]]], None, None]:
-            stream_generator = StreamGenerator(self.parameter)
-            return stream_generator.generate(
-                    query, 
-                    history=history,
-                    temperature=temperature, 
-                    top_k=top_k, 
-                    top_p=top_p,
-                )
-    
-    def retrieve_generate(
-            self,
-            retrieved,
-            query: str, 
-            history: List[Tuple[str, str]] = None,
-            temperature: float=0.8,
-            top_k: int=50,
-            top_p: float=0.95,
-        ) -> str:
-            generator = RetrievalGenerator(self.parameter)
-            return generator.generate(
-                retrieved,
-                query,
-                history=history,
-                temperature=temperature, 
-                top_k=top_k, 
-                top_p=top_p,
-            )
-    
-    def text_generate(            
-            self, 
-            query: str, 
-            temperature: float=0.8,
-            top_k: int=50,
-            top_p: float=0.95,
-        ) -> str:
-        generator = TextGenerator(self.parameter)
-        
-        return generator.generate(
-                query,
-                temperature=temperature, 
-                top_k=top_k, 
-                top_p=top_p,
-            )
-    
-    def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
-        encoder = EmbeddingEncoder(self.parameter)
-        return encoder.encode(sentence)

khaosz/config/__init__.py

@@ -1,16 +0,0 @@
-from khaosz.config.model_config import ModelConfig
-from khaosz.config.param_config import BaseModelIO, ModelParameter
-from khaosz.config.schedule_config import ScheduleConfig, CosineScheduleConfig, SGDRScheduleConfig
-from khaosz.config.train_config import TrainConfig
-
-
-__all__ = [
-    "BaseModelIO",
-    "ModelParameter",
-    "ModelConfig",
-    "TrainConfig",
-    
-    "ScheduleConfig",
-    "CosineScheduleConfig",
-    "SGDRScheduleConfig",
-]

khaosz/config/param_config.py

@@ -1,80 +0,0 @@
-import torch.nn as nn
-import safetensors.torch as st
-
-from dataclasses import dataclass, field
-from typing import Optional, Self, Union
-from pathlib import Path
-
-from khaosz.data.tokenizer import BpeTokenizer
-from khaosz.config.model_config import ModelConfig
-from khaosz.model.transformer import Transformer
-
-@dataclass
-class BaseModelIO:
-    """Base class for model I/O operations."""
-    
-    model: Optional[nn.Module] = field(
-        default=None,
-        metadata={"help": "Transformer model."}
-    )
-    tokenizer: BpeTokenizer = field(
-        default_factory=BpeTokenizer,
-        metadata={"help": "Tokenizer for the model."}
-    )
-    config: ModelConfig = field(
-        default_factory=ModelConfig,
-        metadata={"help": "Transformer model configuration."}
-    )
-    
-    def _get_file_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
-        """Get standardized file paths for model components."""
-        dir_path = Path(directory)
-        return {
-            "model": dir_path / "model.safetensors",
-            "config": dir_path / "config.json", 
-            "tokenizer": dir_path / "tokenizer.json"
-        }
-    
-    def save_components(self, save_dir: Union[str, Path]):
-        """Save core model components."""
-        paths = self._get_file_paths(save_dir)
-        paths["model"].parent.mkdir(parents=True, exist_ok=True)
-        
-        if self.model is not None:
-            st.save_file(self.model.state_dict(), str(paths["model"]))
-        self.config.save(str(paths["config"]))
-        self.tokenizer.save(str(paths["tokenizer"]))
-    
-    def load_components(self, load_dir: Union[str, Path]) -> Self:
-        """Load core model components."""
-        paths = self._get_file_paths(load_dir)
-        
-        self.config.load(str(paths["config"]))
-        self.tokenizer.load(str(paths["tokenizer"]))
-        
-        if self.model is None:
-            self.model = Transformer(self.config)
-        
-        if paths["model"].exists():
-            state_dict = st.load_file(str(paths["model"]))
-            self.model.load_state_dict(state_dict)
-        
-        return self
-    
-    def to(self, *args, **kwargs) -> "BaseModelIO":
-        """Move model to device."""
-        if self.model is not None:
-            self.model.to(*args, **kwargs)
-        return self
-
-
-@dataclass
-class ModelParameter(BaseModelIO):
-    """Container for model parameters with serialization capabilities."""
-    
-    def save(self, save_dir: Union[str, Path]):
-        self.save_components(save_dir)
-    
-    def load(self, load_dir: Union[str, Path]) -> "ModelParameter":
-        return self.load_components(load_dir)
-

khaosz/config/schedule_config.py

@@ -1,93 +0,0 @@
-from typing import Any, Dict
-from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
-
-
-@dataclass
-class ScheduleConfig(ABC):
-    schedule_type: str = field(
-        default="cosine",
-        metadata={
-            "help": "Type of learning rate schedule.", 
-            "choices": ["cosine", "sgdr"]
-        }
-    )
-    warmup_steps: int = field(
-        default=1000,
-        metadata={"help": "Number of warmup steps."}
-    )
-    min_rate: float = field(
-        default=0.05,
-        metadata={"help": "Minimum learning rate multiplier."}
-    )
-    
-    @abstractmethod
-    def get_kwargs(self) -> Dict[str, Any]:
-        raise NotImplementedError
-    
-    def validate(self) -> None:
-        """Validate configuration parameters."""
-        if self.warmup_steps < 0:
-            raise ValueError(f"warmup_steps must be non-negative, got {self.warmup_steps}")
-        if not 0 <= self.min_rate <= 1:
-            raise ValueError(f"min_rate must be between 0 and 1, got {self.min_rate}")
-
-
-@dataclass
-class CosineScheduleConfig(ScheduleConfig):
-    total_steps: int = field(
-        default=None,
-        metadata={"help": "Total training steps for cosine schedule."}
-    )
-    
-    def __post_init__(self) -> None:
-        self.schedule_type = "cosine"
-        self.validate()
-    
-    def get_kwargs(self) -> Dict[str, Any]:
-        if self.total_steps is None:
-            raise ValueError("total_steps must be specified for cosine schedule")
-            
-        return {
-            "schedule_type": self.schedule_type,
-            "warmup_steps": self.warmup_steps,
-            "lr_decay_steps": self.total_steps - self.warmup_steps,
-            "min_rate": self.min_rate
-        }
-    
-    def validate(self) -> None:
-        super().validate()
-        if self.total_steps is not None and self.total_steps <= self.warmup_steps:
-            raise ValueError(f"total_steps ({self.total_steps}) must be greater than warmup_steps ({self.warmup_steps})")
-
-
-@dataclass
-class SGDRScheduleConfig(ScheduleConfig):
-    cycle_length: int = field(
-        default=1000,
-        metadata={"help": "Length of the first cycle in steps."}
-    )
-    t_mult: int = field( 
-        default=2,
-        metadata={"help": "Multiplier for cycle length growth."}
-    )
-
-    def __post_init__(self) -> None:
-        self.schedule_type = "sgdr"
-        self.validate()
-
-    def get_kwargs(self) -> Dict[str, Any]:
-        return {
-            "schedule_type": self.schedule_type,
-            "warmup_steps": self.warmup_steps,
-            "cycle_length": self.cycle_length,
-            "min_rate": self.min_rate,
-            "t_mult": self.t_mult
-        }
-    
-    def validate(self) -> None:
-        super().validate()
-        if self.cycle_length <= 0:
-            raise ValueError(f"cycle_length must be positive, got {self.cycle_length}")
-        if self.t_mult < 1:
-            raise ValueError(f"t_mult must be >= 1, got {self.t_mult}")

khaosz/config/train_config.py

@@ -1,136 +0,0 @@
-import torch.nn as nn
-from torch.utils.data import Dataset
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LRScheduler
-
-from dataclasses import dataclass, field
-from typing import Callable, List, Optional
-
-
-@dataclass
-class TrainConfig:
-    # 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."}
-    )
-    optimizer_fn: Callable[[nn.Module], Optimizer] = field(
-        default=None,
-        metadata={"help": "Optimizer factory for training."}
-    )
-    scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
-        default=None,
-        metadata={"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(
-        default=1,
-        metadata={"help": "Number of iterations between steps."}
-    )
-    max_grad_norm: float = field(
-        default=1.0,
-        metadata={"help": "Maximum gradient norm."}
-    )
-    
-    # checkpoint setting
-    start_epoch: int = field(
-        default=0,
-        metadata={"help": "Start epoch for training."}
-    )
-    start_batch: int = field(
-        default=0,
-        metadata={"help": "Start batch iteration for training."}
-    )
-    checkpoint_dir: str = field(
-        default="./checkpoint",
-        metadata={"help": "Checkpoint directory."}
-    )
-    checkpoint_interval: int = field(
-        default=5000,
-        metadata={"help": "Number of iterations between checkpoints."}
-    )
-    
-    # dataloader setting
-    random_seed: int = field(
-        default=3407,
-        metadata={"help": "Random seed."}
-    )
-    num_workers: int = field(
-        default=0,
-        metadata={"help": "Number of workers for dataloader."}
-    )
-    prefetch_factor: Optional[int] = field(
-        default=None,
-        metadata={"help": "Prefetch factor for dataloader."}
-    )
-    pin_memory: bool = field(
-        default=False,
-        metadata={"help": "Pin memory for dataloader."}
-    )
-    
-    # distributed training
-    nprocs: int = field(
-        default=1,
-        metadata={"help": "Number of processes for distributed training."}
-    )
-    backend: str = field(
-        default="nccl",
-        metadata={"help": "Distributed training backend."}
-    )
-    master_addr: str = field(
-        default="localhost",
-        metadata={"help": "Master address for distributed training."}
-    )
-    master_port: str = field(
-        default="29500",
-        metadata={"help": "Master port for distributed training."}
-    )
-    parallel_wrapper: Optional[Callable] = field(
-        default=None,
-        metadata={"help": "Parallel function for training."}
-    )
-    state_dict_fn: Optional[Callable] = field(
-         default=None,
-         metadata={"help": "Parallel function for state  dict saving."}
-    )
-
-    # others
-    device_ids: Optional[List[int]] = field(
-        default=None,
-        metadata={"help": "Device ids for distributed training."}
-    )
-    device_type: str = field(
-        default="cuda",
-        metadata={"help": "Device type for distributed training."}
-    )
-    extra_kwargs: dict = field(
-        default_factory=dict,
-        metadata={"help": "Other arguments."}
-    )
-    
-    def __post_init__(self):
-        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.")
-        
-        

khaosz/data/__init__.py

@@ -1,24 +0,0 @@
-from khaosz.data.dataset import (
-    BaseDataset, 
-    SeqDataset, 
-    DpoDataset, 
-    SftDataset, 
-    PpoDataset, 
-    MultiSegmentFetcher,
-    DatasetLoader
-)
-
-from khaosz.data.tokenizer import BpeTokenizer
-from khaosz.data.sampler import ResumableDistributedSampler
-
-__all__ = [
-    "BaseDataset",
-    "SeqDataset",
-    "DpoDataset",
-    "SftDataset",
-    "PpoDataset",
-    "MultiSegmentFetcher",
-    "DatasetLoader",
-    "BpeTokenizer",
-    "ResumableDistributedSampler"
-]

khaosz/data/checkpoint.py

@@ -1,67 +0,0 @@
-import json
-import torch
-import torch.distributed as dist
-
-from pathlib import Path
-from typing import Dict, Any
-from khaosz.parallel.setup import get_rank
-
-
-class Checkpoint:
-    def __init__(
-        self,
-        state_dict: Dict[str, Any],
-        epoch: int = 0,
-        iteration: int = 0,
-    ):
-        self.state_dict = state_dict
-        self.epoch = epoch
-        self.iteration = iteration
-
-    def save(
-        self,
-        save_dir: str,
-    ) -> None:
-        
-        save_path = Path(save_dir)
-        save_path.mkdir(parents=True, exist_ok=True)
-
-        rank = get_rank()
-        if rank == 0:
-            meta = {
-                "epoch": self.epoch,
-                "iteration": self.iteration,
-            }
-            with open(save_path / "meta.json", "w") as f:
-                json.dump(meta, f, indent=2)
-
-            with open(save_path / f"state_dict.pt", "wb") as f:
-                torch.save(self.state_dict, f)
-
-    @classmethod
-    def load(
-        cls,
-        save_dir: str,
-    ) -> "Checkpoint":
-
-        rank = get_rank()
-        save_path = Path(save_dir)
-        
-        meta = {}
-        if rank == 0:
-            with open(Path(save_dir) / "meta.json", "r") as f:
-                meta = json.load(f)
-
-        if dist.is_initialized():
-            meta_list = [meta]
-            dist.broadcast_object_list(meta_list, src=0)
-            meta = meta_list[0]
-
-        with open(save_path / f"state_dict.pt", "rb") as f:
-            state_dict = torch.load(f)
-
-        return cls(
-            state_dict=state_dict,
-            epoch=meta["epoch"],
-            iteration=meta["iteration"],
-        )

khaosz/data/dataset.py

@@ -1,201 +0,0 @@
-import torch
-import bisect
-
-from abc import ABC, abstractmethod
-from torch import Tensor
-from torch.utils.data import Dataset
-from khaosz.data.file import load_h5
-from typing import Callable, List, Dict, Literal, Optional, Union
-
-
-
-class BaseSegmentFetcher:
-    def __init__(self, segments: List[Tensor]):
-        self.segments = segments
-        self.cum_lengths = []
-        
-        total = 0
-        for seg in segments:
-            total += torch.numel(seg)
-            self.cum_lengths.append(total)
-            
-        self.total_length = total
-        
-    def __len__(self) -> int:
-        return self.total_length
-
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
-        if not (0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length):
-            raise ValueError("begin_idx or end_idx out of bounds")
-        if begin_idx >= end_idx:
-            return torch.tensor([], dtype=torch.long)
-        
-        # fix the range index bug
-        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
-        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
-
-        result_segments = []
-
-        for i in range(seg_start_idx, seg_end_idx + 1):
-            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
-            start = max(begin_idx - prev_cum, 0)
-            end = min(end_idx - prev_cum, len(self.segments[i]))
-            data = self.segments[i][start:end]
-            result_segments.append(data)
-
-        return torch.cat(result_segments, dim=0)
-    
-
-class MultiSegmentFetcher:
-    def __init__(self, muti_segments: Dict):
-        self.muti_keys = list(muti_segments.keys())
-        self.muti_fetchers = {
-            key: BaseSegmentFetcher(segments)
-            for key, segments in muti_segments.items()
-        }
-    
-    def __len__(self) -> int:
-        len_list = [len(seg) for seg in self.muti_fetchers.values()]
-        return min(len_list)
-        
-    def key_fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]) -> Dict:
-        fetch_dict = {} 
-        keys = [keys] if isinstance(keys, str) else keys
-        
-        for key in keys:
-            fetcher = self.muti_fetchers[key]
-            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
-            fetch_dict[key] = fetch_tensor
-
-        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
-    
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
-        return self.key_fetch(begin_idx, end_idx, self.muti_keys)
-
-
-class BaseDataset(Dataset, ABC):
-    def __init__(self, window_size: int, stride: int):
-        super().__init__()
-        self.segments = {}
-        self.window_size = window_size
-        self.stride = stride
-        self.total_samples = None
-
-    def load(self, load_path: str):
-        self.segments = load_h5(load_path)
-        self.fetcher = MultiSegmentFetcher(self.segments)
-        self.total_samples = len(self.fetcher)
-        
-    def get_index(self, index: int) -> int:
-        assert self.total_samples > self.window_size
-        
-        begin_idx = min(index * self.stride, self.total_samples - 1 - self.window_size)
-        end_idx = min(begin_idx + self.window_size, self.total_samples - 1)
-        
-        return begin_idx, end_idx
-        
-    @abstractmethod
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        raise NotImplementedError
-        
-    def __len__(self) -> int:
-        assert self.total_samples is not None
-        if self.total_samples <= self.window_size:
-            return 0
-        return (self.total_samples - 1 - self.window_size) // self.stride + 1
-    
-
-class SeqDataset(BaseDataset):
-    def __init__(self, window_size: int, stride: int):
-        super().__init__(window_size, stride)
-        self.fetcher = MultiSegmentFetcher(self.segments)
-
-    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
-    
-    def __getitem__(self, index):
-        # fix the range index bug
-        begin_idx, end_idx = self.get_index(index)
-        
-        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
-        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
-        
-        return {"input_ids": x, "target_ids": y}
-    
-    
-class SftDataset(BaseDataset):
-    def __init__(self, window_size: int, stride: int):
-        super().__init__(window_size, stride)
-        self.fetcher = MultiSegmentFetcher(self.segments)
-
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, key)
-    
-    def __getitem__(self, index):
-        begin_idx, end_idx = self.get_index(index)
-        
-        x = self._fetch_data(begin_idx, end_idx, "sequence").to(dtype=torch.long)
-        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(dtype=torch.long)
-        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask").to(dtype=torch.bool)
-        
-        return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask}
-
-
-class DpoDataset(BaseDataset):
-    def __init__(self, window_size: int, stride: int):
-        super().__init__(window_size, stride)
-        self.fetcher = MultiSegmentFetcher(self.segments)
-
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, key)
-
-    def __getitem__(self, index: int):
-        begin_idx, end_idx = self.get_index(index)
-        
-        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
-        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
-        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(dtype=torch.bool)
-        rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(dtype=torch.bool)
-
-        return {"chosen": chosen, "rejected": rejected, "chosen_mask": chosen_mask, "rejected_mask": rejected_mask}
-
-
-class PpoDataset(BaseDataset):
-    def __init__(self, window_size: int, stride: int):
-        super().__init__(window_size, stride)
-        self.fetcher = MultiSegmentFetcher(self.segments)
-
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
-        return self.fetcher.key_fetch(begin_idx, end_idx, key)
-    
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        begin_idx, end_idx = self.get_index(index)
-
-        input_ids =  self._fetch_data(begin_idx, end_idx, "input_ids"),
-        actions = self._fetch_data(begin_idx, end_idx, "actions"),
-        logprobs = self._fetch_data(begin_idx, end_idx, "logprobs"),
-        rewards =  self._fetch_data(begin_idx, end_idx, "rewards")
-        
-        return {"input_ids": input_ids, "actions": actions, "logprobs": logprobs, "rewards": rewards}
-    
-
-class DatasetLoader:
-    @staticmethod       
-    def load(
-        train_type: Literal["seq", "sft", "dpo"],
-        load_path: str,
-        window_size: int, 
-        stride: Optional[int] = None,
-        ) -> BaseDataset:
-        if stride is None:
-            stride = window_size
-        
-        dataset_router: Dict[str, Callable[[int], BaseDataset]] = {
-            "seq": lambda window_size: SeqDataset(window_size, stride),
-            "sft": lambda window_size: SftDataset(window_size, stride),
-            "dpo": lambda window_size: DpoDataset(window_size, stride),
-        }
-        dataset = dataset_router[train_type](window_size)
-        dataset.load(load_path)
-        
-        return dataset

khaosz/data/file.py

@@ -1,42 +0,0 @@
-import os
-import h5py
-import torch
-
-from pathlib import Path
-from torch import Tensor
-from typing import Dict, List
-
-
-def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
-    os.makedirs(file_path, exist_ok=True)
-    full_file_path = os.path.join(file_path, f"{file_name}.h5")
-    with h5py.File(full_file_path, 'w') as f:
-        for key, tensors in tensor_group.items():
-            grp = f.create_group(key)
-            for idx, tensor in enumerate(tensors):
-                arr = tensor.cpu().numpy()
-                grp.create_dataset(f'data_{idx}', data=arr)
-
-def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
-    tensor_group: Dict[str, List[Tensor]] = {}
-
-    root_path = Path(file_path)
-    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
-    
-    for h5_file in h5_files:
-        with h5py.File(h5_file, 'r') as f:
-            for key in f.keys():
-                grp = f[key]
-                dsets = []
-                for dset_name in grp.keys():
-                    dset = grp[dset_name]
-                    tensor = torch.from_numpy(dset[:])
-                    if share_memory:
-                        tensor = tensor.share_memory_()
-                    dsets.append(tensor)
-            
-                if tensor_group.get(key) is None:
-                    tensor_group[key] = []
-                tensor_group[key].extend(dsets)
-
-    return tensor_group

khaosz/data/tokenizer.py

@@ -1,106 +0,0 @@
-from tokenizers import Tokenizer, Encoding
-from tokenizers import decoders, processors, normalizers, pre_tokenizers
-from tokenizers.models import BPE
-from tokenizers.trainers import BpeTrainer
-from typing import List, Union
-
-
-class BpeTokenizer:
-    def __init__(self, path=None):
-        self._control_tokens = ["<bos>", "<eos>", "<pad>"]
-        self._special_tokens = ["<|im_start|>", "<|im_end|>"]
-        
-        model = BPE()
-        self._tokenizer = Tokenizer(model)
-        self._tokenizer.normalizer = normalizers.Sequence([
-            normalizers.NFC(),
-            normalizers.Strip() 
-        ])
-        
-        self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
-            pre_tokenizers.UnicodeScripts(),
-            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)
-        ])
-        
-        self._tokenizer.decoder = decoders.ByteLevel()
-        self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
-        
-        if path is not None:
-            self._tokenizer = Tokenizer.from_file(path)
-        
-    def _prepare_trainer(self, vocab_size: int, min_freq: int, reserved_token_size: int, max_token_length=18) -> tuple:
-        assert reserved_token_size > len(self._special_tokens)
-        reserved_tokens = [f"<|reserve{i:02d}|>" for i in range(reserved_token_size - len(self._special_tokens))]
-        detail_vocab_size = vocab_size - (len(reserved_tokens) + len(self._special_tokens))
-        
-        alphabet = pre_tokenizers.ByteLevel.alphabet()
-        min_size = len(alphabet) + len(self._control_tokens)
-        assert detail_vocab_size > min_size
-        
-        trainer = BpeTrainer(
-            vocab_size=detail_vocab_size,
-            min_frequency=min_freq,
-            limit_alphabet=detail_vocab_size // 6,
-            max_token_length=max_token_length,
-            special_tokens=self._control_tokens,
-            initial_alphabet=alphabet,
-            show_progress=True,
-        )
-        
-        return trainer, detail_vocab_size, reserved_tokens
-
-    def train(self, files, vocab_size, min_freq, reserved_token_size=100):
-        trainer, _, reserved_tokens = self._prepare_trainer(
-            vocab_size=vocab_size,
-            min_freq=min_freq,
-            reserved_token_size=reserved_token_size
-        )
-        self._tokenizer.train(files=files, trainer=trainer)
-        self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
-            
-    def train_from_iterator(self, iterator, vocab_size, min_freq, reserved_token_size=100):
-        trainer, _, reserved_tokens = self._prepare_trainer(
-            vocab_size=vocab_size,
-            min_freq=min_freq,
-            reserved_token_size=reserved_token_size
-        )
-        self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer)
-        self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
-            
-    def save(self, path):
-        self._tokenizer.save(path)
-        
-    def load(self, path):
-        self._tokenizer = Tokenizer.from_file(path)
-
-    def encode(self, tokens: Union[str, List[str]], out_ids: bool=True, add_special_tokens: bool=False) -> List:
-        if isinstance(tokens, str):
-            encoded: Encoding = self._tokenizer.encode(tokens, add_special_tokens=add_special_tokens)
-            return encoded.ids if out_ids else encoded.tokens
-        elif isinstance(tokens, list):
-            encoded_list: List[Encoding] = self._tokenizer.encode_batch(tokens, add_special_tokens=add_special_tokens)
-            return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
-
-    def decode(self, tokens: List[int], skip_special_tokens: bool=True) -> str:
-        return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
-    
-    def __len__(self) -> int:
-        return self._tokenizer.get_vocab_size()
-    
-    @property
-    def stop_ids(self) -> List[int]:
-        stop_token = self._control_tokens + self._special_tokens
-        stop_ids = [self._tokenizer.token_to_id(token) for token in stop_token]
-        return stop_ids
-    
-    @property
-    def bos_id(self) -> int:
-        return self._tokenizer.token_to_id("<bos>")
-    
-    @property
-    def eos_id(self) -> int:
-        return self._tokenizer.token_to_id("<eos>")
-    
-    @property
-    def pad_id(self) -> int:
-        return self._tokenizer.token_to_id("<pad>")

khaosz/inference/__init__.py

@@ -1 +0,0 @@
-# init file

khaosz/inference/core.py

@@ -1,240 +0,0 @@
-import torch
-from torch import Tensor 
-from typing import Any, Callable, List, Tuple, Union, Optional, Self
-from khaosz.config import ModelParameter, ModelConfig
-
-
-def apply_sampling_strategies(
-    logits: Tensor,
-    temperature: float,
-    top_k: int,
-    top_p: float,
-    filter_value: float = -float("inf")
-) -> Tensor:
-    """ 
-    Apply sampling strategies to the logits tensor.
-    
-    Args:
-        logits (Tensor): The logits tensor.
-        temperature (float): The temperature parameter.
-        top_k (int): The top-k parameter.
-        top_p (float): The top-p parameter.
-        filter_value (float, optional): The filter value. Defaults to -float("inf").
-        
-    Returns:
-        Tensor: The sampled logits tensor.
-        
-    """
-    
-    if temperature != 1.0:
-        logits = logits / temperature
-    
-    if top_k > 0:
-        top_k = min(top_k, logits.size(-1))
-        indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
-        logits[indices_to_remove] = filter_value
-    
-    if top_p < 1.0:
-        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
-        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-        
-        sorted_indices_to_remove = cumulative_probs > top_p
-        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-        sorted_indices_to_remove[..., 0] = 0
-        
-        indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
-        indices_to_remove.scatter_(
-            dim=1,
-            index=sorted_indices,
-            src=sorted_indices_to_remove
-        )
-        
-        logits[indices_to_remove] = filter_value
-    
-    return logits
-
-
-class GeneratorCore:
-    def __init__(self, parameter: ModelParameter):
-        self.model = parameter.model
-        self.tokenizer = parameter.tokenizer
-        self.config = parameter.config
-        
-    def generate_iterator(
-        self,
-        input_ids: Tensor,
-        temperature: float,
-        top_k: int,
-        top_p: float,
-        attn_mask: Optional[Tensor] = None,
-        kv_caches: Optional[List[Tuple[Tensor, Tensor]]] = None,
-        start_pos: int = 0
-    )-> Tuple[Tensor, int]:
-        
-        with torch.inference_mode():
-            outputs = self.model(input_ids, attn_mask, kv_caches, start_pos)
-            logits = outputs["logits"][:, -1, :]
-            cache_increase = input_ids.size(-1)   
-        
-        logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
-        probs = torch.softmax(logits, dim=-1)
-        next_token_id = torch.multinomial(probs, num_samples=1)
-    
-        return next_token_id, cache_increase
-
-    def to(self, *args, **kargs) -> Self:
-        self.model.to(*args, **kargs)
-        return self
-
-    def generate_loop(
-        self,
-        input_ids: Tensor,
-        ids: List[int],
-        temperature: float,
-        top_k: int,
-        top_p: float,
-        attn_mask: Optional[Tensor] = None,
-        kv_caches: Optional[List[Tuple[Tensor, Tensor]]] = None,
-        start_pos: int = 0,
-        callback: Optional[Callable[..., Any]] = None
-    ) -> List[int]:
-        cur_cache_pos = start_pos
-        
-        for _ in range(len(ids), self.config.max_len):
-            next_token_id, cache_increase = self.generate_iterator(
-                input_ids, temperature, top_k, top_p, attn_mask, kv_caches, cur_cache_pos)
-            
-            input_ids = next_token_id
-            ids.append(next_token_id.item())
-            cur_cache_pos += cache_increase
-
-            if callback:
-                callback(next_token_id.item(), ids.copy())
-
-            if next_token_id.item() in self.tokenizer.stop_ids:
-                break
-        
-        return ids
-
-
-class EmbeddingEncoderCore:
-    def __init__(self, parameter: ModelParameter):
-        self.model = parameter.model
-        self.tokenizer = parameter.tokenizer
-        self.config = parameter.config
-    
-    def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
-        with_batch = isinstance(sentence, list)
-        ids = self.tokenizer.encode(sentence)
-        batch_ids = ids if with_batch else [ids]
-        max_model_len = self.config.max_len
-        
-        all_fragments = []
-        fragment_origin_idx = []
-        
-        for i, seq in enumerate(batch_ids):
-            if len(seq) > max_model_len:
-                fragments = [seq[j:j+max_model_len] for j in range(0, len(seq), max_model_len)]
-                all_fragments.extend(fragments)
-                fragment_origin_idx.extend([i] * len(fragments))
-            else:
-                all_fragments.append(seq)
-                fragment_origin_idx.append(i)
-        
-        #if empty fragments
-        if not all_fragments or not ids:
-            return [] if with_batch else torch.tensor([])
-        
-        device = next(self.model.parameters()).device
-        max_len = min(max(len(seq) for seq in all_fragments), max_model_len)
-        
-        padded_ids = []
-        masks = []
-        for seq in all_fragments:
-            pad_len = max_len - len(seq)
-            padded_seq = seq + [self.tokenizer.pad_id] * pad_len
-            mask = [token_id != self.tokenizer.pad_id for token_id in padded_seq]
-            padded_ids.append(padded_seq)
-            masks.append(mask)
-        
-        input_tensor = torch.tensor(padded_ids, device=device, dtype=torch.long)
-        seq_mask = torch.tensor(masks, device=device, dtype=torch.bool)
-        
-        with torch.inference_mode():
-            outputs = self.model(input_tensor, seq_mask)["hidden_states"]
-            # [num_fragments, seq_len, hidden_size]
-            fragment_embs = torch.mul(outputs, seq_mask.unsqueeze(-1))  
-        
-        sentence_embs: List[Tensor] = []
-        for i in range(len(batch_ids)):
-            indices = [idx for idx, orig_idx in enumerate(fragment_origin_idx) if orig_idx == i]
-            if indices is not None:
-                sum_frags = torch.sum(fragment_embs[indices, :, :], dim=1)      # [frags, hidden_size]
-                length = torch.sum(seq_mask[indices, :], dim=1).unsqueeze(1)    # [frags, 1]
-                emb = torch.sum(sum_frags / length, dim=0)                      # [frags, hidden_size]
-                sentence_embs.append(emb.flatten())
-        
-        if with_batch:
-            return [emb.flatten() for emb in sentence_embs]
-        else:
-            return sentence_embs[0].flatten()
-
-    def to(self, *args, **kargs) -> Self:
-        self.model.to(*args, **kargs)
-        return self
-
-
-class KVCacheManager:
-    def __init__(
-        self, 
-        config: ModelConfig, 
-        batch_size: int,
-        device: torch.device = "cuda", 
-        dtype: torch.dtype = torch.bfloat16
-    ):
-        self.batch_size = batch_size
-        self.device = device
-        self.dtype = dtype
-        self.num_layers = config.n_layers
-        self.max_len = config.max_len
-        self.num_heads = config.n_kv_heads
-        self.head_dim = config.dim //config.n_heads
-        
-        self._kv_cache: Tuple[Tensor, Tensor] = None
-        self._seq_mask: Tensor = None
-        self._initialize()
-
-    def _initialize(self):
-        k_cache = torch.zeros(
-            (self.batch_size, self.max_len, self.num_layers, self.num_heads, self.head_dim),
-            device=self.device, dtype=self.dtype
-        )
-        v_cache = torch.zeros(
-            (self.batch_size, self.max_len, self.num_layers, self.num_heads, self.head_dim),
-            device=self.device, dtype=self.dtype
-        )
-        self._kv_cache = (k_cache, v_cache)
-        self._seq_mask = torch.ones((self.batch_size, self.max_len), device=self.device, dtype=torch.bool)
-
-    def update(self, active_mask: Tensor):        
-        k_cache, v_cache = self._kv_cache
-        self._kv_cache = (k_cache[active_mask], v_cache[active_mask])
-        self._seq_mask = self._seq_mask[active_mask]
-
-    def reset(self, full_reset=False):
-        if full_reset:
-            self._kv_cache = None
-            self._seq_mask = None
-        else:
-            self._initialize()
-    
-    def set_seq_mask(self, input_ids: Tensor, pad_id: int):
-        batch_size, seq_len = input_ids.shape
-        bool_mask = (input_ids != pad_id)
-        self._seq_mask[: batch_size, : seq_len] = bool_mask
-
-    def get_kvcache(self) -> Tuple[Tensor, Tensor]:
-        return self._kv_cache
-    
-    def get_seq_mask(self) -> Tensor:        
-        return self._seq_mask

khaosz/inference/cuda_graph.py

@@ -1,98 +0,0 @@
-import torch
-from torch import Tensor
-from functools import wraps
-from inspect import signature
-
-
-class CudaGraphWrapper:
-    def __init__(self, function, device="cuda", cast=False):
-        self.function = function
-        self.cast = cast
-        self.device = device
-        self.static_input = None
-        self.static_output = None
-        self.graph = None
-        self.signature = signature(function)
-        
-    def _update_inplace(self, lhs, rhs):
-        if isinstance(lhs, Tensor) and isinstance(rhs, Tensor):
-            if lhs.shape != rhs.shape:
-                raise ValueError(
-                    f"Tensor shape mismatch! "
-                    f"Expected: {lhs.shape}, Got: {rhs.shape}. "
-                    f"Function: {self.function}"
-                )
-            if self.cast:
-                if lhs.device != rhs.device:
-                    rhs = rhs.to(device=lhs.device)
-                
-                if lhs.dtype != rhs.dtype:
-                    rhs = rhs.to(dtype=lhs.dtype)
-            else: 
-                if lhs.device != rhs.device:
-                    raise ValueError(
-                        f"Tensor device mismatch! "
-                        f"Expected: {lhs.device}, Got: {rhs.device}. "
-                        f"Function: {self.function}"
-                    )
-                if lhs.dtype != rhs.dtype:
-                    raise ValueError(
-                        f"Tensor dtype mismatch! "
-                        f"Expected: {lhs.dtype}, Got: {rhs.dtype}. "
-                        f"Function: {self.function}"
-                    )
-                lhs.copy_(rhs)
-        elif isinstance(lhs, dict):
-            for k in lhs:
-                if k in rhs:
-                    self._update_inplace(lhs[k], rhs[k])
-        elif isinstance(lhs, (list, tuple)):
-            for i in range(len(lhs)):
-                if i < len(rhs):
-                    self._update_inplace(lhs[i], rhs[i])
-        elif isinstance(lhs, (int, float, bool, str, type(None))):
-            if lhs != rhs:
-                raise ValueError("Does not support changing control parameters.")   
-
-    def _update_args(self, input_args, input_kwargs):
-        bound_args = self.signature.bind(*input_args, **input_kwargs)
-        bound_args.apply_defaults()
-        args_dict = bound_args.arguments
-        
-        if self.static_input is None:
-            self.static_input = args_dict
-        else:
-            self._update_inplace(self.static_input, args_dict)
-    
-    def run(self, *args, **kwargs):
-        self._update_args(args, kwargs)
-        
-        if self.graph is None:
-            # warmup
-            _ = torch.matmul(
-                torch.randn(100, 100, device=self.device),
-                torch.randn(100, 100, device=self.device)
-            )
-            torch.cuda.synchronize()
-            
-            # capture graph
-            self.graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(self.graph):
-                self.static_output = self.function(**self.static_input)
-        
-        self.graph.replay()
-        
-        return self.static_output
-
-
-def cuda_graph(device="cuda", cast=False):
-    def decorator(func):
-        wrapper = CudaGraphWrapper(func, device, cast)
-        
-        @wraps(func)
-        def wrapped(*args, **kwargs):
-            return wrapper.run(*args, **kwargs)
-        
-        return wrapped 
-    
-    return decorator

khaosz/inference/generator.py

@@ -1,296 +0,0 @@
-import torch
-from torch import Tensor 
-from typing import List, Tuple, Union, Optional, Generator
-from khaosz.inference.core import GeneratorCore, EmbeddingEncoderCore, KVCacheManager
-from khaosz.config.param_config import ModelParameter
-
-
-def build_prompt(
-    query: str, 
-    init_prompt: Optional[str] = None,
-    history: Optional[List[Tuple[str, str]]] = None
-    ) -> str:
-    """ 
-    Build prompt in ChatML format for query and history
-    
-    Args:
-        query(str): query string
-        history(Optional[List[Tuple[str, str]]]): history list of query and response
-        
-    Returns:
-        str: prompt string in ChatML format
-        
-    """
-    prompt = f"<|im_start|>system\n{init_prompt}<|im_end|>\n" if init_prompt else ""
-    
-    # (convert tuple format to ChatML)
-    if history:
-        for user_msg, assistant_msg in history:
-            prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
-            prompt += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
-    
-    prompt += f"<|im_start|>user\n{query}<|im_end|>\n"
-    prompt += "<|im_start|>assistant\n"
-    
-    return prompt
-
-def pad_sequence(ids_list: List[List[int]], max_ids_len: int, pad_id: int) -> List[List[int]]:
-    """ 
-    Pad a list of sequences to a fixed length.
-    
-    Args:
-        ids_list (List[List[int]]): A list of sequences.
-        max_ids_len (int): The maximum length of sequences.
-        pad_id (int): The id to pad sequences.
-        
-    Returns:
-        List[List[int]]: A list of padded sequences.
-        
-    """
-    new_ids_list = []
-    for ids in ids_list:
-        pad_len = max_ids_len - len(ids)
-        padded_seq = [pad_id] * pad_len + ids
-        new_ids_list.append(padded_seq)
-    
-    return new_ids_list
-
-
-class TextGenerator(GeneratorCore):
-    def __init__(self, parameter: ModelParameter):
-        super().__init__(parameter)
-    
-    def generate(
-            self, 
-            query: str, 
-            temperature: float,
-            top_k: int,
-            top_p: float,
-        ) -> str:
-        assert temperature >= 0.0
-        assert top_k >= 0
-        assert top_p >= 0.0 and top_p <= 1.0
-        
-        device = next(self.model.parameters()).device
-        cache_manager = KVCacheManager(self.config, 1, device=device)
-        
-        ids = self.tokenizer.encode(query)
-        input_ids = torch.tensor([ids], device=device, dtype=torch.long)
-        
-        start_cache_pos = len(ids)
-        cur_cache_pos = 0
-        self.model.eval()
-        kv_caches = cache_manager.get_kvcache()
-        
-        ids = self.generate_loop(
-            input_ids, ids, temperature, top_k, top_p, 
-            kv_caches=kv_caches, 
-            start_pos=cur_cache_pos
-        )
-        
-        response = self.tokenizer.decode(ids[start_cache_pos:])
-        
-        return response
-
-
-class ChatGenerator(GeneratorCore):
-    def __init__(self, parameter: ModelParameter):
-        super().__init__(parameter)
-        
-    def generate(
-            self, 
-            query: str, 
-            history: List[Tuple[str, str]],
-            temperature: float,
-            top_k: int,
-            top_p: float,
-        ) -> str:
-        
-        assert temperature >= 0.0
-        assert top_k >= 0
-        assert top_p >= 0.0 and top_p <= 1.0
-
-        if history is None:
-            history = []
-        
-        device = next(self.model.parameters()).device
-        cache_manager = KVCacheManager(self.config, 1, device=device)
-        
-        ids = self.tokenizer.encode(build_prompt(query, history))
-        input_ids = torch.tensor([ids], device=device, dtype=torch.long)
-        
-        start_cache_pos = len(ids)
-        cur_cache_pos = 0
-        self.model.eval()
-        kv_caches = cache_manager.get_kvcache()
-        
-        ids = self.generate_loop(
-            input_ids, ids, temperature, top_k, top_p, 
-            kv_caches=kv_caches, 
-            start_pos=cur_cache_pos
-        )
-        
-        response = self.tokenizer.decode(ids[start_cache_pos:])
-        
-        return response
-    
-
-class StreamGenerator(GeneratorCore):
-    def __init__(self, parameter: ModelParameter):
-        super().__init__(parameter)
-    
-    def generate(
-            self, 
-            query: str, 
-            history: List[Tuple[str, str]],
-            temperature: float,
-            top_k: int,
-            top_p: float,
-        ) -> Generator[Tuple[str, List[Tuple[str, str]]], None, None]:
-        
-        assert temperature >= 0.0
-        assert top_k >= 0
-        assert top_p >= 0.0 and top_p <= 1.0
-
-        if history is None:
-            history = []
-        
-        device = next(self.model.parameters()).device
-        cache_manager = KVCacheManager(self.config, 1, device=device)
-                
-        ids = self.tokenizer.encode(build_prompt(query, history))
-        input_ids = torch.tensor([ids], device=device, dtype=torch.long)
-        cpy_history = history.copy()
-        
-        start_cache_pos = len(ids)
-        cur_cache_pos = 0
-        self.model.eval()
-        kv_caches = cache_manager.get_kvcache()
-        
-        for _ in range(len(ids), self.config.max_len):
-            next_token_id, cache_increase = self.generate_iterator(
-                input_ids, temperature, top_k, top_p, kv_caches=kv_caches, start_pos=cur_cache_pos)
-            
-            input_ids = next_token_id
-            ids.append(next_token_id.item())
-            cur_cache_pos += cache_increase
-            
-            response = self.tokenizer.decode(ids[start_cache_pos:])
-            yield response, cpy_history + [(query, response)]
-            
-            if next_token_id.item() in self.tokenizer.stop_ids:
-                yield response + "\n", cpy_history + [(query, response)]
-                break
-    
-
-class BatchGenerator(GeneratorCore):
-    def __init__(self, parameter: ModelParameter):
-        super().__init__(parameter)
-    
-    def generate(
-        self, 
-        queries: List[str],
-        histories: List[List[Tuple[str, str]]],
-        temperature: float, 
-        top_k: int, 
-        top_p: float 
-    ) -> List[str]:
-        
-        assert temperature >= 0.0
-        assert top_k >= 0
-        assert top_p >= 0.0 and top_p <= 1.0
-        
-        batch_size = len(queries)
-        if histories is None:
-            histories = [[] for _ in range(batch_size)]
-        
-        prompts = [build_prompt(query, history) for query, history in zip(queries, histories)]
-        ids_list = [self.tokenizer.encode(prompt) for prompt in prompts]
-        max_ids_len = max(len(ids) for ids in ids_list)
-        ids_list = pad_sequence(ids_list, max_ids_len, self.tokenizer.pad_id)
-        
-        device = next(self.model.parameters()).device
-        cache_manager = KVCacheManager(self.config, batch_size, device=device)
-        
-        input_tensor = torch.tensor(ids_list, device=device, dtype=torch.long)
-        cache_manager.set_seq_mask(input_tensor, self.tokenizer.pad_id)
-        activate_task_mask = [True] * batch_size
-        
-        start_cache_pos = max_ids_len
-        cur_cache_pos = 0
-        
-        while max_ids_len < self.config.max_len and sum(activate_task_mask) != 0:
-            kv_caches = cache_manager.get_kvcache()
-            attn_mask =cache_manager.get_seq_mask()
-            
-            next_token_id, cache_increase = self.generate_iterator(
-                input_tensor, temperature, top_k, top_p, attn_mask=attn_mask, kv_caches=kv_caches, start_pos=cur_cache_pos)
-            
-            cur_cache_pos += cache_increase
-            active_mask = []
-            c_ids = 0
-            
-            for i in range(batch_size):
-                if activate_task_mask[i]:
-                    token = next_token_id[c_ids, :].item()
-                    ids_list[i].append(token)
-                    c_ids += 1
-                    
-                    is_active = not token in self.tokenizer.stop_ids
-                    activate_task_mask[i] = is_active 
-                    active_mask.append(is_active)
-            
-            active_mask = torch.tensor(active_mask, device=device, dtype=torch.bool)
-            cache_manager.update(active_mask)
-            input_tensor = next_token_id[active_mask, :]
-
-            max_ids_len += 1
-        
-        
-        responses = [str()] * batch_size
-        for i in range(batch_size):
-            responses[i] = self.tokenizer.decode(ids_list[i][start_cache_pos:])
-            histories[i].append((queries[i], responses[i]))
-        
-        return responses
-
-
-class RetrievalGenerator(GeneratorCore):
-    def __init__(self, retriever_parameter: ModelParameter):
-        super().__init__(retriever_parameter)
-    
-    def generate(
-        self,
-        retrieved: List[str],
-        query: str, 
-        history: List[Tuple[str, str]],
-        temperature: float,
-        top_k: int,
-        top_p: float,
-    ) -> str:
-        assert temperature >= 0.0
-        assert top_k >= 0
-        assert top_p >= 0.0 and top_p <= 1.0
-        
-        if history is None:
-            history = []
-            
-        retrieved = "\n".join([f"{idx + 1}. {key}" for idx, key in enumerate(retrieved)]) if retrieved else ""
-        retrieved_query = f"{retrieved}\n\n{query}" if retrieved else query
-        parameter = ModelParameter(self.model, self.tokenizer, self.config)
-        
-        return ChatGenerator(parameter).generate(
-            retrieved_query, 
-            history,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-        )
-
-class EmbeddingEncoder(EmbeddingEncoderCore):
-    def __init__(self, parameter: ModelParameter):
-        super().__init__(parameter)
-        
-    def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
-        return super().encode(sentence)
-        

khaosz/model/__init__.py

@@ -1,17 +0,0 @@
-from khaosz.model.module import ( 
-    Linear,
-    RMSNorm, 
-    MLP,
-    GQA,
-    DecoderBlock,
-)
-from khaosz.model.transformer import Transformer
-
-__all__ = [
-    "Linear",
-    "RMSNorm", 
-    "MLP",
-    "GQA",
-    "DecoderBlock",
-    "Transformer"
-]

khaosz/parallel/__init__.py

@@ -1,27 +0,0 @@
-from khaosz.parallel.setup import (
-    get_world_size, 
-    get_rank,
-    get_current_device,
-    
-    only_on_rank,
-    setup_parallel, 
-    spawn_parallel_fn
-)
-
-from khaosz.parallel.module import (
-    RowParallelLinear,
-    ColumnParallelLinear
-)
-
-__all__ = [
-    "get_world_size",
-    "get_rank",
-    "get_current_device",
-    
-    "only_on_rank",
-    "setup_parallel",
-    "spawn_parallel_fn",
-    
-    "RowParallelLinear",
-    "ColumnParallelLinear"
-]

khaosz/trainer/__init__.py

@@ -1,29 +0,0 @@
-from khaosz.trainer.trainer import Trainer
-from khaosz.trainer.strategy import StrategyFactory
-from khaosz.trainer.schedule import SchedulerFactory
-
-from khaosz.trainer.train_callback import (
-    TrainCallback,
-    ProgressBarCallback,
-    CheckpointCallback,
-    TrainCallback,
-    SchedulerCallback,
-    MetricLoggerCallback
-)
-
-__all__ = [
-    # trainer
-    "Trainer",
-    
-    # factory
-    "StrategyFactory",
-    "SchedulerFactory",
-    
-    # callback
-    "TrainCallback",
-    "ProgressBarCallback",
-    "CheckpointCallback",
-    "TrainCallback",
-    "SchedulerCallback",
-    "MetricLoggerCallback"
-]

khaosz/trainer/schedule.py

@@ -1,164 +0,0 @@
-import math
-from abc import abstractmethod, ABC
-from typing import Any, Dict, List
-from torch.optim.lr_scheduler import LRScheduler
-from khaosz.config.schedule_config import ScheduleConfig
-
-
-class BaseScheduler(LRScheduler, ABC):
-    """
-    Base scheduler class for all other schedulers.
-    """
-    
-    def __init__(self, optimizer, last_epoch: int = -1):
-        super().__init__(optimizer, last_epoch)
-    
-    @abstractmethod
-    def get_lr(self) -> List[float]:
-        raise NotImplementedError
-    
-    def state_dict(self) -> Dict[str, Any]:
-        return super().state_dict()
-    
-    def load_state_dict(self, state_dict: Dict[str, Any]):
-        super().load_state_dict(state_dict)
-
-
-class CosineScheduler(BaseScheduler):
-    """
-    Cosine decay scheduler with warmup, implemented as PyTorch LRScheduler.
-    """
-    
-    def __init__(
-        self, 
-        optimizer, 
-        warmup_steps: int, 
-        lr_decay_steps: int, 
-        min_rate: float = 0.05, 
-        last_epoch: int = -1
-    ):
-        self.warmup_steps = warmup_steps
-        self.lr_decay_steps = lr_decay_steps
-        self.min_rate = min_rate
-        self.total_steps = warmup_steps + lr_decay_steps
-        super().__init__(optimizer, last_epoch)
-    
-    
-    def get_lr(self) -> List[float]:
-        # warmup
-        if self.last_epoch < self.warmup_steps:
-            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
-            return [base_lr * warmup_factor for base_lr in self.base_lrs]
-
-        # cosine decay
-        decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
-        decay_progress = min(decay_progress, 1.0)
-        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
-        decay_factor = max(self.min_rate, cosine_decay)
-        return [base_lr * decay_factor for base_lr in self.base_lrs]
-    
-    def state_dict(self):
-        state = super().state_dict()
-        state.update({
-            'warmup_steps': self.warmup_steps,
-            'lr_decay_steps': self.lr_decay_steps,
-            'min_rate': self.min_rate,
-            'total_steps': self.total_steps,
-        })
-        return state
-    
-    def load_state_dict(self, state_dict):
-        self.warmup_steps = state_dict.pop('warmup_steps')
-        self.lr_decay_steps = state_dict.pop('lr_decay_steps')
-        self.min_rate = state_dict.pop('min_rate')
-        self.total_steps = state_dict.pop('total_steps')
-        super().load_state_dict(state_dict)
-
-
-class SGDRScheduler(BaseScheduler):
-    """
-    SGDR (Stochastic Gradient Descent with Warm Restarts) scheduler, 
-    """
-    
-    def __init__(
-        self, 
-        optimizer, 
-        warmup_steps: int, 
-        cycle_length: int, 
-        min_rate: float = 0.05, 
-        t_mult: int = 2, 
-        last_epoch: int = -1, 
-    ):
-        self.warmup_steps = warmup_steps
-        self.cycle_length = cycle_length
-        self.min_rate = min_rate
-        self.t_mult = t_mult
-        
-        super().__init__(optimizer, last_epoch)
-
-    
-    def get_lr(self):
-        # warmup
-        if self.last_epoch < self.warmup_steps:
-            warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
-            return [base_lr * warmup_factor for base_lr in self.base_lrs]
-
-        # SGDR
-        steps_since_warmup = self.last_epoch - self.warmup_steps
-        
-        # 1. Calculate current cycle and position within cycle
-        current_cycle_length = self.cycle_length
-        total_cycles_length = 0
-        cycle_num = 0
-        
-        while total_cycles_length + current_cycle_length <= steps_since_warmup:
-            total_cycles_length += current_cycle_length
-            current_cycle_length *= self.t_mult
-            cycle_num += 1
-        
-        steps_in_cycle = steps_since_warmup - total_cycles_length
-        
-        # 2. Cosine annealing within the current cycle
-        cosine_factor = 0.5 * (1 + math.cos(math.pi * steps_in_cycle / current_cycle_length))
-        learning_rate_factor = self.min_rate + (1 - self.min_rate) * cosine_factor
-        
-        return [base_lr * learning_rate_factor for base_lr in self.base_lrs]
-    
-    def state_dict(self):
-        """Returns the state of the scheduler as a dict."""
-        state = super().state_dict()
-        state.update({
-            'warmup_steps': self.warmup_steps,
-            'cycle_length': self.cycle_length,
-            'min_rate': self.min_rate,
-            't_mult': self.t_mult
-        })
-        return state
-    
-    def load_state_dict(self, state_dict):
-        """Loads the scheduler's state."""
-        self.warmup_steps = state_dict.pop('warmup_steps')
-        self.cycle_length = state_dict.pop('cycle_length')
-        self.min_rate = state_dict.pop('min_rate')
-        self.t_mult = state_dict.pop('t_mult')
-        super().load_state_dict(state_dict)
-
-
-    
-class SchedulerFactory:
-    """
-    Factory class for creating learning rate schedulers.
-    """
-
-    @staticmethod
-    def load(optimizer, schedule_config: ScheduleConfig) -> BaseScheduler:
-        kwargs = schedule_config.get_kwargs()
-        schedule_type = kwargs.pop("schedule_type")
-        
-        if schedule_type == "cosine":
-            return CosineScheduler(optimizer, **kwargs)
-        elif schedule_type == "sgdr":
-            return SGDRScheduler(optimizer, **kwargs)
-        else:
-            raise ValueError(f"Unsupported schedule type: {schedule_type}")
-        

khaosz/trainer/strategy.py

@@ -1,137 +0,0 @@
-import copy
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from torch import Tensor
-from typing import Any, Callable, Dict, Union
-from abc import ABC, abstractmethod
-
-
-def get_logprobs(
-    model: Union[nn.Module, Callable[..., Dict[str, Tensor]]], 
-    input_ids: Tensor, 
-    mask: Tensor, 
-    pad_token_id: int
-):
-    input_mask =  input_ids.ne(pad_token_id)
-    logits = model(input_ids, input_mask)["logits"]
-    log_probs = torch.log_softmax(logits.float(), dim=-1)
-    
-    shifted_log_probs = log_probs[:, :-1, :] 
-    shifted_input_ids = input_ids[:, 1:]
-    shifted_response_mask = mask[:, 1:]
-    
-    token_logprobs = torch.gather(
-        shifted_log_probs, 
-        dim=-1, 
-        index=shifted_input_ids.unsqueeze(-1)
-    ).squeeze(-1)
-    
-    prompt_mask = input_mask[:, 1:]
-    valid_mask = (prompt_mask & shifted_response_mask).float()
-    
-    return (token_logprobs * valid_mask).sum(dim=-1)
-
-def move_to_device(batch:Dict[str, Tensor], device: str) -> Any:
-    return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
-
-
-class BaseStrategy(ABC):
-    def __init__(self, model: Union[nn.Module, Callable[..., Dict[str, Tensor]]], device: str):
-        self.model = model
-        self.device = device
-    
-    @abstractmethod
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        raise NotImplementedError
-    
-    def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
-        return self.compute_loss(batch)
-
-
-class SeqStrategy(BaseStrategy):
-    def __init__(self, model, device):
-        super().__init__(model, device)
-    
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        batch = move_to_device(batch, self.device)
-        input_ids, target_ids = batch["input_ids"], batch["target_ids"]
-        logits = self.model(input_ids=input_ids)["logits"]
-        
-        loss = F.cross_entropy(
-           input=logits.flatten(0, 1).float(),
-            target=target_ids.flatten()
-        )
-        
-        return loss
-    
-
-class SftStrategy(BaseStrategy):
-    def __init__(self, model, device):
-        super().__init__(model, device)
-    
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        batch = move_to_device(batch, self.device)
-        input_ids, target_ids, loss_mask = batch["input_ids"], batch["target_ids"], batch["loss_mask"]
-        
-        ignore_index = -100
-        logits = self.model(input_ids=input_ids)["logits"]
-        target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
-        
-        loss = F.cross_entropy(
-            input=logits.flatten(0, 1).float(),
-            target=target_ids.flatten(),
-            ignore_index=ignore_index
-        )
-        
-        return loss
-
-
-class DpoStrategy(BaseStrategy):
-    def __init__(self, model, device, pad_token_id, beta):
-        super().__init__(model, device)
-        ref_model = copy.deepcopy(self.model)
-        ref_model.requires_grad_(False)
-        ref_model.eval()
-        
-        self.ref_model = ref_model
-        self.pad_token_id = pad_token_id
-        self.beta = beta
-        
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        batch = move_to_device(batch, self.device)
-        good_ids, bad_ids = batch["chosen"], batch["rejected"]
-        good_mask, bad_mask = batch["chosen_mask"], batch["rejected_mask"]
-        
-        log_pi_good = get_logprobs(self.model, good_ids, good_mask, self.pad_token_id)
-        log_pi_bad = get_logprobs(self.model, bad_ids, bad_mask, self.pad_token_id)
-        
-        with torch.no_grad():
-            log_ref_good = get_logprobs(self.ref_model, good_ids, good_mask, self.pad_token_id)
-            log_ref_bad = get_logprobs(self.ref_model, bad_ids, bad_mask, self.pad_token_id)
-        
-        pi_log_ratio = log_pi_good - log_pi_bad
-        ref_log_ratio = log_ref_good - log_ref_bad
-
-        ratio_diff = pi_log_ratio - ref_log_ratio
-        
-        dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
-        return dpo_loss
-
-
-class StrategyFactory:
-    
-    def load(model, train_type, device, **kwargs):
-        train_strategy: Dict[str, Callable[[], BaseStrategy]] = {
-            "seq": lambda: SeqStrategy(model, device),
-            "sft": lambda: SftStrategy(model, device),
-            "dpo": lambda: DpoStrategy(
-                model,
-                device,
-                kwargs.get("pad_token_id"), 
-                kwargs.get("dpo_beta")
-            )
-        }
-        strategy = train_strategy[train_type]()
-        return strategy        

khaosz/trainer/trainer.py

@@ -1,102 +0,0 @@
-import logging
-from typing import Optional, List
-from khaosz.config import TrainConfig
-from khaosz.trainer.train_callback import (
-    TrainCallback, 
-    ProgressBarCallback, 
-    CheckpointCallback,
-    MetricLoggerCallback,
-    GradientClippingCallback,
-    SchedulerCallback
-)
-from khaosz.trainer.train_context import TrainContext, TrainContextBuilder
-from khaosz.data.checkpoint import Checkpoint
-from khaosz.parallel.setup import spawn_parallel_fn
-
-logger = logging.getLogger(__name__)
-
-
-class Trainer:
-    def __init__(
-        self,
-        train_config: TrainConfig,
-        callbacks: Optional[List[TrainCallback]] = None
-    ):
-        self.train_config = train_config
-        default_callbacks = self._get_default_callbacks()
-        self.callbacks = default_callbacks + callbacks if callbacks else default_callbacks
-
-    def _get_default_callbacks(self) -> List[TrainCallback]:
-        train_config = self.train_config
-        return [
-            ProgressBarCallback(train_config.n_epoch),
-            CheckpointCallback(train_config.checkpoint_dir, train_config.checkpoint_interval),
-            MetricLoggerCallback(train_config.checkpoint_dir, train_config.checkpoint_interval),
-            GradientClippingCallback(train_config.max_grad_norm),
-            SchedulerCallback(),
-        ]
-    
-    def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
-        return (TrainContextBuilder(self.train_config)
-                .with_checkpoint(checkpoint)
-                .with_dataloader()
-                .with_strategy()
-                .build())
-    
-    def _call_callbacks(self, method_name: str, context: TrainContext):
-        for callback in self.callbacks:
-            method = getattr(callback, method_name, None)
-            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,
-            checkpoint=checkpoint
-        )
-
-    def _train_impl(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
-        context = self._build_context(checkpoint)
-        self._call_callbacks('on_train_begin', context)
-        
-        try:
-            context.model.train()
-            # 1.epoch
-            for epoch in range(context.epoch, self.train_config.n_epoch):
-                context.epoch = epoch
-                self._call_callbacks('on_epoch_begin', context)
-                
-                for batch in context.dataloader:
-                    if context.iteration % self.train_config.accumulation_steps == 0:
-                        # 2. step
-                        self._call_callbacks('on_step_begin', context)
-                        context.optimizer.step()
-                        context.optimizer.zero_grad()
-                        self._call_callbacks('on_step_end', context)
-                    
-                    # 3. batch
-                    self._call_callbacks('on_batch_begin', context)
-                    loss = context.strategy(batch)
-                    context.loss = loss.item()
-                    context.iteration += 1
-                    
-                    # to make the loss normalized by accumulation steps
-                    stand_batch = self.train_config.accumulation_steps * self.train_config.nprocs
-                    stand_loss = loss / stand_batch
-                    stand_loss.backward()
-                    
-                    self._call_callbacks('on_batch_end', context)
-
-                self._call_callbacks('on_epoch_end', context)
-        
-        except Exception as e:
-            logger.error(f"Training failed: {str(e)}", exc_info=True)
-            self._call_callbacks('on_error', context)
-            raise
-        finally:
-            self._call_callbacks('on_train_end', context)

khaosz/utils/__init__.py

@@ -1 +0,0 @@
-# init file

khaosz/utils/retriever.py

@@ -1,88 +0,0 @@
-import torch
-import sqlite3
-import numpy as np
-from torch import Tensor
-from typing import Dict, List, Tuple
-
-
-class Retriever:
-    def __init__(self, db_path=None):
-        self.data: Dict[str, Tensor] = {}
-        self.embedding_cache: Tensor = None
-        self.is_caculated: bool = False
-        
-        if db_path is not None:
-            self.load(db_path)
-                
-    def retrieve(self, query: Tensor, top_k: int) -> List[Tuple[str, float]]:
-        if not self.data:
-            return []
-        
-        query = query.flatten().unsqueeze(1)                            # [dim, 1]
-        norm_embeddings = self._embeddings.to(
-            device=query.device,
-            dtype=query.dtype
-        )                                                               # [n_vectors, dim]
-        sim_scores = torch.matmul(norm_embeddings, query).squeeze()     # [n_vectors]
-        
-        top_k = min(top_k, len(self.data))
-        indices = sim_scores.topk(top_k).indices
-        keys = list(self.data.keys())
-        
-        return [(keys[i], sim_scores[i].item()) for i in indices]
-    
-    def add_vector(self, key: str, vector_data: Tensor):
-        self.is_caculated = False
-        self.data[key] = vector_data.flatten().float().cpu()
-        
-    def delete_vector(self, key: str):
-        self.is_caculated = False
-        self.data.pop(key, None)
-    
-    def save(self, db_path):
-        conn = sqlite3.connect(db_path)
-        cursor = conn.cursor()
-        self._init_db(cursor)
-        cursor.execute('DELETE FROM vectors')
-        
-        for item, vec in self.data.items():
-            vec_bytes = vec.numpy().tobytes()
-            cursor.execute('INSERT OR REPLACE INTO vectors (key, vector) VALUES (?, ?)', 
-                           (item, vec_bytes))
-        
-        conn.commit()
-        conn.close()
-
-    def load(self, db_path):
-        conn = sqlite3.connect(db_path)
-        cursor = conn.cursor()
-        self._init_db(cursor)
-        cursor.execute('SELECT key, vector FROM vectors')
-        rows = cursor.fetchall()
-        self.data = {}
-        
-        for row in rows:
-            key, vec_bytes = row
-            vec_numpy = np.frombuffer(vec_bytes, dtype=np.float32).copy()
-            vec = torch.from_numpy(vec_numpy)
-            self.data[key] = vec
-        
-        conn.close()
-        
-    def _init_db(self,cursor: sqlite3.Cursor):
-        # Create table if not exists (in case loading from a new database)
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS vectors (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                key TEXT UNIQUE NOT NULL,
-                vector BLOB NOT NULL
-            )''')
-    
-    @property
-    def _embeddings(self) -> Tensor:
-        if not self.is_caculated:
-            embeddings = torch.stack(list(self.data.values())) 
-            norm_embeddings = embeddings / torch.norm(embeddings, dim=-1, keepdim=True)
-            self.embedding_cache = norm_embeddings
-        
-        return self.embedding_cache

khaosz/utils/splitter.py

@@ -1,127 +0,0 @@
-import re
-import torch
-import torch.nn.functional as F
-
-from abc import ABC, abstractmethod
-from torch import Tensor
-from typing import List, Callable, Optional
-
-
-class BaseTextSplitter(ABC):
-    def __init__(
-        self,
-        max_len: int = 512,
-        chunk_overlap: int = 0,
-    ):
-        if max_len <= 0:
-            raise ValueError("max_len must be > 0")
-        if chunk_overlap < 0:
-            raise ValueError("chunk_overlap must be >= 0")
-
-        self.max_len = max_len
-        self.chunk_overlap = chunk_overlap
-
-    @abstractmethod
-    def split(self, text: str, **kwargs) -> List[str]:
-        raise NotImplementedError
-
-    def preprocess(self, text: str) -> str:
-        return text.strip()
-
-    def postprocess(self, chunks: List[str]) -> List[str]:
-        return [chunk.strip() for chunk in chunks if chunk.strip()]
-
-
-class PriorityTextSplitter(BaseTextSplitter):
-    def __init__(
-        self,
-        separators: List[str],
-        max_len: int = 512,
-        chunk_overlap: int = 0,
-    ):
-        super().__init__(max_len=max_len, chunk_overlap=chunk_overlap)
-        if not separators:
-            raise ValueError("separators must be a non-empty list")
-        self.separators = separators
-
-    def split(self, text: str) -> List[str]:
-        text = self.preprocess(text)
-        for sep in self.separators:
-            parts = text.split(sep)
-
-            valid_parts = [p.strip() for p in parts if p.strip()]
-            if len(valid_parts) > 1:
-                return self.postprocess(valid_parts)
-        return [text]
-
-
-class SemanticTextSplitter(BaseTextSplitter):
-
-    DEFAULT_PATTERN = r'(?<=[。！？!?])(?=(?:[^"\'‘’“”]*["\'‘’“”][^"\'‘’“”]*["\'‘’“”])*[^"\'‘’“”]*$)'
-
-    def __init__(
-        self,
-        embedding_func: Callable[[List[str]], List[Tensor]],
-        pattern: Optional[str] = None,
-        max_len: int = 512,
-        chunk_overlap: int = 0,
-    ):
-        super().__init__(max_len=max_len, chunk_overlap=chunk_overlap)
-        if not callable(embedding_func):
-            raise TypeError("embedding_func must be callable")
-        self.embedding_func = embedding_func
-        self.pattern = pattern or SemanticTextSplitter.DEFAULT_PATTERN
-
-    def split(
-        self,
-        text: str,
-        threshold: float = 0.5,
-        window_size: int = 1,
-    ) -> List[str]:
-        text = self.preprocess(text)
-        sentences = [s.strip() for s in re.split(self.pattern, text) if s.strip()]
-
-        if len(sentences) <= 1:
-            return self.postprocess(sentences)
-
-        try:
-            sentence_embs = self.embedding_func(sentences)
-        except Exception as e:
-            raise RuntimeError(f"Embedding generation failed: {e}")
-
-        if len(sentence_embs) != len(sentences):
-            raise ValueError("Embedding function must return one vector per sentence")
-
-        chunks = []
-        emb_tensor = torch.stack(sentence_embs)  # shape: [N, D]
-        current_chunk: List[str] = [sentences[0]]
-
-        for i in range(1, len(sentences)):
-            start_prev = max(0, i - window_size)
-            end_prev = i
-            start_next = i
-            end_next = min(len(sentences), i + window_size)
-            
-            prev_window_emb = emb_tensor[start_prev:end_prev].mean(dim=0)
-            next_window_emb = emb_tensor[start_next:end_next].mean(dim=0)
-
-            similarity = F.cosine_similarity(
-                prev_window_emb.unsqueeze(0),
-                next_window_emb.unsqueeze(0),
-                dim=1
-            ).item()
-
-            dynamic_threshold = max(threshold * (1 - 0.03 * (end_next - start_prev)), 0.2)
-
-            if similarity < dynamic_threshold:
-                chunks.append(" ".join(current_chunk))
-                overlap_start = max(0, len(current_chunk) - self.chunk_overlap)
-                current_chunk = current_chunk[overlap_start:]
-                current_chunk.append(sentences[i])
-            else:
-                current_chunk.append(sentences[i])
-
-        if current_chunk:
-            chunks.append(" ".join(current_chunk))
-
-        return self.postprocess(chunks)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 dynamic = ["version"]
-name = "khaosz"
+name = "astrai"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
@@ -15,7 +15,11 @@ dependencies = [
     "tqdm==4.67.1",
     "safetensors==0.5.3",
     "huggingface-hub==0.34.3",
-    "pytest==9.0.2"
+    "jinja2>=3.0.0",
+    "fastapi",
+    "uvicorn[standard]",
+    "httpx",
+    "requests",
 ]
 keywords = ["nlp", "datasets", "language-models", "machine-learning"]
 license = { text = "GPL-3.0" }
@@ -24,13 +28,25 @@ classifiers = [
     "License :: OSI Approved :: GPL-3.0",
     "Operating System :: OS Independent",
 ]
-urls = { Homepage = "https://github.com/khaosz/khaosz" }
+urls = { Homepage = "https://github.com/ViperEkura/AstrAI" }
+
+[project.optional-dependencies]
+dev = ["pytest==9.0.2", "ruff"]
 
 [tool.setuptools.packages.find]
 where = ["."]
 
 [tool.pip]
-extra-index-url = "https://download.pytorch.org/whl/cu126" 
+extra-index-url = "https://download.pytorch.org/whl/cu126"
 
 [tool.setuptools.dynamic]
-version = { attr = "khaosz.__version__" }
+version = { attr = "astrai.__version__" }
+
+[tool.ruff]
+target-version = "py312"
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"

scripts/demo/download.py

@@ -0,0 +1,41 @@
+import argparse
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+DEFAULT_LOCAL_DIR = Path(PROJECT_ROOT, "params")
+DEFAULT_REPO_ID = "ViperEk/KHAOSZ"
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Download model parameters from HuggingFace"
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default=DEFAULT_REPO_ID,
+        help=f"HuggingFace repo ID (default: {DEFAULT_REPO_ID})",
+    )
+    parser.add_argument(
+        "--local-dir",
+        type=Path,
+        default=DEFAULT_LOCAL_DIR,
+        help=f"Local directory to save model (default: {DEFAULT_LOCAL_DIR})",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Force download even if files exist",
+    )
+    args = parser.parse_args()
+
+    print(f"Downloading model from {args.repo_id} to {args.local_dir}")
+
+    snapshot_download(
+        repo_id=args.repo_id,
+        local_dir=args.local_dir,
+        force_download=args.force,
+    )
+
+    print("Download complete!")

scripts/demo/generate_ar.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+
+import torch
+
+from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+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)
+
+    query = input(">> ")
+
+    engine = InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+    )
+    response = engine.generate(
+        prompt=query,
+        stream=False,
+        max_tokens=2048,
+        temperature=0.8,
+        top_p=0.95,
+        top_k=50,
+    )
+
+    print(response)
+
+
+if __name__ == "__main__":
+    generate_text()

scripts/demo/generate_batch.py

@@ -0,0 +1,45 @@
+from pathlib import Path
+
+import torch
+
+from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+
+
+def batch_generate():
+    # Load model using AutoModel
+    model = AutoModel.from_pretrained(PARAMETER_ROOT)
+    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
+    model.to(device="cuda", dtype=torch.bfloat16)
+
+    inputs = [
+        "你好",
+        "请问什么是人工智能",
+        "今天天气如何",
+        "我感到焦虑， 请问我应该怎么办",
+        "请问什么是显卡",
+    ]
+
+    engine = InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+    )
+    responses = engine.generate(
+        prompt=inputs,
+        stream=False,
+        max_tokens=2048,
+        temperature=0.8,
+        top_p=0.95,
+        top_k=50,
+    )
+
+    for q, r in zip(inputs, responses):
+        print((q, r))
+
+
+if __name__ == "__main__":
+    batch_generate()

scripts/demo/stream_chat.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import torch
+
+from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+
+
+def chat():
+    model = AutoModel.from_pretrained(PARAMETER_ROOT)
+    tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
+    model.to(device="cuda", dtype=torch.bfloat16)
+
+    messages = []
+    engine = InferenceEngine(model=model, tokenizer=tokenizer)
+
+    while True:
+        query = input(">> ")
+        if query == "!exit":
+            break
+
+        # Add user message
+        messages.append({"role": "user", "content": query})
+
+        # Generate response
+        full_response = ""
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+
+        for token in engine.generate(
+            prompt=prompt,
+            stream=True,
+            max_tokens=2048,
+            temperature=0.8,
+            top_p=0.95,
+            top_k=50,
+        ):
+            print(token, end="", flush=True)
+            full_response += token
+
+        print()
+        # Add assistant response to messages
+        messages.append({"role": "assistant", "content": full_response.strip()})
+
+
+if __name__ == "__main__":
+    chat()

scripts/docker.sh

@@ -0,0 +1,253 @@
+#!/bin/bash
+
+# AstrAI Docker Script
+# Build and manage Docker images
+
+set -e
+
+# Colors
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+# Default values
+IMAGE_NAME="astrai"
+IMAGE_TAG="latest"
+REGISTRY=""
+
+# Print colored messages
+print_info() {
+    echo -e "${BLUE}[INFO]${NC} $1"
+}
+
+print_success() {
+    echo -e "${GREEN}[SUCCESS]${NC} $1"
+}
+
+print_error() {
+    echo -e "${RED}[ERROR]${NC} $1"
+}
+
+print_warning() {
+    echo -e "${YELLOW}[WARNING]${NC} $1"
+}
+
+# Check if Docker is installed
+check_docker() {
+    if ! command -v docker &> /dev/null; then
+        print_error "Docker is not installed"
+        exit 1
+    fi
+    print_success "Docker version: $(docker --version)"
+}
+
+# Build Docker image
+build_image() {
+    local dockerfile="${1:-Dockerfile}"
+    local context="${2:-.}"
+
+    if [ ! -f "$dockerfile" ]; then
+        print_error "Dockerfile not found: $dockerfile"
+        exit 1
+    fi
+
+    print_info "Building Docker image: ${IMAGE_NAME}:${IMAGE_TAG}"
+    docker build -t "${IMAGE_NAME}:${IMAGE_TAG}" -f "$dockerfile" "$context"
+    print_success "Image built successfully"
+}
+
+# Run container
+run_container() {
+    local port="${1:-8000}"
+    local gpu="${2:-false}"
+
+    print_info "Running container on port $port..."
+
+    if [ "$gpu" = true ]; then
+        docker run --gpus all -p "${port}:8000" "${IMAGE_NAME}:${IMAGE_TAG}"
+    else
+        docker run -p "${port}:8000" "${IMAGE_NAME}:${IMAGE_TAG}"
+    fi
+}
+
+# Push image to registry
+push_image() {
+    if [ -z "$REGISTRY" ]; then
+        print_error "Registry not set. Use --registry option"
+        exit 1
+    fi
+
+    local full_tag="${REGISTRY}/${IMAGE_NAME}:${IMAGE_TAG}"
+    print_info "Tagging image: ${full_tag}"
+    docker tag "${IMAGE_NAME}:${IMAGE_TAG}" "$full_tag"
+
+    print_info "Pushing image to registry..."
+    docker push "$full_tag"
+    print_success "Image pushed successfully"
+}
+
+# Remove image
+remove_image() {
+    print_info "Removing image: ${IMAGE_NAME}:${IMAGE_TAG}"
+    docker rmi "${IMAGE_NAME}:${IMAGE_TAG}" 2>/dev/null || print_warning "Image not found"
+    print_success "Image removed"
+}
+
+# Show image info
+show_info() {
+    print_info "Image information:"
+    docker images "${IMAGE_NAME}"
+}
+
+# Show logs
+show_logs() {
+    local container_id="$1"
+    if [ -z "$container_id" ]; then
+        print_error "Container ID required"
+        exit 1
+    fi
+    docker logs "$container_id"
+}
+
+# Main function
+main() {
+    echo "========================================"
+    echo "      AstrAI Docker Management"
+    echo "========================================"
+    echo ""
+
+    COMMAND=""
+    DOCKERFILE="Dockerfile"
+    CONTEXT="."
+    PORT="8000"
+    GPU=false
+
+    # Parse arguments
+    while [[ $# -gt 0 ]]; do
+        case $1 in
+            build)
+                COMMAND="build"
+                shift
+                ;;
+            run)
+                COMMAND="run"
+                shift
+                ;;
+            push)
+                COMMAND="push"
+                shift
+                ;;
+            remove|rm)
+                COMMAND="remove"
+                shift
+                ;;
+            info)
+                COMMAND="info"
+                shift
+                ;;
+            logs)
+                COMMAND="logs"
+                shift
+                ;;
+            --image)
+                IMAGE_NAME="$2"
+                shift 2
+                ;;
+            --tag)
+                IMAGE_TAG="$2"
+                shift 2
+                ;;
+            --registry)
+                REGISTRY="$2"
+                shift 2
+                ;;
+            --dockerfile)
+                DOCKERFILE="$2"
+                shift 2
+                ;;
+            --context)
+                CONTEXT="$2"
+                shift 2
+                ;;
+            --port)
+                PORT="$2"
+                shift 2
+                ;;
+            --gpu)
+                GPU=true
+                shift
+                ;;
+            --help)
+                echo "Usage: $0 <command> [options]"
+                echo ""
+                echo "Commands:"
+                echo "  build     Build Docker image"
+                echo "  run       Run container"
+                echo "  push      Push image to registry"
+                echo "  remove    Remove image"
+                echo "  info      Show image information"
+                echo "  logs      Show container logs"
+                echo ""
+                echo "Options:"
+                echo "  --image NAME       Image name (default: astrai)"
+                echo "  --tag TAG          Image tag (default: latest)"
+                echo "  --registry URL     Registry URL for push"
+                echo "  --dockerfile FILE  Dockerfile path (default: Dockerfile)"
+                echo "  --context PATH     Build context (default: .)"
+                echo "  --port PORT        Port for run (default: 8000)"
+                echo "  --gpu              Enable GPU support"
+                echo "  --help             Show this help message"
+                echo ""
+                echo "Examples:"
+                echo "  $0 build"
+                echo "  $0 build --tag v1.0.0"
+                echo "  $0 run --port 8080"
+                echo "  $0 run --gpu"
+                echo "  $0 push --registry ghcr.io/username"
+                exit 0
+                ;;
+            *)
+                if [ -z "$COMMAND" ]; then
+                    print_error "Unknown command: $1"
+                    exit 1
+                fi
+                shift
+                ;;
+        esac
+    done
+
+    check_docker
+
+    case "$COMMAND" in
+        build)
+            build_image "$DOCKERFILE" "$CONTEXT"
+            ;;
+        run)
+            run_container "$PORT" "$GPU"
+            ;;
+        push)
+            push_image
+            ;;
+        remove)
+            remove_image
+            ;;
+        info)
+            show_info
+            ;;
+        logs)
+            show_logs "$2"
+            ;;
+        "")
+            print_error "No command specified. Use --help for usage"
+            exit 1
+            ;;
+        *)
+            print_error "Unknown command: $COMMAND"
+            exit 1
+            ;;
+    esac
+}
+
+main "$@"

scripts/pre_commit.sh

@@ -0,0 +1,159 @@
+#!/bin/bash
+
+# AstrAI Pre-commit Check Script
+# Runs code format check and tests before committing
+
+set -e
+
+# Colors
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+NC='\033[0m' # No Color
+
+# Print colored messages
+print_info() {
+    echo -e "${YELLOW}[INFO]${NC} $1"
+}
+
+print_success() {
+    echo -e "${GREEN}[SUCCESS]${NC} $1"
+}
+
+print_error() {
+    echo -e "${RED}[ERROR]${NC} $1"
+}
+
+# Check if in project root directory
+check_project_root() {
+    if [ ! -f "pyproject.toml" ]; then
+        print_error "Please run this script from the project root directory"
+        exit 1
+    fi
+}
+
+# Check if Python is installed
+check_python() {
+    if ! command -v python &> /dev/null; then
+        print_error "Python is not installed"
+        exit 1
+    fi
+    print_info "Python version: $(python --version)"
+}
+
+# Install development dependencies
+install_dependencies() {
+    print_info "Installing development dependencies..."
+    pip install --upgrade pip
+    pip install .[dev]
+    print_success "Dependencies installed"
+}
+
+# Run code format check
+run_lint() {
+    print_info "Running code format check (ruff format)..."
+    if ruff format --check .; then
+        print_success "Code format check passed"
+    else
+        print_error "Code format check failed. Please run 'ruff format .' to fix formatting issues"
+        exit 1
+    fi
+}
+
+# Run code style check (linter - import sorting)
+run_ruff_lint_import() {
+    print_info "Running import sorting check (ruff check --select I)..."
+    if ruff check . --select I; then
+        print_success "Import sorting check passed"
+    else
+        print_error "Import sorting check failed. Please run 'ruff check --select I --fix .' to fix import issues"
+        exit 1
+    fi
+}
+
+# Run tests
+run_tests() {
+    print_info "Running tests..."
+    if python -m pytest tests/ -v; then
+        print_success "All tests passed"
+    else
+        print_error "Tests failed"
+        exit 1
+    fi
+}
+
+# Main function
+main() {
+    echo "========================================"
+    echo "    AstrAI Pre-commit Check Script"
+    echo "========================================"
+    echo ""
+
+    check_project_root
+    check_python
+
+    # Parse arguments
+    SKIP_DEPS=false
+    SKIP_LINT=false
+    SKIP_TESTS=false
+
+    while [[ $# -gt 0 ]]; do
+        case $1 in
+            --skip-deps)
+                SKIP_DEPS=true
+                shift
+                ;;
+            --skip-lint)
+                SKIP_LINT=true
+                shift
+                ;;
+            --skip-tests)
+                SKIP_TESTS=true
+                shift
+                ;;
+            --help)
+                echo "Usage: $0 [options]"
+                echo ""
+                echo "Options:"
+                echo "  --skip-deps    Skip dependency installation"
+                echo "  --skip-lint    Skip code checks"
+                echo "  --skip-tests   Skip tests"
+                echo "  --help         Show this help message"
+                exit 0
+                ;;
+            *)
+                print_error "Unknown option: $1"
+                exit 1
+                ;;
+        esac
+    done
+
+    # Install dependencies
+    if [ "$SKIP_DEPS" = false ]; then
+        install_dependencies
+    else
+        print_info "Skipping dependency installation"
+    fi
+
+    # Run code checks
+    if [ "$SKIP_LINT" = false ]; then
+        run_lint
+        run_ruff_lint_import
+    else
+        print_info "Skipping code checks"
+    fi
+
+    # Run tests
+    if [ "$SKIP_TESTS" = false ]; then
+        run_tests
+    else
+        print_info "Skipping tests"
+    fi
+
+    echo ""
+    echo "========================================"
+    print_success "All checks passed! Ready to commit."
+    echo "========================================"
+}
+
+main "$@"

scripts/tools/benchmark.py

@@ -1,7 +1,9 @@
-import torch
-from typing import Dict, Any
 from dataclasses import dataclass
-from khaosz.model.transformer import ModelConfig, Transformer
+from typing import Any, Dict
+
+import torch
+
+from astrai.model.transformer import ModelConfig, Transformer
 
 
 @dataclass
@@ -17,41 +19,47 @@ class GenerationBenchmark:
         self,
         config: ModelConfig,
         device: str = "cuda",
-        dtype: torch.dtype = torch.float16
+        dtype: torch.dtype = torch.float16,
     ):
         self.config = config
         self.device = device
         self.dtype = dtype
         self.model = Transformer(config).to(device=device, dtype=dtype)
         self.model.eval()
-    
+
     def _initialize_kv_cache(self, batch_size: int) -> list:
         """初始化KV缓存"""
         config = self.config
-        shape = (batch_size, config.max_len, config.n_layers, config.n_kv_heads, config.dim // config.n_heads)
+        shape = (
+            batch_size,
+            config.max_len,
+            config.n_layers,
+            config.n_kv_heads,
+            config.dim // config.n_heads,
+        )
         k_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
         v_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
         return (k_cache, v_cache)
-    
+
     def _prepare_inputs(self, batch_size: int, prompt_length: int, total_length: int):
         prompt_ids = torch.randint(
             low=0,
             high=self.config.vocab_size,
             size=(batch_size, prompt_length),
             device=self.device,
-            dtype=torch.long
+            dtype=torch.long,
         )
-        
+
         gen_ids = torch.randint(
             low=0,
             high=self.config.vocab_size,
             size=(batch_size, total_length - prompt_length),
             device=self.device,
-            dtype=torch.long
+            dtype=torch.long,
         )
-        
+
         return prompt_ids, gen_ids
-    
+
     @torch.inference_mode()
     def run_prefill_benchmark(
         self,
@@ -59,32 +67,38 @@ class GenerationBenchmark:
         prompt_length: int = 512,
         num_trials: int = 10,
     ) -> BenchmarkResult:
-        
+
         for _ in range(3):
-            prompt_ids, _ = self._prepare_inputs(batch_size, prompt_length, prompt_length)
+            prompt_ids, _ = self._prepare_inputs(
+                batch_size, prompt_length, prompt_length
+            )
             _ = self.model(prompt_ids)
-        
+
         torch.cuda.synchronize()
-        
+
         total_time = 0.0
         total_tokens = batch_size * prompt_length * num_trials
-        
+
         for trial in range(num_trials):
-            prompt_ids, _ = self._prepare_inputs(batch_size, prompt_length, prompt_length)
+            prompt_ids, _ = self._prepare_inputs(
+                batch_size, prompt_length, prompt_length
+            )
             start_event = torch.cuda.Event(enable_timing=True)
             end_event = torch.cuda.Event(enable_timing=True)
-            
+
             start_event.record()
             _ = self.model(prompt_ids)
             end_event.record()
             torch.cuda.synchronize()
-            
+
             trial_time = start_event.elapsed_time(end_event) / 1000
             total_time += trial_time
-            
+
-            print(f"Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
+            print(
-                  f"({prompt_length / trial_time:.1f} tokens/s)")
+                f"Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s "
-        
+                f"({prompt_length / trial_time:.1f} tokens/s)"
+            )
+
         return BenchmarkResult(
             total_tokens=total_tokens,
             total_time=total_time,
@@ -95,9 +109,9 @@ class GenerationBenchmark:
                 "prompt_length": prompt_length,
                 "dtype": self.dtype,
                 "device": self.device,
-            }
+            },
         )
-    
+
     @torch.inference_mode()
     def run_decoding_benchmark(
         self,
@@ -106,39 +120,43 @@ class GenerationBenchmark:
         gen_length: int = 128,
         num_trials: int = 5,
     ) -> BenchmarkResult:
-        
+
         total_time = 0.0
         total_tokens = batch_size * gen_length * num_trials
-        
+
         for trial in range(num_trials):
-            
+            prompt_ids, gen_ids = self._prepare_inputs(
-            prompt_ids, gen_ids = self._prepare_inputs(batch_size, prompt_length, prompt_length + gen_length)
+                batch_size, prompt_length, prompt_length + gen_length
+            )
             kv_cache = self._initialize_kv_cache(batch_size)
             _ = self.model(prompt_ids, persistent_key_values=kv_cache, start_pos=0)
-            
+
             torch.cuda.synchronize()
 
             start_event = torch.cuda.Event(enable_timing=True)
             end_event = torch.cuda.Event(enable_timing=True)
-            
+
             start_event.record()
-            
+
             current_pos = prompt_length
             for i in range(gen_length):
-                input_token = gen_ids[:, i:i+1]
+                input_token = gen_ids[:, i : i + 1]
-                _ = self.model(input_token, persistent_key_values=kv_cache, start_pos=current_pos)
+                _ = self.model(
+                    input_token, persistent_key_values=kv_cache, start_pos=current_pos
+                )
                 current_pos += 1
-            
+
             end_event.record()
             torch.cuda.synchronize()
-            
+
             trial_time = start_event.elapsed_time(end_event) / 1000
             total_time += trial_time
-            
-            print(f"Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
-                  f"({gen_length / trial_time:.1f} tokens/s)")
 
-        
+            print(
+                f"Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
+                f"({gen_length / trial_time:.1f} tokens/s)"
+            )
+
         return BenchmarkResult(
             total_tokens=total_tokens,
             total_time=total_time,
@@ -150,24 +168,28 @@ class GenerationBenchmark:
                 "gen_length": gen_length,
                 "dtype": self.dtype,
                 "device": self.device,
-            }
+            },
         )
 
 
 def print_benchmark_result(result: BenchmarkResult):
     """打印基准测试结果"""
     benchmark_type = result.metadata["benchmark_type"]
-    
+
     print(f"\n{' ' + benchmark_type.upper().replace('_', ' ') + ' Benchmark ':-^80}")
     print(f"Total Tokens Processed: {result.total_tokens:,}")
     print(f"Time Consumed: {result.total_time:.3f}s")
     print(f"Throughput: {result.tokens_per_second:,.1f} tokens/s")
-    
+
     if benchmark_type == "prefill":
-        print(f"Batch Size: {result.metadata['batch_size']} | Prompt Length: {result.metadata['prompt_length']}")
+        print(
+            f"Batch Size: {result.metadata['batch_size']} | Prompt Length: {result.metadata['prompt_length']}"
+        )
     elif benchmark_type == "decoding":
-        print(f"Batch Size: {result.metadata['batch_size']} | Gen Length: {result.metadata['gen_length']}")
+        print(
-    
+            f"Batch Size: {result.metadata['batch_size']} | Gen Length: {result.metadata['gen_length']}"
+        )
+
     print(f"Device: {result.metadata['device']} | Dtype: {result.metadata['dtype']}")
     print("-" * 80)
 
@@ -183,16 +205,19 @@ if __name__ == "__main__":
         n_layers=24,
         norm_eps=1e-5,
     )
-    
+
     benchmark = GenerationBenchmark(config)
-    
+
     print("=" * 80)
     print("Running Transformer Generation Benchmark")
     print("=" * 80)
-    
+
-    prefill_result = benchmark.run_prefill_benchmark(batch_size=4, prompt_length=512, num_trials=5)
+    prefill_result = benchmark.run_prefill_benchmark(
+        batch_size=4, prompt_length=512, num_trials=5
+    )
     print_benchmark_result(prefill_result)
-    
+
-    gen_result = benchmark.run_decoding_benchmark(batch_size=4, prompt_length=512, gen_length=128, num_trials=5)
+    gen_result = benchmark.run_decoding_benchmark(
+        batch_size=4, prompt_length=512, gen_length=128, num_trials=5
+    )
     print_benchmark_result(gen_result)
-    

scripts/tools/generate.py

@@ -0,0 +1,129 @@
+import argparse
+import json
+
+import torch
+
+from astrai.inference import InferenceEngine
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+
+def processor(
+    model_dir: str,
+    input_json_file: str,
+    output_json_file: str,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    question_key: str,
+    response_key: str,
+    max_tokens: int,
+):
+    # Load model and tokenizer
+    model = AutoModel.from_pretrained(model_dir)
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model.to(device="cuda", dtype=torch.bfloat16)
+
+    # Create inference engine
+    engine = InferenceEngine(model=model, tokenizer=tokenizer)
+
+    with open(input_json_file, "r", encoding="utf-8") as f:
+        input_data = [json.loads(line) for line in f]
+
+    # Check input format: chat messages or raw text
+    if input_data and "messages" in input_data[0]:
+        # Chat format: [{"messages": [...]}]
+        prompts = [
+            tokenizer.apply_chat_template(item["messages"], tokenize=False)
+            for item in input_data
+        ]
+    else:
+        # Raw text format: [{"question": "..."}]
+        prompts = [item[question_key] for item in input_data]
+
+    # Use provided max_tokens or default to model config max_len
+    if max_tokens is None:
+        max_tokens = model.config.max_len
+
+    # Generate responses (batch)
+    responses = engine.generate(
+        prompt=prompts,
+        stream=False,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+    )
+
+    # Write results
+    with open(output_json_file, "w", encoding="utf-8") as f:
+        for prompt, response in zip(prompts, responses):
+            if input_data and "messages" in input_data[0]:
+                output_item = {"response": response}
+            else:
+                output_item = {question_key: prompt, response_key: response}
+            f.write(json.dumps(output_item, ensure_ascii=False) + "\n")
+
+    # Cleanup
+    engine.shutdown()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run generate with a Khaosz model.")
+
+    parser.add_argument(
+        "--model_dir", type=str, required=True, help="Path to the model directory."
+    )
+    parser.add_argument(
+        "--input_json_file",
+        type=str,
+        required=True,
+        help="Path to the input JSONL file.",
+    )
+    parser.add_argument(
+        "--output_json_file",
+        type=str,
+        required=True,
+        help="Path to the output JSONL file.",
+    )
+    parser.add_argument(
+        "--question_key",
+        type=str,
+        default="question",
+        help="Key for the question in the input JSON.",
+    )
+    parser.add_argument(
+        "--response_key",
+        type=str,
+        default="response",
+        help="Key for the response in the output JSON.",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.60,
+        help="Temperature for generating responses.",
+    )
+    parser.add_argument(
+        "--top_k", type=int, default=30, help="Top-k value for generating responses."
+    )
+    parser.add_argument(
+        "--top_p",
+        type=float,
+        default=0.95,
+        help="Top-p value for generating responses.",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=1, help="Batch size for generating responses."
+    )
+    parser.add_argument(
+        "--max_tokens",
+        type=int,
+        default=2048,
+        help="Maximum tokens to generate (default: model config max_len).",
+    )
+
+    args = parser.parse_args()
+
+    with torch.inference_mode():
+        processor(**vars(args))

scripts/tools/perplexity.py

@@ -0,0 +1,111 @@
+import argparse
+import json
+
+import torch
+import torch.nn.functional as F
+import tqdm
+
+from astrai.model import AutoModel
+from astrai.tokenize import AutoTokenizer
+
+
+def process_file(
+    model_dir: str, input_file: str, output_file: str, batch_size: int, text_key: str
+):
+    # Load model and tokenizer
+    model = AutoModel.from_pretrained(model_dir)
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model.to(device="cuda", dtype=torch.bfloat16)
+
+    with open(input_file, "r", encoding="utf-8") as f:
+        input_data = [json.loads(line) for line in f]
+
+    texts = [item[text_key] for item in input_data]
+
+    # Encode all texts
+    print(f"Encoding {len(texts)} texts...")
+    encoded_texts = [tokenizer.encode(text) for text in texts]
+
+    output_data = []
+    total_batches = (len(encoded_texts) + batch_size - 1) // batch_size
+
+    for i in tqdm.tqdm(
+        range(0, len(encoded_texts), batch_size),
+        total=total_batches,
+        desc="Computing perplexity",
+    ):
+        batch_encoded = encoded_texts[i : i + batch_size]
+        batch_texts = texts[i : i + batch_size]
+
+        # Find max length in batch and pad
+        max_len = max(len(seq) for seq in batch_encoded)
+        padded_ids = []
+        masks = []
+
+        for seq in batch_encoded:
+            pad_len = max_len - len(seq)
+            padded_seq = [tokenizer.pad_id] * pad_len + seq
+            mask = [False] * pad_len + [True] * len(seq)
+            padded_ids.append(padded_seq)
+            masks.append(mask)
+
+        # Convert to tensors
+        input_ids = torch.tensor(padded_ids, device="cuda", dtype=torch.long)
+        input_mask = torch.tensor(masks, device="cuda", dtype=torch.bool)
+
+        # Compute perplexity
+        output = model(input_ids, input_mask=input_mask)
+        logits = output["logits"]
+
+        # Shift for causal language modeling
+        shifted_logits = logits[:, :-1, :]  # [batch_size, seq_len-1, vocab_size]
+        shifted_input_ids = input_ids[:, 1:]  # [batch_size, seq_len-1]
+        shifted_mask = input_mask[:, 1:]  # [batch_size, seq_len-1]
+
+        # Compute cross entropy loss
+        loss = F.cross_entropy(
+            shifted_logits.flatten(0, 1),
+            shifted_input_ids.flatten(0, 1),
+            reduction="none",
+        )
+
+        loss = loss.view(shifted_input_ids.shape)  # [batch_size, seq_len-1]
+        loss = loss * shifted_mask
+        sentence_loss = loss.sum(dim=1) / shifted_mask.sum(dim=1).clamp(min=1)
+        perplexity = torch.exp(sentence_loss)  # [batch_size]
+
+        for text, ppl in zip(batch_texts, perplexity):
+            output_data.append({text_key: text, "ppl": float(ppl.item())})
+
+    # Write results
+    with open(output_file, "w", encoding="utf-8") as f:
+        for item in output_data:
+            f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+    print(f"Perplexity computation complete. Results saved to {output_file}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run perplexity with a Khaosz model.")
+    parser.add_argument(
+        "--model_dir", type=str, required=True, help="Path to the model directory."
+    )
+    parser.add_argument(
+        "--input_file", type=str, required=True, help="Path to the input file."
+    )
+    parser.add_argument(
+        "--output_file", type=str, required=True, help="Path to the output file."
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=4, help="Batch size for evaluation."
+    )
+    parser.add_argument(
+        "--text_key",
+        type=str,
+        default="text",
+        help="Key for the text field in the input data.",
+    )
+    args = parser.parse_args()
+
+    with torch.inference_mode():
+        process_file(**vars(args))

scripts/tools/server.py

@@ -0,0 +1,72 @@
+import argparse
+from pathlib import Path
+
+import torch
+
+from astrai.inference.server import run_server
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Start AstrAI inference HTTP server")
+    parser.add_argument(
+        "--host", default="0.0.0.0", help="Host address (default: 0.0.0.0)"
+    )
+    parser.add_argument(
+        "--port", type=int, default=8000, help="Port number (default: 8000)"
+    )
+    parser.add_argument(
+        "--reload", action="store_true", help="Enable auto-reload for development"
+    )
+    parser.add_argument(
+        "--param-path",
+        type=Path,
+        default=None,
+        help="Path to model parameters (default: project_root/params)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Device to load model on (default: cuda)",
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["bfloat16", "float16", "float32"],
+        help="Data type for model weights (default: bfloat16)",
+    )
+    parser.add_argument(
+        "--max_batch_size",
+        type=int,
+        default=16,
+        help="Maximum batch size for continuous batching (default: 16)",
+    )
+    args = parser.parse_args()
+
+    # Convert dtype string to torch dtype
+    dtype_map = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }
+    dtype = dtype_map[args.dtype]
+
+    project_root = Path(__file__).parent.parent.parent
+    param_path = args.param_path or (project_root / "params")
+    print(f"Starting AstrAI inference server on http://{args.host}:{args.port}")
+    print(f"Model parameters expected at: {param_path}")
+    print(f"Device: {args.device}, Dtype: {args.dtype}")
+    run_server(
+        host=args.host,
+        port=args.port,
+        reload=args.reload,
+        device=args.device,
+        dtype=dtype,
+        param_path=param_path,
+        max_batch_size=args.max_batch_size,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/tools/train.py

@@ -0,0 +1,277 @@
+import argparse
+import os
+from functools import partial
+
+import safetensors.torch as st
+import torch
+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.dataset import DatasetFactory
+from astrai.model import Transformer
+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.add_argument(
+        "--train_type",
+        type=str,
+        required=True,
+        choices=["seq", "sft", "dpo"],
+        help="Train type.",
+    )
+    parser.add_argument(
+        "--data_root_path",
+        type=str,
+        required=True,
+        help="Path to the root directory of the dataset.",
+    )
+    parser.add_argument(
+        "--param_path",
+        type=str,
+        required=True,
+        help="Path to the model parameters or resume checkpoint.",
+    )
+
+    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 for training."
+    )
+    parser.add_argument(
+        "--accumulation_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 iters between warnings.",
+    )
+    parser.add_argument(
+        "--max_lr", type=float, default=3e-4, help="Max learning rate for training."
+    )
+    parser.add_argument(
+        "--max_grad_norm",
+        type=float,
+        default=1.0,
+        help="Max gradient norm for clipping.",
+    )
+    parser.add_argument(
+        "--adamw_beta1",
+        type=float,
+        default=0.9,
+        help="Beta values for AdamW optimizer.",
+    )
+    parser.add_argument(
+        "--adamw_beta2",
+        type=float,
+        default=0.95,
+        help="Beta values for AdamW optimizer.",
+    )
+    parser.add_argument(
+        "--adamw_weight_decay",
+        type=float,
+        default=0.01,
+        help="Weight decay for AdamW optimizer.",
+    )
+    parser.add_argument(
+        "--random_seed", type=int, default=3407, help="Random seed for reproducibility."
+    )
+    parser.add_argument(
+        "--num_workers", type=int, default=4, help="Number of workers for data loading."
+    )
+    parser.add_argument(
+        "--no_pin_memory",
+        action="store_false",
+        dest="pin_memory",
+        help="Disable pin memory",
+    )
+    parser.add_argument(
+        "--window_size",
+        type=int,
+        default=None,
+        help="the max length of the input sequence.",
+    )
+    parser.add_argument(
+        "--stride", type=int, default=None, help="the step size of the input sequence."
+    )
+    parser.add_argument("--dpo_beta", type=float, default=0.1, help="DPO beta value.")
+    parser.add_argument(
+        "--label_smoothing",
+        type=float,
+        default=0.1,
+        help="cross_entropy function label smoothing parameter",
+    )
+
+    parser.add_argument(
+        "--ckpt_interval",
+        type=int,
+        default=5000,
+        help="Number of iters between checkpoints.",
+    )
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        default="checkpoint",
+        help="Directory to save checkpoints.",
+    )
+    parser.add_argument(
+        "--start_epoch", type=int, default=0, help="Start epoch for training."
+    )
+    parser.add_argument(
+        "--start_batch", type=int, default=0, help="Start batch for training."
+    )
+
+    parser.add_argument("--nprocs", type=int, default=1, help="Number of GPUs to use.")
+    parser.add_argument(
+        "--device_type", type=str, default="cuda", help="Device type to use."
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def ddp_wrap(model: nn.Module):
+    local_rank = get_rank()
+    model = model.to(device=f"cuda:{local_rank}", dtype=torch.bfloat16)
+    ddp_model = DDP(
+        model,
+        device_ids=[local_rank],
+        output_device=local_rank,
+        find_unused_parameters=False,
+    )
+    return ddp_model
+
+
+def create_optimizer(model: nn.Module, **kwargs) -> optim.Optimizer:
+    return optim.AdamW(model.parameters(), **kwargs)
+
+
+def create_scheduler(
+    optimizer: optim.Optimizer, **kwargs
+) -> optim.lr_scheduler.LRScheduler:
+    return SchedulerFactory.create(optimizer, **kwargs)
+
+
+def prepare_checkpoint(model: nn.Module) -> dict:
+    return model.module.state_dict()
+
+
+def train(
+    train_type: str,
+    param_path: str,
+    data_root_path: str,
+    max_lr: float,
+    n_epoch: int,
+    batch_size: int,
+    start_epoch: int,
+    start_batch: int,
+    accumulation_steps: int,
+    warmup_steps: int,
+    ckpt_interval: int,
+    ckpt_dir: str,
+    dpo_beta: float,
+    adamw_beta1: float,
+    adamw_beta2: float,
+    adamw_weight_decay: float,
+    max_grad_norm: float,
+    label_smoothing: float,
+    random_seed: int,
+    num_workers: int,
+    pin_memory: bool,
+    window_size: int,
+    stride: int,
+    nprocs: int,
+    device_type: str,
+):
+    assert train_type in ["seq", "sft", "dpo"]
+    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)
+
+    if window_size is None:
+        window_size = config.max_len
+
+    # Create bare Transformer (for training, no tokenizer needed)
+    model = Transformer(config)
+
+    # Load weights if available
+    weights_path = os.path.join(param_path, "model.safetensors")
+    if os.path.exists(weights_path):
+        state_dict = st.load_file(weights_path)
+        model.load_state_dict(state_dict, strict=False)
+
+    strategy_kwargs = {"dpo_beta": dpo_beta, "label_smoothing": label_smoothing}
+
+    dataset = DatasetFactory.load(
+        train_type=train_type,
+        load_path=data_root_path,
+        window_size=window_size,
+        stride=stride,
+    )
+
+    optimizer_fn = partial(
+        create_optimizer,
+        **{
+            "lr": max_lr,
+            "betas": (adamw_beta1, adamw_beta2),
+            "weight_decay": adamw_weight_decay,
+        },
+    )
+
+    total_steps = len(dataset) * n_epoch // (batch_size * nprocs)
+    scheduler_fn = partial(
+        create_scheduler,
+        **{
+            "schedule_type": "cosine",
+            "warmup_steps": warmup_steps,
+            "lr_decay_steps": total_steps - warmup_steps,
+        },
+    )
+
+    train_config = TrainConfig(
+        model=model,
+        strategy=train_type,
+        dataset=dataset,
+        optimizer_fn=optimizer_fn,
+        scheduler_fn=scheduler_fn,
+        ckpt_dir=ckpt_dir,
+        n_epoch=n_epoch,
+        batch_size=batch_size,
+        start_epoch=start_epoch,
+        start_batch=start_batch,
+        ckpt_interval=ckpt_interval,
+        accumulation_steps=accumulation_steps,
+        max_grad_norm=max_grad_norm,
+        random_seed=random_seed,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        nprocs=nprocs,
+        parallel_wrapper=ddp_wrap,
+        state_dict_fn=prepare_checkpoint,
+        device_type=device_type,
+        extra_kwargs=strategy_kwargs,
+    )
+
+    trainer = Trainer(train_config)
+    trainer.train()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    train(**vars(args))

tests/conftest.py

@@ -1,42 +1,64 @@
-import os
 import json
-import numpy as np
+import os
-import tempfile
 import shutil
-import torch
+import tempfile
-import pytest
 
+import numpy as np
+import pytest
+import safetensors.torch as st
+import torch
+from tokenizers import Tokenizer, models, pre_tokenizers, trainers
 from torch.utils.data import Dataset
-from khaosz.config.model_config import ModelConfig
+
-from khaosz.data.tokenizer import BpeTokenizer
+from astrai.config.model_config import ModelConfig
-from khaosz.model.transformer import Transformer
+from astrai.model.transformer import Transformer
+from astrai.tokenize import AutoTokenizer
+
+
+def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
+    """Create a simple tokenizer for testing purposes."""
+    tokenizer = Tokenizer(models.BPE())
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size, min_frequency=1, special_tokens=["<unk>", "<pad>"]
+    )
+    # Train on empty iterator with single character
+    tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer)
+    auto_tokenizer = AutoTokenizer()
+    auto_tokenizer._tokenizer = tokenizer
+    auto_tokenizer._special_token_map = {"unk_token": "<unk>", "pad_token": "<pad>"}
+    return auto_tokenizer
 
 
 class RandomDataset(Dataset):
+    """Random dataset for testing purposes."""
+
     def __init__(self, length=None, max_length=64, vocab_size=1000):
         self.length = length or int(np.random.randint(100, 200))
         self.max_length = max_length
         self.vocab_size = vocab_size
-        
+
     def __len__(self):
         return self.length
-        
+
     def __getitem__(self, idx):
         return {
             "input_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
-            "target_ids": torch.randint(0, self.vocab_size, (self.max_length,))
+            "target_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
         }
 
 
 class MultiTurnDataset(Dataset):
+    """Multi-turn dataset with loss mask for SFT training tests."""
+
     def __init__(self, length=None, max_length=64, vocab_size=1000):
         self.length = length or int(np.random.randint(100, 200))
         self.max_length = max_length
         self.vocab_size = vocab_size
-        
+
     def __len__(self):
         return self.length
-        
+
     def __getitem__(self, idx):
         input_ids = torch.randint(0, self.vocab_size, (self.max_length,))
         target_ids = torch.randint(0, self.vocab_size, (self.max_length,))
@@ -50,34 +72,37 @@ class MultiTurnDataset(Dataset):
 
 
 class EarlyStoppingDataset(Dataset):
+    """Dataset that triggers early stopping after a specified number of iterations."""
+
     def __init__(self, length=10, stop_after=5):
         self.length = length
         self.stop_after = stop_after
         self.count = 0
-        
+
     def __len__(self):
         return self.length
-        
+
     def __getitem__(self, idx):
         self.count += 1
         if self.count == self.stop_after:
             raise RuntimeError("Simulated early stopping")
-        
+
         return {
             "input_ids": torch.randint(0, 1000, (64,)),
-            "target_ids": torch.randint(0, 1000, (64,))
+            "target_ids": torch.randint(0, 1000, (64,)),
         }
 
 
 @pytest.fixture
 def base_test_env(request: pytest.FixtureRequest):
+    """Create base test environment with randomly configured model and tokenizer"""
     func_name = request.function.__name__
     test_dir = tempfile.mkdtemp(prefix=f"{func_name}_")
     config_path = os.path.join(test_dir, "config.json")
-    
+
     n_dim_choices = [8, 16, 32]
     n_head_choices = [2, 4]
-    
+
     dim = int(np.random.choice(n_dim_choices))
     n_heads = int(np.random.choice(n_head_choices))
     n_kv_heads = n_heads // 2
@@ -91,16 +116,16 @@ def base_test_env(request: pytest.FixtureRequest):
         "dim_ffn": dim_ffn,
         "max_len": 1024,
         "n_layers": 4,
-        "norm_eps": 1e-5
+        "norm_eps": 1e-5,
     }
-    
+
-    with open(config_path, 'w') as f:
+    with open(config_path, "w") as f:
         json.dump(config, f)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     transformer_config = ModelConfig().load(config_path)
     model = Transformer(transformer_config).to(device=device)
-    tokenizer = BpeTokenizer()
+    tokenizer = create_test_tokenizer()
-    
+
     yield {
         "device": device,
         "test_dir": str(test_dir),
@@ -109,20 +134,63 @@ def base_test_env(request: pytest.FixtureRequest):
         "model": model,
         "tokenizer": tokenizer,
     }
-    
+
     shutil.rmtree(test_dir)
 
+
 @pytest.fixture
 def random_dataset():
     dataset = RandomDataset()
     yield dataset
 
+
 @pytest.fixture
 def multi_turn_dataset():
     dataset = MultiTurnDataset()
     yield dataset
-    
+
+
 @pytest.fixture
 def early_stopping_dataset():
     dataset = EarlyStoppingDataset()
-    yield dataset
+    yield dataset
+
+
+@pytest.fixture
+def test_env(request: pytest.FixtureRequest):
+    """Create a test environment with saved model and tokenizer files."""
+
+    func_name = request.function.__name__
+    test_dir = tempfile.mkdtemp(prefix=f"{func_name}_")
+    config_path = os.path.join(test_dir, "config.json")
+    tokenizer_path = os.path.join(test_dir, "tokenizer.json")
+    model_path = os.path.join(test_dir, "model.safetensors")
+
+    config = {
+        "vocab_size": 1000,
+        "dim": 128,
+        "n_heads": 4,
+        "n_kv_heads": 2,
+        "dim_ffn": 256,
+        "max_len": 64,
+        "n_layers": 2,
+        "norm_eps": 1e-5,
+    }
+    with open(config_path, "w") as f:
+        json.dump(config, f)
+
+    tokenizer = create_test_tokenizer(vocab_size=config["vocab_size"])
+    tokenizer.save(tokenizer_path)
+
+    transformer_config = ModelConfig().load(config_path)
+    model = Transformer(transformer_config)
+    st.save_file(model.state_dict(), model_path)
+
+    yield {
+        "test_dir": test_dir,
+        "model": model,
+        "tokenizer": tokenizer,
+        "transformer_config": transformer_config,
+    }
+
+    shutil.rmtree(test_dir)

tests/data/test_checkpoint.py

@@ -1,11 +1,13 @@
-import torch
 import tempfile
-import torch.distributed as dist
 
+import torch
+import torch.distributed as dist
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from khaosz.data.checkpoint import Checkpoint
+
-from khaosz.parallel.setup import get_rank, spawn_parallel_fn
+from astrai.parallel.setup import get_rank, spawn_parallel_fn
+from astrai.serialization import Checkpoint
+
 
 def test_single_process():
     model = torch.nn.Linear(10, 5)
@@ -14,38 +16,33 @@ def test_single_process():
 
     for epoch in range(3):
         for iteration in range(10):
-
             x = torch.randn(32, 10)
-            y = torch.randn(32, 5)
             loss = model(x).mean()
             loss.backward()
             optimizer.step()
             optimizer.zero_grad()
-        
+
         scheduler.step()
-    
+
-    checkpoint = Checkpoint(
+    checkpoint = Checkpoint(state_dict=model.state_dict(), epoch=3, iteration=30)
-        state_dict=model.state_dict(),
+
-        epoch=3,
-        iteration=30
-    )
-    
     with tempfile.TemporaryDirectory() as tmpdir:
         checkpoint.save(tmpdir)
-        
+
         loaded_checkpoint = Checkpoint.load(tmpdir)
-        
+
         assert loaded_checkpoint.epoch == 3
         assert loaded_checkpoint.iteration == 30
+
+
 def simple_training():
     model = torch.nn.Linear(10, 5)
     optimizer = AdamW(model.parameters(), lr=1e-3)
     scheduler = CosineAnnealingLR(optimizer, T_max=10)
-    
+
     for epoch in range(2):
         for iteration in range(5):
             x = torch.randn(16, 10)
-            y = torch.randn(16, 5)
             loss = model(x).mean()
             loss.backward()
             optimizer.step()
@@ -57,28 +54,23 @@ def simple_training():
         epoch=2,
         iteration=10,
     )
-    
+
     rank = get_rank()
-    
+
     if rank == 0:
         shared_dir = tempfile.mkdtemp()
         checkpoint.save(shared_dir)
     else:
         shared_dir = None
 
-    
     if dist.is_initialized():
         dir_list = [shared_dir]
         dist.broadcast_object_list(dir_list, src=0)
         shared_dir = dir_list[0]
-    
+
- 
     loaded = Checkpoint.load(shared_dir)
     assert loaded.epoch == 2
 
+
 def test_multi_process():
-    spawn_parallel_fn(
+    spawn_parallel_fn(simple_training, world_size=2, backend="gloo")
-        simple_training,
-        world_size=2,
-        backend="gloo"
-    )

tests/data/test_dataset.py

@@ -1,34 +1,36 @@
-import torch
 import numpy as np
+import torch
 
-from khaosz.data.file import save_h5
+from astrai.dataset.dataset import DatasetFactory
-from khaosz.data.dataset import *
+from astrai.serialization import save_h5
-
 
 
 def test_dataset_loader_random_paths(base_test_env):
     """Test dataset loader with multiple random paths"""
     test_dir = base_test_env["test_dir"]
-    
+
     # Create multiple mmap dataset directories with random data
     num_files = np.random.randint(2, 5)
-    
+
     for i in range(num_files):
         seq_length = np.random.randint(200, 400)
         dummy_data = {
-            "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64) for _ in range(10)],
+            "sequence": [
+                torch.randint(0, 1000, (seq_length,), dtype=torch.int64)
+                for _ in range(10)
+            ],
         }
         save_h5(test_dir, f"data_{i}", dummy_data)
-    
+
         # Test loading with multiple paths
-        loaded_dataset = DatasetLoader.load(
+        loaded_dataset = DatasetFactory.load(
-            train_type="seq", 
+            train_type="seq",
-            load_path=test_dir, 
+            load_path=test_dir,
-            window_size=64, 
+            window_size=64,
         )
         assert loaded_dataset is not None
         assert len(loaded_dataset) > 0
-    
+
     # Test that we can get items without errors
     for i in range(len(loaded_dataset)):
         item = loaded_dataset[i]
@@ -41,30 +43,30 @@ def test_dataset_loader_random_paths(base_test_env):
 def test_dpo_strategy_with_random_data(base_test_env):
     """Test DPO strategy with randomized preference data"""
     test_dir = base_test_env["test_dir"]
-    
+
     # Create DPO-style data with memory mapping format
     seq_length = np.random.randint(100, 200)
-    
+
     dummy_data = {
         "chosen": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
         "rejected": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
         "chosen_mask": [torch.ones(seq_length, dtype=torch.bool)],
-        "rejected_mask": [torch.ones(seq_length, dtype=torch.bool)]
+        "rejected_mask": [torch.ones(seq_length, dtype=torch.bool)],
     }
-    
+
     save_h5(test_dir, "dpo_data", dummy_data)
-    
+
     # Load DPO dataset
-    dpo_dataset = DatasetLoader.load(
+    dpo_dataset = DatasetFactory.load(
-        train_type="dpo", 
+        train_type="dpo",
-        load_path=test_dir, 
+        load_path=test_dir,
-        window_size=64, 
+        window_size=64,
     )
-    
+
     assert dpo_dataset is not None
-    assert hasattr(dpo_dataset, 'fetcher')
+    assert hasattr(dpo_dataset, "fetcher")
     assert len(dpo_dataset) > 0
-    
+
     # Test that we can get DPO items without errors
     for i in range(min(3, len(dpo_dataset))):
         item = dpo_dataset[i]
@@ -79,28 +81,28 @@ def test_dpo_strategy_with_random_data(base_test_env):
 def test_sft_dataset_with_random_data(base_test_env):
     """Test SFT dataset with random data"""
     test_dir = base_test_env["test_dir"]
-    
+
     # Create SFT-style data with memory mapping format
     seq_length = np.random.randint(100, 200)
-    
+
     dummy_data = {
         "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
-        "loss_mask": [torch.ones(seq_length, dtype=torch.bool)]
+        "loss_mask": [torch.ones(seq_length, dtype=torch.bool)],
     }
-    
+
     save_h5(test_dir, "sft_data", dummy_data)
-    
+
     # Load SFT dataset
-    sft_dataset = DatasetLoader.load(
+    sft_dataset = DatasetFactory.load(
-        train_type="sft", 
+        train_type="sft",
-        load_path=test_dir, 
+        load_path=test_dir,
-        window_size=64, 
+        window_size=64,
     )
-    
+
     assert sft_dataset is not None
-    assert hasattr(sft_dataset, 'fetcher')
+    assert hasattr(sft_dataset, "fetcher")
     assert len(sft_dataset) > 0
-    
+
     # Test that we can get SFT items without errors
     for i in range(min(3, len(sft_dataset))):
         item = sft_dataset[i]
@@ -114,33 +116,30 @@ def test_sft_dataset_with_random_data(base_test_env):
 def test_dataset_with_custom_stride(base_test_env):
     """Test dataset with custom stride parameter"""
     test_dir = base_test_env["test_dir"]
-    
+
     # Create test data
     seq_length = 200
     dummy_data = {
         "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
     }
-    
+
-    save_h5(test_dir,"stride_test_data", dummy_data)
+    save_h5(test_dir, "stride_test_data", dummy_data)
-    
+
     # Test with custom stride
     custom_stride = 32
-    dataset = DatasetLoader.load(
+    dataset = DatasetFactory.load(
-        train_type="seq", 
+        train_type="seq", load_path=test_dir, window_size=64, stride=custom_stride
-        load_path=test_dir, 
-        window_size=64, 
-        stride=custom_stride
     )
-    
+
     assert dataset is not None
     assert len(dataset) > 0
-    
+
     # With stride 32 and window 64 on 200 length data, we should get more samples
     # than with default stride (which equals window size)
-    default_stride_dataset = DatasetLoader.load(
+    default_stride_dataset = DatasetFactory.load(
-        train_type="seq", 
+        train_type="seq",
-        load_path=test_dir, 
+        load_path=test_dir,
-        window_size=64, 
+        window_size=64,
     )
-    
+
     assert len(dataset) > len(default_stride_dataset)

tests/data/test_sampler.py

@@ -1,30 +1,31 @@
-from khaosz.trainer import *
+from astrai.dataset import ResumableDistributedSampler
-from khaosz.data import *
+
 
 def test_random_sampler_consistency(random_dataset):
     """Test RandomSampler produces consistent results with same seed"""
     dataset = random_dataset
-    
+
     # Create two samplers with same seed
     sampler1 = ResumableDistributedSampler(dataset, seed=42)
     sampler2 = ResumableDistributedSampler(dataset, seed=42)
-    
+
     indices1 = list(iter(sampler1))
     indices2 = list(iter(sampler2))
-    
+
     assert indices1 == indices2
 
+
 def test_random_sampler_different_seeds(random_dataset):
     """Test RandomSampler produces different results with different seeds"""
     dataset = random_dataset
-    
+
     # Create two samplers with different seeds
     sampler1 = ResumableDistributedSampler(dataset, seed=42)
     sampler2 = ResumableDistributedSampler(dataset, seed=123)
-    
+
     indices1 = list(iter(sampler1))
     indices2 = list(iter(sampler2))
-    
+
     # Very high probability they should be different
     assert indices1 != indices2
 
@@ -33,20 +34,20 @@ def test_sampler_across_epochs(random_dataset):
     """Test sampler behavior across multiple epochs"""
     dataset = random_dataset
     n = len(dataset)
-    
+
     sampler = ResumableDistributedSampler(dataset, seed=42)
-    
+
     # Get indices for first epoch
     epoch1_indices = list(iter(sampler))
     assert len(epoch1_indices) == n
-    
+
     # Get indices for second epoch
     epoch2_indices = list(iter(sampler))
     assert len(epoch2_indices) == n
-    
+
     # Check that epochs have different order (should be random)
     assert epoch1_indices != epoch2_indices
-    
+
     # Check that all indices are present in each epoch
     assert set(epoch1_indices) == set(range(n))
-    assert set(epoch2_indices) == set(range(n))
+    assert set(epoch2_indices) == set(range(n))