Partilhar via


Ajuste os modelos Hugging Face para uma única GPU

Este artigo descreve como ajustar um modelo Hugging Face com a biblioteca Hugging Face transformers em uma única GPU. Este também inclui recomendações específicas do Databricks para carregar dados a partir do lakehouse e registar modelos no MLflow, que o habilitam a usar e gerir os seus modelos no Azure Databricks.

A biblioteca Hugging Face transformers fornece o utilitário Trainer e as classes Auto Model que permitem carregar e ajustar modelos Transformers.

Estas ferramentas estão disponíveis para as seguintes tarefas com modificações simples:

  • A carregar modelos para afinamento.
  • Construção da configuração para o utilitário Hugging Face Transformers Trainer.
  • Realização de treinamento em uma única GPU.

Veja O que são Hugging Face Transformers?

Requisitos

Tokenizar um conjunto de dados do Hugging Face

Os modelos Hugging Face Transformers esperam entrada tokenizada, em vez do texto nos dados baixados. Para garantir a compatibilidade com o modelo base, use um AutoTokenizer carregado do modelo base. O Hugging Face datasets permite que você aplique diretamente o tokenizador de forma consistente aos dados de treinamento e teste.

Por exemplo:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(base_model)
def tokenize_function(examples):
    return tokenizer(examples["text"], padding=False, truncation=True)

train_test_tokenized = train_test_dataset.map(tokenize_function, batched=True)

Definir a configuração de treinamento

As ferramentas de configuração de treinamento Hugging Face podem ser usadas para configurar um Trainer. As aulas do Trainer exigem que o usuário forneça:

  • Métricas
  • Um modelo base
  • Uma configuração de treinamento

Você pode configurar métricas de avaliação além da métrica padrão loss que o Trainer calcula. O exemplo a seguir demonstra a adição accuracy como uma métrica:

import numpy as np
import evaluate
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

Use as classes Auto Model para NLP para carregar o modelo apropriado para a sua tarefa.

Para classificação de texto, use AutoModelForSequenceClassification para carregar um modelo base para classificação de texto. Ao criar o modelo, forneça o número de classes e os mapeamentos de rótulo criados durante a preparação do conjunto de dados.

from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(
        base_model,
        num_labels=len(label2id),
        label2id=label2id,
        id2label=id2label
        )

Em seguida, crie a configuração de treinamento. A classe TrainingArguments permite especificar o diretório de saída, a estratégia de avaliação, a taxa de aprendizagem e outros parâmetros.

from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(output_dir=training_output_dir, evaluation_strategy="epoch")

Usar um coletor de dados para agrupar entradas em conjuntos de dados de treino e avaliação. DataCollatorWithPadding oferece um bom desempenho de linha de base para classificação de texto.

from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer)

Com todos esses parâmetros construídos, agora você pode criar um Trainer.

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_test_dataset["train"],
    eval_dataset=train_test_dataset["test"],
    compute_metrics=compute_metrics,
    data_collator=data_collator,
)

Treinar e registrar no MLflow

O Hugging Face faz uma boa interface com o MLflow e registra automaticamente as métricas durante o treinamento do modelo usando o MLflowCallback. No entanto, você mesmo deve registrar o modelo treinado.

Envolva o treinamento em uma execução MLflow. Isso constrói um pipeline dos Transformers a partir do tokenizador e do modelo treinado e armazena esse no disco local. Finalmente, registre o modelo no MLflow com mlflow.transformers.log_model.

from transformers import pipeline

with mlflow.start_run() as run:
  trainer.train()
  trainer.save_model(model_output_dir)
  pipe = pipeline("text-classification", model=AutoModelForSequenceClassification.from_pretrained(model_output_dir), batch_size=1, tokenizer=tokenizer)
  model_info = mlflow.transformers.log_model(
        transformers_model=pipe,
        artifact_path="classification",
        input_example="Hi there!",
    )

Se você não precisar criar um pipeline, poderá enviar os componentes usados no treinamento para um dicionário:

model_info = mlflow.transformers.log_model(
  transformers_model={"model": trainer.model, "tokenizer": tokenizer},
  task="text-classification",
  artifact_path="text_classifier",
  input_example=["MLflow is great!", "MLflow on Databricks is awesome!"],
)

Carregue o modelo para inferência

Quando o modelo estiver registado e pronto, carregar o modelo para inferência é o mesmo que carregar o modelo pré-treinado encapsulado pelo MLflow.

logged_model = "runs:/{run_id}/{model_artifact_path}".format(run_id=run.info.run_id, model_artifact_path=model_artifact_path)

# Load model as a Spark UDF. Override result_type if the model does not return double values.
loaded_model_udf = mlflow.pyfunc.spark_udf(spark, model_uri=logged_model, result_type='string')

test = test.select(test.text, test.label, loaded_model_udf(test.text).alias("prediction"))
display(test)

Consulte Implantar modelos usando o Mosaic AI Model Serving para obter mais informações.

Solucionar erros comuns do CUDA

Esta seção descreve erros CUDA comuns e orientações sobre como resolvê-los.

OutOfMemoryError: CUDA sem memória

Ao treinar modelos grandes, um erro comum que pode encontrar é o erro de falta de memória CUDA.

Exemplo:

OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 14.76 GiB total capacity; 666.34 MiB already allocated; 17.75 MiB free; 720.00 MiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF.

Tente as seguintes recomendações para resolver este erro:

  • Reduza o tamanho do lote para treinamento. Você pode reduzir o per_device_train_batch_size valor em TrainingArguments.

  • Use treinamento de baixa precisão. Você pode definir fp16=True em TrainingArguments.

  • Utilize gradient_accumulation_steps em TrainingArguments para aumentar efetivamente o tamanho geral do lote.

  • Use o otimizador de 8 bits Adam.

  • Limpe a memória da GPU antes do treino. Às vezes, a memória da GPU pode ser ocupada por algum código não utilizado.

    from numba import cuda
    device = cuda.get_current_device()
    device.reset()
    

Erros do kernel CUDA

Ao executar o treinamento, você pode obter erros de kernel CUDA.

Exemplo:

CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.

For debugging, consider passing CUDA_LAUNCH_BLOCKING=1.

Para resolver os problemas:

  • Tente executar o código na CPU para ver se o erro é reproduzível.

  • Outra opção é obter um melhor rastreio definindo CUDA_LAUNCH_BLOCKING=1:

    import os
    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
    

Notebook: aperfeiçoe a classificação de texto em uma única GPU

Para começar rapidamente com o código de exemplo, este bloco de anotações de exemplo fornece um exemplo de ponta a ponta para ajustar um modelo para classificação de texto. As seções subsequentes deste artigo entram em mais detalhes sobre o uso do Hugging Face para ajuste fino no Azure Databricks.

Ajuste fino de modelos de classificação de texto Hugging Face notebook

Obter bloco de notas

Recursos adicionais

Saiba mais sobre o Hugging Face no Azure Databricks.