Ajuster Llama 3.2 1B à l’aide d’AI Runtime

Ouvrir la version du bloc-notes de cette page

Ajustez entièrement un modèle de langage volumineux avec un réglage précis supervisé (SFT) sur Databricks AI Runtime. Le notebook utilise la bibliothèque Transformers Reinforcement Learning (TRL) avec l’optimisation DeepSpeed ZeRO Stage 3 pour entraîner efficacement Llama 3.2 1B sur un nœud unique avec 8 GPU H100.

Concepts clés :

  • TRL (Transformers Reinforcement Learning) : bibliothèque qui fournit des outils pour l’apprentissage des modèles linguistiques avec un apprentissage par renforcement et un réglage précis supervisé.
  • DeepSpeed ZeRO Stage 3 : technique d’optimisation de la mémoire qui partitionne les paramètres du modèle, les dégradés et les états d’optimiseur entre les GPU pour permettre l’entraînement de grands modèles.
  • IA Runtime : calcul GPU géré par Databricks qui provisionne et ajuste automatiquement les ressources GPU pour les charges de travail de formation.

Pour plus d’informations, consultez AI Runtime.

Note

Cet exemple nécessite l’environnement IA Databricks version 5 ou supérieure.

Exigences

Ce notebook nécessite les éléments suivants :

  • AI Runtime : Le notebook utilise Databricks AI Runtime avec 8 GPU H100 pour l'entraînement par répartition. Aucune configuration de cluster n’est nécessaire.
  • Catalogue Unity : catalogue et schéma Unity pour stocker des points de contrôle de modèle et inscrire le modèle entraîné.
  • Jeton HuggingFace : jeton d’accès HuggingFace stocké dans les secrets Databricks pour télécharger le modèle de base et le jeu de données.
  • Python packages : AI Runtime préinstalle la plupart des bibliothèques requises. La section d’installation ci-dessous installe deepspeed.

Connecter au calcul GPU sans serveur

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

  1. Cliquez sur le sélecteur de calcul du notebook en haut à droite et sélectionnez GPU serverless.
  2. Sur le côté droit, cliquez sur le bouton Environnement.
  3. Sélectionnez 8xH100 comme accélérateur.
  4. Choisissez l’environnement AI v5 dans le panneau droit qui contient toutes les bibliothèques requises pour exécuter cet exemple de notebook.
  5. Cliquez sur Appliquer.

La fonction d’entraînement provisionne automatiquement 8 GPU H100 pour l’entraînement distribué.

Installer les packages requis

AI Runtime dispose déjà de la plupart des bibliothèques requises préinstallées. Pour cet exemple, vous devez installer uniquement deepspeed.

%pip install deepspeed==0.19.1
%restart_python

Configurer des variables de catalogue et d’environnement Unity

Configurez les emplacements du catalogue Unity pour stocker des points de contrôle de modèle et inscrire le modèle entraîné. Le notebook utilise des paramètres de requête pour configurer :

  • Catalogue et schéma : espace de noms Catalogue Unity pour organiser des modèles et des points de contrôle
  • Nom du modèle : nom du modèle inscrit dans le catalogue Unity
  • Volume : Volume catalogue Unity pour le stockage des points de contrôle de modèle pendant l’entraînement

La configuration récupère également le jeton HuggingFace à partir des secrets Databricks et configure l’expérience MLflow pour le suivi des métriques d’apprentissage.

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "llama3_2-1b")
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")

# Get HuggingFace token and username
hf_token = dbutils.secrets.get(scope="sgc-nightly-notebook", key="hf_token")
username = spark.sql("SELECT session_user()").collect()[0][0]

REGISTERED_MODEL_NAME = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"
CHECKPOINT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"
MLFLOW_EXPERIMENT_NAME = f"/Users/{username}/{UC_MODEL_NAME}"

# Create the Unity Catalog volume if it doesn't exist
spark.sql(f"CREATE VOLUME IF NOT EXISTS {UC_CATALOG}.{UC_SCHEMA}.{UC_VOLUME}")

print(f"👤 Username: {username}")
print("🔑 HuggingFace token configured")
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"CHECKPOINT_DIR: {CHECKPOINT_DIR}")
print(f"MLFLOW_EXPERIMENT_NAME: {MLFLOW_EXPERIMENT_NAME}")


