模型压缩零基础一、以Cosmos Policy为例的世界动作模型时延测量端到端推理时延、VAE Encoder / Diffusion / VAE Decoder时延1. Cosmos Policy模型架构2. 时延的分类2.1 端到端推理时延2.2 VAE Encoder时延2.3 Diffusion模块时延2.4 VAE Decoder时延3. 时延的测量3.1 理解CPU和GPU的异步工作机制3.2 CUDA synchronization3.3 单独测某一个GPU模块的时延Encoder / Diffusion / VAE Decoder时延1. Cosmos Policy模型架构Cosmos Policy是NVIDIA在2026年提出的一种世界动作模型World Action Model简称WAM是一种基于预训练视频生成模型 Cosmos-Predict2 微调得到的统一机器人策略模型其核心架构是Diffusion Transformer。下面是Cosmos Policy的模型架构图图源论文Cosmos Policy 整体由 VAE Encoder、Text Encoder、Diffusion Transformer 和 VAE Decoder 四个部分组成。如果是做压缩工作可以先从测量各部分时延做起。由于这里Text Encoder所占时间较少并且测量时延的方式与VAE Encoder类似故本文只以VAE Encoder、Diffusion、VAE Decoder三部分的时延测量为例进行记录。2. 时延的分类做压缩工作的时候常常需要统计各部分的推理时延以此判断某个模块还有没有压缩的必要。这里主要测量了端到端推理时延、VAE Encoder时延、Diffusion 时延内部又包括AdaLN、Self-Attention、Cross-Attention、MLP、VAE Decoder时延。2.1 端到端推理时延端到端推理时延end-to-end inference latency是模型加速工作中常用的效率指标之一表示模型从接收输入到产生最终输出所需的总时间。在相同硬件、batch size、精度和测量设置下端到端 speedup 通常可通过 Baseline 与压缩模型的端到端时延之比计算即 Speedup Latency_baseline / Latency_compressed。相比仅统计单个模块或算子的运行时间端到端时延能够更全面地反映模型在实际部署中的整体推理效率。2.2 VAE Encoder时延VAE Encoder的时延就是输入的图像、机器人本体状态、动作噪声等等帧经过VAE Encoder的计算变成低维Latent帧的所需要的时间简单来说就是一堆向量经过VAE encoder计算变成另一堆向量的时间。2.3 Diffusion模块时延Diffusion TransformerDiT是 Cosmos Policy 推理计算的核心部分。进一步细分每个去噪步都会依次经过多个 DiT Block而每个 Block 主要包含 Adaptive Layer NormalizationAdaLN、Self-Attention、Cross-Attention和MLP等计算模块。在当前实验的推理配置下Cosmos Policy 采用 5 个去噪步完成动作生成。因此这里分别统计单个去噪步的平均推理时延以及去噪过程中 AdaLN、Self-Attention、Cross-Attention 和 MLP 等主要模块的计算时延。2.4 VAE Decoder时延VAE Decoder 时延是指将 Diffusion Transformer 生成的图像 latent 表征解码回像素空间所需要的时间。在 Cosmos Policy 中VAE Decoder 主要用于将预测的未来视觉状态Future Imagelatent 解码为未来相机图像而 Action、Proprioception 和 Value 等非图像模态无需经过 VAE Decoder。3. 时延的测量3.1 理解CPU和GPU的异步工作机制运行深度学习模型的 Python 程序时CPU 和 GPU 通常协同完成推理CPU 主要负责 Python 代码执行、数据准备、任务调度以及 GPU kernel 的发起而 GPU 主要负责矩阵乘法、Attention、卷积等大规模并行计算两者通过数据传输、kernel launch 和同步机制进行协作。先以一个简单的例子介绍CPU和GPU各自负责的内容以及具体工作机制。noisetorch.randn(...,devicecuda)outputmodel(noise)CPU 做的是执行 Python调用 PyTorch向 CUDA stream 提交 kernel继续执行后面的 Python。GPU则稍后按照 stream 顺序执行GPU stream:[randn kernel] [DiT kernel 1] [DiT kernel 2] [attention] [MLP]…CUDA 计算通常采用异步执行机制CPU 向 GPU 提交计算任务kernel后通常不会等待 GPU 执行完成而会继续执行后续代码。因此在测量 GPU 推理时延时如果直接使用 time.time() 对代码段计时结束计时点可能发生在 GPU 计算真正完成之前从而低估实际执行时间。为获得准确的墙钟时延wall-clock latency需要在计时边界进行 CPU-GPU 同步这就引出了CUDA synchronization。3.2 CUDA synchronization举个例子讲解synchronization的作用先看一个错误例子t0time.perf_counter()outputmodel(x)t1time.perf_counter()latency_ms(t1-t0)*1000可能发生CPU: t0 – 提交许多CUDA任务 – t1GPU: 尚未执行完 --------------------因此测到的主要是 CPU 提交时间结果可能只有几毫秒而 GPU 实际运行了几十毫秒。正确的端到端测量应该是torch.cuda.synchronize()t0time.perf_counter()outputmodel(x)torch.cuda.synchronize()t1time.perf_counter()latency_ms(t1-t0)*1000时间线变成了CPU: sync – t0 – 提交任务 – 等待GPU完成 – t1GPU: [真正计算过程]两个同步分别有不同作用t0 前同步清空之前遗留的 GPU 工作防止前一个样本污染本次测量。t1 前同步等待本次 GPU 工作真正完成。讲到这里其实你已经理解了端到端时延测量的方式。3.3 单独测某一个GPU模块的时延前面的代码用于测量模型的端到端推理时延。如果进一步希望分别统计 VAE Encoder、Diffusion Transformer 和 VAE Decoder 等模块的执行时延就需要在各模块的调用位置分别设置计时点。例如测VAE Encoderstarttorch.cuda.Event(enable_timingTrue)endtorch.cuda.Event(enable_timingTrue)start.record()latentvae.encode(image)end.record()end.synchronize()encoder_msstart.elapsed_time(end)为什么只在end后做了一次同步先梳理一下代码逻辑start.record()# CPU向当前CUDA stream提交一个“执行到这里时记录GPU时间戳”的事件。此时CPU通常不会等待GPU也可能还没有执行到start.CPU提交 start.record() ------→ 立即继续GPU可能还在执行前面的任务……稍后GPU执行到该事件时才会记录GPU时间戳。latentvae.encode(image)# 1. CPU执行 vae.encode() 中的 Python 代码# 2. CPU向 CUDA stream 提交卷积、attention 等 GPU kernel# 3. CPU提交完后继续向下运行# 4. GPU异步执行这些 kernel。CPU执行Python → 提交encode kernels → 继续执行GPU [encode kernel 1] [kernel 2] [kernel 3]end.record()# CPU向同一个CUDA stream提交另一个“执行到这里时记录GPU时间戳”的事件。GPU任务队列变成[start event] [VAE encode kernels] [end event]因为位于同一个 streamGPU必然依次执行。因此GPU执行到 end 时VAE encode 已经完成。不需要担心需要清空GPU的遗留工作因为GPU stream是依次执行的所以不需要同步。但调用 end.record() 返回时GPU可能还没执行完 encode也可能还没执行到 end.这里才需要CPU停下来等待GPU的执行到endCPUe1.synchronize() ───── 等待 ─────→ 返回GPU [encode kernels][e1]---------------------------------------- ↑--------------------------------CPU可以继续换种说法需要让CPU等待的时候我们就调用cuda.synchronization().当 e1.synchronize() 返回后两个事件的 GPU 时间戳都已经有效encode_msstart.elapsed_time(end)得到的是encode_ms GPU执行到end的时间 - GPU执行到start的时间也就是这段 GPU stream 时间。
