Inizializzazione ritardata dei parametri

L’inizializzazione di un modello di grandi dimensioni per l’addestramento non sempre è possibile se la memoria GPU è limitata. Per risolvere un problema di memoria GPU insufficiente, puoi inizializzare il modello sulla memoria della CPU. Tuttavia, per i modelli più grandi con più di 20 o 40 miliardi di parametri, anche la memoria della CPU potrebbe risultare insufficiente. In tal caso, ti consigliamo di inizializzare il modello su quello che PyTorch chiama metadispositivo, che consente la creazione di tensori senza alcun dato ad essi allegato. Un tensore su un metadispositivo necessita solo delle informazioni sulla forma e questo permette di creare un modello di grandi dimensioni con i relativi parametri su metadispositivi. Hugging Face Accelerate fornisce il gestore di contesto init_empty_weights che facilita la creazione di tale modello su metadispositivi 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 di SMP v2 ritarda la creazione dei parametri del modello dopo che FSDP ha eseguito la suddivisione dei parametri. PyTorch PyTorch FSDP accetta una funzione di inizializzazione dei parametri (param_init_fn) durante la suddivisione dei moduli e chiama ogni modulo. param_init_fn L’API param_init_fn accetta un modulo come argomento e inizializza tutti i parametri in esso contenuti, esclusi i parametri di tutti i moduli secondari. Nota che questo comportamento è diverso dalla versione nativa PyTorch 2.0.1, che presenta un bug che causa l'inizializzazione dei parametri più volte.

SMP v2 fornisce l’API torch.sagemaker.delayed_param.DelayedParamIniter per applicare l’inizializzazione ritardata dei parametri.

I seguenti frammenti di codice mostrano come applicare l’API torch.sagemaker.delayed_param.DelayedParamIniter allo script di addestramento.

Supponiamo di avere uno script di addestramento PyTorch FSDP 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.

Tieni presente che l’approccio di inizializzazione ritardata dei parametri non è indipendente dal modello. Per risolvere questo problema, è necessario scrivere una funzione init_weights, come illustrato nell’esempio precedente, in modo che corrisponda all’inizializzazione nella definizione originale del modello e che ne includa tutti i parametri. Per semplificare il processo di preparazione della funzione init_weights, SMP v2 implementa questa funzione di inizializzazione per i seguenti modelli: GPT-2, GPT-J, GPT-NeoX e Llama di Hugging Face Transformers. L’API torch.sagemaker.delayed_param.DelayedParamIniter funziona anche con l’implementazione parallela del tensore SMP, modello torch.sagemaker.tensor_parallel.transformer.TransformerLMHead, che può essere chiamato dopo chiamata API torch.sagemaker.transform.

Utilizzando l'torch.sagemaker.delayed_param.DelayedParamIniterAPI, è possibile adattare lo script PyTorch FSDP come segue. Dopo aver creato un modello con pesi vuoti, registra l’API torch.sagemaker.delayed_param.DelayedParamIniter sul modello e definiscine un oggetto. Passa l'oggetto alla classe 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()
    )

Note sui pesi condivisi

Quando si allenano modelli con pesi uguali, è necessario prestare particolare attenzione a legare i pesi dopo aver inizializzato i pesi con l'inizializzazione ritardata dei parametri. PyTorch FSDP non dispone di un meccanismo per legare i pesi dopo averli inizializzati utilizzando quanto sopra. param_init_fn Per risolvere questi casi, abbiamo aggiunto un’API per consentire a post_init_hook_fn, da utilizzare per condividere i pesi. È possibile inserire qualsiasi funzione che accetti il modulo come argomento, ma è presente anche una definizione predefinita post_param_init_fn definita in DelayedParamIniter che chiama il metodo tie_weights del modulo, se esistente. Tieni presente che è sempre sicuro passare in post_param_init_fn anche se non esiste un metodo tie_weights 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()
    )