Ajustar Llama 3.2 1B con AI Runtime

Abrir la versión 'notebook' de esta página

Ajuste completo de un modelo de lenguaje de gran tamaño con ajuste fino supervisado (SFT) en Databricks AI Runtime. El notebook utiliza la biblioteca Transformers Reinforcement Learning (TRL) con la optimización DeepSpeed ZeRO Stage 3 para entrenar Llama 3.2 1B de forma eficaz en un solo nodo con 8 GPU H100.

Conceptos clave:

  • TRL (Transformers Reinforcement Learning): biblioteca que proporciona herramientas para entrenar modelos de lenguaje con aprendizaje de refuerzo y ajuste fino supervisado.
  • DeepSpeed ZeRO Stage 3: técnica de optimización de memoria que particiona los parámetros del modelo, los degradados y los estados del optimizador entre GPU para habilitar el entrenamiento de modelos grandes.
  • Entorno de ejecución de IA: proceso de GPU administrado por Databricks que aprovisiona y escala automáticamente los recursos de GPU para cargas de trabajo de entrenamiento.

Para más información, consulte Entorno de ejecución de IA.

Nota

Este ejemplo requiere el entorno de IA Databricks versión 5 o superior.

Requisitos

Este cuaderno requiere lo siguiente:

  • AI Runtime: el cuaderno usa Databricks AI Runtime con 8 GPU H100 para entrenamiento distribuido. No se necesita ninguna configuración de clúster.
  • Catálogo de Unity: catálogo y esquema de Unity Catalog para almacenar puntos de control de modelo y registrar el modelo entrenado.
  • Token HuggingFace: un token de acceso HuggingFace almacenado en secretos de Databricks para descargar el modelo base y el conjunto de datos.
  • Python paquetes: AI Runtime preinstala las bibliotecas más necesarias. En la sección de instalación siguiente se instala deepspeed.

Conectar al cómputo de GPU sin servidor

Este notebook requiere computación GPU sin servidor. Para conectarse:

  1. Haga clic en el selector de proceso del cuaderno en la parte superior derecha y seleccione GPU sin servidor.
  2. A la derecha, haga clic en el botón de entorno.
  3. Seleccione 8xH100 como Acelerador.
  4. Elija entorno de AI v5 en el panel derecho que contenga todas las bibliotecas necesarias para ejecutar este ejemplo de cuaderno.
  5. Haga clic en Aplicar.

La función de entrenamiento aprovisiona automáticamente 8 GPU H100 para el entrenamiento distribuido.

Instalación de paquetes necesarios

Ai Runtime ya tiene la mayoría de las bibliotecas necesarias preinstaladas. Para este ejemplo, solo necesita instalar deepspeed.

%pip install deepspeed==0.19.1
%restart_python

Configuración de variables de entorno y catálogo de Unity

Configure las ubicaciones del catálogo de Unity para almacenar puntos de control del modelo y registrar el modelo entrenado. El cuaderno interactivo usa parámetros de consulta para configurar:

  • Catálogo y esquema: espacio de nombres del catálogo de Unity para organizar modelos y puntos de control
  • Nombre del modelo: nombre del modelo registrado en el catálogo de Unity
  • Volumen: Volumen del catálogo de Unity para almacenar checkpoints del modelo durante el entrenamiento

La configuración también obtiene el token HuggingFace de los secretos de Databricks y establece el experimento de MLflow para monitorear las métricas de entrenamiento.

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

Creación de una configuración de DeepSpeed ZeRO Stage 3

Las particiones de la fase 3 de DeepSpeed ZeRO (Optimizador de redundancia cero) modelan parámetros, gradientes y estados del optimizador entre todas las GPU para reducir el consumo de memoria de cada GPU. Esto permite el entrenamiento de modelos grandes que no caben en la memoria de una sola GPU.

