141 lines
20 KiB
Markdown
141 lines
20 KiB
Markdown
# Fine-tuning LoRA de Qwen3.6-35B-A3B para MCPs y skills de Claude Code
|
|
|
|
## Contexto
|
|
|
|
El usuario sirve `RedHatAI/Qwen3.6-35B-A3B-NVFP4` en producción vía vLLM en su DGX Spark, y quiere un LoRA que le enseñe al modelo a:
|
|
1. Usar correctamente el MCP de Penpot (diseño) y los otros 4 MCPs conectados (gitea, github-personal, docmost, atlassian).
|
|
2. Seguir sus "skills" de Claude Code (en español) sin saltarse pasos ni reglas — su queja principal es que el modelo olvida reglas/pasos al ejecutar tareas largas.
|
|
|
|
Se investigó a fondo antes de diseñar: se leyó el `chat_template.jinja` real del modelo (7764 bytes), se hizo SSH a spark para verificar hardware/disco/contenedores en vivo, se documentaron a fondo las 5 tools del Penpot MCP (arquitectura, schemas, reglas no-obvias), se inventariaron los 5 skills reales y los 5 MCP servers, y se lanzaron dos diseños independientes (uno priorizando un framework turnkey, otro priorizando control total con `transformers`+`peft`+`trl`) que convergieron en casi todos los puntos técnicos. Este plan sintetiza ambos.
|
|
|
|
**Resultado esperado:** un adapter LoRA entrenado, mergeado y re-cuantizado a NVFP4, que reemplaza el checkpoint actual en `~/models/` en spark como *drop-in* — sin tocar el `docker-compose.yaml` de producción.
|
|
|
|
---
|
|
|
|
## Hallazgos que determinan el diseño
|
|
|
|
- **El modelo NO es un MoE clásico.** `config.json` de `Qwen/Qwen3.6-35B-A3B` declara `model_type: "qwen3_5_moe"`, `architectures: ["Qwen3_5MoeForConditionalGeneration"]` — ya soportado en `transformers` ≥5.2 (verificado, no requiere `trust_remote_code`). De 40 capas: **30 son `linear_attention`** (Gated DeltaNet) y solo **10 son `full_attention`** (`full_attention_interval: 4`). Además tiene 256 expertos ruteados (top-8) + 1 shared expert siempre-activo, y un vision tower (SigLIP, ~27 bloques).
|
|
- **Los expertos ruteados son `nn.Parameter` 3D, no `nn.Linear`.** PEFT/bitsandbytes no los puede envolver como LoRA/QLoRA estándar. Son ~32B de los ~35B params totales (91%).
|
|
- **GB10 (SM121) no tiene kernels de linear attention** (`causal_conv1d`/`fla` sin build sm121 según la doc de HF) → Gated DeltaNet corre en el fallback lento de PyTorch. Es el cuello de botella de velocidad y hay que aceptarlo (vale la pena un experimento acotado de 20 min probando `flash-linear-attention` por si el build sí funciona vía Triton JIT, ya que es sospechoso que una librería sin extensiones CUDA "no tenga build" — pero no bloqueante).
|
|
- **Existe precedente casi idéntico**: [kreuzhofer/dgx-spark-unsloth-qwen3.5-training](https://github.com/kreuzhofer/dgx-spark-unsloth-qwen3.5-training) hace LoRA BF16 de Qwen3.5-35B-A3B (misma arquitectura) en un solo DGX Spark, y NVIDIA publica un playbook oficial de Unsloth para DGX Spark.
|
|
- **Existe la receta exacta de cuantización NVFP4** usada por RedHatAI (`recipe.yaml` del repo): `QuantizationModifier(scheme="NVFP4", ignore=["re:.*lm_head","re:visual.*","re:model.visual.*","re:.*mlp.gate$","re:.*embed_tokens$","re:.*shared_expert_gate$","re:.*linear_attn.*"])`. Nota clave: **las capas Gated DeltaNet (`linear_attn`) quedan en BF16 en producción** — si el LoRA se aplica ahí, el merge no pierde precisión de cuantización.
|
|
- **El chat template usa formato XML tipo Hermes para tool calls** (`<tool_call><function=X><parameter=Y>valor</parameter></function></tool_call>`, NO JSON), con `role: "tool"` como rol de primera clase que el template convierte a `<tool_response>`, y `reasoning_content` como campo separado (no meter `<think>` a mano en `content`). Con `preserve_thinking=true` (como en producción), todos los turnos assistant previos preservan su `<think>` al re-serializar.
|
|
- **Los transcripts reales de Claude Code (`~/.claude/projects/*/*.jsonl`) están casi vacíos de uso de MCPs** (10 archivos, 735 líneas, 0 de 216 `tool_use` son `mcp__*`) — no sirven como fuente principal. En cambio `~/.claude/plans/*.md` (28 archivos, ~200KB) sí tiene español real del usuario y es útil para estilo/plantillas de reporte.
|
|
- **Discrepancia sin resolver entre las dos investigaciones**: no coinciden en si las proyecciones de Gated DeltaNet están fusionadas (`in_proj_qkvz`, `in_proj_ba`, por analogía con Qwen3-Next) o separadas (`in_proj_qkv`, `in_proj_z`, `in_proj_a`, `in_proj_b`, según lectura directa del índice de tensores real). **Esto se resuelve empíricamente en la Fase 0** inspeccionando el checkpoint real descargado — no se asume ninguna de las dos.
|
|
- **Hallazgo de seguridad, fuera del alcance de este plan pero a resolver antes/en paralelo**: el mismo token (`<<SPARK_PASSWORD_1>>`) está reusado como Bearer de gitea, github-personal y penpot MCP, y parece ser la misma password de spark. Rotar esos tokens y dejar de reusar la password de spark como bearer HTTP.
|
|
|
|
---
|
|
|
|
## Decisiones de diseño
|
|
|
|
1. **Framework**: stack primario `transformers` + `peft` + `trl` (control total, sin dependencia de que un framework de alto nivel ya soporte esta arquitectura tan nueva). Probar `unsloth` como acelerador opcional (puede dar 1.5-2x velocidad) una vez que el pipeline base funcione — no bloqueante, con fallback documentado al stack puro si falla la carga.
|
|
2. **Memoria de entrenamiento**: **LoRA BF16 puro, no QLoRA**. Los expertos ruteados (91% de los params) no son cuantizables por bnb, así que QLoRA ahorraría ~4-5% de memoria a cambio de fidelidad numérica y una dependencia con bugs abiertos. No vale la pena.
|
|
3. **Módulos objetivo de LoRA**: atención completa (`q/k/v/o_proj`, 10 capas) + proyecciones de Gated DeltaNet (30 capas, nombres exactos a confirmar en Fase 0) + `shared_expert.{gate,up,down}_proj` (40 capas — el único FFN que ve el 100% de los tokens). **Excluidos**: router (`mlp.gate`), expertos ruteados, vision tower. Esto cubre 40/40 capas para aprender formato/comportamiento sin tocar conocimiento factual ni arriesgar colapso de expertos. r=32, alpha=64, dropout=0.05.
|
|
4. **Chat template / masking**: construir el dataset con `messages` + `tool_calls` estructurados (formato TRL/OpenAI, `arguments` como dict) y **dejar que `apply_chat_template` genere el XML** — nunca escribirlo a mano. Verificar primero si TRL reconoce nativamente el chat template de Qwen3.6 y aplica el masking de loss vía `assistant_only_loss=True`; si no, fallback a copiar el `.jinja` con `{% generation %}...{% endgeneration %}` manual alrededor de la salida assistant.
|
|
5. **Ubicación de archivos**: todo el código (pipeline de dataset, scripts de training/merge/cuantización, docker-compose del entorno de training, configs, dataset generado) vive en **`~/projects/ai-projects/qwen3-6-lora/`** (= `/mnt/docker-nas/projects/ai-projects/qwen3-6-lora/` visto desde spark, mismo storage NFS) — una subcarpeta nueva dentro de `ai-projects`, no `ai-projects` en sí. Esa subcarpeta es la raíz de **un repo git nuevo** creado al arrancar la Fase 0 (`git init` dentro de `qwen3-6-lora/`, con `.gitignore` excluyendo binarios de modelo, `data/raw/` sin sanitizar, y cualquier checkpoint). El dataset final ya sanitizado (`data/train.jsonl`, `data/eval.jsonl`) y los schemas de MCP sí se versionan — son el activo más valioso del proyecto y lo que justifica tener historia de git. Los **binarios del modelo** (checkpoint base BF16 ~72GB, merge, cuantizado) se mantienen en disco **local rápido de spark** (no en el NFS ni en git — no tiene sentido versionar cientos de GB de pesos) en un directorio nuevo dedicado; solo el checkpoint NVFP4 final se copia a `~/models/` en spark (requisito del `docker-compose.yaml` de vLLM). Publicar el repo a gitea/github (con la skill `aleleba-pr`) queda como paso posterior, ya en fase de ejecución, no parte de este plan.
|
|
6. **Ventana operativa**: entrenar localmente en Spark, aceptando parar vLLM ~1-2 noches (confirmado con el usuario). No se renta GPU en la nube.
|
|
7. **Alcance**: solo texto por ahora (confirmado con el usuario) — el vision tower queda congelado; un bucket con capturas de Penpot puede agregarse después sin rehacer el pipeline.
|
|
|
|
---
|
|
|
|
## Estructura de directorios
|
|
|
|
```
|
|
~/projects/ai-projects/qwen3-6-lora/ # repo git nuevo (git init en Fase 0), subcarpeta de ai-projects
|
|
├── .gitignore # excluye data/raw/, out/*.safetensors, cualquier checkpoint
|
|
├── docker-compose.yml # contenedor de training (NGC pytorch, con volúmenes persistentes)
|
|
├── scripts/
|
|
│ ├── 00_verify_hardware.py # smoke tests: flash_attn, sdpa, fla (opcional), sm121
|
|
│ ├── 01_inspect_modules.py # resuelve la discrepancia de nombres in_proj_* contra el checkpoint real
|
|
│ ├── 02_dump_mcp_schemas.py # tools/list contra los 5 MCP servers -> data/schemas/*.json
|
|
│ ├── 03_build_replay.py # genera ~700 ejemplos de replay contra el vLLM de producción (antes de pararlo)
|
|
│ ├── 04_sanitize.py # scrubbing de secretos (gitleaks/detect-secrets + regex) sobre fuentes crudas
|
|
│ ├── 05_build_dataset.py # arma los buckets A-K -> data/train.jsonl, data/eval.jsonl
|
|
│ ├── 06_validate_dataset.py # render con apply_chat_template, assert de máscara de loss, filtro de longitud
|
|
│ ├── 10_train.py # LoRA con transformers+peft+trl (+ unsloth opcional)
|
|
│ ├── 20_merge_lora.py # merge streaming shard-a-shard, preserva tensores MTP
|
|
│ ├── 21_quantize_nvfp4.py # llm-compressor, receta clonada de RedHatAI, reinjerta MTP
|
|
│ └── 30_eval_suite.py # las 4 puertas de evaluación (parser real, checklists de skills, etc.)
|
|
├── data/
|
|
│ ├── schemas/ # {penpot,gitea,github-personal,docmost,atlassian}.json
|
|
│ ├── raw/ # fuentes sanitizadas (skills, agents, plans, replay)
|
|
│ ├── train.jsonl / eval.jsonl
|
|
│ └── chat_template_train.jinja # solo si hace falta el fallback manual de masking
|
|
└── out/ # adapters, logs de eval (NO los checkpoints de modelo completo)
|
|
```
|
|
|
|
En spark (fuera del NFS, disco local rápido — confirmar ruta exacta en Fase 0, p.ej. `/home/aleleba/ft-models/`):
|
|
```
|
|
Qwen--Qwen3.6-35B-A3B/ # checkpoint base BF16 descargado de HF (~72GB)
|
|
Qwen3.6-35B-A3B-mcp-bf16/ # merge del LoRA (~72GB)
|
|
Qwen3.6-35B-A3B-mcp-NVFP4/ # resultado final, se copia a ~/models/ para servir
|
|
```
|
|
|
|
---
|
|
|
|
## Fases de implementación
|
|
|
|
### Fase 0 — Infraestructura y verificación (sin tocar producción)
|
|
|
|
0. Crear la subcarpeta `~/projects/ai-projects/qwen3-6-lora/` y hacer `git init` ahí dentro (no en `ai-projects` directamente), con el `.gitignore` de arriba, y un primer commit con la estructura base de carpetas. Todo el código/scripts/dataset sanitizado de las fases siguientes se commitea a medida que se produce.
|
|
1. El contenedor `jupyter-pyt` que ya está corriendo en spark (imagen `nvcr.io/nvidia/pytorch:25.12-py3`) **no tiene volumen persistente** — no lo parches in-place. Crear un servicio nuevo (`docker-compose.yml` de arriba) con bind mounts a `~/projects/ai-projects` y al directorio local de modelos, misma imagen base (trae `flash_attn` y `nvidia-modelopt` preinstalados).
|
|
2. Instalar `transformers>=5.2` (verificar versión exacta que declare soporte de `qwen3_5_moe`), `peft`, `trl`, `accelerate`, `datasets`, `bitsandbytes` (solo para `adamw_8bit`), con `PIP_CONSTRAINT` fijando la versión de `torch` de la imagen NGC (`--no-deps` en transformers) para no romper el build ARM64/Blackwell.
|
|
3. `scripts/00_verify_hardware.py`: probar `flash_attn` real (forward simple), y si falla usar `attn_implementation="sdpa"`. Probar (opcional, 20 min) si `flash-linear-attention` compila y da backward en SM121 — usarlo si funciona, descartarlo si no.
|
|
4. Descargar el checkpoint base: `hf download Qwen/Qwen3.6-35B-A3B` (~72GB). Verificar que `chat_template.jinja` pese igual (7764 bytes) que el de producción — confirma que son el mismo template.
|
|
5. **`scripts/01_inspect_modules.py`** (resuelve la discrepancia): cargar el modelo con `device_map="meta"` e imprimir todos los `nn.Linear`/`nn.Parameter` de una capa `linear_attention` (p.ej. layer 0) y una `full_attention` (p.ej. layer 3). Esto define los nombres reales a usar en `target_modules` — no proceder a la Fase 3 sin este dato confirmado.
|
|
|
|
### Fase 1 — Extracción de schemas y datos de replay (vLLM sigue corriendo)
|
|
|
|
6. `scripts/02_dump_mcp_schemas.py`: conectar a los 5 MCP servers (`mcp__penpot__*`, `gitea`, `github-personal`, `docmost`, `atlassian`) y volcar `tools/list` real a JSON. Esto confirma también si el deployment remoto de Penpot tiene filesystem deshabilitado (si `import_image` no aparece, o `export_shape` pierde `filePath`, hay que entrenar solo sobre las tools que sí existen — no inventar).
|
|
7. `scripts/03_build_replay.py`: generar ~700 ejemplos de "replay" (anti-forgetting) pidiéndole al **modelo en producción actual** (antes de pararlo) que responda prompts genéricos en español/código/razonamiento con `preserve_thinking=true`. Esto ancla el replay a la distribución real del modelo base, mejor que mezclar un corpus externo.
|
|
|
|
### Fase 2 — Construcción del dataset
|
|
|
|
8. `scripts/04_sanitize.py`: gate de scrubbing de secretos sobre las 5 SKILL.md, los 7 agents, los `~/.claude/plans/*.md`, y cualquier transcript usado — regex + `gitleaks`/`detect-secrets`, sustituyendo (no borrando) tokens/passwords/IPs/emails por placeholders estables. El build debe fallar si sobrevive algún patrón conocido.
|
|
9. `scripts/05_build_dataset.py`: construir ~2500-3000 ejemplos en buckets: Penpot MCP (~300, cubriendo cada regla no-obvia documentada: `insertChild` no `appendChild`, orden invertido de `children` en flex, re-fijar `growType` tras `resize()`, `storage` entre llamadas, no loguear lo que se retorna, leer `"Tool execution failed: ..."` como texto y autocorregirse), los otros 4 MCPs (~800 combinados), adherencia a skills — el bucket más importante para la queja principal del usuario, con checklist explícito de reglas en el `<think>` inicial, auto-corrección a mitad de trayectoria, fidelidad de la plantilla de "Reporte final", y escenarios que tientan a violar cada regla `NUNCA/SIEMPRE/OBLIGATORIO` real de los 5 skills (~450), delegación a subagentes (~180), negativos/no-tool (~200), manejo de errores (~150), y el replay de la Fase 1 (~700, ~25% del total). Dejar 1 de las 5 skills completamente fuera del training como held-out de generalización.
|
|
10. `scripts/06_validate_dataset.py`: renderizar cada ejemplo con `apply_chat_template(tools=..., tokenize=False)` y verificar que no lanza excepción, que el texto contiene `<function=...>` bien formado, y que la máscara de `assistant_masks` (con `return_assistant_tokens_mask=True`) no está vacía y no cubre system/tool/user. Filtrar (no truncar) ejemplos que excedan la longitud máxima (8192, con ventana deslizante para skills largas).
|
|
|
|
### Fase 3 — Entrenamiento
|
|
|
|
11. Configurar `LoraConfig` con los `target_modules` confirmados en la Fase 0 (attention + GDN + shared_expert), `r=32, alpha=64, dropout=0.05`. Verificar con `model.print_trainable_parameters()` (~0.11% esperado) y un conteo por módulo antes de arrancar cualquier run largo.
|
|
12. `docker compose stop vllm` (no `down`, conserva volúmenes). Dry run de 20 pasos: confirmar memoria pico <90GB, loss bajando. Luego run completo (2 épocas, LR 1e-4 cosine, gradient checkpointing, `save_steps=50`). Mitigar el riesgo de OOM al cargar el checkpoint (mmap + tensores CUDA compitiendo por la misma pool unificada) con `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` y purgando page cache antes de cargar.
|
|
13. `docker compose start vllm` en cuanto termine el training — producción vuelve, el adapter queda en disco.
|
|
|
|
### Fase 4 — Merge y evaluación
|
|
|
|
14. `scripts/20_merge_lora.py`: merge shard-a-shard (streaming, pico de memoria ~2GB en vez de cargar el modelo completo dos veces) del adapter sobre el checkpoint BF16. Preservar los tensores `mtp.*` (que `from_pretrained` descarta al cargar) copiándolos del checkpoint original — si se pierden, `--speculative-config` de producción no arranca.
|
|
15. Evaluación en 4 puertas, sin pasar a la siguiente sin cerrar la anterior: (1) eval-loss offline por bucket, especialmente el bucket de replay separado del resto; (2) validez de tool-calls usando el **parser real de producción** (`qwen3_coder`/`qwen3` de vLLM, no una regex propia) contra ~200 prompts held-out; (3) checklists de adherencia a reglas por skill (baseline medido con el modelo actual antes de comparar) + test de no-activación (frases cercanas a triggers que no deben disparar la skill); (4) prueba en vivo sirviendo el **BF16 mergeado directo** (no `--enable-lora` — hay incompatibilidades de nombres de módulos entre PEFT y vLLM para esta arquitectura, documentadas como riesgo abierto) contra los 5 MCP reales y las 5 skills de punta a punta.
|
|
16. **Puerta de decisión**: si no mejora sobre el baseline en la puerta 3, volver a la Fase 2 antes de gastar tiempo en cuantizar.
|
|
|
|
### Fase 5 — Re-cuantización y despliegue
|
|
|
|
17. `scripts/21_quantize_nvfp4.py`: clonar exactamente la receta de RedHatAI (`llm-compressor`, mismo `ignore` list) sobre el modelo mergeado, con `moe_calibrate_all_experts=True` (obligatorio, si no la mayoría de los 256 expertos quedan sin calibrar) e idealmente calibrando con una muestra del propio dataset de fine-tuning en vez de un corpus genérico.
|
|
18. Reinjertar tensores MTP en el checkpoint cuantizado. Copiar `chat_template.jinja` original (no el de training).
|
|
19. Mover el checkpoint actual a `.bak` (35GB, sobra espacio) antes de reemplazar — rollback = revertir el nombre y `docker compose up -d`. Copiar el nuevo checkpoint a `~/models/` en spark.
|
|
20. Arrancar vLLM con la config de producción **completa y sin modificar**, incluido `--speculative-config`. Si no arranca, el problema son los tensores MTP. Si arranca, correr el smoke test de las 5 skills contra los MCPs reales una vez más.
|
|
|
|
---
|
|
|
|
## Riesgos principales
|
|
|
|
| Riesgo | Mitigación |
|
|
|---|---|
|
|
| Copiar `target_modules` mal (nombres fusionados vs separados) → adapter vacío, run de horas perdido | Fase 0 verifica empíricamente antes de configurar LoRA; assert de `print_trainable_parameters()` antes de cada run |
|
|
| Masking de loss mal hecho (entrenar sobre system/tool_response) | Validación obligatoria de la máscara sobre 20 ejemplos antes de arrancar (Fase 2, paso 10) |
|
|
| Cuelgue del host por OOM en memoria unificada (no da error limpio) | `expandable_segments:True`, purga de page cache, considerar `systemd-run --scope -p MemoryMax` |
|
|
| Pérdida de tensores MTP en el merge → rompe `--speculative-config` en producción | Merge streaming los preserva por diseño; verificación explícita antes de desplegar |
|
|
| `--enable-lora` no soporta bien esta arquitectura en vLLM (issue abierto) | No usarlo como validación principal; servir el BF16 mergeado directo para evaluar |
|
|
| Fuga de secretos al dataset (tokens/passwords reales en skills/plans/transcripts) | Gate de sanitización que falla el build si sobrevive algún patrón; rotar los tokens reusados de una vez |
|
|
| Overfitting a 5 skills / tool-call spam | Skill held-out, ~25% replay auto-destilado, LR conservador, medir no-activación |
|
|
| Entrenamiento lento (fallback PyTorch de Gated DeltaNet) | Aceptar 10-15h como run nocturno; experimento opcional con `flash-linear-attention` |
|
|
|
|
---
|
|
|
|
## Verificación end-to-end
|
|
|
|
1. Fase 0 completa cuando `01_inspect_modules.py` confirma los nombres reales y `00_verify_hardware.py` corre sin excepciones.
|
|
2. Fase 2 completa cuando `06_validate_dataset.py` pasa sobre el 100% del dataset (render limpio + máscara no vacía) y el gate de secretos no encuentra nada.
|
|
3. Fase 3 completa cuando el dry-run de 20 pasos no OOMea y el loss baja de forma consistente en el run completo.
|
|
4. Fase 4: comparar explícitamente las 4 puertas contra el baseline del modelo actual (medido *antes* de entrenar) — la mejora debe ser medible, no asumida. La prueba real de éxito es servir el BF16 mergeado y ejecutar de punta a punta al menos una tarea real por cada uno de los 5 MCPs y las 5 skills contra el propio Claude Code apuntando a ese endpoint.
|
|
5. Fase 5: `docker compose up -d` de producción arranca sin errores con la config completa (incluido MTP/speculative decoding) sirviendo el checkpoint nuevo, con el `.bak` disponible para rollback inmediato.
|