本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

# 延迟参数初始化
<a name="model-parallel-core-features-v2-delayed-param-init"></a>

在 GPU 内存有限的情况下，不一定能初始化一个大型模型进行训练。要解决 GPU 内存不足的问题，可以在 CPU 内存上初始化模型。不过，对于参数超过 200 亿或 400 亿的大型模型，即使 CPU 内存也可能不够用。在这种情况下，我们建议您在所 PyTorch 谓的*元设备*上初始化模型，这样就可以在不附加任何数据的情况下创建张量。元设备上的张量只需要形状信息，这样就可以在元设备上创建一个带有参数的大型模型。[Hugging Face Accelerate](https://huggingface.co/docs/accelerate/index) 提供了上下文管理器 `init_empty_weights`，以帮助在元设备上创建此类模型，同时在普通设备上初始化缓冲区。在训练开始之前， PyTorch FSDP 会初始化模型参数。SMP v2 的延迟参数初始化功能延迟了模型参数的创建，使其在 PyTorch FSDP 执行参数分片之后发生。 PyTorch FSDP 在对模块进行分片时接受参数初始化函数 (`param_init_fn`)，它会调`param_init_fn`用每个模块。`param_init_fn` API 将一个模块作为参数，并初始化其中的所有参数，不包括任何子模块的参数。请注意，此行为*与原生 PyTorch v2.0.1 不同*，后者存在导致参数多次初始化的错误。

SMP v2 提供了用于延迟参数初始化的 [`torch.sagemaker.delayed_param.DelayedParamIniter`](distributed-model-parallel-v2-reference.md#model-parallel-v2-torch-sagemaker-reference-delayed-param-init) API。

以下代码片段显示了如何在训练脚本中应用 `torch.sagemaker.delayed_param.DelayedParamIniter` API。

假设你有一个 PyTorch FSDP 训练脚本，如下所示。

```
# Creation of model on meta device
from accelerate import init_empty_weights
with init_empty_weights():
    model = create_model()

# Define a param init fn, below is an example for Hugging Face GPTNeoX.
def init_weights(module):
    d = torch.cuda.current_device()
    # Note that below doesn't work if you have buffers in the model
    # buffers will need to reinitialized after this call
    module.to_empty(device=d, recurse=False)
    if isinstance(module, (nn.Linear, Conv1D)):
        module.weight.data.normal_(mean=0.0, std=args.initializer_range)
        if module.bias:
            module.bias.data.zero_()
    elif isinstance(module, nn.Embedding):
        module.weight.data.normal_(mean=0.0, std=args.initializer_range)
        if module.padding_idx:
            module.weight.data[module.padding_idx].zero_()
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

# Changes to FSDP wrapper.
model = FSDP(
    model,
    ...,
    param_init_fn=init_weights
)

# At this point model is initialized and sharded for sharded data parallelism.
```

请注意，延迟参数初始化方法与模型无关。要解决此问题，您需要编写一个 `init_weights` 函数，如前面的示例所示，以匹配原始模型定义中的初始化，并且此函数应涵盖模型的所有参数。为了简化此类 `init_weights` 函数的准备过程，SMP v2 为以下模型实现了此初始化函数：GPT-2、GPT-J、GPT-NeoX 和 Hugging Face 转换器中的 Llama。`torch.sagemaker.delayed_param.DelayedParamIniter` API 还可与 SMP 张量并行实现 `torch.sagemaker.tensor_parallel.transformer.TransformerLMHead` 模型配合使用，您可以在调用 [`torch.sagemaker.transform`](distributed-model-parallel-v2-reference.md#model-parallel-v2-torch-sagemaker-reference-transform) API 后再调用该模型。

使用该 `torch.sagemaker.delayed_param.DelayedParamIniter` API，您可以按如下方式调整您的 PyTorch FSDP 脚本。创建权重为空的模型后，将 `torch.sagemaker.delayed_param.DelayedParamIniter` API 注册到此模型，并定义其对象。将对象传递给 PyTorch FSDP 类的。`param_init_fn`

```
from torch.sagemaker.delayed_param import DelayedParamIniter
from accelerate import init_empty_weights

with init_empty_weights():
    model = create_model()
    
delayed_initer = DelayedParamIniter(model)

with delayed_initer.validate_params_and_buffers_inited():
    model = FSDP(
        model,
        ...,
        param_init_fn=delayed_initer.get_param_init_fn()
    )
```

**关于并列权重的注意事项**

在训练带有绑定权重的模型时，我们需要特别注意在使用延迟参数初始化权重后绑定权重。 PyTorchFSDP 没有在使用上述方法初始化权重后绑定权重`param_init_fn`的机制。为了解决这种情况，我们添加了 API，允许使用 `post_init_hook_fn` 来绑定权重。您可以在其中传递任何接受模块作为参数的函数，但我们也在 `DelayedParamIniter` 中定义了一个预定义的 `post_param_init_fn`，如果模块中存在 `tie_weights` 方法，它就会调用此方法。请注意，即使模块没有 `tie_weights` 方法，在 `post_param_init_fn` 中传递也是安全的。

```
with delayed_initer.validate_params_and_buffers_inited():
    model = FSDP(
        model,
        ...,
        param_init_fn=delayed_initer.get_param_init_fn(),
        post_param_init_fn=delayed_initer.get_post_param_init_fn()
    )
```