Aller au contenu principal
Dernière mise à jour le

Affinement distribué d'OpenAI gpt-oss-20b

Ouvrir dans Databricks

Ajustez le modèle gpt-oss-20b d'OpenAI avec l'entraînement distribué sur AI Runtime. Vous apprendrez comment :

  • Appliquez **LoRA (Low-Rank Adaptation)** pour régler finement un modèle de 20 milliards de paramètres.
  • Utilisez la quantification MXFP4 pour réduire les besoins en mémoire pendant l'entraînement
  • Tirez parti du parallélisme distribué des données sur 8 GPU H100
  • Enregistrez le modèle affiné dans Unity Catalog pour le déploiement.

Concepts clés :

  • gpt-oss-20b: le modèle de langage open source de 20 milliards de paramètres d'OpenAI.
  • LoRA: affinement efficace des paramètres qui entraîne de petites couches d'adaptateur tout en figeant le modèle de base
  • Quantification MXFP4: Format de virgule flottante 4 bits à micro-échelle qui réduit l'utilisation de la mémoire
  • TRL: bibliothèque Transformer Reinforcement Learning pour l'affinement supervisé
  • AI Runtime: une compute gérée par Databricks qui monte en charge automatiquement les ressources GPU

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 provisionnera automatiquement 8 GPU H100 pour l'entraînement distribué.

Configurer Unity Catalog et les paramètres du modèle​

Configurez les paramètres de configuration pour l'enregistrement Unity Catalog et l'entraînement de modèle. Vous pouvez personnaliser ces paramètres à l'aide des widgets ci-dessus :

  • uc_catalog , uc_schema , uc_model_name : Emplacement Unity Catalog pour l'enregistrement du modèle
  • uc_volume : Nom du volume pour le stockage des points de contrôle du modèle
  • Modèle : identifiant de modèle Hugging Face (par default : openai/gpt-oss-20b)
  • dataset_path : Dataset à utiliser pour l'affinement (default : HuggingFaceH4/Multilingual-Thinking)
Python
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "gpt-oss-20b-peft")
dbutils.widgets.text("uc_volume", "checkpoints")
dbutils.widgets.text("model", "openai/gpt-oss-20b")
dbutils.widgets.text("dataset_path", "HuggingFaceH4/Multilingual-Thinking")

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")
HF_MODEL_NAME = dbutils.widgets.get("model")
DATASET_PATH = dbutils.widgets.get("dataset_path")

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"HF_MODEL_NAME: {HF_MODEL_NAME}")
print(f"DATASET_PATH: {DATASET_PATH}")

OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"
print(f"OUTPUT_DIR: {OUTPUT_DIR}")

Choisissez votre dataset​

Par default, ce notebook utilise « HuggingFaceH4/Multilingual-Thinking », qui a été spécifiquement organisé avec des chaînes de pensées traduites dans plusieurs langues. Vous pouvez modifier le paramètre « Chemin du jeu de données » ci-dessus pour utiliser un autre jeu de données.

Définir l'utilitaire de journalisation de la mémoire GPU​

Cette fonction utilitaire aide à surveiller l'utilisation de la mémoire GPU lors de l'entraînement distribué. Il logs la mémoire allouée et réservée pour chaque rang GPU, ce qui est utile pour le debugging des problèmes de mémoire.

Python
import os
import torch
import torch.distributed as dist

def log_gpu_memory(tag=""):
    if not torch.cuda.is_available():
        return

    # rank info (if distributed is initialized)
    if dist.is_available() and dist.is_initialized():
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    else:
        rank = 0
        world_size = 1

    device = torch.cuda.current_device()  # current GPU for this process
    torch.cuda.synchronize(device)

    allocated = torch.cuda.memory_allocated(device) / 1024**2
    reserved  = torch.cuda.memory_reserved(device) / 1024**2

    print(
        f"[{tag}] rank={rank}/{world_size-1}, "
        f"device={device}, "
        f"allocated={allocated:.1f} MB, reserved={reserved:.1f} MB"
    )

