Ajuste fino Llama 3.1 8B usando Mosaic LLM Foundry em GPU sem servidor Databricks

Open in Databricks

Ajuste fino de um modelo Llama 3.1 8B no AI Runtime usando Mosaic LLM Foundry, uma base de código para treinamento, ajuste fino, avaliação e implantação de grandes modelos de linguagem com suporte para estratégias de treinamento distribuído.

O Notebook utiliza:

• Mosaic LLM Foundry : Uma estrutura para treinamento e ajuste fino de LLMs com suporte integrado para FSDP, carregamento eficiente de dados e integração MLflow
• FSDP (Fully Sharded Data Parallel) : Distribui parâmetros do modelo, gradientes e estados do otimizador entre GPUs.
• Databricks Serverless GPU : Executa treinamento distribuído em compute de GPU serverless conectado
• Unity Catalog : Armazena pontos de verificação de modelos e registra modelos treinados.
• MLflow : rastreia experimentos e logs treinamentos detalhados

Conectar a compute de GPU serverless

Este notebook requer compute de GPU serverless. Para conectar:

1. Clique no seletor de compute do notebook no canto superior direito e selecione **GPU Serverless**.
2. Abra o painel lateral Ambientes no lado direito do notebook
3. Defina o **acelerador** para **8xH100**
4. Selecione o ambiente base **Padrão** e defina a **Versão do ambiente** como **5**, que contém as bibliotecas necessárias para executar este exemplo
5. Selecione **Aplicar** e clique em **Confirmar** para aplicar este ambiente ao seu Notebook

Instale a biblioteca necessária.

Instale o Mosaic LLM Foundry e suas dependências para treinamento distribuído. A flash-attn Python wheel pré-compilada é instalada primeiro para que pip a reutilize em vez de compilar o flash-attention a partir do código-fonte (o que é lento) ao llm-foundry[gpu] resolver:

• flash-attn: Implementação otimizada de atenção, instalada de uma wheel pré-construída
• llm-foundryEstrutura central para treinamento e aprimoramento do LLM
• hf_transferdownloads de modelos mais rápidos da Hugging Face
• yamlmagicPermite a configuração YAML em células do Notebook.

Python

%pip install --no-deps "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl"
%pip install llm-foundry[gpu]==0.20.0
%pip install hf_transfer
%pip install git+https://github.com/josejg/yamlmagic.git

Reinicie o ambiente Python.

Reinicie o kernel Python para garantir que os pacotes recém-instalados estejam disponíveis.

Python

dbutils.library.restartPython()

Configure os caminhos Unity Catalog para armazenamento de modelos.

Configure os locais Unity Catalog para armazenar os pontos de verificação do modelo e registre o modelo treinado. A configuração utiliza parâmetros de consulta que podem ser personalizados sem editar o código.

Python

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "llama3_1-8b")
dbutils.widgets.text("uc_volume", "checkpoints")

UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_MODEL_NAME = dbutils.widgets.get("uc_model_name")
UC_VOLUME = dbutils.widgets.get("uc_volume")

MLFLOW_EXPERIMENT_NAME = '/Workspace/Shared/llm-foundry-sgc' # TODO: update this name

print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_MODEL_NAME: {UC_MODEL_NAME}")
print(f"UC_VOLUME: {UC_VOLUME}")
print(f"EXPERIMENT_NAME: {MLFLOW_EXPERIMENT_NAME}")

# Model selection - Choose based on your compute constraints
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}" # Save checkpoint to UC Volume

print(f"OUTPUT_DIR: {OUTPUT_DIR}")

Defina a configuração de treinamento usando YAML.

Carregue a configuração de ajuste fino a partir do formato YAML. A configuração especifica:

• Arquitetura do modelo e pesos pré-treinados (Llama 3.1 8B)
• Configurações FSDP para treinamento distribuído
• hiperparâmetros de treinamento (taxa de aprendizado, tamanho do lote, otimizador)
• Configuração do conjunto de dados (mosaicml/dolly_hhrlhf)
• Registro de logs e checkpoints de modelos do MLflow
• Callbacks para monitoramento e otimização

Python

%load_ext yamlmagic

Python

%%yaml config
seed: 17
model:
  name: hf_causal_lm
  pretrained: true
  init_device: mixed
  use_auth_token: true
  use_flash_attention_2: true
  pretrained_model_name_or_path: meta-llama/Llama-3.1-8B
loggers:
  mlflow:
    resume: true
    tracking_uri: databricks
    rename_metrics:
      time/token: time/num_tokens
      lr-DecoupledLionW/group0: learning_rate
    log_system_metrics: true
    experiment_name: "mlflow_experiment_name"
    run_name: llama3_8b-finetune
    model_registry_uri: databricks-uc
    model_registry_prefix: main.linyuan
