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.
49 lines
1.9 KiB
Python
49 lines
1.9 KiB
Python
"""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()
|