Ajuster Llama 3.1 8B en utilisant Mosaic LLM Foundry sur le GPU Serverless de Databricks

Ouvrir la version du bloc-notes de cette page

Ajustez un modèle Llama 3.1 8B sur AI Runtime à l’aide de Mosaïque LLM Foundry, une base de code pour l’apprentissage, l’optimisation, l’évaluation et le déploiement de modèles de langage volumineux avec prise en charge des stratégies d’entraînement distribuées.

Le notebook utilise :

Mosaic LLM Foundry : un cadre pour l’apprentissage et l’ajustement des modèles de langage (LLM) avec prise en charge intégrée de FSDP, du chargement efficace des données et de l’intégration avec MLflow.
FSDP (Fully Sharded Data Parallel) : distribue les paramètres de modèle, les gradients et les états de l’optimiseur entre les GPU
Databricks Serverless GPU: exécute un entraînement distribué sur des ressources de calcul GPU serverless connectées
Catalogue Unity : stocke les points de contrôle de modèle et inscrit des modèles entraînés
MLflow : effectue le suivi des expériences et consigne les métriques de formation

Connecter au calcul GPU sans serveur

Ce bloc-notes nécessite une puissance de calcul GPU sans serveur. Pour vous connecter :

Cliquez sur le sélecteur de calcul du notebook en haut à droite et sélectionnez GPU serverless
Ouvrez le panneau latéral Environment à droite du notebook
Définir l’accélérateursur 8xH100
Sélectionnez l’environnement de base standard et définissez la version d’environnementsur 5, qui contient les bibliothèques nécessaires pour exécuter cet exemple
Sélectionnez Appliquer , puis cliquez sur Confirmer pour appliquer cet environnement à votre bloc-notes

Installer les bibliothèques requises

Installez Mosaïque LLM Foundry et ses dépendances pour l’entraînement distribué. Le wheel précompilé flash-attn est installé en premier afin que pip le réutilise au lieu de compiler FlashAttention à partir du code source (ce qui est lent) lors de la résolution de llm-foundry[gpu] :

flash-attn : Implémentation optimisée du mécanisme d’attention, installée à partir d’un wheel précompilé
llm-foundry: Cadre principal pour l’apprentissage et l’affinage de LLM
hf_transfer: Téléchargements de modèles plus rapides à partir de Hugging Face
yamlmagic: active la configuration YAML dans les cellules du bloc-notes

%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émarrer l’environnement Python

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

dbutils.library.restartPython()

Configurer les chemins du catalogue Unity pour le stockage de modèles

Configurez les emplacements de Unity Catalog pour stocker des points de contrôle de modèle et inscrire le modèle entraîné. La configuration utilise des paramètres de requête qui peuvent être personnalisés sans modifier le code.

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 une configuration d’entraînement à l’aide de YAML

Chargez la configuration de réglage précis à partir du format YAML. La configuration spécifie :

Architecture de modèle et poids préentraînés (Llama 3.1 8B)
Paramètres FSDP pour l’entraînement distribué
Hyperparamètres d’apprentissage (taux d’apprentissage, taille de lot, optimiseur)
Configuration de dataset (mosaicml/dolly_hhrlhf)
Journalisation MLflow et point de contrôle de modèle
Fonctions de rappel pour la surveillance et l’optimisation

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

Définir la fonction d’entraînement distribuée

Cette cellule définit la fonction d’entraînement qui fonctionnera sur 8 GPU H100 en utilisant le décorateur @distributed. La fonction :

** Configure le jeton Hugging Face pour l’accès au modèle
Active les téléchargements de modèles rapides avec hf_transfer
Appelle la fonction LLM Foundry train() avec la configuration YAML
Retourne l’ID d’exécution MLflow pour le suivi de l’expérience

Le décorateur @distributed exécute la fonction sur la ressource de calcul GPU sans serveur connectée et orchestre l’entraînement distribué.

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 la tâche d'entraînement distribuée

Exécutez la fonction d’entraînement sur 8 GPU H100. La fonction retourne l’ID d’exécution de MLflow, qui peut être utilisé pour suivre les paramètres, afficher les logs, et accéder au modèle entraîné dans l’interface utilisateur de MLflow.

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

Étapes suivantes

Exemple de notebook

Ajuster Llama 3.1 8B en utilisant Mosaic LLM Foundry sur le GPU Serverless de Databricks

Obtenir un ordinateur portable

Commentaires

Cette page a-t-elle été utile ?

Last updated on 2026-08-18

Ajuster Llama 3.1 8B en utilisant Mosaic LLM Foundry sur le GPU Serverless de Databricks

Connecter au calcul GPU sans serveur

Installer les bibliothèques requises

Redémarrer l’environnement Python

Configurer les chemins du catalogue Unity pour le stockage de modèles

Définir une configuration d’entraînement à l’aide de YAML

Définir la fonction d’entraînement distribuée

Exécuter la tâche d'entraînement distribuée

Étapes suivantes

Exemple de notebook

Ajuster Llama 3.1 8B en utilisant Mosaic LLM Foundry sur le GPU Serverless de Databricks

Commentaires

Ressources supplémentaires