引言

sat（SwissArmyTransformer）是一个灵活且功能强大的库，用于开发您自己的 Transformer 变体。

sat 的名称源自“瑞士军刀”，意指所有模型（例如 BERT、GPT、T5、GLM、CogView、ViT 等）共享相同的骨干代码，并通过一些轻量级的混入模块来满足多种用途。

sat 由 deepspeed-ZeRO 和模型并行技术驱动，旨在为大规模模型（1 亿至 200 亿参数）的预训练和微调提供最佳实践。

安装

    pip install SwissArmyTransformer

特性

• 添加与模型无关的组件，例如前缀调优，只需 一行 代码！

  • 前缀调优（或 P-tuning）通过在每个注意力层中添加可训练参数来改进微调。使用我们的库，将它应用于 GLM 分类模型（或任何其他模型）非常容易。

      class ClassificationModel(GLMModel): # 也可以是 BertModel、RobertaModel 等。
          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))
  

  • GPT 和其他自回归模型在训练和推理时的行为不同。在推理过程中，文本是逐 token 生成的，我们需要缓存之前的状态以提高效率。借助我们的库，您只需考虑训练阶段的行为（教师强制），并通过添加一个混入模块将其转换为缓存式自回归模型：

      model, args = AutoModel.from_pretrained('glm-10b-chinese', args)
      model.add_mixin('auto-regressive', CachedAutoregressiveMixin())
      # 使用束搜索生成序列
      from sat.generation.autoregressive_sampling import filling_sequence
      from sat.generation.sampling_strategies import BeamSearchStrategy
      output, *mems = filling_sequence(model, input_seq,
                      batch_size=args.batch_size,
                      strategy=BeamSearchStrategy(args.batch_size))
  

• 用最少的代码构建基于 Transformer 的模型。我们提到了 GLM，它与标准 Transformer（称为 BaseModel）的区别仅在于位置嵌入（以及训练损失）。我们在编码时只需关注相关部分。

  扩展完整定义：

  class BlockPositionEmbeddingMixin(BaseMixin):
      # 在这里定义混入模块的参数
      def __init__(self, max_sequence_length, hidden_size, init_method_std=0.02):
          super(BlockPositionEmbeddingMixin, self).__init__()
          self.max_sequence_length = max_sequence_length
          self.hidden_size = hidden_size
          self.block_position_embeddings = torch.nn.Embedding(max_sequence_length, hidden_size)
          torch.nn.init.normal_(self.block_position_embeddings.weight, mean=0.0, std=init_method_std)
      
      # 在这里定义混入模块的方法
      def position_embedding_forward(self, position_ids, **kwargs):
          position_ids, block_position_ids = position_ids[:, 0], position_ids[:, 1]
          position_embeddings = self.transformer.position_embeddings(position_ids)
          block_position_embeddings = self.block_position_embeddings(block_position_ids)
          return position_embeddings + block_position_embeddings
  
  class GLMModel(BaseModel):
      def __init__(self, args, transformer=None):
          super().__init__(args, transformer=transformer)
          self.add_mixin('block_position_embedding', 
              BlockPositionEmbeddingMixin(args.max_sequence_length, args.hidden_size)
          ) # 为 GLM 添加混入模块
  

• 全面的训练支持。sat 致力于提供预训练和微调的最佳实践，您只需完成 forward_step 和 create_dataset_function，并通过超参数调整有用的训练配置。

  • 通过指定 --num_nodes、--num_gpus 和简单的 hostfile，可以将训练扩展到多台 GPU 或节点。
  • DeepSpeed 和模型并行技术。
  • 更好地集成 ZeRO-2 和激活检查点。
  • 自动扩展和打乱训练数据，并支持 memmap。
  • 成功支持 CogView2 和 CogVideo 的训练。
  • 目前唯一支持在 GPU 上微调 T5-10B 的开源代码库。

快速游览

在 sat 中使用 Bert 进行推理的典型 Python 文件如下：

# @文件: inference_bert.py
from sat import get_args, get_tokenizer, AutoModel
# 解析参数，初始化环境。这是必要的。
args = get_args() 
# 自动下载并加载模型。同时会将模型相关的超参数存储到 args 中。
model, args = AutoModel.from_pretrained('bert-base-uncased', args) 
# 根据 args.tokenizer_type（自动设置）获取 BertTokenizer。
tokenizer = get_tokenizer(args) 
# 现在您可以按需使用 Bert 了！
# ...

然后我们可以通过以下命令运行代码：

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

所有官方支持的模型名称都列在 urls.py 中。

对 Transformer 进行微调或预训练也非常简单！

# @文件: 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) # 来自 create_dataset_function 加载的数据集。
    loss, *others = model(inputs)
    return loss
    
# 解析参数，初始化环境。这是必要的。
args = get_args() 
model, args = AutoModel.from_pretrained('bert-base-uncased', args) 
tokenizer = get_tokenizer(args) 
# 现在您可以按需使用 Bert 了！
model.del_mixin('bert-final')
model.add_mixin('classification_head', MLPHeadMixin(args.hidden_size, 2048, 1))
# 仅需一行代码即可开始训练！

# args 已经包含了学习率、训练迭代次数、zero-stage 等超参数...
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

这里我们在 GPU 0 和 1 上使用数据并行。我们也可以通过 --hostfile /path/to/hostfile 在多台互联的机器上启动训练。更多详情请参阅教程。

