PyTorch 2.0 đã ra mắt: Tất cả những gì chúng ta biết

PyTorch là một framework học sâu mã nguồn mở do cộng đồng dẫn dắt, cung cấp cách linh hoạt và hiệu quả để xây dựng các mô hình máy học. Nó có giao diện thân thiện, được cộng đồng hỗ trợ rộng rãi và tích hợp mượt mà với hệ sinh thái Python.

PyTorch 2.0 đã giới thiệu những thay đổi căn bản đối với các thao tác biên dịch cốt lõi, đồng thời vẫn giữ nguyên mức độ quen thuộc và dễ dùng cho nhà phát triển. Bản nâng cấp mới này hứa hẹn hiệu năng được tăng tốc và mở rộng hỗ trợ cho Dynamic Shapes và Distributed.

Có gì mới trong PyTorch 2.0?

PyTorch đang chuyển một số phần từ C++ trở lại Python, giúp chạy nhanh hơn và dễ tùy biến. Với phiên bản 2.0, họ đã giới thiệu `torch.compile`, thay đổi cách PyTorch vận hành ở tầng trình biên dịch. Tính năng này là tùy chọn và không ảnh hưởng đến mã cũ của bạn.

PyTorch 2.0 Compile

Để tạo nền tảng vững chắc cho tính năng `torch.compile`, các công nghệ mới đã được giới thiệu:

• TorchDynamo. Trình biên dịch JIT ở cấp độ Python được thiết kế riêng để tăng tốc PyTorch. Bằng cách tích hợp với API đánh giá khung (frame evaluation) trong CPython, nó động thời sửa đổi bytecode Python khi chạy, cho phép thực thi mã nhanh hơn.
• AOTAutograd. Bộ công cụ hỗ trợ nhà phát triển tăng tốc quá trình huấn luyện mô hình trên PyTorch. Nó truy vết trước cả đồ thị forward và backward. Ngoài ra, AOTAutograd cung cấp cơ chế đơn giản để biên dịch mượt mà các đồ thị đã trích xuất bằng những trình biên dịch học sâu tiên tiến.
• PrimTorch. Bằng cách giảm mạnh số lượng toán tử của PyTorch từ hơn 2000 xuống khoảng 250 toán tử nguyên thủy, PrimTorch đã đơn giản hóa đáng kể việc phát triển tính năng hoặc backend cho PyTorch.
• TorchInductor. Trình biên dịch học sâu thuần PyTorch, tự động ánh xạ các mô hình PyTorch sang mã sinh cho nhiều bộ tăng tốc và backend. TorchInductor sử dụng OpenAI Triton làm khối xây dựng cho tăng tốc GPU.

Tất cả các công nghệ mới đều được viết bằng Python và hỗ trợ Dynamic shapes. Điều này giúp PyTorch mới chạy mã nhanh hơn, linh hoạt hơn và dễ tùy biến, hạ thấp rào cản gia nhập.

Ví dụ mã

Hãy cùng xem cách triển khai PyTorch Compiler nhanh gọn.

Không dùng torch.compile

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")

Không dùng torch.compile

Để tăng hiệu năng mô hình, chỉ cần thêm wrapper torch.compile quanh mô hình để nhận một mô hình đã biên dịch. Cắm là chạy.

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")
compiled_model = torch.compile(model)

Học cách xây dựng các mô hình học sâu với thư viện PyTorch qua khóa học Deep Learning with PyTorch.

Bạn có thể huấn luyện mô hình của mình mà không cần thay đổi gì.

import torch
model = torch.compile(model)
for batch in dataloader:
    run_epoch(model, batch)

Hoặc bạn có thể chạy suy luận mô hình.

model = torch.compile(model)
model(**input)

torch.compile() đi kèm các tham số bổ sung:

1. mode: bạn có thể chỉ định trình biên dịch nên tối ưu điều gì khi biên dịch.
2. dynamic: dùng để bật đường chạy cho Dynamic Shapes.
3. fullgraph: biên dịch chương trình thành một đồ thị duy nhất.
4. backend: mặc định dùng TorchInductor, nhưng bạn có thể chỉ định các backend trình biên dịch khác.

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  **kwargs
) -> torch._dynamo.NNOptimizedModule

Tìm hiểu các kiến thức cơ bản về API PyTorch và tự xây dựng một mạng nơ-ron đơn giản từ đầu bằng Hướng dẫn PyTorch của chúng tôi.

