PyTorch/XLA 文档¶

PyTorch/XLA 是一个 Python 软件包，它使用 XLA 深度学习编译器来连接 PyTorch 深度学习框架和 Cloud TPU。

文档

XLA 设备上的 PyTorch¶

PyTorch 在具有 torch_xla 软件包的 XLA 设备（如 TPU）上运行。本文档介绍如何在这些设备上运行模型。

创建 XLA 张量¶

PyTorch/XLA 向 PyTorch 添加了新的设备类型。此设备类型仅适用于与其他 PyTorch 设备类型一样。例如，下面介绍如何创建打印 XLA 张量：xla

import torch
import torch_xla
import torch_xla.core.xla_model as xm

t = torch.randn(2, 2, device=xm.xla_device())
print(t.device)
print(t)

此代码应该看起来很熟悉。PyTorch/XLA 使用与常规 PyTorch 中添加了一些内容。导入会初始化 PyTorch/XLA，并返回当前 XLA 设备。这可能是 CPU 或 TPU 具体取决于您的环境。torch_xlaxm.xla_device()

XLA 张量是 PyTorch 张量¶

PyTorch作可以像 CPU 或 CUDA 张量一样对 XLA 张量执行。

例如，可以将 XLA 张量相加：

t0 = torch.randn(2, 2, device=xm.xla_device())
t1 = torch.randn(2, 2, device=xm.xla_device())
print(t0 + t1)

或矩阵相乘：

print(t0.mm(t1))

或与神经网络模块一起使用：

l_in = torch.randn(10, device=xm.xla_device())
linear = torch.nn.Linear(10, 20).to(xm.xla_device())
l_out = linear(l_in)
print(l_out)

与其他设备类型一样，XLA 张量仅适用于同一设备。所以像

l_in = torch.randn(10, device=xm.xla_device())
linear = torch.nn.Linear(10, 20)
l_out = linear(l_in)
print(l_out)
# Input tensor is not an XLA tensor: torch.FloatTensor

将引发错误，因为模块在 CPU 上。torch.nn.Linear

在 XLA 设备上运行模型¶

构建新的 PyTorch 网络或将现有 PyTorch 网络转换为在 XLA 上运行 devices 只需要几行特定于 XLA 的代码。以下代码片段在具有 XLA 的单个设备和多个设备上运行时，突出显示这些行多处理。

在单个 XLA 设备上运行¶

以下代码片段显示了单个 XLA 设备上的网络训练：

import torch_xla.core.xla_model as xm

device = xm.xla_device()
model = MNIST().train().to(device)
loss_fn = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

for data, target in train_loader:
  optimizer.zero_grad()
  data = data.to(device)
  target = target.to(device)
  output = model(data)
  loss = loss_fn(output, target)
  loss.backward()

  optimizer.step()
  xm.mark_step()

此代码段重点介绍了将模型切换为在 XLA 上运行是多么容易。这模型定义、DataLoader、Optimizer 和 Training Loop 可以在任何设备上工作。唯一特定于 XLA 的代码是获取 XLA 设备和标记步骤。在每次培训结束时调用 iteration 使 XLA 执行其当前图形并更新模型的参数。有关更多信息，请参阅 XLA Tensor Deep Dive XLA 如何创建图形和运行作。xm.mark_step()

在具有多重处理功能的多个 XLA 设备上运行¶

PyTorch/XLA 通过在多个 XLA 上运行来轻松加速训练设备。以下代码段显示了如何作：

import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl

def _mp_fn(index):
  device = xm.xla_device()
  mp_device_loader = pl.MpDeviceLoader(train_loader, device)

  model = MNIST().train().to(device)
  loss_fn = nn.NLLLoss()
  optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

  for data, target in mp_device_loader:
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    xm.optimizer_step(optimizer)

if __name__ == '__main__':
  torch_xla.launch(_mp_fn, args=())

此多设备代码段与之前的单个设备代码段。让我们一个一个地回顾一下。

torch_xla.launch()
- 创建每个运行 XLA 设备的进程。
- 此函数是多线程 spawn 的包装器，允许用户也可以使用 torchrun 命令行运行脚本。每个进程将只能访问分配给当前进程的设备。例如，在 TPU v4-8 上，将生成 4 个进程，每个进程将拥有一个 TPU 设备。
- 请注意，如果您在每个进程上打印，您将在所有设备上看到。这是因为每个进程只能看到一个设备。这并不意味着多进程不起作用。唯一的执行是在 TPU v2 和 TPU v3 上使用 PJRT 运行时，因为会有进程，每个进程将有 2 个线程（有关更多详细信息，请查看此文档）。xm.xla_device()xla:0#devices/2
MpDeviceLoader
- 将训练数据加载到每个设备上。
- MpDeviceLoader可以包装在 Torch 数据加载器上。它可以将数据预加载到设备中，并将数据加载与设备执行重叠以提高性能。
- MpDeviceLoader还要调用 per （default to 1） batch being yield。xm.mark_stepbatches_per_execution
xm.optimizer_step(optimizer)
- 合并设备之间的梯度并发出 XLA 设备步长计算。
- 它几乎是一个 + + 并返回减少的损失。all_reduce_gradientsoptimizer.step()mark_step