Définir la fonction de formation distribuée​

La cellule suivante définit la fonction d'entraînement en utilisant le décorateur @distributed de la bibliothèque serverless_gpu. Ce décorateur :

  • Provisionne 8 GPU H100 à la demande pour l'entraînement distribué
  • Gère automatiquement le parallélisme des données sur plusieurs GPU.

La fonction comprend :

  • Chargement et tokenisation des datasets
  • Initialisation du modèle avec quantification MXFP4
  • Configuration de l'adaptateur LoRA
  • Entraînement avec point de contrôle du gradient et précision mixte
  • Enregistrement de modèles dans les volumes Unity Catalog
Python
from serverless_gpu import distributed

@distributed(gpus=8, gpu_type="h100")
def run_train():
    import logging
    import os
    import torch

    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    torch.cuda.set_device(local_rank)
    world_size = int(os.environ.get("WORLD_SIZE", str(torch.cuda.device_count())))

    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

    is_main = rank == 0
    if is_main:
        logging.info("DDP environment")
        logging.info(f"\tWORLD_SIZE={world_size}  RANK={rank}  LOCAL_RANK={local_rank}")
        logging.info(f"\tCUDA device count (this node): {torch.cuda.device_count()}")

    from datasets import load_dataset
    dataset = load_dataset(DATASET_PATH, split="train")

    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)

    from transformers import AutoModelForCausalLM, Mxfp4Config

    quantization_config = Mxfp4Config(dequantize=True)
    model_kwargs = dict(
        attn_implementation="eager", # Use eager attention implementation for better performance
        dtype=torch.bfloat16,
        quantization_config=quantization_config,
        use_cache=False, # Since using gradient checkpointing
    )

    model = AutoModelForCausalLM.from_pretrained(HF_MODEL_NAME, **model_kwargs)

    from peft import LoraConfig, get_peft_model

    peft_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules="all-linear",
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )
    peft_model = get_peft_model(model, peft_config)
    if is_main:
        peft_model.print_trainable_parameters()

    from trl import SFTConfig

    training_args = SFTConfig(
        learning_rate=2e-4,
        num_train_epochs=1,
        logging_steps=1,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=2,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        max_length=2048,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine_with_min_lr",
        lr_scheduler_kwargs={"min_lr_rate": 0.1},
        output_dir=OUTPUT_DIR,
        report_to="mlflow",  # No reporting to avoid Gradio issues
        push_to_hub=False,  # Disable push to hub to avoid authentication issues
        logging_dir=None,  # Disable tensorboard logging
        disable_tqdm=False,  # Keep progress bars for monitoring
        ddp_find_unused_parameters=False,
    )

    from trl import SFTTrainer

    trainer = SFTTrainer(
        model=peft_model,
        args=training_args,
        train_dataset=dataset,
        processing_class=tokenizer,
    )
    #torch.cuda.empty_cache()
    #log_gpu_memory("before model training")
    result = trainer.train()
    #log_gpu_memory("after model loading")

    if is_main:
        logging.info("Training complete!")
        logging.info(f"Final training loss: {result.training_loss:.4f}")
        logging.info(f"Train runtime (s): {result.metrics.get('train_runtime', 'N/A')}")
        logging.info(f"Samples/sec: {result.metrics.get('train_samples_per_second', 'N/A')}")
        logging.info("\nSaving trained model...")
        trainer.save_model(OUTPUT_DIR)
        logging.info("✓ LoRA adapters saved - use with base model for inference")
        tokenizer.save_pretrained(OUTPUT_DIR)
        logging.info("✓ Tokenizer saved with model")
        logging.info(f"\n🎉 All artifacts saved to: {OUTPUT_DIR}")

    import mlflow
    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écutez la formation distribuée​

Cette cellule exécute la fonction d’entraînement sur 8 GPU H100. La formation prend généralement de 30 à 60 minutes selon la taille du dataset et la disponibilité du compute. La fonction renvoie l’ID d’exécution MLflow pour l’enregistrement du modèle.