Opciones de configuración clave:

  • bf16 habilitado: usa la precisión bfloat16 para un entrenamiento más rápido y un uso reducido de memoria
  • Optimización de la fase 3: crea particiones de todos los estados del modelo entre GPU
  • Sin descarga de CPU: mantiene todos los datos en GPU para obtener el máximo rendimiento en el hardware H100
  • Comunicación superpuesta: superpone la comunicación de gradiente con el cálculo para mejorar la eficacia.
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")

Definición de parámetros de entrenamiento

Configure los parámetros de ajuste fino supervisados:

  • Modelo: Llama 3.2 1B Indica, un modelo compacto adecuado para GPU H100
  • Conjunto de datos Capybara de la biblioteca TRL para entrenar inteligencia artificial conversacional
  • Tamaño del lote: 2 por dispositivo con 4 pasos de acumulación de degradado para un tamaño de lote efectivo de 64
  • Velocidad de aprendizaje: 2e-4 con programador de coseno y calentamiento
  • Pasos de entrenamiento: 60 pasos para la demostración (aumento para el entrenamiento completo)

La configuración utiliza precisión bfloat16 y el punto de control de gradiente para optimizar el uso de memoria.

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

Definición de la función de entrenamiento distribuido

El decorador @distributed de la biblioteca serverless_gpu permite ejecutar cargas de trabajo de GPU en Databricks AI Runtime. El decorador aprovisiona 8 GPU H100 y controla la configuración de entrenamiento distribuida automáticamente.

Parámetros clave:

  • gpus=8: solicita 8 GPU para el entrenamiento distribuido
  • gpu_type='H100': especifica hardware de GPU H100

La función de entrenamiento:

  1. Carga el modelo base y el tokenizador desde HuggingFace
  2. Establece la configuración del formato de chat para la inteligencia artificial conversacional
  3. Carga el conjunto de datos de entrenamiento
  4. Inicializa el TRL SFTTrainer con la optimización "DeepSpeed".
  5. Entrena el modelo y guarda los puntos de control
  6. Devuelve los resultados de entrenamiento y el identificador de ejecución de MLflow.

Para más información, consulte la documentación de la API de entorno de ejecución de IA.

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

Ejecuta el trabajo de entrenamiento distribuido

Ejecute la función de entrenamiento llamando a .distributed() en la función decorada. Esto aprovisiona los recursos de AI Runtime, ejecuta el entrenamiento en 8 GPU H100 con optimización DeepSpeed y devuelve los resultados.

El proceso de entrenamiento:

  • Aprovisiona 8 GPU H100 automáticamente
  • Descarga el modelo y el conjunto de datos de HuggingFace
  • Entrena el modelo con ajuste completo
  • Guarda los puntos de control en el volumen del Catálogo de Unity.
  • Registra métricas en MLflow
  • Devuelve el estado de entrenamiento, la pérdida final y el identificador de ejecución de 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']}")

Guardar el modelo ajustado y la inferencia de prueba

Este paso opcional carga el modelo ajustado y lo prueba con un mensaje de ejemplo para comprobar los resultados.

El proceso:

  1. Carga el modelo guardado
  2. Prueba del modelo con un mensaje de conversación de ejemplo
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)

Registro del modelo en el catálogo de Unity

Registre el modelo optimizado en MLflow y regístrelo en el Catálogo de Unity para su implementación y servicio. El modelo se registra con:

  • Modelo y tokenizador: ambos componentes necesarios para la inferencia
  • Tipo de tarea: configurado como llm/v1/chat para la inteligencia artificial conversacional
  • Ejemplo de entrada: formato de mensaje de chat de ejemplo para pruebas
  • Registro del catálogo de Unity: registra automáticamente el modelo en el catálogo y el esquema configurados.

Una vez registrado, el modelo se puede desplegar en los endpoints de servicio del modelo o usarse para la inferencia por lotes.

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

Pasos siguientes

Cuaderno de ejemplo

Ajustar Llama 3.2 1B con AI Runtime

Obtención del cuaderno

  • Last updated on