repository·main·Indexed 24 days ago
https://github.com/nx-ai/xlstmA novel Recurrent Neural Network architecture (Extended Long Short-Term Memory) that utilizes Exponential Gating and Matrix Memory to achieve performance competitive with Transformers and State Space Models. The library provides implementations of xLSTMLarge, xLSTMLMModel for language modeling, and xLSTMBlockStack for mixing mLSTM and sLSTM blocks. It supports optimized TFLA Triton kernels for NVIDIA GPUs and native PyTorch implementations for other hardware.
You can install the xlstm package via pip or by cloning the repository. For the xLSTM Large 7B model, you must also install the mlstm_kernels package to provide fast kernels.
pip install mlstm_kernels
pip install xlstmgit clone https://github.com/NX-AI/xlstm.git
cd xlstm
pip install -e .The xLSTM Large 7B architecture is optimized for training throughput and stability. Its implementation is located in xlstm/xlstm_large/model.py.
To use it on NVIDIA GPUs, configure the model with TFLA Triton kernels using chunkwise--triton_xl_chunk, native_sequence__triton, and triton kernels.
For non-NVIDIA hardware (like Apple Metal), use the native PyTorch implementations by setting the kernels to chunkwise--native_autograd, native_sequence__native, and native respectively.
import torch
from xlstm.xlstm_large.model import xLSTMLargeConfig, xLSTMLarge
# configure the model with TFLA Triton kernels
xlstm_config = 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", # xl_chunk == TFLA kernels
sequence_kernel="native_sequence__triton",
step_kernel="triton",
)
# instantiate the model
xlstm = xLSTMLarge(xlstm_config)
xlstm = xlstm.to("cuda")
# create inputs
input = torch.randint(0, 2048, (3, 256)).to("cuda")
# run a forward pass
out = xlstm(input)To load a pretrained xLSTMLarge model, use the load_from_pretrained function from the xlstm.xlstm_large.from_pretrained module. You must provide the local path to the checkpoint directory via the checkpoint_path argument. Once loaded, the model can be moved to a device (e.g., cuda) and its configuration can be modified.
from xlstm.xlstm_large.from_pretrained import load_from_pretrained
CHECKPOINT_PATH = "/path/to/xlstm_large_checkpoint"
model = load_from_pretrained(checkpoint_path=CHECKPOINT_PATH)
model = model.to("cuda")To use the CUDA version of sLSTM, you need a GPU with Compute Capability >= 8.0.
If you encounter compilation issues, set the TORCH_CUDA_ARCH_LIST environment variable:
export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"To ensure the correct CUDA libraries are included, you can use the XLSTM_EXTRA_INCLUDE_PATHS environment variable:
export XLSTM_EXTRA_INCLUDE_PATHS='/usr/local/include/cuda/:/usr/include/cuda/'Or within Python:
import os
os.environ['XLSTM_EXTRA_INCLUDE_PATHS']='/usr/local/include/cuda/:/usr/include/cuda/'You can run synthetic experiments (like Parity or Multi-Query Associative Recall) using the main.py script in the experiments folder. Use the --config flag to specify the experiment configuration.
Example commands:
# xLSTM[0:1], sLSTM only
PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm01.yaml
# xLSTM[1:0], mLSTM only
PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm10.yaml
# xLSTM[1:1], mLSTM and sLSTM
PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm11.yamlPYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm01.yamlThe xLSTMBlockStack can be used as an alternative backbone in existing projects, similar to a stack of Transformer blocks. It allows you to mix mLSTM and sLSTM blocks. You can specify which layers should be sLSTM using the slstm_at parameter.
import torch
from xlstm import (
xLSTMBlockStack,
xLSTMBlockStackConfig,
mLSTMBlockConfig,
mLSTMLayerConfig,
sLSTMBlockConfig,
sLSTMLayerConfig,
FeedForwardConfig,
)
cfg = xLSTMBlockStackConfig(
mlstm_block=mLSTMBlockConfig(
mlstm=mLSTMLayerConfig(
conv1d_kernel_size=4, qkv_proj_blocksize=4, num_heads=4
)
),
slstm_block=sLSTMBlockConfig(
slstm=sLSTMLayerConfig(
backend="cuda",
num_heads=4,
conv1d_kernel_size=4,
bias_init="powerlaw_blockdependent",
),
feedforward=FeedForwardConfig(proj_factor=1.3, act_fn="gelu"),
),
context_length=256,
num_blocks=7,
embedding_dim=128,
slstm_at=[1],
)
xlstm_stack = xLSTMBlockStack(cfg)
x = torch.randn(4, 256, 128).to("cuda")
xlstm_stack = xlstm_stack.to("cuda")
y = xlstm_stack(x)The xLSTMLMModelConfig class defines the configuration schema for the xLSTMLMModel. It inherits from xLSTMBlockStackConfig and adds parameters specific to the Language Model head and embeddings.
Configuration Keys:
vocab_size (int): The size of the vocabulary. Defaults to -1.tie_weights (bool): If True, the weights of the lm_head are tied to the token_embedding weights. Defaults to False.weight_decay_on_embedding (bool): Determines if weight decay should be applied to the token embeddings. Defaults to False.add_embedding_dropout (bool): If True, applies a dropout layer to the embeddings. Defaults to False.The xLSTMBlockStackConfig dataclass defines the architecture of an xLSTM block stack. You can specify whether to use mLSTMBlock or sLSTMBlock at specific positions using slstm_at or by providing specific block configurations.
Key configuration options:
mlstm_block: An optional mLSTMBlockConfig instance.slstm_block: An optional sLSTMBlockConfig instance.num_blocks: Total number of blocks in the stack.embedding_dim: The dimensionality of the embeddings.context_length: The maximum sequence length (set to -1 for variable length).slstm_at: Determines which block indices use sLSTM. Can be a list[int] of specific indices or the literal `For a type-safe approach, you can instantiate xLSTMBlockStack by passing nested configuration objects directly. This uses classes like mLSTMBlockConfig, sLSTMBlockConfig, mLSTMLayerConfig, sLSTMLayerConfig, and FeedForwardConfig to define the model architecture.
import torch
from xlstm import (
xLSTMBlockStack,
xLSTMBlockStackConfig,
mLSTMBlockConfig,
mLSTMLayerConfig,
sLSTMBlockConfig,
sLSTMLayerConfig,
FeedForwardConfig,
)
cfg = xLSTMBlockStackConfig(
mlstm_block=mLSTMBlockConfig(
mlstm=mLSTMLayerConfig(
conv1d_kernel_size=4, qkv_proj_blocksize=4, num_heads=4
)
),
slstm_block=sLSTMBlockConfig(
slstm=sLSTMLayerConfig(
backend="cuda" if torch.cuda.is_available() else "vanilla",
num_heads=4,
conv1d_kernel_size=4,
bias_init="powerlaw_blockdependent",
),
feedforward=FeedForwardConfig(proj_factor=1.3, act_fn="gelu"),
),
context_length=256,
num_blocks=7,
embedding_dim=128,
slstm_at=[1],
)
xlstm_stack = xLSTMBlockStack(cfg)You can run a forward pass on xLSTMLarge by passing a tensor of token IDs. If return_last_states is set to True in the config, the model returns a tuple containing the output tensor and the final state dictionary.
Input Format:
A tensor of shape (batch_size, sequence_length) containing integer token IDs.
Output Format:
If return_last_states=True, the output is (out, state):
out: Tensor of shape (batch_size, sequence_length, vocab_size).state: A dictionary representing the model's internal states.import torch
# Assuming xlstm is already instantiated and moved to device
# input shape: (batch_size, sequence_length)
input = torch.randint(0, 2048, (3, 256)).to("cuda")
out = xlstm(input)
if len(out) == 2:
out, state = out
# out.shape will be (3, 256, 2048)The xLSTMBlockStack processes input tensors of shape (batch_size, sequence_length, embedding_dim). The output shape matches the input shape.
# Assuming xlstm_stack is already instantiated and moved to device
x = torch.randn(4, 256, 128).to(device=device)
xlstm_stack = xlstm_stack.to(device=device)
y = xlstm_stack(x)
print(y.shape) # Expected: torch.Size([4, 256, 128])You can configure and instantiate an xLSTMBlockStack using a YAML-formatted string combined with OmegaConf and dacite. This approach is useful for managing complex nested configurations for mlstm_block, slstm_block, and global parameters like context_length and embedding_dim.
Key configuration parameters include:
mlstm_block: Configuration for the mLSTM component.slstm_block: Configuration for the sLSTM component, including the backend (e.g., 'cuda' or 'vanilla') and feedforward settings.slstm_at: A list of indices specifying which blocks in the stack should use the sLSTM architecture. An empty list [] results in a stack of only mLSTM blocks.from omegaconf import OmegaConf
from dacite import from_dict
from dacite import Config as DaciteConfig
from xlstm.xlstm_block_stack import xLSTMBlockStack, xLSTMBlockStackConfig
xlstm_cfg = f"""
mlstm_block:
mlstm:
conv1d_kernel_size: 4
qkv_proj_blocksize: 4
num_heads: 4
slstm_block:
slstm:
backend: {'cuda' if torch.cuda.is_available() else 'vanilla'}
num_heads: 4
conv1d_kernel_size: 4
bias_init: powerlaw_blockdependent
feedforward:
proj_factor: 1.3
act_fn: gelu
context_length: 256
num_blocks: 7
embedding_dim: 128
slstm_at: [1]
"""
cfg = OmegaConf.create(xlstm_cfg)
cfg = from_dict(data_class=xLSTMBlockStackConfig, data=OmegaConf.to_container(cfg), config=DaciteConfig(strict=True))
xlstm_stack = xLSTMBlockStack(cfg)