InfiniTransformer

InfiniTransformer 项目介绍

InfiniTransformer是一个非官方的基于PyTorch和🤗Transformers的项目，旨在实现名为“Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention”的技术，并支持Llama3和Gemma模型（同时也支持Llama 2和1）。该项目通过Infini-attention技术提供了更高效的上下文处理能力。

Infini-Attention 实现类型

InfiniTransformer提供了两种Infini-Attention的实现方式，分别是“Model-wise, Trainer-wise”类型和仅限“Attention-Layer”的类型。

类型 I：Model-wise, Trainer-wise

这种实现方式需要重写模型和配置的Python文件。由于进行了全面的修改，它无法与基础的HF训练器兼容，需要使用自定义的训练代码。尽管如此，它具有很高的内存效率：

可以使用AdamW优化器在不进行梯度检查点的情况下，在2块H100 80G显卡上训练Gemma-2B模型，序列长度为32768。
可以使用Adafactor优化器在不进行梯度检查点的情况下，在2块H100 80G显卡上训练Llama-3-8B模型，序列长度高达1M。
甚至可以在单块H100 80G显卡上进行“无限”上下文的训练。

类型 II：仅限 Attention-Layer

这种实现方式仅重写了Attention层的模型文件，属于最小修改，完全兼容HF训练器等。其内存使用量与默认的舍入注意力（SDPA）接近：

可以在2块H100 80G显卡上使用Adafactor优化器和梯度检查点训练Gemma-2B模型，序列长度为8192。

如何使用类型 I 的 Infini Attention

第一步：克隆仓库

使用以下命令克隆InfiniTransformer的代码：

git clone https://github.com/Beomi/InfiniTransformer

第二步：安装依赖项

需要通过源码安装最新版的🤗Transformers：

pip install -r requirements.txt
pip install -e git+https://github.com/huggingface/transformers.git@b109257f4f#egg=transformers

第三步：运行示例（推理及简单的前向/后向测试）

python test_basic.infini.py

第四步：使用自定义数据进行训练

以下是训练示例代码：

训练Llama-3模型，序列长度为1M，使用MiniPile数据集：

./train.llama.infini.noclm.1Mseq.sh

训练Gemma-2B模型，序列长度为32K，使用WikiText2数据集：

./train.gemma.infini.noclm.sh

如何使用类型 II 的 Infini Attention

第一步：克隆仓库

使用以下命令克隆InfiniTransformer的代码：

git clone https://github.com/Beomi/InfiniTransformer

第二步：安装依赖项

需要通过源码安装最新版的🤗Transformers：

pip install -r requirements.txt
pip install -e git+https://github.com/huggingface/transformers.git@b109257f4f#egg=transformers

第三步：运行示例（推理及简单的前向/后向测试）

python test_basic.py

第四步：使用自定义数据进行训练

以下是训练示例代码：

python test_train.small.gemma.py

项目提供的示例代码使用了来自wikitext数据集的数据，以便于快速上手。

推理示例

项目在WikiText2数据集上进行了一轮训练后，生成了一些样本，展示了模型的语言生成能力。以下是两个输入输出示例：

输入：

This work introduces an efficient method to scale Transformer-based

输出示例1：

models for denoising, denoising denoising, and deep denoising of images…（输出继续描述模型在图像去噪方面的应用）

输出示例2：

vision models across platforms using a custom architecture optimized…（输出继续描述模型在跨平台上的视觉模型应用）

这些示例显示了InfiniTransformer项目在推理和生成任务中可能的应用场景。