SAM3提速8倍优化记录

SAM3——Meta的Segment Anything第三代——发布了,突然你的团队需要将它投入生产。有人从仓库里拿来了推理notebook,把它变成了脚本,几小时后:它能运行,能分割,结果看起来不错。上线。但是……它在H200上吃掉87GB显存,速度很慢,GPU集群账单开始有自己的意见了。

这大致就是我走进来的情况。一个能用的基线,在时间压力下快速搭建——做了氛围编码基线最擅长的事:正确运行(大部分情况下)且低效。没什么丢人的。但"它能工作"和"它已准备好用于生产"是两件非常不同的事。

所以这是我们在几轮剖析和优化后的结果——没有使用ONNX、TensorRT或任何量化。

默认批次大小下吞吐量提升8倍。内存从87GB降到23GB——足以完全摆脱H200,转向更便宜的GPU类别。单次前向传播中推断更多提示,检测性能没有任何下降。

图1: 各优化阶段的平均每批次推理时间(毫秒),从基线到最终。每一步都建立在前一步之上。

让我们看看是如何到达那里的。

1. 什么是SAM3?

SAM3是Meta最新的Segment Anything模型,它添加了前两个版本没有的东西:概念理解。SAM1和SAM2要求你显式提示每个对象——在这里点个点,在那里画个框,得到掩码。SAM3接收"企鹅"这样的名词短语,返回图像中每个匹配的实例,或在视频中跟踪所有实例。相同的基础理念,完全不同的任务。

图2: SAM3在图像和视频上的应用。

动画展示了两种操作模式。在图像模式下,感知编码器主干将文本和图像特征送入基于DETR的检测器,返回实例掩码、边界框和置信度分数。在视频模式下,检测器与基于记忆的跟踪器配合:在第一帧检测到的实例被向前传播,检测器定期重新触发以捕获新对象,匹配步骤保持跨帧身份一致性。系统还可以接受图像示例——正面或负面边界框——以交互式细化匹配内容。

这是一个能力很强的模型,有很多组件设计用于覆盖完整范围的用例。其中有多少适用于任何给定的部署正是第3节要讨论的内容。

2. 打开剖析器

在动手之前,我做两件事:阅读足够的代码以了解推理循环的结构,然后设置剖析。你需要前者来做好后者——在没有理解架构的情况下放置NVTX标记只能告诉你一半的故事。在我们的案例中,我们包装了顶级阶段(数据加载、预处理、模型前向、后处理、磁盘写入),并在模型前向内部添加了子标记以单独捕获每个主要子模块。然后:预热5个批次,剖析接下来的4个。批次大小16,每幅图像10个文本提示。这是我们得到的结果。

图3: Nsight Systems基线时间线——4个剖析批次,批次大小16,每幅图像10个文本提示。

首先跃入眼帘的是一切都是多么顺序化。加载→预处理→推理→后处理→写入,一个接一个,没有重叠。批次之间的间隙比推理本身还宽——GPU空闲着,而CPU加载和预处理下一批次。我们在为没有使用的计算时间付费。让我们放大单个批次。

图4: 放大到批次6——数据加载和磁盘写入阶段

数据加载阶段是突发性的——磁盘读取成块到达,块之间有可见间隙,反映了集群存储系统的延迟和带宽(这也随集群负载变化)。在另一端,磁盘写入是完全阻塞的:循环等待所有内容写入后才能开始下一次迭代。写入阶段有一组DtoH峰值——结果被拉回CPU后才序列化。现在进入GPU本身。

图5: 放大到批次6——GPU推理和后处理。

三件事立刻脱颖而出。

cudaStreamSynchronize调用频繁且昂贵。每个调用都会阻塞主机线程直到GPU清空——CPU等待期间无法提交新的内核。同步解决后,在工作恢复之前会有一个调度间隙。没人故意放在那里——它们是在开发过程中悄悄潜入的,因为没有破坏任何东西所以从未被捕获。

