Ajuste fino de um modelo de incorporação usando aprendizado contrastivo.

Este Notebook demonstra como ajustar um modelo de incorporação no estilo BERT em compute GPU serverless usando aprendizado contrastivo. Você usará o modelo gte-large-en-v1.5 e o treinará em uma única GPU A10G com o treinador MosaicML Composer para salvar pontos de verificação, retomar o treinamento e log os resultados no MLflow.

Os modelos de incorporação são amplamente utilizados em bancos de dados vetoriais e em aplicações de geração aumentada de recuperação (RAG). Ajustar um modelo de incorporação aos seus dados personalizados é uma maneira eficaz de melhorar a precisão da recuperação de informações para o seu domínio específico.

Neste caderno, você aprenderá como:

• Instale as dependências e configure seu ambiente.
• Carregar dados de treinamento de tabelas Delta
• Converter dados para o formato Mosaic transmissão dataset (MDS)
• Configure o modelo, o otimizador e os parâmetros de treinamento.
• Ensine o modelo usando aprendizagem contrastiva com negativos de entrada.
• Acompanhe os experimentos com modelos MLflow e registro no Unity Catalog
• Analise o desempenho e disponibilize o modelo otimizado.

Este exemplo utiliza uma versão pré-processada do dataset MS Marco , mas você pode adaptá-lo para funcionar com seus próprios dados.

Visão geral da aprendizagem contrastiva

A aprendizagem contrastiva utiliza a perda contrastiva para aprender representações de dados onde instâncias semelhantes estão mais próximas umas das outras no espaço latente. Para modelos de incorporação, isso significa tratar consultas como "o que é um affenpinscher" e "dachshund fofo" como mais semelhantes do que "o que é um affenpinscher" e "como usar o dbsql". O modelo compara uma consulta com trechos positivos (textos relevantes) e trechos negativos (textos irrelevantes) para aprender as diferenças semânticas.

Duas abordagens para selecionar trechos negativos:

• Negativos dentro dos lotes : As passagens negativas são selecionadas aleatoriamente dentro dos lotes. Para um determinado par de texto-consulta, todos os outros textos nos lotes tornam-se exemplos negativos. Com um tamanho de lote de 8, você obtém 7 passagens negativas e 1 passagem positiva por consulta. Lotes maiores fornecem mais exemplos negativos , tornando essa abordagem mais eficaz.

• Negativas difíceis : Passagens negativas predefinidas que são semanticamente desafiadoras — elas podem estar relacionadas à consulta, mas são ligeiramente incorretas ou irrelevantes. O código llm-foundry suporta valores negativos rígidos para ajustes mais precisos.

Este caderno utiliza negativos embutidos. Caso não sejam fornecidas passagens negativas nos seus dados, o código llm-foundry as inferirá automaticamente, tratando as passagens positivas de outras consultas como negativas.

Requisitos

Antes de executar este Notebook, você precisa configurar vários parâmetros e conectar-se ao compute GPU serverless .

Configurar parâmetros do Notebook

Este Notebook utiliza parâmetros de consulta (widgets) para configurar caminhos e definições. Atualize os seguintes parâmetros antes de executar:

• catalog: Nome do Unity Catalog (ex: main)
• schema: Nome do esquema Unity Catalog
• train_delta_tableNome da tabela Delta de treinamento (sem prefixo de catálogo/esquema)
• val_delta_table: Nome da tabela Delta de validação (opcional)
• uc_checkpoint_folder: Nome da pasta de volume Unity Catalog para pontos de verificação
• register_toNome do modelo para registro Unity Catalog
• experiment_name: Caminho do experimento MLflow (formato: /Users/<username>/<run_name>)

Conecte-se à computeGPU serverless

Este notebook requer uma única GPU A10G:

1. Clique no dropdown Conectar na parte superior do Bloco de Anotações.
2. Selecione GPU sem servidor .
3. Abra o painel lateral Ambiente , localizado no lado direito do Notebook.
4. Configure o acelerador para A10 para esta demonstração.
5. Selecione Aplicar e clique em Confirmar para aplicar este ambiente ao seu Notebook.

Para obter mais informações, consulte computeGPU sem servidor.

Instalar dependências

Para começar, instale todas as bibliotecas necessárias e certifique-se de que seu ambiente esteja pronto para uso.

Python

%pip install llm-foundry[gpu]==0.20.0
%pip uninstall flash_attn -y
%pip install transformers==4.46.0
%pip install hf_transfer
%restart_python

Configurar variável de ambiente

Configure variável de ambiente para treinamento distribuído e armazenamento temporário de arquivos.

Python

