Metadata-Version: 2.4
Name: anthracite
Version: 1.5.5
Summary: Universal AI training & fine-tuning framework with native GPT-2 architecture, GPT-2 tokenizer support, smart SFT, and HF streaming.
License: MIT
Keywords: ai,deep-learning,llm,training,fine-tuning,tokenizer,nlp
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: numpy>=1.24
Requires-Dist: safetensors>=0.4
Requires-Dist: tqdm>=4.65
Requires-Dist: transformers<6,>=4.36
Requires-Dist: huggingface_hub>=0.20
Provides-Extra: hf
Requires-Dist: datasets>=2.14; extra == "hf"
Requires-Dist: huggingface_hub>=0.20; extra == "hf"
Requires-Dist: tokenizers>=0.15; extra == "hf"
Requires-Dist: transformers>=4.36; extra == "hf"
Provides-Extra: tpu
Requires-Dist: torch_xla>=2.1; extra == "tpu"
Provides-Extra: all
Requires-Dist: datasets>=2.14; extra == "all"
Requires-Dist: huggingface_hub>=0.20; extra == "all"
Requires-Dist: tokenizers>=0.15; extra == "all"
Requires-Dist: transformers>=4.36; extra == "all"
Requires-Dist: torch_xla>=2.1; extra == "all"
Dynamic: license-file

# 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:

```bash
pip install -e .
```

Hugging Face support:

```bash
pip install -e ".[hf]"
```

TPU support is optional:

```bash
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:

```python
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:

```python
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:

```python
from anthracite.interface.text import generate_text
```

But new code can consistently use `from anthracite import generate`.

## Train a text model

```python
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:

```text
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](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) and [`torchrun`](https://docs.pytorch.org/docs/stable/elastic/run.html).

### Already inside torchrun

Anthracite also works inside a normal PyTorch distributed launch:

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

Inside `train.py`, keep the same Python call:

```python
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](https://docs.pytorch.org/docs/stable/elastic/run.html).

## 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](https://docs.pytorch.org/docs/stable/distributed.html) 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](https://docs.pytorch.org/xla/master/learn/pytorch-on-xla-devices.html).

Example:

```python
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](https://docs.pytorch.org/xla/master/learn/pytorch-on-xla-devices.html) and [XLA AMP guide](https://docs.pytorch.org/xla/master/perf/amp.html).

### 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:

```python
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:

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

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

```json
{"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

```python
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:

```python
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

```python
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:

```python
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

```python
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

```bash
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

```text
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](https://docs.pytorch.org/docs/stable/distributed.html).

## Version

```text
Anthracite 1.5.5
```

License: MIT.