聚集在这些同步点周围的是HtoD(主机到设备)和DtoH(设备到主机)峰值,没有明显理由存在。模型在GPU上运行——有些东西正在悄悄把数据拉回CPU再发送回去。可疑,但我们会等到深入调查后才找根本原因。

后处理阶段比它应该做的要重,而且有自己的DtoH峰值爆发。那里也有问题。

子模块分解为我们提供了一个有用的地图,知道该关注什么。图像编码器需要268毫秒,transformer融合编码器249毫秒,transformer解码器130毫秒——重计算。但通用分割头——负责实例掩码和语义分割输出——时钟216毫秒,几何编码器16毫秒。相比之下,文本编码器只有11毫秒。后处理又增加了57毫秒。那个216毫秒的分割头很快就会变得非常相关。

我们有足够清晰的画面。在开始拉单个线程之前,了解模型本身会有帮助——它是为了什么而构建的,以及我们实际需要它做什么。

3. 了解你的模型,了解你的用例

SAM3是为了处理很多而构建的。单幅图像、视频序列、多种提示类型——点、边界框、掩码、文本——跨帧记忆注意力用于在视频中跟踪对象,以及用于不同分割任务的多个输出头。这是一个通用的研究模型,旨在灵活应对广泛的用例。

我们的用例要窄得多:单帧预测,仅文本提示(而且总是相同的静态提示——运行时没有动态内容),不需要语义分割。没有视频,没有跨帧记忆,没有其他提示类型。

模型不知道这一点。它运行所有这些。

这不是对Meta工程的批评——这只是研究代码的本质。它携带SAM1和SAM2的DNA,以向后兼容的遗留组件形式存在。它有训练产物:仅为损失项计算的中间输出,传递但从没被下游消费者读取的隐藏状态,只在开发期间重要的断言。它有你不适用的用例的多任务输出头。当你剖析模型时,所有这些都出现在时间线中,消耗计算和内存,却不产生任何你需要的东西。

很多工程师在这里停下:"这是大实验室的代码,我不该碰它。"可以理解,但错了。你有特定的用例和特定的约束——这足以让你打开模型并修改它。我们的方法:为你的部署用例添加子目录,实现相关模块的修改版本,导入这些代替原始版本。上游代码保持完整,你的更改保持隔离。

什么值得修改、什么不值得是一个判断调用,剖析器帮助你做出决定。例如,文本编码器运行在从不改变的静态提示上——我们可以缓存输出并完全跳过编码器。但在11毫秒的情况下,它在时间线中几乎不可见,手术的成本将超过运行时节省的工程时间。(剧透:在所有优化后它降到3.3毫秒——仍然不值得。)216毫秒的通用分割头则是另一回事。

我们会在优化章节讨论具体细节。现在的关键点是:理解模型能做什么和你需要它做什么之间的差距——然后给自己许可去关闭它。

图6: SAM3单帧推理的粗略架构草图。

草图绘制了单帧推理的主要组件。文本提示通过分词器和文本编码器。视觉提示——点、框或掩码——通过几何编码器。图像编码器处理原始图像。所有三个送入transformer编码器,融合成条件表示。transformer解码器接收这些以及学习的查询令牌,产生通过头传递的检测结果,给出框、分数和实例掩码。通用分割上下文——像素解码器和语义分割头——位于该路径旁,处理语义输出,像素解码器还送入实例掩码生成。

其中哪些你真正需要取决于你的用例。框图看起来像是翻转开关——实践中,每次移除都意味着追踪代码中的数据流并针对基线验证输出。没有视觉提示?几何编码器不产生有意义的信号。不需要语义分割?语义分割头可以去掉,但像素解码器保留——它仍然送入实例掩码输出。只需要框和分数而完全不需要掩码?那么整个通用分割头都可以移除。剖析器告诉你每个组件的成本;理解你的用例告诉你哪些可以安全地砍掉。

有了那张地图在手,让我们开始修复。

4. 优化之旅

