Update builder.sh

builder.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-echo "🚀 Builder (FlashAttn LayerNorm + Apex + Q8) — runtime com GPU visível"
+echo "🚀 Builder (FlashAttn LayerNorm extra + Apex + Q8) — runtime com GPU visível"
 
 # ===== Config e diretórios =====
 export SELF_HF_REPO_ID="${SELF_HF_REPO_ID:-euIaxs22/Aduc-sdr}"  # Repo no HF para wheels
@@ -15,7 +15,7 @@ mkdir -p /app/wheels /app/cuda_cache "$HF_HOME" "$TORCH_HOME" /app/wheels/src
 chmod -R 777 /app/wheels || true
 export CUDA_CACHE_PATH="/app/cuda_cache"
 
-# Preserve licença NGC em wheels (se presente no container base)
+# Preserva licença NGC (se existir)
 if [ -f "/NGC-DL-CONTAINER-LICENSE" ]; then
   cp -f /NGC-DL-CONTAINER-LICENSE /app/wheels/NGC-DL-CONTAINER-LICENSE || true
 fi
@@ -49,16 +49,23 @@ echo "[env] PY_TAG=${PY_TAG} TORCH_VER=${TORCH_VER} CU_TAG=${CU_TAG}"
 #                               CHECKERS
 # ============================================================================
 
