首页 > 其他分享 >Google的Jax框架的JAX-Triton目前只能成功运行在TPU设备上(使用Pallas为jax编写kernel扩展)—— GPU上目前无法正常运行,目前正处于 experimental 阶段

Google的Jax框架的JAX-Triton目前只能成功运行在TPU设备上(使用Pallas为jax编写kernel扩展)—— GPU上目前无法正常运行,目前正处于 experimental 阶段

时间:2024-01-18 21:24:29浏览次数:35  
标签:Triton kernel Pallas Jax JAX TPU GPU

使用Pallas为jax编写kernel扩展,需要使用JAX-Triton扩展包。由于Google的深度学习框架Jax主要是面向自己的TPU进行开发的,虽然也同时支持NVIDIA的GPU,但是支持力度有限,目前JAX-Triton只能在TPU设备上正常运行,无法保证在GPU上正常运行。


该结果使用kaggle上的TPU和GPU进行测试获得。


测试时间:

2024-01-18 21:12:09 星期四



image

image

标签:Triton,kernel,Pallas,Jax,JAX,TPU,GPU
From: https://www.cnblogs.com/devilmaycry812839668/p/17973422

相关文章

  • kaggle上的jax框架的环境配置(TPU版本)
    导出时间:2024-01-1821:00:37星期四python版本:Python3.10.13absl-py==1.4.0accelerate==0.25.0aiofiles==22.1.0aiosqlite==0.19.0anyio==4.2.0argon2-cffi==23.1.0argon2-cffi-bindings==21.2.0array-record==0.5.0arrow==1.3.0astroid==3.0.2asttokens==2.4......
  • Google的jax框架在TPU上的循环控制 —— 向量计算设备的循环结构控制
    相关:https://jax.readthedocs.io/en/latest/pallas/tpu.html向量计算设备,如:GPU、TPU等,都是通过向量计算来进行加速的,因此在这类设备中进行向量计算的计算单元是成百上千的,但是进行结构控制的电路单元比较少,可以基本认为在向量设备中进行流程控制是标量的,而不是向量的,也就是说......
  • jax框架的 Pallas 方式的GPU扩展不可用
    说下深度学习框架的GPU扩展功能的部分,也就是使用个人定制化的GPU代码编写方式来为深度学习框架做扩展。深度学习框架本身就是一种对GPU功能的一种封装和调用,但是由于太high-level,因此就会摒弃掉一些原有的GPU底层的编程功能,为此可以使用GPU原始功能的代码来为深度学习编写扩展函......
  • $.ajax()方法详解
    1.url要求为String类型的参数,(默认为当前页地址)发送请求的地址。2.type要求为String类型的参数,请求方式(post或get)默认为get。注意其他http请求方法,例如put和delete也可以使用,但仅部分浏览器支持。3.timeout要求为Number类型的参数,设置请求超时时间(毫秒)。此设置将覆盖$.ajaxSetup()方......
  • Ajax注册表单用户名实时验证
    Ajax注册表单用户名实时验证 1.1.1摘要很多时候在网站上注册时,我们会发现,注册表单通常需要检查用户名和电子邮件地址的可用性;从而确保用户之间不拥有相同的用户名和电子邮件地址;一些网站喜欢在用户提交填写的用户信息时,做信息可用性的检查,而一些网站会做实时的用户名和电子......
  • 通过 KernelUtil.dll 劫持 QQ / TIM 客户端 QQClientkey / QQKey 详细教程(附源码)
    前言由于QQ9.7.20版本后已经不能通过模拟网页快捷登录来截取QQClientkey/QQKey,估计是针对访问的程序做了限制,然而经过多方面测试,诸多的地区、环境、机器也针对这种获取方法做了相应的措施,导致模拟网页快捷登录来截取数据被彻底的和谐,为了解决这个问题我们只能更改思路对......
  • ajax实现文件上传
    ajax上传文件上传文件我一般采用form表单上传,form简单form上传文件<formid="jsSubmit"method="post"action="url"enctype="multipart/form-data"><inputtype="file"name="file"/></form><......
  • Ajax异步交互技术
    在HTTP协议下,前后端交互有两种技术,一是同步,二是异步Ajax异步交互技术简介“AsynchronousJavascriptAndXML”(异步JavaScript和XML),是指一种创建交互式、快速动态网页应用的网页开发技术,无需重新加载整个网页的情况下,能够更新部分网页的技术。通过在后台与服务器进行少量数据......
  • Ajax异步交互技术
    在HTTP协议下,前后端交互有两种技术,一是同步,二是异步Ajax异步交互技术简介“AsynchronousJavascriptAndXML”(异步JavaScript和XML),是指一种创建交互式、快速动态网页应用的网页开发技术,无需重新加载整个网页的情况下,能够更新部分网页的技术。通过在后台与服务器进行少量数据......
  • Ajax异步交互技术
    在HTTP协议下,前后端交互有两种技术,一是同步,二是异步Ajax异步交互技术简介“AsynchronousJavascriptAndXML”(异步JavaScript和XML),是指一种创建交互式、快速动态网页应用的网页开发技术,无需重新加载整个网页的情况下,能够更新部分网页的技术。通过在后台与服务器进行少量......