Treinar um modelo de linguagem com uma arquitetura de transformador profundo consome muito tempo. No entanto, existem técnicas que você pode usar para acelerar o treinamento. Neste artigo, você aprenderá sobre:
- Usando
torch.compile()para acelerar o modelo - Usando acumulação de gradiente para treinar um modelo com um tamanho de lote efetivo maior
Vamos começar!
Treine um modelo mais rapidamente com torch.compile e Gradient Accumulation
Foto de François Genon. Alguns direitos reservados.
Visão geral
Este artigo está dividido em duas partes; eles são:
- Usando
torch.compile() - Acumulação de gradiente
Usando torch.compile
Quando você escreve o código do seu modelo e o executa com PyTorch, o código é executado no modo ansioso. Isso significa que o código é executado linha por linha e os resultados são armazenados na memória. Isso é nativo do Python, pois é uma linguagem interpretada. Você sabe que esse é o caso porque, ao cometer um erro em seu código, você não verá o erro até executar essa linha de código.
Executar um modelo no modo ansioso é lento. Começando com PyTorch 2.0, você pode usar torch.compile() para compilar um modelo para melhorar o desempenho. Isso gera um novo objeto de modelo que é otimizado. Não é o mesmo objeto de modelo que você criou usando nn.Modulemas compartilha os mesmos tensores do modelo unique. Você pode usar este modelo compilado para avanço, retrocesso e atualizações do otimizador normalmente.
Construir um modelo e compilá-lo como um gráfico de computação é como o TensorFlow 1.0 deveria funcionar. Isso torna a depuração mais difícil, pois o modelo executado não pode corresponder linha por linha ao código que você escreveu. Portanto, você não deve compilar seu modelo antes de executar uma avaliação e confirmar que ele está livre de erros.
Nem todos os modelos podem ser compilados. No entanto, se o seu modelo suportar compilação, você se beneficiará imediatamente da aceleração. Para compilar um modelo, tudo o que você precisa fazer é substituir o objeto do modelo antes de estar pronto para usá-lo:
... modelo = LlamaForPré-treinamento(modelo_config).para(dispositivo) modelo.carregar_estado_dict(posto de controle) modelo = tocha.compilar(modelo) ... |
Não carregue os pesos do modelo após a compilação. Isso ocorre porque o modelo compilado é um objeto que compartilha os mesmos pesos do modelo unique. Durante a compilação, o gráfico de computação é construído referenciando os tensores de peso do modelo unique. Se você carregar os pesos após a compilação, o modelo poderá não funcionar conforme o esperado.
Da mesma forma, para salvar o modelo compilado, você deve consultar o ditado de estado do modelo unique, como segue:
tocha.salvar(getattr(modelo, “_orig_mod”, modelo).estado_dict(), “modelo.pth”) |
O modelo unique pode ser acessado a partir do modelo compilado usando mannequin._orig_mod. No código acima, usamos getattr(mannequin, "_orig_mod", mannequin) para obter o modelo unique, se existir, ou usar mannequin em si, se isso não acontecer. Esta linha de código funciona tanto para modelos compilados quanto para modelos originais.
Acumulação de gradiente
Ao treinar um modelo, você provavelmente gasta duas a três vezes mais tempo na passagem para trás do que na passagem para frente. Isso ocorre porque a passagem para trás é mais intensiva em termos computacionais e usa mais memória.
Um truque fácil para acelerar o treinamento é realizar menos passes para trás. Isto pode ser conseguido aumentando o tamanho do lote: com o mesmo número de amostras de dados, um tamanho de lote maior significa menos lotes para processar.
No entanto, um tamanho de lote maior requer mais memória. Em um ambiente com restrição de memória, você pode imitar um tamanho de lote maior executando várias passagens de encaminhamento e acumulando gradientes. Isso é chamado acumulação de gradiente.
É mais fácil explicar essa ideia com código:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | .. acumular_passos = 4 para época em faixa(num_épocas): otimizador.zero_grad() para eu, lote em enumerar(carregador de dados): #obtém dados em lote input_ids, alvo_ids = lote #criar máscara de atenção: máscara causal + máscara de preenchimento máscara_de_atenção = criar_máscara_causal(input_ids.forma(1), dispositivo) + create_padding_mask(input_ids, PAD_TOKEN_ID, dispositivo) # extrai a saída do modelo logitos = modelo(input_ids, máscara_de_atenção) # perda de computação: entropia cruzada entre logits e alvo, ignorando tokens de preenchimento perda = perda_fn(logitos.visualizar(–1, logitos.tamanho(–1)), alvo_ids.visualizar(–1)) perda = perda / acumular_passos # Executa para trás, mas atualiza apenas uma vez a cada etapa `accumulate_steps` perda.para trás() se (eu + 1) % acumular_passos == 0: tocha.nn.utilitários.clip_grad_norm_(modelo.parâmetros(), 1,0) otimizador.etapa() otimizador.zero_grad() agendador.etapa() |
O ciclo de treinamento acima é um trecho do artigo anterior para treinar um modelo Llama em sua GPU native.
Normalmente, quando você executa um passe para frente, você calcula a perda. Então você liga loss.backward() para retropropagar o gradiente de perda através dos parâmetros do modelo. No PyTorch, o backward() O método é cumulativo, o que significa que os gradientes são somados. Portanto, você precisa ligar optimizer.zero_grad() explicitamente para limpar os gradientes antes de executar a passagem para trás.
No código acima, você deliberadamente não chama optimizer.zero_grad() em cada iteração. Em vez disso, você executa a retropropagação para a perda dividida por accumulate_steps. Dessa forma, os gradientes são reduzidos, mas acumulados ao longo accumulate_steps iterações. Uma vez a cada accumulate_steps iterações, você executa o otimizador para ajustar os parâmetros do modelo.
Essa abordagem produz resultados comparáveis ao uso de um tamanho de lote maior. No entanto, como você executa menos atualizações do otimizador, a programação da taxa de aprendizagem deve ser ajustada adequadamente. Isso significa que você precisa inicializar o agendador com um número diferente de etapas:
... num_training_steps = (lento(carregador de dados) // acumular_passos) * num_épocas cosine_scheduler = lr_scheduler.CosenoAnnealingLR( otimizador, T_máx=num_training_steps – num_warmup_steps, eta_min=0 ) |
Leitura adicional
Abaixo estão alguns materiais que você pode achar interessantes:
Resumo
Neste artigo, você aprendeu que usar torch.compile() pode ajudá-lo a acelerar o modelo compilando o gráfico de computação. Você também aprendeu que o acúmulo de gradiente é uma técnica para treinar com um tamanho de lote efetivo maior, acumulando gradientes de vários minilotes. Como você executa menos atualizações do otimizador dessa forma, você economiza tempo em passagens reversas e atualizações de parâmetros.