Pytorch 学习笔记(4) : torch.backends

PyTorch torch.backends 模块详解

一、模块概述

torch.backends 是 PyTorch 中用于控制各种后端行为的模块，支持 CPU、CUDA、cuDNN、MPS、MKL、OpenMP 等多种硬件和库后端。

二、各后端详细介绍

1. torch.backends.cpu - CPU 后端

函数/属性	说明
`get_cpu_capability()`	获取 CPU 能力，返回值：`DEFAULT`、`VSX`、`Z VECTOR`、`NO AVX`、`AVX2`、`AVX512`、`SVE256`

2. torch.backends.cuda - CUDA 后端

函数/属性	说明
`is_built()`	检查 PyTorch 是否以 CUDA 支持编译
`matmul.allow_tf32`	控制是否使用 TensorFloat-32（Ampere 及以上 GPU）
`matmul.allow_fp16_reduced_precision_reduction`	控制 fp16 GEMM 是否允许降低精度归约
`matmul.allow_bf16_reduced_precision_reduction`	控制 bf16 GEMM 是否允许降低精度归约
`cufft_plan_cache`	cuFFT 计划缓存管理（size/max_size/clear）
`preferred_blas_library()`	选择 BLAS 库：cuBLAS/cuBLASLt/CK
`preferred_rocm_fa_library()`	ROCm 下选择 Flash Attention 后端
`preferred_linalg_library()`	选择线性代数库：cuSOLVER/MAGMA

SDPA (Scaled Dot Product Attention) 相关：

函数	功能
`flash_sdp_enabled()` / `enable_flash_sdp()`	Flash Attention 开关
`mem_efficient_sdp_enabled()` / `enable_mem_efficient_sdp()`	内存高效注意力开关
`math_sdp_enabled()` / `enable_math_sdp()`	数学实现注意力开关
`cudnn_sdp_enabled()` / `enable_cudnn_sdp()`	cuDNN 注意力开关
`sdp_kernel()`	上下文管理器，临时启用/禁用指定后端
`is_flash_attention_available()`	检查 FlashAttention 是否可用
`can_use_flash_attention()`	检查给定参数是否能使用 FlashAttention
`can_use_efficient_attention()`	检查是否能使用 efficient attention
`can_use_cudnn_attention()`	检查是否能使用 cuDNN attention

3. torch.backends.cudnn - cuDNN 后端

函数/属性	说明
`version()`	返回 cuDNN 版本
`is_available()`	检查 cuDNN 是否可用
`enabled`	控制 cuDNN 是否启用
`allow_tf32`	控制 TensorFloat-32 使用
`deterministic`	是否仅使用确定性卷积算法
`benchmark`	是否对多种卷积算法进行基准测试
`benchmark_limit`	基准测试时尝试的最大算法数量

4. torch.backends.cusparselt - cuSPARSELt 后端

函数	说明
`version()`	返回 cuSPARSELt 版本
`is_available()`	检查 cuSPARSELt 是否可用

5. torch.backends.mha - 多头注意力后端

函数	说明
`get_fastpath_enabled()`	获取 TransformerEncoder/MHA 快速路径是否启用
`set_fastpath_enabled()`	设置快速路径启用状态

6. torch.backends.miopen - MIOpen 后端 (ROCm)

属性	说明
`immediate`	控制是否使用 Immediate Mode

7. torch.backends.mps - Apple MPS 后端

函数	说明
`is_available()`	检查 MPS 是否可用
`is_built()`	检查是否以 MPS 支持编译

8. torch.backends.mkl - Intel MKL 后端

函数	说明
`is_available()`	检查是否以 MKL 支持编译
`verbose(level)`	按需启用 MKL 详细日志（上下文管理器）

日志级别： VERBOSE_OFF、VERBOSE_ON

9. torch.backends.mkldnn - oneDNN 后端

函数	说明
`is_available()`	检查是否以 MKL-DNN 支持编译
`verbose(level)`	按需启用 oneDNN 详细日志

日志级别： VERBOSE_OFF、VERBOSE_ON、VERBOSE_ON_CREATION

10. torch.backends.nnpack - NNPACK 后端

函数	说明
`is_available()`	检查是否以 NNPACK 支持编译
`flags()`	上下文管理器设置 NNPACK 启用状态
`set_flags()`	全局设置 NNPACK 启用状态

11. torch.backends.openmp - OpenMP 后端

函数	说明
`is_available()`	检查是否以 OpenMP 支持编译

12. torch.backends.opt_einsum - opt_einsum 优化后端

函数/属性	说明
`is_available()`	检查 opt_einsum 是否可用
`get_opt_einsum()`	获取 opt_einsum 包（如果可用）
`enabled`	控制是否启用 opt_einsum（默认 True）
`strategy`	设置收缩策略：`auto`（默认）、`greedy`、`optimal`

安装方式：pip install torch[opt-einsum] 或 pip install opt-einsum

13. torch.backends.xeon - Intel Xeon 后端

文档中未列出具体 API。

三、实用代码示例

python 复制代码

import torch

# 1. 查看 CPU 能力
print(torch.backends.cpu.get_cpu_capability())  # 如: "AVX512"

# 2. 检查 CUDA 相关
if torch.backends.cuda.is_built():
    print(f"CUDA 可用: {torch.backends.cuda.is_available()}")
    print(f"cuDNN 版本: {torch.backends.cudnn.version()}")

# 3. 控制 cuDNN 行为
torch.backends.cudnn.benchmark = True  # 自动寻找最优算法
torch.backends.cudnn.deterministic = True  # 使用确定性算法

# 4. 控制 TF32（Ampere 及以上 GPU）
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# 5. SDPA 后端控制
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)

# 使用上下文管理器临时切换
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
    # 在此范围内只使用 Flash Attention
    output = torch.nn.functional.scaled_dot_product_attention(q, k, v)

# 6. MKL/oneDNN 详细日志
with torch.backends.mkl.verbose(torch.backends.mkl.VERBOSE_ON):
    # 只在此范围内输出 MKL 详细日志
    model(data)

with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON):
    # 只在此范围内输出 oneDNN 详细日志
    model(data)

# 7. 优化 einsum
torch.backends.opt_einsum.enabled = True
torch.backends.opt_einsum.strategy = "greedy"  # 或 "optimal", "auto"

四、环境变量速查

环境变量	作用
`TORCH_BLAS_PREFER_CUBLASLT=1`	优先使用 cuBLASLt
`TORCH_LINALG_PREFER_CUSOLVER=1`	优先使用 cuSOLVER
`TORCH_ROCM_FA_PREFER_CK=1`	ROCm 下优先使用 CK 作为 Flash Attention 后端
`MKL_VERBOSE`	全局启用 MKL 详细日志
`DNNL_VERBOSE`	全局启用 oneDNN 详细日志

五、总结

torch.backends 模块是 PyTorch 性能调优的重要工具，通过它可以：

查询硬件能力 - 了解 CPU/GPU 支持特性
控制计算精度 - TF32、fp16/bf16 归约精度
选择计算后端 - cuBLAS/cuBLASLt、cuSOLVER/MAGMA 等
优化注意力计算 - Flash Attention、Memory Efficient Attention 等
调试性能问题 - MKL/oneDNN 按需详细日志！