Affinement complet de Qwen3-4B

Affiner entièrement le grand modèle de langage Qwen3-4B sur un seul GPU H100. Ce tutoriel montre comment :

Exécutez **un affinement complet**, qui met à jour chaque paramètre du modèle pour une adaptation maximale à vos données
Utilisez l'**environnement Databricks AI v5** sans installer de bibliothèques supplémentaires.
Tirer parti de **TRL (Transformer Reinforcement Learning)** pour l'affinement supervisé
Enregistrez le modèle affiné dans Unity Catalog pour la gouvernance et le déploiement

Concepts clés :

Affinement complet : met à jour tous les poids du modèle, lui conférant la plus grande capacité à apprendre de votre dataset au prix d’une mémoire et d’un compute plus élevés que les méthodes à parameter efficace
TRL: Une bibliothèque pour l'entraînement de modèles de langage avec l'apprentissage par renforcement et l'affinement supervisé
Entraînement économe en mémoire : utilise la précision mixte BF16 et le pointage de gradient pour un affinage complet de 4 milliards de parameters sur un seul GPU H100

Matrice de décision Affinement complet vs LoRA

Ce Notebook utilise l'**affinement complet**, qui met à jour tous les parameters du modèle. L'alternative, LoRA (Low-Rank Adaptation), fige le modèle de base et n'entraîne que de petites couches d'adaptateur.

Scénario	Recommandation	Raison
Changement majeur de comportement du modèle	Affinement complet	Met à jour tous les paramètres pour les modifications fondamentales du comportement du modèle
Qualité la plus élevée possible pour une seule tâche	Affinement complet	Pas d'approximation de rang faible, donc le modèle a une pleine capacité d'adaptation
Mémoire GPU limitée	LoRA	Adapte les modèles plus grands en mémoire en n’entraînant qu’environ 1 % des paramètres.
Plusieurs adaptateurs spécifiques aux tâches	LoRA	Échangez différents adaptateurs sur le même modèle de base.

Scénario	Recommandation	Raison
Changement majeur de comportement du modèle	Affinement complet	Met à jour tous les paramètres pour les modifications fondamentales du comportement du modèle
Qualité la plus élevée possible pour une seule tâche	Affinement complet	Pas d'approximation de rang faible, donc le modèle a une pleine capacité d'adaptation
Mémoire GPU limitée	LoRA	Adapte les modèles plus grands en mémoire en n’entraînant qu’environ 1 % des paramètres.
Plusieurs adaptateurs spécifiques aux tâches	LoRA	Échangez différents adaptateurs sur le même modèle de base.

L'affinement complet d'un modèle à 4B parameter nécessite beaucoup plus de mémoire GPU que LoRA, car l'état de l'optimiseur et les gradients sont maintenus pour chaque parameter. Ce notebook utilise la précision mixte BF16 et le gradient checkpointing afin que l'entraînement puisse s'exécuter sur un seul GPU H100 (80 Go).

Se connecter au compute GPU serverless

Pour vous connecter au compute GPU Serverless :

Cliquez sur le menu déroulant Connecter dans le Notebook et sélectionnez Serverless GPU .
Choisissez un GPU **1x H100** comme accélérateur.
Ouvrez le panneau Environnement et choisissez IA v5 comme environnement de base.
Cliquez sur Appliquer .

Pour plus d'informations, consultez la documentation sur le compute GPU.

Importer des bibliothèques

L'environnement Databricks AI v5 inclut déjà toutes les bibliothèques requises pour cet exemple (telles que trl, transformers, datasets et mlflow), donc aucune installation supplémentaire n'est nécessaire.

La cellule suivante importe les bibliothèques requises pour l'entraînement du modèle, la gestion du dataset et le suivi MLflow.

Python
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import (
    SFTConfig,
    SFTTrainer,
    setup_chat_format
)
import torch
import mlflow

Configuration de l'installation

Intégration de Unity Catalog

La cellule suivante configure l'emplacement où votre modèle ajusté sera stocké et enregistré :