import os
import tempfile
import mlflow

os.environ["TMPDIR"] = os.path.join(os.getcwd(), tempfile.mkdtemp())
os.environ["NCCL_DEBUG"] = "WARN"
os.environ["WORLD_SIZE"] = "1"

Criar widgets de configuração

Crie widgets de entrada para os parâmetros do Notebook. Preencha esses valores no painel de widgets na parte superior do Notebook antes de prosseguir.

Python

# Create widgets for configuration
dbutils.widgets.text(
    "train_delta_table", "ms_marco_v_1_1_train_processed", "Training Delta Table Name"
)
dbutils.widgets.text(
    "val_delta_table", "ms_marco_v_1_1_val_processed", "Validation Delta Table Name"
)
dbutils.widgets.text("register_to", "sgc_ft_embedding", "Model Registry Path")
dbutils.widgets.text("experiment_name", "/Users/<EMAIL>/Embedding_finetuning", "MLflow Experiment Name")
dbutils.widgets.text("uc_checkpoint_folder", "checkpoints", "UC Checkpoint Folder")
dbutils.widgets.text("catalog", "main", "catalog")
dbutils.widgets.text("schema", "default", "schema")

Python

# Validate widget inputs
assert dbutils.widgets.get("train_delta_table")
assert dbutils.widgets.get("register_to")
assert dbutils.widgets.get("experiment_name")
assert dbutils.widgets.get("catalog")
assert dbutils.widgets.get("schema")
assert dbutils.widgets.get("uc_checkpoint_folder")

Python

# Build env paths
catalog = dbutils.widgets.get("catalog")
schema = dbutils.widgets.get("schema")
train_delta_table = dbutils.widgets.get("train_delta_table")
val_delta_table = dbutils.widgets.get("val_delta_table")
uc_checkpoint_folder = dbutils.widgets.get("uc_checkpoint_folder")
register_to = dbutils.widgets.get("register_to")
experiment_name = dbutils.widgets.get("experiment_name")

train_delta_table = f"{catalog}.{schema}.{train_delta_table}"
val_delta_table = f"{catalog}.{schema}.{val_delta_table}" if val_delta_table else None
uc_checkpoint_path = f"{catalog}.{schema}.{uc_checkpoint_folder}"
register_to_path = f"{catalog}.{schema}.{register_to}"

Carregar dados de treinamento de tabelas Delta

Carregue seus dados de treinamento e validação das tabelas Delta. Os dados devem conter as seguintes colunas:

• query_text: O texto da consulta ou pergunta
• positive_passageTrecho de texto relevante para a consulta
• negative_passages (opcional): Conjunto de passagens negativas rígidas predefinidas

Este exemplo utiliza uma versão pré-processada do dataset MS Marco . Se você não fornecer negative_passages, o código de treinamento usará automaticamente números negativos em lotes.

Python

df_train = spark.table(train_delta_table)
df_val = None

if val_delta_table:
    df_val = spark.table(val_delta_table)

MODEL_REGISTRY_PREFIX = f"{catalog}.{schema}"
REGISTERED_MODEL_NAME = register_to

EXPERIMENT_NAME = experiment_name
UC_CHECKPOINT_PATH = f"{catalog}.{schema}.{uc_checkpoint_folder}"

Configurar credenciais do Databricks

Configure as credenciais CLI Databricks para acessar Unity Catalog e MLflow durante o treinamento. Este arquivo de configuração permite que o processo de treinamento se autentique com o serviço Databricks .

Python

%sh echo -e "[DEFAULT]\nhost=$DATABRICKS_HOST\ntoken=$DATABRICKS_TOKEN" > ~/.databrickscfg

Converter dados para o formato de conjunto de dados de transmissão Mosaic

Converta suas tabelas Delta para o formato Mosaic Transmissão Dataset (MDS), que é otimizado para treinamento distribuído em compute GPU serverless . A MDS oferece:

• Carregamento de dados mais rápido durante o treinamento
• Compressão e armazenamento eficientes
• Integração perfeita com o treinador Composer

A função de conversão realiza as transformações de esquema necessárias e salva os arquivos MDS em um volume Unity Catalog . Se você tiver nomes de coluna diferentes em seus dados, atualize a função convert_x de acordo.

Para mais informações, consulte a documentação do StreamingDataset.

Python

import os
import gc
from streaming import MDSWriter, StreamingDataset
from pyspark.sql import DataFrame
import json
import warnings
warnings.filterwarnings("ignore", module="threadpoolctl")


