xLLM

xLLM is a PyTorch-based framework for long-context language modeling, with distributed training, online data preparation, evaluation, and model export in one repository.

Start with PyTorch >= 2.11 and CUDA >= 12.8. Follow the PyTorch installation guide for your environment before building xLLM.

Building the native extensions requires a CUDA development environment. PyTorch must be installed before installing xLLM: requirements.txt does not include torch, and setup.py imports it to build the extensions (hence --no-build-isolation).

git clone https://github.com/ifm-ai/xllm.git
cd xllm

python -m pip install -r requirements.txt
python -m pip install -e . --no-build-isolation --config-settings editable_mode=compat

Install FlashAttention for the Quick Start below:

python -m pip install flash-attn --no-build-isolation

# FlashAttention 3
export ENABLE_FLASH_ATTENTION_3=TRUE

# Alternatively, use FlashAttention 4 instead:
# export ENABLE_FLASH_ATTENTION_4=TRUE

git clone https://github.com/ifm-ai/xattn.git
cd xattn

git submodule update --init --recursive
python -m pip install -r requirements-build.txt
python -m pip install --no-build-isolation .

python -m pip install flash-linear-attention

For Hugging Face tokenizers (required for the Quick Start below):

python -m pip install "transformers[torch]"

Start a short, single-GPU training run with the K2 Horizon 0.9B configuration. Replace the paths below with a directory containing *.chunk*.jsonl text data and a Hugging Face tokenizer with BOS/EOS tokens. See the data format examples.

DATA_DIR=/path/to/text-data
TOKENIZER=/path/to/tokenizer

torchrun --standalone --nproc_per_node=1 train.py \
  --model k2-horizon-0.9B \
  --model.causal_attn_backend flash --model.chunk_size 1024 \
  --data "${DATA_DIR}:1.0:text:text" \
  --tokenizer.type huggingface --tokenizer.path "$TOKENIZER" \
  --dataloader.packing_type bestfit --dataloader.buffer_size 512 \
  --batch_size 1 --seq_len 2048 --dtype bf16 \
  --steps 100 --optim.warmup 10 \
  --log_freq 10 --dump_freq 100 --eval_freq -1 \
  --keep_eval_checkpoints false \
  --dump_dir saved_models/quickstart

batch_size is per data-parallel rank. Alternatively, set global_batch_size, divisible by the data-parallel size; do not set both. Use a fresh dump_dir for a new run; an existing checkpoint in that directory is resumed automatically.

When changing these values, keep the constraints checked in xllm/config.py:

• seq_len must be divisible by model.chunk_size (times context_parallel_size when using context parallelism). The default chunk_size is 2048; the Quick Start uses 1024 so that seq_len 2048 splits into two chunks.
• dump_freq must be divisible by log_freq, and a positive eval_freq must be divisible by log_freq.
• With keep_eval_checkpoints true (the default), a positive eval_freq must also be divisible by dump_freq.

Conversion and serving checkpoints of xLLM is supported by xBridges.

xllm/
  configuration/  Config dataclass and command-line parsing
  csrc/           Custom C++/CUDA kernels
  data/           Online data loading, tokenization, packing, and resume
  models/         Model architectures and fused blocks
  modules/        Attention, expert layers, normalization, and operators
  distributed/    Parallelism and distributed training utilities
  optim/          Optimizers and learning-rate schedulers
  eval/           Evaluation tasks and runnersxllm_bridges/     Hugging Face conversion and vLLM integration
examples/         Training launch scripts and benchmarks
tests/            Tests and validation utilities

Bug reports and focused pull requests are welcome. For runtime issues, include the relevant configuration, package versions, hardware, and a minimal reproducer. Include tests for behavior changes.

xLLM is released under the Apache 2.0 License.