延遲參數初始化

有限GPU記憶體不一定可以初始化大型模型以進行訓練。若要解決GPU記憶體不足的問題，您可以在CPU記憶體上初始化模型。不過，對於參數超過 200 億或 400 億的較大模型，即使是CPU記憶體也可能不夠。在這種情況下，我們建議您在 PyTorch 呼叫中繼裝置 的 上初始化模型，這允許建立張量，而不會連接任何資料。中繼裝置上的張量只需要形狀資訊，這允許在中繼裝置上建立具有其參數的大型模型。Hugging Face Accelerate 提供內容管理員init_empty_weights，可協助在中繼裝置上建立此類模型，同時在一般裝置上初始化緩衝區。訓練開始前， PyTorch FSDP 會初始化模型參數。vSMP2 的延遲參數初始化功能會延遲在執行參數分割後 PyTorch FSDP建立模型參數。 PyTorch FSDP 在分割模組時接受參數初始化函數 （param_init_fn），並param_init_fn呼叫每個模組。param_init_fn API 會將模組作為引數，並初始化其中的所有參數，不包括任何子模組的參數。請注意，此行為與原生 v2 PyTorch .0.1 不同，其存在導致參數多次初始化的錯誤。

SMP v2 提供用於套用延遲參數初始化torch.sagemaker.delayed_param.DelayedParamIniterAPI的 。

下列程式碼片段說明如何將 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函數的程序，SMPv2 會針對下列模型實作此初始化函數：Hugging Face Transformer 中的 GPT-2、GPT-J、GPT-NeoX 和 Llama。torch.sagemaker.delayed_param.DelayedParamIniter API 也可以使用SMP張量平行實作torch.sagemaker.tensor_parallel.transformer.TransformerLMHead模型，您可以在torch.sagemaker.transformAPI呼叫後呼叫。

您可以使用 torch.sagemaker.delayed_param.DelayedParamIniter API調整指令碼 PyTorch FSDP，如下所示。建立具有空權重的模型後，請將 註冊torch.sagemaker.delayed_param.DelayedParamIniterAPI到模型，並定義其物件。將物件傳遞至 類別param_init_fn的 PyTorch FSDP 。

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()
    )

綁定權重注意事項

使用綁定權重訓練模型時，我們需要特別小心在初始化權重後與延遲參數初始化權重綁定權重。 PyTorch FSDP沒有機制在初始化權重後param_init_fn依上述方式綁定權重。為了解決此類案例，我們新增API了允許 post_init_hook_fn，可用於繫結權重的 。您可以在其中傳遞接受模組作為引數的任何函數，但我們也預先post_param_init_fn定義了模組存在時的DelayedParamIniter呼叫tie_weights方法。請注意，post_param_init_fn即使沒有模組tie_weights的方法，始終傳入是安全的。

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()
    )