Chapter 7: 2:4 结构化稀疏 (Structured Sparsity)：收益条件与网络约束

1. 开篇与学习目标

NVIDIA DRIVE Orin 的 GPU 基于 Ampere 架构，引入了一项革命性的硬件特性：2:4 结构化稀疏（2:4 Structured Sparsity）。这项技术允许在保持网络精度基本不变的情况下，将 Tensor Core 的矩阵运算理论吞吐量翻倍（2x Math Throughput），并显著减少显存带宽占用。

在自动驾驶场景中，Vision Transformer (ViT) 和大型 BEV 网络正变得越来越重。传统的剪枝（Pruning）会导致内存访问不连续，难以在 GPU 上加速；而 2:4 稀疏 是 NVIDIA 给出的“硬件-算法协同设计”的标准答案。

然而，“天下没有免费的午餐”。启用 2:4 稀疏不是简单的“开关”，它要求网络设计者在 维度对齐、算子选择 和 训练流程 上必须遵守严格的物理法则。如果在错误的层（如 DWConv）强行使用，不仅不会加速，反而会因为元数据（Metadata）读取开销导致性能倒退。

本章学习目标：

1. 透彻理解硬件机制：Ampere Tensor Core 如何处理稀疏矩阵与元数据。
2. 精准识别收益层：分辨哪些层是“黄金候选者”，哪些是“性能陷阱”。
3. 掌握训练与落地流程：从 Dense 训练到 Sparse Fine-tuning 再到 QAT 的完整链路。
4. Orin 异构策略：在 GPU（支持计算稀疏）与 DLA（仅支持存储稀疏）之间做出正确的架构决策。

2. 核心概念论述

2.1 硬件原理：Ampere Tensor Core 的魔法

传统的非结构化剪枝（Unstructured Pruning）虽然能大幅减少参数，但留下的权重位置随机，导致 GPU 线程束（Warp）出现严重的 Divergence分支发散） 和非合并内存访问，实际推理往往更慢。

Ampere 架构定义了强制性的 2:4 模式：

Rule of Thumb: 在权重矩阵的每 4 个连续元素（沿着输入通道 $K$ 维度）中，必须严格有 2 个元素为零。

硬件执行流程图解

当 Tensor Core 执行稀疏矩阵乘法（Sparse GEMM）时，它不再执行 $1 \times 4$ 的点积，而是执行压缩后的 $1 \times 2$ 点积。

原始权重 (Dense Weights) - 1x4 Slice:
Idx: [ 0,     1,     2,     3   ]
Val: [ 0.1,  -0.5,   0.8,   0.2 ]

-----------------------------------------------------------
Step 1: 稀疏化 (Pruning) - 保留绝对值最大的两个
-----------------------------------------------------------
Mask:[ 0,     1,     0,     1   ]
Val: [ 0.0,  -0.5,   0.0,   0.2 ]  <-- 索引 0 和 2 被置零

-----------------------------------------------------------
Step 2: 压缩存储 (Compression) - 存入显存
-----------------------------------------------------------
Data Stream (FP16):  [ -0.5,  0.2 ]      (仅存储非零值，显存占用 50%)
Metadata (2-bit):    [  01,   11  ]      (记录非零值在原始 4 元素组中的位置)

-----------------------------------------------------------
Step 3: 硬件执行 (Hardware Execution)
-----------------------------------------------------------
Input Feature Map (Activation):
[ x0, x1, x2, x3 ]

Tensor Core 的稀疏多路选择器 (Mux):
根据 Metadata [01, 11]，自动抓取 Input 中的 x1 和 x3。

ALU 计算:
Result += (-0.5 * x1) + (0.2 * x3)
*注意：硬件在一个时钟周期内完成的有效计算量翻倍，因为它跳过了 x0 and x2 的乘法。*

2.2 收益分析：哪些层适合做 2:4？

稀疏化能否带来加速，取决于 算术强度 (Arithmetic Intensity) 和 维度大小。只有当计算节省的时间 > 读取元数据的额外开销时，加速才会发生。

A. 黄金候选层 (High Yield Layers)

这些层通常是 Compute-bound 的，且维度足以满 Tensor Core 的 Pipeline。

1. Transformer Linear Layers (Projections):
   • qkv_proj, out_proj, fc1, fc2。
   • 特征：输入输出维度通常较大（$C \ge 512$），且完全是矩阵乘法（GEMM）。
   • 收益：极高。这是 Orin 上跑 ViT/BEVFormer 的关键优化点。
