Affiner Llama 3.2 1B à l'aide du Runtime IA

Ouvrir dans Databricks

Affinez entièrement un grand modèle de langage avec l'affinement supervisé (SFT) sur le 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 seul nœud avec 8 GPU H100.

Concepts clés :

• TRL (Transformers Reinforcement Learning) : une bibliothèque qui fournit des outils pour entraîner des modèles linguistiques avec l'apprentissage par renforcement et l'affinement supervisé.
• DeepSpeed ZeRO Stage 3 : une technique d’optimisation de la mémoire qui partitionne les parameters, les gradients et les états de l’optimiseur du modèle sur plusieurs GPU pour permettre l’entraînement de grands modèles.
• Runtime IA : compute GPU géré par Databricks qui provisionne et met à l’échelle automatiquement les ressources GPU pour les charges de travail d’entraînement.

Pour plus d'informations, consultez AI Runtime.

Exigences

Ce Notebook nécessite ce qui suit :

• AI Runtime : le Notebook utilise Databricks AI Runtime avec 8 GPU H100 pour l'entraînement distribué. Aucune configuration de cluster n'est nécessaire.
• Unity Catalog : un catalogue et un schéma Unity Catalog pour stocker les points de contrôle du modèle et enregistrer le modèle entraîné.
• Jeton HuggingFace : un jeton d’accès HuggingFace stocké dans les secrets Databricks pour download le modèle de base et le dataset.
• Packages Python : l'AI Runtime préinstalle la plupart des bibliothèques requises. La section de configuration ci-dessous installe deepspeed.

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. Sur le côté droit, cliquez sur le bouton d’environnement.
3. Sélectionnez 8xH100 comme Accélérateur .
4. Choisissez l'environnement **AI v5** dans le panneau de droite 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é.

Installez les packages requis

AI Runtime dispose déjà de la plupart des bibliothèques requises préinstallées. Pour cet exemple, vous n’avez besoin d’installer que deepspeed.

Python

%pip install deepspeed==0.19.1
%restart_python

Configurez Unity Catalog et les variables d'environnement

Configurez les emplacements Unity Catalog pour le stockage des points de contrôle de modèle et l'enregistrement du modèle entraîné. Le Notebook utilise les paramètres de query pour configurer :

• Catalogue et schéma : espace de noms Unity Catalog pour l'organisation des modèles et des points de contrôle
• Nom du modèle : Nom du modèle enregistré dans Unity Catalog
• **Volume** : volume Unity Catalog pour le stockage des points de contrôle du modèle pendant l'entraînement

La configuration récupère également le jeton HuggingFace des secrets Databricks et configure l'expérimentation MLflow pour le suivi des métriques d'entraînement.

Python

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

Python

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 la configuration DeepSpeed ZeRO Stage 3

DeepSpeed ZeRO (Zero Redundancy Optimizer) Stage 3 partitionne les parameters, les dégradés et les états de l'optimiseur du modèle sur tous les GPU afin de réduire la consommation de mémoire par GPU. Cela permet l'entraînement de grands modèles qui ne tiendraient pas dans la mémoire d'un seul GPU.

Paramètres clés de configuration :

• bf16 activé : utilise la précision bfloat16 pour un entraînement plus rapide et une utilisation de la mémoire réduite.
• Optimisation de l'étape 3 : Partitionne tous les états du modèle sur l'ensemble des GPU.
• Pas de déchargement CPU : Conserve toutes les données sur les GPU pour une performance maximale sur le matériel H100
• Communication par chevauchement : superpose la communication de gradient avec le calcul pour l'efficacité

Python

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 les paramètres d'entraînement

Configurez les paramètres d'affinement supervisé :

• **Modèle** : Llama 3.2 1B Instruct, un modèle compact adapté aux GPU H100
• Dataset : Dataset Capybara de la bibliothèque TRL pour la formation en IA conversationnelle
• **Taille du batch** : 2 par appareil avec 4 étapes d’accumulation du gradient pour une taille de batch effective de 64.
• Taux d'apprentissage : 2e-4 avec un planificateur cosinus et un préchauffage
• Étapes d'entraînement : 60 étapes pour la démonstration (à augmenter pour un entraînement complet)

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

Python

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 de formation distribuée

Le décorateur @distributed de la bibliothèque serverless_gpu permet l’exécution de charges de travail GPU sur Databricks AI Runtime. Le décorateur prend en charge le provisionnement de 8 GPU H100 et gère automatiquement la configuration de la formation distribuée.

Key parameters :

• gpus=8 : Requiert 8 GPU pour l'entraînement distribué
• gpu_type='H100' : Spécifie le matériel GPU H100

La fonction d'entraînement :

1. Charge le modèle de base et le tokenizer depuis HuggingFace.
2. Configure le formatage du chat pour l'IA conversationnelle.
3. Charge le dataset 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.

Python

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 le job de formation distribuée

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

Le processus d'entraînement :

• Provisionne automatiquement 8 GPU H100.
• Télécharge download le modèle et le dataset depuis HuggingFace
• Entraîne le modèle avec un affinement complet.
• Enregistre les points de contrôle dans le volume Unity Catalog
• Logs metrics dans MLflow
• Renvoie l'état de l'entraînement, la perte finale et l'ID d'exécution MLflow

Python

# 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']}")

Enregistrez le modèle affiné et testez l'inférence

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

Le processus :

1. Charge le modèle enregistré
2. Teste le modèle avec un exemple de prompt conversationnel

Python

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)

Enregistrer le modèle dans Unity Catalog

Enregistrez le modèle affiné dans MLflow et enregistrez-le dans Unity Catalog pour le déploiement et le service. Le modèle est enregistré avec :

• **Modèle et tokenizer** : Les deux composants nécessaires à l'inférence
• **Type de tâche** : Configuré en tant que llm/v1/chat pour l’IA conversationnelle
• Exemple d'entrée : Format d'exemple de message de chat pour les tests
• Unity Catalog registration : Enregistre 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 Endpoint de service de modèles ou utilisé pour l'inférence par batch.

Python

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

• AI Runtime
• Bonnes pratiques pour l’AI Runtime
• Dépannage des problèmes sur l'AI Runtime
• Formation distribuée multi-GPU et multi-nœuds
• Entraîner des modèles avec MLflow
• Déployer des modèles avec Model Serving

Exemple de Notebook

Affiner Llama 3,2 1B à l'aide de l'AI Runtime

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks