Optimización de Hiperparámetros de Modelos de LLM con Agentes Autónomos y Computación Distribuida

La optimización de hiperparámetros y la búsqueda de arquitecturas de modelos de Machine Learning, especialmente en el contexto de Large Language Models (LLMs), es un problema computacionalmente intensivo. Tradicionalmente, este proceso se ha abordado mediante búsquedas secuenciales o heurísticas que, si bien son efectivas, son lentas y propensas a quedar atrapadas en óptimos locales. La tesis central de este trabajo es que la aplicación de agentes autónomos con acceso a computación distribuida y heterogénea puede transformar la estrategia de búsqueda de hiperparámetros de un proceso secuencial y glotón a una exploración paralela y factorial, permitiendo descubrir interacciones complejas y optimizaciones que de otro modo serían inaccesibles en un tiempo razonable.

Este enfoque no solo acelera la experimentación, sino que también cambia fundamentalmente la naturaleza de la investigación, pasando de una serie de hipótesis individuales a la evaluación simultánea de rejillas factoriales. Esto es crucial en el diseño de arquitecturas de redes neuronales, donde la interacción entre parámetros como el ancho del modelo, la tasa de aprendizaje y los parámetros del optimizador puede tener efectos no lineales significativos. La capacidad de un agente para gestionar y explotar recursos computacionales heterogéneos de forma autónoma representa un avance hacia sistemas de IA más eficientes y auto-optimizables.

Arquitectura del Sistema

El sistema se compone de tres elementos principales: el agente de codificación (Claude Code), el entorno de experimentación (autoresearch de Andrej Karpathy) y la infraestructura de computación distribuida (SkyPilot sobre Kubernetes). El agente interactúa con el entorno de autoresearch, que consta de tres archivos: prepare.py (datos y funciones de evaluación, inmutable), train.py (modelo GPT, optimizador, bucle de entrenamiento, modificable por el agente) y program.md (instrucciones para el agente).

El agente modifica train.py, luego utiliza SkyPilot para lanzar experimentos en paralelo. SkyPilot actúa como una capa de abstracción para la gestión de recursos en la nube y Kubernetes. Cada experimento se define en un archivo experiment.yaml que especifica el tipo de GPU (H100, H200), la imagen Docker (nvcr.io/nvidia/pytorch), el backend de infraestructura (k8s), el directorio de trabajo y las variables de entorno. El bloque setup instala dependencias (pip install uv, uv sync, uv run prepare.py), mientras que el bloque run ejecuta uv run train.py con un presupuesto fijo de 5 minutos, capturando métricas como val_bpb y peak_vram_mb.

Para la ejecución paralela, el agente lanza múltiples clústeres (sky launch) y somete diferentes experimentos a cada uno. El comando sky exec permite reutilizar un clúster existente para experimentos subsiguientes, lo que minimiza el tiempo de inactividad al omitir el setup inicial. El agente monitorea los resultados con sky logs y utiliza esta información para planificar la siguiente ola de experimentos, adaptando su estrategia de búsqueda y explotando las diferencias de rendimiento entre GPUs H100 y H200 de forma autónoma.

resources:
  accelerators: {H100:1, H200:1}
image_id: docker:nvcr.io/nvidia/pytorch:24.07-py3
infra: k8s
workdir: .
envs:
  EXPERIMENT_ID: baseline
  EXPERIMENT_DESC: "baseline run"
setup: |
  pip install uv
  uv sync
  uv run prepare.py
run: |
  uv run train.py 2>&1 | tee run.log
  EXIT_CODE=${PIPESTATUS[0]}
  if [ $EXIT_CODE -ne 0 ]; then
    echo "EXPERIMENT_STATUS: crash"
  else
    VAL_BPB=$(grep "^val_bpb:" run.log | awk '{print $2}')
    PEAK_VRAM=$(grep "^peak_vram_mb:" run.log | awk '{print $2}')
    MEMORY_GB=$(echo "scale=1; ${PEAK_VRAM} / 1024" | bc)
    echo "EXPERIMENT_STATUS: done"
    echo "EXPERIMENT_RESULT: ${EXPERIMENT_ID} val_bpb=${VAL_BPB} memory_gb=${MEMORY_GB}"
  fi
  echo "EXPERIMENT_DESC: ${EXPERIMENT_DESC}"

# Launch a cluster with the first experiment
sky launch gpu-01 experiment.yaml -d -y \
  --env EXPERIMENT_ID=exp-01 \
  --env EXPERIMENT_DESC="baseline run"

# Reuse the same cluster for the next experiment (skips setup)
sky exec gpu-01 experiment.yaml -d \
  --env EXPERIMENT_ID=exp-02 \
  --env EXPERIMENT_DESC="higher LR"

Fundamentos Teóricos

La problemática de la optimización de hiperparámetros y la búsqueda de arquitecturas de redes neuronales se remonta a los primeros días del Machine Learning. Conceptos como la búsqueda en cuadrícula (grid search) y la búsqueda aleatoria (random search), popularizados por Bergstra y Bengio (2012), son métodos fundamentales para explorar el espacio de hiperparámetros. Sin embargo, estos métodos son computacionalmente costosos y no escalan bien a espacios de búsqueda de alta dimensionalidad o con interacciones complejas. La idea de un agente autónomo que aprende a optimizar su propio proceso de aprendizaje resuena con los principios del meta-aprendizaje (meta-learning) y la optimización bayesiana, donde se construye un modelo de la función objetivo para guiar la búsqueda de manera más eficiente.

La explotación de hardware heterogéneo y la adaptación de estrategias de búsqueda en función de las características de los recursos disponibles, como se observa en el agente que distingue entre H100 y H200, se alinea con la investigación en sistemas distribuidos y planificación de tareas (task scheduling) en entornos heterogéneos. Algoritmos de scheduling como el de Johnson o el de list-scheduling buscan optimizar el throughput en sistemas con diferentes capacidades de procesamiento. La capacidad del agente para inferir y explotar estas diferencias de rendimiento de forma autónoma es un ejemplo de inteligencia adaptativa en sistemas distribuidos, un área activa de investigación en la intersección de la IA y la computación de alto rendimiento.