**Catalogue et Schéma** : Organisez les modèles dans votre espace de noms Unity Catalog (par défautmain.default:)
Nom du modèle : nom du modèle enregistré dans Unity Catalog pour la gouvernance et le déploiement.
**Volume** : volume Unity Catalog pour le stockage des points de contrôle du modèle pendant l'entraînement

Ces widgets vous permettent de personnaliser l'emplacement de stockage sans modifier le code. Le modèle sera enregistré en tant que {catalog}.{schema}.{model_name} pour un accès facile et le contrôle de version.

Hyperparamètres d'entraînement

La cellule définit également les principaux paramètres de formation :

**Modèle & Dataset** : Qwen3-4B avec le dataset conversationnel Capybara
Taille de batch (1) : nombre d'exemples par GPU par étape d'entraînement, maintenu petit pour permettre un affinement complet en mémoire.
**Accumulation de gradient (8)** : accumule les gradients sur 8 batchs pour une taille de batch effective de 8
Taux d'apprentissage (2e-5) : taux conservateur adapté à l'affinement complet
Étapes max. (50) : Limite la formation à 50 étapes pour une exécution de démonstration rapide
Journalisation et points de contrôle : Enregistre la progression toutes les 25 étapes, journalise les métriques toutes les 10 étapes

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

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

# MLflow and Unity Catalog configuration

# Model selection
MODEL_NAME = "Qwen/Qwen3-4B"
DATASET_NAME = "trl-lib/Capybara"
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"

# Training hyperparameters
BATCH_SIZE = 1
GRADIENT_ACCUMULATION_STEPS = 8
LEARNING_RATE = 2e-5
MAX_STEPS = 50
EVAL_STEPS = 25
LOGGING_STEPS = 10
SAVE_STEPS = 25

Chargez et préparez le dataset

La cellule suivante charge le dataset d’entraînement et le prépare pour l’affinement :

Dataset : trl-lib/Capybara – données conversationnelles de haute qualité optimisées pour le suivi des instructions
Fractionnement entraînement/validation : Crée un fractionnement 90/10 si aucun ensemble de test n'existe.
Validation des données : garantit un formatage approprié pour l’affinement conversationnel

Python
dataset = load_dataset(DATASET_NAME)
print(f"✓ Dataset loaded: {dataset}")

if "test" not in dataset:
    print("Creating validation split from training data...")
    dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)
    print("✓ Data split: 90% train, 10% validation")

Initialiser le modèle et le tokeniseur

La cellule suivante charge le modèle de base et le tokenizer, puis les configure pour l'affinement conversationnel :

**Chargement du modèle** : download Qwen3-4B de Hugging Face avec une précision BF16
Configuration du tokenizer : Configure le tokenizer rapide avec un remplissage approprié
Chat formatting : Applies a chat Template for structured conversations if the tokenizer doesn't already define one
Configuration du jeton : définit le jeton de remplissage sur le jeton EOS pour une gestion correcte des séquences.

Python
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_fast=True
)

# Chat template formatting for conversational fine-tuning
if tokenizer.chat_template is None:
    print("Adding chat template for proper conversation formatting...")
    model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")
    print("✓ ChatML format applied for structured conversations")

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    print("✓ Padding token set to EOS token")

print("✓ Model and tokenizer loaded successfully")

Entraîner le modèle

La cellule suivante configure et exécute le processus d’affinement complet :

Configuration de la formation

Configuration de batch : 1 échantillon par appareil avec 8 étapes d'accumulation de gradient (taille de batch effective : 8)
Optimisation : Étapes d’initialisation, dégradation du poids et sélection du meilleur modèle basée sur la perte d’évaluation
Journalisation : Rapporte les métriques à MLflow pour le suivi des expérimentations

Optimisations clés activées

Précision mixte BF16 : calcul plus rapide avec une empreinte mémoire réduite, bien adapté aux GPU H100
Point de contrôle du gradient : Permet de troquer un compute supplémentaire contre une réduction importante de la mémoire d’activation, ce qui permet d’adapter un réglage fin complet de 4 B sur un seul H100.
Accumulation de gradient : simule des tailles de batch plus grandes pour un entraînement stable
Points de contrôle : Enregistre le modèle toutes les 25 étapes avec une limite de 2 points de contrôle

