Github
论文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项目在推理和生成任务中可能的应用场景。
