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.

Prerequisites

• Python 3.12+
• PyTorch 2.11+ (CUDA 12.8 recommended for GPU support)
• NVIDIA GPU with CUDA (optional but recommended; CPU works for inference)

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)
# CSRC_KERNELS=true pip install -e . --no-build-isolation

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

CUDA kernels are opt-in. They are not built by default. When built, they can be activated via with attn_backend(ATTN_BACKEND.CUDA): for accelerated decode/prefill. You can skip them for normal 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 multi-turn interactive chat session with streaming output.

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 (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 -m scripts.tools.server --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-30