以下是五类我们找到并修复的问题。五类并不详尽——还有分散各处的小调整,不足以单独成节——但这五类涵盖了最高影响模式,代表了研究代码遇到生产环境时出现的那类问题的典型。

值得在开始前标记一件事:这些类别不是孤立实验的日志。几个修复在同一迭代中落地,所以后期旅程中的"之前"捕获可能已经包含早期类别的更改。将每个前后理解为:这是这个特定问题看起来的样子,这是当我们解决它时发生的情况。聚合数字在第6节。

4.1. 类别1:数据管道

在触碰模型之前,有一件事值得先修复:确保GPU真的在被投喂。两个原因。明显的一个——如果GPU一半时间空闲,我们对模型做的任何事情都几乎不重要。不那么明显的一个——快速的数据管道使每次后续剖析轮次更快更干净,看着GPU利用率在你优化时攀升是一个有用的信号,表明事情正朝着正确方向前进。

问题,正如基线时间线显示的,是一切都顺序运行。主进程从磁盘加载批次,预处理,复制到GPU,运行推理,将结果移回CPU,写入磁盘——然后重复。GPU拿到接力棒,跑完它的一段,然后等待其他所有人完成他们的。

图7: 基线管道——每个阶段阻塞下一个。GPU在每个批次开始的数据加载和结束时的磁盘写入期间空闲。

修复数据端:异步预取

PyTorch的DataLoader已经有异步预取所需的一切——你只需要使用它:

dataloader = DataLoader(
    dataset,
    batch_size=args.batch_size,
    num_workers=WORKERS_DATALOADER,
    pin_memory=True,
    persistent_workers=True,
    in_order=True,
)

num_workers生成后台工作进程,在GPU忙于当前批次时加载和预处理下一批次。pin_memory=True在非分页区域分配主机内存——这启用非阻塞HtoD传输,GPU通过DMA直接拉取数据而无需CPU参与。persistent_workers=True保持工作进程在批次之间存活,而不是每次迭代重新生成。

加载得到了异步处理。写入通常没有。每批次结束时的磁盘写入同样阻塞。主循环等待每个结果写入后才能开始下一批次。同样的问题,管道的另一端。

修复写入端:异步后处理和I/O

修复方法是将后处理和磁盘写入卸载到通过队列连接的专用工作进程。推理循环将结果入队并立即移动到下一批次——它从不等待任何写入。

# spawn required for CUDA + multiprocessing — fork doesn't work
mp.set_start_method("spawn", force=True)

post_write_queue = mp.Queue(maxsize=WORKERS_POSTPROCESSING * 2)
writer_process = mp.Process(target=writer_worker, args=(...), daemon=True)
writer_process.start()

for batch in dataloader:
    batch = copy_data_to_device(batch, device="cuda", non_blocking=True)
    
    output = model(batch)
    boxes, scores = postprocessor(output)  # GPU-side post-processing
    # Blocking copy - the data is needed before enqueuing.
    # After enqueuing, the worker runs in parallel with the next batch on GPU.
    
    post_write_queue.put((batch.text, boxes.cpu(), scores.cpu()))

post_write_queue.put(None)  # Sentinel: signals worker to shut down cleanly
writer_process.join()

几个值得解释的决定。daemon=True意味着工作进程在主进程退出时自动被杀死——不需要手动清理。None哨兵信号工作进程排空队列并停止。

写入工作进程本身在内部使用多线程。磁盘写入通过阻塞系统调用经过操作系统,释放GIL——所以当一个线程等待I/O时,其他线程继续处理已经在队列中的结果。

图8: 优化后的管道——批次N+1的数据加载和批次N-1的后处理/写入与批次N在GPU上的推理并行运行。

GPU现在接近连续运行,下一批次在推理完成时已经等待,结果被卸载而不阻塞主循环。

有一点需要注意:我们在这里将后处理分割在GPU和CPU之间——其中一些仍然在postprocessor()内在GPU上运行,然后结果才入队,其余在工作进程中运行。时间线中那个分割的样子,以及为什么GPU端后处理本身仍有改进空间,是类别4要讨论的内容。