callbacks:
  lr_monitor: {}
  run_timeout:
    timeout: 7200
  scheduled_gc:
    batch_interval: 1000
  speed_monitor:
    window_size: 10
  memory_monitor: {}
  runtime_estimator: {}
  hf_checkpointer:
    save_folder: "dbfs:/Volumes/main/sgc/checkpoints/llama3_1-8b-hf"
    save_interval: "1ep"
    precision: "bfloat16"
    overwrite: true

    mlflow_registered_model_name: "main.sgc.llama3_1_8b_full_ft"
    mlflow_logging_config:
      task: "llm/v1/completions"
      metadata:
        pretrained_model_name: "meta-llama/Llama-3.1-8B-Instruct"
optimizer:
  lr: 5.0e-07
  name: decoupled_lionw
  betas:
  - 0.9
  - 0.95
  weight_decay: 0
precision: amp_bf16
scheduler:
  name: linear_decay_with_warmup
  alpha_f: 0
  t_warmup: 10ba
tokenizer:
  name: meta-llama/Llama-3.1-8B
  kwargs:
    model_max_length: 1024
algorithms:
  gradient_clipping:
    clipping_type: norm
    clipping_threshold: 1
autoresume: false
log_config: false
fsdp_config:
  verbose: false
  mixed_precision: PURE
  state_dict_type: sharded
  limit_all_gathers: true
  sharding_strategy: FULL_SHARD
  activation_cpu_offload: false
  activation_checkpointing: true
  activation_checkpointing_reentrant: false
max_seq_len: 1024
save_folder: "output_folder"
dist_timeout: 600
max_duration: 20ba
progress_bar: false
train_loader:
  name: finetuning
  dataset:
    split: test
    hf_name: mosaicml/dolly_hhrlhf
    shuffle: true
    safe_load: true
    max_seq_len: 1024
    packing_ratio: auto
    target_prompts: none
    target_responses: all
    allow_pad_trimming: false
    decoder_only_format: true
  timeout: 0
  drop_last: false
  pin_memory: true
  num_workers: 8
  prefetch_factor: 2
  persistent_workers: true
eval_interval: 1
save_interval: 1h
log_to_console: true
save_overwrite: true
python_log_level: debug
save_weights_only: false
console_log_interval: 10ba
device_eval_batch_size: 1
global_train_batch_size: 32
device_train_microbatch_size: 1
save_num_checkpoints_to_keep: 1

Python

config["loggers"]["mlflow"]["experiment_name"] = MLFLOW_EXPERIMENT_NAME
config["save_folder"] = OUTPUT_DIR
config["callbacks"]["hf_checkpointer"]["save_folder"] = OUTPUT_DIR
config["callbacks"]["hf_checkpointer"]["mlflow_registered_model_name"] = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

Defina a função de treinamento distribuído

Esta célula define a função de treinamento que será executada em 8 GPUs H100 utilizando o decorador @distributed . A função:

• Configura os tokens Hugging Face para acesso ao modelo.
• Permite downloads rápidos de modelos com hf_transfer
• Chama a função train() do LLM Foundry com a configuração YAML.
• Retorna o ID de execução MLflow para acompanhamento do experimento

O decorador @distributed executa a função no compute de GPU serverless conectado e gerencia a orquestração de treinamento distribuído.

Python

from serverless_gpu import distributed
from llmfoundry.command_utils.train import train
from omegaconf import DictConfig
import mlflow
from huggingface_hub import constants

HF_TOKEN = dbutils.secrets.get(scope="sgc-nightly-notebook", key="hf_token")

@distributed(gpus=8, gpu_type='H100')
def run_llm_foundry():
    import os
    import logging
    os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
    constants.HF_HUB_ENABLE_HF_TRANSFER = True
    train(DictConfig(config))

    logging.info("\n✓ Training completed successfully!")

    mlflow_run_id = None
    if mlflow.last_active_run() is not None:
        mlflow_run_id = mlflow.last_active_run().info.run_id

    return mlflow_run_id

execução do treinamento distribuído Job

Execute a função de treinamento em 8 GPUs H100. A função retorna o ID de execução MLflow , que pode ser usado para rastrear métricas, view logs e acessar o modelo treinado na interface do usuário MLflow .

Python

mlflow_run_id = run_llm_foundry.distributed()[0]
print(mlflow_run_id)

Próximos passos

• Treinamento distribuído com múltiplas GPUs e múltiplos nós
• Melhores práticas para computede GPU sem servidor
• Solucionar problemas em computeGPU serverless
• Documentação da Mosaic LLM Foundry
• Unity Catalog registro de modelo

Exemplo de caderno

Ajuste fino Llama 3.1 8B usando Mosaic LLM Foundry em GPU sem servidor Databricks

Abrir notebook em uma nova abaOpen in Databricks