Skip to content

English · Español

Lab 01 — Experimento de scaling laws

🇪🇸 Tres corridas pequeñas (5M, 25M, 100M) sobre el mismo presupuesto fijo de FLOPs. Ajustamos la curva de Hoffmann y predecimos el conteo de tokens compute-óptimo para un hipotético modelo de 1B. El objetivo no es reproducir Chinchilla — es medir el principio en miniatura con datos propios.

Objetivo

Entrenar tres modelos pequeños (5M, 25M, 100M parámetros) sobre un presupuesto fijo de FLOPs. Ajustar la forma paramétrica de Hoffmann \(L(N, D) = E + A/N^\alpha + B/D^\beta\) a los triples \((N, D, L)\) resultantes. Usar el ajuste para predecir el conteo de tokens compute-óptimo de Chinchilla para un modelo hipotético de 1B parámetros.

¿Por qué tres tamaños?

Hoffmann 2022 usó >50 puntos IsoFLOP para ajustar cinco parámetros. Nosotros tenemos tres puntos y queremos ajustar cinco parámetros — el sistema está sub-determinado.

Esa es la lección: los ajustes de scaling laws (scaling laws) a escala pequeña producen barras de error lo suficientemente anchas como para meter un camión. El ejercicio es medir eso — cita la predicción con un intervalo de confianza honesto (que será ~2-3× ancho).

Configuración

Corrida N (no-embed) d_model n_layer n_head Objetivo de reloj Tokens (D)
A 5M 256 6 4 3 h 4.5B
B 25M 512 6 8 4 h 3.0B
C 100M 1024 8 16 5 h 1.5B

Reloj de pared total: ~12 horas en el mismo A100 80GB. Coste total: ~$13.

Presupuesto fijo de FLOPs por corrida. Las tres corridas apuntan a \(C \approx 1.35 \times 10^{18}\) FLOPs (\(\approx 2\) horas de cómputo a MFU 0.40; los objetivos de reloj de arriba incluyen overhead de warmup / val / checkpoint). Verifica:

  • Corrida A: \(6 \cdot 5 \times 10^6 \cdot 4.5 \times 10^9 = 1.35 \times 10^{17}\) — espera, eso es solo \(10^{17}\).

Permíteme rehacerlo con los números correctos. Para que las corridas sean IsoFLOP a \(C\) fijo:

  • Corrida A: \(6 \cdot 5 \times 10^6 \cdot D = C \Rightarrow D = C / (3 \times 10^7)\)
  • Corrida B: \(D = C / (1.5 \times 10^8)\)
  • Corrida C: \(D = C / (6 \times 10^8)\)

Poniendo \(C = 1.35 \times 10^{17}\) FLOPs (alrededor de 15 min de cómputo a 150 TFLOP/s):

  • Corrida A: \(D \approx 4.5 \times 10^9 = 4.5\)B tokens.
  • Corrida B: \(D \approx 0.9 \times 10^9 = 900\)M tokens.
  • Corrida C: \(D \approx 0.225 \times 10^9 = 225\)M tokens.

Eso nos da tres puntos sobre una curva IsoFLOP. En el mínimo IsoFLOP esperamos ver el \(N\) óptimo para este \(C\). La corrida A está sub-parametrizada (D/N = 900, muy sobre-entrenada); la corrida C está sobre-parametrizada (D/N = 2.25, severamente sub-entrenada); la corrida B (D/N = 36) debería estar cerca del óptimo.

El perfil IsoFLOP es aproximadamente una parábola en log-\(N\); el fondo de la parábola es \(N_\text{opt}\).

Estimaciones revisadas de reloj

A 150 TFLOP/s sostenidos: - \(C = 1.35 \times 10^{17}\) FLOPs → 900 s ≈ 15 min de cómputo por corrida. - Añade ~5 min de warmup + ~5 min de val + ~5 min de checkpoint = 30 min de reloj por corrida. - Tres corridas: ~90 min de cómputo en total. - Más setup inicial / carga de datos / eval final: ~2 horas de reloj en total.

Coste revisado: 2 h × \(1.10/h = ~\)2.20 para el triple IsoFLOP. (La estimación original de $13 es para un presupuesto de cómputo más grande por punto; lo bajamos para mantener el lab 01 barato.)

Presupuesto actualizado

