Metadata-Version: 2.4
Name: composelm
Version: 0.3.1
Summary: One-line configurable, production-ready modular Transformer library
Author: ComposeLM Contributors
License: Apache-2.0
Project-URL: Homepage, https://github.com/DW-dev-UE/ComposeLM
Project-URL: Repository, https://github.com/DW-dev-UE/ComposeLM
Project-URL: Documentation, https://github.com/DW-dev-UE/ComposeLM/tree/main/docs
Project-URL: Issues, https://github.com/DW-dev-UE/ComposeLM/issues
Keywords: llm,transformer,pytorch,modular
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
License-File: LICENSE
License-File: NOTICE
Requires-Dist: torch>=2.1
Requires-Dist: numpy>=1.24
Requires-Dist: pyyaml>=6.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0; extra == "dev"
Requires-Dist: ruff>=0.5; extra == "dev"
Requires-Dist: mypy>=1.8; extra == "dev"
Requires-Dist: black>=24.0; extra == "dev"
Provides-Extra: einops
Requires-Dist: einops>=0.7; extra == "einops"
Provides-Extra: datasets
Requires-Dist: datasets>=2.14; extra == "datasets"
Provides-Extra: all
Requires-Dist: composelm[datasets,dev,einops]; extra == "all"
Dynamic: license-file

<div align="right">
  <a href="README.md"><img src="https://img.shields.io/badge/English-24292f?style=for-the-badge" alt="English"></a>
  <a href="README.ko.md"><img src="https://img.shields.io/badge/%ED%95%9C%EA%B5%AD%EC%96%B4-d0d7de?style=for-the-badge" alt="한국어"></a>
  <a href="README.ja.md"><img src="https://img.shields.io/badge/%E6%97%A5%E6%9C%AC%E8%AA%9E-d0d7de?style=for-the-badge" alt="日本語"></a>
</div>

# ComposeLM

**A decoder-only Transformer library you assemble from a single config**

