ComposeLM

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

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.

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). Multi-node is untested and no pretrained weights ship. Read Caveats and known limitations before wiring it into production or publishing benchmark numbers.

Table of contents

• What this library is for
• Installation
• Quick start
• Configuration: ModelConfig
• Architecture
• Component catalog
• Training
• Inference
• HuggingFace weight conversion
• Distributed training
• Benchmarks
• Caveats and known limitations ⚠️
• Development
• Roadmap
• License

What this library is for

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

The design document is Architecture.md; API summaries live in docs/.

Installation

pip install composelm

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

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

[!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 for details.

Verify the install:

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

Quick start

Build a model

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.

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

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 below. Here is the minimal form.

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

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/.

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

The complete field list and validation rules are in Architecture.md §4.

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.

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

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

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

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

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.

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

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

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

rope_scaling accepts either a string or a dict.

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

You can also specify it granularly as a dict.

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:

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 below.

Training

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.

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

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

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

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.

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

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/): 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.

torchrun --standalone --nproc_per_node=2 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. shard_expert_indices() computes round-robin shard indices, but the actual token all-to-all dispatch is not implemented yet. See Caveats 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/, and the detailed write-up is H100_REPORT.md.

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.

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.

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/.

Reproducing

# 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/; the write-up is 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.

[!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.

# 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

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.

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.

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.

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

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

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 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.

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
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 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.

ComposeLM is an independent implementation of publicly published architectures; the implementations and papers consulted are listed in 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.

@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.