Hugging Face ra mắt Tokenizers v1: Tối ưu hóa hiệu năng vượt trội

Trước đây, tokenizer (bộ tách từ) chưa bao giờ là nút thắt cổ chai trong các quy trình làm việc ML. Về mặt tính toán, quá trình tokenization rất nhẹ so với các tác vụ mô hình hóa nặng nề diễn ra trong phần còn lại của pipeline. Tuy nhiên, trong một số trường hợp, nó đã nhanh chóng trở thành yếu tố then chốt để tăng tốc (hoặc làm chậm) công việc học máy của bạn.

Khi các mô hình trở nên nhanh hơn và khối lượng công việc tăng lên, sự cân bằng đó bắt đầu thay đổi. Việc huấn luyện trên các tập dữ liệu khổng lồ, phục vụ nhiều yêu cầu đồng thời hoặc xử lý lặp đi lặp lại các đầu vào dài có thể gây áp lực đủ lớn lên tokenizer khiến nó làm mô hình bị "đói" dữ liệu.

Đây là lý do tại sao chúng tôi tập trung mạnh mẽ vào hiệu năng cho phiên bản 1 sắp tới của tokenizers. Tokenization cần phải nhẹ và có khả năng mở rộng theo quy trình làm việc của bạn. GPU của bạn không bao giờ nên ở trạng thái nhàn rỗi chỉ để chờ CPU hoàn tất quá trình tokenization.

Trong bài viết này, chúng ta sẽ xem xét những yếu tố giúp v1 nhanh hơn v0.23, thường là nhanh hơn gấp hàng chục lần.

Công việc này hoàn toàn có thể thực hiện được nhờ vào phần còn lại của hệ sinh thái. Tokenization là một lĩnh vực rất sôi động trong cộng đồng mã nguồn mở, và các thư viện như gigatoken, tiktoken, kitoken, tokie, fastokens, wordchipper và ai-tokenizer, cùng nhiều thư viện khác, mỗi cái đều đã thúc đẩy giới hạn về tốc độ của một tokenizer. Chúng tôi đã nghiên cứu các công trình đó, và một vài ý tưởng dưới đây đã đến với chúng tôi vì một dự án khác đã chứng minh rằng chúng đáng để thử nghiệm.

Trước khi tái cấu trúc này, tokenizers không hề đạt được hiệu năng như mong đợi, vì vậy việc đóng góp cho nó có vẻ không đáng giá. Với lần tái cấu trúc này, chúng tôi hy vọng sẽ làm rõ rằng chúng tôi dự định biến tokenizers thành một thư viện xứng đáng để đóng góp.

Chúng tôi cũng cảm ơn IBM, NVIDIA và đội ngũ ExecuTorch vì đã đóng góp các bản vá và giúp chúng tôi thử nghiệm trên nhiều loại phần cứng để mở rộng khả năng hỗ trợ nền tảng.

Kết quả

Chúng tôi trình bày kết quả cho bản ứng viên phát hành (release candidate) của tokenizers v1 so với các giải pháp thay thế được sử dụng rộng rãi khác. Chúng tôi xem xét các khía cạnh: đơn luồng, đa luồng, khả năng mở rộng qua các luồng, so sánh theo từng mô hình, so sánh theo từng ngôn ngữ, độ trễ, thông lượng giải mã, bộ nhớ heap, cũng như kích thước crate.

Chúng tôi chạy thử nghiệm này từ kho lưu trữ tokbench và thêm một lệnh để bạn có thể chạy lại các benchmark trên phần cứng của mình nếu muốn.

V1 là gì

v1 sẽ tạo ra các token ID giống hệt như v0.23. Mục tiêu là giữ nguyên đầu ra, API, từ vựng và các merge rank, đồng thời cải thiện mọi thứ có thể cải thiện. Điều đó bao gồm cả tính bao quát. Thư viện vẫn giữ tính tổng quát trên các họ tokenizer thay vì chỉ chuyên biệt cho BPE, vì vậy v1 tải được mọi thứ mà v0.23 đã tải.