模型定义、优化器定义和训练循环保持不变。

注意：需要注意的是，当使用多处理时，用户可以启动仅从的目标函数（或在调用中作为父函数的任何函数）中检索和访问 XLA 设备堆栈）。torch_xla.launch()torch_xla.launch()

有关使用多处理在多个 XLA 设备上训练网络的更多信息，请参阅完整的多处理示例。

在 TPU Pod 上运行¶

不同加速器的多主机设置可能大不相同。本文档将讨论多主机训练中独立于设备的位，并将使用 TPU + PJRT 运行时（目前在 1.13 和 2.x 版本中可用）作为示例。

在您开始之前，请查看我们的用户指南，其中将解释一些 Google Cloud 基础知识，例如如何使用命令以及如何设置您的项目。你也可以在这里查看所有 Cloud TPU Howto。本文档将重点介绍 Setup 的 PyTorch/XLA 透视图。gcloud

假设您在上一节中有上述 mnist 示例。如果是单机多设备训练，你可以通过 ssh 到 TPUVM 并运行类似train_mnist_xla.py

PJRT_DEVICE=TPU python3 train_mnist_xla.py

现在，为了在 TPU v4-16（有 2 个主机，每个主机有 4 个 TPU 设备）上运行相同的模型，您需要

确保每个主机都可以访问训练脚本和训练数据。这通常通过使用 command 或 command 将训练脚本复制到所有主机来完成。gcloud scpgcloud ssh
同时在所有主机上运行相同的训练命令。

gcloud alpha compute tpus tpu-vm ssh $USER-pjrt --zone=$ZONE --project=$PROJECT --worker=all --command="PJRT_DEVICE=TPU python3 train_mnist_xla.py"

上述命令将ssh到TPUVM Pod中的所有主机，并同时运行相同的命令..gcloud ssh

注意：您需要在 TPUVM 虚拟机外部运行上述命令。gcloud

多进程训练和多主机训练的模型代码和训练脚本相同。PyTorch/XLA 和底层基础设施将确保每个设备都知道全局拓扑以及每个设备的本地和全局序号。跨设备通信将跨所有设备进行，而不是在本地设备上进行。

有关 PJRT 运行时以及如何在 Pod 上运行它的更多详细信息，请参阅此文档。有关 PyTorch/XLA 和 TPU Pod 的更多信息，以及在 TPU Pod 上运行带有 fakedata 的 resnet50 的完整指南，请参阅本指南。

XLA Tensor 深入探究¶

使用 XLA 张量和设备只需要更改几行代码。但尽管 XLA 张量的行为与 CPU 和 CUDA 张量非常相似，但它们的内部结构是不同。本节介绍了 XLA 张量的唯一之处。

XLA 张量是惰性的¶

CPU 和 CUDA 张量立即或急切地启动作。XLA 张量 / 另一方面，他们是懒惰的。它们在图表中记录作，直到需要结果。像这样延迟执行可以让 XLA 对其进行优化。A 图表的多个单独作可能会融合到单个优化的作。

延迟执行通常对调用方不可见。自动使用 PyTorch/XLA 构建图形，将它们发送到 XLA 设备，并在在 XLA 设备和 CPU 之间复制数据。在以下情况下插入屏障执行优化器步骤会显式同步 CPU 和 XLA 设备。为有关我们的 Lazy Tensor 设计的更多信息，您可以阅读此论文。

内存布局¶

XLA 张量的内部数据表示形式对用户来说是不透明的。他们不公开它们的存储，并且它们看起来总是连续的，这与 CPU 和 CUDA 张量。这允许 XLA 调整张量的内存布局更好的性能。

将 XLA 张量移入和移出 CPU¶

XLA 张量可以从 CPU 移动到 XLA 设备，也可以从 XLA 设备移动添加到 CPU。如果移动视图，则其查看的数据也会复制到 other device 的 device 和视图关系。换句话说，一旦数据被复制到另一台设备，它就与其 previous device 或其上的任何张量。同样，根据你的代码的运行方式，欣赏和适应这种转变可能很重要。

保存和加载 XLA 张量¶

在保存之前，应将 XLA 张量移动到 CPU，如下所示片段：

import torch
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()

t0 = torch.randn(2, 2, device=device)
t1 = torch.randn(2, 2, device=device)

tensors = (t0.cpu(), t1.cpu())

torch.save(tensors, 'tensors.pt')

tensors = torch.load('tensors.pt')

t0 = tensors[0].to(device)
t1 = tensors[1].to(device)

这样，您就可以将加载的张量放在任何可用设备上，而不仅仅是初始化它们的设备上。

根据上述有关将 XLA 张量移动到 CPU 的说明，在以下情况下必须小心使用视图。建议您重新创建视图，而不是保存视图它们在 Tensor 加载并移动到其目标设备之后。

提供了一个实用程序 API，通过之前移动数据来保存数据到 CPU：

import torch
import torch_xla
import torch_xla.core.xla_model as xm

xm.save(model.state_dict(), path)

在多个设备的情况下，上述 API 只会保存 master 的数据 device 序号（0）。