2. Conv2d 1x1 (Pointwise Conv):
   • ResNet 的 Bottleneck 结构、EfficientNet 的 Expansion/Projection 层。
   • 特征：计算模式等同于 GEMM。
   • 收益：高。
3. Large Conv2d 3x3:
   • 当 Batch Size 较大（如 BEV 场景下的多摄 Batch）或 Feature Map 分辨率较高时。

B. 陷阱层 (The "Do Not Touch" Layers)

这些层通常是 Memory-bound 的，或者无法映射到 Tensor Core。

1. Depthwise Convolution (DWConv):
   • 原因：DWConv 的计算密度极低，主要瓶颈在显存带宽。稀疏化引入了 Metadata 读取，反而加剧了带宽压力。切勿对 DWConv 启用 2:4。
2. 第一与最后一层 (Stem & Head):
   • 原因：第一层输入通道少（如 3），最后一层输出通道少（如检测头 Class=80）。维度不足以掩盖 overhead。
3. 非 GEMM 算子:
   • LayerNorm, Softmax, Add, Concat。这些算子不使用 Tensor Core，稀疏化无效。

2.3 2:4 稀疏对网络设计的约束

要在 Orin 上成功部署，架构设计阶段就必须考虑以下约束：

1. 维度对齐 (Alignment):
   • 为了启用 Tensor Core 的稀疏指令，矩阵的维度（$M, N, K$）最好是 32 的倍数（至少是 16 的倍数）。
   • 设计建议：如果是检测头，类别数是 80，建议 Padding 到 96 或 128。如果是 Hidden Size，尽量使用 768, 1024, 2048 等标准值。
2. 结构化剪枝的破坏性:
   • 2:4 属于强约束。如果某一层对精度极其敏感（如某些 Backbone 的下采样层），强制 2:4 可能会导致特征崩塌。需要预留“跳过稀疏化”的机制（Mixed Sparsity）

3. 训练与部署策略

2:4 稀疏模型不能通过简单的后处理得到，必须经过稀疏感知训练 (Sparsity-Aware Training)。

3.1 推荐的工作流 (The ASP Workflow)

NVIDIA 提供了 ASP (Automatic SParsity) 工具（包含在 TensorRT PyTorch Quantization Toolkit 或 Apex 中）。

1. Step 1: Dense Training:
   • 先训练一个收敛的高精度稠密模型。
2. Step 2: Sparsity Masking & Permutation:
   • Magnitude Pruning: 对权重进行分析，生成 2:4 Mask。
   • Channel Permutation (关键技巧):
     • 在生成 Mask 前，算法可以寻找最佳的通道排列组合。
     • 原理：如果某一行是 [0.9, 0.1, 0.9, 0.1]，直接剪枝会丢掉信息。如果交换第2、3列变成 [0.9, 0.9, 0.1, 0.1]，则两个 0.9 都能保留。这种数学上等价的变换能显著提升稀疏后的精度。
3. Step 3: Sparse Fine-tuning:
   • 固定 Mask（或使用 SR-STE 动态调整），以较小的学习率进行重训练。
   • Distillation (蒸馏)：强烈建议使用原始 Dense 模型作为 Teacher，Sparse 模型作为 Student。损失函数通常是 Loss = Hard_Label_Loss + Alpha * KL_Div(Teacher, Student)。
4. Step 4: QAT (Quantization Aware Training):
   • 在稀疏化的基础上，插入 Fake Quantization 节点，进行 INT8 微调。
   • 最终形态：Sparse + INT8。这是 Orin 性能的终极形态。

3.2 GPU vs. DLA 的决策

在 Orin 异构计算中，这是一个常见的混淆点。

| 特性 | Orin GPU (Ampere) | Orin DLA (Deep Learning Accelerator) |

Rule of Thumb:

• 如果块部署在 GPU (如 ViT, Head, BEV Pool)：做 2:4 稀疏训练。
• 如果模块部署在 DLA (如 ResNet/EfficientNet Backbone)：不做 2:4 结构化稀疏。DLA 更喜欢结构化的通道剪枝（Channel Pruning）或单纯的 INT8。

4. 本章小结

1. 机制：2:4 稀疏通过每 4 个权重丢弃 2 个，配合 Tensor Core 的稀疏指令，实现 2x 理论算力。
2. 选型：仅对 Conv1x1 和 Linear 层收益显著；严禁用于 DWConv 和小维度层。
3. 对齐：输入输出通道数应设计为 32 的倍数 以利用硬件 Tensor Core Tiling。
4. 流程：必须进行 Fine-tuning，最好配合 Channel Permutation 和 蒸馏。
5. INT8 耦合：Sparse + INT8 是 Orin GPU 上吞吐量最大的模式，但训练难度最高。