Một tokenizer chuyển đổi văn bản thành danh sách các số nguyên mà mô hình có thể đọc được. tokenizers thực hiện quá trình chuyển đổi đó qua bốn giai đoạn. Normalization áp dụng các thao tác như viết thường hoặc chuẩn hóa Unicode cho văn bản thô. Pre-tokenization chia văn bản thành các phần nhỏ hơn gọi là pre-token. Mô hình biến mỗi pre-token thành các token và ánh xạ chúng tới các ID trong từ vựng của nó. Post-processing thêm bất kỳ token đặc biệt nào mà mô hình yêu cầu.

Giai đoạn mô hình là nơi hầu hết công việc được mô tả ở đây diễn ra. Tám trong số mười họ mô hình được đo lường trong bài viết này sử dụng byte pair encoding, hay BPE. BPE bắt đầu từ các byte của một pre-token và liên tục kết hợp cặp liền kề có thứ hạng cao nhất cho đến khi không còn cặp nào. Thứ hạng được học khi tokenizer được huấn luyện và đi kèm với nó, vì vậy cùng một văn bản luôn tạo ra cùng một ID. Một phép hợp (merge) không bao giờ vượt qua ranh giới pre-token. Hai họ còn lại sử dụng WordPiece và Unigram, hai loại mô hình khác mà thư viện hỗ trợ.

Trang tokenization pipeline ghi lại bốn giai đoạn này. Tokenization algorithms ghi lại BPE, WordPiece và Unigram.

Mỗi giai đoạn đều đã được cải tiến. Đây là những thay đổi quan trọng:

Phân tách: Bitstream thay vì Regex

Các mô hình BPE sử dụng biểu thức chính quy (regex) để chia văn bản đầu vào thành các đoạn nhỏ hơn, dễ xử lý hơn gọi là pre-token. Các phép hợp xảy ra bên trong một pre-token và không bao giờ vượt qua ranh giới giữa hai pre-token, vì vậy việc phân tách này quyết định những gì phần còn lại của pipeline nhìn thấy.

Biểu thức chính quy đó là một tham số cố định của mô hình. Nó đi kèm với tokenizer và không bao giờ thay đổi trong thời gian chạy, vì vậy không cần một công cụ regex đa năng để diễn giải nó trong mỗi lần encode. Một hàm phân tách tương đương có thể được viết thủ công, một lần duy nhất, cho mô hình cụ thể đó.

Một hàm viết tay sau đó có thể sử dụng các chỉ dẫn SIMD (single instruction, multiple data) của CPU hiện đại, áp dụng một thao tác cho nhiều byte cùng lúc và rất phù hợp với văn bản UTF-8. bitcannon xem các byte của đầu vào như các luồng bit song song, vì vậy các ranh giới được xác định từ các phép toán boolean trên toàn bộ thanh ghi thay vì quét từng ký tự một. Nó quyết định 64 byte cho mỗi thao tác thanh ghi. Ý tưởng tương tự thúc đẩy Parabix cho xử lý văn bản và simdjson cho JSON.

Điều này phụ thuộc vào việc nhận diện mẫu. Một số ít các ngữ pháp bao phủ hầu hết các mô hình BPE cấp byte, và một tokenizer có mẫu không nằm trong số đó sẽ giữ lại đường dẫn regex và không nhận được tốc độ này. Đó là lý do tại sao mức tăng trưởng ở trên lại khác biệt nhiều đến vậy.

Bộ nhớ đệm từ (Word Cache)

Văn bản thực tế chứa nhiều từ lặp lại. Vì BPE luôn tạo ra cùng một token ID cho một pre-token nhất định, v1 có thể lưu kết quả sau khi xử lý nó một lần. Một bộ nhớ đệm cục bộ theo luồng (thread-local cache) ánh xạ các byte của mỗi pre-token tới các token ID của nó, cho phép các lần xuất hiện sau đó bỏ qua quá trình hợp.

Đương nhiên, khi đầu vào tăng lên, số lượng từ duy nhất có thể tăng chậm hơn tổng số từ. Các từ lặp lại sau đó chiếm tỷ trọng ngày càng lớn trong đầu vào. Các từ mới vẫn xuất hiện, điều này giải thích cho những lần bỏ lỡ (misses) không thường xuyên trong hình ảnh động bên dưới.

Tái tạo kết quả chia sẻ tiền tố (shared-prefix) với:

Bộ nhớ đệm hoạt động tốt nhất khi đầu vào chứa các pre-token lặp lại. Đầu vào với ít pre-token lặp lại có thể phải trả chi phí cho việc tra cứu mà không nhận được nhiều kết quả trùng khớp.

