Ajuste os modelos Hugging Face para uma única GPU

  • Artigo

Este artigo descreve como ajustar um modelo Hugging Face com a biblioteca Hugging Face transformers em uma única GPU. Ele também inclui recomendações específicas do Databricks para carregar dados do lakehouse e registrar modelos para MLflow, o que permite que você use e governe 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:

  • Carregando modelos para ajustar.
  • 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 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 aprendizado e outros parâmetros.

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

Usando um coletor de dados lotes de entrada em conjuntos de dados de treinamento 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 Trainerarquivo .

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 do Transformers a partir do tokenizador e do modelo treinado e o grava 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 é registrado e pronto, carregar o modelo para inferência é o mesmo que carregar o modelo pré-treinado encapsulado 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 Modelo de serviço com o Azure Databricks 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 você pode encontrar é o CUDA erro 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 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.

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

  • Use o otimizador Adam de 8 bits.

  • 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: ajuste 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 Abraço Face modelos de classificação de texto notebook

Obter o bloco de notas

Recursos adicionais

Saiba mais sobre o Hugging Face no Azure Databricks.