SwissArmyTransformer

项目介绍

SwissArmyTransformer（简称 sat）是一个灵活且强大的库，旨在帮助用户开发自己的 Transformer 变体。这一库的命名取意于瑞士军刀，意味着各类模型（如 BERT、GPT、T5、GLM、CogView、ViT 等）都可以通过相同的基础代码实现自身的多功能性，只需添加一些轻量级的组件。

sat 由 deepspeed-ZeRO 和模型并行技术提供支持，目标是为大规模模型（参数范围在 100M 到 20B 之间）的预训练和微调提供最佳实践。

安装方式

pip install SwissArmyTransformer

主要特点

模型无关的组件添加：例如，只需一行代码便可实现前缀调优 (Prefix-tuning)。这种方法通过在每个注意力层中加入可训练参数来提高微调效果。无论是用于 GLM 分类模型还是其他模型，这一过程都非常简单。

class ClassificationModel(GLMModel): 
    def __init__(self, args, transformer=None, **kwargs):
        super().__init__(args, transformer=transformer, **kwargs)
        self.add_mixin('classification_head', MLPHeadMixin(args.hidden_size, 2048, 1))
        self.add_mixin('prefix-tuning', PrefixTuningMixin(args.num_layers, args.hidden_size // args.num_attention_heads, args.num_attention_heads, args.prefix_len))

最少的代码构建 Transformer 模型：开发者可以很轻松地指定不同的 Transformer 模型的特性。例如，GLM 与标准 Transformer 的区别只在于位置嵌入和训练损失。开发者只需专注于相关部分的编码。
全面的训练支持：sat 旨在为预训练和微调提供最佳实践，用户只需完成 forward_step 和 create_dataset_function，更改超参数即可调整训练配置。
- 通过指定 --num_nodes 和 --num_gpus，以及一个简单的 hostfile，即可扩展至多 GPU 或节点的训练。
- 支持 DeepSpeed 和模型并行。
- 更好地集成 ZeRO-2 和激活检查点。
- 自动扩展和打乱训练数据及内存映射。
- 成功支持 CogView2 和 CogVideo 的训练。
- 当前唯一一个在 GPU 上支持微调 T5-10B 的开源代码库。

快速入门

以下是一个在 sat 中使用 Bert 进行推理的典型 Python 文件：

# @File: inference_bert.py
from sat import get_args, get_tokenizer, AutoModel
args = get_args() 
model, args = AutoModel.from_pretrained('bert-base-uncased', args) 
tokenizer = get_tokenizer(args) 
# 使用 Bert 进行其他操作
# ...

要运行上述代码，只需执行：

SAT_HOME=/path/to/download python inference_bert.py --mode inference

微调或预训练同样非常简单：

# @File: finetune_bert.py
from sat import get_args, get_tokenizer, AutoModel
from sat.model.mixins import MLPHeadMixin

def create_dataset_function(path, args):
    # 加载数据集
    # ...
    assert isinstance(dataset, torch.utils.data.Dataset)
    return dataset

def forward_step(data_iterator, model, args, timers):
    inputs = next(data_iterator)
    loss, *others = model(inputs)
    return loss

args = get_args() 
model, args = AutoModel.from_pretrained('bert-base-uncased', args) 
tokenizer = get_tokenizer(args) 
model.del_mixin('bert-final')
model.add_mixin('classification_head', MLPHeadMixin(args.hidden_size, 2048, 1))
training_main(args, model_cls=model, forward_step_function=forward_step, create_dataset_function=create_dataset_function)

运行上述代码：

deepspeed --include localhost:0,1 finetune_bert.py \
    --experiment-name ftbert \
    --mode finetune --train-iters 1000 --save /path/to/save \
    --train-data /path/to/train --valid-data /path/to/valid \
    --lr 0.00002 --batch-size 8 --zero-stage 1 --fp16

教程

如何使用 sat 中收集的预训练模型？
为什么以及如何在 sat 中训练模型？

引用

目前我们还没有正式的论文，如果这个项目对你的研究或工程有所帮助，可以通过 \footnote{https://github.com/THUDM/SwissArmyTransformer} 进行说明，并向他人推荐 SwissArmyTransformer。项目基于 DeepSpeed、Megatron-LM 和 Huggingface transformers，感谢他们的精彩工作。