import os
import json
import tempfile
import torch
import mlflow
from huggingface_hub import constants
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer

Créer une configuration DeepSpeed ZeRO Stage 3

DeepSpeed ZeRO (Optimiseur de redondance zéro) Étape 3 partitionne les paramètres du modèle, les dégradés et les états d’optimiseur sur tous les GPU pour réduire la consommation de mémoire par GPU. Cela permet d’entraîner des modèles volumineux qui ne tiennent pas dans la mémoire d’un seul GPU.

Paramètres de configuration de clé :

  • bf16 activé : utilise la précision bfloat16 pour une formation plus rapide et une utilisation réduite de la mémoire
  • Optimisation de l’étape 3 : partitionne tous les états de modèle entre les GPU
  • Aucun déchargement du processeur : conserve toutes les données sur les GPU pour des performances maximales sur le matériel H100
  • Communication de chevauchement : chevauche la communication de gradient avec le calcul pour optimiser l’efficacité
def create_deepspeed_config(stage: int):
    """Create a DeepSpeed ZeRO configuration for single-node 8xH100 training."""

    deepspeed_config = {
        "fp16": {
            "enabled": False
        },
        "bf16": {
            "enabled": True
        },
        "zero_optimization": {
            "stage": stage,
            "offload_optimizer": {
                "device": "none"
            },
            "offload_param": {
                "device": "none"
            },
            "overlap_comm": True,
            "contiguous_gradients": True,
            "sub_group_size": 1e9,
            "reduce_bucket_size": "auto",
            "stage3_prefetch_bucket_size": "auto",
            "stage3_param_persistence_threshold": 0,
            "stage3_max_live_parameters": 1e9,
            "stage3_max_reuse_distance": 1e9,
            "stage3_gather_16bit_weights_on_model_save": True
        },
        "gradient_accumulation_steps": 1,
        "gradient_clipping": "auto",
        "steps_per_print": 2000,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": "auto",
        "wall_clock_breakdown": False
    }

    return deepspeed_config


# Create DeepSpeed configuration
zero_stage = 3
deepspeed_config = create_deepspeed_config(zero_stage)
print(f"⚙️  DeepSpeed ZeRO Stage {zero_stage} configuration created")

Définir des paramètres d’entraînement

Configurez les paramètres de réglage précis supervisés :

  • Modèle : Llama 3.2 1B Instruct, modèle compact adapté aux GPU H100
  • Jeu de données : Jeu de données Capybara à partir de la bibliothèque TRL pour l’apprentissage de l’IA conversationnelle
  • Taille du lot : 2 par appareil avec 4 étapes d’accumulation des gradients pour une taille de lot effective de 64
  • Taux d’apprentissage : 2e-4 avec programmateur cosinus et échauffement
  • Étapes de formation : 60 étapes de démonstration (augmentation de la formation complète)

La configuration utilise la précision bfloat16 et le point de contrôle des gradients pour optimiser l’utilisation de la mémoire.

def create_training_config():
    """Create training configuration for TRL SFT."""

    # Model and dataset configuration (not part of TrainingArguments)
    model_config = {
        "model_name": "meta-llama/Llama-3.2-1B-Instruct",  # Small Llama model suitable for 8xH100
        "dataset_name": "trl-lib/Capybara"
    }

    # Training arguments that will be passed directly to TrainingArguments
    training_args_config = {
        "output_dir": CHECKPOINT_DIR,
        "per_device_train_batch_size": 2,
        "per_device_eval_batch_size": 2,
        "gradient_accumulation_steps": 1,
        "learning_rate": 2e-4,
        "max_steps": 60,   # TO DO remove when fine-tuning on full dataset. Demo purposes only.
        "logging_steps": 10,
        "save_steps": 30,
        "eval_steps": 30,
        "eval_strategy": "steps",
        "warmup_steps": 10,
        "lr_scheduler_type": "cosine",
        "gradient_checkpointing": False,
        "fp16": False,
        "bf16": True,
        "optim": "adamw_torch",
        "remove_unused_columns": False,
        "run_name": f"llama3.2-1b-fft-zero3",
        "report_to": "mlflow",
        "save_total_limit": 2,
        "load_best_model_at_end": True,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
    }

    return model_config, training_args_config

# Create training configuration
model_config, training_args_config = create_training_config()

