Compartilhar via

Ajustar modelos do Hugging Face para uma única GPU

  • Artigo

Este artigo descreve como fazer ajuste fino em um modelo do Hugging Face com a biblioteca transformers do Hugging Face em uma única GPU. Ele também inclui recomendações específicas do Databricks para carregar dados do lakehouse e modelos de registro em log para MLflow, o que habilita você a usar e governar seus modelos no Azure Databricks.

A biblioteca transformers do Hugging Face fornece o utilitário Trainer e as classes de Modelo Automático que habilitam o carregamento e o ajuste fino de modelos Transformers.

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

  • Carregando modelos para ajuste fino.
  • Construindo a configuração para o utilitário Hugging Face Transformers Trainer.
  • Realizando o treinamento em uma única GPU.

Consulte O que são Hugging Face Transformes?

Requisitos

Tokenize um conjunto de dados de Hugging Face

Os modelos de Hugging Face Transformers esperam uma entrada tokenizada, em vez do texto nos dados baixados. Para garantir a compatibilidade com o modelo base, use um AutoTokenizer carregado a partir 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 do Hugging Face podem ser usadas para configurar um Trainer. As classes do Trainer exigem que o usuário forneça:

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

Você pode configurar as métricas de avaliação além da métrica loss padrão que Trainer computa. O seguinte exemplo demonstra a adição de accuracy como 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 de Auto Modelo para NLP para carregar o modelo apropriado para sua tarefa.

Para classificação de texto, use AutoModelForSequenceClassification para carregar um modelo base para a 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 que você especifique o diretório de saída, a estratégia de avaliação, a taxa de aprendizado e outros parâmetros.

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

Usando um agrupador de dados, a entrada de lotes em conjuntos de dados de treinamento e avaliação. O DataCollatorWithPadding oferece um bom desempenho de linha de base para a 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,
)

Treinamento e registro para MLflow

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

Encapsulamento do treinamento em uma execução do MLflow. Isso constrói um pipeline de Transformers a partir do tokenizador e do modelo treinado e o grava no disco local. Por fim, 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 de criar um pipeline, você pode enviar os componentes que são 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!"],
)

Carregar o modelo para inferência

Quando seu modelo estiver registrado e pronto, carregar o modelo para inferência é o mesmo que carregar o modelo pré-treinado encapsulado do 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 Serviço de Modelo do Azure Databricks para obter mais informações.

Solucionar problemas comuns de erros CUDA

Esta seção descreve os erros comuns de CUDA e as diretrizes sobre como resolvê-los.

OutOfMemoryError: CUDA sem memória

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

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.

Experimente as seguintes recomendações para resolver esse erro:

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

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

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

  • Use Otimizador Adam de 8 bits.

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

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

Erro do kernel CUDA

Ao executar o treinamento, você pode receber 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 solucionar problemas:

  • Tente executar o código na CPU para ver se o erro pode ser reproduzido.

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

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

Notebook: ajuste fino de classificação de texto em uma única GPU

Para começar a usar rapidamente o código de exemplo, este notebook de exemplo fornece um exemplo completo de ajuste fino de um modelo para classificação de texto. As seções subsequentes deste artigo entram em mais detalhes sobre como usar o Hugging Face para ajuste fino no Azure Databricks.

Notebook de ajuste fino de classificação de texto com Hugging Face

Obter notebook

Recursos adicionais

Saiba mais sobre Hugging Face no Azure Databricks.