Bantam

Bantam is a configurable decoder-only transformer packaged with a CLI for tokenizer training, pre-training, SFT, and chat inference. It runs on pure PyTorch (torch>=2.1) and uses the scaled-dot-product/FlashAttention stack for grouped-query attention, sliding windows, and sink tokens without custom kernels.

Key Features

• Flex-attention based GQA with optional sliding windows and sink tokens.
• Per-layer overrides let you mix head counts, windows, sinks, and MoE routing per block via layer_configs.
• SwiGLU dense blocks or Switch/Top-k MoE with router/aux loss controls.
• Rotary embeddings with linear, NTK, and YaRN-style scaling.
• Regex-aware BPE trainer (bantam-cli trainbpe) that isolates punctuation, chunks long numbers, and optionally prefills 0–N.
• Unified CLI for model init, pre-training, SFT (full or LoRA), resume, and chat streaming.

Installation

python3 -m pip install bantam-cli

CLI Quick Start

bantam-cli --help          # list commands

# Train a tokenizer from JSONL
bantam-cli trainbpe --input data/pretrain.jsonl --out tokenizers/run1

# Generate starter configs (model + train + SFT)
bantam-cli define --out configs

# Initialise a model directory from split configs
bantam-cli init --mconf configs/models/pretrain_small.yaml --config configs/training/pretrain_small.yaml --out models/bantam-init

# Launch pre-training run
bantam-cli pretrain --mconf configs/models/pretrain_small.yaml --config configs/training/pretrain_small.yaml

# Resume a checkpointed run (model config auto-detected)
bantam-cli continue-pretrain --config configs/training/pretrain_small.yaml --model models/run1/checkpoints/step_20000

# Supervised fine-tuning (LoRA or full)
bantam-cli sft --config configs/training/sft_small.yaml --model models/bantam-init

# Chat with a trained checkpoint (pretrain format by default)
bantam-cli chat --model models/run1/best

# base model inference
bantam-cli chat --model Theoistic/Bantam-285m

# Conversational chat retaining the last two turns
bantam-cli chat --model Theoistic/Bantam-285m-Instruct --format chat --history-turns 2

Model configs live under configs/models/ with a bantam_config block. Training/SFT configs live under configs/training/ with training_args, optional sft_args, and a run seed. Relative paths resolve against the YAML file location.

Example per-layer tweak:

bantam_config:
  hidden_size: 1024
  num_hidden_layers: 12
  num_attention_heads: 16
  num_key_value_heads: 4
  layer_configs:
    - { window: 2048, num_attention_sinks: 2 }
    - { expert_type: switch, num_experts: 8, moe_capacity_factor: 1.0 }

More Info

• Source code & full documentation: https://github.com/theo-bntm/bantam
• Release checklist: publish.md inside the repository

Report issues or ideas via GitHub. Enjoy experimenting with Bantam!

LLaMAFactory Integration

You can fine-tune Bantam models with LLaMAFactory once bantam-cli is installed

you need to register the architecture in LLaMA-Factory/src/llamafactory/cli.py

import bantam  # lazy imports
import bantam.tokenization_bantam  # registers BantamFastTokenizer with AutoTokenizer
import bantam.modeling_bantam      # registers config/model with AutoConfig/AutoModel

and add this somewhere in the bottom of LLaMA-Factory/src/llamafactory/data/template.py

register_template(
    name="bantam_chat",
    format_user=StringFormatter(
        slots=["<|USER_START|>{{content}}<|USER_END|><|AGENT_START|>"]
    ),
    format_assistant=StringFormatter(slots=["{{content}}<|AGENT_END|>"]),
    format_system=StringFormatter(slots=["<|SYSTEM_START|>{{content}}<|SYSTEM_END|>"]),
    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
    stop_words=["<|AGENT_END|>", "<|EOS|>"],
    replace_eos=False,
)

• Use --template bantam_chat in LLaMAFactory for both SFT and preference/DPO runs.
• The template leaves user turns in-place but only supervises the assistant span (mirrors Bantam’s mask_user_queries=true default).
• Set include_agent_end=true and include_eos=true in your Bantam SFT config so completions end with <|AGENT_END|><|EOS|>, giving the stop tokens above reliable cut points.
• For LoRA vs. full SFT, follow LLaMAFactory’s standard CLI flags—Bantam checkpoints expose all weights so “full” fine-tunes work out of the box.

After training, export the adapter or merged weights from LLaMAFactory and load them with bantam-cli chat --model <path> to validate generations.