print("📊 Training Configuration:")
print(f"  🤖 Model: {model_config['model_name']}")
print(f"  📚 Dataset: {model_config['dataset_name']}")
print(f"  🎯 Batch size: {training_args_config['per_device_train_batch_size']}")
print(f"  📈 Learning rate: {training_args_config['learning_rate']}")

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

Le décorateur @distributed de la bibliothèque serverless_gpu permet d’exécuter des charges de travail GPU sur Databricks AI Runtime. Le décorateur provisionne automatiquement 8 GPU H100 et gère automatiquement la configuration de l’entraînement distribué.

Paramètres clés :

  • gpus=8 : Demandes de 8 GPU pour l’entraînement distribué
  • gpu_type='H100' : spécifie le matériel GPU H100

Fonction d’entraînement :

  1. Charge le modèle de base et le tokenizer à partir de HuggingFace
  2. Configurer la mise en forme de chat pour l'intelligence artificielle conversationnelle
  3. Charge le jeu de données d’entraînement
  4. Initialise le TRL SFTTrainer avec l’optimisation DeepSpeed
  5. Entraîne le modèle et enregistre les points de contrôle
  6. Retourne les résultats d’entraînement et l’ID d’exécution MLflow

Pour plus d’informations, consultez la documentation de l’API AI Runtime.

from serverless_gpu import distributed

mlflow.set_experiment(MLFLOW_EXPERIMENT_NAME)
@distributed(
    gpus=8,
    gpu_type='H100',
)
def run_distributed_trl_sft():
    """
    Distributed TRL SFT training function using AI Runtime.

    This function will be executed on the H100 GPU with DeepSpeed optimization.
    """

    # Set up environment variables for remote jobs
    import os
    import tempfile
    import json
    from huggingface_hub import constants
    from datasets import load_dataset
    from transformers import AutoTokenizer
    from trl import SFTTrainer, SFTConfig

    # HuggingFace configuration
    os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token
    os.environ['HF_TOKEN'] = hf_token
    constants.HF_HUB_ENABLE_HF_TRANSFER = True

    # Set up temporary directories
    temp_dir = tempfile.mkdtemp()

    print("🚀 Starting TRL SFT training on H100 GPU...")

    try:
        # Load tokenizer
        print(f"📥 Loading tokenizer from model: {model_config['model_name']}")
        tokenizer = AutoTokenizer.from_pretrained(model_config['model_name'])

        # Add pad token if not present
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        # Load dataset
        print(f"📚 Loading dataset: {model_config['dataset_name']}")
        dataset = load_dataset(model_config['dataset_name'])

        # Create temporary DeepSpeed config file
        deepspeed_config_path = os.path.join(temp_dir, "deepspeed_config.json")
        with open(deepspeed_config_path, "w") as f:
            json.dump(deepspeed_config, f, indent=2)

        # Training arguments - dynamically pass all config parameters
        training_args = SFTConfig(
            **training_args_config,
            deepspeed=deepspeed_config_path,  # Override deepspeed with the config file path
        )

        # Initialize SFT Trainer
        print("🏋️ Initializing SFT Trainer with DeepSpeed...")
        trainer = SFTTrainer(
            model=model_config["model_name"],
            args=training_args,
            train_dataset=dataset["train"],
            eval_dataset=dataset["test"] if "test" in dataset else None,
            processing_class=tokenizer,
        )

        # Start training
        print("🎯 Starting training...")
        trainer.train()

        # Save the model
        print("💾 Saving trained model...")
        trainer.save_model()

        # Get training results
        train_results = trainer.state.log_history
        final_loss = train_results[-1].get('train_loss', 'N/A') if train_results else 'N/A'

        print("✅ Training completed successfully!")
        print(f"📊 Final training loss: {final_loss}")

        mlflow_run_id = None
        if mlflow.last_active_run() is not None:
            mlflow_run_id = mlflow.last_active_run().info.run_id

        return {
            "status": "success",
            "final_loss": final_loss,
            "output_dir": training_args_config['output_dir'],
            "model_name": model_config['model_name'],
            "mlflow_run_id": mlflow_run_id,
        }

    except Exception as e:
        print(f"❌ Training failed: {e}")
        import traceback
        traceback.print_exc()
        return {
            "status": "failed",
            "error": str(e)
        }

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

Exécutez la fonction d’entraînement en appelant .distributed() sur la fonction décorée. Cela provisionne les ressources AI Runtime, exécute l’entraînement sur 8 GPU H100 avec l’optimisation DeepSpeed et retourne les résultats.

