aetherlm 0.1.0
A DeepSpeed-like training library for Google TPUs, built on JAX/Flax.
License: Apache-2.0
Requires Python: >=3.10
Keywords: jax, flax, tpu, deep-learning, distributed-training, bert, transformers

# AetherLM

A DeepSpeed-like training library for Google TPUs, built on JAX/Flax.

AetherLM abstracts away the complexity of distributed training on TPU pods, providing a simple `initialize()` API for pretraining and fine-tuning transformer models.

## Features

- **One-call initialization** — `aetherlm.initialize(config)` sets up mesh, model, optimizer, and metrics
- **TPU-native** — Built on JAX with automatic TPU topology detection and optimal sharding
- **Multiple training modes** — MLM pretraining, contrastive learning, causal LM
- **Built-in models** — `AetherBERT` (bidirectional) and `AetherCausalLM` (autoregressive with generation)
- **YAML/dict configuration** — DeepSpeed-style config system
- **MTEB evaluation** — Integrated embedding evaluation with 56 tasks
- **Checkpoint management** — Automatic saving, rotation, and restore with Orbax

## Quick Start

```python
import aetherlm

# Configure and initialize (auto-detects TPU)
engine = aetherlm.initialize(config={
    "model": {"model_type": "bert", "embed_dim": 768, "num_layers": 12},
    "training": {"mode": "mlm", "batch_size": 32, "learning_rate": 1e-4},
    "tpu": {"precision": "bf16"},
})

# Load data
from aetherlm.data import load_mlm_datasets
train_data, val_data = load_mlm_datasets(maxlen=512, mask_prob=0.15, vocab_size=50265)

# Train (handles sharding, logging, checkpointing)
engine.train(train_data, val_data)
```

## Installation

```bash
# From PyPI
pip install aetherlm

# With eval support (MTEB, sklearn, scipy)
pip install aetherlm[eval]

# From source
pip install -e ".[all]"
```

## Training Modes

### MLM Pretraining (BERT-style)

```bash
aetherlm --mode mlm --config configs/default.yaml
```

### Causal Language Modeling (GPT-style)

```bash
aetherlm --mode causal --config configs/causal_small.yaml
```

### Contrastive Learning (requires checkpoint)

```bash
aetherlm --mode contrastive --checkpoint ./checkpoints/step_10000
```

### Evaluation

```bash
# Quick MTEB eval (3 tasks, ~1 min)
aetherlm --mode eval --checkpoint ./checkpoints/step_10000 --mteb_preset quick

# Full leaderboard (56 tasks)
aetherlm --mode eval --checkpoint ./checkpoints/step_10000 --mteb_preset leaderboard
```

## Configuration

Aether uses a dataclass-based config system. Create configs from YAML, JSON, or Python dicts:

```python
from aetherlm import AetherConfig

# From YAML
config = AetherConfig.from_yaml("configs/default.yaml")

# From dict
config = AetherConfig.from_dict({
    "model": {"embed_dim": 768, "num_layers": 12},
    "training": {"mode": "mlm", "batch_size": 32},
})

# Save for reproducibility
config.to_yaml("my_experiment.yaml")
```

See `configs/` for example configurations.

## Project Structure

```
aetherlm/
    __init__.py           # Top-level API: initialize(), AetherConfig, models
    core/
        config.py         # Dataclass configuration system
        engine.py         # Training engine (the heart of the library)
        sharding.py       # Automatic mesh sharding
    models/
        base.py           # Abstract model interface + utilities
        transformer.py    # Transformer blocks, embeddings
        bert.py           # AetherBERT (bidirectional MLM)
        causal.py         # AetherCausalLM (autoregressive + generation)
    optim/
        optimizers.py     # Optimizer factory (AdamW, Adam, SGD)
        schedules.py      # LR schedules (warmup-cosine, linear, constant)
        switching.py      # Plateau detection + optimizer switching
    data/
        pipeline.py       # Tokenizer caching, batch iterators
        mlm.py            # MLM masking and dataset processing
        contrastive.py    # Contrastive pair creation (self-supervised + AllNLI)
        causal.py         # Causal LM next-token prediction format
    losses/
        mlm.py            # Efficient gather-based MLM loss
        contrastive.py    # SimCLR-style contrastive loss
        causal.py         # Causal LM cross-entropy loss
    training/
        steps.py          # JIT-compiled train/eval steps for all modes
    checkpoint/
        manager.py        # Orbax checkpoint save/load/rotation
    metrics/
        tracker.py        # Throughput, loss, ETA, WandB logging
    eval/
        tasks.py          # MTEB task lists and presets
        mteb.py           # MTEB EncoderProtocol wrapper
        runner.py         # Evaluation orchestrators
    tpu/
        topology.py       # TPU detection, mesh creation
        precision.py      # bfloat16/mixed precision config
    cli/
        main.py           # CLI entry point
configs/                  # Example YAML configurations
notebooks/                # Tutorial notebooks
```

## Notebooks

| Notebook | Description |
|----------|-------------|
| `01_quickstart.ipynb` | Core features: models, config, initialize |
| `02_mlm_pretraining.ipynb` | Full MLM pretraining pipeline |
| `03_causal_lm.ipynb` | Causal LM training + text generation |
| `04_evaluation.ipynb` | MTEB and custom evaluation |

## License

Apache 2.0