La boucle d'entraînement Logs la progression toutes les 10 étapes et évalue toutes les 25 étapes.

Python
with mlflow.start_run(run_name=f"{MODEL_NAME}_full-fine-tuning", log_system_metrics=True):
    try:
        print(f"Learning rate: {LEARNING_RATE}")

        training_args_dict = {
            "output_dir": OUTPUT_DIR,
            "per_device_train_batch_size": BATCH_SIZE,
            "per_device_eval_batch_size": BATCH_SIZE,
            "gradient_accumulation_steps": GRADIENT_ACCUMULATION_STEPS,
            "learning_rate": LEARNING_RATE,
            "max_steps": MAX_STEPS,
            "eval_steps": EVAL_STEPS,
            "logging_steps": LOGGING_STEPS,
            "save_steps": SAVE_STEPS,
            "save_total_limit": 2,
            "report_to": "mlflow",  # Log to MLflow
            "warmup_steps": 10,
            "weight_decay": 0.01,
            "metric_for_best_model": "eval_loss",
            "greater_is_better": False,
            "eval_strategy": "steps",  # Run evaluation every eval_steps
            "save_strategy": "steps",  # Checkpoint on the same cadence as eval
            "load_best_model_at_end": True,  # Register the best-eval checkpoint, not the last
            "dataloader_pin_memory": False,
            "remove_unused_columns": False,
            "bf16": True,  # Mixed precision training
            "gradient_checkpointing": True,  # Reduce activation memory for full fine-tuning
            "gradient_checkpointing_kwargs": {"use_reentrant": False},
        }

        training_args = SFTConfig(**training_args_dict)

        trainer = SFTTrainer(
            model=model,
            args=training_args,
            train_dataset=dataset["train"],
            eval_dataset=dataset["test"],
            processing_class=tokenizer,
        )

        print("\n" + "="*50)
        print("STARTING TRAINING")
        print("="*50)

        print("🚀 Full fine-tuning Qwen3-4B on a single H100 GPU")

        trainer.train()
        print("\n✓ Training completed successfully!")

    except Exception as e:
        print(f"✗ Training failed: {e}")
        raise

Enregistrer les artefacts de modèle

La cellule suivante enregistre le modèle entraîné et le tokenizer dans le volume Unity Catalog :

**Poids complets du modèle** : Enregistre le modèle affiné complet, prêt à être chargé directement pour l'inférence.
Tokenizer : Enregistre la configuration du tokenizer pour l'inférence
Emplacement de stockage : Enregistre dans /Volumes/{catalog}/{schema}/{volume}/{model_name}

Python
try:
    print("\nSaving trained model...")

    trainer.save_model(training_args.output_dir)
    print("✓ Full model weights saved")

    tokenizer.save_pretrained(training_args.output_dir)
    print("✓ Tokenizer saved with model")
    print(f"\n🎉 All artifacts saved to: {training_args.output_dir}")

except Exception as e:
    print(f"✗ Model saving failed: {e}")
    raise

Enregistrer un modèle dans Unity Catalog

La cellule suivante enregistre le modèle affiné dans Unity Catalog pour la gouvernance et le déploiement :

Workflow d'enregistrement du modèle

Charger le modèle entraîné : charge le modèle à poids complet et le tokenizer enregistrés
Préparation de la journalisation : Crée un dictionnaire de modèle de transformeurs avec le modèle et le tokeniseur.
Enregistrer dans Unity Catalog : Enregistre les logs dans MLflow et s'enregistre dans Unity Catalog.
Ajouter des métadonnées : Inclut le type de tâche, la famille de modèles et les informations de taille.

Avantages de l'enregistrement Unity Catalog