要编写自己的模型，你只需要考虑它与标准 Transformer 的不同之处。例如，如果你有一个改进注意力机制的想法：

from sat.model import BaseMixin
class MyAttention(BaseMixin):
    def __init__(self, hidden_size):
        super(MyAttention, self).__init__()
        # MyAttention 可能需要一些新的参数，比如一个可学习的 alpha。
        self.learnable_alpha = torch.nn.Parameter(torch.ones(hidden_size))
    
    # 这是一个钩子函数，函数名 `attention_fn` 是特殊的。
    def attention_fn(q, k, v, mask, dropout=None, **kwargs):
        # 我的注意力机制实现代码。
        # ...
        return attention_results

这里的 attention_fn 是一个钩子函数，它会替换默认的注意力计算过程。所有可用的钩子函数都在 transformer_defaults.py 中。现在我们可以使用 add_mixin 将我们的修改应用到所有的 Transformer 模型中，比如 BERT、Vit 和 CogView。更多细节请参阅教程。

教程

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

引用

目前我们还没有发表论文，所以不需要正式引用我们！~

如果这个项目对你的研究或工程工作有所帮助，请使用 \footnote{https://github.com/THUDM/SwissArmyTransformer} 来提及我们，并向他人推荐 SwissArmyTransformer。

关于如何为 sat 贡献代码的教程正在准备中！

该项目基于 DeepSpeed、Megatron-LM 和 Huggingface Transformers（由其用户开发）。感谢这些优秀团队的工作。

SwissArmyTransformer 快速上手指南

SwissArmyTransformer (简称 sat) 是一个灵活且强大的库，旨在帮助开发者构建自己的 Transformer 变体。其核心理念是“瑞士军刀”，即所有模型（如 BERT, GPT, T5, GLM, CogView, ViT 等）共享同一套骨干代码，并通过轻量级的 Mixin 机制实现多样化功能。该工具基于 deepspeed-ZeRO 和模型并行技术，专为大规模模型（1 亿至 200 亿参数）的预训练和微调提供最佳实践。

环境准备

• 系统要求：Linux 操作系统（推荐 Ubuntu）。
• Python 版本：建议 Python 3.8 或更高版本。
• 前置依赖：
  • PyTorch (需与 CUDA 版本匹配)
  • DeepSpeed (用于分布式训练和 ZeRO 优化)
  • NVIDIA GPU (用于加速训练和推理，支持多卡及多机分布式)

提示：确保已正确安装 CUDA 驱动和对应的 PyTorch 版本，以发挥 sat 在大规模模型训练上的性能优势。

安装步骤

通过 pip 直接安装官方发布版本：

pip install SwissArmyTransformer

国内加速建议：如果下载速度较慢，推荐使用国内镜像源进行安装：

pip install SwissArmyTransformer -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 模型推理 (Inference)

以下示例展示了如何加载预训练的 BERT 模型并进行推理。sat 会自动处理模型下载和环境初始化。

# @File: inference_bert.py
from sat import get_args, get_tokenizer, AutoModel

# 解析参数并初始化环境（必需步骤）
args = get_args() 

# 自动下载并加载模型，同时将模型相关的超参数写入 args
model, args = AutoModel.from_pretrained('bert-base-uncased', args) 

# 根据 args.tokenizer_type 获取对应的分词器
tokenizer = get_tokenizer(args) 

# 在此处使用 bert 进行你的任务！
# ...

运行上述代码：

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

注：SAT_HOME 环境变量用于指定模型文件的下载保存路径。所有官方支持的模型名称可在 sat/resources/urls.py 中查看。

2. 模型微调 (Finetuning)

sat 让微调变得极其简单。你只需定义数据加载函数和前向传播步骤，其余训练配置（如学习率、Zero 阶段等）均可通过命令行参数控制。

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

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) # 来自 create_dataset_function 定义的数据集
    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))

# 一行代码启动训练！
# args 已包含 lr, train-iters, zero-stage 等超参数
training_main(args, 
    model_cls=model, 
    forward_step_function=forward_step, # 用户自定义
    create_dataset_function=create_dataset_function # 用户自定义
)

运行微调脚本（示例使用 2 张 GPU 进行数据并行训练）：

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

• 分布式扩展：如需在多机集群上训练，可通过 --hostfile /path/to/hostfile 指定主机文件。
• 混合精度训练：示例中已开启 --fp16 以节省显存并加速训练。

3. 自定义模型组件 (Mixin)

sat 的核心特性是通过 Mixin 机制轻松扩展模型。例如，若要改进注意力机制，只需定义一个继承自 BaseMixin 的类并实现对应的钩子函数：

from sat.model import BaseMixin
import torch

class MyAttention(BaseMixin):
    def __init__(self, hidden_size):
        super(MyAttention, self).__init__()
        # 定义新参数，例如可学习的 alpha
        self.learnable_alpha = torch.nn.Parameter(torch.ones(hidden_size))
    
    # 钩子函数，名称 `attention_fn` 是特殊的，将替换默认行为
    def attention_fn(q, k, v, mask, dropout=None, **kwargs):
        # 在此编写自定义的注意力逻辑
        # ...
        return attention_results

定义完成后，使用 model.add_mixin('my_attention', MyAttention(...)) 即可将该改进应用到 BERT、ViT 或 CogView 等任何基于 sat 的模型中。