量化

限制与约束

仅Atlas 800I A2 推理服务器支持此特性。

操作步骤

MindIE SD量化需要先使用工具导出权重，再使用推理框架的接口进行量化推理。

1. 安装大模型压缩工具，详情请参见链接。
2. 对于包含激活值的量化算法，参考对应的示例和接口，对量化权重进行导出。对于仅权重量化，参考对应的示例和接口，对量化权重进行导出。
   

   • 对于单卡进行量化权重导出，使用工具默认量化权重和描述符命名即可。
   • 对于多卡并行量化，推理框架限制命名规则。
   • 量化权重命名规则为quant_model_weight_{quant_algo.lower()}_{rank}.safetensors。
   • 描述符命名规则为quant_model_description_{quant_algo.lower()}_{rank}.json。
3. 使用def quantize的接口，对浮点模型进行量化转换，该接口会处理量化权重和修改计算图，示例如下所示。

   from mindiesd import quantize
   
   model = from_pretrain()
   model = quantize(model, "步骤2导出的quant json path")

   

   • 对于fa量化，要求被fa量化的类中持有inner_dim和heads属性。
   • 模型自行加载原始权重，并完成实例初始化，quantize由插件提供，在接口中对相应层进行量化转换。
   • 模型可以选择在quantize转换完后再使用to npu。

   • 如果使用时间步量化，在quantize中还需要传入class TimestepPolicyConfig，量化转换后还需要使用class TimestepManager在模型中设置时间步信息，示例如下：

     with self.progress_bar(total=num_inference_steps) as progress_bar:
         for i,t in enumerate(timesteps):
             if self.interrupt:
                 continue
             # -----------新增代码-----------
             from mindiesd import TimestepManager
             TimestepManager.set_timestep_idx(i)
             # -----------新增代码-----------
             latent_model_input = (
                 torch.cat([latents] * 2)
                 if self.do_classifier_free_guidance
                 else latents
             )

   • 如果使用FA量化，量化转换后还需要手动使用class QuantFA在模型的transformer block中修改FA的调用逻辑，例如fa3会在量化转换后自动生成：

     if hasattr(attn, "fa3") and query.shape[0] == key.shape[0]:
         seq_len = [query.shape[0]]*query.shape[1]
         query = rearrange(query, "s b (n d) -> (b s) n d", d=head_dim, b=1)
         key = rearrange(key, "s b (n d) -> (b s) n d", d=head_dim, b=1)
         value = rearrange(value, "s b (n d) -> (b s) n d", d=head_dim, b=1)
         hidden_states = attn.fa3.forward(query, key, value, seq_len)
         hidden_states = rearrange(hidden_states, "(b s) n d -> s b (n d)", d=head_dim, b=1)
     else:
         query = rearrange(query, "s b (n d) -> b n s d", d=head_dim) # BNSD
         key = rearrange(key, "s b (n d) -> b n s d", d=head_dim) # BNSD
         value = rearrange(value, "s b (n d) -> b n s d", d=head_dim) # BNSD
         hidden_states = torch_npu.npu_fusion_attention(query, key, value,
                                                     atten_mask=attention_mask,
                                                     input_layout="BNSD",
                                                     scale=1 / math.sqrt(head_dim),
                                                     head_num=attn.heads // sp_size)[0]