ComposeLM

ComposeLM is a PyTorch library for assembling, training, resuming, and locally running decoder-only language models. Architecture choices live in one ModelConfig, so experiments do not require a fork of the model code.

The current package version is 1.0.0. Stable and Preview boundaries, and the scope of the verification behind this release, are listed in the 1.0 release guide.

ComposeLM does not ship pretrained weights or a tokenizer. It also does not include an HTTP inference server. Bring a tokenizer from a library such as Transformers, and use the optional vLLM export/adapter when you need a separate high-throughput inference runtime.

Install

Python 3.10 or newer and PyTorch 2.8 or newer are required.

pip install composelm

To work on ComposeLM itself, install from a checkout:

pip install -e .

For development:

pip install -e ".[dev]"
python -m pytest

Optional extras are deliberately small:

pip install -e ".[datasets]"       # Hugging Face datasets examples
pip install -e ".[tracking]"       # TensorBoard and Weights & Biases
pip install -e ".[stateful-data]"  # exact streaming resume with workers

flash-attn, Transformers, vLLM, and DeepSpeed are platform-dependent and are not installed by ComposeLM. Install them separately only when you use the corresponding integration.

Build a model

Presets provide sensible defaults; explicit keyword arguments always win.

import torch
from composelm import build_model

model = build_model("llama3", d_model=256, n_layers=4, n_heads=8, n_kv_heads=2, vocab_size=32_000,
    max_seq_len=512, precision="fp32")

input_ids = torch.randint(0, 32_000, (2, 64))
logits = model(input_ids)
print(logits.shape) # [2, 64, 32000]

Available presets:

from composelm import list_archs

print(list_archs())

The presets are gpt, llama/llama1, llama2, llama3, mistral, gemma, gemma2, qwen, qwen2, deepseek, phi, phi3, and custom. They describe architecture defaults, not exact replicas or pretrained model releases.

You can also keep the complete configuration as data:

from composelm import ModelConfig, build_model, save_config_yaml

config = ModelConfig.from_arch("custom", d_model=512, n_layers=8, n_heads=8, n_kv_heads=2,
    vocab_size=32_000, attention_type="gqa", ffn_type="swiglu", pos_emb="rope",
    precision="bf16_mixed")
save_config_yaml(config, "model.yaml")
model = build_model(config)

Architecture options

These values are implemented by ModelConfig and validated before the model is built.

Examples:

# Partial RoPE, QK RMSNorm, and local/global attention
model = build_model("custom", d_model=512, n_layers=8, n_heads=8, n_kv_heads=2, vocab_size=32_000,
    attention_type="gqa", qk_norm="rmsnorm", rope_dim=32, sliding_window=256,
    global_attention_every_n_layers=4)

# Four routed experts and one shared expert
moe = build_model("custom", d_model=512, n_layers=8, n_heads=8, vocab_size=32_000, ffn_type="moe",
    num_experts=4, num_experts_per_tok=2, moe_num_shared_experts=1, moe_router_type="softmax",
    moe_load_balance="aux_loss")

MLA, distributed expert parallelism, block-sparse attention, Mamba/SSM, MTP, FlashAttention-3/4 selection, and true FP8 compute are not part of the stable 1.0 runtime. fp8_mixed currently warns and falls back to BF16.

See Architecture.md for component boundaries and field details.

Train

TrainingConfig is the public training configuration. Existing Trainer keyword arguments and train() remain as deprecated 1.x compatibility paths.

from composelm import Trainer, TrainingConfig

# A real dataset should yield input_ids and may also yield labels and
# attention_mask. Passing None creates synthetic data for a smoke run.
training = TrainingConfig(batch_size=4, gradient_accumulation_steps=8, learning_rate=3e-4,
    max_steps=200, logging_steps=10, save_steps=50, output_dir="runs/first", precision="bf16_mixed")

trainer = Trainer(model, train_dataset=None, config=training)
result = trainer.fit()
print(result.steps, result.history[-1])

A fresh fit() refuses to overwrite existing training artifacts. Use a new directory or resume explicitly:

result = trainer.fit(resume_from="latest")
# A checkpoint path is accepted as well:
# result = trainer.fit(resume_from="runs/first/checkpoint-100")

Exact resume restores the model, optimizer, scheduler, scaler, optimizer-step progress, per-rank RNG, and data position. It rejects a mismatch in model, training configuration, data fingerprint, or distributed topology. Checkpoints are committed only at optimizer-step boundaries.

For transfer learning, load weights into a new model and start a new run:

from composelm import load_model_weights, save_model_weights

save_model_weights("weights/model.safetensors", model)
load_model_weights("weights/model.safetensors", fresh_model)

This does not restore optimizer, RNG, or data state.

Logs

Rank zero writes:

• run.json: redacted configuration, environment, seed, Git revision, and model/data fingerprints;
• events.jsonl: start, step, checkpoint, resume, warning, error, and end events.

Step events include global loss, learning rate, gradient norm, step/data time, tokens per second, GPU memory, loss scale, and overflow state. JSONL write failure stops training; a failing optional callback is logged and disabled.

from composelm.train.callbacks import TensorBoardCallback

trainer = Trainer(model, dataset, config=training,
    callbacks=[TensorBoardCallback("runs/first/tensorboard")])

See docs/api.md for the event schema and streaming dataset contract.

Distributed training

Stable 1.0 strategies are single, ddp, and fsdp2 on one node. A process group must already exist for DDP or FSDP2.

torchrun --standalone --nproc_per_node=8 examples/train_distributed.py \
  --strategy fsdp2 --output-dir runs/fsdp2

The example initializes and destroys the process group itself. Loss and throughput are aggregated across ranks.

Multi-node training, DeepSpeed, and FSDP1 are Preview. expert_parallel=True with more than one process is rejected because token all-to-all is not implemented. See docs/distributed.md.

Inference

Local generation supports greedy decoding, top-k/top-p sampling, padded batches, and KV caching.

from composelm import generate

tokens = generate(model.eval(), input_ids, max_new_tokens=32, do_sample=True, temperature=0.8,
    top_p=0.9)

ComposeLM also includes a small continuous batcher and speculative decoding. They are local utilities, not a production server. For a separately installed vLLM runtime:

from composelm import (
    VLLMAdapter,
    VLLMSamplingConfig,
    export_vllm_checkpoint,
)

export_vllm_checkpoint(model, "export/model", tokenizer=tokenizer)
engine = VLLMAdapter.from_pretrained("export/model")
outputs = engine.generate(["Hello"], sampling=VLLMSamplingConfig(max_tokens=64, temperature=0.7))

The exporter accepts only exact GPT-2, Llama, Mistral, and Qwen2-compatible layouts. See docs/infer.md and docs/convert.md.

Checkpoint safety

Publish model weights as SafeTensors. Full training checkpoints contain trusted-only runtime state that may use pickle serialization. Manifest hashes detect incomplete or corrupted files; they do not make an untrusted checkpoint safe to load. See SECURITY.md.

Verification

python -m pytest
ruff check composelm tests examples
mypy composelm
python -m build

Hardware results are kept under bench_results/. Historical results describe the exact source and host used for that run; they should not be treated as measurements of an edited working tree. Multi-GPU qualification commands are documented in bench_multigpu/README.md.

Documentation

• API and logging
• Architecture
• Distributed training
• Inference
• HF and vLLM conversion
• 1.0 release and migration guide
• Changelog
• Security policy

License

Apache-2.0. See LICENSE and NOTICE.