Línea $-coste Notas
1× A100 80GB spot, Lambda, 2 h × $1.10/h $2.20 Tres corridas IsoFLOP
Disco persistente (reutilizado del lab 00) $0 Ya adjunto
Buffer $0.80 Margen de reinicio
Tope del lab 01 $3 Barato

Si has ejecutado el lab 00 primero y reutilizas la misma instancia + datos + dir de checkpoints, el lab 01 cabe en $3.

Los datos

Reutiliza el binario tokenizado de FineWeb-Edu del lab 00 (/workspace/data/tokenized/train.bin).

Para los 4.5B tokens de la corrida A, excedemos los tokens distintos de la muestra de 10B si queremos evitar duplicación. Permite ~1.5 épocas (4.5B / 10B = 0.45 épocas de la muestra completa; bien por debajo de 1 época — está bien).

El comando de entrenamiento (por corrida)

Mismo script que el lab 00, config diferente:

# Corrida A (5M parámetros, 4.5B tokens)
python -m x1_pretrain.train \
  --config configs/x1-isoflop-A.yaml \
  --total-steps 8800 \
  --eval-every 880 \
  --mlflow-uri file:///workspace/mlruns \
  --mlflow-experiment isoflop-A

# Corrida B (25M, 900M tokens)
python -m x1_pretrain.train \
  --config configs/x1-isoflop-B.yaml \
  --total-steps 1760 \
  --eval-every 176 \
  --mlflow-experiment isoflop-B

# Corrida C (100M, 225M tokens)
python -m x1_pretrain.train \
  --config configs/x1-isoflop-C.yaml \
  --total-steps 440 \
  --eval-every 44 \
  --mlflow-experiment isoflop-C

Los conteos de steps asumen tokens efectivos de batch = 512K del lab 00: - Corrida A: 4.5B / 512K ≈ 8800 steps. - Corrida B: 900M / 512K ≈ 1760 steps. - Corrida C: 225M / 512K ≈ 440 steps.

Val losses esperados

Predicciones del ajuste de Hoffmann (theory/01-scaling-laws.md, con \(E=1.69\), \(A=406.4\), \(\alpha=0.34\), \(B=410.7\), \(\beta=0.28\)):

Corrida N D L predicho Rango realista
A 5M 4.5B 4.66 [4.4, 5.0]
B 25M 900M 4.40 [4.2, 4.6]
C 100M 225M 4.48 [4.3, 4.7]

Nótese que el loss predicho de la corrida B es el más bajo — eso coincide con la intuición del mínimo IsoFLOP. Los rangos realistas son anchos porque (a) el ajuste se hizo a escala mucho mayor y extrapola peor a \(N\) pequeño, y (b) FineWeb-Edu no son los datos sobre los que ajustó Hoffmann (él usó MassiveText).

El ajuste

Tras completar las tres corridas, ajusta la forma de Hoffmann a los triples (N, D, L).

python -m x1_pretrain.fit_scaling_law \
  --input experiments/X1-pretraining/scaling-law/runs.csv \
  --output experiments/X1-pretraining/scaling-law/fit.json \
  --bootstrap 1000

El script: 1. Lee (N, D, L) de runs.csv. 2. Ajusta \(L = E + A/N^\alpha + B/D^\beta\) vía scipy.optimize.curve_fit. 3. Bootstrappea 1000 remuestreos (remuestrea residuos; con solo 3 puntos el CI está dominado por el ruido residual). 4. Saca fit.json con estimaciones puntuales y CIs al 95% para cada uno de (E, A, α, B, β). 5. Renderiza iso-flop.png mostrando los tres puntos y la curva ajustada.

Output esperado (ilustrativo): 

{
  "E":     {"point": 2.1, "ci95": [1.4, 2.9]},
  "A":     {"point": 380, "ci95": [120, 950]},
  "alpha": {"point": 0.32, "ci95": [0.20, 0.46]},
  "B":     {"point": 290, "ci95": [80, 700]},
  "beta":  {"point": 0.26, "ci95": [0.15, 0.42]}
}

Los CIs son anchos porque tres puntos no pueden fijar cinco parámetros. Eso es esperado y es la lección.

La predicción

Dado el ajuste, predice \(D_\text{opt}\) para \(N = 1\)B:

\[D_\text{opt}(N=10^9) = \left( \frac{B \beta N^\alpha}{A \alpha} \right)^{1/\beta} \cdot 10^9\]

