A lightweight Transformer training & inference framework

📖 Table of Contents

• Overview
• Getting Started
• Demo
• Documentation
• Contributing
• Community
• License

English

Overview

AstrAI is an end-to-end Transformer framework for building, training, evaluating, and serving models. It provides a compact PyTorch codebase for the complete model lifecycle, from declarative data preprocessing and distributed training to continuous-batching inference and OpenAI/Anthropic-compatible APIs.

| Area | Capabilities | |---|---| | Models | Autoregressive language models and embedding models with GQA, MLA, MoE, RoPE, and extensible attention/FFN components | | Training | Pre-training (seq), supervised fine-tuning (sft), DPO, and online RL (GRPO/PPO) with in-process or cross-GPU rollout backends, gradient accumulation, checkpointing, DDP, and FSDP | | Data | Declarative JSON preprocessing, configurable masking and packing, binary/JSONL storage, and streaming datasets | | Inference | Continuous batching, paged KV cache, radix prefix caching, streaming generation, and Torch/CUDA/FlashAttention backends | | Serving | FastAPI server with OpenAI and Anthropic chat completion protocols, including SSE streaming and tool calls | | Evaluation | Perplexity, MMLU, HumanEval, IFEval, IFD, ROUGE, and weight-analysis evaluation tools | | Extensibility | Factory and registry architecture for models, datasets, training strategies, callbacks, kernels, and protocol components |

Getting Started

End-to-end walkthrough in 5 steps:

1. Install

AstrAI requires Python 3.12+ and pins PyTorch exactly to 2.11.0. Training, scripts/tools/generate.py, generation evaluations, and the generation demos require CUDA; CPU support is limited to components with an explicit CPU device path, such as the HTTP server and direct-scoring evaluations.

git clone https://github.com/ViperEkura/AstrAI.git
cd AstrAI
pip install -e .                                          # kernels auto-build when nvcc + CUDA are detected
CSRC_KERNELS=false pip install -e .                     # skip kernels (pure PyTorch)
CSRC_KERNELS=true pip install -e . --no-build-isolation  # force the fused CUDA kernel build
pip install -e ".[dev]"                                  # dev dependencies (pytest, ruff)

2. Download model

python scripts/demo/download.py    # downloads 1B checkpoint to params/

3. Preprocess data

Create pretrain.json (preprocessing config for seq strategy):

{
    "version": 1,
    "input": {"sections": [{"field": "text", "action": "train"}]},
    "preprocessing": {"max_seq_len": 2048},
    "output": {"storage_format": "bin"}
}

python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json

4. Train

export CUDA_VISIBLE_DEVICES=0,1,2,3
nohup python scripts/tools/train.py \
    --dp_size=4 \
    --dp_mode=ddp \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --weight_decay=0.1 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &

5. Serve & query

# Terminal 1: start server
python scripts/tools/server.py --param_path ./params --device cuda
Terminal 2: query
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'

Demo

Check out the demos in the scripts/demo/ folder:

# Download model weights (required before running demos)
python scripts/demo/download.py                      # model → params/
Single-turn interactive streaming prompt loop (no conversation history)
python scripts/demo/stream_chat.py
Type your message after >>, type !exit to quit
Batch generation (5 hardcoded prompts, non-streaming)
python scripts/demo/generate_batch.py
Single-prompt autoregressive streaming
python scripts/demo/generate_ar.py

All generation demos use temperature=0.8, top_p=0.95, top_k=50, max_tokens=2048 by default and require params/ to contain model weights (run download.py first).

Watch a video walkthrough on bilibili.

See Documentation for full references beyond the examples above.

Text Generation

Batch generation from a JSONL file:

python scripts/tools/generate.py \
    --param_path ./params \
    --input_json_file input.jsonl \
    --output_json_file output.jsonl

Docker

Build and run with Docker (recommended for GPU environments):

# Build image
docker build -t astrai:latest .
Run with GPU support
docker run --gpus all -it astrai:latest
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
docker run --gpus all -v /path/to/data:/data -it astrai:latest
Docker Compose (GPU, default)
docker compose up -d
Docker Compose CPU server profile (CUDA-only generation scripts/demos are unavailable)
docker compose --profile cpu up -d
YAML-driven serving (see serve.yaml; up/run/down/logs/status...)
bash scripts/serve.sh up

Note: --gpus all is required for CUDA support. Without it, torch.cuda.is_available() will return False.

HTTP API Examples

Additional request examples beyond the Getting Started flow:

# OpenAI-compatible streaming
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}'
Anthropic-compatible
curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{"model":"astrai","system":"You are a helpful assistant.","messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
Anthropic-compatible streaming with stop sequences
curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{"model":"astrai","messages":[{"role":"user","content":"Write a story"}],"max_tokens":500,"stream":true,"stop_sequences":["The end"]}'
Health check
curl http://localhost:8000/health

See Inference Guide for SSE streaming format, error codes, and stats endpoint.

Documentation

| Document | Description | |----------|-------------| | Get Started | Installation and quickstart | | CLI Reference | Parameters for all CLI tools (train, server, generate, preprocess) | | Preprocessing | Declarative JSON-driven data preprocessing | | Training | Training loop, strategies & formulas | | Inference | KVCache, continuous batching, sampling & HTTP API | | Evaluation | HumanEval, MMLU, PPL, ROUGE, IFD, IFEval | | Distributed | Multi-GPU DDP / FSDP training | | Architecture | System architecture, class diagram & design patterns | | Data Flow | Data pipeline, storage backends & dataset architecture | | Internals | Training internals: loss formulas, callback lifecycle, KV cache | | CUDA Kernels | Custom CUDA attention kernels & benchmarks | | Docker Serving | YAML-driven containerized serving (serve.yaml, serve.sh) | | Docker Training | YAML-driven containerized training (train.yaml, train.sh) |

Contributing

We welcome contributions! Please see our Contributing Guidelines for details.

1. Fork the repository.
2. Create a feature branch.
3. Commit your changes.
4. Open a Pull Request.

Community

• GitHub Issues: Issue Tracker
• Discussions: GitHub Discussions
• HuggingFace: Model Hub

License

This project is licensed under the Apache-2.0 License.