支持 Hugging Face 转换器模型 - Amazon SageMaker
现成支持的模型
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

支持 Hugging Face 转换器模型

SageMaker 模型并行性库的张量并行性为以下 Hugging Face 转换器模型提供了现成可用的支持：

  • GPT-2、BERT 和 RoBERTa（在 SageMaker 模型并行性库 v1.7.0 及更高版本中可用）

  • GPT-J（在 SageMaker 模型并行性库 v1.8.0 及更高版本中可用）

  • GPT-Neo（在 SageMaker 模型并行性库 v1.10.0 及更高版本中可用）

注意

对于任何其他转换器模型，您需要使用 smdistributed.modelparallel.torch.tp_register_with_module() API 来应用张量并行性。

注意

要使用张量并行性来训练 Hugging Face 转换器模型，对于所使用的适用于 PyTorch 的 Hugging Face 深度学习容器，请确保采用 SageMaker 模型并行库 v1.7.0 及更高版本。有关更多信息，请参阅 SageMaker 模型并行性库发行说明。

现成支持的模型

对于库现成支持的 Hugging Face 转换器模型，您无需手动实施钩子以将转换器 API 转换为 smdistributed 转换器层。您可以使用上下文管理器 smdistributed.modelparallel.torch.tensor_parallelism() 激活张量并行性，使用 smdistributed.modelparallel.torch.DistributedModel() 包装模型。您无需使用 smp.tp_register API 手动注册钩子用于张量并行性。

Hugging Face Transformers 与 smdistributed.modelparallel 之间的 state_dict 转换函数可以如下所示访问。

  • smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_state_dict_to_hf_gpt2(state_dict, max_seq_len=None)

  • smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_hf_state_dict_to_smdistributed_gpt2(state_dict)

  • smdistributed.modelparallel.torch.nn.huggingface.bert.translate_state_dict_to_hf_bert(state_dict, max_seq_len=None)

  • smdistributed.modelparallel.torch.nn.huggingface.bert.translate_hf_state_dict_to_smdistributed_bert(state_dict)

  • smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_state_dict_to_hf_roberta(state_dict, max_seq_len=None)

  • smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_hf_state_dict_to_smdistributed_roberta(state_dict)

  • smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_state_dict_to_hf_gptj(state_dict, max_seq_len=None)（在 SageMaker 模型并行性库 v1.8.0 及更高版本中可用）

  • smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_hf_gptj_state_dict_to_smdistributed_gptj（在 SageMaker 模型并行性库 v1.8.0 及更高版本中可用）

  • smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_state_dict_to_hf_gptneo(state_dict, max_seq_len=None)（在 SageMaker 模型并行性库 v1.10.0 及更高版本中可用）

  • smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_hf_state_dict_to_smdistributed_gptneo(state_dict)（在 SageMaker 模型并行性库 v1.10.0 及更高版本中可用）

GPT-2 转换函数的使用示例

首先包装模型，如以下代码所示。

from transformers import AutoModelForCausalLM

with smp.tensor_parallelism():
    model = AutoModelForCausalLM.from_config(hf_gpt2_config)

model = smp.DistributedModel(model)

对于 DistributedModel 对象中的 state_dict，您可以使用 translate_state_dict_to_hf_gpt2 函数将权重加载到原始 Hugging Face GPT-2 模型中，如以下代码所示。

from smdistributed.modelparallel.torch.nn.huggingface.gpt2 \
                                      import translate_state_dict_to_hf_gpt2
max_seq_len = 1024

# [... code block for training ...]

if smp.rdp_rank() == 0:
    state_dict = dist_model.state_dict()
    hf_state_dict = translate_state_dict_to_hf_gpt2(state_dict, max_seq_len)

    # can now call model.load_state_dict(hf_state_dict) to the original HF model

RoBERTa 转换函数的使用示例

与此类似，对于支持的 HuggingFace 模型 state_dict，您可以使用 translate_hf_state_dict_to_smdistributed 函数将其转换为可由 smp.DistributedModel 读取的格式。这在迁移学习使用场景中非常有用，此时预训练的模型加载到 smp.DistributedModel 中用于模型并行微调：

from smdistributed.modelparallel.torch.nn.huggingface.roberta \
                                      import translate_state_dict_to_smdistributed

model = AutoModelForMaskedLM.from_config(roberta_config)
model = smp.DistributedModel(model)

pretrained_model = AutoModelForMaskedLM.from_pretrained("roberta-large")
translated_state_dict =
        translate_state_dict_to_smdistributed(pretrained_model.state_dict())

# load the translated pretrained weights into the smp.DistributedModel
model.load_state_dict(translated_state_dict)

# start fine-tuning...