AstrAI/docs/get-started.md

Getting Started

This guide walks you through installing AstrAI, downloading a model, running inference, preprocessing data, and launching your first training job.

Contents

• Prerequisites
• 1. Install
• 2. Download Model Weights
• 3. Run Inference
• 4. Preprocess Data
• 5. Train
• 6. Evaluate
• 7. Docker
• Next Steps

Prerequisites

• Python 3.12+
• PyTorch 2.11.0 (the exact version pinned by AstrAI; CUDA 12.8 build recommended for GPU support)
• NVIDIA GPU with CUDA for training, scripts/tools/generate.py, generation evaluations, and demos. The HTTP server and direct-scoring evaluations can run on CPU where their CLI exposes a CPU device.

1. Install

git clone https://github.com/ViperEkura/AstrAI.git
cd AstrAI

# Basic install (pure PyTorch, no custom CUDA kernels)
pip install -e .

# With CUDA kernels (optional, for fused attention and rotary embedding)
# CSRC_KERNELS=true pip install -e . --no-build-isolation

# With dev dependencies (pytest, ruff)
# pip install -e ".[dev]"

CUDA kernels are opt-in at build time (CSRC_KERNELS=true). Once built, CudaBackend is the default attention backend on GPU (cuda > flash > torch priority). Override via ASTR_BACKEND env var or attn_backend() context manager. Fused rotary embedding kernel is auto-dispatched when available. Skip for CPU-only usage.

2. Download Model Weights

AstrAI uses HuggingFace-style model directories. Download the default 1B instruction-tuned model:

python scripts/demo/download.py
# → Downloads to params/

To use a different model:

python scripts/demo/download.py --repo-id <HF_REPO_ID> --local-dir ./my_model

The model directory contains:

• config.json — model architecture configuration
• model.safetensors — model weights
• tokenizer.json + tokenizer_config.json — tokenizer files (including chat template)

3. Run Inference

Interactive Chat (Simplest)

python scripts/demo/stream_chat.py
# Type your message after >>, type !exit to quit

This starts a single-turn interactive prompt loop with streaming output. Each prompt is independent; conversation history is not retained.

Start an HTTP Server

# Terminal 1: start server
python scripts/tools/server.py --param_path ./params --device cuda

# Terminal 2: query (OpenAI-compatible API)
curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'

The server also supports the Anthropic API at /v1/messages. See Inference Guide for full API documentation.

Batch Generation from a File

Create an input JSONL file (one JSON object per line):

{"question": "What is machine learning?"}
{"question": "Explain gradient descent."}

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

4. Preprocess Data

AstrAI uses a declarative JSON config to define the preprocessing pipeline. Create a config file for your training type:

Pretraining (seq)

Input JSONL:

{"text": "Artificial intelligence is..."}

Config (pretrain.json):

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

SFT (Supervised Fine-Tuning)

Input JSONL:

{"messages": [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]}

Config (sft.json):

{
    "input": {
        "sections": [{"field": "messages", "action": "$role", "template": true}]
    },
    "mask": {
        "system": "mask",
        "user": "mask",
        "assistant": "train"
    },
    "mask_default": "mask",
    "preprocessing": {"max_seq_len": 2048},
    "output": {"storage_format": "bin", "dtype": {"loss_mask": "bool"}}
}

Run Preprocessing

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

See Preprocessing Guide for DPO/GRPO configs and all options.

5. Train

Single GPU

python scripts/tools/train.py \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=./params \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --max_lr=1e-4 \
    --window_size=2048 \
    --ckpt_dir=./checkpoint \
    --nprocs=1 \
    --parallel_mode=none

Multi-GPU (DDP)

export CUDA_VISIBLE_DEVICES=0,1,2,3
export NCCL_P2P_DISABLE=1
export NCCL_NET_GDR_LEVEL=0

python scripts/tools/train.py \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=./params \
    --parallel_mode=ddp \
    --nprocs=4 \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --max_lr=1e-4 \
    --window_size=2048 \
    --ckpt_dir=./checkpoint

Training Types

--train_type	Description	Data Keys
seq	Pre-training (next-token prediction)	sequence
sft	Supervised fine-tuning (masked loss)	sequence, loss_mask
dpo	Direct Preference Optimization	chosen, rejected, *_mask
grpo	Group Relative Policy Optimization	prompts, responses, masks, rewards

See Training Guide for loss formulas and strategies. See Distributed Guide for DDP/FSDP details.

6. Evaluate

HumanEval and MMLU download their benchmark data through HuggingFace datasets, which is not part of the base install:

pip install datasets

# HumanEval (code generation, auto-downloads dataset)
python scripts/eval/evaluate_humaneval.py --param_path ./params --num_samples 20

# MMLU (knowledge, auto-downloads dataset)
python scripts/eval/evaluate_mmlu.py --param_path ./params --n_shot 5

# Perplexity on custom data
python scripts/eval/evaluate_ppl.py --param_path ./params --input_path data.jsonl --output_dir ppl_results/

See Evaluation Guide for all benchmarks.

7. Docker

# Build
docker build -t astrai:latest .

# Run inference server with GPU
docker run --gpus all -p 8000:8000 astrai:latest \
  python scripts/tools/server.py --port 8000 --device cuda

# Docker Compose (GPU)
docker compose up -d

Next Steps

Topic	Document
CLI parameters (train, server, generate, preprocess)	CLI Reference
Preprocessing pipeline details	Preprocessing Guide
Training loop, strategies, schedulers	Training Guide
KV cache, continuous batching, HTTP API	Inference Guide
Evaluation benchmarks	Evaluation Guide
Multi-GPU DDP / FSDP	Distributed Guide
System architecture	Architecture
Data pipeline internals	Data Flow

Document Update Time: 2026-07-31