Mở rộng các trình xuất

torch.export theo dõi mô hình để tạo thành một đồ thị, các giai đoạn sau đó biến đổi đồ thị này, và giai đoạn cuối cùng hạ cấp hoặc xuất nó cho runtime đích. Hầu hết các mô hình đi qua mà không bị thay đổi. Khi có một pattern của PyTorch mà backend không xử lý được, trình xuất sẽ áp dụng một giải pháp thay thế nhỏ ở giai đoạn đó thay vì chỉnh sửa mô hình.

Thêm một giải pháp thay thế bằng cách viết một hàm và đăng ký nó với một decorator. Mỗi giải pháp thay thế nên được đặt ở giai đoạn thấp nhất có thể diễn đạt nó một cách rõ ràng.

Bản vá và sửa lỗi

Một giải pháp tạm thời có thể là một bản vá (patch) hoặc một bản sửa lỗi (fix). Hai cách này khác nhau ở chỗ chúng có thể bị hoàn tác hay không.

Cả hai đều nằm trong một registry trong exporters/utils.py, và exporter sẽ cài đặt mọi thứ đã được đăng ký cho backend của nó ở giai đoạn phù hợp.

Hãy dùng một bản vá khi vấn đề là lỗi hạ tầng (lowering) của một backend duy nhất: thiếu bản dịch ONNX, một đặc điểm kiểm tra của ORT, hoặc một phép phân rã FX tạo ra một toán tử chết. Biện pháp xử lý nằm trong trình xuất (exporter), và mã mô hình hóa vẫn được giữ sạch sẽ.

Khi các khối pattern xuất ra trên nhiều backend khác nhau, chẳng hạn như các vòng lặp phụ thuộc vào dữ liệu, bộ nhớ đệm có trạng thái nằm ngoài Cache, hoặc chú ý tách vòng lặp viết tay, hãy sửa model. Việc sửa một lần trong model sẽ giúp ích cho mọi trình xuất.

Thêm một bản vá

Giả sử một phương thức của model thực hiện một thao tác mà torch.export không thể trace được. Phương thức NllbMoeTop2Router._cast_classifier của NLLB-MoE chuyển đổi trọng số của classifier sang một dtype khác, đây là thao tác không thể trace được. Hãy thay thế nó bằng một thao tác không làm gì (no-op) trong suốt quá trình export.

Viết một factory nhận phương thức gốc và trả về phương thức thay thế, sau đó đăng ký factory với đường dẫn chấm của phương thức:

from transformers.exporters.utils import register_patch

@register_patch("dynamo", "transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier")
def _patch_classifier_cast(_original):
    # Replace the untraceable dtype cast with a no-op during export.
    return lambda self, *args, **kwargs: None

Bộ xuất (exporter) thay thế phương thức trước khi theo dõi (tracing) và khôi phục nó sau đó, do đó bản vá chỉ ảnh hưởng đến quá trình xuất. Một vài biến thể:

• Truyền thêm các đường dẫn để chia sẻ một nhà máy giữa các vị trí gọi, ví dụ @register_patch("dynamo", path_a, path_b).
• Chỉnh sửa một op torch bằng cách trỏ đường dẫn đến nó, ví dụ @register_patch("onnx", "torch.where"). Hàm factory nhận op thực tế làm đối số, do đó bản thay thế có thể gọi xuyên qua đến nó.
• Viết một bản sửa lỗi thay vì một bản vá khi bạn cần viết lại đồ thị sau khi theo dõi. Cơ chế là giống nhau, một hàm được trang trí trong registry khớp.

Tham chiếu giai đoạn

Mã nguồn của mỗi exporter gắn nhãn các giai đoạn của nó bằng các khối chú thích # ── Stage N: … ──, nên tệp và tài liệu tham khảo này khớp nhau. Hãy xem ở đó để biết chính xác các thao tác và lớp mà mỗi giai đoạn xử lý.

DynamoExporter

Trình xuất cơ bản chạy một giai đoạn vá và bốn hàm trợ giúp, theo thứ tự, bên trong DynamoExporter.export (xem exporter_dynamo.py).

1. Patch cho chữ ký phía trước: cung cấp cho model.forward một chữ ký tham số phẳng để torch.export không gộp các đầu vào vào một bộ **kwargs duy nhất. Đây là phần nội bộ và không phải là điểm mở rộng.
2. Các bản vá mô hình: thay thế các phương thức mô hình không thể truy vết bằng các phương thức tương đương an toàn cho xuất trong quá trình truy vết. Mở rộng bằng @register_patch("dynamo", ...).
3. Đăng ký Pytree: đăng ký từng Cache và ModelOutput để torch.export có thể làm phẳng và xây dựng lại chúng (thường xảy ra tự động). Thêm một nhánh vào _flatten_to_context / _unflatten_from_context cho một kiểu dữ liệu mà việc duyệt thuộc tính không thể tiếp cận được.
4. Hình dạng động: gán Dim.AUTO cho mọi tensor và cache lá khi dynamic=True. Ghi đè bằng DynamoConfig.dynamic_shapes.
5. Dọn dẹp trạng thái: đặt lại các thuộc tính tensor mà mô hình thiết lập bên trong forward mà torch.export để lại dưới dạng tensor giả. Mở rộng bằng cách thêm tên thuộc tính vào _STATEFUL_CACHE_ATTRS.

