Ottimizzare Llama 3.1 8B utilizzando Mosaic LLM Foundry su GPU Serverless di Databricks

Apri la versione notebook di questa pagina

Ottimizzare un modello Llama 3.1 8B in AI Runtime usando Mosaic LLM Foundry, una codebase per il training, l'ottimizzazione, la valutazione e la distribuzione di modelli linguistici di grandi dimensioni con supporto per le strategie di training distribuite.

Il notebook usa:

Mosaic LLM Foundry: framework per il training e l'ottimizzazione di LLM con supporto predefinito per FSDP, caricamento efficiente dei dati e integrazione MLflow
FSDP (Fully Sharded Data Parallel): distribuisce i parametri del modello, i gradienti e gli stati dell'ottimizzatore tra GPU
Databricks Serverless GPU: consente l'addestramento distribuito su risorse di calcolo serverless GPU connesse
Unity Catalog: archivia i checkpoint dei modelli e registra i modelli addestrati
MLflow: tiene traccia degli esperimenti e dei log delle metriche di training

Connessione alla computazione GPU senza server

Questo notebook richiede l'elaborazione GPU in modalità serverless. Per connettersi:

Fare clic sul selettore di calcolo del notebook in alto a destra e selezionare GPU serverless
Aprire il pannello laterale Ambiente sul lato destro del notebook
Imposta Acceleratore su 8xH100
Selezionare l'ambiente di base Standard e impostare Ambiente versionesu 5, che contiene le librerie necessarie per eseguire questo esempio
Selezionare Applica e fare clic su Conferma per applicare questo ambiente al notebook

Installare le librerie necessarie

Installare Mosaic LLM Foundry e le relative dipendenze per il training distribuito. Il wheel precompilato flash-attn viene installato per primo in modo che pip lo riutilizzi invece di compilare flash-attention dal sorgente (operazione lenta) quando risolve llm-foundry[gpu]:

flash-attn: implementazione dell'attenzione ottimizzata, installata da una rotellina predefinita
llm-foundry: framework di base per il training e l'ottimizzazione LLM
hf_transfer: download più veloci dei modelli da Hugging Face
yamlmagic: abilita la configurazione YAML nelle celle 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

Riavviare l'ambiente Python

Riavviare il kernel Python per assicurarsi che i pacchetti appena installati siano disponibili.

dbutils.library.restartPython()

Configurare i percorsi del catalogo Unity per l'archiviazione dei modelli

Configurare i percorsi di Unity Catalog per archiviare i checkpoint del modello e registrare il modello addestrato. La configurazione usa parametri di query che possono essere personalizzati senza modificare il codice.

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

Definire la configurazione di training con YAML

Caricare la configurazione di ottimizzazione dal formato YAML. La configurazione specifica:

Architettura del modello e pesi pre-addestrati (Llama 3.1 8B)
Impostazioni FSDP per il training distribuito
Iperparametri di training (tasso di apprendimento, dimensione del batch, ottimizzatore)
Configurazione del set di dati (mosaicml/dolly_hhrlhf)
Registrazione di MLflow e salvataggio di checkpoint del modello
Callback per il monitoraggio e l'ottimizzazione

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

Definire la funzione di training distribuita

Questa cella definisce la funzione di training che verrà eseguita su 8 GPU H100 usando il @distributed decorator. La funzione :

Configura il token Hugging Face per l'accesso al modello
Abilita download rapidi dei modelli con hf_transfer
Chiama la funzione LLM Foundry train() con la configurazione YAML
Restituisce l'ID di esecuzione MLflow per tenere traccia dell'esperimento

Il decorator @distributed esegue la funzione sulla risorsa di calcolo GPU serverless connessa e gestisce l'orchestrazione dell'addestramento distribuito.

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

Eseguire il processo di training distribuito

Eseguire la funzione di training su 8 GPU H100. La funzione restituisce l'ID di esecuzione MLflow, che può essere usato per tenere traccia delle metriche, visualizzare i log e accedere al modello sottoposto a training nell'interfaccia utente di MLflow.

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

Passaggi successivi

Notebook di esempio

Ottimizzare Llama 3.1 8B utilizzando Mosaic LLM Foundry su GPU Serverless di Databricks

Ottieni il notebook

Commenti e suggerimenti

Questa pagina è stata utile?

Last updated on 2026-06-30

Ottimizzare Llama 3.1 8B utilizzando Mosaic LLM Foundry su GPU Serverless di Databricks

Connessione alla computazione GPU senza server

Installare le librerie necessarie

Riavviare l'ambiente Python

Configurare i percorsi del catalogo Unity per l'archiviazione dei modelli

Definire la configurazione di training con YAML

Definire la funzione di training distribuita

Eseguire il processo di training distribuito

Passaggi successivi

Notebook di esempio

Ottimizzare Llama 3.1 8B utilizzando Mosaic LLM Foundry su GPU Serverless di Databricks

Commenti e suggerimenti

Risorse aggiuntive