หน้าแรก>การเขียนโปรแกรมที่เกี่ยวข้อง>ซอร์สโค้ดอื่น ๆ
ดาวน์โหลด

pytorch/xla

สถานะ CI ปัจจุบัน:

Pytorch/XLA เป็นแพ็คเกจ Python ที่ใช้คอมไพเลอร์การเรียนรู้ลึก XLA เพื่อเชื่อมต่อเฟรมเวิร์กการเรียนรู้ลึกของ Pytorch และคลาวด์ TPU คุณสามารถลองได้ทันทีฟรีบนคลาวด์ TPU VM เดียวกับ Kaggle!

ลองดูสมุดบันทึก Kaggle ของเราเพื่อเริ่มต้น:

  • การแพร่กระจายที่มั่นคงด้วย pytorch/xla 2.0
  • พื้นฐาน Pytorch/XLA แบบกระจาย

การติดตั้ง

TPU

ในการติดตั้ง pytorch/xla สร้างความเสถียรใน TPU VM ใหม่: 

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

ในการติดตั้ง pytorch/xla สร้างยามค่ำคืนใน TPU VM ใหม่: 

 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

ปลั๊กอิน GPU

ตอนนี้ 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
  • การทำโปรไฟล์บน TPU VM
  • คู่มือ GPU

รูปภาพและล้อที่มีอยู่

แพ็คเกจ Python

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 สาธารณะของเรา

รุ่น ล้อคลาวด์ GPU VM
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
ใช้งานสร้างทุกคืนก่อน 08/13/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 Wheel เวอร์ชัน 2.6.0.dev20240925+cpu สามารถดูได้ที่ https://download.pytorch.org/whl/nightly/torch/

ใช้งานสร้างทุกคืนหลังจาก 08/20/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 Wheel เวอร์ชัน 2.6.0.dev20240925+cpu สามารถดูได้ที่ https://download.pytorch.org/whl/nightly/torch/

รุ่นเก่ากว่า
รุ่น ล้อคลาวด์ TPU VMS
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

รุ่น ล้อ GPU
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
Nightly + 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

นักเทียบท่า

รุ่น Cloud 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

ในการทำงานบนอินสแตนซ์การคำนวณด้วย GPU

การแก้ไขปัญหา

หาก Pytorch/XLA ไม่ได้ทำงานตามที่คาดไว้ให้ดูคู่มือการแก้ไขปัญหาซึ่งมีคำแนะนำสำหรับการดีบักและเพิ่มประสิทธิภาพเครือข่ายของคุณ

ให้ข้อเสนอแนะ

ทีม Pytorch/XLA ยินดีที่จะได้ยินจากผู้ใช้และผู้สนับสนุน OSS! วิธีที่ดีที่สุดในการเข้าถึงคือการยื่นปัญหาเกี่ยวกับ GitHub นี้ คำถาม, รายงานข้อผิดพลาด, คำขอคุณสมบัติ, การสร้างปัญหา ฯลฯ ยินดีต้อนรับ!

การบริจาค

ดูคู่มือการบริจาค

คำเตือน

พื้นที่เก็บข้อมูลนี้ดำเนินการร่วมกันและดูแลโดย Google, Meta และผู้มีส่วนร่วมจำนวนหนึ่งที่ระบุไว้ในไฟล์ผู้มีส่วนร่วม สำหรับคำถามที่กำกับที่ Meta โปรดส่งอีเมลไปที่ [email protected] สำหรับคำถามที่กำกับที่ Google โปรดส่งอีเมลไปที่ [email protected] สำหรับคำถามอื่น ๆ โปรดเปิดปัญหาในที่เก็บนี้ที่นี่

อ่านเพิ่มเติม

คุณสามารถค้นหาสื่อการอ่านที่มีประโยชน์เพิ่มเติมได้ใน

  • ประสิทธิภาพการดีบักบนคลาวด์ TPU VM
  • อินโทร Lazy Tensor
  • ปรับขนาดเวิร์กโหลดการเรียนรู้ลึกด้วย pytorch / xla และ Cloud TPU VM
  • ปรับขนาดโมเดล pytorch บนคลาวด์ TPU ด้วย FSDP

โครงการที่เกี่ยวข้อง

  • Openxla
  • กอด
  • jetstream
ขยาย
ข้อมูลเพิ่มเติม
แอปที่เกี่ยวข้อง
แนะนำสำหรับคุณ
ข้อมูลที่เกี่ยวข้อง ทั้งหมด