Tối ưu hóa tinh chỉnh mô hình video và hình ảnh quy mô lớn với NVIDIA NeMo Automodel và Diffusers

Quay lại danh sách bài viết

Một bài đăng chung từ NVIDIA và Hugging Face. Đặc biệt cảm ơn Sayak Paul từ Hugging Face vì những đóng góp cho công việc tích hợp và đồng tác giả bài blog này.

Các mô hình khuếch tán (diffusion models) là động lực đằng sau một số bản phát hành mã nguồn mở thú vị nhất trong hai năm qua — như FLUX.1-dev cho chuyển đổi văn bản thành hình ảnh, cùng với Wan 2.1 và HunyuanVideo cho chuyển đổi văn bản thành video. Thư viện 🤗 Diffusers đã trở thành "ngôi nhà" mặc định cho các mô hình này, cung cấp cho các nhà nghiên cứu và nhà phát triển một giao diện nhất quán, duy nhất để thực hiện suy luận (inference), tùy chỉnh và kết hợp các pipeline.

Ngoài ra, việc huấn luyện và tinh chỉnh (fine-tuning) các mô hình khuếch tán cũng đang gia tăng, đòi hỏi các tiện ích cung cấp khả năng phân mảnh (sharding) tiết kiệm bộ nhớ, lưu trữ cache không gian tiềm ẩn (latent caching), phân nhóm đa độ phân giải (multiresolution bucketing) và các cấu hình có khả năng mở rộng linh hoạt từ một GPU lên đến hàng trăm GPU.

Để đáp ứng những nhu cầu kỹ thuật này, chúng tôi cung cấp thư viện mã nguồn mở NVIDIA NeMo Automodel. Hôm nay, chúng tôi làm nổi bật sự hợp tác giữa NVIDIA và Hugging Face nhằm mang đến khả năng huấn luyện khuếch tán phân tán, đạt chuẩn sản xuất cho bất kỳ mô hình nào ở định dạng Diffusers trên Hugging Face Hub — mà không cần chuyển đổi checkpoint hay viết lại mô hình cho bất kỳ mô hình mới nào. Việc tích hợp này đã được ghi lại trong hướng dẫn huấn luyện Diffusers và hoàn toàn là mã nguồn mở theo giấy phép Apache 2.0.

Mục lục

NeMo Automodel là gì?

NeMo Automodel là một thư viện huấn luyện mã nguồn mở dựa trên PyTorch DTensor, thuộc framework NVIDIA NeMo, được xây dựng dựa trên hai nguyên tắc thiết kế quan trọng đối với hệ sinh thái Diffusers:

AutoModel hiện chỉ hỗ trợ các mô hình flow-matching. Về cơ bản, nó sử dụng flow matching làm mục tiêu huấn luyện, kết hợp với huấn luyện không gian tiềm ẩn (thông qua các đầu ra VAE đã mã hóa trước) và tải dữ liệu theo nhóm đa độ phân giải để tăng tốc độ xử lý (throughput).

Các mô hình khuếch tán được hỗ trợ

Việc tích hợp NeMo Automodel đi kèm với các công thức tinh chỉnh sẵn sàng sử dụng cho các mô hình khuếch tán mã nguồn mở dưới đây. Danh sách này phản ánh các công thức hiện có trong examples/diffusion/finetune.

Những lợi ích từ sự hợp tác này

Đối với người dùng Diffusers, những lợi ích thực tế được chia thành một vài khả năng cụ thể.

Không cần chuyển đổi checkpoint. Các trọng số đã được huấn luyện trước (pretrained weights) từ Hub hoạt động ngay lập tức. Không cần "định dạng huấn luyện" riêng biệt để chuyển đổi qua lại. Checkpoint sau khi tinh chỉnh của bạn có thể tải trực tiếp vào DiffusionPipeline để suy luận hoặc tải ngược lên Hub để chia sẻ. Các công cụ hạ nguồn (downstream tools) — như lượng tử hóa, biên dịch, LoRA adapters, bộ lấy mẫu tùy chỉnh — đều tiếp tục hoạt động bình thường.

Hỗ trợ nhanh chóng cho các mô hình mới. Khi một mô hình khuếch tán mới xuất hiện trong Diffusers, việc kích hoạt nó trong NeMo Automodel chỉ cần một phần bổ sung mã nhỏ, gọn — một trình xử lý tiền xử lý dữ liệu và một bộ điều hợp mô hình (model adapter) — thay vì phải viết một tập lệnh huấn luyện tùy chỉnh hoàn toàn. Phần còn lại của ngăn xếp công thức (FSDP2, tải dữ liệu theo nhóm, checkpointing, tạo nội dung) vẫn giữ nguyên và quy trình làm việc dựa trên YAML tương tự vẫn được áp dụng.