-check_flashln () {
+# Checa especificamente o módulo nativo requerido pelo layer_norm (sem checar 'flash-attn' geral)
+check_flash_layer_norm_bin () {
 python - <<'PY'
-import importlib.util, pkgutil
-cands = [
-    "flash_attn_layer_norm",     # empacotamentos recentes
-    "dropout_layer_norm",        # nome histórico do módulo C++
-    "flash_attn.ops.layer_norm", # caminho dentro do pacote flash_attn
-]
-ok = any(importlib.util.find_spec(n) is not None for n in cands)
-ok = ok or any(("flash" in m.name and "norm" in m.name) for m in pkgutil.iter_modules())
+import importlib
+ok = False
+# extensões conhecidas produzidas por csrc/layer_norm
+for name in [
+    "dropout_layer_norm",                 # nome do módulo nativo
+    "flash_attn.ops.layer_norm",          # wrapper python que usa o nativo
+    "flash_attn.ops.rms_norm",            # pode depender do mesmo backend em alguns empacotamentos
+]:
+    try:
+        importlib.import_module(name)
+        ok = True
+        break
+    except Exception:
+        pass
 raise SystemExit(0 if ok else 1)
 PY
 }
@@ -66,7 +73,7 @@ PY
 check_apex () {
 python - <<'PY'
 try:
-    from apex.normalization import FusedLayerNorm, FusedRMSNorm
+    from apex.normalization import FusedLayerNorm
     import importlib; importlib.import_module("fused_layer_norm_cuda")
     ok = True
 except Exception:
@@ -87,13 +94,12 @@ PY
 #                         DOWNLOAD DO HUB (GENÉRICO)
 # ============================================================================
 
-# Instala uma wheel do HF por padrão de nome (prefixo) e tags.
-# Uso: install_from_hf_by_prefix <prefixo>
+# Instala uma wheel do HF por prefixo simples (ex.: apex-, q8_kernels-)
 install_from_hf_by_prefix () {
   local PREFIX="$1"
   echo "[hub] Procurando wheels '${PREFIX}-*.whl' em ${SELF_HF_REPO_ID} com tags ${PY_TAG}/${CU_TAG}"
   python - "$PREFIX" "$PY_TAG" "$CU_TAG" <<'PY' || exit 0
-import os, sys, re
+import os, sys
 from huggingface_hub import HfApi, hf_hub_download, HfFolder
 
 prefix, py_tag, cu_tag = sys.argv[1], sys.argv[2], sys.argv[3]
@@ -104,18 +110,14 @@ try:
 except Exception:
     raise SystemExit(0)
 
-def match_tags(name: str) -> bool:
-    # Requer Python tag; CUDA tag preferencial se existir na wheel
-    if py_tag not in name:
-        return False
-    return True
+def match(name: str) -> bool:
+    return name.endswith(".whl") and name.rsplit("/",1)[-1].startswith(prefix + "-") and (py_tag in name)
 
-cands = [f for f in files if f.endswith(".whl") and f.rsplit("/",1)[-1].startswith(prefix + "-") and match_tags(f)]
+cands = [f for f in files if match(f)]
 pref = [f for f in cands if cu_tag and cu_tag in f] or cands
 if not pref:
     raise SystemExit(0)
 
-# Seleciona o "maior" por ordenação lexicográfica como heurística de versão/compat
 target = sorted(pref, reverse=True)[0]
 print(target)
 path = hf_hub_download(repo_id=repo, filename=target, repo_type="model", local_dir="/app/wheels")
@@ -123,9 +125,9 @@ print(path)
 PY
 }
 
-# Instala uma wheel do HF por padrão flexível para FlashAttn LayerNorm (aceita várias variantes de nome)
-install_flashln_from_hf () {
-  echo "[hub] Procurando wheels de FlashAttention LayerNorm em ${SELF_HF_REPO_ID}"
+# Instala wheels do submódulo layer_norm aceitando variantes de nome
+install_flash_layer_norm_from_hf () {
+  echo "[hub] Procurando wheels FlashAttention LayerNorm em ${SELF_HF_REPO_ID}"
   python - "$PY_TAG" "$CU_TAG" <<'PY' || exit 0
 import os, sys, re
 from huggingface_hub import HfApi, hf_hub_download, HfFolder
@@ -138,15 +140,13 @@ try:
 except Exception:
     raise SystemExit(0)
 
+pats = [
+    r"^flash[_-]?attn[_-]?.*layer[_-]?norm-.*\.whl$",
+    r"^dropout[_-]?layer[_-]?norm-.*\.whl$",
+]
 def ok(fn: str) -> bool:
     name = fn.rsplit("/",1)[-1]
-    if not name.endswith(".whl"): return False
     if py_tag not in name: return False
-    # Padrões comuns de empacotamento
-    pats = [
-        r"^flash[_-]?attn[_-]?.*layer[_-]?norm-",
-        r"^dropout[_-]?layer[_-]?norm-",
-    ]
     return any(re.search(p, name, flags=re.I) for p in pats)
 
 cands = [f for f in files if ok(f)]
@@ -165,11 +165,29 @@ PY
 #                                BUILDERS
 # ============================================================================
 
-build_flashln () {
+# Passo extra: SIEMPRE tenta instalar o submódulo layer_norm via wheel do HF;
+# se não houver wheel compatível, compila a partir de csrc/layer_norm e gera wheel.
+build_or_install_flash_layer_norm () {
+  echo "[flow] === FlashAttn LayerNorm (passo extra) ==="
+
+  # 1) Tentar instalar wheel do HF primeiro (evita recompilar)
+  HF_OUT="$(install_flash_layer_norm_from_hf || true)"
+  if [ -n "${HF_OUT:-}" ]; then
+    WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
+    echo "[hub] Baixado: ${WHEEL_PATH}"
+    python -m pip install -v -U --no-build-isolation --no-deps "${WHEEL_PATH}" || true
+    if check_flash_layer_norm_bin; then
+      echo "[flow] FlashAttn LayerNorm: OK via wheel do Hub"
+      return 0
+    fi
+    echo "[flow] Wheel do Hub não resolveu import; seguirá com build"
+  else
+    echo "[hub] Nenhuma wheel compatível encontrada para FlashAttn LayerNorm"
+  fi
+
+  # 2) Build from source do submódulo csrc/layer_norm -> wheel
   local SRC="/app/wheels/src/flash-attn"
   echo "[build] Preparando fonte FlashAttention (layer_norm) em ${SRC}"
-
-  # Sync/clone
   if [ -d "$SRC/.git" ]; then
     git -C "$SRC" fetch --all -p || true
     git -C "$SRC" reset --hard origin/main || true
@@ -179,41 +197,49 @@ build_flashln () {
     git clone --depth 1 https://github.com/Dao-AILab/flash-attention "$SRC"
   fi
 
-  # Arquitetura CUDA alvo baseada na GPU ativa (uma só CC por build para wheel específica da máquina)
+  # Define CC alvo a partir da GPU ativa (reduz tempo/ruído de build)
   export TORCH_CUDA_ARCH_LIST="$(python - <<'PY'
 import torch
 try:
-    major, minor = torch.cuda.get_device_capability(0)
-    print(f"{major}.{minor}")
+    cc = "%d.%d" % torch.cuda.get_device_capability(0)
+    print(cc)
 except Exception:
-    # fallback genérico caso não haja GPU visível no build host
-    print("8.9")
+    print("8.9")  # fallback p/ Ada (L40S) caso build sem GPU visível
 PY
   )"
   echo "[build] TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}"
 
-  # Compila apenas o submódulo csrc/layer_norm para wheel
   pushd "$SRC/csrc/layer_norm" >/dev/null
   export MAX_JOBS="${MAX_JOBS:-90}"
+  # Gera wheel reutilizável
   python -m pip wheel -v --no-build-isolation --no-deps . -w /app/wheels || true
   popd >/dev/null
 
-  # Instala a wheel recém-criada
+  # Instala a wheel gerada
   local W="$(ls -t /app/wheels/*flash*attn*layer*norm*-*.whl 2>/dev/null | head -n1 || true)"
   if [ -z "${W}" ]; then
     W="$(ls -t /app/wheels/*dropout*layer*norm*-*.whl 2>/dev/null | head -n1 || true)"
   fi
   if [ -z "${W}" ]; then
-    # fallback para qualquer wheel produzida no diretório
+    # fallback para qualquer .whl recém gerado
     W="$(ls -t /app/wheels/*.whl 2>/dev/null | head -n1 || true)"
   fi
+
   if [ -n "${W}" ]; then
     python -m pip install -v -U --no-deps "${W}" || true
-    echo "[build] FlashAttention LayerNorm instalado da wheel: ${W}"
+    echo "[build] FlashAttn LayerNorm instalado da wheel: ${W}"
   else
-    echo "[build] Nenhuma wheel gerada; instalando do source (fallback)"
+    echo "[build] Nenhuma wheel gerada; instalando direto do source (último recurso)"
     python -m pip install -v --no-build-isolation "$SRC/csrc/layer_norm" || true
   fi
+
+  # Checagem final do binário
+  if check_flash_layer_norm_bin; then
+    echo "[flow] FlashAttn LayerNorm: import OK após build"
+    return 0
+  fi
+  echo "[flow] FlashAttn LayerNorm: falhou import após build"
+  return 1
 }
 
 build_apex () {
@@ -235,7 +261,7 @@ build_apex () {
     python -m pip install -v -U --no-deps "${W}" || true
     echo "[build] Apex instalado da wheel recém-compilada: ${W}"
   else
-    echo "[build] Nenhuma wheel Apex gerada; instalando do source (pode falhar)"
+    echo "[build] Nenhuma wheel Apex gerada; instalando do source"
     python -m pip install -v --no-build-isolation "$SRC" || true
   fi
 }
@@ -255,94 +281,52 @@ build_q8 () {
     python -m pip install -v -U --no-deps "${W}" || true
     echo "[build] Q8 instalado da wheel recém-compilada: ${W}"
   else
-    echo "[build] Nenhuma wheel q8_kernels gerada; instalando do source (pode falhar)"
+    echo "[build] Nenhuma wheel q8_kernels gerada; instalando do source"
     python -m pip install -v --no-build-isolation "$SRC" || true
   fi
 }
 
 # ============================================================================
-#                          PIPELINES DE GARANTIA
+#                               EXECUÇÃO
 # ============================================================================
 
-ensure_pkg () {
-  local PKG="$1"      # apex | q8_kernels
-  local CHECK_FN="$2" # check_apex | check_q8
-  local BUILD_FN="$3" # build_apex | build_q8
-  echo "[flow] === ${PKG} ==="
-  if ${CHECK_FN}; then
-    echo "[flow] ${PKG}: já instalado (import OK)"
-    return 0
-  fi
-  echo "[flow] ${PKG}: tentando wheel do Hub (${SELF_HF_REPO_ID})"
-  HF_OUT="$(install_from_hf_by_prefix "$PKG" || true)"
-  if [ -n "${HF_OUT:-}" ]; then
-    WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
-    echo "[hub] Baixado: ${WHEEL_PATH}"
-    python -m pip install -v -U --no-build-isolation "${WHEEL_PATH}" || true
-    if ${CHECK_FN}; then
-      echo "[flow] ${PKG}: sucesso via Hub (${WHEEL_PATH})"
-      return 0
-    else
-      echo "[flow] ${PKG}: import falhou após wheel do Hub; compilando"
-    fi
-  else
-    echo "[hub] Nenhuma wheel compatível encontrada para ${PKG}"
+# Passo adicional SEM depender de "flash-attn" já instalado: trata somente o layer_norm
+build_or_install_flash_layer_norm || true
+
+# Apex (mantido)
+# Tenta primeiro via wheel no HF e, se não houver, compila e instala em wheel
+echo "[flow] === apex ==="
+HF_OUT="$(install_from_hf_by_prefix "apex" || true)"
+if [ -n "${HF_OUT:-}" ]; then
+  WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
+  echo "[hub] Baixado: ${WHEEL_PATH}"
+  python -m pip install -v -U --no-build-isolation "${WHEEL_PATH}" || true
+  if ! check_apex; then
+    echo "[flow] apex: import falhou após wheel; compilando"
+    build_apex || true
   fi
-  echo "[flow] ${PKG}: compilando (fallback)"
-  ${BUILD_FN}
-  if ${CHECK_FN}; then
-    echo "[flow] ${PKG}: sucesso após compilação"
-    return 0
-  fi
-  echo "[flow] ${PKG}: falhou após build; seguindo"
-  return 1
-}
-
-ensure_flashln () {
-  echo "[flow] === flash_attn_layer_norm ==="
-  if check_flashln; then
-    echo "[flow] FlashAttn LayerNorm: já instalado (import OK)"
-    return 0
-  fi
-  echo "[flow] FlashAttn LayerNorm: tentando wheel do Hub (${SELF_HF_REPO_ID})"
-  HF_OUT="$(install_flashln_from_hf || true)"
-  if [ -n "${HF_OUT:-}" ]; then
-    WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
-    echo "[hub] Baixado: ${WHEEL_PATH}"
-    python -m pip install -v -U --no-build-isolation "${WHEEL_PATH}" || true
-    if check_flashln; then
-      echo "[flow] FlashAttn LayerNorm: sucesso via Hub (${WHEEL_PATH})"
-      return 0
-    else
-      echo "[flow] FlashAttn LayerNorm: import falhou após wheel do Hub; compilando"
-    fi
-  else
-    echo "[hub] Nenhuma wheel compatível encontrada para FlashAttn LayerNorm"
-  fi
-  echo "[flow] FlashAttn LayerNorm: compilando (fallback)"
-  build_flashln
-  if check_flashln; then
-    echo "[flow] FlashAttn LayerNorm: sucesso após compilação"
-    return 0
-  fi
-  echo "[flow] FlashAttn LayerNorm: falhou após build; seguindo"
-  return 1
-}
-
-# ============================================================================
-#                                 EXECUÇÃO
-# ============================================================================
-
-# 1) FlashAttention LayerNorm (novo cache em wheel no HF)
-ensure_flashln || true
-
-# 2) Apex (mantém pipeline existente)
-ensure_pkg "apex" check_apex build_apex || true
-
-# 3) Q8 kernels (opcional)
-# ensure_pkg "q8_kernels" check_q8 build_q8 || true
+else
+  echo "[hub] Nenhuma wheel apex compatível; compilando"
+  build_apex || true
+fi
 
-# 4) Upload de wheels produzidas para o HF (cache cross-restarts)
+# Q8 (opcional)
+# echo "[flow] === q8_kernels ==="
+# HF_OUT="$(install_from_hf_by_prefix "q8_kernels" || true)"
+# if [ -n "${HF_OUT:-}" ]; then
+#   WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
+#   echo "[hub] Baixado: ${WHEEL_PATH}"
+#   python -m pip install -v -U --no-build-isolation "${WHEEL_PATH}" || true
+#   if ! check_q8; then
+#     echo "[flow] q8_kernels: import falhou após wheel; compilando"
+#     build_q8 || true
+#   fi
+# else
+#   echo "[hub] Nenhuma wheel q8_kernels compatível; compilando"
+#   build_q8 || true
+# fi
+
+# Upload de wheels produzidas para o HF (cache cross-restarts)
 python - <<'PY'
 import os
 from huggingface_hub import HfApi, HfFolder