Inizializzazione ritardata dei parametri

L'inizializzazione di un modello di grandi dimensioni per l'addestramento non è sempre possibile con una memoria limitata. GPU Per risolvere questo problema di GPU memoria insufficiente, è possibile inizializzare il modello in CPU memoria. Tuttavia, per i modelli più grandi con più di 20 o 40 miliardi di parametri, anche la CPU memoria potrebbe non essere sufficiente. In tal caso, si consiglia di inizializzare il modello su un PyTorch cosiddetto meta-dispositivo, che consente la creazione di tensori senza alcun dato ad essi allegato. Un tensore su un metadispositivo necessita solo delle informazioni sulla forma e ciò consente di creare un modello di grandi dimensioni con i relativi parametri sui metadispositivi. Hugging Face Accelerate fornisce il gestore di contesto che aiuta init_empty_weights a creare tale modello su meta dispositivi mentre inizializza i buffer su un dispositivo normale. Prima dell'inizio dell'addestramento, PyTorch FSDP inizializza i parametri del modello. Questa funzionalità di inizializzazione ritardata dei parametri della SMP v2 ritarda la creazione dei parametri del modello dopo aver eseguito la suddivisione dei parametri. PyTorch FSDP PyTorch FSDPaccetta una funzione di inizializzazione dei parametri (param_init_fn) durante la suddivisione dei moduli e chiama ogni modulo. param_init_fn param_init_fnAPIPrende un modulo come argomento e inizializza tutti i parametri in esso contenuti, esclusi i parametri di alcun modulo figlio. Nota che questo comportamento è diverso dalla versione nativa PyTorch 2.0.1 che presenta un bug che causa l'inizializzazione dei parametri più volte.

SMPv2 fornisce l'applicazione dell'inizializzazione ritardata dei torch.sagemaker.delayed_param.DelayedParamIniter API parametri.

I seguenti frammenti di codice mostrano come applicarlo allo script di addestramento torch.sagemaker.delayed_param.DelayedParamIniterAPI.

Si supponga di disporre di uno script PyTorch FSDP di allenamento come segue.

# 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.

Si noti che l'approccio di inizializzazione ritardata dei parametri non è indipendente dal modello. Per risolvere questo problema, è necessario scrivere una init_weights funzione, come illustrato nell'esempio precedente, in modo che corrisponda all'inizializzazione nella definizione originale del modello e che copra tutti i parametri del modello. Per semplificare questo processo di preparazione di tale init_weights funzione, SMP v2 implementa questa funzione di inizializzazione per i seguenti modelli: GPT -2, GPT -J, GPT -NeoX e Llama di Hugging Face Transformers. Funziona torch.sagemaker.delayed_param.DelayedParamIniter API anche con l'implementazione SMP parallela del tensore, torch.sagemaker.tensor_parallel.transformer.TransformerLMHead modello, che puoi chiamare dopo la torch.sagemaker.transform API chiamata.

Usando torch.sagemaker.delayed_param.DelayedParamIniterAPI, puoi adattare PyTorch FSDP lo script come segue. Dopo aver creato un modello con pesi vuoti, registratelo nel modello e ne definite un oggetto. torch.sagemaker.delayed_param.DelayedParamIniter API Passate l'oggetto param_init_fn alla PyTorch FSDP classe.

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

Note sui pesi legati

Quando si allenano modelli con pesi legati, dobbiamo prestare particolare attenzione a legare i pesi dopo aver inizializzato i pesi con l'inizializzazione ritardata dei parametri. PyTorchFSDPnon dispone di un meccanismo per legare i pesi dopo averli inizializzati utilizzando quanto sopra. param_init_fn Per risolvere questi casi abbiamo aggiunto API a allow apost_init_hook_fn, che può essere usato per legare i pesi. È possibile inserire qualsiasi funzione che accetti il modulo come argomento, ma abbiamo anche una definizione post_param_init_fn predefinita in DelayedParamIniter cui chiama il tie_weights metodo del modulo se esiste. Nota che è sempre sicuro passare post_param_init_fn anche se non esiste un tie_weights metodo per il modulo.

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