Đo hiệu năng

Khi bắt đầu với PyTorch 2.0, các nhà phát triển đã dùng 163 mô hình mã nguồn mở (46 HuggingFace Transformers, 61 TIMM và 56 TorchBench) để tạo bộ đo hiệu năng cho tính năng compile mới. Bộ đo bao gồm các tác vụ như phân loại ảnh, tạo ảnh, mô hình ngôn ngữ, hệ thống gợi ý và học tăng cường.

Kết quả cho thấy hiệu năng cải thiện đáng kể khi huấn luyện trên GPU NVIDIA A100.

Lưu ý: hiện tại, backend mặc định chỉ hỗ trợ CPU và dòng GPU Nvidia Volta và Ampere.

Đo hiệu năng trình biên dịch PyTorch trên GPU NVIDIA A100

Đây mới là khởi đầu, và trong các bản cập nhật sắp tới, bạn sẽ thấy thêm nhiều cải tiến về hiệu năng và khả năng mở rộng.

Cách cài đặt PyTorch 2.0

Bạn có thể cài đặt phiên bản mới hơn của PyTorch chỉ với pip.

Sao chép và dán lệnh dưới đây vào terminal của bạn.

Cho GPU: CUDA 11.8

Các phiên bản GPU mới hơn cho thấy hiệu năng cải thiện rõ rệt.

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu118

Cho GPU: CUDA 11.7

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117

Cho CPU:

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu

Để xác minh:

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

Tăng tốc Hugging Face với PyTorch 2.0

Hãy thử hàm torch.compile để tăng tốc các transformer của Hugging Face. Bạn có thể làm cho mã Hugging Face chạy nhanh hơn chỉ với một decorator một dòng.

Lưu ý: Với torch.compile(), chúng tôi đã thấy hiệu năng huấn luyện tăng 30%–200% - Bảng điều khiển hiệu năng TorchDynamo.

Trong ví dụ này, chúng ta sẽ áp dụng torch.compile lên mô hình ngôn ngữ lớn “dolly-v2-3b" để suy luận nhanh hơn. Để chạy mã trên Google Colab, trước tiên cần cài đặt các thư viện Python cần thiết.

%%capture
%pip install transformers accelerate xformers

Tiếp theo, chúng ta sẽ tải và nạp tokenizer và mô hình ngôn ngữ bằng Hugging Face transformers. Sau đó, thêm nn.Module vào hàm torch.compile().

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    GenerationConfig,
    pipeline,
)


tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-3b")
model = AutoModelForCausalLM.from_pretrained(
    "databricks/dolly-v2-3b", device_map="auto", torch_dtype=torch.bfloat16
)
model = torch.compile(model) #only line of code is required

Ở bước cuối, chúng ta sẽ chuyển văn bản thành token bằng tokenizer, đưa vào model.generate, rồi giải mã đầu ra sinh được thành văn bản bằng tokenizer.batch_decode.

prompt = "I love you because"
inputs = tokenizer(prompt, return_tensors="pt").to(device="cuda:0")


# Generate
generate_ids = model.generate(inputs.input_ids, max_length=50)
tokenizer.batch_decode(
    generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)[0]

Như ta thấy, “dolly-v2” đã hoàn thành câu bằng cách thêm, “you are a good person…”

"I love you because you are a good person.
You are kind, you help others, you are honest, you are loyal, you are humble, you are humble, you are humble.
You are a good person."

Nó cũng hoạt động với pipeline của Hugging Face. Chỉ cần cung cấp loại tác vụ, mô hình và tokenizer.

generator = pipeline("text-generation", model= model,tokenizer=tokenizer)
generator("What is the name of Germany's Capital?")

Đầu ra:

[{'generated_text': "What is the name of Germany's Capital? 
The name of Germany's Capital is Berlin."}]

Hàm compile hoạt động ngay với các thư viện Python transformers, accelerate và TIMM.

Nếu bạn chưa chắc nên bắt đầu sự nghiệp học sâu và AI từ đâu, lộ trình Machine Learning Scientist with Python là điểm khởi đầu tuyệt vời. Nó cung cấp tổng quan toàn diện về các kỹ năng Python then chốt giúp bạn có thể giành được công việc nhà khoa học máy học.