AVTensor: Bộ giải mã đa phương tiện hiệu năng cao bằng Rust dành cho quy trình huấn luyện AI

AVTensor: Một bộ giải mã phương tiện (media decoder) hiệu năng cao viết bằng Rust cho các quy trình huấn luyện (training pipelines)

Ngày 10 tháng 7 năm 2026

Bởi

Rik Heijdens, Kỹ sư Nghiên cứu

Các container phương tiện lưu giữ thời gian theo những cách rất dễ gây ra sai sót tinh vi—ví dụ như dấu thời gian (timestamp) âm hoặc các luồng âm thanh và video không bắt đầu cùng lúc. Sau khi một bộ giải mã có sẵn vô tình làm lệch âm thanh và video trong dữ liệu huấn luyện của chúng tôi, chúng tôi đã tự viết một bộ giải mã riêng. Hôm nay, chúng tôi mở mã nguồn dự án này: AVTensor, một bộ giải mã phương tiện bằng Rust giúp tách luồng (demux) cả hai luồng trong một lần quét duy nhất và đã cải thiện MFU (Model Flops Utilization) trong quá trình huấn luyện của chúng tôi thêm 1,8 điểm phần trăm.

Việc huấn luyện các mô hình nghe nhìn (audio-visual) đòi hỏi âm thanh và video phải đồng bộ. Một quy trình giải mã một đoạn clip thường mặc định rằng khung hình i của video tương ứng với cửa sổ mẫu i của âm thanh. Chỉ cần lệch vài chục mili giây là đủ để nhận ra, đặc biệt là với khẩu hình (lip sync)—tuy nhiên các container lại không mã hóa sự đồng bộ đó theo bất kỳ cách trực diện nào.

Để hiểu lý do, hãy xem xét thực chất một tệp video là gì. Một tệp MP4 hoặc MOV là một container chứa nhiều luồng độc lập: thường là một luồng video, một hoặc nhiều luồng âm thanh và đôi khi là phụ đề. Mỗi luồng được lưu trữ dưới dạng một chuỗi các gói nén và duy trì đồng hồ riêng. Khi bạn phát tệp, một bộ tách luồng (demuxer) sẽ tách các gói và chuyển chúng đến các bộ giải mã (codec), nơi chúng được giải mã thành các khung hình video và các mẫu âm thanh.

Bên trong một container phương tiện: các luồng độc lập, mỗi luồng chạy trên đồng hồ riêng.

Ba đặc tính của một tệp MP4 thông thường gây ra hầu hết các rắc rối:

Chúng tôi đã gặp phải các tài nguyên (assets) có gói video đầu tiên với PTS (Presentation Timestamp) là −3,337 giây trong khi gói âm thanh đầu tiên nằm ở −0,023 giây. Hầu hết các trình phát đều hiển thị điều này một cách hợp lý (bằng cách áp dụng siêu dữ liệu thời gian của container như edit lists), đó là lý do tại sao nó không bị phát hiện—trong các tập dữ liệu và trong một số công cụ chúng tôi dùng để kiểm tra chúng.

Vào thời điểm đó, quy trình của chúng tôi giải mã video bằng torchcodec và âm thanh bằng một thư viện hoàn toàn khác—StreamingMediaDecoder của torchaudio—sau đó ghép kết quả lại với nhau. Trên các tài nguyên có PTS âm, khung hình đầu tiên của torchcodec xuất hiện ở 1,6 giây (nó loại bỏ mọi thứ trước mốc 0), trong khi trình đọc âm thanh lại bắt đầu từ một gốc thời gian khác. Khung hình đầu tiên bị đóng băng trong hơn một giây trong khi âm thanh vẫn phát và phần còn lại của clip bị chạy lệch nhịp.

Một tài nguyên thực tế từ dữ liệu của chúng tôi: ba điểm bắt đầu khác nhau cho cùng một clip.

Đó là một lỗi âm thầm: quy trình huấn luyện vẫn chạy bình thường. Nhưng mô hình lại học được rằng khẩu hình và lời nói chỉ khớp với nhau một cách lỏng lẻo, và bạn bắt đầu đi tìm nguyên nhân trong kiến trúc—bởi vì bản năng đầu tiên luôn là nghĩ đó là vấn đề của RoPE (Rotary Positional Embeddings), phải không?

Đó là lý do tại sao chúng tôi liên tục nhắc lại câu thần chú: luôn luôn kiểm tra dữ liệu. Chúng tôi đã truy vết các vấn đề về chất lượng trong quá trình huấn luyện nghe nhìn và phát hiện video và âm thanh được giải mã trả về độ dài hiệu dụng và điểm gốc khác nhau. Giải pháp tạm thời của chúng tôi là cắt xén phòng vệ (cắt cả hai luồng theo luồng ngắn hơn và khẳng định chúng khớp nhau trong phạm vi 10 mili giây), cách này chỉ xử lý triệu chứng mà không sửa được bộ giải mã.

