aetherlm 0.1.0__tar.gz

aetherlm-0.1.0/PKG-INFO

@@ -0,0 +1,201 @@
+Metadata-Version: 2.4
+Name: aetherlm
+Version: 0.1.0
+Summary: A DeepSpeed-like training library for Google TPUs, built on JAX/Flax.
+Author: aetherBERT contributors
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/aetherBERT/aetherlm
+Project-URL: Repository, https://github.com/aetherBERT/aetherlm
+Keywords: jax,flax,tpu,deep-learning,distributed-training,bert,transformers
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+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.10
+Description-Content-Type: text/markdown
+Requires-Dist: jax[tpu]>=0.4.20
+Requires-Dist: flax>=0.10.0
+Requires-Dist: optax>=0.2.0
+Requires-Dist: orbax-checkpoint>=0.6.0
+Requires-Dist: transformers>=4.36.0
+Requires-Dist: datasets>=2.16.0
+Requires-Dist: wandb>=0.16.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: pyyaml>=6.0
+Provides-Extra: eval
+Requires-Dist: mteb>=1.12.0; extra == "eval"
+Requires-Dist: torch>=2.0.0; extra == "eval"
+Requires-Dist: scikit-learn>=1.3.0; extra == "eval"
+Requires-Dist: scipy>=1.11.0; extra == "eval"
+Provides-Extra: nmn
+Requires-Dist: nmn; extra == "nmn"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: jupyter>=1.0.0; extra == "dev"
+Provides-Extra: all
+Requires-Dist: aetherlm[dev,eval,nmn]; extra == "all"
+
+# 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

aetherlm-0.1.0/README.md

@@ -0,0 +1,159 @@
+# 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

aetherlm-0.1.0/aetherlm/__init__.py

@@ -0,0 +1,55 @@
+"""
+Aether - A DeepSpeed-like training library for Google TPUs
+
+Built on JAX/Flax for high-performance distributed training
+of transformer models on TPU pods.
+
+Quick Start:
+    import aetherlm
+
+    engine = aetherlm.initialize(config={
+        "model": {"model_type": "bert", "embed_dim": 768},
+        "training": {"mode": "mlm", "batch_size": 32},
+    })
+
+    for batch in train_data:
+        loss = engine.train_step(batch)
+"""
+
+__version__ = "0.1.0"
+
+from aetherlm.core.config import (
+    AetherConfig,
+    ModelConfig,
+    TrainingConfig,
+    OptimizerConfig,
+    ContrastiveConfig,
+    CausalConfig,
+    CheckpointConfig,
+    LoggingConfig,
+    EvalConfig,
+    TPUConfig,
+)
+from aetherlm.core.engine import AetherEngine, initialize
+from aetherlm.models.bert import AetherBERT
+from aetherlm.models.causal import AetherCausalLM
+
+__all__ = [
+    # Core
+    "initialize",
+    "AetherEngine",
+    "AetherConfig",
+    # Config
+    "ModelConfig",
+    "TrainingConfig",
+    "OptimizerConfig",
+    "ContrastiveConfig",
+    "CausalConfig",
+    "CheckpointConfig",
+    "LoggingConfig",
+    "EvalConfig",
+    "TPUConfig",
+    # Models
+    "AetherBERT",
+    "AetherCausalLM",
+]

aetherlm-0.1.0/aetherlm/checkpoint/__init__.py

@@ -0,0 +1,3 @@
+"""Aether checkpoint management with Orbax."""
+
+from aetherlm.checkpoint.manager import CheckpointManager, save_checkpoint, load_checkpoint

aetherlm-0.1.0/aetherlm/checkpoint/manager.py