4.2. 类别2:模型中的死代码

研究代码是为了探索而构建的,探索留下痕迹。以前模型版本的遗留组件,在推理时毫无意义训练产物,你永远不会碰的多任务头——随着时间推移,代码库积累。它仍然正确运行。它只是运行得比需要的多。

在SAM3的案例中,三类死代码值得处理。

仅单帧模式。 SAM3支持具有跨帧记忆注意力和对象跟踪器的视频序列。我们不需要任何这些。幸运的是,SAM3已经将视频与单帧的代码路径分开,所以记忆库和跟踪器在单帧模式下根本不执行。但支持两条路径的共享组件保持通用——充满条件判断,更难阅读,更难优化。

仅文本提示。 SAM3支持点、边界框、掩码和文本作为提示类型。我们只使用文本。移除未使用的提示类型比听起来更复杂。未使用的提示并不总是"不运行"——有时它们产生仍然通过torch.cat的零元素张量,有时它们产生注意力掩码全为False的批次条目。计算接触了它们;只是没有贡献任何东西。你必须仔细追踪数据流才能知道什么是安全切割的。

仅推理路径。 模型计算训练期间损失计算需要的中间输出,运行开发时断言,执行只在梯度流动时有意义的检查。在推理时这些都是纯开销——有些,如我们将在类别3中看到的,在此过程中悄悄导致CPU往返。

几何编码器:一个案例研究

有些移除是简单的。其他只有在简化代码足够多以真正看到发生了什么后才变得可见。几何编码器是后者。

编码器的任务是对视觉几何特征进行编码——在完整用例中,来自图像特征和视觉提示(框、点、掩码)的组合。为了理解它是否可以为我们的用例移除,我用仅文本输入追踪了执行:没有框,没有点,没有掩码。在没有提供视觉提示的情况下,简化的前向路径简化为:

cls_mask = torch.zeros(N, 1, dtype=torch.bool, device=device)

返回的键填充掩码全为零——提供没有视觉提示的直接后果。在下游Transformer融合编码器中,注意力应用于文本、几何和视觉提示特征,这意味着几何令牌被关注但除了学习的CLS嵌入外没有携带输入特定的信号。结合visual_prompt_embed是零元素张量——torch.zeros((0, ...)),仍然通过torch.cat传递,没有贡献——很明显几何编码器的输出对我们的用例是结构性惰性的。经验验证确认:移除它对检测结果零影响。

对调用代码的影响是立竿见影的。_encode_prompt之前:

def _encode_prompt(self, backbone_out, img_feats, img_pos_embeds, find_input):
    txt_ids = find_input.text_ids
    txt_feats = backbone_out["language_features"][:, txt_ids]
    txt_masks = backbone_out["language_mask"][txt_ids]
    
    geo_feats, geo_masks = self.geometry_encoder(img_feats=img_feats, ...)
    
    visual_prompt_embed = torch.zeros((0, *geo_feats.shape[1:]), ...)
    visual_prompt_mask = torch.zeros((*geo_masks.shape[:-1], 0), ...)
    
    prompt = torch.cat([txt_feats, geo_feats, visual_prompt_embed], dim=0)
    prompt_mask = torch.cat([txt_masks, geo_masks, visual_prompt_mask], dim=1)
    
    return prompt, prompt_mask

之后:

def _encode_prompt(self, backbone_out, find_input):
    txt_ids = find_input.text_ids
    txt_feats = backbone_out["language_features"][:, txt_ids]
    txt_masks = backbone_out["language_mask"][txt_ids]
    return txt_feats, txt_masks

相同的模式——追踪、简化、发现、移除——应用于语义头、像素解码器和一堆训练产物。累积效应是整个优化中最大的胜利:峰值GPU内存从87GB降到约23GB,这使我们完全摆脱H200,转向更便宜的GPU类别。

关于机制:这一切都不需要触碰原始SAM3源代码。每个修改后的模块都存在于部署特定的子目录中,并代替原始版本导入。上游代码保持完整;更改是隔离的,当模型发布更新时可审计。

