aleleba
d121c4465e
00_verify_hardware.py: flash_attn, sdpa y flash-linear-attention
funcionan los tres en el GB10 (SM121) via Triton JIT -- mejor de lo
esperado, ya no hace falta el fallback lento de PyTorch para Gated
DeltaNet. attn_implementation recomendado: flash_attention_2.
01_inspect_modules.py: resuelve la discrepancia de nombres de las
proyecciones de Gated DeltaNet inspeccionando la arquitectura real
(device_map=meta, sin pesos). Resultado: parcialmente fusionado --
in_proj_qkv (q+k+v en un solo Linear) pero in_proj_z, in_proj_a e
in_proj_b por separado. Confirma tambien que mlp.experts.{gate_up,down}_proj
son nn.Parameter 3D (no LoRA-ables) y que shared_expert.{gate,up,down}_proj
son nn.Linear normales, como anticipaba el plan.
constraints.txt: fija la version exacta de torch de la imagen NGC
(2.10.0a0+b4e4ee81d3.nv25.12) para instalar transformers/peft/trl/
accelerate/bitsandbytes sin romper el build ARM64/Blackwell.
113 lines
4.3 KiB
Python
113 lines
4.3 KiB
Python
"""Fase 0: verifica qué attn_implementation usar y si flash-linear-attention compila en SM121.
|
|
|
|
No bloqueante: flash_attn y flash-linear-attention son "nice to have"; si fallan,
|
|
el resultado documentado es usar el fallback correspondiente (sdpa / PyTorch puro).
|
|
"""
|
|
import sys
|
|
import time
|
|
|
|
import torch
|
|
|
|
RESULTS = {}
|
|
|
|
|
|
def test_flash_attn():
|
|
try:
|
|
from flash_attn import flash_attn_func
|
|
except ImportError as e:
|
|
return False, f"import failed: {e}"
|
|
|
|
try:
|
|
device = "cuda"
|
|
dtype = torch.bfloat16
|
|
batch, seqlen, nheads, headdim = 2, 128, 8, 64
|
|
q = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype)
|
|
k = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype)
|
|
v = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype)
|
|
out = flash_attn_func(q, k, v, causal=True)
|
|
torch.cuda.synchronize()
|
|
assert out.shape == (batch, seqlen, nheads, headdim)
|
|
return True, f"forward ok, output shape {tuple(out.shape)}"
|
|
except Exception as e:
|
|
return False, f"forward failed: {type(e).__name__}: {e}"
|
|
|
|
|
|
def test_sdpa():
|
|
try:
|
|
device = "cuda"
|
|
dtype = torch.bfloat16
|
|
batch, nheads, seqlen, headdim = 2, 8, 128, 64
|
|
q = torch.randn(batch, nheads, seqlen, headdim, device=device, dtype=dtype)
|
|
k = torch.randn(batch, nheads, seqlen, headdim, device=device, dtype=dtype)
|
|
v = torch.randn(batch, nheads, seqlen, headdim, device=device, dtype=dtype)
|
|
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
|
|
torch.cuda.synchronize()
|
|
assert out.shape == (batch, nheads, seqlen, headdim)
|
|
return True, f"forward ok, output shape {tuple(out.shape)}"
|
|
except Exception as e:
|
|
return False, f"forward failed: {type(e).__name__}: {e}"
|
|
|
|
|
|
def test_flash_linear_attention(timeout_s=1200):
|
|
start = time.time()
|
|
try:
|
|
import fla # noqa: F401
|
|
from fla.ops.gated_delta_rule import chunk_gated_delta_rule
|
|
except ImportError as e:
|
|
return False, f"import failed (package not installed): {e}"
|
|
|
|
try:
|
|
device = "cuda"
|
|
dtype = torch.bfloat16
|
|
batch, seqlen, nheads, headdim = 2, 64, 4, 64
|
|
q = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=True)
|
|
k = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=True)
|
|
v = torch.randn(batch, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=True)
|
|
beta = torch.rand(batch, seqlen, nheads, device=device, dtype=dtype).sigmoid()
|
|
g = -torch.rand(batch, seqlen, nheads, device=device, dtype=torch.float32)
|
|
|
|
out, _ = chunk_gated_delta_rule(q, k, v, g, beta)
|
|
loss = out.sum()
|
|
loss.backward()
|
|
torch.cuda.synchronize()
|
|
|
|
if time.time() - start > timeout_s:
|
|
return False, "compiled/ran but exceeded time budget"
|
|
return True, f"forward+backward ok, output shape {tuple(out.shape)}"
|
|
except Exception as e:
|
|
elapsed = time.time() - start
|
|
return False, f"failed after {elapsed:.0f}s: {type(e).__name__}: {e}"
|
|
|
|
|
|
def main():
|
|
if not torch.cuda.is_available():
|
|
print("CUDA not available — aborting.")
|
|
sys.exit(1)
|
|
|
|
print(f"torch {torch.__version__}, CUDA {torch.version.cuda}, device: {torch.cuda.get_device_name(0)}")
|
|
print()
|
|
|
|
ok, msg = test_flash_attn()
|
|
RESULTS["flash_attn"] = {"ok": ok, "detail": msg}
|
|
print(f"[flash_attn] {'OK' if ok else 'FAIL'} — {msg}")
|
|
|
|
ok, msg = test_sdpa()
|
|
RESULTS["sdpa"] = {"ok": ok, "detail": msg}
|
|
print(f"[sdpa] {'OK' if ok else 'FAIL'} — {msg}")
|
|
|
|
recommended = "flash_attention_2" if RESULTS["flash_attn"]["ok"] else "sdpa"
|
|
print(f"\n>>> attn_implementation recomendado para el training: {recommended}\n")
|
|
|
|
print("Probando flash-linear-attention (opcional, acotado)...")
|
|
ok, msg = test_flash_linear_attention()
|
|
RESULTS["flash_linear_attention"] = {"ok": ok, "detail": msg}
|
|
print(f"[flash-linear-attention] {'OK' if ok else 'FAIL (se descarta, no bloqueante)'} — {msg}")
|
|
|
|
print("\n=== Resumen ===")
|
|
for k, v in RESULTS.items():
|
|
print(f" {k}: {'OK' if v['ok'] else 'FAIL'} — {v['detail']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|