Skip to content

Commit

Permalink
doc(xla): modify xla doc
Browse files Browse the repository at this point in the history
  • Loading branch information
cqchu authored Jan 3, 2024
1 parent cff275a commit 247c81d
Showing 1 changed file with 17 additions and 0 deletions.
17 changes: 17 additions & 0 deletions source/user-guide/model-development/jit/xla.rst
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,23 @@ XLA executable 加速训练过程。
print (xla_fused_softmax(inp)) # run in xla
如果我们想看到 mge 和 xla 优化的一些中间 IR 表示,可以通过设置环境变量 MGE_VERBOSE_XLA_IR 来打印相关结果。MGE_VERBOSE_XLA_IR 为 1 时,会打印 mge trace 出来的图 IR,MGE_VERBOSE_XLA_IR 为 2 时,会打印xla 的 hlo 图结构,在 MGE_VERBOSE_XLA_IR 为 3 时会打印 xla 编译优化后的图结构。如果我们 export MGE_VERBOSE_XLA_IR=1 后再执行上述代码,则可以看到:

.. code-block:: python
please_realize_func_name_system_1(
0%:<256x1000x1000,f32>
) {
1%:<256x1000x1000,f32> = io_mark_var(0%:<256x1000x1000,f32>)
2%:<256x1000x1,f32> = ReduceMAX(1%:<256x1000x1000,f32>)
3%:<256x1000x1000,f32> = SUB(1%:<256x1000x1000,f32>, 2%:<256x1000x1,f32>)
4%:<256x1000x1000,f32> = EXP(3%:<256x1000x1000,f32>)
5%:<256x1000x1,f32> = ReduceSUM(4%:<256x1000x1000,f32>)
6%:<256x1000x1000,f32> = TRUE_DIV(4%:<256x1000x1000,f32>, 5%:<256x1000x1,f32>)
7%:<256x1000x1000,f32> = io_mark_var(6%:<256x1000x1000,f32>)
return 1 7%:<256x1000x1000,f32>
}
当模型训练迭代(Iteration)完全静态的情况下, 您也可以使用 jit.xla_trace 装饰器将训练迭代全部交由XLA执行。
需要将 optimizer, module 作为train_func 参数传入,同时 train_func 中需包含包含模型前向、 反向
、 参数更新等代码,
Expand Down

0 comments on commit 247c81d

Please sign in to comment.