如果内存与模型参数的大小相比有限，则提供了 API，以减少主机上的内存占用：

import torch_xla.utils.serialization as xser

xser.save(model.state_dict(), path)

此 API 一次将 XLA 张量流式传输到 CPU，从而减少主机数量内存，但它需要匹配的加载 API 才能恢复：

import torch_xla.utils.serialization as xser

state_dict = xser.load(path)
model.load_state_dict(state_dict)

可以直接保存 XLA 张量，但不建议这样做。XLA 张量总是被加载回保存它们的设备，如果该设备不可用，则加载将失败。PyTorch/XLA 与所有 PyTorch 一样，正在积极开发中，此行为将来可能会更改。

编译缓存¶

XLA 编译器将跟踪的 HLO 转换为在设备。编译可能很耗时，并且在 HLO 不会在执行之间更改，则编译结果可以持久化到 disk 进行重用，从而显著缩短开发迭代时间。

请注意，如果 HLO 在执行之间发生更改，则重新编译仍将发生。

目前，这是一个实验性的 opt-in API，必须先激活执行所有计算。初始化是通过 API 完成的：initialize_cache

import torch_xla.runtime as xr
xr.initialize_cache('YOUR_CACHE_PATH', readonly=False)

这将在指定路径处初始化持久编译缓存。该参数可用于控制 worker 是否能够写入缓存，这在共享缓存挂载用于 SPMD 工作负载。readonly

如果要在多进程训练中使用持久化编译缓存（带或），应该为不同的进程使用不同的路径。torch_xla.launchxmp.spawn

def _mp_fn(index):
  # cache init needs to happens inside the mp_fn.
  xr.initialize_cache(f'/tmp/xla_cache_{index}', readonly=False)
  ....

if __name__ == '__main__':
  torch_xla.launch(_mp_fn, args=())

如果您无权访问，则可以使用。在此处查看 runnable 示例。indexxr.global_ordinal()

延伸阅读¶

PyTorch/XLA 存储库中提供了其他文档。更多跑步示例 TPU 上的网络可在此处获得。

PyTorch/XLA API¶

torch_xla¶

torch_xla。device（index： Optional[int] = None） → device[源]¶

返回 XLA 设备的给定实例。

如果启用 SPMD，则返回一个虚拟设备，该虚拟设备包装所有可用设备到这个过程。

参数: index – 要返回的 XLA 设备的索引。对应于 torch_xla.devices（）中的 index。
返回: 一个 XLA torch.device。

torch_xla。devices（） → List[设备][源代码]¶

返回当前进程中可用的所有设备。

返回: XLA torch.devices 的列表。

torch_xla。device_count（） → int[来源]¶: 返回当前进程中可寻址设备的数量。

torch_xla。sync（等待：bool = False）[来源]¶

启动所有待处理的图形作。

参数: wait （bool） – 是否阻止当前进程直到执行完成。

torch_xla。compile（f：可选[可调用] = 无， full_graph：可选[bool] = False， 名称： Optional[str] = None， num_different_graphs_allowed： Optional[int] = 无）[来源]¶

使用 torch_xla 的 LazyTensor 跟踪模式优化给定的模型/函数。 PyTorch/XLA 将使用给定的输入跟踪给定的函数，然后生成表示 PyTorch作的图形发生在此函数中。这 graph 将由 XLA 编译并在加速器上执行（由 Tensor 的设备）。Eager 模式将对 funciton 的编译区域禁用。

参数

model （Callable） - 要优化的模块/函数，如果未传递，此函数将充当上下文管理器。
full_graph （Optional[bool]） – 此编译是否应生成单个图形。如果设置为 True 并且将生成多个图形torch_xla 将引发带有调试信息的错误并退出。
name （Optional[name]） – 已编译程序的名称。将使用函数 f 的名称如果未指定。此名称将用于 PT_XLA_DEBUG 消息以及 HLO/IR 转储文件。
num_different_graphs_allowed （Optional[python：int]） – 给定的不同跟踪图的数量我们被允许拥有的 model/function 的 model/function 来执行。如果此限制超出。

例：

# usage 1
@torch_xla.compile()
def foo(x):
  return torch.sin(x) + torch.cos(x)

def foo2(x):
  return torch.sin(x) + torch.cos(x)
# usage 2
compiled_foo2 = torch_xla.compile(foo2)

# usage 3
with torch_xla.compile():
  res = foo2(x)

torch_xla。manual_seed（seed， device=None）[来源]¶

设置用于为当前 XLA 设备生成随机数的种子。

参数

seed （python：integer） – 要设置的状态。
device （torch.device，可选） – 需要设置 RNG 状态的设备。如果缺少，将设置默认设备种子。

运行¶

torch_xla.runtime 中。device_type（） → 可选[str][源代码]¶

返回当前 PjRt 设备类型。

如果未配置任何设备，则选择默认设备

返回: 设备的字符串表示形式。

torch_xla.runtime 中。local_process_count（） → int[来源]¶: 返回此主机上运行的进程数。

torch_xla.runtime 中。local_device_count（） → int[来源]¶

返回此主机上的设备总数。

