Ajustar Llama 3.1 8B mediante Mosaic LLM Foundry en Databricks Serverless GPU

Abrir la versión 'notebook' de esta página

Ajusta un modelo Llama 3.1 8B en AI Runtime con Mosaic LLM Foundry, una base de código para entrenar, realizar el ajuste fino, evaluar y desplegar modelos de lenguaje de gran tamaño, con soporte para estrategias de entrenamiento distribuido.

El cuaderno usa:

  • Mosaic LLM Foundry: un marco para entrenar y ajustar las MÁQUINAS VIRTUALES con compatibilidad integrada con FSDP, carga eficaz de datos y integración de MLflow
  • FSDP (datos totalmente particionados paralelos): distribuye los parámetros del modelo, los degradados y los estados del optimizador entre GPU.
  • Databricks Serverless GPU: ejecuta el entrenamiento distribuido en recursos de cómputo GPU serverless conectados
  • Catálogo de Unity: almacena puntos de control de modelo y registra modelos entrenados
  • MLflow: realiza un seguimiento de experimentos y registra métricas de entrenamiento

Conectar al cómputo de GPU sin servidor

Este notebook requiere computación GPU sin servidor. Para conectarse:

  1. Haga clic en el selector de proceso del cuaderno en la parte superior derecha y seleccione GPU sin servidor.
  2. Abra el panel lateral Entorno en el lado derecho del cuaderno.
  3. Establecer acelerador en 8xH100
  4. Seleccione el entorno base estándar y establezca Versión del entorno en 5, que contiene las bibliotecas necesarias para ejecutar este ejemplo.
  5. Seleccione Aplicar y haga clic en Confirmar para aplicar este entorno al cuaderno.

Instalación de bibliotecas necesarias

Instale Mosaic LLM Foundry y sus dependencias para el entrenamiento distribuido. El paquete wheel precompilado flash-attn se instala primero para que pip lo reutilice en lugar de compilar flash-attention desde el código fuente (lo cual es lento) cuando resuelve llm-foundry[gpu]:

  • flash-attn: implementación de atención optimizada, instalada desde una rueda precompilada
  • llm-foundry: marco básico para el entrenamiento y ajuste de LLM
  • hf_transfer: descargas de modelos más rápidas de Hugging Face
  • yamlmagic: Permite configuración YAML en las celdas del notebook.
%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

Reinicio del entorno de Python

Reinicie el kernel de Python para asegurarse de que los paquetes recién instalados están disponibles.

dbutils.library.restartPython()

Configuración de rutas de acceso del catálogo de Unity para el almacenamiento de modelos

Configure las ubicaciones del catálogo de Unity para almacenar puntos de control del modelo y registrar el modelo entrenado. La configuración usa parámetros de consulta que se pueden personalizar sin editar el código.

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}")

Definición de la configuración de entrenamiento mediante YAML

Cargue la configuración de ajuste fino desde el formato YAML. La configuración especifica:

  • Arquitectura del modelo y pesos previamente entrenados (Llama 3.1 8B)
  • Configuración de FSDP para el entrenamiento distribuido
  • Hiperparámetros de entrenamiento (velocidad de aprendizaje, tamaño de lote, optimizador)
  • Configuración del conjunto de datos (mosaicml/dolly_hhrlhf)
  • Registro de eventos y control de puntos de control de modelos en MLflow
  • Callbacks para la supervisión y optimización
%load_ext yamlmagic

%%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

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}"

Definición de la función de entrenamiento distribuido

Esta celda define la función de entrenamiento que se ejecutará mediante el decorador @distributed en 8 GPUs H100. La función :

  • Configura el token de Hugging Face para el acceso al modelo.
  • Habilita las descargas rápidas de modelos con hf_transfer
  • Llama a la función LLM Foundry train() con la configuración YAML.
  • Devuelve el identificador de ejecución de MLflow para realizar el seguimiento del experimento.

El decorador @distributed ejecuta la función en el entorno de computación GPU sin servidor conectado y se encarga de la orquestación del entrenamiento distribuido.

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

Ejecuta el trabajo de entrenamiento distribuido

Ejecute la función de entrenamiento en 8 GPU H100. La función devuelve el identificador de ejecución de MLflow, que se puede usar para realizar un seguimiento de las métricas, ver los registros y acceder al modelo entrenado en la interfaz de usuario de MLflow.

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

Pasos siguientes

Cuaderno de ejemplo

Ajustar Llama 3.1 8B mediante Mosaic LLM Foundry en Databricks Serverless GPU

Obtención del cuaderno

  • Last updated on