延遲參數初始化

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

SMP v2 提供用於套用延遲參數初始化的 torch.sagemaker.delayed_param.DelayedParamIniter 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 Transformer 的 Llama。torch.sagemaker.delayed_param.DelayedParamIniter API 也適用於 SMP 張量平行實作torch.sagemaker.tensor_parallel.transformer.TransformerLMHead模型，您可以在 torch.sagemaker.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()
    )

綁定權重的備註

使用綁定權重訓練模型時，在初始化權重和延遲參數初始化之後，我們需要特別注意綁定權重。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()
    )