5. 练习题

基础题 (Basic)

Q1. 2:4 掩码计算与数据排布 给定一组 FP16 权重（单通道切片），请手动执行 2:4 剪枝，并计算压缩后的存储节省比率（考虑元数据）。 权重：W = [1.2, 0.1, -0.9, 0.05, 0.8, 0.9, -0.1, 0.2] 提示：FP16=16bit，Metadata通常每对非零值需要 4bit (2bit2)*

Q2. 适用性判断 对于一个部署在 Orin GPU 上的目标检测网络，以下哪些层应该开启 2:4 稀疏？（多选） A. ImageNet Pretrained ResNet50 的 conv1 (7x7, stride=2, Cin=3, Cout=64) B. FPN (Feature Pyramid Network) 中的横向连接层 (1x1 Conv, Cin=2048, Cout=256) C. EfficientNet-B0 的 MBConv 中的 Depthwise Conv (3x3, Cin=144, Groups=144) D. Vision Transformer 的 MLP 层 (Linear, Cin=768, Cout=3072)

挑战题 (Challenge)

Q3. 性能回归的根本原 (Root Cause Analysis) 你将一个 batch size = 1 的实时语义分割模型（基于 HRNet）进行了 2:4 稀疏化并部署在 Orin 上。虽然模型参数量减半，但推理延迟（Latency）几乎没有变化，甚至某些 Layer 变慢了。请结合 GPU 架构特性分析原因。

Q4. INT8 与 Sparse 的联合量化误差 在进行 Sparse + INT8 部署时，你发现某一层（Attention Output Projection）的 MSE 误差异常高。经检查，该层权重的分布呈现极端的长尾形状（大部分值很小，极少数指极大）。请解释为什么 2:4 稀疏加剧了 INT8 的量化误差，并提出解决方案。

6. 常见陷阱与错误 (Gotchas)

6.1 "Ghost Sparsity" (虚假稀疏)

• 现象: 你辛辛苦苦训练了稀疏模型，导出 ONNX 后权重确实有很多 0，但在 TensorRT 并没有触发 Sparse Kernel，速度没变。
• 原因: TensorRT 对 2:4 的检查极为严格。
  1. Flags: 构建 Engine 时未开启 --sparsity=enable。
  2. 结构违规: 只要矩阵中有一个 $1 \times 4$ 块不满足 "2个0" 的条件，TensorRT 就会自动回退到 Dense 模式，且往往不会报错，只是默默变慢。
• 避坑: 在导出 ONNX 前，务必编写脚本遍历所有权重，断言 check_sparsity_2_4(weight) == True。

6.2 维度未对齐导致的 Padding 惩罚

• 现象: 将一个输出为 1000 类（ImageNet）的全连接层稀疏化，性能反而下降。
• 原因: Tensor Core 稀疏计算通常要求 $M, N, K$ 为 16 或 32 的倍数。1000 不是 16 的倍数（$1000 / 16 = 62.5$）。TensorRT 会自动 Padding 到 1008 或 1024。
• 代价: 对于这种并未对齐的层，Padding 带来的额外计算量 + 内存拷贝可能超过了稀疏化带来的收益。
• 避坑: 模型设计之初就定义 num_classes = 1024 (剩下的 24 个作为 dummy class)，从源头解决对齐问题。

6.3 忽视 Permutation 导致精度无法恢复

• 现象: 直接对训练好的 Dense 模型做 Magnitude Pruning，精度下降 5% 以上，Fine-tuning 怎么也救不回来。
• 原因: 刚性地每 4 个删 2 个，打破了特征的组合性。
• 避坑: 必须在剪枝前进行 Channel Permutation。许多现代库（如 NVIDIA ASP）支持搜索最优的通道排列矩阵，将大权重聚拢到同一个 block 中，从而在稀疏化时保留更多信息。

6.4 DLA 的“稀疏”误解

• 陷阱: 看到 DLA 文档里写了 "Weight Compression" 和 "Sparse Support"，就以为 DLA 也能像 GPU 一样获得 2x 算力。
• 真相: DLA 的稀疏主要为了省带宽。如果你的模型主要瓶颈在 DLA 的 MAC 利用率（算力），做稀疏化毫无帮助，反而可能因为特殊的存储格式变得麻烦。在 DLA 上，Structured Pruning (减少通道数) 永远 2:4 Sparsity 有效。

< 上一章：量化设计（INT8/FP16） | 下一章：推荐的 Orin 友好模块库 >