Inicialização do parâmetro atrasada

Nem sempre é possível a inicialização de um modelo grande para treinamento com a memória limitada da GPU. Para resolver esse problema de memória insuficiente de GPU, você pode inicializar o modelo na memória da CPU. Porém, para modelos maiores com mais de 20 ou 40 bilhões de parâmetros, até a memória da CPU pode não ser suficiente. Nesse caso, recomendamos que você inicialize o modelo no que PyTorch chama um meta-dispositivo, o que permite a criação de tensores sem nenhum dado anexado a eles. Um tensor em um dispositivo meta precisa apenas das informações de forma, e permite criar um modelo grande com parâmetros em dispositivo metas. O Hugging Face Accelerate fornece o gerenciador de contexto init_empty_weights para auxiliar a criação desse modelo no dispositivo meta, enquanto inicializa os buffers em um dispositivo comum. Antes do início do treinamento, o PyTorch FSDP inicializa os parâmetros do modelo. Esse recurso de inicialização retardada de parâmetros do SMP v2 atrasa a criação dos parâmetros do modelo após o PyTorch FSDP realizar a fragmentação de parâmetros. PyTorch O FSDP aceita uma função de inicialização de parâmetros (param_init_fn) ao fragmentar os módulos e chama param_init_fn cada módulo. A API param_init_fn usa um módulo como argumento e inicializa todos os parâmetros, sem incluir os parâmetros de nenhum módulo filho. Observe que esse comportamento difere da PyTorch versão 2.0.1 nativa, que tem um bug que faz com que os parâmetros sejam inicializados várias vezes.

O SMP v2 fornece a API torch.sagemaker.delayed_param.DelayedParamIniter para aplicar a inicialização atrasada de parâmetros.

Os trechos de código a seguir demonstram como aplicar a API torch.sagemaker.delayed_param.DelayedParamIniter ao script de treinamento.

Suponha que você tenha um script de treinamento PyTorch do FSDP da seguinte forma.

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

Observe que a abordagem de inicialização atrasada de parâmetros não é independente do modelo. Para resolver esse problema, você precisa escrever uma função init_weights, conforme demonstrado no exemplo anterior, para corresponder à inicialização na definição do modelo original e ela deve abranger todos os parâmetros do modelo. Para simplificar o processo de preparação dessa função init_weights, o SMP v2 implementa essa função de inicialização para os seguintes modelos: GPT-2, GPT-J, GPT-NeoX e Llama da Hugging Face Transformers. A API torch.sagemaker.delayed_param.DelayedParamIniter também funciona com a implementação paralela do tensor do SMP, modelo torch.sagemaker.tensor_parallel.transformer.TransformerLMHead, que você pode chamar após a chamada da API torch.sagemaker.transform.

Usando a torch.sagemaker.delayed_param.DelayedParamIniter API, você pode adaptar seu script PyTorch FSDP da seguinte forma. Depois de criar um modelo com pesos vazios, registre a API torch.sagemaker.delayed_param.DelayedParamIniter no modelo e defina um objeto. Passe o objeto para o param_init_fn da classe 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()
    )

Notas sobre pesos iguais

Ao treinar modelos com pesos empatados, precisamos tomar cuidado especial para empatar os pesos após inicializar os pesos com a inicialização atrasada dos parâmetros. PyTorch O FSDP não tem um mecanismo para amarrar os pesos após inicializá-los usando as instruções acima. param_init_fn Para resolver esses casos, adicionamos uma API para permitir post_init_hook_fn, que pode ser usada para igualar os pesos. Você pode passar qualquer função que aceite o módulo como argumento, mas também temos um método predefinido post_param_init_fn no DelayedParamIniter, que chama o método tie_weights do módulo, se houver. Observe que é seguro sempre passar post_param_init_fn mesmo que não haja um método tie_weights para o módulo.

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