[![CI](https://github.com/DW-dev-UE/ComposeLM/actions/workflows/ci.yml/badge.svg)](https://github.com/DW-dev-UE/ComposeLM/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/composelm.svg)](https://pypi.org/project/composelm/)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![PyTorch 2.1+](https://img.shields.io/badge/PyTorch-2.1+-ee4c2c.svg)](https://pytorch.org/)
[![License: Apache-2.0](https://img.shields.io/badge/License-Apache%202.0-green.svg)](LICENSE)
[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)

GPT · Llama · Mistral · Gemma · Qwen · DeepSeek · Phi all look different on the surface, but take
them apart and the differences come down to four axes: **attention variant / FFN gate / norm type /
positional encoding**. ComposeLM splits those axes into independent modules and lets you pick a
combination through a single `ModelConfig`. Training (`Trainer`) and inference (`generate`) inherit
the very same config.

```python
from composelm import build_model

model = build_model("llama3", d_model=2048, n_layers=16, n_heads=16, n_kv_heads=4)
```

> [!IMPORTANT]
> ComposeLM is at **0.3.1 · alpha**. Verified on CPU, a single H100, and 2× H100 (FSDP data parallel, 13 archs — see [Benchmarks](#benchmarks)). Multi-node is untested and no pretrained weights ship. Read [Caveats and known limitations](#caveats-and-known-limitations) before wiring it into production or publishing benchmark numbers.

---

## Table of contents

- [What this library is for](#what-this-library-is-for)
- [Installation](#installation)
- [Quick start](#quick-start)
- [Configuration: ModelConfig](#configuration-modelconfig)
- [Architecture](#architecture)
- [Component catalog](#component-catalog)
- [Training](#training)
- [Inference](#inference)
- [HuggingFace weight conversion](#huggingface-weight-conversion)
- [Distributed training](#distributed-training)
- [Benchmarks](#benchmarks)
- [Caveats and known limitations](#caveats-and-known-limitations) ⚠️
- [Development](#development)
- [Roadmap](#roadmap)
- [License](#license)

---

## What this library is for

| | |
|---|---|
| **What it does** | Assembles a decoder-only Transformer from a one-line config, then trains and runs inference from that same config |
| **Why** | So you stop forking and patching model code every time you want to try an architecture |
| **What it does not do** | Ship pretrained weights, train tokenizers, serve at vLLM scale, or handle encoders / multimodal |

The design principles fit in four lines.

- **SSOT** — every hyperparameter lives in `ModelConfig`. No magic numbers buried inside layers.
- **Swappable** — attention / FFN / norm / positional embedding / residual share one contract and interchange freely.
- **Explicit fallbacks** — when a kernel is missing, nothing dies quietly; it falls back and logs.
- **Mainstream defaults** — use it without thinking and you get `llama3 + SwiGLU + RMSNorm + RoPE + GQA`.

### Current status

| Item | Status |
|------|--------|
| Package | **0.3.1** — the P0–P15 path plus the hardware-validated multi-GPU training path |
| Tests | **516** CPU `pytest` cases (~5 s). Unit · integration · Trainer · infer · convert · all-arch smoke |
| Static checks | `ruff` · `mypy` clean |
| Benchmarks | **Measured on an H100 and on 2× H100** — speed, accuracy, training and multi-GPU scaling across 13 archs; see [Benchmarks](#benchmarks) |
| Distribution | `pip install composelm` — released from the [publish workflow](.github/workflows/publish.yml) |
| Validation | All archs measured on a single H100 plus real-weight PPL parity; FSDP data parallel validated on 2× H100 (13 archs, WikiText-103). EP all-to-all and multi-node remain unvalidated |

The design document is [Architecture.md](Architecture.md); API summaries live in [docs/](docs/).

---

## Installation

```bash
pip install composelm
```

If you plan to contribute or modify the code, use an editable install.

```bash
git clone https://github.com/DW-dev-UE/ComposeLM.git
cd ComposeLM
pip install -e ".[dev]"
```

| Group | Packages | Without them |
|-------|----------|--------------|
| Required | `torch>=2.1`, `numpy`, `pyyaml` | — |
| Dev | `pytest`, `ruff`, `mypy`, `black` (`[dev]`) | Cannot run tests or checks |
| Performance (optional) | `flash-attn`, `liger-kernel` | Falls back to pure PyTorch SDPA / RMSNorm |
| Distributed (optional) | `deepspeed` | ImportError when using `deepspeed_config` |
| Conversion (optional) | `transformers`, `safetensors` | Only the offline `state_dict=` path works |
| Examples (optional) | `datasets`, `einops` | Some examples are skipped |

> [!IMPORTANT]
> Optional dependencies are **not installed automatically.** `flash-attn` in particular requires you to pick a wheel whose CUDA, PyTorch, ABI and SM support all line up; if any one of them is off, the install fails or quietly breaks at import time. See [Caveats](#kernels--acceleration) for details.

Verify the install:

```bash
python -c "import composelm; print(composelm.__version__)"
python -m pytest tests -q
```

---

## Quick start

### Build a model

```python
from composelm import build_model

model = build_model(arch="llama3", d_model=256, n_layers=4, n_heads=8,
                    n_kv_heads=2, vocab_size=32000, max_seq_len=512, precision="fp32")
print(f"{model.num_parameters():,} params")
```

`build_model` also accepts a `ModelConfig` object directly.

```python
from composelm import ModelConfig, build_model

cfg = ModelConfig.from_arch("mistral", d_model=512, n_layers=8, n_heads=8, vocab_size=32000)
model = build_model(cfg)
model = build_model(cfg, n_layers=12)
```

### Generation

```python
import torch
from composelm import generate

ids = torch.randint(0, 32000, (1, 16))
out = generate(model, ids, max_new_tokens=32, do_sample=True, temperature=0.8, top_p=0.9)
```

### A taste of training

The full option set is covered in [Training](#training) below. Here is the minimal form.

```python
from composelm import Trainer

# train_dataset=None generates synthetic tokens, so this runs with no data prepared
Trainer(model=model, train_dataset=None, batch_size=4, gradient_accumulation_steps=8,
        learning_rate=3e-4, max_steps=2000, output_dir="./composelm-output").train()
```

### A few other combinations

```python
mla = build_model("deepseek", d_model=256, n_layers=2, n_heads=8, vocab_size=1000, precision="fp32")
gpt = build_model("gpt", d_model=768, n_layers=12, n_heads=12, vocab_size=50257)

moe = build_model("custom", d_model=512, n_layers=4, n_heads=8, vocab_size=8000,
                  ffn_type="moe", num_experts=8, num_experts_per_tok=2, precision="fp32")
```

Runnable scripts live in [`examples/`](examples/).

```bash
python examples/custom_arch.py
python examples/train_wikitext.py
python examples/generate_sample.py
python examples/moe_train.py
python examples/speculative_sample.py
python examples/convert_roundtrip.py
```

---

## Configuration: ModelConfig

Every hyperparameter converges here. The merge order is:

```
ARCH_PRESETS[arch] → user overrides → derived fields → validation → assembly
```

Derived fields are `head_dim` (= `d_model // n_heads`), `intermediate_size` (= `ffn_multiplier`
applied, then rounded up to a multiple of `ffn_multiple_of`), the `n_kv_heads` default, and the MLA
rank defaults.

### Key fields

| Field | Default | Description |
|-------|---------|-------------|
| `arch` | `"llama3"` | Preset key. Unknown keys raise |
| `d_model` / `n_layers` / `n_heads` | preset | `custom` requires you to set them |
| `n_kv_heads` | `n_heads` | GQA/MQA. `n_heads % n_kv_heads == 0` |
| `head_dim` | `d_model // n_heads` | Must be even for rope/yarn |
| `vocab_size` | preset | Required |
| `ffn_type` | `"swiglu"` | `relu` `gelu` `silu` `geglu` `swiglu` `reglu` `moe` |
| `intermediate_size` | auto | If unset: `ffn_multiplier` × `d_model` → multiple of `ffn_multiple_of` |
| `ffn_multiplier` / `ffn_multiple_of` | `8/3` / `256` | Llama-family convention |
| `norm` / `norm_placement` / `norm_eps` | `rmsnorm` / `pre` / `1e-5` | `post` is experimental |
| `pos_emb` | `"rope"` | `absolute` `sinusoidal` `rope` `yarn` `alibi` `none` |
| `rope_theta` / `rope_scaling` | `10000.0` / `None` | See the scaling table below |
| `max_seq_len` / `original_max_seq_len` | `4096` / `None` | Length before YaRN extension |
| `attention_type` | preset | `mha` `gqa` `mqa` `mla` |
| `sliding_window` | `None` | `None` means full causal |
| `attention_dropout` / `residual_dropout` / `embed_dropout` | `0.0` | `[0, 1]` |
| `use_bias` / `qkv_bias` / `attention_output_bias` / `ffn_bias` | `False` / inherited | Fine-grained control |
| `use_flash_attn` | `"auto"` | `True` / `False` / `"auto"` |
| `precision` | `"bf16_mixed"` | Preset string or dict |
| `tie_word_embeddings` | `True` | Share input and output embeddings |
| `num_experts` / `num_experts_per_tok` / `moe_aux_loss_coef` | `None` / `None` / `1e-2` | Required for `ffn_type="moe"` |
| `mla_kv_lora_rank` / `mla_q_lora_rank` / `mla_rope_head_dim` | `max(32, d_model//8)` / `None` / `head_dim` | MLA only |
| `loss_chunk_size` | `None` | Chunked CE — lm_head+loss per N-token chunk, full logits never materialized |
| `activation_checkpointing` | `False` | Recompute each block in backward — large activation-memory savings for extra backward compute |
| `use_compile` / `compile_mode` | `False` / `"default"` | `torch.compile` |
| `debug` | `False` | Asserts plus a readable attention path |

The complete field list and validation rules are in [Architecture.md §4](Architecture.md).

### arch presets

Each preset carries representative values for `d_model`, `n_layers` and `vocab_size`, and in
practice you almost always override them. Here are just the structural axes.

| `arch` | FFN | Norm | PosEmb | Attention | Notes |
|--------|-----|------|--------|-----------|-------|
| `gpt` | gelu | layernorm | absolute | mha | `use_bias=True`, the classic stack |
| `llama` / `llama1` / `llama2` | swiglu | rmsnorm | rope | mha | theta 10k |
| `llama3` | swiglu | rmsnorm | rope | **gqa** (16/4) | theta **500k**, the recommended default |
| `mistral` | swiglu | rmsnorm | rope | gqa (16/4) | `sliding_window=4096` |
| `gemma` / `gemma2` | **geglu** | rmsnorm | rope | mha / gqa | gemma2 includes a sliding window |
| `qwen` / `qwen2` | swiglu | rmsnorm | rope | mha / gqa (16/2) | qwen2 uses `qkv_bias=True`, theta 1M |
| `deepseek` | swiglu | rmsnorm | rope | **mla** | `mla_kv_lora_rank=256` |
| `phi` / `phi3` | swiglu | rmsnorm | rope | mha | phi3 uses `sliding_window=2047` |
| `custom` | swiglu | rmsnorm | rope | mha | You supply every dimension field |

```python
from composelm import list_archs, resolve_arch
list_archs()
resolve_arch("qwen2")
```

### YAML

```yaml
# configs/mistral_small.yaml
arch: mistral
d_model: 2048
n_layers: 16
n_heads: 16
n_kv_heads: 4
vocab_size: 32000
sliding_window: 4096
max_seq_len: 8192
precision: bf16_mixed
```

```python
from composelm import load_config_yaml, save_config_yaml, build_model

cfg = load_config_yaml("configs/mistral_small.yaml")
model = build_model(cfg)
save_config_yaml(cfg, "my_config.yaml")
```

---

## Architecture

### Directory layout

```text
composelm/
├── config.py          ModelConfig — SSOT schema · validation · serialization
├── build.py           build_model factory
├── types.py           shared types · allowed enum sets
├── models/
│   ├── registry.py    arch → default field dict
│   ├── transformer.py ComposeTransformer · ComposeForCausalLM
│   └── block.py       TransformerBlock (pre / post-norm)
├── layers/
│   ├── attention.py   MHA · GQA · MQA · KV cache · causal/sliding masks
│   ├── mla.py         Multi-Latent Attention (DeepSeek family)
│   ├── ffn.py         dense · gated FFN
│   ├── moe.py         top-k router · experts · aux loss
│   ├── norm.py        RMSNorm · LayerNorm
│   ├── pos_emb.py     absolute · sinusoidal · RoPE · YaRN · ALiBi
│   ├── residual.py    residual + dropout
│   ├── embedding.py   token embedding · tie
│   └── linear.py      bias policy · init hooks
├── kernels/           flash detection · liger fused · torch.compile
├── precision/         dtype policy · autocast · GradScaler
├── train/             Trainer · optimizer/scheduler · checkpoint
├── infer/             generate · continuous batching · speculative
├── convert/           HF key mapping · load/save
├── distributed/       FSDP · DeepSpeed · expert parallel
└── bench/             benchmark runner · metrics
```

### Forward flow

```text
input_ids [B, T]
   │
   ├─ TokenEmbedding          (+ absolute / sinusoidal PE)
   ├─ embed_dropout
   │
   ├─ for layer in layers:            ← shown for pre-norm
   │     h = Norm(x)
   │     h = Attention(h)             QKV → RoPE/YaRN(q,k) or ALiBi bias
   │                                  → causal [+ sliding window] mask
   │                                  → flash / SDPA / manual
   │     x = x + Dropout(h)
   │
   │     h = Norm(x)
   │     h = FFN(h) or MoE(h)
   │     x = x + Dropout(h)
   │
   ├─ FinalNorm
   └─ lm_head → logits [B, T, vocab_size]
                (CE loss if labels are given; aux_loss added for MoE)
```

With `norm_placement="post"` this becomes `x = Norm(x + Attn(x))`.

### Forward return contract

`ComposeForCausalLM.forward` returns either a tensor or a dict depending on the situation.

| Condition | Returns |
|-----------|---------|
| No labels, no cache, not MoE | `logits` tensor `[B, T, V]` |
| Otherwise | `{"logits", "loss"?, "aux_loss"?, "past_key_values"?}` |

With `use_cache=True` you get a per-layer list of `(k, v)` as `past_key_values`. When a sliding
window is in play the cache is trimmed to the window size, and a `KVCache` (a tuple subclass that
also carries `seen_tokens` for absolute position tracking) is returned instead. It unpacks like an
ordinary tuple.

With `loss_chunk_size` set **and labels given**, lm_head + cross-entropy run in N-token chunks with
recompute, so the full `[B, T, V]` logits are never materialized — the dict then carries
`logits=None`. At large vocabularies this removes the dominant activation peak (measured: a ~1B
llama3 dropped 17+ GB of peak train memory going from 128k- to 50k-vocab logits). Without labels
(generation) the path is unchanged.

---

## Component catalog

### Attention

| Type | Q heads | KV heads | Notes |
|------|---------|----------|-------|
| `mha` | `n_heads` | `n_heads` | Standard |
| `gqa` | `n_heads` | `n_kv_heads` | KV repeated per group |
| `mqa` | `n_heads` | 1 | Single KV head |
| `mla` | `n_heads` | low-rank latent | K/V compressed into a `kv_lora_rank` latent space, then reconstructed per head |

Setting `mla_rope_head_dim < head_dim` gives you the nope/rope split where RoPE applies to only part
of the head dimension (the DeepSeek approach).

With `sliding_window=W`, each query attends to the most recent W keys only. Combined with the KV
cache, the cache is also trimmed to W, so decoding memory stays constant.

### FFN

| `ffn_type` | Structure |
|------------|-----------|
| `relu` `gelu` `silu` | `Linear → act → Linear` |
| `geglu` `swiglu` `reglu` | `down(act(gate) * up)` — gate/up are fused into a single `gate_up` matrix |
| `moe` | top-k routing + per-expert gated FFN + load-balancing aux loss |

The MoE aux loss follows the Switch Transformer form (`E × Σ f_i·P_i`) and evaluates to exactly 1.0
under uniform routing. It is scaled by `moe_aux_loss_coef` and added to the task loss.

### Positional encoding

| `pos_emb` | Applied at |
|-----------|------------|
| `absolute` | Learned `nn.Embedding`, added to the token embedding |
| `sinusoidal` | Fixed sin/cos table, added to the token embedding |
| `rope` | Rotates q, k |
| `yarn` | Rotates q, k + NTK-by-parts ramp + attention temperature |
| `alibi` | Per-head slope added as attention bias |
| `none` | None (experimental) |

`rope_scaling` accepts either a string or a dict.

| Type | Required fields | Behavior |
|------|-----------------|----------|
| `linear` (= `pi`, `position_interpolation`) | `factor` | `inv_freq /= factor` |
| `ntk` (= `ntk_aware`) | `factor` | `base' = θ · factor^(d/(d-2))` |
| `dynamic` (= `dynamic_ntk`) | `factor` | Recomputes base from sequence length; short inputs get the original back |
| `llama3` | `factor`, `low_freq_factor`, `high_freq_factor`, `original_max_position_embeddings` | Smooth interpolation per frequency band |
| yarn | select via `pos_emb="yarn"` | `factor`, `beta_fast`, `beta_slow`, `attention_factor`, … |

```python
build_model("llama3", d_model=256, n_layers=2, n_heads=8, vocab_size=1000,
            rope_scaling={"type": "ntk", "factor": 4.0}, max_seq_len=16384)
```

### Precision presets

| Preset | compute | weights | master | GradScaler |
|--------|---------|---------|--------|------------|
| `fp32` | fp32 | fp32 | fp32 | — |
| `bf16_mixed` | bf16 | bf16 | fp32 | — |
| `fp16_mixed` | fp16 | fp16 | fp32 | ✔ |
| `fp8_mixed` | **falls back to bf16** | bf16 | fp32 | — |

You can also specify it granularly as a dict.

```python
precision = {"compute": "bf16", "weights": "bf16", "master_weights": "fp32", "gradients": "bf16"}
```

> [!WARNING]
> **`fp8_mixed` is not real FP8 yet.** Internally it maps to bf16, and running it actually emits this warning:
>
> ```text
> fp8 dtype requested without an FP8 backend; mapping to bf16
> ```
>
> Why it works this way, and which paths deserve extra care, is covered in [Caveats](#precision) below.

---

## Training

```python
from composelm import build_model, Trainer

model = build_model("llama3", d_model=512, n_layers=8, n_heads=8, n_kv_heads=2, vocab_size=32000)

# my_dataset yields dicts with input_ids; labels and attention_mask are optional
result = Trainer(model=model, train_dataset=my_dataset,
                 batch_size=8, gradient_accumulation_steps=4, learning_rate=3e-4,
                 weight_decay=0.1, warmup_steps=100, max_steps=10000, max_grad_norm=1.0,
                 precision="bf16_mixed", save_steps=1000,
                 output_dir="./composelm-output").train()

print(result["history"][-1])   # {"step": ..., "loss": ..., "lr": ...}
```

Included: gradient accumulation, autocast, fp16 GradScaler, gradient clipping, warmup + cosine
schedule, checkpoints (model, optimizer, scheduler, scaler and RNG state), and FSDP/DeepSpeed hooks.

Weight decay groups are split automatically. **1-D tensors (all biases and norm gains) and
embeddings** (token embedding plus learned absolute PE) are excluded from decay; only projection
matrices get it. Add name-based rules with `build_optimizer(model, no_decay_keywords=(...))`.

> [!TIP]
> Leaving `train_dataset=None` builds a synthetic token dataset matching the config's `vocab_size` and runs a smoke training loop. It is a fast way to confirm the model structure and the training loop work before you prepare any data.

Checkpoints can also be used standalone.

```python
from composelm.train import save_checkpoint, load_checkpoint

save_checkpoint("ckpt.pt", model, optimizer=opt, scheduler=sched, step=1000)
meta = load_checkpoint("ckpt.pt", model, optimizer=opt, scheduler=sched)
```

> [!NOTE]
> If a checkpoint is higher precision than the model, `load_checkpoint` promotes the model dtype first and then loads — meaning fp32 weights loaded into a bf16 model are not truncated. Pass `preserve_source_dtype=False` if you do not want this.

---

## Inference

### generate

```python
from composelm import generate

out = generate(model, input_ids, max_new_tokens=128,
               do_sample=True, temperature=0.8, top_k=50, top_p=0.95,
               eos_token_id=2, pad_token_id=0, attention_mask=mask, use_kv_cache=True)
```

`do_sample=False` or `temperature=0` gives greedy decoding. With `use_kv_cache=True` it decodes
incrementally, and if the model does not accept cache arguments it automatically drops back to full
re-forward.

> [!NOTE]
> When every sequence in the batch emits EOS, generation stops early. That means the returned sequence can be shorter than `max_new_tokens` — be careful if your post-processing assumes a fixed length.

### Continuous batching

```python
from composelm.infer import GenerateRequest, continuous_generate

reqs = [
    GenerateRequest(request_id="a", input_ids=[1, 2, 3], max_new_tokens=32),
    GenerateRequest(request_id="b", input_ids=[4, 5], max_new_tokens=64, temperature=0.7),
]
results = continuous_generate(model, reqs, max_batch_size=8)
```

Finished requests leave the batch and waiting ones take their slot. Use `ContinuousBatcher`
directly if you want `step()`-level control.

### Speculative decoding

```python
from composelm.infer import speculative_generate

out = speculative_generate(target_model, draft_model, input_ids,
                           max_new_tokens=128, draft_tokens=4,
                           do_sample=True, temperature=0.8)
```

The draft proposes γ tokens and the target verifies them in one pass. Each token is accepted with
`min(1, p_target/p_draft)`; on rejection a token is resampled from the residual distribution. If all
are accepted you get one bonus token. Draft and target must share a vocabulary.

---

## HuggingFace weight conversion

This handles key mapping only. It works from a `state_dict` alone, with no network access.

```python
from composelm.convert import from_hf, to_hf_state_dict, save_hf_state_dict

hf_state = to_hf_state_dict(model)
save_hf_state_dict(model, "model.safetensors")

model = from_hf(state_dict=hf_state, arch="llama3", d_model=2048, n_layers=16,
                n_heads=16, n_kv_heads=4, vocab_size=32000)
model = from_hf("./llama-3-1b", arch="llama3")
```

When a `config.json` is present, `hidden_size` / `num_key_value_heads` / `rope_theta` /
`rope_scaling` / `sliding_window` and friends are read to build a `ModelConfig` automatically.

> [!NOTE]
> Only four `model_type` values support automatic config construction: `gpt2`, `llama`, `mistral`, `qwen2`. Anything else raises explicitly instead of guessing something close enough.

Fused SwiGLU is split/joined automatically between `gate_proj` + `up_proj` ↔ `gate_up`, and GPT-2's
`c_attn` fused QKV and Conv1D transposes are handled as well.

---

## Distributed training

```python
Trainer(model=model, use_fsdp=True, ...).train()
Trainer(model=model, deepspeed_config="ds_config.json", ...).train()

# FSDP tuning knobs — all optional; leaving them unset keeps the historical wrap
Trainer(model=model, use_fsdp=True,
        fsdp_auto_wrap=True,           # one FSDP unit per TransformerBlock
        fsdp_mixed_precision="bf16",   # bf16 param/grad gather + reduce (default: fp32)
        fsdp_sharding="full_shard",    # full_shard | shard_grad_op | no_shard
        fsdp_kwargs={...},             # raw FSDP kwargs, override anything above
        ...).train()
```

Measured on 2× H100 ([`bench_results/multigpu_20260725/`](bench_results/multigpu_20260725/)):
`fsdp_mixed_precision="bf16"` gave **+5.8% throughput and −6.1 GB peak** per rank on a ~1B model,
and `fsdp_auto_wrap=True` gathers/frees parameters per block — lowering peaks further and
unblocking `no_shard` replicas that OOM under the default flat wrap.

Launched under `torchrun`, it reads `LOCAL_RANK` to pick the device, and attaches a
`DistributedSampler` automatically for map-style datasets. If the process group is not initialized
yet, the FSDP wrap is skipped with a warning — which is what lets the same code run unchanged in a
single process.

> [!WARNING]
> `torchrun` only sets environment variables — **it does not create the process group.** The entry
> script must call `torch.distributed.init_process_group` (the bundled
> `composelm.distributed.init_distributed_from_env()` does exactly that); without it the FSDP wrap
> skips and every rank silently trains an independent, unsynchronized copy.

```bash
torchrun --standalone --nproc_per_node=2 examples/train_distributed.py
```

[`examples/train_distributed.py`](examples/train_distributed.py) owns the process-group lifecycle
(init on entry, destroy on exit) and runs unchanged as plain
`python examples/train_distributed.py` in a single process.

MoE expert-parallel hooks live in [`composelm/distributed/ep.py`](composelm/distributed/ep.py).
`shard_expert_indices()` computes round-robin shard indices, but **the actual token all-to-all
dispatch is not implemented yet.** See [Caveats](#moe) for details.

---

## Benchmarks

### H100 measurements (2026-07-25)

The 0.3.0 code was measured across every architecture on an **NVIDIA H100 80GB** (torch
2.11.0+cu128, Ubuntu 22.04, Python 3.10). Raw CSVs, run logs and reproduction scripts are all
preserved under [`bench_results/h100_20260725/`](bench_results/h100_20260725/), and the detailed
write-up is [H100_REPORT.md](bench_results/h100_20260725/H100_REPORT.md).

| Check | Result |
|-------|--------|
| Functionality (pytest, H100 environment) | **483 passed / 0 failed** (as measured; a cuDNN-fix regression test and later additions have since grown the suite to 516) |
| Numerical accuracy (real-weight PPL) | **≤ 0.03 % relative error** against the transformers reference for all four families |
| Trainability (13 archs) | **13/13 succeeded** at 500 steps under identical WikiText-103 settings, loss decreasing for every arch |
| Speed (13 archs) | All 145 measurements clean (0 errors; 7 OOMs at fp32@8192 on large models were recorded) |

#### Accuracy — perplexity on real HF checkpoints (WikiText-2 test, bf16)

Public checkpoints loaded with `from_hf` and A/B compared against a `transformers` reference model
under identical settings. Sliding window 1024 / stride 512.

| Model | ComposeLM | transformers | Δ |
|-------|----------:|-------------:|---:|
| GPT-2 (124M) | 24.606 | 24.613 | -0.007 |
| TinyLlama-1.1B | 7.413 | 7.413 | +0.0002 |
| Qwen2-0.5B | 12.732 | 12.732 | ±0.0000 |
| Mistral-7B-v0.1 | 5.015 | 5.016 | -0.0003 |

#### Speed — 13 arch presets (bf16, batch 8, tok/s)

Preset default dimensions (0.1B–3.7B). Prefill is forward throughput; decode is single-token
generation with the KV cache.

| arch | params | prefill@2048 | prefill@8192 | decode@2048 |
|------|-------:|-------------:|-------------:|------------:|
| gpt | 0.13B | 540k | 512k | 1,142 |
| llama / llama2 | 0.89B | 185k | 166k | 609–630 |
| llama3 | 0.98B | 182k | 163k | 591 |
| mistral | 0.79B | 179k | 123k | 554 |
| gemma | 1.45B | 134k | 124k | 561 |
| gemma2 | 2.11B | 65k | 34k | 322 |
| qwen | 1.54B | 115k | 103k | 331 |
| qwen2 | 1.37B | 126k | 113k | 388 |
| deepseek (MLA) | 0.92B | 181k | 163k | 589 |
| phi | 1.34B | 121k | 104k | 435 |
| phi3 | 3.72B | 46k | 32k | 218 |
| custom MoE (8E top2) | 0.61B | 290k | 450k | 254 |

With `torch.compile` we measured a further **×1.5–×3.0** speedup on prefill (gpt ×3.04,
llama3/mistral ×1.53, deepseek ×1.59; MoE only ×1.02, limited by the naive dispatch).

> [!NOTE]
> - **There are no FP8 numbers.** `fp8_mixed` falls back to bf16, so it was excluded from the matrix.
> - Measured without the `flash-attn` package installed (no wheel for torch 2.11). Attention went
>   through PyTorch SDPA (its built-in flash kernel), recorded in the CSVs as `flash_available=False`.
> - fp32 was measured with TF32 disabled.
> - PPL parity covers only the four families `from_hf` supports. For the other archs
>   (gemma, phi, deepseek, MoE) the evidence is unit tests plus measured training, not real-weight
>   equivalence.
> - These are single-GPU measurements; 2-GPU scaling results are in the next section.
>
> This measurement campaign surfaced two real bugs — transformers 5.x's nested `rope_parameters`
> not being supported (Qwen2 PPL off by +3.07), and a decode bottleneck from torch 2.11's cuDNN SDPA
> rebuilding its plan per KV length (up to ×36 slower) — both fixed, and the numbers above are from
> the re-run. The diagnostic logs and the pre-fix data are kept as-is under
> [`bench_results/h100_20260725/logs/`](bench_results/h100_20260725/logs/).

#### Reproducing

```bash
# 1) Speed: all archs × precisions × seq lengths (prefill + KV decode) — no corpus needed
python scripts/bench_all_archs.py --device cuda --profile preset \
  --precisions bf16_mixed fp16_mixed fp32 --seq-lens 2048 8192 --batch-size 8 \
  --output-csv bench_results/arch_sweep.csv

# 2) Accuracy: real-weight PPL A/B against transformers (WikiText-2 downloaded automatically)
python scripts/eval_hf_ppl.py --device cuda --dtype bf16 --compare-hf

# 3) Training: same-scale training curves for every arch (WikiText-103, GPT-2 BPE)
python scripts/train_curves.py --device cuda --steps 500 --batch-size 16
```

All three scripts write append-only CSVs, record failures such as OOM as status rows, and keep going.

### 2× H100 multi-GPU measurements (2026-07-25)

Training scaling for every arch preset on **2× NVIDIA H100 80GB** (NVLink, torch 2.7.0+cu128),
measured on **WikiText-103** (`Salesforce/wikitext`, GPT-2 BPE, vocab 50257 — parameter counts
therefore differ from preset defaults). Conditions: bf16_mixed, batch 8 per GPU, seq 2048,
30 measured steps. The full test suite and a 9-check distributed correctness gate (sampler
sharding, gradient sync, FSDP checkpoint save/resume, fp16 sharded scaler, EP guard) ran on the
host before any measurement. Raw CSV, per-arch logs, the corpus manifest (dataset id, tokenizer,
token count, cache SHA256) and repro scripts are under
[`bench_results/multigpu_final_20260725/`](bench_results/multigpu_final_20260725/); the write-up is
[REPORT.md](bench_results/multigpu_final_20260725/REPORT.md).

"recipe" below = the opt-in memory options: `activation_checkpointing` + `loss_chunk_size=4096` +
`fsdp_auto_wrap` + `fsdp_mixed_precision="bf16"`. Defaults stay exactly as before.

| arch | params | 1 GPU tok/s | 2 GPU tok/s | efficiency | recipe tok/s | recipe peak | default peak |
|------|-------:|------------:|------------:|-----------:|-------------:|------------:|-------------:|
| gpt | 0.13B | 140k | 265k | 0.95 | 220k | 3.7 GB | 20.1 GB |
| llama / llama2 | 0.93B | 39.7k | 78.0k | 0.98 | 63.8k | 9.9 GB | 50.4 GB |
| llama3 | 0.82B | 41.9k | 82.6k | 0.98 | 67.4k | 9.1 GB | 49.2 GB |
| mistral | 0.82B | 34.0k | 67.0k | 0.99 | 54.4k | 9.2 GB | 49.3 GB |
| gemma | 1.03B | 35.8k | 70.6k | 0.99 | 57.6k | 10.8 GB | 55.3 GB |
| gemma2 | 1.63B | OOM | 22.1k | — | 21.6k | 16.5 GB | 82.7 GB |
| qwen | 1.34B | 28.8k | 56.9k | 0.99 | 46.6k | 13.4 GB | 69.9 GB |
| qwen2 | 1.16B | 31.0k | 61.0k | 0.98 | 50.0k | 11.8 GB | 67.8 GB |
| deepseek (MLA) | 0.82B | 41.8k | 82.1k | 0.98 | 66.3k | 9.1 GB | 48.2 GB |
| phi | 1.34B | 28.1k | 55.2k | 0.98 | 45.6k | 13.4 GB | 69.9 GB |
| phi3 | 3.78B | OOM | OOM | — | **16.1k** | 37.9 GB | OOM |
| custom MoE (8E top2) | 0.64B | 64.6k | 117.5k | 0.91 | 87.6k | 10.3 GB | 30.2 GB |

> [!NOTE]
> - OOM rows are the library defaults hitting one card's 80 GB (fp32 master + Adam states +
>   stored activations). They are kept on purpose — the recipe column is the fix. With the
>   recipe, phi3 (3.78B) also trains on a **single** H100 (7.9k tok/s, 75.0 GB peak).
> - The recipe costs a consistent 17–19% throughput for a 75–82% lower per-rank peak, with
>   matching loss curves.
> - gpt's learned absolute PE is shared with every attention module, so per-block wrapping
>   falls back to the flat wrap automatically (logged); the rest of the recipe still applies.
> - MoE numbers use the sort-based grouped dispatch (×2.1 over the previous mask-loop dispatch
>   under identical settings). DeepSpeed and EP all-to-all were not part of this campaign.

```bash
# reproduce on any 2-GPU host (sync/install per scripts/ssh_bench.md)
bash bench_results/multigpu_final_20260725/repro_scripts/run_final.sh
```

### Single-run benchmark runner

```bash
python -m composelm.bench.runner --device auto --d-model 512 --n-layers 4 --seq-len 512 --batch-size 2
composelm-bench --device cuda --precision bf16_mixed --output-csv bench_results/run.csv
```

It measures forward latency, tokens/s, peak memory and parameter count, and appends to a CSV. It can
also be configured through environment variables such as `COMPOSELM_BENCH_D_MODEL`. The procedure for
syncing to a remote GPU instance, installing, running and pulling the CSV back is documented in
[`scripts/ssh_bench.md`](scripts/ssh_bench.md).


---

## Caveats and known limitations

> [!IMPORTANT]
> **FP8, MoE weight conversion and device movement** are the three that fail silently — no exception,
> just wrong numbers. Everything else raises or logs.

### Precision

> [!WARNING]
> **Despite the name, `fp8_mixed` does not perform real FP8 arithmetic.** With no FP8 GEMM backend (transformer-engine or similar) wired up, it **falls back to bf16** and logs a warning. PyTorch having a dtype like `torch.float8_e4m3fn` and PyTorch being able to actually run FP8 math are two different things.
>
> **Do not publish speed or memory numbers measured with `precision="fp8_mixed"` as "FP8 performance."** Those numbers are bf16 numbers. Putting `compute: "fp8"` directly into the precision dict lands on bf16 all the same.

- Calling `build_model` with `bf16_mixed` **creates the parameters already cast to bf16.**
  Afterwards `load_checkpoint` / `from_hf` promote the model when the source is higher precision to
  avoid truncation, but calling `model.load_state_dict()` yourself skips that protection.
- `Trainer` restores parameters to the master weight dtype (usually fp32) and then trains under
  autocast. So the dtype right after `build_model(precision="bf16_mixed")` differs from the dtype
  during training.
- On CUDA devices without BF16 support it falls back to `fp16_mixed` automatically. On CPU only the
  casting is allowed — do not expect an actual speedup.

### Kernels / acceleration

> [!NOTE]
> `flash-attn` is **not installed automatically.** You have to pick a wheel whose CUDA version, PyTorch version, ABI and supported SMs all match — for an H100 (SM90) you need a wheel that supports SM90 — and a source build can take anywhere from tens of minutes to over an hour. If it is missing or incompatible, it logs the following once and quietly falls back to PyTorch SDPA.
>
> ```text
> FlashAttention unavailable, falling back to SDPA: <original exception message>
> ```

- The flash path is **fairly narrow.** It is taken only when all of the following hold: CUDA tensors
  / no additive bias (i.e. no sliding window, ALiBi, padding mask or KV cache) / q length == k
  length. In other words it works **only for full-sequence causal prefill**; incremental decoding and
  windowed attention use SDPA. **MLA has no flash path at all.**
- With `debug=True` both flash and SDPA are skipped in favor of a manual matmul + softmax path. It
  reads well and carries asserts, but it is slow. Do not benchmark with it.
- `use_compile=True` actually compiles a trivial function at wrap time to check backend availability
  up front. Without a toolchain (e.g. Windows without MSVC `cl.exe`) it returns the eager model as
  is. If the backend is fine but the model itself fails to compile, **the first call propagates the
  exception** and subsequent calls fall back to eager — a deliberate choice so a partially executed
  forward is not applied twice.
- When `liger-kernel` is present, RMSNorm switches to the fused implementation. Otherwise pure
  PyTorch RMSNorm is used.

### Memory / masking

> [!NOTE]
> The causal mask is **materialized as a `[1, 1, T, KV]` tensor.** It is not block-sparse, so at long context the mask itself consumes memory. Sliding window works the same way — it computes the full `T × KV` scores and then masks out what falls outside the range — so **prefill FLOPs and mask memory do not go down.** What actually shrinks is only the KV cache during decoding (fixed at W).

- The RoPE cos/sin cache **grows automatically** when a requested length exceeds `max_seq_len`.
  Because it re-registers the buffer, this can trigger recompilation under `torch.compile`.
- `absolute` / `sinusoidal` positional encodings, by contrast, raise when you exceed `max_seq_len`.
- Using a sliding window with the KV cache trims the cache to W and tracks absolute positions through
  `seen_tokens`. Preserve that field if you assemble a cache by hand.

### MoE

- Dispatch **groups tokens by expert** (sort + bincount) and runs each expert once per forward on
  one contiguous batch, with a single host sync — the earlier `top_k × num_experts` mask loop with
  a device sync per iteration is gone. Experts remain separate modules (no grouped GEMM /
  stacked-weight kernel), so a very large expert count still costs one kernel launch each.

> [!WARNING]
> **`expert_parallel=True` together with `world_size > 1` raises `NotImplementedError`.**
>
> ```text
> expert_parallel=True with world_size > 1 requires token all-to-all, which is not implemented
> ```
>
> Silently replicating experts while token all-to-all is unimplemented would produce wrong results, so this is blocked on purpose. In a single process it runs all experts locally and just logs.

- The aux loss is overwritten on `MoELayer.last_aux_loss` on every forward and collected by the
  block. Because it is computed during inference too, a MoE model's forward **always returns a dict**
  rather than a tensor under the default `return_dict=True`. The return type differs from dense
  models, so branch accordingly at the call site.

### Inference

- `continuous_generate` is **right-padding** based, with no paged KV cache and no prefix sharing. It
  re-forwards everything on every step. It is for testing small models, not a serving stack.
  (`generate`, by contrast, reads `attention_mask` and **re-aligns valid tokens into left-padded
  form** before feeding the model. The two paths use different padding conventions, so take care when
  handling logit positions directly.)
- Speculative decoding **ends the block for the entire batch** as soon as one row rejects. Results
  stay correct, but acceptance rates drop on heterogeneous batches.

> [!CAUTION]
> Value validation for `attention_mask` and `position_ids` runs **only on CPU tensors**, because calling `.item()` on GPU at every decode step would force a synchronization. So passing out-of-range `position_ids` on GPU **dies on a device-side assert** instead of raising a friendly `ValueError`. Validate masks and positions once on CPU, then move them up.

### Conversion / compatibility

- **Presets do not guarantee HF checkpoint compatibility.** `arch="llama3"` means "Llama 3-style
  structure," not "eats Meta's Llama 3 weights as is." Weight compatibility is `convert/`'s job, and
  dimensions and layout have to match.

> [!WARNING]
> `to_hf_state_dict` **simply drops** parameters with no corresponding HF key (it logs a warning, but does not raise). MoE routers/experts, MLA projections and learned absolute PE all fall into this bucket. In other words, **exporting a MoE or MLA model to HF format silently loses those weights.** Check the logs before you export.

- The only `model_type` values `from_hf` can auto-configure from config.json are `gpt2`, `llama`,
  `mistral` and `qwen2`.
- In YaRN settings, the DeepSeek-family `mscale` and `mscale_all_dim` fields are unsupported and raise.
- MLA and YaRN are **reasonable approximations rather than 1:1 paper reproductions.** In particular,
  with partial RoPE (`mla_rope_head_dim < head_dim`), the YaRN attention temperature is an
  approximation applied to the whole score.

### Distributed

> [!WARNING]
> The FSDP path is validated on real hardware — 2× H100, all 13 archs, plus a 9-check correctness gate (sampler sharding, gradient sync, checkpoint save/resume, fp16 sharded scaler); see [Benchmarks](#benchmarks). **DeepSpeed integration remains at the hook level with no hardware validation**, and multi-node (more than one host) is untested. Verify those yourself before committing to a large run.

- `use_fsdp` and `deepspeed_config` cannot be used together.
- The FSDP path does not promote dtype when loading checkpoints (deliberately left out so it does not
  conflict with FSDP's MixedPrecision policy).

### Device movement happens in place

`Trainer` and `ContinuousBatcher` **move the model you hand them in place, without copying.** Get
the order wrong and it bites immediately:

```python
model = build_model("llama3", ...)          # cpu
Trainer(model=model, ...).train()           # model moves to cuda:0 if CUDA is available
generate(model, cpu_ids)                    # 💥 device mismatch

continuous_generate(model, reqs)            # device defaults to "cpu" → model is pulled back to cpu
```

> [!CAUTION]
> - Without an explicit `device` argument, `Trainer` grabs `cuda:{LOCAL_RANK}` when CUDA is available, and also changes parameter dtype to the master weight dtype.
> - `device` defaults to `"cpu"` for `ContinuousBatcher` / `continuous_generate`. Hand it a model that lives on the GPU without passing `device`, and the model gets dragged back down to CPU.

When running inference right after training, match inputs to
`next(model.parameters()).device`, and get into the habit of passing `device=` to the batcher
explicitly.

### Miscellaneous

- **This is still alpha.** The public API may change before it stabilizes.
- **No pretrained weights are distributed.** Everything starts randomly initialized.
- **There is no tokenizer.** Bring one from `transformers`.

> [!NOTE]
> `tie_word_embeddings` defaults to `True`. The lm_head and the token embedding **share the same tensor**, so modifying one changes the other. Be careful when inspecting or editing weights directly.

- A preset's `sliding_window` (mistral 4096, phi3 2047, …) is clamped down to `max_seq_len`
  automatically when you shrink `max_seq_len`. A `sliding_window` **you set explicitly**, on the
  other hand, raises rather than being clamped if it exceeds `max_seq_len`.

---

## Development

```bash
pip install -e ".[dev]"            # pytest · ruff · mypy · black

python -m pytest tests -q          # 516 cases, ~5 s (CPU)
python -m pytest tests -q -k moe   # a subset

ruff check composelm tests
mypy composelm --ignore-missing-imports
black composelm tests
```

Line length is set to **100** for both ruff and black in `pyproject.toml`. The same checks run in
[CI](.github/workflows/ci.yml) on every push and PR — tested on Linux × Python 3.10/3.11/3.12 plus
Windows × 3.11, including sdist/wheel builds and `twine check --strict`.

Tests are split by layer.

| File | Covers |
|------|--------|
| `test_config.py` `test_registry.py` `test_build.py` | Config validation · preset merging · assembly |
| `test_attention.py` `test_mla.py` `test_pos_emb.py` `test_yarn.py` | Attention · RoPE/YaRN numerics |
| `test_ffn.py` `test_moe.py` `test_norm.py` `test_embedding.py` | Per-layer |
| `test_block_forward.py` `test_all_arch_smoke.py` | Blocks · all-arch smoke |
| `test_precision.py` `test_kernels.py` | dtype policy · kernel fallbacks |
| `test_trainer_smoke.py` `test_trainer_edges.py` `test_checkpoint.py` | Training · checkpoints |
| `test_generate.py` `test_batching.py` `test_speculative.py` | Inference |
| `test_convert_hf.py` `test_distributed_hooks.py` `test_bench.py` | Conversion · distributed hooks · bench |
| `test_chunked_loss.py` `test_activation_checkpointing.py` | Chunked CE and activation-checkpoint parity vs the default paths |

Follow this order when adding a new component.

1. Implement it as a standalone module under `layers/` (no coupling to the training loop or I/O)
2. Add the `ModelConfig` field, the allowed enum, and the `validate_config` rule
3. Wire it into the `registry` / `build_*` factories
4. Write unit and shape tests
5. Update the tables in the README and [Architecture.md](Architecture.md)
6. Check the default path (`llama3`) for regressions

If a change breaks the public API, bump the major version and include migration notes.

---

## Roadmap

Everything listed under [Current status](#current-status) is implemented. Not there yet: EP token
all-to-all dispatch, multi-node training, and a paged KV cache.

---

## License

**Apache License 2.0** — use, modify and distribute freely, including commercially. The full text is
in [LICENSE](LICENSE).

ComposeLM is an independent implementation of publicly published architectures; the implementations
and papers consulted are listed in [NOTICE](NOTICE). When redistributing, please include `LICENSE`
and `NOTICE` (that is the Apache-2.0 §4 requirement, and there is nothing else to observe).

Architecture preset names (`llama3`, `mistral`, …) are descriptive labels identifying those
structures. ComposeLM is not affiliated with any organization, and no pretrained weights are
distributed.

```bibtex
@software{ComposeLM2026,
  title  = {ComposeLM: One-line Configurable Transformer Library},
  author = {ComposeLM Contributors},
  year   = {2026},
  url    = {https://github.com/DW-dev-UE/ComposeLM}
}
```

Issues and PRs are always welcome. Including the output of `ModelConfig.to_dict()` along with your
PyTorch / CUDA versions in a bug report makes reproduction much faster.
