Skip to content

English · Español

Break 00 — Tamaños de shard desajustados entre workers DP; el all-reduce se atasca

🇪🇸 El all-reduce de NCCL/gloo asume que cada worker contribuye exactamente el mismo número de bytes. Si un worker tiene un shard 1 elemento más pequeño que los demás, la operación o bien cuelga indefinidamente o devuelve resultados sin sentido. Este /break introduce ese desajuste y muestra cómo se ve el stall.


Qué vas a hacer

Haz que el split data-parallel asigne una muestra menos al worker 0 que a los demás. El forward pass todavía funciona (los gradientes están bien definidos), pero el paso de all-reduce espera que los buffers de gradientes tengan tamaños idénticos y o bien cuelga o devuelve valores silenciosamente erróneos.

Paso 1 — Localiza el splitter de shards

src/minidp/dataloader.py          # the per-worker shard slicing (Phase 35 lab 01)
src/minidp/allreduce.py           # the gloo all-reduce wrapper

Paso 2 — Introduce el bug

En src/minidp/dataloader.py, el sharding usa actualmente np.array_split que mantiene los shards balanceados. Reemplaza con un slice desbalanceado:

# OLD — balanced shards, total batch = N * B per step
shards = np.array_split(np.arange(global_batch), world_size)

# NEW (the broken version)
# Worker 0 gets one fewer sample than the others.
shards = np.array_split(np.arange(global_batch), world_size)
if rank == 0:
    shards = shards[:-1]   # one fewer sample on worker 0

El bucle de entrenamiento sigue produciendo gradientes. Las shapes divergen en 1 solo en la dimensión batch que se promedia dentro de la pérdida, así que los tensores de gradientes siguen bien definidos y tienen la misma shape de parámetro — el bug no dispara un chequeo de shape.

Pero la norma del gradiente en el worker 0 es incorrecta (se promedia sobre B-1 muestras en lugar de B), y dependiendo de la build de gloo, el all-reduce o bien:

  • cuelga (gloo con aserciones estrictas de tamaño),
  • completa silenciosamente con una suma errónea (paths antiguos de gloo),
  • o dispara un RuntimeError: size mismatch desde el wrapper si el buffer se redimensiona dinámicamente.

Paso 3 — Registra el break

learners/borja/phase-35/notes/breaks.md:

- bug-id: 35-01
  concept: collective ops require matching contributions
  symptom: training step at iter 1 either hangs (no progress, no logs after
           "starting all-reduce") OR loss diverges within 10 steps with a
           wrong gradient magnitude on worker 0.
  hidden_cause: dataloader.py drops one sample on rank 0; effective batch
                differs across workers; the gradient average is wrong.
  hint_1: "Print len(batch) on each worker at iter 1. Are they equal?"
  hint_2: "Compare loss on worker 0 vs worker 1 for the same global step."
  hint_3: "grep 'array_split' in dataloader.py. What's the slice doing?"
  fix_diff: remove the `if rank == 0: shards = shards[:-1]` post-split tweak.

Paso 4 — Verifica que es observable

Ejecuta just dp-train (o el entrypoint que sea del lab 01). Esperado con el bug:

[rank=0] step 0 loss=2.40
[rank=1] step 0 loss=2.39
[rank=0] step 1 starting all-reduce ...
[rank=1] step 1 starting all-reduce ...
[STALL: 60 seconds, no further output]

o, en la variante que no cuelga:

[rank=0] step 10 loss=2.38
[rank=1] step 10 loss=1.92   <-- divergent!

Cualquiera es observable. El test tests/phase35/test_dp_consistency.py::test_ranks_agree_at_step_zero se pone rojo.

Paso 5 — El momento de enseñanza

Las operaciones colectivas son síncronas y estrictas en tamaño por diseño. El argumento de rendimiento (Fase 35 teoría 03/05) depende de que cada worker contribuya el mismo número de bytes; una contribución asimétrica o bien provoca deadlock en el protocolo o produce una suma sin sentido.

Este es uno de los bugs más comunes en entrenamiento distribuido real: un dataloader que da al último worker el batch sobrante termina con B - r items donde r = global_batch % world_size, y a menos que hagas pad-o-drop consistentemente, los gradientes divergen. La solución estándar es DistributedSampler(drop_last=True) o padding.

La lección: el dataloader es parte del contrato distribuido, no una preocupación local.

Reglas duras respetadas

  • Un solo bug; un solo condicional.
  • Reversible en 2 líneas.
  • Observable (cuelgue o divergencia; cualquiera de los dos es un test que falla).
  • Sin impacto de seguridad.
  • Tests no modificados.

Siguiente: cuando esté en verde, lee ../theory/05-ring-allreduce-derivation-and-strategy-choice.md.