Главная страница>Связанные с программированием>Другой исходный код
Скачать

Pytorch/xla

Текущий статус CI:

Pytorch/XLA - это пакет Python, в котором используется компилятор Deep Learning для глубокого обучения для подключения структуры глубокого обучения Pytorch и облачных TPU. Вы можете попробовать это прямо сейчас, бесплатно, на одном облачном виртуальной машине TPU с Kaggle!

Взгляните на одну из наших записей Kaggle, чтобы начать:

  • Стабильная диффузия с pytorch/xla 2.0
  • Распределенные основы Pytorch/XLA

Установка

TPU

Чтобы установить стабильную сборку Pytorch/XLA в новую виртуальную машину TPU: 

 pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html

Чтобы установить ночную сборку Pytorch/XLA в новой VM TPU: 

 pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev-cp310-cp310-linux_x86_64.whl' -f https://storage.googleapis.com/libtpu-releases/index.html

Плагин графического процессора

Pytorch/XLA теперь обеспечивает поддержку GPU через пакет плагинов, аналогичный libtpu : 

 pip install torch~=2.5.0 torch_xla~=2.5.0 https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla_cuda_plugin-2.5.0-py3-none-any.whl

Начиная

Чтобы обновить существующий петлю обучения, внесите следующие изменения: 

 - import torch.multiprocessing as mp
+ import torch_xla as xla
+ import torch_xla.core.xla_model as xm

 def _mp_fn(index):
   ...

+  # Move the model paramters to your XLA device
+  model.to(xla.device())

   for inputs, labels in train_loader:
+    with xla.step():
+      # Transfer data to the XLA device. This happens asynchronously.
+      inputs, labels = inputs.to(xla.device()), labels.to(xla.device())
       optimizer.zero_grad()
       outputs = model(inputs)
       loss = loss_fn(outputs, labels)
       loss.backward()
-      optimizer.step()
+      # `xm.optimizer_step` combines gradients across replicas
+      xm.optimizer_step(optimizer)

 if __name__ == '__main__':
-  mp.spawn(_mp_fn, args=(), nprocs=world_size)
+  # xla.launch automatically selects the correct world size
+  xla.launch(_mp_fn, args=())

Если вы используете DistributedDataParallel , внесите следующие изменения: 

 import torch.distributed as dist
- import torch.multiprocessing as mp
+ import torch_xla as xla
+ import torch_xla.distributed.xla_backend

 def _mp_fn(rank):
   ...

-  os.environ['MASTER_ADDR'] = 'localhost'
-  os.environ['MASTER_PORT'] = '12355'
-  dist.init_process_group("gloo", rank=rank, world_size=world_size)
+  # Rank and world size are inferred from the XLA device runtime
+  dist.init_process_group("xla", init_method='xla://')
+
+  model.to(xm.xla_device())
+  # `gradient_as_bucket_view=True` required for XLA
+  ddp_model = DDP(model, gradient_as_bucket_view=True)

-  model = model.to(rank)
-  ddp_model = DDP(model, device_ids=[rank])

   for inputs, labels in train_loader:
+    with xla.step():
+      inputs, labels = inputs.to(xla.device()), labels.to(xla.device())
       optimizer.zero_grad()
       outputs = ddp_model(inputs)
       loss = loss_fn(outputs, labels)
       loss.backward()
       optimizer.step()

 if __name__ == '__main__':
-  mp.spawn(_mp_fn, args=(), nprocs=world_size)
+  xla.launch(_mp_fn, args=())

Дополнительная информация о Pytorch/XLA, включая описание его семантики и функций, доступна на pytorch.org. См. Руководство по API для лучших практик при написании сети, которые работают на устройствах XLA (TPU, CUDA, CPU и ...).

Наши всеобъемлющие пользовательские руководства доступны по адресу:

Документация для последнего выпуска

Документация для главного филиала

