Optimización de Hiperparámetros de LLM con Agentes Autónomos y Computación Paralela Heterogénea

El problema fundamental que aborda este trabajo es la ineficiencia inherente a la optimización secuencial de hiperparámetros en el entrenamiento de modelos de Machine Learning, especialmente LLMs, donde cada experimento es costoso en tiempo y recursos. La búsqueda de hiperparámetros es un problema de optimización combinatoria en un espacio de búsqueda de alta dimensionalidad, donde las interacciones entre parámetros pueden ser no lineales y difíciles de predecir. Tradicionalmente, esto se ha abordado con métodos como la búsqueda en cuadrícula (grid search), la búsqueda aleatoria (random search) o la optimización bayesiana, que a menudo son secuenciales o requieren una planificación manual intensiva.

La tesis central es que la combinación de agentes de código autónomos con una infraestructura de computación distribuida y heterogénea puede transformar radicalmente la eficiencia y la estrategia de esta búsqueda. Al permitir que el agente no solo modifique el código, sino que también gestione dinámicamente sus propios recursos de computación paralela, se supera el cuello de botella de la ejecución secuencial. Esto no solo acelera el proceso, sino que también permite estrategias de búsqueda más sofisticadas, como la evaluación de rejillas factoriales completas en una sola 'ola' de experimentos, lo que facilita la detección de efectos de interacción que serían invisibles para un enfoque secuencial de 'hill-climbing'. La relevancia actual es crítica dada la creciente complejidad y el costo computacional de entrenar modelos de IA de vanguardia.

Arquitectura del Sistema

El sistema se compone de tres elementos principales: un agente de código autónomo (Claude Code), un entorno de entrenamiento de modelos de lenguaje (Autoresearch de Karpathy) y una capa de orquestación de infraestructura (SkyPilot). El agente de código interactúa con el entorno de Autoresearch, que consta de prepare.py (datos y tokenizador, inmutable), train.py (modelo GPT, optimizador, bucle de entrenamiento, modificable por el agente) y program.md (instrucciones para el agente). El objetivo del agente es minimizar val_bpb (validation bits per byte) dentro de un presupuesto de tiempo fijo de 5 minutos por experimento.

La pieza clave para la paralelización es SkyPilot, una herramienta de código abierto que abstrae el aprovisionamiento y la gestión de recursos de computación en la nube (Kubernetes, AWS, GCP, Azure, etc.) a través de archivos YAML. El agente, instruido con una 'skill' de SkyPilot, genera y ejecuta archivos experiment.yaml que definen cada experimento. Estos YAML especifican el tipo de GPU (H100:1, H200:1), la imagen Docker (nvcr.io/nvidia/pytorch), el backend de infraestructura (k8s), el directorio de trabajo y los comandos de setup y run. El agente utiliza sky launch -d para provisionar nuevos clústeres de GPU y sky exec -d para enviar experimentos a clústeres existentes, lo que permite la ejecución en pipeline y el reuso de recursos. La salida de los experimentos se monitoriza a través de sky logs, permitiendo al agente leer val_bpb y peak_vram_mb para tomar decisiones. La arquitectura distribuida permite al agente mantener múltiples clústeres de GPU ocupados simultáneamente, transformando la estrategia de búsqueda de hiperparámetros de un proceso secuencial a uno masivamente paralelo.

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

Fundamentos Teóricos

Este trabajo se conecta con los fundamentos de la optimización de algoritmos y la computación distribuida. El problema de la búsqueda de hiperparámetros es una instancia del problema de optimización global, donde el espacio de búsqueda es a menudo no convexo y de alta dimensionalidad. Métodos como la búsqueda en cuadrícula y la búsqueda aleatoria son bien conocidos en la literatura de optimización de Machine Learning, con trabajos como Bergstra y Bengio (2012) demostrando la superioridad de la búsqueda aleatoria sobre la búsqueda en cuadrícula para espacios de alta dimensionalidad. La capacidad de ejecutar experimentos en paralelo y explorar rejillas factoriales completas se alinea con los principios de la computación paralela y distribuida, que buscan reducir el tiempo de ejecución de tareas computacionalmente intensivas dividiéndolas entre múltiples procesadores o nodos. La estrategia emergente del agente de explotar hardware heterogéneo (H100s para screening, H200s para validación) recuerda a los principios de la programación de tareas en sistemas distribuidos, donde las cargas de trabajo se asignan a recursos con capacidades adecuadas para optimizar el rendimiento y el costo, un concepto explorado en trabajos sobre scheduling en entornos de clúster heterogéneos.