状态: 活跃（处于积极开发中，可能会有破坏性变更）

Blocksparse

blocksparse 包含用于块稀疏矩阵乘法的 TensorFlow 操作及相应的 GPU 内核。此外，还包含与之相关的操作，如边偏置、稀疏权重归一化和层归一化。

欲了解更多信息，请参阅 OpenAI 博客上的发布文章。

前提条件

首先，您需要至少一块 Nvidia GPU。为获得最佳性能，我们建议使用 Pascal 或 Maxwell 架构的 GPU——以下是按 GPU 类型划分的功能列表：

GPU 系列	BSMatMul-ASM	BSMatMul-CudaC	BSConv
Kepler	-	X	-
Maxwell	X（最快）	X	X
Pascal	X（最快）	X	X
Volta	-	X（最快）	-

请注意，BSMatMul-CudaC 仅支持 feature_axis=0，而 BSMatMul-ASM 仅支持 feature_axis=1。

此外，您还需要：

• 一个可用的 Linux 系统（我们使用 Ubuntu 16.04），并安装适用于您的 GPU 的 Nvidia 驱动程序。
• CUDA 8（位于 /usr/local/cuda）。
• Python 3.5 或更高版本，或 2.7 或更高版本。
• TensorFlow 1.4.0 或更高版本，并具备 GPU 支持（例如：pip install tensorflow-gpu）。
• 如果您使用 CUDA 9 和 Volta 架构的 GPU，需更新构建目标（-gencode=arch=compute_70,code=sm_70），并从源码编译 TensorFlow。

安装

pip install blocksparse

使用方法

以下示例执行了一个块稀疏矩阵乘法：

from blocksparse.matmul import BlocksparseMatMul
import tensorflow as tf
import numpy as np

hidden_size = 4096
block_size = 32
minibatch_size = 64

