微调

微调是一个持续训练预训练模型以提高特定使用场景性能的过程。

微调完全适合单个 GPU 的小型模型，或者那些完全适合 8 个模型副本的模型非常 CPUs 简单。它不需要对常规的 FSDP 训练进行特别修改。如果模型大于此范围，就需要考虑使用延迟参数初始化功能，这可能比较麻烦。

为了解决这个问题，SMP 库在其中一个等级上加载完整模型，而其他等级则在元设备上创建空权重模型。然后， PyTorch FSDP 使用init_weights函数初始化非零等级上的权重，并将所有等级的权重与设置为 0 的权重同步。sync_module_states True下面的代码片段显示了如何在训练脚本中进行设置。

import torch.distributed as dist
from transformers import AutoModelForCasalLM
from accelerate import init_empty_weights
from torch.sagemaker.delayed_param import DelayedParamIniter

if dist.get_rank() == 0:
    model = AutoModelForCasalLM.from_pretrained(..., low_cpu_mem_usage=True)
else:
    with init_empty_weights():
        model = AutoModelForCasalLM.from_config(AutoConfig.from_pretrained(...))
    delayed_initer = DelayedParamIniter(model)

model = FSDP(
    model,
    ...,
    sync_module_states=True,
    param_init_fn=delayed_initer.get_param_init_fn() if dist.get_rank() > 0 else None
)

利用 SMP 张量并行微调预训练的 Hugging Face 转换器模型

本节讨论针对两个使用场景加载转换器模型：微调小型转换器模型和微调大型转换器模型。对于没有延迟参数初始化的较小模型，请先torch.sagemaker.transform用 API 封装模型，然后再用 PyTorch FSDP 封装模型。

import functools
from transformers import AutoModelForCausalLM
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.sagemaker import transform

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", low_cpu_mem_usage=True)

# Transform model while loading state dictionary from rank 0.
tp_model = transform(model, load_state_dict_from_rank0=True)

# Wrap with FSDP.
model = FSDP(
    tp_model, 
    ...
    sync_module_states=True,
)

对于较大的模型，上述方法会导致 CPU 内存不足。我们建议您使用延迟参数初始化来避免此类 CPU 内存问题。在这种情况下，您可以应用 torch.sagemaker.transform API 和 torch.sagemaker.delayed_param.DelayedParamIniter API，如下代码示例所示。

from transformers import AutoModelForCausalLM
from torch.sagemaker import transform
from torch.sagemaker.delayed_param import DelayedParamIniter

# Create one instance of model without delayed param
# on CPU, on one rank.
if dist.get_rank() == 0:
    model = AutoModelForCasalLM.from_pretrained(...,low_cpu_mem_usage=True)
else:
    with init_empty_weights():
        model = AutoModelForCasalLM.from_config(AutoConfig.from_pretrained(...))

# Transform model while loading state dictionary from rank 0
model = transform(model, load_state_dict_from_rank0=True)

if dist.get_rank() != 0: # For fine-tuning, delayed parameter on non-zero ranks
    delayed_initer = DelayedParamIniter(model)
else:
    delayed_initer = None

with (
    delayed_initer.validate_params_and_buffers_inited() if delayed_initer else nullcontext()
):
    # Wrap the model with FSDP
    model = FSDP(
        model, 
        ..., 
        sync_module_states=True,
        param_init_fn=delayed_initer.get_param_init_fn() if delayed_initer else None
    )