OnnxExporter

OnnxExporter thêm năm giai đoạn bao quanh torch.onnx.export (xem exporter_onnx.py). Tìm kiếm trong tệp để có danh sách đầy đủ các bản vá:

grep -nE "^def (_patch_|_fix_|_aten_)" src/transformers/exporters/exporter_onnx.py

1. Torch patches: thay thế các op torch mà trình xuất ONNX không thể chuyển đổi trực tiếp. Mở rộng bằng @register_patch("onnx", ...).
2. Các bản vá ONNX: chạy lại các bản sửa lỗi node sau run_decompositions để các node shape-guard mới được sửa trước khi hạ cấp. Sử dụng cùng một registry @register_patch("onnx", ...).
3. Sửa lỗi node FX: viết lại các node đồ thị mà trình xuất ONNX không thể hạ cấp, chẳng hạn như các op alias, view in-place và các assert chết. Mở rộng bằng @register_fx_node_fix("onnx").
4. Bản dịch ONNX: cung cấp một quy trình hạ cấp (lowering) tùy chỉnh cho một op aten khi quy trình mặc định bị thiếu hoặc có lỗi (ví dụ aten.index_put hoặc aten._grouped_mm). Thêm một hàm _aten_* vào _ONNX_TRANSLATION_TABLE.
5. Sửa lỗi ONNX IR: viết lại chương trình ONNX đã hoàn thành để khắc phục các lỗi của ONNX Runtime (ví dụ ép buộc TopK(sorted=True)). Thêm một hàm _fix_ir_* vào _IR_FIXES.

ExecutorchExporter

ExecutorchExporter thêm năm giai đoạn bao quanh to_edge_transform_and_lower và to_executorch, bắt đầu với việc chuẩn bị backend (xem exporter_executorch.py).

1. Chuẩn bị backend: di chuyển mô hình sang thiết bị và dtype đích và chọn bộ phân vùng của nó (prepare_for_xnnpack, prepare_for_cuda). Thêm một backend bằng cách đăng ký prepare_for_<name> trong _BACKEND_PREPARE.
2. Torch patches: thay thế các op torch mà các backend của ExecuTorch không thể chấp nhận, chẳng hạn như split_copy, chunk và topk(k>dim). Mở rộng bằng @register_patch("executorch", ...).
3. Các bản vá ExecuTorch: thay thế các thành phần bên trong của ExecuTorch gây lỗi trên các đồ thị hình dạng động hợp lệ. Sử dụng cùng một bộ đăng ký @register_patch("executorch", ...).
4. Sửa lỗi chương trình FX: sửa chương trình đã xuất khi việc sửa lỗi cần ngữ cảnh của toàn bộ chương trình, chẳng hạn như mở rộng ràng buộc phạm vi hoặc điền vào metadata placeholder bị thiếu. Mở rộng bằng @register_fx_program_fix("executorch").
5. Sửa lỗi nút FX: viết lại các nút riêng lẻ, chẳng hạn như ánh xạ các toán tử sym Python sang executorch_prim.* hoặc viết lại pow thành chuỗi mul. Mở rộng bằng @register_fx_node_fix("executorch").

Các giải pháp thay thế đã biết từ nguồn

Một số lớp model gặp lỗi đã được xác nhận trong trình tối ưu hóa đồ thị của onnxscript (constant folding bị crash trên SplitToSequence, các initializer của FPN bị loại bỏ). ONNX_DISABLE_OPTIMIZE tắt tối ưu hóa onnxscript cho các model đó. Mỗi mục ghi lại vấn đề upstream ngay cạnh tên model. Danh sách này dự kiến sẽ thu hẹp khi các lỗi upstream được khắc phục, do đó một mục mới phải tham chiếu đến một lỗi upstream cụ thể thay vì tắt tối ưu hóa một cách tùy tiện.

EXPORT_SKIPS, loại trừ một số lớp model khỏi quy trình export hoàn toàn khi model về cơ bản không thể export được ở dạng hiện tại (luồng điều khiển phụ thuộc dữ liệu không thể vector hóa, hoặc các module được xử lý như các tham số forward). Mỗi mục trong danh sách này đều kèm theo lý do nêu rõ thay đổi cần thực hiện phía model. Danh sách này cũng được kỳ vọng sẽ thu hẹp lại, không mở rộng thêm.