Nội tình PyTorch (2019)

• Bài viết giải thích cấu trúc bên trong của PyTorch, như một hướng dẫn dành cho những ai muốn đóng góp vào codebase C++ của PyTorch
• Mục tiêu của bài viết là giúp hiểu cấu trúc thư viện tensor của PyTorch và kỹ thuật vi phân tự động (autograd), đồng thời hỗ trợ định hướng trong codebase

Cấu trúc cơ bản của tensor trong PyTorch

• Trong PyTorch, tensor là cấu trúc dữ liệu cơ bản nhất
• Tensor là cấu trúc dữ liệu n chiều, có thể lưu các giá trị vô hướng như số thực dấu phẩy động (float), số nguyên (int), v.v.
• Tensor bao gồm các metadata sau:
  • kích thước (size): thông tin về các chiều của tensor
  • dtype: kiểu dữ liệu được lưu trữ (ví dụ: float32, int64, v.v.)
  • device: nơi dữ liệu được lưu trữ (CPU, CUDA, v.v.)
  • stride: thông tin offset của dữ liệu trong bộ nhớ vật lý

• Vai trò của stride

  • stride được dùng để chuyển chỉ số logic thành vị trí bộ nhớ vật lý
  • stride thiết lập offset cho từng chiều và xác định vị trí bộ nhớ vật lý bằng cách nhân chỉ số với giá trị stride
  • nhờ stride, có thể xem cùng một dữ liệu theo cách khác dưới dạng view mà không cần tạo tensor mới

Khái niệm tensor và bộ lưu trữ (Storage)

• Trong PyTorch, tensor không trực tiếp giữ dữ liệu thực tế → dữ liệu được quản lý trong bộ lưu trữ (Storage)
• Tensor = kích thước + dtype + device + stride + offset
• Nhiều tensor có thể chia sẻ một bộ lưu trữ → hỗ trợ khái niệm view
• Nhờ việc tách biệt storage và tensor, bộ nhớ có thể được sử dụng hiệu quả hơn

Quy trình dispatch của phép toán tensor

• Trong PyTorch, phép toán đi qua hai bước dispatch:
  1. Dispatch dựa trên loại thiết bị và layout
     • Tùy theo tensor CPU hay tensor CUDA mà thực thi mã cài đặt khác nhau
  2. Dispatch dựa trên dtype
     • Tùy theo kiểu dữ liệu như float hay int mà gọi kernel khác nhau

Mô hình mở rộng tensor của PyTorch

• Ba yếu tố mở rộng chính của tensor:

  • Device: định nghĩa cách cấp phát bộ nhớ trên CPU, GPU, TPU, v.v.
  • Layout: định nghĩa cách tensor được lưu trong bộ nhớ (ví dụ: lưu liên tục, lưu thưa (sparse), v.v.)
  • dtype: định nghĩa kiểu dữ liệu được lưu trong từng phần tử của tensor

• Các tùy chọn mở rộng:

  • Có thể mở rộng tensor bằng cách sửa trực tiếp mã PyTorch
  • Có thể viết lớp wrapper bao quanh tensor hiện có
  • Nếu cần wrapper trong quá trình autograd thì phải tự mở rộng trực tiếp

Cách hoạt động của vi phân tự động (Autograd)

• PyTorch thực hiện vi phân tự động dựa trên reverse-mode differentiation
• Khi tính forward sẽ tạo đồ thị → khi backpropagation sẽ duyệt đồ thị để tính đạo hàm
• Autograd quản lý thêm các thông tin sau:
  • AutogradMeta: metadata gắn với tensor, được dùng cho backpropagation
  • Ghi lại kết quả phép toán và thực hiện đạo hàm khi backpropagation

Cấu trúc mã nguồn PyTorch và vị trí các tệp

• Các thư mục chính trong codebase PyTorch:
  • torch/ → mô-đun Python (mã Python)
  • torch/csrc/ → mã binding giữa Python và C++, engine autograd, trình biên dịch JIT, v.v.
  • aten/ → định nghĩa các phép toán tensor (bao gồm phần lớn các phép toán cốt lõi)
  • c10/ → định nghĩa các cấu trúc lõi như tensor và storage

Quy trình thực thi phép toán trong PyTorch

• Ví dụ: quy trình khi gọi torch.add():
  1. Chuyển đổi đối số từ Python sang mã C++
  2. Thực hiện dispatch trong VariableType
  3. Thực hiện dispatch dựa trên Device/layout
  4. Chạy kernel cuối cùng

Quy trình viết kernel và công cụ

• Trong PyTorch, kernel được viết theo các bước sau:
  1. Viết metadata cho phép toán: định nghĩa chữ ký hàm, thiết bị hỗ trợ và kiểu dữ liệu hỗ trợ
  2. Kiểm tra đầu vào: xác thực đầu vào như số chiều, kiểu dữ liệu, v.v.
  3. Cấp phát tensor đầu ra
  4. Dispatch theo dtype: chạy kernel theo kiểu dữ liệu
  5. Xử lý song song: trên CPU dùng OpenMP, trên CUDA dùng cơ chế song song tích hợp
  6. Truy cập dữ liệu và tính toán: dùng TensorAccessor, TensorIterator, v.v.

Các macro dispatch chính

• AT_DISPATCH_ALL_TYPES → thực hiện dispatch theo dtype
• Có hỗ trợ macro cho nhiều kiểu dữ liệu khác nhau → có thể tối ưu hiệu năng

Mẹo tối ưu hiệu năng và nâng cao hiệu quả làm việc

• Giảm tối đa việc sửa tệp header → sửa sẽ khiến toàn bộ mã phải build lại
• Thiết lập môi trường phát triển cục bộ → giảm thời gian lãng phí khi dùng CI
• Dùng ccache → có thể tiết kiệm thời gian biên dịch lại
• Dùng máy chủ mạnh → có thể rút ngắn thời gian biên dịch C++ và build CUDA

Hướng dẫn đóng góp cho PyTorch

• Những hạng mục phù hợp để bắt đầu đóng góp:
  • Issue có nhãn triaged → issue đã được nhà phát triển PyTorch xem xét
  • Cải thiện tài liệu và hỗ trợ tái hiện lỗi
  • Đưa ra ý kiến về các RFC (đề xuất tính năng) của PyTorch
• PyTorch đã phát triển nhờ các cộng tác viên mã nguồn mở và luôn chào đón sự tham gia của cộng đồng