# 创建一个（随机）稀疏模式
sparsity = np.random.randint(2, size=(hidden_size//block_size,hidden_size//block_size))

# 初始化稀疏矩阵乘法对象
bsmm = BlocksparseMatMul(sparsity, block_size=block_size)

# 图中的输入
x = tf.placeholder(tf.float32, shape=[None, hidden_size])

# 初始化块稀疏权重
w = tf.get_variable("w", bsmm.w_shape, dtype=tf.float32)

# 执行块稀疏矩阵乘法
y = bsmm(x, w)

# 运行
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
result = sess.run([y], feed_dict = {x: np.ones((minibatch_size,hidden_size), dtype='float32')})
print(result)

有关使用块稀疏操作训练语言模型的更复杂示例，请参阅 examples/。

开发

如果您有兴趣对这些操作和内核进行修改，可以直接从源码构建：

git clone git@github.com:openai/blocksparse.git
cd blocksparse

make compile
pip install dist/*.whl

# 如果需要，可以运行测试
test/blocksparse_matmul_test.py
test/blocksparse_conv_test.py

如果您的 CUDA 不在 /usr/local/cuda 目录下，或者同时安装了多个版本，例如 /usr/local/cuda-8.0 和 /usr/local/cuda-9.0，请将 CUDA_HOME 设置为编译时使用的基础路径，以便在执行 make compile 时正确引用。

API 文档：

blocksparse.matmul

class BlocksparseMatMul(object)

    def __init__(self, layout, block_size=32, feature_axis=1)
    """
    layout：一个由 1 和 0 组成的二维数组，用于指定块布局。
    block_size：支持 32、16 和 8。
    feature_axis：当 block_size 小于 32 时，采用 (C,N) 激活布局会使内存访问效率显著提高。
    """

    # 用于生成张量的形状辅助函数（N 表示小批量大小）
    self.w_shape
    def i_shape(self, N)
    def o_shape(self, N)

    # 根据给定的块 ID 返回布局中对应的 (c,k) 坐标
    def block_coord(self, block)

    # 实验性正交初始化
    def ortho_init(self)

    # 在实践中，只需使用 identity_init 加上层归一化即可完成初始化。
    # 当 gpu=True 时，初始化操作由设备上的内核直接执行。
    def identity_init(self, gpu=False)

    # 用于实现权重归一化。但在实际应用中，层归一化效果更好。
    def l2_normalize(self, W，gain=None，epsilon=1e-6，dtype=np.float32)

    def __call__(self, I，W，dw_dtype=tf.float32)
    """
    执行该操作。请注意，权重变量与 bsmm 对象是独立的，这允许多个权重绑定到同一个 bsmm 布局。

    dw_dtype：用于控制 dw 的精度格式。
    """


def group_param_grads(param_grad, group_size=8, cast32=True)
"""
param_grad：给定 bsmm 权重变量的 TensorFlow 参数梯度（由 tf.gradients 返回）。
group_size：期望的分组大小，最多支持 8 组。

此函数会重写 TensorFlow 图，将来自不同时间步的权重梯度矩阵乘法（以及跨时间共享的权重）合并为一个更高效的单次矩阵乘法。
"""


class SparseProj(object):
    def __init__(self, nhidden, nproj=None，proj_stride=None，block_size=32，gather_lut=None)
    """
    实验性类，用于支持稠密到稀疏以及稀疏到稠密的投影操作。其功能与 TensorFlow 操作基本相同，但速度更快，并支持其他精度格式。此类假设存在唯一的 1 对 1 映射，因此在反向传播操作中无需使用原子操作。
    """

    def gather(self, x)
    def scatter(self, x)
    def scatter_add(self, x，y)
    def scatter_mul(self, x，y)

blocksparse.conv

class BlocksparseConv(object):
    def __init__(self, BCK, TRS, DHW, MPQ=None, strides=(1,1,1), dilates=(1,1,1), padding="SAME", edge_bias=False)
    """
    BCK: (                                             # 块(B)/输入(C)/输出(K)特征维度
             ( (c0, c1, c2, ...), (k0, k1, k2, ...) ), # 块 0 的 c,k 是 C,K 维度的索引
             ( (c0, c1, c2, ...), (k0, k1, k2, ...) ), # 块 1
             ( (c0, c1, c2, ...), (k0, k1, k2, ...) ), # 块 2 ...
         )
    TRS: (T,R,S) 或 (R,S) 或 (S,)         - 卷积核的空间尺寸维度
    DHW: (D,H,W) 或 (H,W) 或 (W,)         - 输入图像的空间尺寸维度
    MPQ: (M,P,Q) 或 (P,Q) 或 (Q,) 或 None - 输出图像的空间尺寸维度（用于步幅转置卷积中存在歧义的情况）
    strides: (1,1,1) 或 (1,1) 或 (1,)
    dilates: (1,1,1) 或 (1,1) 或 (1,)
    padding: (1,1,1) 或 (1,1) 或 (1,) 或 "SAME" 或 "VALID"
    edge_bias: True/False
    """

    # 用于设置变量或测试张量的形状辅助函数
    def edge_bias_shape(self)
    def f_shape(self, block=None)
    def i_shape(self, N)
    def o_shape(self, N)

    # 执行操作，传入参数变量和输入
    def __call__(vs, F, I, edge_bias=None):

    # 用于实现权重归一化
    def l2_normalize(self, F, gain=None, epsilon=1e-6, dtype=np.float32):

class BlocksparseDeconv(BlocksparseConv)
    def __init__(self, BCK, TRS, DHW, MPQ=None, strides=(1,1,1), dilates=(1,1,1), padding="SAME", edge_bias=False)
    """
    反卷积。参数与上述相同。
    """

def cwise_linear(x, a=None, b=None)
"""
在 NCHW 张量格式下，TensorFlow 对中间 C 维度上的简单广播操作实现得非常缓慢。
这个函数可以用来执行：
    y = a*x + b
    y = a*x
    y = x + b

其中 a 和 b 的形状为 (1,C,1,1)。
这对于权重归一化等操作非常有用。

blocksparse.ew

# 与 tf 操作相同，但通常更高效，并支持自定义精度格式
def        add(x, y, name=None)
def   multiply(x, y, name=None)
def   subtract(x, y, name=None)
def     divide(x, y, name=None)
def    maximum(x, y, name=None)
def    minimum(x, y, name=None)

def   negative(x,    name=None)
def reciprocal(x,    name=None)
def     square(x,    name=None)
def       sqrt(x,    name=None)
def        exp(x,    name=None)
def        log(x,    name=None)
def    sigmoid(x,    name=None)
def       tanh(x,    name=None)
def       relu(x,    name=None)
def       elu (x, alpha=1.0, name=None)

# 这里参数可以是 4 个独立的门控张量，或者
# 一个合并后的门控张量（内部会拆分成 4 个）
def fused_lstm_gates(c, *args, name=None)

def split4(x)
def concat4(x0, x1, x2, x3)

# 一个自定义的类型转换操作，用于探索新的精度格式
def float_cast(x, dtype, dx_dtype=None)

# 一个更快（且非确定性）的 dropout 操作
# 同时也支持新的精度格式
def dropout(x, keep_prob=0.8, mask=None)

# 一个在 tf.gradients 中使用，用于将多个梯度贡献相加的操作。
# 注意，只支持最多 8 个输入，因为不希望单个操作在图中开始执行之前就消耗掉所有可能的输入，
# 从而减少内存占用。
def add_n8(xs, name=None)

blocksparse.norms

def layer_norm(x, g, b, axis=1, epsilon=1e-6, relu=False)
"""
非常快速的层归一化，支持两种 bsmm 特征轴激活布局。
同时还包括可选的集成 ReLU（在最后应用）。
"""

# NCHW 布局下的基本批归一化操作
def batch_norm(x, g, b, epsilon=1e-6)
def batch_norm_inference(x, g, b, m, v, epsilon=1e-6)

Blocksparse 快速上手指南

blocksparse 是 OpenAI 开源的 TensorFlow 扩展包，提供了针对块稀疏矩阵乘法（Block-Sparse Matrix Multiplication）优化的 GPU 算子及内核。它旨在显著提升稀疏神经网络在 GPU 上的训练和推理速度。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux（推荐 Ubuntu 16.04 或更高版本）。
• 硬件要求：至少一张 NVIDIA GPU。
  • 推荐架构：Maxwell 或 Pascal 系列（性能最佳，支持所有特性）。
  • 其他支持：Kepler（部分支持）、Volta（需特殊编译配置）。
• 软件依赖：
  • NVIDIA 显卡驱动（已正确安装）。
  • CUDA Toolkit 8.0（默认路径 /usr/local/cuda）。
    • 注：若使用 Volta 架构及 CUDA 9.0+，需从源码编译 TensorFlow 并修改构建目标。
  • Python 3.5+ 或 2.7+。
  • TensorFlow 1.4.0+（必须包含 GPU 支持，即 tensorflow-gpu）。

安装步骤

1. 安装 TensorFlow GPU 版

如果您尚未安装支持 GPU 的 TensorFlow，请使用以下命令（推荐使用国内镜像源加速）：

pip install tensorflow-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装 Blocksparse

直接通过 pip 安装官方发布的 wheel 包：

pip install blocksparse -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：如果您的 CUDA 不在默认路径 /usr/local/cuda，或者需要自定义编译选项，请参考原文档 "Development" 部分进行源码编译。

基本使用

以下是一个最简单的示例，演示如何创建一个块稀疏矩阵乘法对象并执行计算。

from blocksparse.matmul import BlocksparseMatMul
import tensorflow as tf
import numpy as np

# 定义超参数
hidden_size = 4096
block_size = 32
minibatch_size = 64

# 1. 创建随机稀疏模式 (0 和 1 组成的二维数组)
sparsity = np.random.randint(2, size=(hidden_size//block_size, hidden_size//block_size))

# 2. 初始化块稀疏矩阵乘法对象
bsmm = BlocksparseMatMul(sparsity, block_size=block_size)

# 3. 定义输入占位符
x = tf.placeholder(tf.float32, shape=[None, hidden_size])

# 4. 初始化块稀疏权重变量
w = tf.get_variable("w", bsmm.w_shape, dtype=tf.float32)

# 5. 执行块稀疏矩阵乘法
y = bsmm(x, w)

# 6. 运行会话
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())

# 构造测试数据并运行
result = sess.run([y], feed_dict={x: np.ones((minibatch_size, hidden_size), dtype='float32')})

print(result)

核心 API 简述

• BlocksparseMatMul(layout, block_size, feature_axis): 核心类。layout 定义稀疏结构，block_size 支持 32/16/8。
• __call__(I, W): 执行乘法操作，其中 I 为输入激活，W 为权重。
• 辅助功能: 包内还包含优化的 layer_norm、conv（卷积）、ew（逐元素运算）等算子，可直接替换标准 TensorFlow 算子以提升性能。

更多复杂示例（如语言模型训练）可参考项目源码中的 examples/ 目录。