为何及何时需要手动使用torch.cuda.Stream()？相关最佳实践有哪些？

PyTorch官方提到的「通常不需要显式创建流」是针对通用场景的表述，默认流的串行执行逻辑足够满足大部分常规训练、推理需求。手动创建自定义流的核心价值是打破默认流的串行限制，实现无依赖CUDA操作的并行执行，提升GPU资源利用率。

核心适用场景 #

• 数据流水线重叠：这是最常见的使用场景，像问题示例中拆分输入、模型、输出三个流的写法，就是为了让「下一批次数据从CPU拷贝到GPU」「当前批次模型计算」「上一批次结果回拷/后处理」三个无依赖的操作并行执行，完全消除GPU等待数据IO的空闲时间，在小模型推理、数据预处理开销占比高的场景下，能直接将吞吐量提升30%以上。
• 独立模型分支并行：如果你的模型存在多个无数据依赖的计算分支，或者单卡上同时跑多个互不影响的任务，可以把不同分支/任务放到不同流中执行，让GPU同时调度多个独立计算任务，避免单流串行导致的算力浪费。
• 避免全局同步阻塞：默认流会和同设备上的所有其他流隐式同步，如果你有大张量初始化、批量数据拷贝这类高开销但不影响主流程的操作，可以放到自定义流中执行，避免这类操作卡住主训练/推理链路的正常执行。

最佳实践 #

• 「先定位瓶颈再加流」：不要在开发初期就盲目引入自定义流，首先用PyTorch Profiler或者Nsight Systems工具确认当前GPU利用率低、存在明显的空闲等待窗口，再针对性用自定义流做并行优化，避免无谓增加代码复杂度和调试成本。
• 必须显式处理数据依赖：如果两个流中的操作存在数据依赖（比如A流的输出要作为B流的输入），必须通过torch.cuda.Event或者wait_stream()方法显式同步，否则会出现数据竞争，导致计算结果异常且极难排查。参考示例：

# 输入流加载数据、模型流执行计算的配合逻辑
input_stream = torch.cuda.Stream()
model_stream = torch.cuda.current_stream()

with torch.cuda.stream(input_stream):
    # 异步加载下一批数据到GPU
    next_batch = next(data_iter).to('cuda', non_blocking=True)

# 模型流等待输入流数据加载完成后再执行计算
model_stream.wait_stream(input_stream)
with torch.cuda.stream(model_stream):
    output = model(next_batch)

• 控制流的数量：每个自定义流都有额外的调度和显存开销，单卡同时活跃的流不要超过8个，绝大多数场景下2~3个流即可覆盖优化需求，过多的流反而会导致调度开销上升，抵消并行收益。
• 优先使用torch.cuda.stream()上下文管理器切换流，不要手动调用stream.set_current()切换，避免忘记切回默认流导致其他逻辑异常。

内容的提问来源于stack exchange，提问作者aerin