4.3. 类别3:绕远路的张量

推理中期的HtoD和DtoH峰值默认可疑。它们并不总是错的——一些数据传输是合法的——但当它们出现在模型前向传递中间时,值得停下来问问为什么。我们的基线有四个,已经标记。

图9: 推理时间线放大——同步停顿、DtoH/HtoD峰值和GPU活动间隙。四个问题已标记。

问题1:分词器

文本分词器在CPU上运行,结果需要移到GPU——没问题。缺少的non_blocking=True使其成为阻塞传输:

# Before
tokenized = self.tokenizer(...).to(device)

# After
tokenized = self.tokenizer(...).to(device, non_blocking=True)

一个词。一个同步停顿消失。

问题2、3和4:没人预料到的GPU张量

这三个是相连的。问题3和4出现在transformer解码器的_get_rpb_matrix中——乍一看没什么问题:

def _get_coords(self, H, W, device):
    coords_h = torch.arange(0, H, device=device) / H  # 问题3:arange使用GPU终止值→同步
    coords_w = torch.arange(0, W, device=device) / W
    return coords_h, coords_w

def _get_rpb_matrix(self, reference_boxes, feat_size):
    H, W = feat_size
    
    if self.compilable_stored_size == (H, W):  # 问题4:与GPU标量比较→同步
        # [...]
    
    self.compilable_cord_cache = self._get_coords(H, W, reference_boxes.device)
    
    # [...]
    assert coords_h.shape == (H,)  # 问题4:使用GPU标量的断言→同步
    assert coords_w.shape == (W,)

HW看起来像是普通整数,但feat_size是GPU张量——所以解包它给你GPU标量。torch.arange以GPU标量为终止值触发CPU往返来评估它。条件比较和断言也是如此。每一个都会停顿主机线程,GPU排空,在同步解决之前无法调度新的内核。

问题就变成了:为什么feat_size是GPU张量?向上游追踪导致问题2,回到编码器:

def _prepare_multilevel_features(...):
    # 问题2:spatial_shapes不必要地放在GPU上
    spatial_shapes = torch.tensor(spatial_shapes, dtype=torch.long, device=src_flatten.device)
    level_start_index = torch.cat((
        spatial_shapes.new_zeros((1,)),
        spatial_shapes.prod(1).cumsum(0)[:-1],
    ))

spatial_shapes被放在GPU上,作为feat_size向下游流动,当它到达解码器时H和W是GPU标量。由于level_start_index在我们的管道中从未使用,GPU张量创建根本不需要——将spatial_shapes保持为普通元组列表打破链条。问题3和4随之消失,作为副作用编码器中的另一个同步点也消失了。

图10: 修复后的相同推理区域——同步停顿消失,DtoH峰值消除,GPU活动稳定。

清理后的时间线中可见的一件事:CPU现在完成后处理的速度足够快,主线程在入队前短暂同步结果的DtoH——一个信号表明瓶颈已从GPU停顿转移到推理和写入工作进程之间的交接。

下次值得记住:当你看到推理中的设备传输时,追踪张量的来源几乎比仅仅修复它浮出水面的那一行更有趣。

4.4. 类别4:后处理开销

在类别1中,我们将后处理移到后台工作进程,这样GPU就不用等待它。但我们忽略了某事:在结果入队之前GPU上发生什么。那部分也不是免费的。

原始后处理的核心看起来像这样:

keep = scores > self.detection_threshold
boxes_out = [b[k] for b, k in zip(boxes, keep)]
scores_out = [s[k] for s, k in zip(scores, keep)]

紧凑、可读,而且相当昂贵。框张量的形状是[批次大小×类别数, 查询数, 4]——对于批次大小16和10个文本提示,这是160个图像-类别对,每对最多200个候选框。对于每对,我们应用布尔掩码以仅保留高于阈值的检测。结果是可变长度张量的列表——每图像-类别对一个,大小由多少检测通过过滤器决定。

