Skip to content

English · Español

Lab 00 — LoRALinear a mano

🇪🇸 Implementar LoRALinear desde cero. Tres invariantes innegociables: (1) en init la salida es idéntica a la del Linear base; (2) sólo A y B reciben gradiente; (3) merge_ produce pesos equivalentes byte-a-byte (dentro de 1e-5). Sin estos tres, todo lo demás falla en silencio.

Anclas: src/minituner/BLUEPRINT.md; theory/02-parameter-count.md.

Qué produces

Un módulo LoRALinear funcional en src/minituner/lora.py que envuelve un nn.Linear existente. Los tests fallidos en tests/test_minituner.py vienen pre-scaffolded por Claude; tu trabajo es rellenar lora.py hasta que pasen.

Entregables en disco:

  • src/minituner/lora.py — LoRALinear, wrap_minigpt_with_lora, lora_state_dict, load_lora_state_dict.
  • tests/test_minituner.py — en verde.
  • Una nota breve en markdown experiments/28-lora-by-hand/notes.md registrando: diff inicial de salida vs base, comprobación del flujo de gradientes, equivalencia tras merge.

TODOs (esbozo)

Bloque A — LoRALinear identidad en init

  1. Construir LoRALinear(base: nn.Linear, r=8, alpha=16.0, dropout=0.0, freeze_base=True).
  2. Reservar dos matrices de parámetros: A ∈ ℝ^{r × in_features}, B ∈ ℝ^{out_features × r}.
  3. Inicializar: A vía kaiming_uniform_(A, a=math.sqrt(5)) (coincide con el default de PyTorch para pesos de Linear); B = zeros.
  4. Opcionalmente aplicar nn.Dropout(p=dropout) a la entrada de la rama LoRA.
  5. Forward: y = base(x) + (alpha / r) * F.linear(F.linear(dropout(x), A), B). (Equivalente: (α/r) · x @ Aᵀ @ Bᵀ.)
  6. Si freeze_base: poner base.weight.requires_grad = False y base.bias.requires_grad = False (si existe bias).

Bloque B — Envolver un MiniGPT

  1. wrap_minigpt_with_lora(model, r=8, alpha=16.0, target_modules=("q_proj", ..., "mlp.fc2"), dropout=0.05) recorre model.named_modules().
  2. Para cada nn.Linear cuyo nombre con puntos termine en uno de los target_modules, reemplázalo in-place por un LoRALinear que envuelve el original.
  3. Tras envolver: comprobar que sum(p.requires_grad for p in model.parameters() if "lora_" in name) > 0 y que todos los params base tienen requires_grad=False.

Bloque C — Round-trip de checkpoint del adapter

  1. lora_state_dict(model) devuelve un dict que mapea nombres totalmente cualificados de parámetros → tensores, restringido a A y B de LoRA (y cualquier escala por módulo si la guardaste como buffer).
  2. load_lora_state_dict(model, state) hace lo inverso, validando compatibilidad de formas.
  3. Test de round-trip: guardar → recargar → la salida del forward coincide con la original dentro de 1e-5.

Bloque D — Merge para inferencia

  1. merge_(self) calcula W_new = base.weight + (alpha/r) * B @ A in-place, y luego pone A, B a cero.
  2. Post-merge: forward(x) coincide con el forward(x) no fusionado dentro de 1e-5 (la única diferencia es el redondeo en coma flotante).
  3. Después de fusionar (merge), las matrices LoRA pueden re-aleatorizarse para un nuevo adapter (no lo exijas — sólo no falles si ocurre).

Restricciones

  • PyTorch está permitido (la Fase 24+ lo desbloqueó). Usa torch.nn, torch.nn.functional. Sin import peft.
  • El forward debe permanecer numéricamente idéntico al base en init — dentro de epsilon de máquina. El test B = 0 ⇒ LoRA(x) == base(x) es innegociable.
  • No uses nn.Linear como matrices LoRA — instancia nn.Parameter directamente. Por qué: hace explícito el conjunto de parámetros; evita bias sorpresa.
  • No uses register_module("base", base) a ciegas — asegúrate de que el requires_grad=False del base se preserva al guardar/cargar.
  • Reproducibilidad: aplica la semilla vía el fixture de conftest antes de inicializar A.

Condiciones de parada

Has terminado cuando:

  1. pytest tests/test_minituner.py -k "lora_init_identity or param_count or merge or freeze or roundtrip" está en verde.
  2. mypy --strict src/minituner/lora.py pasa.
  3. notes.md registra valores medidos: diff inicial de salida (debe ser 0.0), diff pre/post-merge (debe ser < 1e-5), conteo de params entrenables para un Linear(64, 64) con r=8 (debe ser 1024).

Trampas (específicas de esta práctica)

  1. B = zeros no B ~ small_random. Un B aleatorio hace que el paso 0 no sea identidad; este es el bug de implementación de LoRA más común. Compruébalo explícitamente: assert (LoRALinear(x) - base(x)).abs().max() == 0 tras __init__.
  2. Olvidar la escala α/r. Sin ella, duplicar r duplica la LR efectiva — experimentos confusos. Pon alpha=16.0 por defecto; nunca dejes silenciosamente alpha=r.
  3. Modificación in-place de los pesos base. merge_ debe producir un nuevo peso de nn.Linear vía data.copy_ o construir un LoRALinear fresco con el nuevo base congelado. Cualquiera vale; no retengas accidentalmente los params LoRA tras el merge.
  4. Fuga de requires_grad. Si creas el nn.Parameter de LoRA después de poner base.requires_grad = False, pero una llamada posterior a model.train() o to(device) resetea cosas — fácil de pasar por alto. Comprueba la propiedad de freeze explícitamente tras .to(device).
  5. Dropout en la rama LoRA con model.eval(). El dropout debe estar apagado en eval; si lo implementaste vía F.dropout(..., self.training) crudo, verifica. Más fácil con una instancia de nn.Dropout.

Cuándo consultar las soluciones

Cuando los tests fallidos dejen de decirte algo nuevo (has mirado uno más de 20 minutos), abre solutions/00-lora-by-hand-ref.md. Las soluciones se escriben después del primer intento de Borja.

Tiempo estimado

3-5 horas. La dificultad conceptual es baja (ya has leído theory 02); la dificultad de implementación es moderada (sutilezas de nn.Module y el registro de parámetros en PyTorch).

Siguiente: lab/01-lora-counts.md.