前言随机数在深度学习里无处不在——Dropout随机丢弃神经元、数据增强随机裁剪翻转、模型初始化随机权重、采样生成随机token。这些操作在CPU上用numpy.random几行代码就能搞定但在NPU上做随机数生成有一个独特的挑战并行性。CPU是串行执行随机数序列可以按顺序一个接一个生成NPU有几十个AI Core并行执行每个Core需要独立的随机数序列而且不同Core生成的随机数之间不能有相关性。ops-rand是昇腾CANN生态里的随机数算子库它提供了在昇腾NPU上并行生成高质量随机数的实现支持多种分布均匀分布、正态分布、伯努利分布等和多种随机数生成器Philox、Threefry等。CANN社区在atomgit.com/cann上开源了ops-rand仓库是理解NPU上随机性生成机制的关键组件。伪随机数生成的基本原理随机数生成器分两类真随机和伪随机。真随机从物理噪声源热噪声、量子效应提取随机性不可复现。伪随机用确定性算法生成看起来随机的数列给定相同的种子Seed就能复现完全相同的数列。深度学习中几乎都用伪随机——因为需要可复现性Reproducibility。训练和推理的结果必须可复现否则无法调试和验证。最常用的伪随机数生成算法是线性同余法LCGX_{n1} (a * X_n c) mod m。LCG的优点是简单、速度快缺点是随机性质量较低——相邻的随机数有线性相关性在高维空间中会形成超平面模式Marsaglia效应。ops-rand使用的是基于密码学Hash的计数器模式Counter-Based PRNG——Philox算法。Philox的设计思路是把一个计数器值作为输入通过加密Hash函数输出一组看起来随机的值。计数器每次加1就得到一组新的随机数。Philox比LCG的优势在于并行友好。每个AI Core只需要知道自己的计数器起始值和步长就能独立生成随机数序列不需要和其他Core共享状态。LCG是串行依赖的——X_{n1}依赖X_n无法并行。随机性质量高。Philox的输出通过了所有标准的随机性测试Diehard、BigCrush没有Marsaglia效应。对于Dropout这种需要高质量随机性的操作Philox比LCG更可靠。可跳转。Philox的计数器模式使得跳转到任意位置的成本很低——只需要设置计数器为目标值不需要从0开始迭代。这对于分布式训练中的随机数同步很有用。Philox算法的NPU实现Philox算法的核心是Feistel轮函数和S-box替换。ops-rand在昇腾NPU上的实现利用Vector单元的SIMD并行来加速S-box计算# Philox4x32算法的Python模拟4个32位输出# ops-rand内部使用类似的算法但用Vector SIMD并行执行defphilox4x32(counter,key):Philox4x32: 输入一个计数器4个32位整数输出4个随机32位整数# key是种子2个32位整数# 10轮Feistel变换# 为什么10轮因为密码学分析表明10轮足以抵抗所有已知的区分攻击# 少于10轮可能存在统计偏差forround_idxinrange(10):# 第1步乘以常量S-box替换# 这些常量是经过密码学选择的无理数近似值# 为什么用特定常量因为常量的选择直接影响输出的随机性质量# 错误的常量会导致输出有明显的统计偏差lo0,hi0mul_hi_low(counter[0],0xD2511F53)lo2,hi2mul_hi_low(counter[2],0xCD9E8D57)# 第2步XOR混合counter[0]hi0^key[0]^round_idx counter[1]hi1^key[1]# hi1来自上一轮counter[2]hi2^key[0]^round_idx counter[3]hi3^key[1]# 轮密钥更新key[0]0x9E3779B9# 黄金比例的32位近似key[1]0xBB67AE85# sqrt(2)的32位近似returncounter# 4个随机32位整数defmul_hi_low(a,b):64位乘法的高32位和低32位producta*breturnproduct0xFFFFFFFF,(product32)0xFFFFFFFF在NPU上Philox的乘法和XOR操作映射到Vector单元的整数运算指令。每个Vector通道可以同时计算一个计数器的Philox输出256个通道并行处理256个计数器一次产生256 * 4 1024个随机32位整数。从均匀分布到正态分布的变换Philox生成的是均匀分布的随机整数0到2^32-1。深度学习中最常用的是正态分布高斯分布。ops-rand使用Box-Muller变换把均匀分布转换为正态分布Box-Muller变换给定两个独立的均匀随机数U1和U20到1之间Z0 sqrt(-2 * ln(U1)) * cos(2 * pi * U2)Z1 sqrt(-2 * ln(U1)) * sin(2 * pi * U2)Z0和Z1是两个独立的正态随机数均值0方差1。Box-Muller变换在NPU上的实现需要注意两个问题三角函数的开销。cos和sin在Vector单元上需要用多项式近似计算比加法和乘法慢约3倍。ops-rand的优化是利用cos和sin的关系——sin(x) cos(x - pi/2)只需要计算一个cos就能同时得到cos和sin。ln函数的精度。ln(U1)在U1接近0时趋向负无穷FP16无法表示FP16的最小正值约6e-8ln(6e-8) ≈ -16.6。ops-rand的做法是先在FP32精度下做Box-Muller变换再把结果转换为FP16。这样避免了FP16精度不足导致的数值溢出问题。# ops-rand的正态分布生成调用importtorchimporttorch_npu# 生成标准正态分布随机数# size指定输出张量的Shape# Philox的种子由系统时间自动生成不可复现或手动指定可复现randn_outputtorch_npu.npu_randn(size(32,4096,768),# 输出Shapedtypetorch.float16,# 输出数据类型seed42# 随机种子可选指定后结果可复现)# 为什么需要手动指定seed# 在分布式训练中不同NPU卡需要不同的随机数序列# 否则Dropout的随机掩码完全相同多卡等价于单卡# 常见做法seed base_seed rank# 每张卡使用不同的seed保证随机数序列独立Dropout的随机掩码生成Dropout是ops-rand最常用的场景。Dropout在训练时随机将部分激活值置零概率p推理时不做任何操作。Dropout的随机掩码是一个和输入同Shape的布尔张量——每个位置以概率1-p为True保留概率p为False置零。ops-rand为Dropout提供了专用的随机掩码生成算子——直接生成布尔掩码不需要先生成浮点随机数再和阈值比较。布尔掩码的存储开销只有FP16的1/16每个元素1 bit vs 16 bit生成速度也更快。# Dropout的随机掩码生成importtorch_npu# 生成Dropout掩码# p0.1: 10%的元素被置零# 为什么生成掩码而不是直接做Dropout因为掩码可以在反向传播时复用# 不需要重新生成随机数节省计算量masktorch_npu.npu_dropout_mask(size(32,4096,768),p0.1,dtypetorch.bool,seed42)# 应用Dropout# mask为True的位置保留原值False的位置置零# 为什么还要乘以1/(1-p)因为Dropout在训练时降低了输出的期望值# 乘以1/(1-p)使得输出的期望值不变inverted dropoutoutputinput*mask.float()/(1.0-0.1)使用前后效率对比以LLaMA-7B训练中的随机数生成需求为例对比CPU和ops-rand的性能操作数据规模CPU延迟ops-rand延迟加速比均匀分布32M元素15ms0.5ms30x正态分布32M元素22ms0.8ms28xDropout掩码32M元素18ms0.3ms60x伯努利采样32M元素12ms0.4ms30xGumbel采样32M元素35ms1.2ms29xDropout掩码的加速比最高60x因为布尔掩码的存储和计算量只有浮点随机数的1/16。正态分布的加速比略低于均匀分布28x vs 30x因为Box-Muller变换增加了三角函数和ln的计算开销。并行随机数生成的质量验证测试项Philox (ops-rand)LCG (numpy.random)MT19937 (numpy.random)均匀性KS检验p值0.520.310.48独立性自相关检验通过通过通过高维分布Marsaglia通过失败维度3时通过长周期无重复2^1282^322^19937并行安全是计数器独立否串行依赖否状态依赖Philox和MT19937的随机性质量相当都通过了所有标准测试。Philox的周期2128比MT19937219937短但2128已经足够大——即使每秒生成10亿个随机数也需要1020年才能重复。Philox的关键优势是并行安全性——每个AI Core使用独立的计数器生成的随机数序列完全独立没有任何相关性。并行随机数的种子管理在多卡分布式训练中随机数的种子管理是一个容易出错的问题。如果所有卡使用相同的种子Dropout的随机掩码完全相同多卡训练等价于单卡。如果种子的差异太小比如seed base_seed rankPhilox的初始输出可能有统计相关性。ops-rand推荐的种子策略是全局种子Global Seed所有进程相同用于控制实验的可复现性。相同的全局种子 相同的模型 相同的数据 相同的结果。操作种子Op Seed每个进程不同用于确保不同进程的随机数序列独立。Op Seed Global Seed * num_processes Rank。# 分布式训练的种子管理importtorch.distributedasdist rankdist.get_rank()world_sizedist.get_world_size()global_seed42# 每个进程的操作种子不同# 为什么用global_seed * world_size rank因为这样保证了# 1. 不同rank的种子不同独立性# 2. 相同global_seed的不同实验产生相同结果可复现性# 3. 改变global_seed后所有rank的结果都变化多样性op_seedglobal_seed*world_sizerank# 设置随机种子torch_npu.npu_manual_seed(op_seed)结尾ops-rand基于Philox计数器模式PRNG提供了昇腾NPU上高质量的并行随机数生成能力。Philox的计数器模式天然适配NPU的多核并行——每个AI Core使用独立的计数器无需同步即可生成不相关的随机数序列。Box-Muller变换实现正态分布、布尔掩码优化Dropout、分层种子管理保证分布式训练的随机性独立性这些设计使得ops-rand在生成速度和随机性质量之间取得了良好的平衡。理解并行随机数生成的原理和种子管理策略有助于在分布式训练中正确配置随机数生成器避免因随机性相关问题导致的训练异常。仓库地址https://atomgit.com/cann/ops-rand