def process_embedding_data(
    df: DataFrame,
    output_path: str,
    compression: str,
    hashes: list[str],
    limit: str,
):
    def convert_x(x: dict) -> dict:
        return {
            "query_text": x["query_text"],
            "positive_passage": x["positive_passage"],
            "negative_passages": (
                json.dumps(x["negative_passages"])
                if x.get("negative_passages") is not None
                else "[]"
            ),
        }

    try:
        dtypes = {
            "query_text": "str",
            "positive_passage": "str",
            "negative_passages": "str",
        }

        print(f"Starting conversion to MDS at {output_path}")

        # Clear memory before processing
        gc.collect()

        row_count = 0
        with MDSWriter(
            out=output_path,
            columns=dtypes,
            compression=compression,
            hashes=hashes,
            size_limit=limit,
        ) as out:
            for row in df.toLocalIterator():
                record = convert_x(row.asDict())
                out.write(record)
                row_count += 1
                if row_count % 10000 == 0:
                    print(f"Processed {row_count} records...")

        print(f"Successfully wrote {row_count} records to {output_path}")

    except Exception as e:
        print(f"Error during data conversion: {e}")
        raise e

Python

compression = "zstd:7"
hashes = ["sha1"]
limit = "10mb"

import os

# FUSE path for checking existence
uc_train_folder = f"/Volumes/{catalog}/{schema}/embedding_temp_data/train"
uc_val_folder = f"/Volumes/{catalog}/{schema}/embedding_temp_data/val"

# dbfs path for MDSWriter and StreamingDataset
train_folder = f"dbfs:/Volumes/{catalog}/{schema}/embedding_temp_data/train"
val_folder = f"dbfs:/Volumes/{catalog}/{schema}/embedding_temp_data/val"

train_index = os.path.join(uc_train_folder, "index.json")
val_index = os.path.join(uc_val_folder, "index.json")

if os.path.exists(train_index):
    print("Train MDS data already exists, skipping conversion.")
else:
    process_embedding_data(df_train, train_folder, compression, hashes, limit)

if df_val is not None:
    if os.path.exists(val_index):
        print("Validation MDS data already exists, skipping conversion.")
    else:
        process_embedding_data(df_val, val_folder, compression, hashes, limit)

Configure o modelo de incorporação

Defina o modelo e a configuração do tokenizador para ajustes finos. Este exemplo utiliza o modelo gte-large-en-v1.5 da Hugging Face.

Parâmetros de configuração principais:

• temperatureHiperparâmetro (0-1) que escala as pontuações de similaridade na perda contrastiva. Ajuste este parâmetro se os valores de perda forem extremamente altos ou baixos (default: 0,5).

• pos_step_size: Posicione o tamanho do passo para amostragem negativa. Defina como 2 para números negativos dentro de lotes, ou 1 + number of hard negatives se estiver usando números negativos predefinidos.

• vector_representationComo representar embeddings

  • avgIncorporações médias de tokens (recomendadas para a maioria dos modelos)
  • eos: Use a incorporação de tokens de fim de sequência

• gather_in_batch_negativesDefina como true para negativos dentro de lotes, false para negativos rígidos predefinidos.

• pretrained_model_name_or_pathIdentificador do modelo Hugging Face

A configuração do tokenizador define o comprimento máximo da sequência e os tokens especiais para o modelo.

Python

model_cfg = {
    "name": "finetune_embedding_model",
    "trust_remote_code": True,
    "contrastive_config": {
        "temperature": 0.5,
        "pos_step_size": 2,  # set to 2 when not using predefined hard negatives. Otherwise use 1 + number of hard negatives
        "normalize_output": True,
        "vector_representation": "avg",  # or eos, depending on the model default
        "gather_in_batch_negatives": True,  # set to true when not using predefined hard negatives
    },
    "pretrained_model_name_or_path": "Alibaba-NLP/gte-large-en-v1.5",
    "loss_fn": "torch_crossentropy"
}

tokenizer_cfg = {
    "name": "Alibaba-NLP/gte-large-en-v1.5",
    "kwargs": {
        "eos_token": "</s>",  # this is the standard eos token for gte-large-en-v1.5
        "model_max_length": 128,
        "trust_remote_code": True,
    },
}

Configure o registro de logs do MLflow.

Configure o registro de logs do MLflow para rastrear métricas de treinamento, parâmetros e artefatos. A configuração do registrador especifica:

• O experimento MLflow onde a execução será feita por meio de registros.
• Unity Catalog como destino de registro de modelo
• O prefixo de catálogo e esquema para registro de modelo

Python

logger_cfg = {
    "mlflow": {
        "run_name": "finetune_embedding",
        "tracking_uri": "databricks",
        "experiment_name": EXPERIMENT_NAME,
        "model_registry_uri": "databricks-uc",
        "model_registry_prefix": MODEL_REGISTRY_PREFIX,
    }
}

