Skip to content

English · Español

Lab 03 — Lectura anotada: tensor-parallel de Megatron-LM + FSDP de PyTorch

Objetivo: leer dos ficheros de entrenamiento distribuido de calidad de producción. Escribir notas anotadas señalando decisiones de diseño. Coste cloud cero.

Tiempo estimado: 3–5 horas (carga de lectura).

Prerrequisito: teoría 01–04 hecha; labs 00–02 hechos. Borja puede ahora leer estos ficheros con el vocabulario correcto cargado.

Qué produces

Un directorio experiments/35-reading-notes/ con dos notas de lectura anotadas:

  • megatron-tp-layers.md — lectura anotada de megatron-lm/megatron/core/tensor_parallel/layers.py (específicamente ColumnParallelLinear y RowParallelLinear).
  • fsdp-prefetch.md — lectura anotada de la lógica de prefetch de flat-param en torch/distributed/fsdp/_runtime_utils.py.

Cada nota: ≥5 decisiones de diseño señaladas con citas de líneas, ≤300 palabras cada una (para que sea un resumen, no una transcripción).

TODOs

Bloque A — clona Megatron-LM en un commit fijo

git clone https://github.com/NVIDIA/Megatron-LM.git /tmp/megatron-lm
cd /tmp/megatron-lm
git rev-parse HEAD > /home/overdrive/claude/lynx-cortex/experiments/35-reading-notes/megatron-sha.txt

Fija el SHA en las notas — el código de Megatron cambia; el Borja del futuro necesita saber qué versión leyó este lab.

Bloque B — lee la linear TP de Megatron

Fichero objetivo: megatron/core/tensor_parallel/layers.py. Clases foco: ColumnParallelLinear, RowParallelLinear. Ficheros laterales sugeridos para contexto: megatron/core/tensor_parallel/mappings.py (las primitivas _copy_to_tensor_model_parallel_region, _reduce_from_tensor_model_parallel_region, _gather_from_tensor_model_parallel_region).

