简介
sat
(SwissArmyTransformer
)是一个灵活强大的库,用于开发自己的Transformer变体。
sat
名称来源于"瑞士军刀",意味着所有模型(如BERT、GPT、T5、GLM、CogView、ViT等)共享相同的核心代码,并通过一些额外的轻量级混入来满足多样化的使用需求。
sat
由deepspeed-ZeRO
和模型并行支持,旨在为大型模型(100M~20B参数)的预训练和微调提供最佳实践。
安装
pip install SwissArmyTransformer
特性
-
只需一行代码即可添加与模型无关的组件,例如前缀调优!
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和其他自回归模型在训练和推理过程中表现不同。在推理过程中,文本是逐个生成的,我们需要缓存先前的状态以提高效率。使用我们的库,你只需考虑训练(教师强制)期间的行为,并通过添加混入将其转换为缓存的自回归模型:
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
,同时通过超参数来修改有用的训练配置。
快速上手
在sat中使用Bert
进行推理的最典型Python文件如下:
# @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(自动设置)获取BertTokenizer。
tokenizer = get_tokenizer(args)
# 在这里按你的需求使用bert!
# ...
然后我们可以通过以下方式运行代码:
SAT_HOME=/path/to/download python inference_bert.py --mode inference
所有官方支持的模型名称都在urls.py中。
微调或预训练Transformer也非常简单!
# @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) # 来自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阶段等超参数
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。更多详情请参阅教程。
教程
引用
目前我们还没有论文,所以你不需要正式引用我们!
如果这个项目对你的研究或工程有帮助,请使用\footnote{https://github.com/THUDM/SwissArmyTransformer}
来提及我们,并向他人推荐SwissArmyTransformer
。
贡献sat的教程正在编写中!
该项目基于(作为用户的)DeepSpeed、Megatron-LM和Huggingface transformers。感谢他们的出色工作。