Tinh chỉnh toàn phần và tinh chỉnh hiệu quả tham số (PEFT). Cả tinh chỉnh toàn phần (full fine-tuning) và PEFT kiểu LoRA đều được hỗ trợ, vì vậy bạn có thể chọn giữa chất lượng tối đa (tinh chỉnh toàn phần trên cụm máy chủ lớn) hoặc hiệu quả tối đa (LoRA trên một node đơn lẻ). Cấu trúc công thức tương tự xử lý cả hai trường hợp.

Huấn luyện có khả năng mở rộng vượt xa các tập lệnh tích hợp sẵn. NeMo Automodel bổ sung các sơ đồ phân mảnh như FSDP2, song song hóa tensor, ngữ cảnh và pipeline, điều phối đa node (hiện tại là SLURM, sắp tới là Kubernetes) và phân nhóm đa độ phân giải. Những khả năng này giúp việc huấn luyện các mô hình lớn hơn như FLUX.1-dev (12B) và HunyuanVideo (13B) trở nên khả thi.

Tổng quan về quy trình tinh chỉnh

Trong phần này, chúng ta sẽ đi qua quy trình điển hình để tinh chỉnh bất kỳ mô hình nào được hỗ trợ. Cách cài đặt Automodel được khuyến nghị là sử dụng container Docker NeMo Automodel (nvcr.io/nvidia/nemo-automodel:26.06), vốn đã đi kèm với PyTorch, TransformerEngine và các phụ thuộc đã biên dịch CUDA được xây dựng sẵn. Ngoài ra, bạn có thể cài đặt bằng lệnh pip3 install nemo-automodel hoặc từ mã nguồn (pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git); xem hướng dẫn cài đặt để biết tất cả các tùy chọn.

Hướng dẫn này đi qua quá trình tinh chỉnh full-transformer của FLUX.1-dev trên bộ dữ liệu bài tarot Rider–Waite gồm 78 lá, sau đó tạo hình ảnh từ checkpoint thu được. Nó sử dụng lại các cấu hình YAML đã có và áp dụng các cài đặt cụ thể cho lần chạy thông qua ghi đè dòng lệnh, vì vậy không cần tệp cấu hình mới.

1. Tiền mã hóa bộ dữ liệu

Công thức khuếch tán tiêu thụ các VAE latents đã lưu vào cache và các nhúng văn bản (text embeddings) thay vì mã hóa hình ảnh nguồn trong mỗi bước huấn luyện. Truyền phát 78 hình ảnh Rider–Waite trực tiếp từ Hugging Face và phân phối quá trình tiền xử lý trên tất cả các GPU khả dụng:

Các chú thích (captions) đã chứa token kích hoạt trtcrd. Với ngân sách pixel này và tỷ lệ khung hình chân dung của bộ dữ liệu, quá trình tiền xử lý gán các mẫu vào nhóm 384×640 được sử dụng cho lần chạy trình diễn.

Đối với huấn luyện hình ảnh, quá trình tiền xử lý tạo ra các tệp cache.pt và siêu dữ liệu đã phân mảnh:

2. Khởi chạy huấn luyện với YAML FLUX hiện có

Sử dụng trực tiếp examples/diffusion/finetune/flux_t2i_flow.yaml. Tệp YAML đã chọn FLUX.1-dev, tinh chỉnh full transformer, bộ điều hợp FLUX flow-matching, kích thước batch hiệu dụng là 32 và FSDP2 tám chiều.

Cung cấp các đường dẫn và cài đặt cụ thể cho tarot dưới dạng ghi đè dòng lệnh:

Quá trình chạy tạo ra các checkpoint tại các bước 50, 100, 150 và 200. Checkpoint cuối cùng được gắn nhãn epoch_66_step_199; nhãn này bắt đầu từ số 0 mặc dù nó đại diện cho bước tối ưu hóa thứ 200 đã hoàn thành.

3. Tạo hình ảnh từ checkpoint đã tinh chỉnh

Sử dụng YAML tạo hình ảnh FLUX hiện có và trỏ model.checkpoint vào checkpoint huấn luyện hoàn chỉnh:

Bao gồm trtcrd để kích hoạt phong cách tarot đã học. Để so sánh đối chứng, hãy giữ nguyên seed và bối cảnh nhưng bỏ qua token kích hoạt:

Kết quả

Tại bước 200, các câu lệnh (prompts) về phi hành gia có chứa token kích hoạt vẫn giữ được nội dung yêu cầu trong khi có thêm bảng màu cổ điển kem, đỏ và đen, các đường viền mực đậm, các mảng màu phẳng, tông màu giấy cũ và bố cục thẻ bài ngụ ngôn. Phi hành gia không có token kích hoạt vẫn giữ nguyên tính chất ảnh chụp, chứng minh rằng hiệu ứng đã học gắn liền với trtcrd thay vì thay thế mô hình cơ sở trên toàn bộ.

4. Hiệu suất