torchcodec, bộ giải mã phương tiện chính thức của PyTorch, rất tốt cho mục đích mà nó được xây dựng: giải mã video truy cập ngẫu nhiên với dấu thời gian cho từng khung hình. Nhưng khi chúng tôi đánh giá nó (phiên bản 0.11 vào thời điểm đó), bộ giải mã âm thanh của nó vẫn còn sơ khai—phía upstream đã làm lại logic tìm kiếm (seeking) trong các bản phát hành sau này—vì vậy đường dẫn âm thanh của chúng tôi phải chạy qua StreamingMediaDecoder của torchaudio. Không thư viện nào được thiết kế để giải mã âm thanh và video cùng nhau trên một dòng thời gian duy nhất, và trên các tài nguyên có PTS âm, chúng không thống nhất được đâu là mốc thời gian 0. Sự bất đồng đó chính là nguyên nhân gây ra khung hình đầu tiên bị đóng băng. Vì vậy, chúng tôi đã tự viết bộ giải mã riêng.

AVTensor

AVTensor là một thư viện Rust giao tiếp trực tiếp với API C của FFmpeg và cung cấp một lệnh gọi duy nhất cho Python: cung cấp cho nó một URI và một khoảng thời gian, bạn sẽ nhận lại một tensor video và một tensor âm thanh có chung một điểm gốc.

Chúng tôi đã và đang chuyển đổi cơ sở hạ tầng nền tảng tại Runway sang Rust. Rust không phải là ngôn ngữ hiển nhiên cho một bộ giải mã, vì hệ sinh thái phương tiện và ML chủ yếu là C và C++, và mọi lệnh gọi vào FFmpeg hoặc libtorch đều phải vượt qua ranh giới FFI (Foreign Function Interface). Trên thực tế, chi phí này là không đáng kể: các vòng lặp nóng (hot loops) dù sao cũng chạy bên trong FFmpeg và tch-rs bao bọc libtorch rất tốt đến mức các khung hình đã giải mã đi thẳng vào các tensor của torch. Chúng tôi chọn Rust vì loại mã này—quản lý thủ công các bộ đệm và vòng đời của FFmpeg—chính là nơi các dự án C++ dễ bị lỗi, còn trình biên dịch Rust sẽ bắt được các lỗi bộ nhớ và luồng ngay tại thời điểm biên dịch thay vì khi đang vận hành.

Bộ tách luồng (demuxer) chia luồng byte của container thành các gói; chúng tôi định tuyến từng gói đến đúng bộ giải mã (hoặc loại bỏ nếu nó nằm ngoài phạm vi yêu cầu). Các khung hình đã giải mã sau đó đi qua một đồ thị bộ lọc (filter graph) để chuẩn hóa tốc độ khung hình, độ phân giải và định dạng pixel cho video, cũng như tốc độ lấy mẫu, bố cục kênh và độ lớn cho âm thanh. Các khung hình đã lọc sẽ nằm trong các tensor đầu ra. Vì cả hai luồng đều đến từ một lần tách luồng duy nhất trên một tệp với dòng thời gian được chia sẻ, âm thanh và video luôn thống nhất về mốc thời gian 0 theo cấu trúc.

Việc cấp phát bộ nhớ diễn ra một lần ngay từ đầu. AVTensor tính toán các hình dạng đầu ra—[khung hình, chiều cao, chiều rộng, 3] cho video, [kênh, mẫu] cho âm thanh—sau đó ghi các khung hình đã giải mã trực tiếp vào các lát cắt (slices) của chúng, vì vậy các mẫu vẫn liên tục và không có gì phải cấp phát lại giữa chừng quá trình giải mã. Sự đánh đổi là độ dài đầu ra chỉ là dự đoán tốt nhất từ dấu thời gian của luồng và siêu dữ liệu container, vốn có thể sai trên các tài nguyên bị lỗi. Toàn bộ quá trình giải mã cũng chạy với Global Interpreter Lock của Python được giải phóng, điều này rất quan trọng vì các trình tải dữ liệu (dataloaders) của chúng tôi là các quy trình dựa trên luồng (xây dựng trên SPDL) thay vì các nhóm tiến trình worker: các luồng Python vẫn tiếp tục hoạt động trong khi mã gốc thực hiện giải mã.

Một ràng buộc mà chúng tôi cố tình giữ lại: API tìm kiếm theo dấu thời gian, không phải theo chỉ số khung hình. Các bộ tách luồng được xây dựng để tìm kiếm dựa trên thời gian—khái niệm duy nhất bao trùm cả hai luồng—trong khi "khung hình 500" là khái niệm chỉ dành cho video và có thể yêu cầu quét toàn bộ luồng bit để xác định. Các quy trình suy nghĩ theo chỉ số khung hình (ranh giới cảnh, căn chỉnh phụ đề) sẽ tính toán trước ánh xạ từ khung hình sang PTS một lần và mọi thứ phía sau sẽ sử dụng PTS.

