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 que a criação dos parâmetros do modelo ocorra 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 é diferente do nativo PyTorch v2.0.1, que tem um bug que faz com que os parâmetros sejam inicializados várias vezes.

O SMP v2 fornece a API torch.sagemaker.param_delayed_. 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 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 esse processo de preparação dessa init_weights função, 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 amarrar os pesos após inicializar os pesos com atraso na inicialização dos parâmetros. PyTorch O FSDP não tem um mecanismo para amarrar os pesos após inicializá-los usando como descrito 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()
    )