From d121c4465ef6d15eabc6efdb02b7a53242f8f56e Mon Sep 17 00:00:00 2001 From: Alejandro Lembke Barrientos Date: Tue, 28 Jul 2026 03:52:10 +0000 Subject: [PATCH] Fase 0: scripts de verificacion de hardware e inspeccion de modulos 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. --- constraints.txt | 1 + scripts/00_verify_hardware.py | 112 ++++++++++++++++++++++++++++++++++ scripts/01_inspect_modules.py | 48 +++++++++++++++ 3 files changed, 161 insertions(+) create mode 100644 constraints.txt create mode 100644 scripts/00_verify_hardware.py create mode 100644 scripts/01_inspect_modules.py diff --git a/constraints.txt b/constraints.txt new file mode 100644 index 0000000..de37aaf --- /dev/null +++ b/constraints.txt @@ -0,0 +1 @@ +torch==2.10.0a0+b4e4ee81d3.nv25.12 diff --git a/scripts/00_verify_hardware.py b/scripts/00_verify_hardware.py new file mode 100644 index 0000000..92c21ae --- /dev/null +++ b/scripts/00_verify_hardware.py @@ -0,0 +1,112 @@ +"""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() diff --git a/scripts/01_inspect_modules.py b/scripts/01_inspect_modules.py new file mode 100644 index 0000000..4fd9f87 --- /dev/null +++ b/scripts/01_inspect_modules.py @@ -0,0 +1,48 @@ +"""Fase 0: resuelve la discrepancia de nombres reales de proyecciones de Gated DeltaNet +(in_proj_qkvz/in_proj_ba fusionadas vs in_proj_qkv/in_proj_z/in_proj_a/in_proj_b separadas) +inspeccionando la arquitectura real (device_map="meta", sin cargar pesos) de una capa +linear_attention (layer 0) y una full_attention (layer 3). +""" +import re +import sys + +import torch +from transformers import AutoConfig, AutoModelForCausalLM + +MODEL_PATH = sys.argv[1] if len(sys.argv) > 1 else "/workspace/ft-models/Qwen--Qwen3.6-35B-A3B" + + +def main(): + config = AutoConfig.from_pretrained(MODEL_PATH) + layer_types = config.text_config.layer_types + print(f"num_hidden_layers: {len(layer_types)}") + print(f"layer_types[:8]: {layer_types[:8]}") + + linear_idx = layer_types.index("linear_attention") + full_idx = layer_types.index("full_attention") + print(f"\nUsando layer {linear_idx} (linear_attention) y layer {full_idx} (full_attention)\n") + + with torch.device("meta"): + model = AutoModelForCausalLM.from_config(config) + + for label, idx in [("linear_attention", linear_idx), ("full_attention", full_idx)]: + print(f"=== capa {idx} ({label}) ===") + pattern = re.compile(rf"\.layers\.{idx}\.") + found = False + for name, param in model.named_parameters(): + if pattern.search(name): + found = True + print(f" {name} {tuple(param.shape)}") + if not found: + print(f" (no se encontraron params para layers.{idx} — revisar prefijo real del modelo)") + print() + + print("=== nombres de nn.Module (no solo parámetros) para la capa linear_attention ===") + pattern = re.compile(rf"\.layers\.{linear_idx}\.") + for name, module in model.named_modules(): + if pattern.search(name) and list(module.children()) == []: + print(f" {name} ({type(module).__name__})") + + +if __name__ == "__main__": + main()