A Simple & Efficient Training Framework for Long-Context LLMs
Optimized for scalability, memory efficiency, and seamless integration β built to unlock the full potential of long-context large language models.
- π
2025-10-07 β π Initial Release: Loom-Train is now live!
β Native support for π€ Hugging Face Trainer
β Optimized attention with π Ring-Flash-Attention
β Lightweight, plug-and-play design for long-sequence training (128K+ tokens)
- π§ Plug-and-Play: Drop-in replacement for HF Trainer β no major code changes needed.
- π Memory-Efficient: Leverages Ring-Flash-Attention to reduce GPU memory footprint by up to 50%.
- π Scalable: Seamlessly scales to 100K+ context lengths without sacrificing speed.
- β‘ Fast Setup: Minimal dependencies, easy installation via
pip install loom-train.
To install theloomtrain package from the gitee repository, run:
git clone https://github.com/LCM-Lab/LOOM-Train.git
conda create -n loom_train python=3.10 -y
conda activate loom_train
cd LOOM-Train/loomtrain
pip install -e .To install flash attention, run the command below to obtain the required flah-attn version:
loomtrain-required-flash-attnDownload the suitable version of flash_attn from https://github.com/Dao-AILab/flash-attention/releases
pip install <path_to_flash_attn_whl_file>
pip install ring_flash_attnThen just swap your Trainer with LoomTrainer:
from loomtrain import LoomTrainer
trainer = LoomTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
# ... rest unchanged!
)We welcome contributions! Whether itβs bug fixes, new features, or documentation improvements β feel free to open an issue or PR.
Letβs build the future of long-context training, together. πͺ
Questions? Suggestions? Reach out at: [email protected] and [email protected]