CANN Hunyuan3D NPU优化实践
NPU Hunyuan3D 模型推理优化实践【免费下载链接】cann-recipes-embodied-intelligence本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-intelligence本文主要介绍Hunyuan3D模型基于NPU的推理优化策略其中包括以下优化点shapegen部分使能融合算子与高性能计算算子PFA算子适配torchair图模式适配DIT 图模式适配VAE-decoder 图模式适配CPU-NPU搬运优化DIT-Cache step-level适配step-level指跳过/预测step范式texgen部分多线程并行光栅化aicpu算子迁移Inpaint计算优化shapegen性能优化介绍融合算子优化优化原因RmsNorm算子常见于LLaMA、LLaMA2、Baichuan等LLM模型中由于torch侧没有提供RmsNorm算子的接口因此在模型中通常是以自定义类的形式出现在forward函数下定义计算逻辑需要的运行算子比较多计算公式为x * torch.rsqrt(x.pow(2).mean(-1, keepdimTrue) self.eps)优化方式将hy3dgen/shapegen/models/denoisers/hunyuan3ddit.py文件中torch.rsqrt函数替换如下替换部分 class RMSNorm(torch.nn.Module): def __init__(self, dim: int): super().__init__() self.scale nn.Parameter(torch.ones(dim)) def forward(self, x: Tensor): x_dtype x.dtype x x.float() rrms torch.rsqrt(torch.mean(x ** 2, dim-1, keepdimTrue) 1e-6) return (x * rrms).to(dtypex_dtype) * self.scale #替换后 class RMSNorm(torch.nn.Module): def __init__(self, dim: int, eps1e-6): super().__init__() self.eps eps self.scale nn.Parameter(torch.ones(dim)) def forward(self, x: Tensor): return torch_npu.npu_rms_norm(x, self.scale, epsilonself.eps)[0]替代高性能算子替代高性能算子优化原因在神经网络中GELU是一个重要的激活函数其灵感来源于relu和dropout在激活中引入了随机正则的思想。为了提升GELU算子在NPU上的运行性能业界提出了FastGelu等版本。本接口FasterGelu是针对FastGelu的化简版本可以大幅度提升计算性能。原理是使用泰勒展开逼近计算结果。优化方式修改hy3dgen/shapegen/models/denoisers/hunyuan3ddit.py文件中GELU函数使用torch_npu.npu_fast_gelu替换nn.functional.gelu的算子实现。替换部分 class GELU(nn.Module): def __init__(self, approximatetanh): super().__init__() self.approximate approximate def forward(self, x: Tensor) - Tensor: return nn.functional.gelu(x, approximateself.approximate) #替换后 class GELU(nn.Module): def __init__(self, approximatetanh): super().__init__() self.approximate approximate def forward(self, x: Tensor) - Tensor: return torch_npu.npu_fast_gelu(x)FIA算子适配和优化优化原因官方使用的scaled_dot_product_attention对应npu算子是npu_fusion_attentionnpu_fusion_attention不支持图模式因此如果将npu_fusion_attention做静态图编译时会导致npu_fusion_attention编译为三个attention核心算子进行运算带来性能劣化。BatchMatMulSoftmaxV2BatchMatMul优化方式使用torch_npu.fused_infer_attention_score代替torch_npu.npu_fusion_attentionFIA同时适配增量全量推理场景的FlashAttention算子能够进行图编译并且同时解决了Softmax运算溢出问题以及内存连续问题。进行如下优化替换部分 scaled_dot_product_attention nn.functional.scaled_dot_product_attention def attention(q: Tensor, k: Tensor, v: Tensor, **kwargs) - Tensor: x scaled_dot_product_attention(q, k, v) x rearrange(x, B H L D - B L (H D)) return x # 替换后 def npu_fia(q, k, v, scale): attn_mask None batch_size, num_head, seq_len, head_dim q.shape out torch_npu.torch_npu.fused_infer_attention_score( q, k, v, num_headsnum_head, input_layoutBNSD, scalescale, atten_maskattn_mask )[0] return out def attention(q: Tensor, k: Tensor, v: Tensor, **kwargs) - Tensor: x npu_fia(q, k, v, scale(1 / math.sqrt(head_dim))) x rearrange(x, B H L D - B L (H D)) return xtorchair图模式适配优化原因在 PyTorch 的单算子执行模式下算子不仅可能面临主机端下发的host bound还可能错过未被识别的算子融合机会而未能充分发挥性能。相较于单算子模式图模式可以通过计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。因此可通过 torch.compile 启用图模式以实现性能优化。TorchAir即 Torch Ascend Intermediate Representation作为 Ascend Extension for PyTorchtorch_npu的图模式功能拓展工具库提供了昇腾设备亲和的torch.compile图模式后端。它能够助力 PyTorch 网络在昇腾 NPU 上开启图模式推理进而实现运算加速与性能层面的优化提升。优化方式在模型推理场景下使能DIT model模块以及VAE decoder模块图编译可以避免标量入图导致的图编译失败从而提高推理性能。同时启用tiling图下沉优化。#增加部分 if args.full_graph: import torchair import torch._dynamo config torchair.CompilerConfig() config.experimental_config.keep_inference_input_mutations True config.experimental_config.frozen_parameter True config.experimental_config.tiling_schedule_optimize True npu_backend torchair.get_npu_backend(compiler_configconfig) #编译入图部分 if args.full_graph: pipeline.model torch.compile(pipeline.model, dynamicFalse, backendnpu_backend, fullgraphTrue) pipeline.vae.geo_decoder torch.compile(pipeline.vae.geo_decoder, dynamicFalse, backendnpu_backend, fullgraphTrue) #图模式指令 python minimal_demo_npu.py --full_graphCPU-NPU搬运优化优化原因从CPU搬运到NPU存在搬运耗时频繁的搬运会导致性能劣化在生成时间步的步骤尤其明显优化方式直接在NPU上生成时间步替换部分 def timestep_embedding(t: Tensor, dim, max_period10000, time_factor: float 1000.0): t time_factor * t half dim // 2 freqs torch.exp(-math.log(max_period) * torch.arange(start0, endhalf, dtypetorch.float32) / half).to( t.device ) #替换后 def timestep_embedding(t: Tensor, dim, max_period10000, time_factor: float 1000.0): t time_factor * t half dim // 2 freqs torch.exp(-math.log(max_period) * torch.arange(start0, endhalf, dtypetorch.float32,devicet.device) / half)DIT-Cache step-level适配基本原理DIT-Cache作为扩散模型推理加速的缓存框架通过复用/预测已有的结果减少冗余前向计算。其加速逻辑可清晰的分为Step-level和Block-level范式Step-level通过判断不同采样步数step间的特定特征差异通过阈值比较决定是否跳过完整的step计算直接复用或者预测缓存结果。Step-level典型方法** 在Step-level加速范畴内FBCache与TeaCache是典型的cache方法。TeaCache 利用模型输入与输出的强相关性通过Timestep Emebdding输入来估计输出差异先利用该输入粗估输出变化再通过多项式拟合修正缩放偏差最终以累积差异作为判断标准动态决定是否复用上一步被Cache的输出避免冗余计算。FBCache的原理是基于First Block L1误差比较第一个Block输出残差与上一步的第一个Block输出残差之间的差异如果首块输出误差与上一轮首块输出误差差异小于指定阈值就跳过当前步计算复用残差对当前步的输出进行估计。block-level典型方法** 在Block-level加速范畴内Taylorseer是典型的以预测范式代替缓存的方法。Taylorseer 解决了Step-level中因扩散模型中的特征相似性显著下降导致的特征缓存引入的错误显著增加损害生成质量。Taylorseer方法基于未来时间步的扩散模型特征可以基于前一时间步的值进行预测这一理论基础和特征在时间步间缓慢且连续变化的现象采用微分方法近似特征的高阶导数并应用Taylor级数预测后续特征。与直接重用缓存特征的方法不同Taylorseer利用特征变化的连续性来预测未来特征使扩散模型能够实现无训练且高比例的加速而不会显著降低生成质量。FBCache优化效果通过63张图像生成任务使用UNI3D指标阈值跳过率采样耗时sUni3d-I加速比baseline019.280.3748x1.000.0450%10.190.3704x1.890.0563%9.620.3687x2.080.0666%7.800.3690x2.47TeaCache优化效果通过63张图像生成任务使用UNI3D/ULIP指标结果如下在跳过率为75%精度损失大于1%因此推荐使用阈值0.1 |阈值|跳过率|采样耗时s|Uni3d-I|ULIP-I|加速比| |:---:|:---:|:---:|:---:|:---:|:---:| |baseline|0|19.28|0.3179|0.2157|x1.00| |0.1|48%|10.32|0.3090|0.2050|x1.86| |0.2|75%|5.83|0.23908|0.1879|x3.30|Taylorseer性能优化效果在推理步数为100预热步数为3截断步数为0进行多轮生成取每次生成时间的平均值 |导数阶数 |跳过间隔|0|1|2|3|4|5|6| |:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:| ||attention次数|4800|2496|1728|1344|1104|960|864| ||MLP次数|6400|3328|2304|1792|1472|1280|1152| ||预测步数|0|48|64|72|77|80|83| |0|总时间|19.9s|10.24s|7.66s|6.17s|5.14s|4.58s|3.98s| |0|加速比|1.00|1.85|2.60|3.23|3.87|4.35|5.00| |0|性能提升|0%|46.0%|61.5%|69.0%|74.2%|77.0%|79.9%| |1|总时间|20.44s|11.14s|8.02s|6.44s|5.51s|4.92s|4.35s| |1|加速比|1.00|1.84|2.55|3.17|3.71|4.15|4.70| |1|性能提升|0%|45.5%|60.8%|68.8%|72.9%|75.9%|78.7%| |2|总时间|20.72s|11.47s|8.30s|6.67s|5.82s|5,23s|4.64s| |2|加速比|1.00|1.81|2.50|3.11|3.56|3.96|4.47| |2|性能提升|0%|44.7%|60.0%|67.9%|71.9%|74.8%|77.6%|导数/阶数跳过率加速比Dit耗时性能提升Uni3d-IULIP-Ibaseline01.0019.26s0%0.3740.221O1N1481.8411.14s45.5%0.3740.222O1N3723.176.4s68.6%0.3720.220O1N5784.154.92s75.9%0.3710.217O2N1481.8111.47s44.7%0.3740.222O2N3723.116.67s67.9%0.3730.221O2N5783.965.23s74.8%0.3680.217启动方式本代码模块通过修改cache_config.json文件决定是否使用CacheCache范式Cache相关参数均在3d_vision\Hunyuan3D\hy3dgen\cache\cache_config.json 中直接修改同时使用如下指令可以自定义cache_config.json位置python minimal_demo_npu.py --cache_config ./hy3dgen/cache/cache_config.json #cache_config.json位置其中参数意义如下{ cache_forward: NoCache, #直接设置Cache方案目前支持FBCache/TeaCache/Taylorseer,默认启动NoCache也就是无ditcache方法只需按照下面的提示代替NoCache即可启动 comment: choose from FBCache/TeaCache/Taylorseer, otherwise use NoCache, enable_separate_cfg: false, #CFG串行开关适配wan2.2等CFG串行的项目开启双分支TeaCache/FBcache,不支持Taylorseer同时在非CFG串行项目保持关闭Hunyuan3D为非CFG串行因此关闭 FBCache:{ cache_name: FBCache, rel_l1_thresh: 0.05, #FBCache阈值阈值越大跳过越多精度损失越大需要平衡性能和精度 latent: latent, judge_input: cache_latent }, TeaCache:{ cache_name : TeaCache, rel_l1_thresh: 0.1, #TeaCache阈值阈值越大跳过越多精度损失越大需要平衡性能和精度 coefficients: [733.226126,-401.131952,67.5869174,-3.149879,0.0961237896], #TeaCache多项式拟合通过输入输出进行拟合 warmup: 2, #预热步数强制前面几步进行计算 latent: latent, judge_input: modulated_inp }, Taylorseer:{ cache_name : Taylorseer, n_derivatives: 3, #Taylorseer导数设置导数阶数越大理论上对原特征拟合越优秀但是过大的导数可能会出现过拟合现象且占用显存会更大 skip_interval_steps: 4, #Taylorseer计算间隔设置注意 计算间隔 跳过计算步数1 跳过计算间隔越大对性能提升越明显但是精度会下降 warmup: 1, #预热步数强制前面几步进行计算避免前期特征变化明显的时候出现跳过计算 cutoff_steps: 1, #截断步数强制最后几步进行计算 offload: false #offload开关开启后可以降低显存占用 }, NoCache:{ cache_name : NoCache } }框架位置使用dit_cache作为自定义库dit_cache方案插入逻辑如下cann-recipes-embodied-intelligence --- module #公共优化模块 | --- dit_cache #Dit_Cache通用方案核心位置 | --- __init__.py #初始化 | --- cache_method.py #Dit_Cache方案选择和方案设计 --- 3d_vision --- Hunyuan3D --- set_env.sh #激活样例环境能够导入 module/dit_cache --- minimal_demo_npu.py #顶层推理入口进行cache方案选择模型层数传入和前向推理替换 --- hy3dgen --- cache #cache适配模型口 --- cache_block.py #Dit_Cache替换模块位置 --- cache_config.json #默认cache参数位置 --- models --- denoisers --- hunyuan3ddit #Cache适配模型核心位置texgen性能优化介绍多线程并行光栅化优化原因CPU侧执行光栅化会带来巨大时延Hunyuan3D texgen在6个不同视角分别进行法向信息、位姿信息、纹理信息投影共进行18次光栅化。由于光栅化相互之间与delighting来自不同数据通路没有数据依赖关系可以采用并行计算的方式隐藏光栅化时延。优化方法将法向计算和位姿计算以及delighting过程中相互间没有依赖关系的计算过程施加并行提升计算效率。对hy3dgen/texgen/pipelines.py中的代码做如下替换替换部分 images_prompt [self.models[delight_model](https://link.gitcode.com/i/019dde2341a10fb216cb63e29cb91b78) for image_prompt in images_prompt] ... normal_maps self.render_normal_multiview( selected_camera_elevs, selected_camera_azims, use_abs_coorTrue) position_maps self.render_position_multiview( selected_camera_elevs, selected_camera_azims) ... texture, mask self.bake_from_multiview(multiviews, selected_camera_elevs,selected_camera_azims,selected_view_weights, methodself.config.merge_method) #替换后 def _delighting_render_multiview(self,images_prompt,selected_camera_elevs,selected_camera_azims ,use_abs_coorTrue): total_tasks len(images_prompt) 2 * len(selected_camera_elevs) max_workers min(20, total_tasks) with ThreadPoolExecutor(max_workersmax_workers) as executor: normal_futures {} position_futures {} for idx,(elev,azim) in enumerate(zip(selected_camera_elevs, selected_camera_azims)): normal_future executor.submit( self.render.render_normal, elev, azim, use_abs_coor use_abs_coor,return_type pl ) normal_futures[normal_future] idx position_future executor.submit( self.render.render_position, elev, azim, use_abs_cooruse_abs_coor, return_typepl ) position_futures[position_future] idx rast_out_maps [None] * len(selected_camera_elevs) normal_maps [None] * len(selected_camera_elevs) position_maps [None] * len(selected_camera_elevs) images_prompt [self.moels[delight_model](https://link.gitcode.com/i/019dde2341a10fb216cb63e29cb91b78) for image_prompt in images_prompt] for future in as_completed(list(normal_futures.keys()) list(position_futures.keys())): try: result future.result() if future in normal_futures: idx normal_futures[future] normal_maps[idx] normal elif future in position_futures: idx position_futures[future] position_maps[idx] result except Exception as e: print(fTask failed with error: {e}) return images_prompt, normal_maps, position_maps def _bake_from_multiview(self, view, camera_elevs, camera_azim, view_weights,methodgraphcut): with ThreadPoolExecutor(max_workers20) as executor: futures [] for idx,(view,camera_elev, camera_azim, weight) in enumerate(zip( views, camera_elevs, camera_azims, view_weights)): future executor.submit(self.render.back_project, view, camera_elev, camera_azim) futures.append(future) project_textures [] project_weighted_cos_maps [] project_boundary_maps [] for future in as_completed(futures): project_texture, project_weighted_cos_map, project_boundary_map future.result() project_textures.append(project_texture) project_weighted_cos_maps.append(project_weighted_cos_map) project_boundary_maps.append(project_boundary_map) def __call__(self,mesh,image): ... # 对应光栅化和delighting部分 images_prompt, normal_maps, position_maps self._delighting_render_multiview(images_prompt, selected_camera_elevs, selected_camera_azims, use_abs_coorTrue) ... # 对应multiview部分 texture, mask self._bake_from_multiview(multiviews, selected_camera_elevs selected_camera_azims, selected_view_weights, method self.config.merge_method)aicpu算子迁移优化原因FP32数据格式会导致torch.mean、torch.std算子运行在aicpu上性能下降。优化方法将FP32类型的数据手动进行转换使其运行在aicore上。# hy3dgen/texgen/pipelines.py 替换部分 texture torch.tensor(texture_np / 255).float().to(texture.device) #替换为 texture torch.tensor(texture_np / 255).to(torch.float16).to(texture.device)# hy3dgen/texgen/utils/dehighlight_utils.py 替换部分 src_mean, src_stddev torch.mean(src_flat[:, i].to(torch.float32)), torch.std(src_flat[:, i].to(torch.float32)) target_mean, target_stddev torch.mean(target_flat[:, i].to(torch.float32)), torch.std(target_flat[:, i].to(torch.float32)) #替换为 src_mean, src_stddev torch.mean(src_flat[:, i].to(torch.float16)), torch.std(src_flat[:, i].to(torch.float16)) target_mean, target_stddev torch.mean(target_flat[:, i].to(torch.float16)), torch.std(target_flat[:, i].to(torch.float16)) 替换部分 image_tensor torch.tensor(np.array(image)/255.0).to(self.device) #替换为 image_tensor torch.tensor(np.array(image) / 255.0).to(torch.float16).to(self.device)Inpaint计算优化优化原因生成的3Dmesh纹理仍存在部分未上色点需要通过对其平滑化补全未上色点。源码实现了python代码和c代码默认使用python代码进行着色点平滑但是c具有更好的性能需要手动进行替换。优化方法手动替换为c代码。# hy3dgen/texgen/differentiable_renderer/mesh_render.py 替换部分 from .mesh_processor import meshVerticeInpaint #替换为 import mesh_processor 替换部分 texture_np, mask meshVerticeInpaint( texture_np, mask, vtx_pos, vtx_uv, pos_idx, uv_idx) #替换为 texture_np, mask mesh_processor.meshVerticeInpaint( texture_np, mask, vtx_pos, vtx_uv, pos_idx, uv_idx)光栅化结果复用优化优化原因源码的18次光栅化计算中只有6种不同输入数据其余12次计算为冗余计算造成性能损失。优化方法将第一次计算的光栅化过程进行保存每当输入相同的相机位姿算法就从缓存中直接读取计算结果。替换部分 rast_out, rast_out_db self.raster_rasterize( pos_clip, self.pos_idx, resolutionresolution) #替换为 if (self.save_render and (elev, azim, camera_distance) in self.render_result.keys()): rast_out self.render_result[(elev, azim, camera_distance)] ... self.render_result[(elev, azim, camera_distance)] rast_out光栅化过程npu迁移优化原因Hunyuan3D源码光栅化没有可在NPU执行的版本会将光栅化过程迁移至CPU侧执行不必要的内存搬运和串行计算会带来巨大时延。为了适配NPU架构提升光栅化性能本仓库引入render_npu将光栅化的计算过程迁移至NPU执行。优化方法修改光栅化算法将原先采用遍历mesh顶点的方式改为对二维平面分块每个区域中的三角网格合并成一个矩阵计算深度、重心坐标等参数最终判断每个像素点对应的三角面片。由NPU侧执行小算子计算不能同时进行光栅化并行。替换部分 rast_out, rast_out_db self.raster_rasterize( pos_clip, self.pos_idx, resolutionresolution) #替换为 if self.use_render_npu: rast_out, _ render_npu_rasterize( resolution, self.pos_idx, pos_clip ) else: rast_out, _ self.raster_rasterize( pos_clip, self.pos_idx, resolutionresolution )性能优化指标shapegen性能优化指标本方案使用1卡Atlas 800I A2推理产品输入hunyuan3D 2.0 提供的样例数据(assets/example_images/004.png)在扩散步数为num_inference_steps100minimal_demo_npu.py第89行与111行情况下,性能指标如下 |使能方法|DIT 扩散耗时s|VAE 生成耗时s| |:---:|:---:|:---:| |baseline|21.63|17.98| |算子优化搬运优化|20.07|18.71| |算子优化图模式编译|19.00|16.68|texgen性能优化指标本方案使用1卡Atlas 800I A2推理产品输入hunyuan3D 2.0 提供的样例数据(assets/example_images/004.png)输入mesh网格面数量为20000开启减少面数量--face_reduce性能指标如下 |使能方法|texgen运行时长s| |:---:|:---:| |baseline|59.83| |光栅化并行|35.43| |光栅化并行算子优化|30.46| |NPU光栅化算子优化Inpaint优化光栅化结果复用|32.86| |光栅化并行算子优化Inpaint优化光栅化结果复用|26.52|【免费下载链接】cann-recipes-embodied-intelligence本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-intelligence创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考