**Gouvernance** : registre de modèles centralisé avec contrôle d'accès et traçabilité du lignage
Gestion des versions : Gestion automatique des versions pour le cycle de vie du modèle.
Déploiement : Déploiement facile vers les endpoints Model Serving
Découvrabilité : les modèles sont consultables et documentés dans Unity Catalog

Python
mlflow_run_id = mlflow.last_active_run().info.run_id
print("\nRegistering model with MLflow and Unity Catalog...")

with mlflow.start_run(run_id=mlflow_run_id):
    try:
        # Load the trained full-weight model for registration
        print("Loading fine-tuned model for registration...")
        trained_model = AutoModelForCausalLM.from_pretrained(
            training_args.output_dir,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True
        )
        tokenizer = AutoTokenizer.from_pretrained(training_args.output_dir)
        model_type = "Full fine-tuning"
        size_params = "4b"

        # Prepare transformers model dictionary
        transformers_model = {
            "model": trained_model,
            "tokenizer": tokenizer
        }

        # Create Unity Catalog model name
        full_model_name = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

        print(f"Registering model as: {full_model_name}")

        # Start MLflow run and log model
        task = "llm/v1/chat"
        model_info = mlflow.transformers.log_model(
            transformers_model=transformers_model,
            task=task,
            registered_model_name=full_model_name,
            metadata={
                &quot;task&quot;: task,
                &quot;pretrained_model_name&quot;: MODEL_NAME,
                &quot;databricks_model_family&quot;: &quot;Qwen3ForCausalLM&quot;,
                &quot;databricks_model_size_parameters&quot;: size_params,
            },
            repo_type="local",  # Fix: specify repo_type for local path
        )

        print(f"✓ Model successfully registered in Unity Catalog: {full_model_name}")
        print(f"✓ MLflow model URI: {model_info.model_uri}")

        # Print deployment information
        print(f"\n📦 Model Registration Complete!")
        print(f"Unity Catalog Path: {full_model_name}")
        print(f"Model Type: {model_type}")

    except Exception as e:
        print(f"✗ Model registration failed: {e}")
        print("Model is still saved locally and can be registered manually")
        print(f"Local model path: {training_args.output_dir}")
        raise

Étapes suivantes

Votre modèle Qwen3-4B a été affiné avec succès avec un affinement à poids complet et enregistré dans Unity Catalog. Ensuite, vous pouvez :

Déployer le modèle : Mettre à disposition des modèles avec Model Serving
En savoir plus sur l'entraînement distribué : entraînement distribué multi-GPU et multi-nœuds
Suivez les Experimentation et surveillez les GPU : Suivi et observabilité des Experimentation
Dépanner les problèmes : Dépanner les problèmes sur le compute GPU serverless

Exemple de Notebook

Affinement complet de Qwen3-4B

Ouvrir le notebook dans un nouvel onglet Ouvrir dans Databricks

Matrice de décision Affinement complet vs LoRA​

Se connecter au compute GPU serverless​

Importer des bibliothèques​

Configuration de l'installation​

Intégration de Unity Catalog​

Hyperparamètres d'entraînement​

Chargez et préparez le dataset​

Initialiser le modèle et le tokeniseur​

Entraîner le modèle​

Configuration de la formation​

Optimisations clés activées​

Enregistrer les artefacts de modèle​

Enregistrer un modèle dans Unity Catalog​

Workflow d'enregistrement du modèle​

Avantages de l'enregistrement Unity Catalog​

Étapes suivantes​

Exemple de Notebook​

Affinement complet de Qwen3-4B

Matrice de décision Affinement complet vs LoRA

Se connecter au compute GPU serverless

Importer des bibliothèques

Configuration de l'installation

Intégration de Unity Catalog

Hyperparamètres d'entraînement

Chargez et préparez le dataset

Initialiser le modèle et le tokeniseur

Entraîner le modèle

Configuration de la formation

Optimisations clés activées

Enregistrer les artefacts de modèle

Enregistrer un modèle dans Unity Catalog

Workflow d'enregistrement du modèle

Avantages de l'enregistrement Unity Catalog

Étapes suivantes

Exemple de Notebook