NX-AI/xlstm

xLSTM – 扩展的长短期记忆

是什么 – xLSTM 是一种基于经典 LSTM 的新型循环神经网络架构。它引入了 指数门控、矩阵记忆 以及多种归一化/稳定化技巧，使其能够扩展到非常大的语言模型（例如，一个 70 亿参数的模型），同时保持推理速度快且内存高效。作者声称它在语言建模基准测试中可与 Transformer 和状态空间模型相媲美。

核心组件

xLSTMBlockStack – 可直接替换 Transformer 块堆栈的组件；内部混合了三种块类型（mLSTM、sLSTM 和前馈层），可通过 dataclasses 配置。
xLSTMLMModel – 在块堆栈之上添加词元嵌入和语言模型头的语言模型包装器。
mlstm_kernels – 提供自定义 CUDA/Triton 内核（"sLSTM" 内核）的独立包，可大幅加速循环运算。
xLSTMLarge – 后续论文中使用的 70 亿参数模型的单文件实现，支持针对不同硬件的可配置内核。

安装

# 可选：创建作者使用的精确 conda 环境
conda env create -f environment_pt240cu124.yaml
conda activate xlstm

# 安装快速内核（7B 模型必需）
pip install mlstm_kernels

# 安装库本身
pip install xlstm   # 或：git clone https://github.com/NX-AI/xlstm && pip install -e .

该包需要 PyTorch ≥ 1.8。要使用 GPU 加速，需配备较新的 NVIDIA GPU（CUDA 计算能力 ≥ 8.0）以使用 Triton 内核；否则，纯 PyTorch 回退可在任何平台运行。

快速开始（7B 推理）

import torch
from xlstm.xlstm_large.model import xLSTMLargeConfig, xLSTMLarge

cfg = xLSTMLargeConfig(
    embedding_dim=512,
    num_heads=4,
    num_blocks=6,
    vocab_size=2048,
    return_last_states=True,
    mode="inference",
    chunkwise_kernel="chunkwise--triton_xl_chunk",
    sequence_kernel="native_sequence__triton",
    step_kernel="triton",
)
model = xLSTMLarge(cfg).to("cuda")
inputs = torch.randint(0, 2048, (3, 256), device="cuda")
out = model(inputs)
print(out.shape)   # (3, 256, 2048)

一个笔记本文件（notebooks/xlstm_large/demo.ipynb）演示了相同的工作流程。

硬件建议

NVIDIA GPU – 使用 Triton 内核可获得最佳性能；仓库提到在 RTX 3080/3090 及更新型号（CC 8.0+）上成功运行。
AMD GPU – Triton 内核可能仍可运行，但作者建议回退到原生 PyTorch 内核。
Apple Silicon – 使用社区维护的 xLSTM-metal 版本（MLX）以获得 Metal 原生实现。

模型

xLSTM-Large 7B – 在 2.3T 个 token 上训练的 70 亿参数循环语言模型；权重托管在 Hugging Face 上（https://huggingface.co/NX-AI/xLSTM-7b）。
更小的研究模型可通过 xLSTMBlockStack 和 xLSTMLMModel 类从 YAML 配置文件实例化（README 中有示例）。

实验仓库包含合成任务（Parity、Multi-Query Associative Recall），展示了两种子块（sLSTM 用于状态追踪，mLSTM 用于记忆）的互补优势。运行方式如下：

PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm11.yaml

（注：训练循环非常简单——没有早停或测试评估。）

引用如果您使用代码或 7B 模型，请引用 README 中列出的两篇论文（NeurIPS 2024 xLSTM 论文和 ICML 2025 xLSTM-7B 论文）。

所有信息均直接来自仓库的 README；未添加任何外部假设。

NX-AI/xlstm

xLSTM – 扩展的长短期记忆

相关