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. Também inclui recomendações específicas do Databricks para carregar dados do lakehouse e registrar modelos no MLflow, o que permite usar e controlar seus modelos no Databricks.

A biblioteca Hugging Face transformers fornece as utilidades do Trainer e as classes Auto Model que permitem carregar e ajustar os modelos Transformers.

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

  • Carregando modelos para ajustar.

  • Construindo a configuração para as utilidades do Hugging Face Transformers Trainer.

  • Realizando treinamento em uma única GPU.

Consulte O que são transformadores Hugging Face ?

Requisitos

Tokenize um conjunto de dados Hugging Face

Os modelos Hugging Face Transformers esperam entrada tokenizada, em vez do texto nos dados downloads . Para garantir a compatibilidade com o modelo base, use um AutoTokenizer carregado do modelo base. 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 do 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 básico

  • Uma configuração de treinamento

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

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 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ótulos criados durante a preparação dataset .

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 aprendizado e outros parâmetros.

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

Usando uma entrada de lotes de agrupamento de dados no dataset de treinamento e avaliação. DataCollatorWithPadding fornece 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,
)

ensina e logs no MLflow

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

Envolva o treinamento em uma execução do MLflow. Isso constrói um pipeline Transformers do tokenizador e do modelo treinado e o grava no disco local. Por fim, logs o modelo no MLflow com mlflow.transformers.logss.

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!"],
)

Carregar o modelo para inferência

Quando seu modelo estiver logs 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)

Solucionar erros comuns de CUDA

Esta seção descreve erros comuns de CUDA e orientações 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.

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

  • Reduza o tamanho dos lotes para treinamento. Você pode reduzir o valor 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 dos lotes.

  • Use o otimizador Adam de 8 bits.

  • Limpe a memória da GPU antes do treinamento. Às vezes, a memória da GPU pode estar 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 solucionar problemas:

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

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

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

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

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

Ajuste fino dos modelos de classificação de texto Hugging Face Notebook

Abra o bloco de anotações em outra guia

Recursos adicionais

Saiba mais sobre Hugging Face no Databricks.