XLA中生成Causal Mask上三角-inf矩阵

transformers生成CausalAttentionMask的上三角-inf矩阵:

参考transformers源码

python 复制代码
import torch
import torch_xla
import torch_xla.core.xla_model as xm
import os

os.environ['PJRT_DEVICE']='IPU'
# os.environ['PJRT_DEVICE']='GPU'
# os.environ['XLA_FLAGS']='--xla_dump_to=gen_AttnFwd-XLA_GPU'

tgt_len = 10
dtype=torch.float32
device = xm.xla_device()

# src/transformers/modeling_attn_mask_utils.py#AttentionMaskConverter::_make_causal_mask
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
mask_cond = torch.arange(mask.size(-1), device=device)
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
mask = mask.to(dtype)
print(mask)
# print(mask.size())
# print(mask[3][3])

"""
2024-11-07 07:16:18.824506: F tensorflow/compiler/xla/service/hlo_computation.cc:70] Check failed: nullptr != root (nullptr vs. 0)
Aborted (core dumped)
"""

'''
module @SyncTensorsGraph.25 {
  func.func @main() -> tuple<tensor<10x10xf32>> {
    %0 = mhlo.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]> : tensor<10xi64>
    %1 = "mhlo.broadcast_in_dim"(%0) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<10xi64>) -> tensor<10x10xi64>
    %2 = mhlo.constant dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]> : tensor<10xi64>
    %3 = "mhlo.broadcast_in_dim"(%2) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<10xi64>) -> tensor<10x10xi64>
    %4 = mhlo.compare  LT, %1, %3 : (tensor<10x10xi64>, tensor<10x10xi64>) -> tensor<10x10xi1>
    %5 = mhlo.constant dense<false> : tensor<i1>
    %6 = "mhlo.broadcast_in_dim"(%5) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<i1>) -> tensor<10x10xi1>
    %7 = mhlo.compare  NE, %4, %6 : (tensor<10x10xi1>, tensor<10x10xi1>) -> tensor<10x10xi1>
    %8 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %9 = "mhlo.broadcast_in_dim"(%8) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<10x10xf32>
    %10 = mhlo.constant dense<-3.40282347E+38> : tensor<f32>
    %11 = "mhlo.broadcast_in_dim"(%10) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<10x10xf32>
    %12 = "mhlo.select"(%7, %9, %11) : (tensor<10x10xi1>, tensor<10x10xf32>, tensor<10x10xf32>) -> tensor<10x10xf32>
    %13 = "mhlo.tuple"(%12) {xla_shape = "(f32[10,10]{1,0})"} : (tensor<10x10xf32>) -> tuple<tensor<10x10xf32>>
    return %13 : tuple<tensor<10x10xf32>>
  }
}
'''

'''
XLA_GPU甚至给出了完整的mhlo实现:
gen_AttnFwd-XLA_GPU/module_0000.SyncTensorsGraph.25.sm_8.0_gpu_after_optimizations.txt

HloModule SyncTensorsGraph.25, entry_computation_layout={(f32[])->(f32[10,10]{1,0})}

fused_computation {
  iota.3 = s64[10,10]{1,0} iota(), iota_dimension=1
  iota.2 = s64[10]{0} iota(), iota_dimension=0
  constant_5 = s64[] constant(1)
  broadcast.7 = s64[10]{0} broadcast(constant_5), dimensions={}
  add.0 = s64[10]{0} add(iota.2, broadcast.7)
  broadcast.6 = s64[10,10]{1,0} broadcast(add.0), dimensions={0}
  compare.1 = pred[10,10]{1,0} compare(iota.3, broadcast.6), direction=LT
  constant_3 = pred[] constant(false)
  broadcast.4 = pred[10,10]{1,0} broadcast(constant_3), dimensions={}
  compare.0 = pred[10,10]{1,0} compare(compare.1, broadcast.4), direction=NE
  constant_0 = f32[] constant(0)
  broadcast.3 = f32[10,10]{1,0} broadcast(constant_0), dimensions={}
  param_0.1 = f32[] parameter(0)
  broadcast.2 = f32[10,10]{1,0} broadcast(param_0.1), dimensions={}
  ROOT select.0 = f32[10,10]{1,0} select(compare.0, broadcast.3, broadcast.2)
}

ENTRY SyncTensorsGraph.25 {
  p0.13 = f32[] parameter(0)
  fusion = f32[10,10]{1,0} fusion(p0.13), kind=kLoop, calls=fused_computation
  ROOT tuple.24 = (f32[10,10]{1,0}) tuple(fusion)
}

-----
INFO:torch_xla:Letting libtpu.so load fail during _XLAC import. libtpu.so will be loaded from `libtpu` Python package when the ComputationClient is created.
2024-11-07 11:50:41.174644: I tensorflow/compiler/xla/service/service.cc:173] XLA service 0x905c190 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2024-11-07 11:50:41.174714: I tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (0): NVIDIA A100-SXM4-80GB, Compute Capability 8.0
2024-11-07 11:50:41.175641: I tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client.cc:194] Using BFC allocator.
2024-11-07 11:50:41.175713: I tensorflow/compiler/xla/pjrt/gpu/gpu_helpers.cc:105] XLA backend allocating 75175958937 bytes on device 0 for BFCAllocator.
2024-11-07 11:50:42.013482: I tensorflow/compiler/xla/service/dump.cc:485] HloModule dump enabled with path prefix: , suffix: before_optimizations
2024-11-07 11:50:42.037845: I tensorflow/tsl/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory
tensor([[ 0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38,
         -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38,
         -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38,
         -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38,
         -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
         -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
          0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
          0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
          0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
          0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00]],
       device='xla:0')

'''
相关推荐
z千鑫4 小时前
【人工智能】PyTorch、TensorFlow 和 Keras 全面解析与对比:深度学习框架的终极指南
人工智能·pytorch·深度学习·aigc·tensorflow·keras·codemoss
平头哥在等你4 小时前
求一个3*3矩阵对角线元素之和
c语言·算法·矩阵
学不会lostfound5 小时前
三、计算机视觉_05MTCNN人脸检测
pytorch·深度学习·计算机视觉·mtcnn·p-net·r-net·o-net
Mr.谢尔比6 小时前
李宏毅机器学习课程知识点摘要(1-5集)
人工智能·pytorch·深度学习·神经网络·算法·机器学习·计算机视觉
做程序员的第一天7 小时前
在PyTorch中,钩子(hook)是什么?在神经网络中扮演什么角色?
pytorch·python·深度学习
2402_871321958 小时前
MATLAB方程组
gpt·学习·线性代数·算法·matlab
Nerinic8 小时前
PyTorch基础2
pytorch·python
曼城周杰伦8 小时前
自然语言处理:第六十二章 KAG 超越GraphRAG的图谱框架
人工智能·pytorch·神经网络·自然语言处理·chatgpt·nlp·gpt-3
Joyner20189 小时前
pytorch训练的双卡,一个显卡占有20GB,另一个卡占有8GB,怎么均衡?
人工智能·pytorch·python
AI视觉网奇9 小时前
pytorch3d linux安装
linux·人工智能·pytorch