假设每个进程具有相同数量的可寻址设备。

torch_xla.runtime 中。addressable_device_count（） → int[来源]¶: 返回此进程可见的设备数。

torch_xla.runtime 中。global_device_count（） → int[来源]¶: 返回所有进程/主机中的设备总数。

torch_xla.runtime 中。global_runtime_device_count（） → int[来源]¶: 返回所有进程/主机中的运行时设备总数，尤其适用于 SPMD。

torch_xla.runtime 中。world_size（） → int[来源]¶: 返回参与作业的进程总数。

torch_xla.runtime 中。global_ordinal（） → int[来源]¶

返回所有进程中此线程的全局序号。

全局序数在 [0， global_device_count）范围内。全局序数不是保证与 TPU 工作程序 ID 有任何可预测的关系，也不是它们保证在每个主机上都是连续的。

torch_xla.runtime 中。local_ordinal（） → int[来源]¶

返回此主机中此线程的局部序号。

局部序数在 [0， local_device_count）范围内。

torch_xla.runtime 中。get_master_ip（） → str[来源]¶

检索运行时的主工作程序 IP。这将调用特定于后端的发现 API。

返回: 主工作程序的 IP 地址（字符串形式）。

torch_xla.runtime 中。use_spmd（auto：可选[bool] = False）[来源]¶

启用 SPMD 模式的 API。这是启用 SPMD 的推荐方法。

如果某些张量已在非 SPMD 上初始化，这将强制使用 SPMD 模式设备。这意味着这些张量将在设备之间复制。

参数: auto （bool） – 是否启用自动分片。阅读 https://github.com/pytorch/xla/blob/master/docs/spmd_advanced.md#auto-sharding 了解更多详情

torch_xla.runtime 中。is_spmd（）[来源]¶: 如果 SPMD 设置为执行，则返回。

torch_xla.runtime 中。initialize_cache（path： str， readonly： bool = False）[来源]¶

初始化持久编译缓存。必须调用此 API 在执行任何计算之前。

参数

path （str） – 存储持久缓存的路径。
readonly （bool） – 此工作程序是否应具有对缓存的写入权限。

xla_model¶

torch_xla.core.xla_model 中。xla_device（n：可选[int] = 无， devkind：可选[str] = 无）→设备[来源]¶

返回 XLA 设备的给定实例。

参数

n （python：int， optional） – 要返回的特定实例（ordinal）。如果指定，则将返回特定的 XLA 设备实例。否则将返回 devkind 的第一个设备。
devkind （字符串...，可选）— 如果指定，则为设备类型，例如 TPU、CUDA、CPU 或自定义 PJRT 设备。荒废的。

返回

具有请求实例的 torch.device。

torch_xla.core.xla_model 中。xla_device_hw（device： Union[str， device]） → str[source]¶

返回给定设备的硬件类型。

参数: device （string 或 torch.device） – 将映射到真实设备。
返回: 给定设备的硬件类型的字符串表示形式。

torch_xla.core.xla_model 中。is_master_ordinal（本地： bool = True） → bool[来源]¶

检查当前进程是否为主序号（0）。

参数: local （bool） – 应检查本地还是全局主序号。在多主机复制的情况下，只有一个全局主序号（主机 0，设备 0），而本地主序号NUM_HOSTS。默认值：True
返回: 一个布尔值，指示当前进程是否为主序号。

torch_xla.core.xla_model 中。all_reduce（reduce_type： str， inputs： Union[Tensor， List[Tensor]]， scale： float = 1.0， groups：可选[List[List[int]]] = 无，pin_layout：bool = true） → Union[Tensor， List[Tensor]][来源]¶

对输入张量执行就地 reduce作。

参数

reduce_type （string） – 、和之一。xm.REDUCE_SUMxm.REDUCE_MULxm.REDUCE_ANDxm.REDUCE_ORxm.REDUCE_MINxm.REDUCE_MAX
inputs （输入） – 单个 Torch。Tensor 或 torch 列表。Tensor 设置为执行所有减少作到。
scale （python：float） – 要在 reduce 之后应用的默认缩放值。默认值：1.0
组（列表，可选） –
一个列表列表，表示的副本组 all_reduce（）作。示例：[[0， 1， 2， 3]， [4， 5， 6， 7]]

定义两个组，一个具有 [0， 1， 2， 3] 副本，另一个具有 [4， 5， 6， 7] 副本。如果为 None，则只有一个组其中的所有副本。
pin_layout （bool， optional） – 是否固定此通信作的布局。布局固定可以防止在每个进程参与通信的程序略有不同，但可能会导致某些 XLA 编译失败。当您看到错误消息时取消固定布局例如“HloModule has a mix of layout constrained”。

返回

如果单个Torch。Tensor 传入，返回值为 torch。Tensor 保存缩减值（跨副本）。如果传递了列表/元组，则此函数对输入张量执行就地 all-reduce 运算，并且返回列表/元组本身。

torch_xla.core.xla_model 中。all_gather（value： Tensor， dim： int = 0， groups： Optional[List[List[int]]] = 无，输出：可选[Tensor] = 无，pin_layout：bool = True）→Tensors量[来源]¶

沿给定维度执行 all-gather作。

参数

值（Torch。Tensor） - 输入张量。
dim （python：int） – 收集维度。默认值：0
组（列表，可选） –
一个列表列表，表示的副本组 all_gather（）作。示例：[[0， 1， 2， 3]， [4， 5， 6， 7]]

定义两个组，一个具有 [0， 1， 2， 3] 副本，另一个具有 [4， 5， 6， 7] 副本。如果为 None，则只有一个组其中的所有副本。
输出（Torch.Tensor） - 可选的输出张量。
pin_layout （bool， optional） – 是否固定此通信作的布局。布局固定可以防止在每个进程参与通信的程序略有不同，但可能会导致某些 XLA 编译失败。当您看到错误消息时取消固定布局例如“HloModule has a mix of layout constrained”。

返回

一个张量，在维度中，具有参与副本。dim

torch_xla.core.xla_model 中。all_to_all（值：张量，split_dimension：int，concat_dimension：int，split_count：int， groups：可选[List[List[int]]] = None， pin_layout： bool = True） → 张量[来源]¶

对输入张量执行 XLA AllToAll（）作。

另请： https://www.tensorflow.org/xla/operation_semantics#alltoall

参数

值（Torch。Tensor） - 输入张量。
split_dimension （python：int） – 应进行拆分的维度。
concat_dimension （python：int） – 连接应发生的维度。
split_count （python：int） – 拆分计数。
组（列表，可选） –
一个列表列表，表示的副本组 all_reduce（）作。示例：[[0， 1， 2， 3]， [4， 5， 6， 7]]

定义两个组，一个具有 [0， 1， 2， 3] 副本，另一个具有 [4， 5， 6， 7] 副本。如果为 None，则只有一个组其中的所有副本。
pin_layout （bool， optional） – 是否固定此通信作的布局。布局固定可以防止在每个进程参与通信的程序略有不同，但可能会导致某些 XLA 编译失败。当您看到错误消息时取消固定布局例如“HloModule has a mix of layout constrained”。

返回

结果Torch。all_to_all（）作的张量。

torch_xla.core.xla_model 中。add_step_closure（closure： Callable[[...]， Any]， args：元组[Any] = （）， run_async： bool = 错误）[来源]¶

将 closure 添加到要在步骤结束时运行的 Merge 列表中。

在模型训练期间，很多时候需要打印/报告（打印到控制台、发布到 TensorBoard 等...需要要检查的中间张量。检查模型代码不同点的不同张量内容需要多次执行，并且通常会导致性能问题。添加步骤闭包将确保它将在 barrier 之后运行，当所有 live 张量都将物化为 Device 数据。 Live Tensors，其中包括 closure 参数捕获的 Tensor。因此，使用 add_step_closure（）将确保单次执行将是执行，即使多个 closure 排队时，也需要多个张量接受检查。步骤闭包将按照其排队顺序按顺序运行。请注意，即使使用此 API 执行将得到优化，它建议每 N 步限制一次打印/报告事件。

参数

closure （callable） – 要调用的函数。
args （tuple） – 要传递给闭包的参数。
run_async – 如果为 True，则异步运行闭包。

torch_xla.core.xla_model 中。wait_device_ops（设备： List[str] = []）[源]¶

等待给定设备上的所有异步作完成。

参数: devices （字符串...，可选）— 需要等待其异步作的设备为。如果为空，则将等待所有本地设备。

torch_xla.core.xla_model 中。optimizer_step（optimizer：优化器， barrier： bool = False， optimizer_args： Dict = {}， groups：可选[List[List[int]]] = None， pin_layout： bool = True）[来源]¶

运行提供的优化器步骤并在所有设备上同步 gradidents。

参数

optimizer （） – Torch。需要调用其 step（）函数的 Optimizer 实例。step（）函数将被调用替换为 optimizer_args 命名参数。torch.Optimizer
barrier （bool，可选） - 是否应在这个 API 的 API 中。如果使用 PyTorch XLA ParallelLoader 或 DataParallel 支持，则不需要这样做，因为屏障将由 XLA 颁发数据加载器迭代器 next（）调用。默认值：False
optimizer_args （dict， optional） – optimizer.step（）调用的命名参数字典。
组（列表，可选） –
一个列表列表，表示的副本组 all_reduce（）作。示例：[[0， 1， 2， 3]， [4， 5， 6， 7]]

定义两个组，一个具有 [0， 1， 2， 3] 副本，另一个具有 [4， 5， 6， 7] 副本。如果为 None，则只有一个组其中的所有副本。
pin_layout （bool， optional） – 是否在减少渐变时固定布局。有关详细信息，请参阅 xm.all_reduce。

返回

与 optimizer.step（）调用返回的值相同。

例

>>> import torch_xla.core.xla_model as xm
>>> xm.optimizer_step(self.optimizer)

torch_xla.core.xla_model 中。save（data： Any， file_or_path： Union[str， TextIO]， master_only： bool = 真，global_master：bool = False）[来源]¶

将输入数据保存到文件中。

保存的数据在保存之前会传输到 PyTorch CPU 设备，因此遵循 torch.load（）将加载 CPU 数据。使用视图时必须小心。它不是保存视图建议您在加载张量后重新创建它们，并且已移动到其目标设备。

参数

data – 要保存的输入数据。Python 对象的任何嵌套组合（列表、元组、集合、字典等）。
file_or_path – 数据保存作的目标。要么是文件 path 或 Python 文件对象。如果 master_only 是路径或 file 对象必须指向不同的目标，否则所有来自同一主机的写入将相互覆盖。False
master_only （bool， optional） – 是否只有主设备应保存数据。如果为 False，则 file_or_path 参数应为不同的文件或 path （路径），否则同一主机上的所有副本都将写入同一位置。默认值：True
global_master （bool， optional） – 何时是此标志控制每个主机的 master （if is ）保存内容，或仅保存全局主控（序号 0）。默认值：Falsemaster_onlyTrueglobal_masterFalse

例

>>> import torch_xla.core.xla_model as xm
>>> xm.wait_device_ops() # wait for all pending operations to finish.
>>> xm.save(obj_to_save, path_to_save)
>>> xm.rendezvous('torch_xla.core.xla_model.save') # multi process context only

torch_xla.core.xla_model 中。rendezvous（标签： str， payload： bytes = b''， 副本： List[int] = []） → List[字节][来源]¶

等待所有网格客户端到达命名的会合点。

注意：PJRT 不支持 XRT 网格服务器，因此这实际上是一个 alias 设置为 xla_rendezvous。

参数

tag （string） – 要加入的集合的名称。
payload （bytes， optional）（payload （bytes， optional）（要发送到集合的负载） – 要发送到集合的负载。
replicas （list， python：int） – 参与会合的副本序号。 Empty 表示网格中的所有副本。默认值：[]

返回

所有其他 core 与 core 的 payload 交换的 payload 序数 i 在返回的元组中的位置 i。

例

>>> import torch_xla.core.xla_model as xm
>>> xm.rendezvous('example')

torch_xla.core.xla_model 中。mesh_reduce（标签： str， data， reduce_fn： Callable[[...]， Any]） → Union[Any， ToXlaTensorArena][来源]¶

执行图外客户端网格缩减。

参数

tag （string） – 要加入的集合的名称。
data （数据） – 要减少的数据。reduce_fn callable 将收到一个列表使用来自所有 Mesh 客户端进程的相同数据的副本（每个内核一个）。
reduce_fn （callable） – 一个函数，它接收类似 data 的列表对象并返回 Reduced 结果。

返回

减少的值。

例

>>> import torch_xla.core.xla_model as xm
>>> import numpy as np
>>> accuracy = xm.mesh_reduce('test_accuracy', accuracy, np.mean)

torch_xla.core.xla_model 中。set_rng_state（seed： int， device： Optional[str] = None）[来源]¶

设置随机数生成器状态。

参数

seed （python：integer） – 要设置的状态。
device （string， optional）（device （string， optional）（设备） - 需要设置 RNG 状态的设备。如果缺少，将设置默认设备种子。

torch_xla.core.xla_model 中。get_rng_state（device： Optional[str] = None） → int[源代码]¶

获取当前正在运行的随机数生成器状态。

参数: device （string， optional）（device （string， optional）（device （字符串，可选）） – 需要检索其 RNG 状态的设备。如果缺少，将设置默认设备种子。
返回: RNG 状态，以整数表示。

torch_xla.core.xla_model 中。get_memory_info（device： Optional[device] = None） → MemoryInfo[source]¶

检索设备内存使用情况。

参数

device – 可选[torch.device] 请求其内存信息的设备。
装置。（如果未通过，将使用默认值） –

返回

MemoryInfo dict 替换为给定设备的内存使用情况。

例

>>> xm.get_memory_info()
{'bytes_used': 290816, 'bytes_limit': 34088157184}

torch_xla.core.xla_model 中。get_stablehlo（张量：可选[List[Tensor]] = None） → str[源]¶

以字符串格式获取计算图的 StableHLO。

如果 tensors 不为空，则将转储以 tensors 作为输出的图形。如果 tensors 为空，则整个计算图将被 dump。

对于推理图，建议将模型输出传递给张量。对于训练图，识别 “outputs” 并不简单。建议使用空张量。

要在 StableHLO 中启用源行信息，请设置 env var XLA_HLO_DEBUG=1。

参数: 张量（list[torch.Tensor]，可选）– 表示 StableHLO 图的输出/根的张量。
返回: 字符串格式的 StableHLO 模块。

torch_xla.core.xla_model 中。get_stablehlo_bytecode（张量：可选[张量] = 无） →字节[来源]¶

以字节码格式获取计算图的 StableHLO。

如果 tensors 不为空，则将转储以 tensors 作为输出的图形。如果 tensors 为空，则整个计算图将被 dump。

对于推理图，建议将模型输出传递给张量。对于训练图，识别 “outputs” 并不简单。建议使用空张量。

参数: 张量（list[torch.Tensor]，可选）– 表示 StableHLO 图的输出/根的张量。
返回: 字节码格式的 StableHLO 模块。

分散式¶

类 torch_xla.distributed.parallel_loader 中。MpDeviceLoader（加载器，设备，**kwargs）[来源]¶

使用后台数据上传包装现有的 PyTorch DataLoader。

此类应仅与多处理数据并行性一起使用。它将包装使用 ParallelLoader 传入的数据加载器，并per_device_loader返回当前设备。

参数

loader （） – PyTorch DataLoader 为包裹。torch.utils.data.DataLoader
device （torch.device...） – 必须发送数据的设备。
kwargs – ParallelLoader 构造函数的命名参数。

例

>>> device = torch_xla.device()
>>> train_device_loader = MpDeviceLoader(train_loader, device)

torch_xla.distributed.xla_multiprocessing 中。spawn（fn， args=（）， nprocs=None， join=True， daemon=False， start_method='spawn'）[来源]¶

支持基于多处理的复制。

参数

fn （callable） – 要为每个设备调用的函数，该函数是复制。该函数将被调用，第一个参数为复制中进程的全局索引，后跟在 args 中传递的参数。
args （tuple） - fn 的参数。默认值：空元组
nprocs （python：int） – 用于复制的进程/设备的数量。在 moment（如果指定）可以是 1 或最大设备数。
join （bool） – 调用是否应阻止等待进程。默认值：True
daemon （bool） – 正在生成的进程是否应设置守护进程标志（请参阅 Python 多处理 API）。默认值：False
start_method （string） – Python 多处理进程创建方法。默认值：spawn

返回

由 torch.multiprocessing.spawn API 返回的相同对象。如果 nprocs 为 1，则直接调用 fn 函数，API 将返回 None。

SPMD¶

torch_xla.distributed.spmd 中。mark_sharding（t： Union[Tensor， XLAShardedTensor]， mesh： Mesh， partition_spec： → XLAShardedTensor 的 Tuple[Optional[Union[Tuple， int， str]]]） ¶

注释 XLA 分区规范提供的张量。内部它将相应的 XLATensor 注释为 XLA SpmdPartitioner 传递的分片。

参数

t （联合[torch.Tensor， XLAShardedTensor]） – 要用 partition_spec 注释的输入张量。
mesh （Mesh） – 描述逻辑 XLA 设备拓扑和底层设备 ID。
partition_spec （Tuple[Tuple[Tuple， python：int， str， None]） – 维度索引或 None device_mesh元组。每个索引是一个 int、str（如果网格轴已命名）或 int 或 str 的元组。这指定了如何对每个输入排名进行分片（索引到 mesh_shape）或复制（None）。指定元组时，相应的输入张量轴将沿 all 分片元组中的逻辑轴。请注意，网格轴在元组中的指定顺序将影响生成的分片。
dynamo_custom_op （bool） – 如果设置为 True，则调用 mark_sharding 的 dynamo 自定义运算变体使自身可被 Dynamo 识别和跟踪。

例

>>> import torch_xla.runtime as xr
>>> import torch_xla.distributed.spmd as xs
>>> mesh_shape = (4, 2)
>>> num_devices = xr.global_runtime_device_count()
>>> device_ids = np.array(range(num_devices))
>>> mesh = Mesh(device_ids, mesh_shape, ('x', 'y'))
>>> input = torch.randn(8, 32).to(xm.xla_device())
>>> xs.mark_sharding(input, mesh, (0, None)) # 4-way data parallel
>>> linear = nn.Linear(32, 10).to(xm.xla_device())
>>> xs.mark_sharding(linear.weight, mesh, (None, 1)) # 2-way model parallel

torch_xla.distributed.spmd 中。clear_sharding（t： Union[Tensor， XLAShardedTensor]） → Tensor[来源]¶

从输入张量中清除分片注释并返回 cpu 强制转换的张量。这是就地作，但也会返回相同的手电筒。Tensor 返回。

参数: t （联合[torch.Tensor， XLAShardedTensor]） – 我们要清除分片的张量
返回: tensor 的 Tensor 进行分片。
返回类型: t （Torch。张量）

例

>>> import torch_xla.distributed.spmd as xs
>>> torch_xla.runtime.use_spmd()
>>> t1 = torch.randn(8,8).to(torch_xla.device())
>>> mesh = xs.get_1d_mesh()
>>> xs.mark_sharding(t1, mesh, (0, None))
>>> xs.clear_sharding(t1)

torch_xla.distributed.spmd 中。set_global_mesh（mesh： Mesh）[来源]¶

设置可用于当前流程的全局网格。

参数: mesh – （Mesh）将成为全局网格的 Mesh 对象。

例

>>> import torch_xla.distributed.spmd as xs
>>> mesh = xs.get_1d_mesh("data")
>>> xs.set_global_mesh(mesh)

torch_xla.distributed.spmd 中。get_global_mesh（） → 可选[网格][来源]¶

获取当前流程的全局网格。

返回: （可选 [网格]）Mesh 对象（如果设置了全局网格），则返回 None。
返回类型: 网孔

例

>>> import torch_xla.distributed.spmd as xs
>>> xs.get_global_mesh()

torch_xla.distributed.spmd 中。get_1d_mesh（axis_name：可选[str] = 无） →网格[源]¶

Helper 函数返回一个维度中所有设备的网格。

参数: axis_name – （Optional[str]）表示网格轴名称的可选字符串
返回: Mesh 对象
返回类型: 网孔

例

>>> # This example is assuming 1 TPU v4-8
>>> import torch_xla.distributed.spmd as xs
>>> mesh = xs.get_1d_mesh("data")
>>> print(mesh.mesh_shape)
(4,)
>>> print(mesh.axis_names)
('data',)

类 torch_xla.distributed.spmd。mesh（device_ids： Union[ndarray， list]， mesh_shape： Tuple[int， ...]， axis_names：可选[Tuple[str， ...]] = 无）[来源]¶

描述逻辑 XLA 设备拓扑网格和底层资源。

参数

device_ids （Union[np.ndarray， List]） – 按自定义顺序排列的设备（ID）列表。列表已重塑添加到mesh_shape数组中，使用类似 C 的索引顺序填充元素。
mesh_shape （Tuple[python：int， ...]） – 描述逻辑拓扑形状的 int 元组的设备网格中，每个元素描述相应的轴。
axis_names （Tuple[str， ...]） – 要分配给维度的资源轴名称序列的设备参数。其长度应与设备的等级相匹配。

例

>>> mesh_shape = (4, 2)
>>> num_devices = len(xm.get_xla_supported_devices())
>>> device_ids = np.array(range(num_devices))
>>> mesh = Mesh(device_ids, mesh_shape, ('x', 'y'))
>>> mesh.get_logical_mesh()
>>> array([[0, 1],
          [2, 3],
          [4, 5],
          [6, 7]])
>>> mesh.shape()
OrderedDict([('x', 4), ('y', 2)])

类 torch_xla.distributed.spmd。HybridMesh（*， ici_mesh_shape：元组[int， ...]，dcn_mesh_shape：可选[Tuple[int， ...]] = 无，axis_names：可选[Tuple[str， ...]] = 无）[来源]¶

创建与 ICI 和 DCN 网络连接的设备的混合设备网格。: 逻辑网格的形状应按增加网络强度进行排序例如，[replica， data， model]，其中 MDL 具有最多的网络通信要求。

参数

ici_mesh_shape – 内部连接设备的逻辑网格形状。
dcn_mesh_shape – 外部连接设备的逻辑网格形状。

例

>>> # This example is assuming 2 slices of v4-8.
>>> ici_mesh_shape = (1, 4, 1) # (data, fsdp, tensor)
>>> dcn_mesh_shape = (2, 1, 1)
>>> mesh = HybridMesh(ici_mesh_shape, dcn_mesh_shape, ('data','fsdp','tensor'))
>>> print(mesh.shape())
>>> >> OrderedDict([('data', 2), ('fsdp', 4), ('tensor', 1)])

实验的¶

torch_xla.实验性。eager_mode（启用：bool）[来源]¶

配置 torch_xla 的默认执行模式。

在 Eager 模式下，只有 'torch_xla.compile' 的函数才会是跟踪和编译。其他Torch作将立即执行。

调试¶

torch_xla.debug.metrics 中。metrics_report（）[来源]¶: 检索包含完整指标和计数器报表的字符串。

torch_xla.debug.metrics 中。short_metrics_report（counter_names：可选[列表] = 无， metric_names：可选[列表] = 无）[来源]¶

检索包含完整指标和计数器报表的字符串。

参数

counter_names （list） – 需要打印其数据的计数器名称的列表。
metric_names （list） – 需要打印其数据的指标名称列表。

torch_xla.debug.metrics 中。counter_names（）[来源]¶: 检索所有当前活动的计数器名称。

torch_xla.debug.metrics 中。counter_value（name）[来源]¶

返回活动计数器的值。

参数: name （string） – 需要检索其值的计数器的名称。
返回: 计数器值为整数。

torch_xla.debug.metrics 中。metric_names（）[来源]¶: 检索所有当前活动的指标名称。

torch_xla.debug.metrics 中。metric_data（name）[来源]¶

返回活动量度的数据。

参数: name （string） – 需要检索其数据的指标的名称。
返回: 度量数据，它是（TOTAL_SAMPLES， ACCUMULATOR， SAMPLES）的元组。 TOTAL_SAMPLES 是已发布到指标。一个量度仅保留给定数量的样本（在循环缓冲区）。 Acculator 是 TOTAL_SAMPLES 的样本之和。 SAMPLES 是（TIME， VALUE）元组的列表。

PyTorch/XLA 文档¶

XLA 设备上的 PyTorch¶

创建 XLA 张量¶

XLA 张量是 PyTorch 张量¶

在 XLA 设备上运行模型¶

在单个 XLA 设备上运行¶

在具有多重处理功能的多个 XLA 设备上运行¶

在 TPU Pod 上运行¶

XLA Tensor 深入探究¶

XLA 张量是惰性的¶

内存布局¶

将 XLA 张量移入和移出 CPU¶

保存和加载 XLA 张量¶

编译缓存¶

延伸阅读¶

PyTorch/XLA API¶

torch_xla¶

运行¶

xla_model¶

分散式¶

SPMD¶

实验的¶

调试¶

文档

教程

资源