config_nvfp4.json

{
  "_comment": "JoyImageEdit 16B DiT 量化配置 — 激进策略：边界层BF16，注意力FP8，MLP用NVFP4",
  "format": "1.0",
  "block_names": ["double_blocks", "condition_embedder", "proj_out"],

  "rules": [
    {
      "_comment": "输入/输出边界层 + 时间/文本嵌入 → BF16，不做量化（参数少但敏感）",
      "match": ["condition_embedder", "proj_out"],
      "policy": null
    },
    {
      "_comment": "注意力 QKV + proj 权重 → FP8（注意力对量化噪声更敏感，保留更多精度）",
      "match": ["img_attn_qkv", "img_attn_proj", "txt_attn_qkv", "txt_attn_proj"],
      "policy": "float8_e4m3fn"
    },
    {
      "_comment": "所有 MLP 权重 → NVFP4（参数量的主体，约60-70%，激进4bit量化）",
      "match": ["img_mlp", "txt_mlp"],
      "policy": "nvfp4"
    }
  ],

  "_tensor_coverage": {
    "description": "各层量化策略汇总",
    "img_in.weight": "5D Conv3d，自动跳过（非2D），保持BF16",
    "proj_out.weight": "非量化（匹配规则1），保持BF16",
    "condition_embedder.*.weight": "非量化（匹配规则1），保持BF16",
    "double_blocks.*.attn.*_qkv.weight": "FP8（匹配规则2）",
    "double_blocks.*.attn.*_proj.weight": "FP8（匹配规则2）",
    "double_blocks.*.*_mlp.net.*.weight": "NVFP4（匹配规则3）",
    "double_blocks.*.*_mlp.net.*.proj.weight": "NVFP4（匹配规则3）",
    "double_blocks.*.*_mod.modulate_table": "1D，自动跳过，保持BF16",
    "double_blocks.*.attn.*_norm.weight": "1D RMSNorm，自动跳过，保持BF16",
    "all *.bias": "1D，自动跳过，保持BF16"
  },

  "_estimated_savings": {
    "bf16_size_gb": "~32 GB（未量化 DiT 16B参数）",
    "quantized_size_gb": "~14 GB（MLP: 4x压缩, Attention: 2x压缩, 边界: 无压缩）",
    "total_estimated_gb": "~30 GB（DiT 14GB + MLLM 16GB BF16）",
    "nvfp4_blocks": "320个MLP权重矩阵（40 blocks × 2 streams × 2 types × 2 layers）",
    "fp8_blocks": "320个Attention权重矩阵（40 blocks × 2 streams × 4 weights）"
  }
}