图11: 过滤问题:固定形状的建议张量(行=图像-类别对,列=候选框)被逐行过滤成可变长度输出。绿色=保留,橙色=丢弃。

可变输出大小是问题所在。对于Python循环的每次迭代,GPU必须执行完整压缩:标记保留元素(DeviceSelect::Flagged),用包含和计算写索引(DeviceReduce::Sum),将结果收集到新的连续分配中——然后在CPU能确定输出形状并移动到下一次迭代之前执行阻塞DtoH传输。该序列每批次重复160次。

图12: 基线时间线的后处理区域——该阶段由DtoH内存传输和密集的小内核主导,每批次花费53毫秒。

放大到几次迭代详细揭示模式:

图13: 放大到4次循环迭代——每一次触发布尔掩码、计算写索引的包含和、收集到新的连续分配中,以及在下一次迭代开始前阻塞DtoH。

修复方法是停止要求GPU首先产生可变大小输出:

# GPU side — elementwise, fixed shape, no compaction, no sync
keep = scores > self.detection_thresholds
boxes_out = boxes.masked_fill(~keep[..., None], 0)  # [B, C, Q, 4]
scores_out = scores.masked_fill(~keep, float("-inf"))  # [B, C, Q]

# Single DtoH for the whole batch - not per loop iteration
boxes_cpu, scores_cpu = boxes_out.cpu(), scores_out.cpu()

# Variable-size index work moves to CPU, runs in the worker
img_ids, cls_ids, _ = torch.where(scores_cpu != float("-inf"))
for img_id, cls_id in zip(img_ids.tolist(), cls_ids.tolist()):
    ...  # write to results structure

masked_fill是逐元素的——固定输出形状,没有包含和,没有压缩,没有每迭代DtoH。GPU一次性完成整个批次,我们执行单次DtoH传输对(框一个,分数一个),所有可变大小工作移到后台工作进程。当那个工作进程迭代结果时,GPU已经在运行下一批次。

图14: 修复后:37微秒的GPU后处理(masked_fill),单次DtoH传输对,CPU后处理在后台工作进程中运行,而批次7已经在GPU上启动。

53毫秒的GPU后处理带有160次同步停顿→GPU上37微秒+CPU上7毫秒并行运行。从GPU的角度来看,后处理基本上是免费的。

4.5. 类别5:混合精度幽灵

我们运行的是torch.autocast——混合精度已启用。理论上,一切都在bf16中计算,GPU的张量核心很高兴。实践中,并非每个操作都有bf16优化内核。例如,LayerNorm为了数值稳定性回退到fp32。PyTorch静默处理:转换为fp32,运行操作,转换回来。没有警告,没有错误——只有内存往返和一个你没有要求的潜在同步点。

LayerNorm在SAM3中大量使用——跨编码器、解码器和注意力块归一化激活——所以那个模式运行很多。

NVIDIA的transformer_engine库为此提供了融合内核。不是在单独的fp32缓冲区中转换、运行操作、再转换回来,精度转换在内核内部的寄存器中即时发生——没有全局内存往返,没有同步点。而且te.LayerNormnn.LayerNorm的即插即用替代:无需重新训练,相同的预训练权重,只需一次导入更改。

from torch import nn
import transformer_engine.pytorch as te

# Before
self.norm = nn.LayerNorm(d_model)

# After
self.norm = te.LayerNorm(d_model)

transformer_engine还提供融合的Linear层,所以我们也测试了——尽管收益更温和。对于图像编码器,我们还试验了FP8精度。每模块分解讲述故事:

表1: 各模块优化收益分解。

LayerNorm替换在图像编码器上节省30毫秒,在transformer编码器上节省40毫秒。Linear层显示收益递减——那里的精度开销本来就较小。图像编码器的FP8是另一步,尽管在那点上仔细输出验证比早期修复更重要。

完整的优化管道——端到端的一切样子——在下一节。

5. 最终时间线

下面的两个时间线覆盖相同的四个批次,相同配置:16幅图像,每幅图像10个文本提示。时间轴相同。顶部是我们开始的地方。底部是我们最终到达的地方。