Configurar retornos de chamada de treinamento

Os callbacks controlam vários aspectos do processo de treinamento. A função de retorno mais importante é hf_checkpointer, que:

• Salva pontos de verificação compatíveis Hugging Faceno Unity Catalog
• Registre o modelo no Unity Catalog para disponibilizá-lo.
• Configura metadados do modelo para provisionamento Taxa de transferência servindo
• Salva pontos de controle em intervalos regulares (a cada 1 hora neste exemplo).

Outras funções de retorno de chamada monitoram a taxa de aprendizado, o uso de memória e realizam a coleta de lixo para otimizar a memória da GPU.

Python

callback_cfg = {
    "lr_monitor": {},
    "scheduled_gc": {"batch_interval": 1000},
    "memory_monitor": {},
    "hf_checkpointer": {
        "precision": "bfloat16",
        "save_folder": UC_CHECKPOINT_PATH,
        "save_interval": "1h",
        "mlflow_logging_config": {
            "task": "llm/v1/embeddings",
            "metadata": {
                "task": "llm/v1/embeddings",
                "source": "huggingface",
                "pretrained_model_name": "Alibaba-NLP/gte-large-en-v1.5",
                "databricks_model_family": "NewModel (gte_v1_5)",
                "databricks_model_size_parameters": "434m",
            },
        },
        "mlflow_registered_model_name": REGISTERED_MODEL_NAME,
    },
}

Configurar hiperparâmetros de treinamento

Defina o otimizador, a taxa de aprendizado do programador, a precisão e os algoritmos de treinamento. Esses são parâmetros padrão de treinamento machine learning que você pode ajustar com base em seus dados e requisitos:

• Otimizador : AdamW com decaimento de peso desacoplado
• Taxa de aprendizado : 3e-5 com programa de aquecimento de cosseno
• Precisão : Precisão mista automática com bfloat16 para treinamento mais rápido.
• Limitação de gradiente : Impede a explosão de gradientes durante o treinamento.

Python

optimizer_cfg = {
    "lr": 0.00003,
    "eps": 1.0e-08,
    "name": "decoupled_adamw",
    "betas": [0.9, 0.95],
    "weight_decay": 0.0001,
}

precision_cfg = "amp_bf16"

scheduler_cfg = {"name": "cosine_with_warmup", "alpha_f": 0.02, "t_warmup": "0.06dur"}

algorithms_cfg = {
    "gradient_clipping": {"clipping_type": "norm", "clipping_threshold": 1}
}

Configurar carregadores de dados

Defina como os dados de treinamento e avaliação serão carregados durante o treinamento. Os carregadores de dados:

• Indique os dados formatados em MDS nos volumes Unity Catalog .
• Configurar pré-processamento de texto (adicionando os prefixos "query: " e "passage: ")
• Defina o comprimento máximo da sequência e o processamento de lotes.
• Ative o embaralhamento para melhor convergência do treinamento.

Certifique-se de que o caminho remote aponte para o local dos seus dados MDS convertidos.

Python

train_loader = {
    "name": "contrastive_pairs",
    "dataset": {
        "local": None,
        "split": None,
        "remote": train_folder,
        "shuffle": True,
        "max_seq_len": 128,
        "shuffle_seed": 42,
        "prepend_query": "query: ",
        "prepend_passage": "passage: ",
        "append_eos_token": True,
    },
    "drop_last": True,
    "num_workers": 8,
}

eval_loader = {
    "name": "contrastive_pairs",
    "dataset": {
        "local": None,
        "split": None,
        "remote": val_folder,
        "shuffle": True,
        "max_seq_len": 128,
        "shuffle_seed": 42,
        "prepend_query": "query: ",
        "prepend_passage": "passage: ",
        "append_eos_token": True,
    },
    "drop_last": True,
    "num_workers": 8,
}

Monte a configuração completa de treinamento.

Combine todos os componentes de configuração em um único objeto de configuração de treinamento. Isso inclui o modelo, o tokenizador, os carregadores de dados, o otimizador, os callbacks e os parâmetros de treinamento.

Python

from omegaconf import DictConfig, OmegaConf

