Seguimiento de experimentos y observabilidad

AI Runtime se integra con MLflow para proporcionar seguimiento de experimentos y observabilidad. MLflow proporciona un único lugar para los parámetros de una ejecución, métricas, métricas del sistema GPU, registros y artefactos, ya sea que desarrolles en un notebook, envíes una carga de trabajo con la databricks air CLI o programes un trabajo.

Lo que ofrece MLflow

Un experimento MLflow agrupa las ejecuciones de entrenamiento relacionadas, y cada ejecución proporciona un registro persistente de la configuración y resultados del entrenamiento.

• Métricas y parámetros: Registra la pérdida de entrenamiento, métricas de evaluación, tasas de aprendizaje e hiperparámetros, y luego compáralos entre las ejecuciones.
• Métricas del sistema: Consulta la utilización de GPU, CPU y memoria para ejecuciones gestionadas en la pestaña de métricas del sistema .
• Logs: Consulta la salida del driver para una carga de trabajo asociada a un trabajo en la pestaña de Logs de la ejecución.
• Artefactos y modelos: Almacenar archivos de modelos, configuraciones y otras salidas junto con la ejecución. Utiliza un volumen del Catálogo de Unity para artefactos grandes o duraderos.
• Compartir y colaborar: Concede acceso a un experimento para que otros usuarios del espacio de trabajo puedan ver y comparar sus ejecuciones. Consulte Organización de ejecuciones de entrenamiento con experimentos de MLflow.
• Integraciones de frameworks: Utiliza integraciones de MLflow con Hugging Face Transformers, PyTorch Lightning y otras bibliotecas compatibles.

Para cargas de trabajo respaldadas por jobs, como las cargas de trabajo enviadas con databricks air run, la ejecución de trabajos de Azure Databricks complementa MLflow registrando el estado de ejecución, los recursos de proceso, los intentos y la salida del driver.

Cómo se crean las ejecuciones de MLflow

Cómo inicias o te unes a una ejecución de MLflow depende de cómo inicies la carga de trabajo:

Después de que tu código empiece o se una a la ejecución correcta, usa la API de seguimiento de MLflow para registrar datos de entrenamiento personalizados. Para frameworks compatibles, puedes usar el autologging en lugar de añadir llamadas individuales de registro. Para patrones de deep learning, consulta el flujo de trabajo de deep learning de MLflow 3.

Utiliza las siguientes prácticas al configurar el seguimiento de MLflow:

• Establece step a un valor significativo como el lote global o la época. Registra en intervalos en vez de en cada lote. MLflow limita el número de pasos métricos por ejecución. Consulte Límites de los recursos.
• Para el entrenamiento distribuido, registra una métrica desde rank 0 solo para que cada valor se registre una vez.
• Para flujos de trabajo de cuaderno, utiliza rutas de experimento absolutas, como /Users/<username>/my-experiment o /Workspace/Shared/<team>/my-experiment. Pon los experimentos que quieras compartir en una carpeta de espacio de trabajo compartido.
• Guarda puntos de control grandes y duraderos en un volumen del Catálogo de Unity. Para patrones de checkpointing, véase Improve training performance and resiliency on AI Runtime.

Supervisar las cargas de trabajo de la CLI de Databricks

Cuando envías una carga de trabajo con databricks air run, AI Runtime crea un trabajo de Azure Databricks para los detalles de ejecución y una ejecución de MLflow para los datos de los experimentos. El envío original y sus reintentos comparten una sola ejecución de trabajo, mientras que cada intento crea una ejecución separada de MLflow en el mismo experimento para que puedas comparar intentos.

Las métricas del sistema de GPU, CPU y memoria se capturan automáticamente para cada ejecución gestionada y aparecen en la pestaña de métricas del Sistema MLflow.

Configurar y registrar en la ejecución

Utiliza el YAML de la carga de trabajo para seleccionar el experimento y, opcionalmente, configura el nombre de la ejecución y la ubicación del artefacto. El siguiente ejemplo muestra solo los campos MLflow:

experiment_name: my-training
mlflow_run_name: baseline-lr3e5
mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

Para requisitos de campo, valores predeterminados y valores admitidos, consulte la referencia Workload YAML.

AI Runtime expone el ID de ejecución gestionado a través de MLFLOW_RUN_ID. Para registrar tus propios parámetros, métricas y artefactos en esa ejecución, ábrela de nuevo con la API de seguimiento de MLflow:

import os

import mlflow

# Log from rank 0 only because all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
    with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
        mlflow.log_param("learning_rate", 3e-4)
        for step, loss in enumerate(training_losses):
            mlflow.log_metric("train_loss", loss, step=step)