Escribe megatron-tp-layers.md señalando ≥5 decisiones de diseño. Candidatas sugeridas (puedes elegir otras si las detectas):

  1. Por qué dividir la matriz de pesos en N column shards vía helpers de tipo init.partial — ¿qué hace Megatron que un torch.nn.Parameter(W[:, rank * o//N:(rank+1) * o//N]) ingenuo hace mal?
  2. Async tensor parallel + flag sequence_parallel — ¿qué hace sequence_parallel=True de manera distinta al TP plano? ¿Cuándo merece la pena la comm extra?
  3. gradient_accumulation_fusion — ¿qué cómputo se fusiona con la escritura del buffer de gradientes? ¿Por qué ayuda?
  4. La autograd Function para el all-reduce (busca _ReduceFromModelParallelRegion) — ¿por qué esto es una Function de autograd explícita, no un nn.Module?
  5. Inicialización en CPU vs GPU (flag use_cpu_initialization) — ¿cuándo quieres inicializar primero en CPU? (Pista: memoria.)
  6. async_tensor_model_parallel_allreduce — ¿con qué solapa el camino async?
  7. Manejo del bias bajo TP — el bias se replica, no se hace shard. ¿Por qué?

Para cada elección: 50–80 palabras. Cita la ruta del fichero + rango de líneas (p. ej., megatron/core/tensor_parallel/layers.py:L142-L168).

Bloque C — lee el prefetch de flat-param de FSDP de PyTorch

Fichero objetivo: torch/distributed/fsdp/_runtime_utils.py. Foco en los caminos _pre_forward y _post_forward, y la lógica _prefetch_handle. Fichero lateral: torch/distributed/fsdp/_flat_param.py para la estructura de flat-param.

Escribe fsdp-prefetch.md señalando ≥5 decisiones de diseño. Candidatas sugeridas:

  1. ¿Qué es un "flat parameter"? ¿Por qué FSDP aplana y no mantiene parámetros individuales?
  2. Prefetch-tras-shard-bound — ¿en qué momento FSDP emite el all-gather para la capa \(\ell + 1\)? ¿Qué se solapa?
  3. USE_ORIG_PARAMS — flag para usar los objetos de parámetro originales en lugar de las vistas flat. ¿Qué rompe? ¿Qué habilita?
  4. CPU_OFFLOAD y el movimiento a/desde CPU — ¿cuándo FSDP paginar los master weights entre GPU y CPU? ¿Cuál es el presupuesto de latencia para esto?
  5. All-gather del backward + reduce-scatter — ¿cuál es la dependencia? ¿Por qué el reduce-scatter no es un "all-reduce normal"?
  6. El set _handles_prefetched — ¿qué máquina de estados evita el double-prefetch o el stale-prefetch?
  7. Precisión mixta en FSDP — MixedPrecision(param_dtype=..., reduce_dtype=..., buffer_dtype=...) — ¿qué controla cada uno?

Para cada una: 50–80 palabras + cita de línea.

Bloque D — sintetiza: el diagrama

En experiments/35-reading-notes/synthesis.md, dibuja dos diagramas mermaid (~10 líneas cada uno):

  • Layout de bloque TP bajo Megatron: un bloque transformer con el patrón MLP column → GELU → row y el patrón attention column-QKV → attn → row-out. Marca los dos all-reduces. Anota "intra-nodo NVLink, ~600 GB/s" en los bordes de comm.
  • Línea temporal de forward de FSDP: el cómputo de capa en el track inferior, los all-gathers de prefetch en el track superior, mostrando el solapamiento.

Los diagramas son mermaid (texto-editables, versionables). Commitéalos.

Bloque E — párrafo de conectar-los-puntos

Termina cada nota con un párrafo de "conectar-los-puntos":

  • megatron-tp-layers.md: ¿cómo explica lo que acabas de leer el slowdown del lab 02 en el tutor de gramática? (Pista: el \(d_{\text{model}}\) del tutor es pequeño, así que el volumen por all-reduce es minúsculo respecto al cómputo por token, así que la comm domina.)
  • fsdp-prefetch.md: ¿cuándo sería FSDP la elección correcta para el entrenamiento del tutor de gramática? (Pista: nunca al tamaño de vocab actual; algún momento alrededor de \(|\theta| \approx 1\text{B}\).)

Restricciones

  • Sin reescrituras de código. No intentes "mejorar" Megatron o FSDP. El ejercicio es leer, no refactorizar.
  • Cita rangos de líneas, no fragmentos completos. El SHA del repo + rango de líneas basta para reconstruir lo que leíste. No pegues 200 líneas de fuente en tus notas.
  • ≤300 palabras por nota. Si se está alargando, estás transcribiendo, no resumiendo.
  • Solo diagramas mermaid. Sin PNGs de herramientas de dibujo — mantén los diagramas diff-ables.
  • Coste cloud cero. Esto es lectura. Solo portátil local.

Condiciones de parada

Has acabado cuando:

  1. experiments/35-reading-notes/{megatron-tp-layers.md, fsdp-prefetch.md, synthesis.md, megatron-sha.txt} existen.
  2. Cada .md tiene ≥5 bullets de decisión-de-diseño con citas de línea y ≤300 palabras.
  3. synthesis.md tiene dos diagramas mermaid (layout de bloque TP + línea temporal FSDP).
  4. Cada nota tiene un párrafo de cierre "conectar-los-puntos" atando de vuelta al lab 02 o al perfil de entrenamiento del tutor de gramática.
  5. Puedes responder, de memoria: "¿qué hace gradient_accumulation_fusion?" y "¿cuándo hace prefetch FSDP de la capa \(\ell+1\)?".

Pista de último recurso

Si la fuente de Megatron parece impenetrable: empieza por los docstrings del módulo arriba de layers.py. NVIDIA los ha mejorado recientemente; explican el split column/row con diagramas. Luego relee el código con los docstrings como mapa.

Para FSDP: docs/source/fsdp.rst de PyTorch (en el repo de PyTorch) es la narrativa oficial. Léelo primero, luego lee _runtime_utils.py. La narrativa te dice qué buscar en el código.

Cuándo consultar solutions/

Tras commitear las notas. La solución vive en solutions/03-reading-ref.md — escrita en la apertura de la fase. La solución es un conjunto de referencia de elecciones de decisión-de-diseño con citas, no "la respuesta correcta". Las elecciones de Borja pueden legítimamente diferir; la comparación debería ser "¿me dejé alguna de estas?" no "¿coincidí exactamente?".

Siguiente fase: docs/phase-36-frontier-architectures/.