Axis

Axis is a from-scratch framework for training and fine-tuning transformer models on the GPU. Instead of composing autograd over a pile of dependencies, Axis lowers the entire training step — embedding, every transformer block, the loss, the full backward pass and the AdamW update — into a single compiled C++/CUDA execution plan, run as one native call or captured as a CUDA graph and replayed.

Every capability is proven numerically identical to a reference implementation before it ships. It is a proprietary framework by Zyora Labs, made in India.

pip install axis-zyora     # engine builds on first use via your local nvcc
axis-zyora check           # report GPU / nvcc / engine status

A modern training step dispatches hundreds of small kernel launches per layer, each crossing the Python boundary. That overhead is invisible at toy scale and dominant at real scale. Axis removes the boundary entirely: the training step is described once as a static plan of typed operations, then executed by a native runtime that owns its own cuBLAS handle, workspace and buffer table.

Because the plan is static, it can be captured as a CUDA graph — so the steady-state training loop replays with almost no host-side cost, and learning-rate schedules, gradient clipping and the optimizer all run device-side, graph-safe.

• A single execution plan: embed → N transformer blocks → norm → head → cross-entropy → full backward → AdamW, executed as one native call.
• bf16 storage with fp32 master weights and fp32 accumulation — tensor-core throughput with no loss scaling.
• Fused flash attention, forward and backward, on WMMA tensor cores — attention probabilities never touch memory, so memory stays flat as sequence length grows.
• A complete training loop: LR schedules under CUDA graphs, global-norm gradient clipping, correct weight decay (1D params excluded), ignore_index / padding masking, and fully resumable checkpoints.

LoRA fine-tuning is where Axis has a clear, defensible win. The frozen base carries no optimizer state and stays in bf16, so Axis fine-tunes a 1B-class model in 38% less GPU memory than HuggingFace + PEFT, at comparable speed.

import axis
from axis import lora

model = axis.from_pretrained("path/to/llama-model")   # HF safetensors
lora.apply_lora(model, rank=16, alpha=32)             # freeze base, train ~0.5%

ct = axis.compile_model(model, batch=4, seq=2048, dtype="bf16", lr=1e-4)
for x, y in loader:            # padding / masked positions in y set to -1
    loss = ct.step(x, y)

Data parallelism runs one process per GPU with the gradient all-reduce fused directly into the plan between backward and the optimizer step — 99% weak-scaling efficiency on two A100s. ZeRO-1 shards the optimizer state (fp32 master + Adam moments) across ranks, bit-identical to replicated data parallel, so you can train a model whose optimizer state is larger than one GPU can hold.

ct = axis.compile_model(model, batch=4, seq=2048, dtype="bf16",
                        grad_sync=True, zero_stage=1, rank=rank, world=world)

We publish the numbers, including where Axis loses. Against the real PyTorch stack (HuggingFace Transformers with SDPA flash attention, and torch.compile), full-model training at 1B scale is slower and uses more memory — a decade of tuning sits behind PyTorch. Axis's advantage is fine-tuning memory, not raw training throughput.

• Compiled engine vs reference, 30 full training steps end to end through the optimizer: max relative drift 3.1e-07.
• 32 / 32 shape-fuzz configurations (varied depth, width, heads, KV groups, vocab, tied embeddings and LoRA).
• 500-step CUDA-graph soak: converges, zero NaN, exact bias correction.
• Multi-GPU and ZeRO-1 weights bit-identical to replicated data parallel; loads real HuggingFace Llama checkpoints and reproduces logits to ~3e-6.

Axis is on PyPI as axis-zyora. The CUDA engine compiles from source on first use with your local nvcc (GPU architecture auto-detected and cached), so a fresh install just works on a CUDA box.

import axis
from axis import nn

model = nn.Transformer(
    vocab_size=32000, dim=1536, n_layers=48,
    n_heads=24, n_kv_heads=8, mlp_hidden=4096, max_seq_len=2048,
)

# Compile the whole training step to the CUDA engine (bf16 + fused flash attn)
ct = axis.compile_model(model, batch=4, seq=2048, dtype="bf16",
                        lr=3e-4, wd=0.1, max_grad_norm=1.0)

for x, y in loader:            # one native call = forward + backward + AdamW
    loss = ct.step(x, y)