Anthracite 1.5.5

Anthracite is a small, transparent PyTorch training framework for causal language models and embedding models. Version 1.5.5 focuses on making the runtime predictable: the public API is consistent, multi-GPU training uses real distributed processes, TPU training uses PyTorch/XLA's multiprocessing path, token budgets are global across replicas, and SFT is wired into the actual training pipeline.

Install

Base install:

pip install -e .

Hugging Face support:

pip install -e ".[hf]"

TPU support is optional:

pip install -e ".[tpu]"

torch_xla must be compatible with the installed PyTorch version. Use the PyTorch/XLA installation matrix for the exact PyTorch/XLA pair for your TPU runtime.

The simple API

All common features are available from the package root:

from anthracite import train, finetune, generate, embed, similarity, search, load_model, create_interface

Generation is intentionally the same in examples and in the real package:

from anthracite import generate

text = generate(
    "./models/MyModel",
    "Hello, my name is",
    max_new_tokens=100,
    temperature=0.8,
    top_p=0.95,
)
print(text)

The lower-level alias remains available for compatibility:

from anthracite.interface.text import generate_text

But new code can consistently use from anthracite import generate.

Train a text model

from anthracite import train

train(
    model_name="MyModel",
    dataset="./data/train.jsonl",
    tokens=100_000_000,
    params="20M",
    context_length=512,
    device="auto",
    precision="auto",
    batch_size="auto",
)

device="auto" selects a single accelerator when only one is visible. When multiple CUDA GPUs are visible, Anthracite automatically starts one worker per GPU and trains through DistributedDataParallel (DDP). Each worker owns one full model replica and receives a different shard of the batch.

This is data parallelism, not model sharding: every GPU has the model, every GPU computes on different samples, and gradients are synchronized. Consequently the global effective batch grows with the number of GPUs:

global_batch = micro_batch_per_gpu × gradient_accumulation × GPU_count

A 4-GPU job with micro_batch_per_gpu=8 and accumulation 2 therefore has a global effective batch of 64.

This behavior is deliberate. PyTorch recommends DistributedDataParallel rather than nn.DataParallel for multi-GPU training, with one process per GPU. See the PyTorch DDP documentation and torchrun.

Already inside torchrun

Anthracite also works inside a normal PyTorch distributed launch:

torchrun --standalone --nproc-per-node=4 train.py

Inside train.py, keep the same Python call:

train(
    model_name="MyModel",
    dataset="./data/train.jsonl",
    device="auto",
    batch_size="auto",
)

When Anthracite starts DDP automatically, normal Python scripts are re-executed through torchrun so a top-level train(...) call does not recursively spawn itself.

PyTorch's torchrun launcher provides the rank/world-size environment, and each process operates on one GPU. See the torchrun documentation.

Batch sizing

batch_size="auto" now means:

1. estimate a micro-batch that fits the memory of one local device;
2. use all distributed replicas in the global batch calculation;
3. use gradient accumulation when the requested/effective batch is larger than one micro-batch;
4. never pretend that the sum of all GPUs' VRAM is usable by a single model replica.

This matters because DDP replicates the model. The per-GPU memory limit is still the limit for the model replica on that process. See the PyTorch distributed overview for the distinction between data parallelism and sharded/model-parallel approaches.

If a distributed worker actually OOMs, Anthracite stops the distributed job with a clear memory error instead of shrinking only one rank and risking a collective-operation hang. Lower batch_size, increase gradient_accumulation, reduce context_length, or reduce model size.

TPU training

TPU training is a separate execution path. Anthracite does not treat a TPU like CUDA or try to wrap it in CUDA-style DataParallel.

For multi-TPU execution it uses the PyTorch/XLA multiprocessing launcher, an MpDeviceLoader, XLA's optimizer step, and parameter broadcast at worker startup. See the PyTorch/XLA multi-device guide.

Example:

from anthracite import train

train(
    model_name="TPUModel",
    dataset="./data/train.jsonl",
    tokens=500_000_000,
    params="50M",
    context_length=512,
    device="tpu",
    precision="bf16",
    batch_size="auto",
)

Why this fixes common TPU failures:

• XLA devices are acquired inside the spawned worker, rather than probing an XLA device too early during generic hardware detection.
• MpDeviceLoader is used for multi-device input delivery.
• xm.optimizer_step() performs the XLA distributed gradient consolidation and device step.
• initial parameters are synchronized across replicas.
• bfloat16 is the default automatic precision for TPU.
• fixed local batch shapes are preferred to reduce recompilation.

PyTorch/XLA documents torch_xla.launch() for per-device workers, MpDeviceLoader for input delivery, and xm.optimizer_step() for the distributed XLA optimizer step. See the PyTorch/XLA multi-device guide and XLA AMP guide.

About the “loss stuck at 10.8” problem

A loss value staying near one number is not, by itself, enough to prove a TPU-specific mathematical bug. In the previous implementation, however, the TPU path did have runtime/design problems: device access occurred during generic detection, there was no proper XLA multi-process input path, and distributed step behavior was not aligned with the XLA execution model.

1.5.5 fixes those execution issues. It does not hard-code a target loss or promise a particular loss curve; whether loss decreases depends on the model, tokenizer, objective, data quality, learning rate, batch size, and token budget.

PyTorch/XLA also notes that compilation is expensive and changing tensor shapes can trigger recompilation, so stable batch/sequence shapes are important for TPU performance. citeturn641880search2

Tokenizer preservation and universal loading

Anthracite 1.5.5 treats a tokenizer as a model artifact, not something that should silently be rebuilt during fine-tuning. When fine-tuning an existing model without an explicit tokenizer override, the base tokenizer bundle is copied and fingerprinted, and the model is checked for vocabulary compatibility before training. Hugging Face/tokenizers JSON files are routed to the external loader, so nested model.vocab data is no longer mistaken for the native Anthracite tokenizer format.

Fine-tuning also writes tokenizer_manifest.json with the vocabulary ids and a SHA-256 fingerprint. The output model therefore carries the same tokenizer contract as the base model instead of unexpectedly shrinking a multi-megabyte tokenizer into a tiny replacement file.

SFT is now connected to the real training path

objective="sft" is supported for text models and reaches the SFT dataset/collator instead of silently falling back to ordinary causal-LM training.

Example:

from anthracite import train

train(
    model_name="MyChatModel",
    dataset="./data/instructions.jsonl",
    tokens=20_000_000,
    params="20M",
    context_length=512,
    objective="sft",
    sft_mask_input=True,
    device="auto",
    batch_size="auto",
)

Supported examples include:

{"instruction": "Summarise this.", "input": "Long article...", "output": "Short summary."}

{"prompt": "2 + 2 =", "completion": "4"}

{"question": "What is water?", "answer": "A chemical compound..."}

The SFT dataset masks the input portion with -100 and trains the supervised target portion. For chat records, system/user/context turns are masked and assistant/model turns are supervised. Anthracite auto-detects common schemas such as messages, text, content, instruction, input, output, prompt, response, context, system, and related aliases. The language-model loss uses the standard PyTorch ignore_index=-100 convention, so the mask is honored by the actual model loss.

Architecture 1.2 and embedding backbones

New text training defaults to architecture="anthracite-1.2". The 1.2 path keeps the existing RoPE + GQA + SwiGLU design while adding optional Q/K RMS normalization, scaled residual connections, and safer generation fallbacks. Legacy anthracite-1 and anthracite-1.0 configs remain loadable for fine-tuning compatibility.

Embedding training is selectable with embedding_backbone="auto", "anthracite", or "bert". The built-in BERT option is a dependency-free BERT-style bidirectional encoder with learned positions, GELU Transformer blocks, mean/CLS/max pooling, masked-language-modelling support, and contrastive training support. It is intentionally implemented natively inside Anthracite rather than requiring a Hugging Face runtime.

For capable PyTorch installations, compile=True enables torch.compile as an optional training optimization; the model's attention continues to use PyTorch's standard scaled-attention path.

Fine-tuning

from anthracite import finetune

finetune(
    model="./models/MyModel",
    dataset="./data/instructions.jsonl",
    tokens=20_000_000,
    objective="sft",
    device="auto",
    batch_size="auto",
)

Fine-tuning supports the same automatic multi-GPU and TPU execution paths as base training.