Processus de formation :

  • Approvisionne automatiquement des GPU H100, 8
  • Télécharge le modèle et le jeu de données à partir de HuggingFace
  • Effectue l’apprentissage du modèle avec un réglage complet
  • Enregistre les points de contrôle dans le volume Unity Catalog
  • Enregistre les métriques sur MLflow
  • Retourne l’état d’entraînement, la perte finale et l’ID d’exécution MLflow
# Execute the distributed training
results = run_distributed_trl_sft.distributed()

print("🏁 Training execution completed!")
print(f"📊 Results: {results}")

if results and results[0].get('status') == 'success':
    print("✅ Training completed successfully!")
    print(f"💾 Model saved to: {results[0].get('output_dir', 'N/A')}")
    print(f"📈 Final loss: {results[0].get('final_loss', 'N/A')}")
    print(f"🎉 MLflow run ID: {results[0].get('mlflow_run_id', 'N/A')}")
else:
    print("❌ Training failed!")
    if results and 'error' in results:
        print(f"🔍 Error: {results['error']}")

Enregistrer le modèle affiné et l’inférence de test

Cette étape facultative charge le modèle affiné et le teste avec un exemple de prompt afin de vérifier les résultats.

Le processus :

  1. Charge le modèle enregistré
  2. Teste le modèle avec un exemple d’invite conversationnelle
def save_and_load_trained_model():
    """Load the fully fine-tuned model from the Unity Catalog volume."""

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"📥 Loading fine-tuned model from: {training_args_config['output_dir']}")

    # Load the fully fine-tuned model directly from the checkpoint directory
    model = AutoModelForCausalLM.from_pretrained(
        training_args_config['output_dir'],
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map={"":0}
    )
    tokenizer = AutoTokenizer.from_pretrained(training_args_config['output_dir'], trust_remote_code=True)

    print("✅ Model loaded successfully!")
    return model, tokenizer

def test_trained_model(model, tokenizer):
    """Test the trained model with simple inference."""

    try:
        import torch
        # Test prompt
        # Create a conversation following the schema
        conversation = [
            {
                "content": "What is machine learning?",
                "role": "user"
            }
        ]

        # Convert conversation to chat format
        prompt = ""
        for message in conversation:
            if message["role"] == "user":
                prompt += f"### User: {message['content']}\n### Response:"
            else:
                prompt += f" {message['content']}\n\n"

        # Tokenize
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

        # Generate
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=500,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )

        # Decode
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        print("🤖 Model Response:")
        print(response)
        return response

    except Exception as e:
        print(f"❌ Model testing failed: {e}")

# Save and load the trained model
model, tokenizer = save_and_load_trained_model()

# Test the trained model
test_trained_model(model, tokenizer)

Inscrire le modèle dans le catalogue Unity

Consignez le modèle affiné sur MLflow et inscrivez-le dans le catalogue Unity pour le déploiement et le service. Le modèle est enregistré avec :

  • Modèle et générateur de jetons : les deux composants nécessaires à l’inférence
  • Type de tâche : configuré comme llm/v1/chat pour l’IA conversationnelle
  • Exemple d’entrée : exemple de format de message de conversation pour le test
  • Inscription du catalogue Unity : inscrit automatiquement le modèle dans le catalogue et le schéma configurés

Une fois enregistré, le modèle peut être déployé sur des points de terminaison de service de modèle ou utilisé pour l’inférence par lots.

run_id = results[0].get('mlflow_run_id')
mlflow.set_registry_uri("databricks-uc")

# log the model to mlflow using the latest run id and register to Unity Catalog
with mlflow.start_run(run_id=run_id) as run:
    components = {
        "model": model,
        "tokenizer": tokenizer
    }
    logged_model = mlflow.transformers.log_model(
        transformers_model=components,
        name="model",
        task="llm/v1/chat",
        input_example={
            "messages": [
                {"role": "user", "content": "What is machine learning?"}
            ]
        },
        registered_model_name=REGISTERED_MODEL_NAME
        )
    print(f"🔍 Model logged to: {logged_model}")

Étapes suivantes

Exemple de notebook

Ajuster Llama 3.2 1B à l’aide d’AI Runtime

Obtenir un ordinateur portable

  • Last updated on