brsx

AutoForge — build, train, and fine-tune models that calibrate themselves to your hardware.

brsx probes your GPU/CPU by actually running a few steps, picks a model size that fits your memory and speed target, and trains it. No config wrestling — pick a mode and go.

Everything is pure PyTorch. No custom CUDA kernels, no compilation step: if torch runs on your machine, brsx runs on your machine.

Install

pip install brsx

Optional extras:

pip install brsx[sound]   # hear your loss while it trains (needs PortAudio)
pip install brsx[qlora]   # 4-bit QLoRA (bitsandbytes, Linux/CUDA only)

Everything else — transformers, peft, accelerate, tokenizers — comes with the base install. You don't need an extra for those.

Usage

from brsx import automodel

automodel.run()

You'll get a menu:

1)  Transformer training        — train a standard Transformer from scratch
2)  MTP Transformer training    — Multi-Token Prediction heads (DeepSeek-V3 style)
3)  HuggingFace model fine-tune — fine-tune a HF model (local path or hub id)
4)  brsx (.pt) model fine-tune  — fine-tune an existing brsx model
5)  Fine-tune an MTP model
6)  Find the right config for your PC
7)  Dataset editor for json, jsonl, parquet...
8)  Convert a brsx model (.pt) to safetensors
9)  Resume training — continue a half-finished run without an optimizer state
10) Hybrid model training — transformer, mamba, gru and cnn in one model
11) Hybrid fine-tune (full or LoRA)
12) Chat with a .safetensors brsx model or a standard HF model

Pick one and follow the prompts — everything (mode, learning rate, steps, checkpointing, data source) is asked interactively.

What's in it

Auto-calibration. Pick min, mid, max, or my (your own parameter target) and brsx measures your actual hardware — running real steps, not guessing from specs — then sizes the model to fit. On OOM it shrinks the batch and keeps going instead of crashing.

Hybrid architectures. Mix four block types in a single model, in any order you like:

mamba,transformer,cnn,gru,transformer,mamba

That list is the architecture. Each type brings something different — mamba for cheap long-range context, transformer for global attention, cnn for local patterns, gru for sequential state — and a small model can borrow from all four. Removing a name from the list and retraining is all an ablation takes.

Multi-Token Prediction. DeepSeek-V3-style sequential MTP heads: the model predicts 3 tokens ahead, and generation emits them a block at a time — roughly 3× fewer forward passes.

Resume without optimizer state. Checkpoints stay small (weights only, no 2× bloat from Adam moments). On resume the LR re-warms over the first few steps and the data stream fast-forwards past what was already seen, so you don't re-train on the same tokens.

Fine-tuning. Full or LoRA, on brsx models, hybrids, MTP models, or HuggingFace checkpoints. LoRA rank is yours to pick; the adapters are merged back into plain weights on save, so the output is just a normal model.

Audio feedback. With brsx[sound] installed, the loss becomes a tone: the pitch falls as the model learns, and a siren tells you it went NaN. Leave a run in the background and hear when something breaks instead of watching the log.

Portable chat. Every trained model folder gets a standalone chat.py that reads config.json, rebuilds whatever architecture it describes, and runs. Move the folder anywhere — it only needs torch and tokenizers.

Design

brsx is built for small models on ordinary hardware. Not a 70B on a cluster: a 50M that trains on the laptop you already own. Everything in it — the calibration, the pure-PyTorch blocks, the small checkpoints — follows from that.

License

Apache-2.0. See LICENSE.