Buscar ejecuciones y registros

Los registros están disponibles tanto a través de la CLI Databricks como de la página de ejecución de MLflow. Utiliza los siguientes comandos para encontrar cargas de trabajo, inspeccionar una carga de trabajo y transmitir o descargar registros:

databricks air list --all-status
databricks air get <job-run-id>
databricks air logs <job-run-id>
databricks air logs <job-run-id> --node 2
databricks air logs <job-run-id> --download-to ./logs/

Supervisar las cargas de trabajo de los notebooks

Desde un cuaderno, puedes ejecutar código directamente en la sesión adjunta o ejecutar entrenamiento distribuido a través del decorador @distributed desde el SDK de IA de Databricks.

Ejecuta código directamente en un notebook

El código ejecutado directamente en un cuaderno adjunto no crea automáticamente un run de MLflow. Selecciona un experimento, luego inicia una ejecución o activa el autologging:

import mlflow

mlflow.set_experiment("/Users/<username>/my-experiment")

with mlflow.start_run(run_name="baseline"):
    train()

Para habilitar el autologging de PyTorch Lightning, utiliza el siguiente patrón:

mlflow.pytorch.autolog()

with mlflow.start_run(run_name="lightning-baseline"):
    trainer.fit(model, datamodule=datamodule)

Para Hugging Face Transformers, configura TrainingArguments de la siguiente manera:

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="/Volumes/<catalog>/<schema>/<volume>/checkpoints",
    report_to="mlflow",
    run_name="llama7b-sft-lr3e5",
    logging_steps=50,
)

Utiliza el SDK de Runtime de IA de Databricks

Para el entrenamiento distribuido desde un notebook, utiliza el @distributed decorador del SDK de IA Runtime de Databricks (databricks-sdk-air). El SDK está disponible en el databricks.air espacio de nombres en el entorno GPU versión 6 y superior, o como serverless_gpu en el entorno GPU versión 5 e inferior.

from databricks.air import distributed

@distributed()
def train():
    # Training code
    ...

train.distributed()

Cada llamada .distributed() crea automáticamente una ejecución de MLflow. El experimento predeterminado es /Users/{WORKSPACE_USER}/{notebook-name}. Para seleccionar otro experimento, configúralo antes de llamar a .distributed():

mlflow.set_experiment("/Users/<username>/my-experiment")
train.distributed()

Cuando llamas .distributed() dentro de una ejecución activa, AI Runtime crea una ejecución hija anidada:

with mlflow.start_run() as parent_run:
    train.distributed()

Para reanudar una ejecución anterior, coloque MLFLOW_RUN_ID antes de .distributed():

import os

os.environ["MLFLOW_RUN_ID"] = "<previous-run-id>"
train.distributed()

Usa el enlace de ejecución en la salida de la celda o el panel Experiment del notebook para abrir la ejecución de MLflow.

Ver registros

• Salida del cuaderno: La salida estándar y los errores del código de entrenamiento aparecen en la salida de la celda del cuaderno.
• Logs de MLflow: Para cargas de trabajo de notebooks respaldadas por trabajos, como las llamadas .distributed(remote=True), la pestaña de Logs de MLflow muestra la salida de la ejecución del trabajo asociada.

Monitorizar los recursos en tiempo real de la GPU

El panel de recursos de GPU del notebook muestra la utilización en tiempo real de la GPU, la memoria y la temperatura para cargas de trabajo de un solo nodo y de varios nodos. Utiliza la pestaña de métricas del Sistema MLflow cuando necesites un registro persistente vinculado a una ejecución.

Para abrir el panel, conecta el cuaderno a AI Runtime y luego haz clic en Recursos de la GPU en el panel lateral derecho.

El panel sondea las métricas cada 10 segundos y conserva hasta 2 horas de historial. Haz clic Actualizar para obtener los valores más recientes. Tras 5 minutos de inactividad, el panel se detiene. Vuelve a abrirlo para reanudar la monitorización.

Solución de problemas para acceder a los registros

La pestaña Logs en una página de ejecución de MLflow muestra la salida en tiempo real de la ejecución del trabajo asociado de Azure Databricks. El acceso al experimento MLflow no implica acceso al trabajo.

Si la pestaña muestra que no tienes acceso a estos logs, pide a un usuario con permiso Puede gestionar o a un administrador del espacio de trabajo que al menos te conceda Puede ver en la tarea. Vea Controlar el acceso a una tarea.

Límites globales en Azure Databricks

Consulte Límites de los recursos.