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

Este notebook demonstra como ajustar um modelo Llama 3.1 8B usando o Mosaic LLM Foundry em uma GPU sem servidor Databricks . 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.
• GPU sem servidorDatabricks : Provisionamento e gerenciamento automático de recurso compute de GPU
• Unity Catalog : Armazena pontos de verificação de modelos e registra modelos treinados.
• MLflow : rastreia experimentos e logs treinamentos detalhados

Instale a biblioteca necessária.

Instale o Mosaic LLM Foundry e suas dependências para treinamento distribuído:

• llm-foundryEstrutura central para treinamento e aprimoramento do LLM
• mlflow: Envio de experimento e registro de modelo
• flash-attentionImplementação otimizada de atenção para treinamento mais rápido.
• hf_transferdownloads de modelos mais rápidos da Hugging Face
• yamlmagicPermite a configuração YAML em células do Notebook.

Python

%pip install llm-foundry[gpu]==0.20.0
%pip install mlflow==3.6
%pip install matplotlib==3.10.0
%pip install --force-reinstall --no-cache-dir --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 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 provisiona automaticamente compute de GPU serverless e gerencia a orquestração de treinamento distribuída.

Python

from serverless_gpu import distributed
from serverless_gpu import runtime as sgc_runtime
from llmfoundry.command_utils.train import train
from omegaconf import DictConfig
import mlflow
import yaml
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 aba