Учебники Pytorch/XLA

  • Cloud TPU VM QuickStart
  • Cloud TPU Pod Slice QuickStart
  • Профилирование на VM TPU
  • GPU Guide

Доступные изображения Docker и колеса

Python Packages

Выпуски Pytorch/XLA, начиная с версии R2.1, будут доступны на PYPI. Теперь вы можете установить основную сборку с помощью pip install torch_xla . Чтобы также установить плагин Cloud TPU, соответствующий вашему установленному torch_xla , установите дополнительные зависимости tpu после установки основной сборки с 

 pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

GPU и ночные сборки доступны в нашем публичном ведре GCS.

Версия Облачные графические колеса VM Wheels
2.5 (CUDA 12.1 + Python 3.9) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0-cp39-cp39-manylinux_2_28_x86_64.whl
2.5 (CUDA 12.1 + Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.5 (CUDA 12.1 + Python 3.11) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0-cp311-cp311-manylinux_2_28_x86_64.whl
2.5 (CUDA 12,4 + Python 3.9) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.4/torch_xla-2.5.0-cp39-cp39-manylinux_2_28_x86_64.whl
2.5 (CUDA 12,4 + Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.4/torch_xla-2.5.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.5 (CUDA 12,4 + Python 3.11) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.4/torch_xla-2.5.0-cp311-cp311-manylinux_2_28_x86_64.whl
Ночью (Python 3.8) https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev-cp38-cp38-linux_x86_64.whl
Ночью (Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev-cp310-cp310-linux_x86_64.whl
Ночью (CUDA 12.1 + Python 3.8) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.6.0.dev-cp38-cp38-linux_x86_64.whl
Используйте ночную сборку до 13.08.2024 Вы также можете добавить `+yyyymmdd` после` torch_xla-nightly`, чтобы получить ночное колесо указанной даты. Вот пример: 
 pip3 install torch==2.6.0.dev20240925+cpu --index-url https://download.pytorch.org/whl/nightly/cpu
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20240925-cp310-cp310-linux_x86_64.whl

Версия колеса Torch 2.6.0.dev20240925+cpu можно найти по адресу https://download.pytorch.org/whl/nightly/torch/.

Используйте ночную сборку после 20.08.2024

Вы также можете добавить yyyymmdd после torch_xla-2.6.0.dev , чтобы получить ночное колесо указанной даты. Вот пример: 

 pip3 install torch==2.5.0.dev20240820+cpu --index-url https://download.pytorch.org/whl/nightly/cpu
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0.dev20240820-cp310-cp310-linux_x86_64.whl

Версия колеса Torch 2.6.0.dev20240925+cpu можно найти по адресу https://download.pytorch.org/whl/nightly/torch/.

старые версии
Версия Облачное колесо VMS TPU
2.4 (Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.4.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.3 (Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.3.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.2 (Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.2.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.1 (xrt + python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/xrt/tpuvm/torch_xla-2.1.0%2Bxrt-cp310-cp310-manylinux_2_28_x86_64.whl
2.1 (Python 3.8) https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.1.0-cp38-cp38-linux_x86_64.whl

Версия Колесо графического процессора
2.5 (CUDA 12.1 + Python 3.9) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0-cp39-cp39-manylinux_2_28_x86_64.whl
2.5 (CUDA 12.1 + Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.5 (CUDA 12.1 + Python 3.11) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.5.0-cp311-cp311-manylinux_2_28_x86_64.whl
2.5 (CUDA 12,4 + Python 3.9) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.4/torch_xla-2.5.0-cp39-cp39-manylinux_2_28_x86_64.whl
2.5 (CUDA 12,4 + Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.4/torch_xla-2.5.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.5 (CUDA 12,4 + Python 3.11) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.4/torch_xla-2.5.0-cp311-cp311-manylinux_2_28_x86_64.whl
2.4 (CUDA 12.1 + Python 3.9) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.4.0-cp39-cp39-manylinux_2_28_x86_64.whl
2.4 (CUDA 12.1 + Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.4.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.4 (CUDA 12.1 + Python 3.11) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.4.0-cp311-cp311-manylinux_2_28_x86_64.whl
2.3 (CUDA 12.1 + Python 3.8) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.3.0-cp38-cp38-manylinux_2_28_x86_64.whl
2.3 (CUDA 12.1 + Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.3.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.3 (CUDA 12.1 + Python 3.11) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.3.0-cp311-cp311-manylinux_2_28_x86_64.whl
2.2 (CUDA 12.1 + Python 3.8) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.2.0-cp38-cp38-manylinux_2_28_x86_64.whl
2.2 (CUDA 12.1 + Python 3.10) https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.2.0-cp310-cp310-manylinux_2_28_x86_64.whl
2.1 + cuda 11,8 https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/11.8/torch_xla-2.1.0-cp38-cp38-manylinux_2_28_x86_64.whl
ночь + cuda 12.0> = 2023/06/27 https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.0/torch_xla-nightly-cp38-cp38-linux_x86_64.whl

Докер

Версия Облачный TPU VMS Docker
2.5 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.5.0_3.10_tpuvm
2.4 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm
2.3 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.3.0_3.10_tpuvm
2.2 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.2.0_3.10_tpuvm
2.1 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.1.0_3.10_tpuvm
Ночной питон us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm

Чтобы использовать приведенные выше докеры, пожалуйста, пройдите --privileged --net host --shm-size=16G . Вот пример: 

docker run --privileged --net host --shm-size=16G -it us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm /bin/bash

Версия GPU CUDA 12.4 Docker
2.5 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.5.0_3.10_cuda_12.4
2.4 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_cuda_12.4

Версия GPU CUDA 12.1 Docker
2.5 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.5.0_3.10_cuda_12.1
2.4 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_cuda_12.1
2.3 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.3.0_3.10_cuda_12.1
2.2 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.2.0_3.10_cuda_12.1
2.1 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.1.0_3.10_cuda_12.1
ночью us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.8_cuda_12.1
Ночью на свидание us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.8_cuda_12.1_YYYYMMDD

Версия GPU CUDA 11,8 + Docker
2.1 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.1.0_3.10_cuda_11.8
2.0 us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.0_3.8_cuda_11.8

Запустить на вычислительные экземпляры с помощью графических процессоров.

Поиск неисправностей

Если Pytorch/XLA не выполняется, как и ожидалось, см. Руководство по устранению неполадок, в котором есть предложения по отладке и оптимизации вашей сети.

Предоставление обратной связи

Команда Pytorch/XLA всегда рада услышать от пользователей и участников OSS! Лучший способ обработать это - подарить проблему на этот GitHub. Вопросы, отчеты об ошибках, запросы на функции, проблемы с построением и т. Д. Все приветствуются!

Внося

См. Руководство по взносу.

Отказ от ответственности

Этот репозиторий совместно эксплуатируется и поддерживается Google, Meta и рядом отдельных участников, перечисленных в файле участников. По вопросам, направленным на Meta, отправьте электронное письмо по адресу [email protected]. По вопросам, указанным в Google, пожалуйста, отправьте электронное письмо по адресу [email protected]. Для всех остальных вопросов, пожалуйста, откройте проблему в этом хранилище здесь.

Дополнительные чтения

Вы можете найти дополнительные полезные материалы для чтения в

  • Отладка производительности на облачной TPU VM
  • Ленивый тензор вступление
  • Масштабирование рабочих нагрузок с глубоким обучением с помощью Pytorch / XLA и Cloud TPU VM
  • Масштабирование моделей Pytorch на облачных TPU с FSDP

Связанные проекты

  • Openxla
  • Объятие
  • JetStream
Расширять
Дополнительная информация
Связанные приложения
Рекомендуем вам
Связанные новости Все