A lightweight Transformer training & inference framework
AstrAI is a PyTorch framework for preparing data, training models, evaluating checkpoints and serving inference. It brings these steps into one codebase, with CLI tools for running experiments and Python components for custom workflows.
Use it to train on your own data, compare training stages, run streaming generation, or develop training objectives and execution backends. The Quick Start walks through the full workflow.
| Area | What you can do |
|---|---|
| Data preparation | Tokenize JSON/JSONL data, configure role-based loss masks, pack sequences and write binary datasets with explicit record boundaries. |
| Training | Run pretraining, SFT, DPO and online GRPO/PPO; accumulate gradients, validate during training and resume from committed checkpoints. |
| Optimization and parallelism | Configure compute, gradient and optimizer-state precision; use CPU optimizer-state offload, DDP/FSDP and supported CP/TP layouts. |
| Inference | Generate synchronously or asynchronously, stream text, batch requests continuously and optionally use paged KV allocation with prefix reuse. |
| Serving | Expose OpenAI and Anthropic chat endpoints with streaming responses, or call the Python inference engine directly. |
| Evaluation | Compare checkpoints with MMLU, HumanEval, MBPP, IFEval and other evaluation tools, including perplexity and weight analysis. |
Training objectives, datasets, optimizers and schedulers can be selected through configuration or extended in Python. Attention and other operations have selectable Torch and accelerated backends; see architecture for module boundaries and extension points.
- Train on your data: preprocess raw records, train or fine-tune, then evaluate saved checkpoints.
- Run a checkpoint: use the demo below for interactive input, or start an HTTP service for application requests.
- Modify the framework: follow the module design and testing guide when changing an objective or backend.
Requires Python 3.12+ and PyTorch 2.13.0. The demo requires CUDA.
git clone https://github.com/ViperEkura/AstrAI.git
cd AstrAI
pip install -e .
python scripts/demo/download.py
python scripts/demo/stream_chat.pyThe download is saved to params/. Each prompt is independent; press Ctrl+C
or Ctrl+D to exit. To use another local checkpoint, pass --model_path /path/to/checkpoint.
# Raw continuation (omit --prompt for interactive input)
python scripts/demo/stream_chat.py --raw --prompt "Once upon a time"
# HTTP server
python scripts/tools/server.py --param_path ./params --device cudaThe optional CUDA kernels build when CUDA and nvcc are detected.
Use CSRC_KERNELS=false pip install -e . to skip the extension.
See kernel build instructions for details.
| Task | Guide |
|---|---|
| Install and run the full workflow | Quick Start |
| Prepare training data | Preprocessing |
| Train, resume and configure precision | Training |
| Run on multiple GPUs | Distributed training |
| Generate text and serve HTTP requests | Inference |
| Evaluate checkpoints | Evaluation |
| Look up CLI options | CLI reference |
| Understand subsystem structure and responsibilities | Architecture |
| Work with tokenizer and checkpoint files | HF file conventions |
| Run checks | Testing |
| Deploy with Docker | Serving, training |
See CONTRIBUTING.md. Discuss major changes in an issue before opening a pull request. Questions and examples belong in Discussions.