图15: 基线(顶部)vs. 优化后(底部)——相同的4个批次,相同比例。总运行时间:17.1秒→2.3秒。峰值内存:87GB→23GB。平均每批次:4263毫秒→564毫秒。

基线有其特征形状:宽的批次,中间有大间隙,每批次末端有嘈杂的后处理纠缠。优化后的时间线相比之下几乎认不出来——批次紧密,背靠背,没有空闲拉伸,也没有我们在前面五节中花费时间追捕的尖峰传输模式。

6. 结果

在默认配置下——16幅图像,10个文本提示——平均每批次推理时间从4263毫秒降到564毫秒。8倍吞吐量,87GB→23GB峰值内存,纯PyTorch。

批次大小扩展

在优化后的管道中扩展批次大小几乎是线性的——批次翻倍,大致墙时间翻倍。这实际上是个好迹象:意味着GPU核心被充分利用,管道不再受数据加载或后处理开销瓶颈。每图像时间保持在所有测试批次大小约35.8毫秒,确认GPU始终饱和。

所有测量都是在4个剖析批次上平均的(跳过前5个进行GPU预热,与第2节的剖析设置一致)。列:每批次时间是一个完整批次的平均墙时间;每图像时间按批次大小归一化;每(图像×提示)时间进一步按提示数归一化,使两个表可直接比较。

表2: 批次大小扩展结果。

提示扩展

更有趣的扩展故事是文本提示。SAM3的图像编码器每幅图像只运行一次,无论你有多少个提示——只有transformer编码器、解码器和后处理随提示数扩展。相对于添加图像,添加提示是便宜的。

这解锁了一个实际的胜利:以前需要多次推理传递的配置——因为一次性放入所有提示会OOM——现在可以在单次传递中运行。N次前向传递变成1次意味着不仅延迟更低,而且编排更简单。

表3: 提示扩展结果。

次线性扩展在每(图像×提示)时间列中可见:随着提示数增长,每图像-提示成本持续下降——从10个提示的3.53毫秒降到100个提示的2.13毫秒。图像编码器做相同的工作;只有较轻的transformer阶段扩展。标题数字:使用优化后的管道每幅图像运行100个提示(3409毫秒)仍然比基线只运行10个(4263毫秒)更快——而且基线根本不能运行100个提示。

内存下降

从87GB到23GB使我们完全摆脱H200。对于大多数生产部署,GPU类别之间的成本差异是实质性的——这可能是立即可 impactful 的结果。

准确性

每个输出都针对基线进行了验证。没有回归。

所有这些都是纯PyTorch——没有ONNX导出,没有TensorRT,没有超过AMP和transformer_engine的量化。对于许多工作负载这就够了。如果你需要进一步推进,TensorRT是自然的下一步——但在那一点上,你从一个比我们开始好得多的基线出发。

关键要点

  • 先修复明显的管道问题——异步数据加载和异步写入。一旦这些清理干净,剖析器对找到真正剩余的内容就变得更有用。
  • 足够了解你的用例以手术式修改模型。 不只是读论文——而是知道"仅文本提示,单帧"对前向传递中每个模块意味着什么,并有信心砍掉不服务的部分。
  • 推理中意外的DtoH和HtoD传输总是可疑的。 它们通常是上游某处设备不匹配的症状——追踪张量回到其源头,而不仅仅是它浮出水面的地方。
  • 热循环中的可变输出形状是昂贵的。 布尔掩码聚集强制包含和、压缩内核和每迭代DtoH——每批次160次。固定形状masked_fill将其压缩到37微秒。
  • AMP不意味着所有操作都在低精度中运行。 LayerNorm静默回退到fp32。transformer_engine的即插即用替代在所有三个模块中恢复约70毫秒——一次导入交换,无需重新训练。

原文链接: We Made SAM3 8x Faster for Production — Here's What the Profiler Actually Showed

汇智网翻译整理,转载请标明出处