Affiner Llama 3.1 8B à l'aide de Mosaic LLM Foundry sur le GPU serverless Databricks

Ouvrir dans Databricks

Affinez un modèle Llama 3.1 8B sur le Runtime AI à l'aide de Mosaic LLM Foundry, une base de code pour la formation, l'affinement, l'évaluation et le déploiement de grands modèles de langage avec prise en charge des stratégies de formation distribuée.

Le notebook utilise :

• Mosaic LLM Foundry : un framework pour l'entraînement et l'affinement des LLM avec prise en charge intégrée de FSDP, chargement efficace des données et intégration MLflow
• FSDP (Fully Sharded Data Parallel) : distribue les paramètres du modèle, les gradients et les états de l'optimiseur sur les GPU.
• **Databricks Serverless GPU** : Exécute une formation distribuée sur le compute GPU serverless connecté.
• Unity Catalog : Stocke les points de contrôle des modèles et enregistre les modèles entraînés
• MLflow : Suit les Experimentations et les Logs des métriques d'entraînement

Se connecter au compute GPU serverless

Ce Notebook nécessite un compute GPU serverless. Connecter :

1. Cliquez sur le sélecteur de compute du notebook en haut à droite et sélectionnez Serverless GPU .
2. Ouvrez le panneau latéral Environnement sur le côté droit du Notebook.
3. Définissez l' Accélérateur sur 8xH100
4. Sélectionnez l'environnement de base **Standard** et définissez la **Version de l'environnement** sur **5**, qui contient les bibliothèques nécessaires pour exécuter cet exemple
5. Sélectionnez Appliquer et cliquez sur Confirmer pour appliquer cet environnement à votre notebook

Installez les bibliothèques requises

Installez Mosaic LLM Foundry et ses dépendances pour l'entraînement distribué. La roue flash-attn précompilée est installée en premier afin que pip la réutilise au lieu de compiler flash-attention à partir de la source (ce qui est lent) lorsqu'elle résout llm-foundry[gpu]:

• flash-attn: implémentation d'attention optimisée, installée à partir d'une roue pré-construite
• llm-foundry: cadre essentiel pour l'entraînement et l'affinement des LLM
• hf_transfer: Téléchargements de modèles plus rapides depuis Hugging Face
• yamlmagic: active la configuration YAML dans les cellules de 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

Redémarrez l'environnement Python

Redémarrez le kernel Python pour vous assurer que les packages nouvellement installés sont disponibles.

Python

dbutils.library.restartPython()

Configurez les chemins Unity Catalog pour le stockage de modèles.

Configurez les emplacements Unity Catalog pour le stockage des points de contrôle de modèle et l'enregistrement du modèle entraîné. La configuration utilise des paramètres de query qui peuvent être personnalisés sans modifier le code.

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

Définir la configuration d'entraînement à l'aide de YAML.

Chargez la configuration d'affinement à partir du format YAML. La configuration spécifie :

• Architecture du modèle et poids pré-entraînés (Llama 3.1 8B)
• Paramètres FSDP pour la formation distribuée
• Hyperparamètres d'entraînement (taux d'apprentissage, taille de batch, optimiseur)
• Configuration du dataset (mosaicml/dolly_hhrlhf)
• Enregistrement MLflow et création de points de contrôle de modèle
• Callbacks pour le monitoring et l'optimisation

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

Définir la fonction de formation distribuée

Cette cellule définit la fonction d'entraînement qui s'exécutera sur 8 GPU H100 à l'aide du décorateur @distributed. La fonction :

• Configure le jeton Hugging Face pour l'accès au modèle
• Permet des téléchargements rapides de modèles avec hf_transfer
• Appelle la fonction train() de LLM Foundry avec la configuration YAML
• Renvoie l'ID d'exécution MLflow pour le suivi de l'expérimentation.

Le décorateur @distributed exécute la fonction sur le compute GPU serverless connecté et gère l’orchestration de l’entraînement distribué.

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

Exécuter le job de formation distribuée

Exécuter la fonction d'entraînement sur 8 GPU H100. La fonction renvoie l'ID d'exécution MLflow, qui peut être utilisé pour suivre les métriques, afficher les logs et accéder au modèle entraîné dans l'interface utilisateur MLflow.

Python

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

Étapes suivantes

• Formation distribuée multi-GPU et multi-nœuds
• Meilleures pratiques pour le compute GPU Serverless
• Dépanner les problèmes sur le compute GPU serverless
• Documentation Mosaic LLM Foundry
• Registre de modèles Unity Catalog

Exemple de Notebook

Affinez Llama 3.1 8B à l'aide de Mosaic LLM Foundry sur Databricks Serverless GPU

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks