问题标签 [tensorflow-xla]

我按照本指南在 Google Colab TPU 上启动了我的 PyTorch Lightning 项目。所以我安装了

然后

然后我

安装项目要求后，我按要求重新启动运行时，并从上面重新运行 cloud-TPU-client 安装、pytorch-lightning 安装和这两个命令。它运行得很顺利。

但就在 TPU 开始使用 PyTorch 1.9 版之后，我收到以下导入错误：

Trainer与国旗一起发射TPU_cores=8。

该模型事先已在 CPU 和 GPU 上运行（即在另一个会话上）。

我试图将 PyTorch 降级到 1.9（与 TPU 启动时显示的相同），因为 Colab 使用了 torch 1.10.0+cu111 并且出现了不同的错误：

我可以做些什么来在 TPU 上训练模型吗？

非常感谢

我正在使用 XLA C++ API，并且我设法运行了一个简单的添加，但我不知道我是否做得对。似乎有很多我没有使用过的类。这是我的例子

XLA 课程XlaBuilder有什么用？文档描述了它的界面，但没有提供动机。

XlaBuilder文档中的演示文稿，以及源代码中的上述注释

表明它只是一个实用程序。然而，这似乎并不能解释它在其他地方的行为。例如，我们可以XlaOp用XlaBuildervia构造一个

在这里，我不清楚扮演什么角色（在已发布的文档中没有记录builder构造 s 的注意函数）。XlaOp此外，当我添加两个XlaOps （使用+or Add）时，似乎必须使用相同的构建器构建操作，否则我会看到

实际上，XlaOp保留了一个XlaBuilder. 这向我表明，它XlaBuilder具有更根本的意义。

除了标题问题之外，是否有使用多个XlaBuilders 的用例，或者您通常会为所有事情使用一个全局实例？

我目前正在修改 TensorFlow 的 XLA 代码，为 GPU 编译器添加优化通道。我添加了一个std::map作为成员变量的类，用于在其他优化过程中存储和使用信息。

添加的代码如下：

该文件位于tensorflow/compiler/xla/service/gpu

错误信息如下：

在添加其他优化通道（例如图形重写通道）时，我没有遇到这种错误。HloInstruciton但是，我在修改类时遇到了类似的错误。

我检查了 Bazel 构建文件，我相信我已经添加了每个 cc_library 进行编译。

我正在使用 TF 2.4.1，搭配 AMD EPYC 7452 32 核处理器，GPU 以 RTX 2080 Ti * 4 为目标。

有没有办法为 TensorFlow 编译额外的实现？

我开始使用 Google JAX 和内置的 jit 和 grad 功能。这些方面在我的机器上运行良好，但是当我增加参数数量时，我收到以下通知：

我很想增加输入参数的数量，所以我想很快我将需要更快的编译时间，所以这个通知很吸引我......但我不明白如何实现它。

我一直在使用 conda 来安装 jax。基本上，我在终端中运行以下命令：

我确定在 conda 中安装时必须有一种方法可以添加一些选项（例如，使用conda install jax=arguments但我在任何地方的文档中都找不到如何操作。堆栈溢出似乎也没有任何内容-搜索只发现以下内容： 使用 jax 时 XLA 的 jit 编译速度非常慢

任何建议将不胜感激！

我正在向 XLA 编写一个纯 C FFI 层，并希望返回一个指向 的指针GlobalData，由xla::LocalClient.TransferToServer(...). 我试图GlobalData在免费商店创建一个新的，因为源代码中的评论说

我还释放了堆栈副本的句柄，如

但它不起作用。它似乎仍在释放句柄，所以当client.ExecuteAndTransfer(...)我看到

我也试过client.Unregister(*global_data)没有帮助。

有许多函数可以XlaOp从原生 C++ 值创建 s。我试图弄清楚如何使用每个来构建图表。我浏览了 xla_builder.h 并挑选了一些候选者，省略了重载和便利包装器。两个最有可能的候选人似乎是

我是否正确地认为Parameter是“符号”，而ConstantLiteral对于常数值？例如，在 中f(x) = x + 1，我们将编码1为 a ConstantLiteral，然后x我们可以

1. 编写f(x)为 C++ 函数，并在应用程序站点使用另一个ConstantLiteral作为我们的值x，或
2. 使用相应的. x_ 也就是说，我不清楚如何实际调用with a ，除了 with which 不适用于多个s afaict。ParameterXlaComputationXlaBuilderXlaComputationLiteralLocalClientXlaComputation

这两种方法有什么区别？这个比那个好吗？我注意到前者对于高阶函数似乎是不可能的：那些接受XlaComputations 的函数。

接下来有

• Infeed，我猜这是Parameter.
• Recv这看起来像是在计算之间传递数据的一种方式，但实际上并没有创建一个全新的XlaOp本身。
• ReplicaId, Iota, 并且XlaOp CreateToken(XlaBuilder* builder);似乎与本次讨论无关。

我做对了吗？我还错过了其他重要的功能吗？

矩阵链是矩阵矩阵乘积的链。我考虑以下矩阵链：

美国广播公司；其中A和B的大小为 3000x3000，C的大小为 3000x600

有两种方法可以评估上述表达式，它们在性能上存在显着差异：

变体 1：(AB)C：6.48e10 FLOPs

变体 2: A(BC) : 2.16e10 FLOPs

矩阵矩阵乘法AB的成本是2xmxnxk，其中A的大小为mxn，B的大小为nxk。使用这个公式，我获得了上述变体的 FLOP 性能。

如果括号没有明确指定，我的 TensorFlow 构建（版本 2.8，禁用 Eager 模式）仅选择变体 1（从左到右括号），其执行时间几乎是变体 2 的三倍。虽然我可以优化这个并且通过显式计算矩阵乘法的 FLOP 手动加括号，我很好奇这是否可以由 TensorFlow 使用的 Grappler 图优化器自动完成？是否有任何其他图形优化器可以自动选择最佳括号？

观察不同括号的性能效果的示例脚本

使用 Intel MKL 和 Python 3.9.7 构建的 TensorFlow 2.8 (CPU) 的执行时间在 Mac book pro 2018 Big sur 上运行

• 变体 1（默认括号）：0.65s
• 变体 2（优化括号）：0.2s

XlaComputation构建另一个时如何使用XlaComputation？例如，我想构建Add一个计算，我正在尝试这样做

但我明白了

试图获取值而不是处理错误无效的参数：带有句柄 3 的 XlaOp 由构建器“Add”构建，但试图在构建器“Mul”中使用它：

我知道我不应该在另一个计算中使用来自一个构建器的，但是当它对应于该计算的结果XlaOp时，我不明白如何避免这种情况。XlaOp阅读源代码中的以下注释后

我猜这仅适用于XlaOps 使用相同的构建器。

我正在比较 XLAxla::Map和元素操作，我想弄清楚何时使用它们。

API

快速浏览一下，似乎所有对标量的计算也可以定义为任意形状，因为所有相关的操作都有逐元素的变体（Eq、Sin等Add）。我不清楚Map除了替代 API 还提供什么。或者只是为了让我们可以使用预先构建的标量计算，在这种情况下它与 有何不同Call？

表现

是否xla::Map具有与元素操作相同的性能特征？也就是说，在伪代码中，是

优化至少大致与