GPT-2 architecture

architecture="gpt2" now selects a native GPT-2 decoder instead of aliasing to Anthracite-1.2. The implementation uses learned token/position embeddings, multi-head self-attention, pre-LayerNorm residual blocks, a 4x GELU MLP, tied input/output embeddings, and an autoregressive KV cache. The default tokenizer for this architecture is the standard GPT-2 tokenizer (gpt2); pass another tokenizer explicitly only when you intentionally want a different vocabulary.

Example:

train(
    model_name="GPT2Model",
    architecture="gpt2",
    dataset="./data/train.jsonl",
    params="117M",
    context_length=1024,
    device="auto",
)

gpt-2 and gpt2-like are compatibility aliases that normalize to the real gpt2 architecture; they no longer map to Anthracite-1.2.

Embeddings

from anthracite import train, embed, similarity, search

train(
    model_name="Embedder",
    model_type="embedding",
    dataset="./data/corpus.jsonl",
    tokens=20_000_000,
    params="30M",
    device="auto",
)

vectors = embed("./models/Embedder", ["hello world", "goodbye world"])
score = similarity("./models/Embedder", "hello", "hi")
hits = search("./models/Embedder", "refund", ["refund policy", "shipping policy"], top_k=2)

Tokenizers

The root API supports:

from anthracite import load_tokenizer

tok = load_tokenizer("auto")
tok = load_tokenizer("bundled")
tok = load_tokenizer("./tokenizer.json")
tok = load_tokenizer("gpt2")

A compatible tokenizer can also be passed directly to train() or finetune().

Checkpoints and resume

train(
    model_name="Resumable",
    dataset="./data/train.jsonl",
    tokens=100_000_000,
    output_dir="./models/Resumable",
    checkpoint_interval=5000,
)

train(
    model_name="Resumable",
    dataset="./data/train.jsonl",
    tokens=100_000_000,
    output_dir="./models/Resumable",
    resume=True,
)

In distributed training, checkpoint and final-model writes are performed by rank 0 only, after a distributed barrier. This prevents multiple workers from overwriting the same model files.

CLI

anthracite --version

anthracite train \
  --model-name MyModel \
  --dataset ./data/train.jsonl \
  --tokens 100M \
  --params 20M \
  --device auto \
  --batch-size auto

anthracite finetune \
  --model ./models/MyModel \
  --dataset ./data/instructions.jsonl \
  --objective sft \
  --device auto

anthracite generate \
  --model ./models/MyModel \
  --prompt "Hello" \
  --max-tokens 100

Package layout

anthracite/
  architectures/   Anthracite-1 transformer implementations
  core/            config, registry, training/fine-tuning orchestration
  datasets/        local/HF/text/pair/SFT loading and collation
  devices/         CPU/CUDA/TPU device backends
  inference/       model loading, generation, embeddings
  interface/       optional Gradio/terminal interface + compatibility aliases
  io/              config, metadata, safetensors
  tokenizer/       native + external tokenizer support
  training/        optimizer, scheduler, memory, distributed loop
  utils/           logging, progress, seed, parameter helpers

Runtime guarantees and boundaries

Anthracite 1.5.5 intentionally makes a few distinctions explicit:

• multi-GPU means replicated data-parallel training with synchronized gradients;
• GPU memory is not pooled into one address space;
• TPU execution uses XLA-specific workers and device loaders;
• token budgets are counted globally across distributed replicas;
• distributed OOM recovery does not mutate the batch plan on only one rank;
• causal text tokenization can be materialized once into a shared cache for all ranks;
• fine-tuning preserves the base tokenizer by default and records a fingerprint;
• runtime failures are recorded as structured error.json / error-rankN.json files before the original exception is re-raised;
• common user functions are exposed from anthracite;
• objective="sft" is connected to the real loss-masking path;
• checkpoint/final writes are single-writer in distributed jobs.

For models that do not fit on one GPU, DDP is the wrong parallelism primitive; PyTorch documents sharded approaches such as FSDP for that case. Anthracite 1.5.5 remains a data-parallel framework and does not claim model sharding. See the PyTorch distributed overview.

Version

Anthracite 1.5.5

License: MIT.