@@ -0,0 +1,123 @@
+"""
+Checkpoint Management
+
+Save and load model checkpoints using Orbax, with support for
+checkpoint rotation (keeping only the N most recent).
+"""
+
+import os
+from typing import Optional
+
+import flax.nnx as nnx
+import orbax.checkpoint as orbax
+
+from aetherlm.core.config import CheckpointConfig
+
+
+class CheckpointManager:
+    """
+    Manages model checkpointing with Orbax.
+
+    Features:
+        - Save/load model parameters
+        - Checkpoint rotation (max_to_keep)
+        - Automatic directory creation
+    """
+
+    def __init__(self, config: Optional[CheckpointConfig] = None):
+        self.config = config or CheckpointConfig()
+        self.checkpoint_dir = os.path.abspath(self.config.dir)
+        self.max_to_keep = self.config.max_to_keep
+        self._saved_checkpoints = []
+
+        os.makedirs(self.checkpoint_dir, exist_ok=True)
+
+    def save(self, model: nnx.Module, path: Optional[str] = None, step: Optional[int] = None):
+        """
+        Save a model checkpoint.
+
+        Args:
+            model: NNX model to save.
+            path: Full path to save to. If None, auto-generates from step.
+            step: Training step number (used for auto-naming).
+        """
+        if path is None:
+            if step is not None:
+                path = os.path.join(self.checkpoint_dir, f"step_{step}")
+            else:
+                path = os.path.join(self.checkpoint_dir, "latest")
+
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+
+        checkpointer = orbax.PyTreeCheckpointer()
+        _, param_state, _ = nnx.split(model, nnx.Param, ...)
+        checkpointer.save(path, item=param_state)
+        checkpointer.close()
+
+        self._saved_checkpoints.append(path)
+        print(f"Checkpoint saved at {path}")
+
+        # Rotate old checkpoints
+        self._rotate()
+
+    def load(self, model: nnx.Module, path: str) -> nnx.Module:
+        """
+        Load model parameters from a checkpoint.
+
+        Args:
+            model: NNX model (used as template for parameter shapes).
+            path: Path to checkpoint directory.
+
+        Returns:
+            Model with restored parameters.
+        """
+        print(f"Loading model from {path}...")
+
+        _, params_template, _ = nnx.split(model, nnx.Param, ...)
+        checkpointer = orbax.PyTreeCheckpointer()
+        restored_params = checkpointer.restore(path, item=params_template)
+
+        if restored_params is None:
+            raise ValueError(f"Could not restore parameters from {path}")
+
+        nnx.update(model, restored_params)
+        print("Model loaded successfully.")
+        return model
+
+    def _rotate(self):
+        """Remove old checkpoints beyond max_to_keep."""
+        if self.max_to_keep <= 0:
+            return
+
+        while len(self._saved_checkpoints) > self.max_to_keep:
+            old_path = self._saved_checkpoints.pop(0)
+            if os.path.exists(old_path):
+                import shutil
+                try:
+                    shutil.rmtree(old_path)
+                    print(f"Removed old checkpoint: {old_path}")
+                except Exception as e:
+                    print(f"Warning: could not remove {old_path}: {e}")
+
+
+def save_checkpoint(model: nnx.Module, path: str):
+    """Convenience function to save a model checkpoint."""
+    os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True)
+    checkpointer = orbax.PyTreeCheckpointer()
+    _, param_state, _ = nnx.split(model, nnx.Param, ...)
+    checkpointer.save(path, item=param_state)
+    checkpointer.close()
+    print(f"Checkpoint saved at {path}")
+
+
+def load_checkpoint(model: nnx.Module, path: str) -> nnx.Module:
+    """Convenience function to load a model checkpoint."""
+    print(f"Loading model from {path}...")
+    _, params_template, _ = nnx.split(model, nnx.Param, ...)
+    checkpointer = orbax.PyTreeCheckpointer()
+    restored_params = checkpointer.restore(path, item=params_template)
+    if restored_params is None:
+        raise ValueError(f"Could not restore parameters from {path}")
+    nnx.update(model, restored_params)
+    print("Model loaded successfully.")
+    return model

aetherlm-0.1.0/aetherlm/cli/__init__.py

@@ -0,0 +1,3 @@
+"""Aether CLI entry point."""
+
+from aetherlm.cli.main import main