Vòng lặp hợp (Merge Loop)

Chi phí lớn tiếp theo đến từ vòng lặp hợp BPE. Đối với mỗi pre-token, vòng lặp liên tục tìm cặp liền kề có ưu tiên cao nhất và hợp nhất nó. Cách triển khai trước đây đã cấp phát bộ nhớ mới cho mỗi lần gọi và xây dựng một hàng đợi ưu tiên mới cho mỗi pre-token.

v1 tái sử dụng một bộ đệm tạm (scratch buffer) do người gọi sở hữu, loại bỏ các lần cấp phát lặp đi lặp lại đó. Nó lưu trữ các ký hiệu trong một mảng phẳng và liên kết các ký hiệu liền kề theo vị trí của chúng trong mảng đó, giúp việc cập nhật trong quá trình hợp trở nên rẻ hơn. Nó cũng xử lý một lô (batch) các pre-token trong một lần gọi mô hình duy nhất.

Mỗi cặp ứng viên cũng được đóng gói thành một giá trị 64-bit duy nhất, với thứ hạng hợp nằm ở các bit cao. Việc so sánh hai ứng viên lúc này chỉ là so sánh hai số nguyên, và "không có phép hợp nào ở đây" là giá trị lớn nhất có thể, vì vậy vòng lặp tìm thấy phép hợp tiếp theo mà không cần rẽ nhánh.

Phương pháp

Những khác biệt nhỏ trong thiết kế benchmark có thể tạo ra sự khác biệt lớn về hiệu năng của tokenizer. Chúng tôi đã sử dụng các quy tắc sau để giữ cho việc so sánh nhất quán giữa các engine.

Việc encode lặp đi lặp lại một tài liệu có thể nhanh hơn so với việc encode một luồng các tài liệu khác nhau trên cùng một bản dựng. Cách tiếp cận đầu tiên đo lường hiệu năng khi toàn bộ tài liệu đã có trong bộ nhớ đệm. Cách thứ hai đo lường hiệu năng trên đầu vào mới trong khi vẫn cho phép các pre-token đã thấy trước đó được lưu trong bộ nhớ đệm.

Cả hai điều kiện đôi khi được mô tả là "warm" (ấm), mặc dù chúng đo lường các khối lượng công việc khác nhau. Các kết quả chính của chúng tôi sử dụng các tài liệu riêng biệt, và toàn bộ kho ngữ liệu quá lớn để vừa với bộ nhớ đệm. Các benchmark của tokenizer nên xác định rõ khối lượng công việc mà chúng sử dụng vì lựa chọn này có thể chi phối kết quả.

Tổng kết lại

Trên mười họ mô hình mà đường dẫn encode của v1 bao phủ, nó encode văn bản nhanh hơn từ 3 đến 30 lần so với v0.23 với một luồng trên Apple M4 Max. Mức thấp nhất là t5-base, mức cao nhất là gpt2. Nó mở rộng ở mức 76% tuyến tính trên tám worker. Trong suốt những thay đổi này, v1 tạo ra chính xác các token ID giống như thư viện đã phát hành.

Sự cải thiện tổng thể đến từ nhiều thay đổi hoạt động cùng nhau: một bộ tách viết tay thay thế cho engine regex, một bộ nhớ đệm trả lời từ lặp lại mà không cần hợp lại, một vòng lặp hợp không bao giờ chạm vào bộ cấp phát, và một lần gọi mô hình cho mỗi lô pre-token thay vì một lần cho mỗi pre-token. Mỗi thay đổi làm giảm khối lượng công việc tại một điểm khác nhau trong pipeline.

Ưu tiên tiếp theo là hỗ trợ thêm nhiều họ mô hình hơn. Chúng tôi sẽ chuyển các mô hình bổ sung sang vòng lặp hợp mới trước phiên bản 1.0.0. Khi các ứng viên phát hành ổn định, bước tiếp theo sẽ là mang lại những cải tiến này vào thư viện transformers và phần còn lại của hệ sinh thái vốn phụ thuộc vào thư viện tokenizers.

Bài viết này được tạo ra từ kết quả của tokbench và sẽ được cập nhật khi khả năng hỗ trợ mở rộng.