(Derivación: minimizar \(L(N, D)\) sobre \(D\) a \(C = 6ND\) fijo — da el escalado IsoFLOP-óptimo.)

Sustituyendo las estimaciones puntuales:

\[D_\text{opt}(10^9) = \left( \frac{290 \cdot 0.26 \cdot (10^9)^{0.32}}{380 \cdot 0.32} \right)^{1/0.26} \cdot 10^9\]

Calcula paso a paso: - \((10^9)^{0.32} = 10^{2.88} \approx 760\). - Numerador: \(290 \cdot 0.26 \cdot 760 = 57{,}300\). - Denominador: \(380 \cdot 0.32 = 121.6\). - Ratio: \(471\). - \(471^{1/0.26} = 471^{3.85} = ?\) — \(\log_{10}(471) \approx 2.673 \cdot 3.85 = 10.29\); \(10^{10.29} \approx 1.95 \times 10^{10}\).

Así que \(D_\text{opt}(10^9) \approx 20 \times 10^9 = 20\)B tokens — da \(D/N \approx 20\), coincidiendo con Chinchilla.

(El número que tú obtengas realmente de tu ajuste puede ser 8B u 80B — cualquier cosa en [3B, 100B] está dentro del CI de 3-puntos. Cita el CI honestamente.)

Qué enseña esto

  1. El número titular (20 tokens/parámetro) es robusto — incluso con tres puntos a escala pequeña y datos ligeramente diferentes, converges cerca del número de Chinchilla. El "20" es una propiedad estructural profunda del modelado de lenguaje, no una idiosincrasia de MassiveText.
  2. Las barras de error son enormes a N pequeño. Predecir de 5M-100M a 1B son 2 órdenes de magnitud de extrapolación. Las versiones de laboratorio frontera usan puntos de 10-100B parámetros y aun así citan incertidumbre.
  3. Las curvas IsoFLOP son casi planas cerca del óptimo. La corrida A y la corrida C están fuera del óptimo, pero su loss es solo ~0.1 nat peor que la corrida B. La penalización por estar 4× off-óptimo es modesta — por eso el campo toleró la receta de Kaplan durante dos años.
  4. La calidad de los datos interactúa con la ley. Los tokens filtrados de FineWeb-Edu son individualmente más informativos que los de MassiveText. Tu \(E\) ajustado debería ser más bajo que el 1.69 de Hoffmann por ~0.1-0.3, porque la entropía irreducible del "texto educativo de alta calidad" está por debajo de la de "toda la web".

Entregables (DoD check 2)

  • experiments/X1-pretraining/scaling-law/runs.csv con filas (run_id, N, D, val_loss, $-gastados).
  • experiments/X1-pretraining/scaling-law/fit.json con estimaciones de parámetros + CIs.
  • experiments/X1-pretraining/scaling-law/iso-flop.png mostrando los tres puntos y la curva ajustada.
  • experiments/X1-pretraining/scaling-law/prediction.md con:
  • \(D_\text{opt}(1\text{B})\) predicho con CI al 95%.
  • El ratio \(D/N\) implicado con CI.
  • Un párrafo comparando con el número publicado de Chinchilla.
  • Un párrafo honesto sobre el riesgo de extrapolación.

Qué puede ir mal y cómo responder

  • La corrida C diverge. 100M sobre 225M tokens está severamente sub-entrenada — el modelo es demasiado grande para los datos y el loss puede ser más ruidoso de lo que el ajuste asume. Ejecútala de todas formas; cita el loss ruidoso; el CI ancho lo absorbe.
  • La corrida A sobre-ajusta al val set. Con 4.5B tokens y solo 1k docs de val de la misma distribución, el val loss puede bajar por debajo del train loss. Usa un shard held-out (último 10% del rango de tokens) — nunca solapes.
  • El throughput difiere entre corridas. Los modelos más pequeños a menudo alcanzan MFU más bajo porque su batch es demasiado pequeño para mantener los matmul cores ocupados. Reduce batch_size y añade grad_accum para la corrida A para recuperar throughput, o acéptalo (el objetivo IsoFLOP se fija por \(C\), no por reloj).
  • El ajuste falla en converger. scipy.optimize.curve_fit puede fallar con guesses iniciales malos. Usa los valores publicados de Hoffmann como inicialización.

Fin del Módulo de Extensión X1. Tras completar el lab 01, escribe learners/borja/extension-X1/reflections.md según el DoD del README.