AVTensor cũng truyền phát trực tiếp từ bộ lưu trữ đối tượng (object storage). Quy trình trước đây bắt đầu bằng giai đoạn tải xuống, kéo từng tài nguyên vào bộ nhớ và chuyển các byte cho bộ giải mã—tạo ra một bản sao thừa của mỗi video trên heap của Python và phải tải xuống toàn bộ ngay cả khi quá trình huấn luyện chỉ cần năm giây của một clip.

FFmpeg có một lối thoát gốc mà CLI không hiển thị: một AVIOContext tùy chỉnh, cho phép bạn cung cấp các lệnh gọi lại (callbacks) đọc và tìm kiếm mà bộ tách luồng sử dụng làm nguồn dữ liệu. AVTensor triển khai một bộ được hỗ trợ bởi các yêu cầu phạm vi byte đồng thời tới GCS hoặc S3 (các URL HTTP(S) thông thường đi qua lớp giao thức riêng của FFmpeg), vì vậy bộ tách luồng đọc từ bộ lưu trữ đối tượng như thể đó là một tệp cục bộ.

Chúng tôi đã ghi lại mọi lệnh gọi lại để xem cách FFmpeg đọc một tệp MP4 trong thực tế: nó đọc 4 KiB đầu tiên cho tiêu đề, hỏi kích thước tệp, sau đó—trên các tệp không được mã hóa với faststart—nhảy thẳng đến phần cuối của tệp, vì đó là nơi chứa moov atom lưu trữ chỉ mục tìm kiếm. Từ đó, nó hạ cánh gần dấu thời gian được yêu cầu và trong quá trình giải mã, nó nhảy qua lại giữa các phạm vi byte gần đó khi luân phiên giữa dữ liệu âm thanh và video xen kẽ. Truyền phát tuần tự xử lý mô hình này rất kém, vì vậy trình đọc sẽ tìm nạp các phân đoạn lớn (64 MiB), lưu vào bộ đệm các phạm vi đã tìm nạp trước đó cho các bước nhảy lùi ngắn và khi tìm kiếm xa, nó sẽ chạy đua một yêu cầu phạm vi mới với các trình đọc đang hoạt động.

Mô hình đọc của FFmpeg trên tệp MP4, từ các lệnh gọi lại AVIO đã ghi lại.

Việc tìm kiếm vào giữa một tài nguyên lớn giờ đây chỉ tải xuống tiêu đề, chỉ mục và các byte xung quanh mục tiêu—thay vì toàn bộ tệp. Điều đó cũng thay đổi cách chúng tôi lưu trữ: việc trích xuất một cảnh quay không còn yêu cầu phải cắt nó thành một đối tượng riêng trước, vì vậy quy trình có thể huấn luyện trên các cảnh quay trực tiếp từ các clip đầy đủ thay vì phải trả phí lưu trữ mọi tài nguyên hai lần.

Trong một bài kiểm tra hiệu năng tải dữ liệu huấn luyện, việc thay thế quy trình tải xuống-rồi-giải mã bằng truyền phát trực tiếp từ bộ lưu trữ đối tượng đã cắt giảm 100 batch từ 176 giây xuống còn 122 giây (khoảng 30%) và loại bỏ giai đoạn tải xuống cũng như các bản sao trên heap. Trong một lần chạy huấn luyện nghe nhìn thực tế, sự thay đổi này mang lại khoảng 1,8 điểm phần trăm MFU so với quy trình dựa trên torchcodec. MFU đo lường tỷ lệ tính toán lý thuyết của một đội tàu GPU đang thực hiện công việc hữu ích, vì vậy ở quy mô của một cụm huấn luyện, ngay cả một phần mười điểm cũng là rất đáng kể.

Và lỗi đồng bộ đã biến mất. Chúng tôi đã lấy tài nguyên gây lỗi từ sự cố trước đó và hiện không còn vấn đề lệch nhịp nữa.

Hiệu năng so với torchcodec (phiên bản 0.14 tại thời điểm viết bài) là tương đương như trong biểu đồ. Sự khác biệt chính xuất hiện ở việc thay đổi kích thước (resize) trong quá trình giải mã, điều phổ biến đối với các trình tải dữ liệu huấn luyện cần giảm tỷ lệ trong khi giải mã. Việc giảm tỷ lệ từ 1080p xuống 256×144 chạy nhanh hơn 1,8 lần trong AVTensor với một luồng FFmpeg và nhanh hơn khoảng 6 lần với tính năng đa luồng tự động của FFmpeg, vì việc thay đổi kích thước của đồ thị bộ lọc chồng lấp với quá trình giải mã và các bản ghi đầu ra nhỏ hơn—trong khi phép biến đổi resize của torchcodec chạy trên luồng giải mã.

Thời gian giải mã trung bình trên một clip 1080p dài 30 giây.

AVTensor hiện đã là mã nguồn mở. Đây là đường dẫn video mặc định trong ngăn xếp xử lý dữ liệu của chúng tôi, cung cấp dữ liệu cho quá trình tiền xử lý và huấn luyện trên các mô hình nghe nhìn của chúng tôi. Chúng tôi hoan nghênh mọi phản hồi và đóng góp.

Khám phá thêm