config = {
  "seed": 42,
  "max_seq_len": 128,

  "model": model_cfg,
  "tokenizer": tokenizer_cfg,

  "loggers": logger_cfg,
  "callbacks": callback_cfg,

  "run_name": "finetune-BERT",

  "optimizer": optimizer_cfg,
  "precision": precision_cfg,
  "scheduler": scheduler_cfg,
  "algorithms": algorithms_cfg,

  "train_loader": train_loader,
  "eval_loader": eval_loader,

  "eval_first": True,
  "save_folder": 'dbfs:/databricks/mlflow-tracking/{mlflow_experiment_id}/{mlflow_run_id}/artifacts/{run_name}/checkpoints',
  "max_duration": "200ba",
  "progress_bar": False,

  "eval_interval": "50ba",
  "save_interval": "50ba",
  "log_to_console": True,
  "load_weights_only": True,
  "console_log_interval": "1ba",
  "device_eval_batch_size": 1,
  "eval_subset_num_batches": 4,
  "global_train_batch_size": 1,
  "device_train_microbatch_size": 1
}
cfg = DictConfig(config)

ensinar o modelo de incorporação

execução do processo de treinamento utilizando a biblioteca MosaicML LLM-Foundry, que fornece treinamento otimizado para incorporação de modelos. A função de treinamento:

• Utiliza o decorador @distributed para provisionamento do recurso GPU
• ensino por 200 lotes com avaliação a cada 50 lotes
• Salva pontos de verificação no Unity Catalog
• logs métricas e artefatos para MLflow
• Registre o modelo final no Unity Catalog

O treinamento pode levar vários minutos, dependendo do tamanho do seu dataset . Você pode acompanhar o progresso do experimento MLflow.

Para obter mais informações, consulte LLM-Foundry no GitHub.

Python

from serverless_gpu import distributed
from llmfoundry.command_utils.train import train
from omegaconf import DictConfig
import mlflow
import torch

@distributed(gpus=1, gpu_type='a10', remote=True)
def run_training():
    mlflow.end_run()
    trainer = train(cfg)

    run = mlflow.active_run()
    mlflow_run_id = run.info.run_id
    run_name = trainer.state.run_name
    del trainer
    mlflow.end_run()

    return mlflow_run_id, run_name

# Run training
result = run_training.distributed()
mlflow_run_id, run_name = result[0]

print(f"Run ID: {mlflow_run_id}")

# Download and load checkpoint
checkpoint_artifact_path = f"{run_name}/checkpoints/ep0-ba200-rank0.pt"
print(f"Downloading: {checkpoint_artifact_path}")
local_path = mlflow.artifacts.download_artifacts(
    artifact_uri=f"runs:/{mlflow_run_id}/{checkpoint_artifact_path}"
)
print(f"Downloaded to: {local_path}")

ckpt = torch.load(local_path, map_location="cpu", weights_only=False)

print("\nTop-level keys:")
print(list(ckpt.keys()))

if "state" in ckpt:
    print("\nState keys:")
    print(list(ckpt["state"].keys()))

    if "model" in ckpt["state"]:
        print(f"\nModel state dict: {len(ckpt['state']['model'])} keys")
        for k in list(ckpt["state"]["model"].keys())[:10]:
            print(f"  {k}")

Analise os resultados do treinamento e disponibilize o modelo.

Após a conclusão do treinamento, revise os resultados e implemente seu modelo otimizado:

Analisar as métricas de treinamento:

1. Navegue até o experimento MLflow especificado em experiment_name

   • Você também pode encontrar experimentos na página Experimentos da interface do usuário workspace

2. Selecione sua execução de treinamento para view:

   • treinamento e avaliação na tab métricas
   • Parâmetros do modelo na tab Parâmetros
   • Pontos de verificação e artefatos na tab Artefatos

3. A tab Detalhes do Modelo exibe o modelo registrado no Unity Catalog

Sirva o modelo:

1. Navegue até o modelo registrado no caminho especificado em register_to
2. Selecione a versão mais recente do modelo.
3. Clique em Servir este modelo para ser implantado usando provisionamento Taxa de transferência
4. Configure o endpoint de serviço com as suas configurações de Taxa de transferência desejadas.

Próximos passos

Agora que você ajustou um modelo de incorporação usando aprendizado contrastivo, explore estes recursos para aprender mais:

• computede GPU sem servidor - Saiba mais sobre recursos e capacidades de GPU serverless
• Melhores práticas para computeGPU serverless - Otimize suas cargas de trabalho de GPU
• APIsFoundation Model - modelos implantados e de serviço com provisionamento Taxa de transferência
• LLMDocumentação da Foundry - Explore recursos e configurações de treinamento avançado
• Documentação do StreamingDataset - Saiba mais sobre o formato MDS e sua otimização.
• Registro de modeloUnity Catalog - Gerencie e versione seus modelos

Exemplo de caderno

Ajuste fino de um modelo de incorporação usando aprendizado contrastivo.

Abrir notebook em uma nova aba