TorchDynamo
Chụp biểu đồ biến mã mô hình do người dùng viết bằng Python thành biểu đồ, đây là nền tảng của tất cả các phần tổng hợp. Sau khi thử rất nhiều giải pháp, PyTorch dường như đã khóa TorchDynamo là hướng chụp đồ thị trong tương lai, vì vậy viết một chút về TorchDynamo chủ yếu để giải thích lý do tại sao điều này được thực hiện (Tôi đã rời FB được một năm và nội dung chủ yếu dựa trên đoán và hiểu của riêng tôi).
Cố gắng giải thích TorchDynamo làm gì trong một câu: sử dụng API của PEP523 ( https://peps.python.org/pep-0523/ ) để lấy mã byte của khung trước khi người dùng thực thi từng khung python và chuyển đổi một phần Sử dụng phương pháp theo dõi để trích xuất biểu đồ (và gửi nó đến phần phụ trợ để biên dịch) và giữ nguyên các phần không xác định. Trả lại mã byte đã sửa đổi cho CPython để chạy.
Vì cả LazyTensor và TorchDynamo đều thực hiện theo dõi và cả hai đều là cách chụp biểu đồ nỗ lực nhất, tức là chúng chỉ biên dịch phần có thể tự chụp và sử dụng Python để chạy (còn gọi là Python dự phòng) nếu không thể chụp được, vì vậy cả hai có thể trông giống nhau về ngoại hình.
Tuy nhiên, điểm khác biệt giữa hai giải pháp là TorchDynamo rất quan trọng:
LazyTensor là một giải pháp hoàn toàn dựa vào truy vết.Vấn đề không thể tránh khỏi là “bạn chỉ có thể nhìn thấy các phần được theo dõi và bạn chỉ có thể biết nơi bạn không thể theo dõi sau khi truy tìm.” Và mỗi khi mô hình được thực hiện, những phần không thể truy tìm được có thể khác nhau. Để đảm bảo tính chính xác, LazyTensor phải theo dõi lại mỗi khi nó được thực thi. Để đưa ra một ví dụ cực đoan, một torch.add(tensor, random.random()) được viết trong mô hình, trong đó ngẫu nhiên là một hàm Python mà LazyTensor không thể nhìn thấy hoặc chạm vào. Chỉ có thể theo dõi lại mới có thể đảm bảo tính chính xác.
Và khi TorchDynamo sửa đổi mã byte, mọi thứ sẽ khác:
- Tất cả các thông tin cần thiết có thể được nhìn thấy trong mã byte, vì vậy có thể chứng minh rằng “mã mẫu này không sử dụng những thứ lạ, vì vậy không cần phải truy tìm lại”.
- Chỉ cần chứng minh rằng “không cần theo dõi” không có nghĩa là không cần theo dõi, bởi vì mã của người dùng vẫn được Python chạy từng dòng một. Nhưng TorchDynamo lại ở đây: CPython chạy mã byte nào thì có thể thay bằng nó!
Do đó, nó có thể làm một việc như vậy: khi người dùng gọi một mô hình đã chụp, hầu hết mã Python trong mô hình không tồn tại, thậm chí cả chi phí thực thi tượng trưng cũng không tồn tại và được thay thế bằng mã gốc đã biên dịch. Điều này không thể thực hiện được trong tất cả các giải pháp chụp đồ thị từng phần trước đây:
- Ngay cả biểu đồ đã biên dịch của LazyTensor cũng phải được truy tìm bằng Python mỗi lần để tìm “Ồ, biểu đồ này tôi đã thấy trước đây”.
- @ torch.jit.script , @ tf.function , @jax.jit có thể thay thế mã python được trang trí bằng mã đã biên dịch, nhưng tất cả phụ thuộc vào người dùng để đưa cấu trúc lại sơ đồ con này vào một chức năng riêng biệt. Và TorchDynamo hoàn toàn tự động và không yêu cầu người dùng thay đổi mã.
- Ngoài việc bổ sung thêm khối lượng công việc, kiểu cấu trúc lại này cũng có thể xung đột với cấu trúc mã của người dùng, vì “ranh giới của biểu đồ được sử dụng để biên dịch” có thể không khớp với ” ranh giới trừu tượng mà mã người dùng yêu cầu” : Ví dụ: người dùng ban đầu muốn viết Ba hàm, nhưng cách tối ưu hóa tốt nhất là biến hai trong số các bán hàm thành một đồ thị, điều này sẽ khiến người dùng rất lúng túng.
Đây chỉ là ví dụ trực tiếp nhất. Do khả năng đọc và ghi bytecode, về mặt lý thuyết, TorchDynamo có thể truy cập nhiều thông tin hơn mà LazyTensor hoàn toàn không có, và làm được nhiều việc hơn (sẽ đề cập sau). Độ khó của việc đọc và ghi mã byte thấp hơn nhiều so với mã nguồn, vì vậy nó đã trở thành một giải pháp khả thi.
Chụp toàn bộ biểu đồ không hữu ích lắm?
Một số người có thể nói rằng những điều được đề cập ở trên không hữu ích lắm cho việc chụp toàn bộ biểu đồ.
Tôi nghĩ rằng đây thực sự là trường hợp: TorchDynamo là một giải pháp theo đuổi mục tiêu cuối cùng là chụp một phần biểu đồ. Nó có thể tăng tốc gần như tất cả các mô hình do Python triển khai ngay lập tức mà không cần thay đổi mã – tiền đề là chạy Python dưới dạng dự phòng . Tuy nhiên, những gì thường được yêu cầu để triển khai là chụp toàn bộ biểu đồ.Không thể sử dụng toàn bộ mô hình trong Python trong một biểu đồ.
Tiền đề của việc sử dụng theo dõi để thực hiện chụp toàn bộ biểu đồ là người dùng nên tránh tất cả những thứ không thể theo dõi trong mã Python. Ba điều phổ biến nhất mà người dùng cần làm là: sử dụng hình dạng tượng trưng, sử dụng luồng điều khiển tượng trưng và vô hiệu hóa thư viện tensor hiện tại Tất cả các thư viện khác ngoại trừ. Nếu người dùng thực hiện việc này, thì chỉ một lần truy tìm biểu tượng thông thường mới có thể chụp được biểu đồ hoàn chỉnh mà không cần cơ chế phức tạp như TorchDynamo. TorchDynamo có thể đơn giản hóa một chút khối lượng công việc để người dùng thực hiện việc này, nhưng tôi không nghĩ nó sẽ khác về cơ bản.
Quan điểm cá nhân của tôi là từ quan điểm thực tế, việc yêu cầu người dùng thực hiện những điều trên không quá phức tạp: không nên đề cập đến việc vô hiệu hóa các thư viện khác; ngay cả ngày nay PyTorch cũng không có biểu tượng hay {shape, control flow }, nhưng miễn là @torch.jit.script_if_tracing được sử dụng để xử lý một lượng nhỏ hình dạng biểu tượng và luồng điều khiển biểu tượng, hầu hết các mô hình đều có thể được chụp chính xác bởi torch.jit.tracecapture. Meta nên có hàng chục hoặc hàng trăm mô hình tầm nhìn được triển khai trong detectron2/d2go và hiện tại về cơ bản chúng được triển khai theo cách này (tôi có một bài viết khác https://ppwwyyxx.com/blog/2022/TorchScript-Tracing-vs -Scripting/ giới thiệu về chi tiết tại đây).
Việc chụp toàn bộ biểu đồ của TensorFlow rất đơn giản: TF có hình dạng biểu tượng tốt và luồng kiểm soát biểu tượng ngay từ ngày đầu tiên và nó sẽ được sử dụng hết. tf.autograph thậm chí còn tự động hóa một phần việc viết lại luồng điều khiển.
Do đó, người dùng vẫn cần thay đổi mã một chút. Tất nhiên, TorchDynamo có siêu năng lực “thay đổi mã byte mà người dùng muốn chạy”. Vì vậy, nếu bạn muốn, về mặt lý thuyết, bạn có thể làm cho công việc chụp toàn bộ biểu đồ của người dùng dễ dàng hơn. Ví dụ:
Một số nhánh ở giữa mô hình như if x.shape[0] > 100 có thể được chuyển tương đương đến phần đầu của mô hình thông qua suy luận hình dạng. Theo cách này, có thể chụp được một đồ thị con lớn hơn không có nhánh. Thứ này hiện được gọi là “bảo vệ” trong TorchDynamo.
Về lý thuyết, có thể tự động thay thế luồng điều khiển python bằng ký hiệu, tương tự như những gì tf.autograph thực hiện, ngoại trừ đầu vào là mã byte thay vì mã nguồn.
Chế độ “nopython” hiện tại của TorchDynamo là chụp toàn bộ biểu đồ. Nhưng có vẻ như đó không phải là trọng tâm của công việc.
PT2 sẽ cung cấp cơ sở hạ tầng cho chế độ xuất không có python cho các trường hợp phân phát nhạy cảm về hiệu suất và cạnh. Nhóm PT2 sẽ không điều khiển ngăn xếp từ đầu đến cuối này, nhưng chúng tôi sẽ duy trì vòng phản hồi với các nhóm phụ trách việc này và đảm bảo các thành phần mà chúng tôi bản dựng có thể tái sử dụng trong những tình huống này.
Nhưng đồng thời, PyTorch 2.0 gần đây đã cải thiện khả năng hỗ trợ cho các hình tượng trưng; một số lượng nhỏ toán tử luồng điều khiển cũng đã được thêm vào functorch. Đây là tin tốt cho việc chụp toàn bộ biểu đồ.
Kết luận
Nói chung, vì TorchDynamo tạo ra sự phức tạp ở cấp mã byte nên nó có thể làm những việc mà các giải pháp khác không làm được. Ưu điểm của nó chủ yếu là để chụp một phần biểu đồ: nó cho phép mã mô hình Python của người dùng được chụp và tăng tốc mà không cần sửa đổi. Điều này phản ánh cam kết của PyTorch đối với triết lý “Python là trên hết”. Cho dù sự gắn bó như vậy là cần thiết là một vấn đề quan điểm.
Ưu điểm chính của TorchDynamo đến từ việc đọc và ghi bytecode. Lỗi của trình biên dịch tập lệnh JIT cho thấy không có nhiều thứ có thể thực hiện được ở cấp mã nguồn.Thật thông minh khi TorchDynamo có thể thực hiện mọi thứ ở cấp mã byte. Tuy nhiên, để tái tạo hoàn chỉnh trình thông dịch bytecode CPython, khối lượng công việc, độ khó bảo trì (và khả năng xảy ra lỗi) của nó là không hề nhỏ.
Ngoài ra, TorchDynamo không hữu ích lắm cho việc chụp toàn bộ biểu đồ. Đối với các mô hình phức tạp, việc viết lại mà người dùng nên làm vẫn phải được thực hiện. Nhưng tôi đoán ít nhất 2.0 có thể có một tuyên bố rõ ràng về “những gì người dùng nên làm”.
Tất nhiên, việc PT2 có thể thực hiện tốt công việc trong trình biên dịch hay không còn phụ thuộc vào nhiều yếu tố khác: cách thiết kế IR, khi nào chuyên biệt hóa/biên dịch lại, các đặc điểm khác nhau của các chương trình phụ trợ khác nhau, v.v. Ví dụ: IR được sử dụng bởi TorchDynamo và LazyTensor thực sự khác nhau. Nhưng bài viết này chỉ đề cập đến việc chụp đồ thị, còn các vấn đề khác sẽ không được đề cập.
