Anthracite 1.5.2

Anthracite is a compact PyTorch framework for training and fine-tuning text-generation and embedding models. This release unifies the public API, fixes version drift, hardens optional TPU execution, and makes multi-GPU training cover text and embedding objectives.

Recommended imports: from anthracite import train, finetune, load_model, generate, stream_text, embed, similarity, search.

Install

pip install -e .
# Optional Hugging Face datasets/tokenizers
pip install -e '.[hf]'

TPU support is optional. Install the torch_xla build matching your PyTorch version on a TPU host; otherwise Anthracite raises a typed TPUNotAvailableError rather than failing later in the training loop.

Quick start

from anthracite import train, load_model, generate

meta = train(
    model_name="MyGPT",
    dataset="data/train.jsonl",
    tokens="100M",
    params="20M",
    tokenizer="bundled",
    device="auto",
    output_dir="models/MyGPT",
)

model = load_model(meta["output_dir"], device="auto")
print(generate(model, "The future of open models is", max_new_tokens=80))

anthracite.interface.text.generate_text remains available for old code, but new code should use the flat top-level API. All saved model metadata and tokenizer manifests report version 1.5.2.

Stable generation API

max_new_tokens is preferred. The older max_tokens keyword is still accepted for compatibility; passing both raises an immediate, clear TypeError.

from anthracite import generate, stream_text

print(generate("models/MyGPT", "Hello", max_new_tokens=64, temperature=0.8, top_p=0.95))
print("".join(stream_text("models/MyGPT", "Write one sentence about GPUs.")))

Training and SFT

Anthracite accepts plain text, JSONL, lists of records, and supported Hugging Face datasets. Instruction, QA, ShareGPT, ChatML, OpenAI-style prompt/completion, and common instruction/input/output records are detected automatically.

from anthracite import train

train(
    model_name="SupportGPT",
    dataset="data/instructions.jsonl",
    objective="sft",
    sft_mask_input=True,
    tokens="50M",
    params="20M",
    context_length=512,
    batch_size="auto",
    device="auto",
    precision="auto",
    output_dir="models/SupportGPT",
)

Input masking means the model is optimized primarily on the answer tokens instead of learning to reproduce the prompt. For large Hugging Face datasets, streaming is enabled by default through hf_streaming=True.

Multi-GPU: what actually happens

Use device="multi-gpu" to explicitly request every visible CUDA GPU, or use device="auto" to select multi-GPU automatically when at least two devices are visible.

train(
    model_name="MultiGPU-GPT",
    dataset="data/train.jsonl",
    params="100M",
    tokens="1B",
    device="multi-gpu",
    batch_size="auto",
)

The single-process runtime uses torch.nn.DataParallel: each GPU receives a replica and a shard of the batch, gradients are gathered on the primary GPU, and gradient accumulation preserves the requested effective batch. The automatic batch planner uses the combined memory pool and scales the effective target with the GPU count. Text forward passes and embedding MLM/contrastive objectives are both dispatched through the parallel wrapper.

Checkpoints are saved from the original model, not from the wrapper, so they load normally on CPU, one GPU, or another multi-GPU machine. An explicit multi-GPU request never silently changes to TPU. For very large production jobs, multi-process DistributedDataParallel remains preferable, but this release provides a reliable no-launcher path.

TPU: safe execution and loss debugging

Use device="tpu", device="xla", or device="multi-tpu" only on a working XLA host. Anthracite performs one XLA optimizer step and one graph boundary per optimizer update, and supports torch_xla versions that expose either the native optimizer_step helper or only the optimizer itself.

If loss appears stuck, check the following before changing the learning rate:

1. The dataset is non-empty and labels contain non-ignored targets.
2. Contrastive training has a micro-batch of at least two.
3. precision="bf16" or precision="auto" matches the TPU hardware.
4. The token budget is large enough to produce real optimizer steps.
5. Logs show increasing optimizer step, not only increasing micro-batches.

XLA does not expose CUDA-equivalent free-memory information, so TPU memory planning is deliberately conservative. Never mix arbitrary torch and torch_xla versions.

Fine-tuning and embeddings

from anthracite import finetune, embed, similarity, search

finetune(
    model="models/SupportGPT",
    dataset="data/new_instructions.jsonl",
    tokens="20M",
    output_dir="models/SupportGPT-v2",
)

encoder = "models/MyEmbedder"
vectors = embed(encoder, ["reset password", "track my order"])
print(vectors.shape)
print(similarity(encoder, "reset password", "forgot password"))
print(search(encoder, "where is my package?", ["order tracking", "refund policy"], top_k=1))

Tokenizers

Supported values include "auto", "bundled", a local directory or tokenizer.json, a Hugging Face repository such as "gpt2", or any object exposing encode() and decode(). The tokenizer is saved beside the model and checked against the model vocabulary during loading.

from anthracite import train

train(model_name="GPT2Tok", dataset="data.txt", tokenizer="gpt2", params="20M")
train(model_name="LocalTok", dataset="data.txt", tokenizer="./tokenizers/mytok", params="20M")

Validation before a long run

python -m compileall -q anthracite
python -c 'import anthracite; print(anthracite.__version__)'
python -m anthracite.cli --help

Start with a small smoke run first:

train(
    model_name="smoke",
    dataset=[{"text": "small training corpus " * 20}],
    tokens=10_000,
    params="500K",
    context_length=64,
    batch_size=2,
    device="cpu",
    precision="fp32",
    output_dir="models/smoke",
)

Common runtime causes are an incompatible tokenizer vocabulary, context length below 8, invalid attention dimensions, insufficient memory, or requesting an accelerator that PyTorch cannot see. Anthracite validates these conditions early and reports the correction.

License

MIT. See LICENSE.