Python
run_id = run_train.distributed()[0]

Enregistrer un modèle dans Unity Catalog​

Vous pouvez désormais enregistrer le modèle affiné avec MLflow et Unity Catalog pour le déploiement.

**Important :** Étant donné la taille du modèle (20B parameters), reconnectez le Notebook à l'accélérateur **H100** avant d'exécuter les cellules d'enregistrement.

Le processus d'inscription sera :

  1. Chargez le modèle de base et Merge-le avec les adaptateurs LoRA affinés.
  2. Créer un pipeline de génération de texte
  3. Enregistrer le modèle dans MLflow avec l'enregistrement Unity Catalog
Python
dbutils.widgets.dropdown("register_model", "False", ["True", "False"])
register_model = dbutils.widgets.get("register_model")
if register_model == "False":
  dbutils.notebook.exit("Skipping model registration...")

Vérifier le paramètre d'inscription​

Cette cellule vérifie le paramètre register_model. Si défini sur False, le Notebook ignorera l'enregistrement du modèle. Vous pouvez modifier ce parameter à l'aide du widget en haut du Notebook.

Python
print("\nRegistering model with MLflow and Unity Catalog...")

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

from peft import PeftModel
import mlflow
import torch

torch.cuda.empty_cache()
# Load the trained model for registration
print("Loading LoRA model for registration...")
# For LoRA models, we need both base model and adapter
base_model = AutoModelForCausalLM.from_pretrained(
    HF_MODEL_NAME,
    trust_remote_code=True
)
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)
adapter_dir = OUTPUT_DIR
peft_model = PeftModel.from_pretrained(base_model, adapter_dir)
# Merge LoRA into base and drop PEFT wrappers
merged_model = peft_model.merge_and_unload()

components = {
    "model": merged_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}")

text_gen_pipe = pipeline(
    task="text-generation",
    model=peft_model,
    tokenizer=tokenizer,
)

input_example = ["Hello, world!"]

with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model=text_gen_pipe,   # 🚨 pass the pipeline, not just the model
        artifact_path="model",
        input_example=input_example,
        # optional: save_pretrained=False for reference-only PEFT logging
        # save_pretrained=False,
    )
# Start MLflow run and log model
print(f"✓ Model successfully registered in Unity Catalog: {full_model_name}")
print(f"✓ MLflow model URI: {model_info.model_uri}")
print(f"✓ Model version: {model_info.registered_model_version}")

# Print deployment information
print(f"\n📦 Model Registration Complete!")
print(f"Unity Catalog Path: {full_model_name}")
print(f"Optimization: Liger Kernels + LoRA")

Tester les capacités de raisonnement multilingues​

Le modèle affiné a été entraîné sur le dataset Multilingual-Thinking, qui inclut un raisonnement en chaîne de pensée dans plusieurs langues.

La cellule suivante illustre cette fonctionnalité en :

  • Définir la langue de raisonnement en allemand
  • Fournir une invite en espagnol (« Quelle est la capitale de l'Australie ? »)
  • Constatant que le raisonnement interne du modèle s'effectue en allemand
Python
REASONING_LANGUAGE = "German"
SYSTEM_PROMPT = f"reasoning language: {REASONING_LANGUAGE}"
USER_PROMPT = "¿Cuál es el capital de Australia?"  # Spanish for "What is the capital of Australia?"

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": USER_PROMPT},
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
).to(merged_model.device)

gen_kwargs = {"max_new_tokens": 512, "do_sample": True, "temperature": 0.6, "top_p": None, "top_k": None}

output_ids = merged_model.generate(input_ids, **gen_kwargs)
response = tokenizer.batch_decode(output_ids)[0]
print(response)

Étapes suivantes​

Maintenant que vous avez affiné et testé votre modèle, vous pouvez :

Exemple de Notebook​

Affinement distribué d'OpenAI gpt-oss-20b

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks
Sur cette page