Khắc phục lỗi hết bộ nhớ khi huấn luyện Flux LoRA trên RTX 4090

Bạn có RTX 4090 với 24GB VRAM, được cho là đủ cho việc huấn luyện Flux cục bộ, nhưng mỗi lần thử đều crash với lỗi CUDA out of memory. Quá trình huấn luyện bắt đầu tốt, có thể chạy vài bước, rồi dừng lại. Bạn đã thử giảm kích thước batch nhưng vẫn crash. Điều gì đang tiêu thụ hết bộ nhớ đó?

Câu trả lời nhanh: Lỗi OOM khi huấn luyện Flux LoRA trên RTX 4090 xảy ra vì kiến trúc lớn của Flux yêu cầu 30-40GB VRAM ở cài đặt huấn luyện mặc định. Khắc phục điều này bằng cách bật gradient checkpointing để đổi tính toán lấy bộ nhớ, giảm kích thước batch xuống 1, sử dụng độ phân giải huấn luyện 512x512 thay vì 1024x1024, bật mixed precision với FP16 hoặc BF16, và sử dụng các triển khai attention tiết kiệm bộ nhớ. Các cài đặt này cho phép huấn luyện hoàn chỉnh trên 24GB trong khi tạo ra LoRA chất lượng.

RTX 4090 là card tuyệt vời cho huấn luyện Flux LoRA khi được cấu hình đúng cách. Vấn đề là các cấu hình huấn luyện mặc định giả định VRAM nhiều hơn 24GB. Với các cài đặt đúng, bạn có thể huấn luyện Flux LoRA chất lượng cao một cách hiệu quả. Hãy cấu hình thiết lập của bạn cho huấn luyện ổn định.

Tại sao huấn luyện Flux cần nhiều VRAM như vậy?

Hiểu những gì tiêu thụ bộ nhớ trong quá trình huấn luyện giúp bạn tối ưu hóa hiệu quả.

Kích thước mô hình

Mô hình cơ sở của Flux lớn hơn đáng kể so với SDXL hoặc SD 1.5. Chỉ riêng trọng số mô hình độ chính xác đầy đủ đã tiêu thụ khoảng 23GB. Đó đã là toàn bộ VRAM của 4090 trước khi huấn luyện bắt đầu.

Trong quá trình huấn luyện, bạn cần bộ nhớ cho mô hình, gradient, trạng thái optimizer và activation. Mỗi thứ trong số này có thể tiếp cận kích thước của chính mô hình.

Đối với huấn luyện LoRA cụ thể, bạn đóng băng mô hình cơ sở và chỉ huấn luyện các lớp adapter nhỏ. Điều này giúp đáng kể nhưng không loại bỏ áp lực bộ nhớ từ activation và gradient chảy qua mô hình đầy đủ.

Bộ nhớ activation

Trong quá trình forward pass, các activation trung gian được lưu trữ để sử dụng trong backward pass. Các activation này tăng theo kích thước batch và độ phân giải.

Ở độ phân giải 1024x1024, bộ nhớ activation có thể vượt quá kích thước mô hình. Một batch huấn luyện đơn có thể yêu cầu 15-20GB chỉ cho activation.

Đây là lý do tại sao huấn luyện crash sau vài bước. Bước đầu tiên có thể vừa, nhưng phân mảnh bộ nhớ và trạng thái tích lũy khiến các bước tiếp theo thất bại.

Trạng thái optimizer

Các optimizer như Adam lưu trữ hai giá trị momentum cho mỗi tham số có thể huấn luyện. Điều này nhân đôi bộ nhớ cần thiết cho các tham số đang được huấn luyện.

Đối với fine-tuning đầy đủ, bộ nhớ trạng thái optimizer bằng 2 lần kích thước mô hình. Huấn luyện LoRA có trạng thái optimizer nhỏ hơn vì ít tham số được huấn luyện hơn, nhưng vẫn đáng kể.

Bộ nhớ gradient

Gradient cho mỗi tham số có thể huấn luyện cần lưu trữ trong backward pass. Điều này thêm một bản sao đầy đủ khác của kích thước tham số có thể huấn luyện vào yêu cầu bộ nhớ.

Kết hợp với mô hình, activation và trạng thái optimizer, tổng yêu cầu bộ nhớ dễ dàng đạt 40-50GB cho các cấu hình huấn luyện Flux mặc định.

Cách cấu hình huấn luyện cho 24GB VRAM?

Các cài đặt này cho phép huấn luyện Flux LoRA ổn định trên RTX 4090.

Bật gradient checkpointing

Gradient checkpointing là tối ưu hóa bộ nhớ có tác động lớn nhất. Nó giảm bộ nhớ activation 60-70% với chi phí 20-30% thời gian tính toán thêm.

Thay vì lưu trữ tất cả activation trong forward pass, checkpointing loại bỏ hầu hết và tính toán lại chúng trong backward pass. Sử dụng bộ nhớ trở nên gần như không đổi bất kể độ sâu mô hình.

Trong Kohya SS, bật gradient checkpointing trong cấu hình huấn luyện. Tùy chọn thường là checkbox đơn giản hoặc tham số boolean.

Trong các script huấn luyện tùy chỉnh sử dụng diffusers, gọi model.enable_gradient_checkpointing() trước khi huấn luyện bắt đầu.

Huấn luyện mất nhiều thời gian hơn với checkpointing được bật nhưng thực sự hoàn thành thay vì crash. Đánh đổi thời gian là xứng đáng.

Đặt kích thước batch thành 1

Kích thước batch trực tiếp nhân bộ nhớ activation. Kích thước batch 4 sử dụng khoảng 4 lần bộ nhớ activation của kích thước batch 1.

Đặt kích thước batch của bạn thành 1. Sử dụng tích lũy gradient để mô phỏng kích thước batch hiệu quả lớn hơn mà không có chi phí bộ nhớ.

Ví dụ, kích thước batch 1 với 4 bước tích lũy gradient cho kích thước batch hiệu quả là 4 trong khi chỉ giữ activation của 1 mẫu trong bộ nhớ.

Tích lũy gradient tích lũy gradient qua nhiều forward pass trước khi cập nhật trọng số. Sử dụng bộ nhớ giữ nguyên ở kích thước batch 1 trong khi động lực huấn luyện xấp xỉ batch lớn hơn.

Giảm độ phân giải huấn luyện

Độ phân giải có tác động bình phương lên bộ nhớ. Nhân đôi độ phân giải gấp bốn lần bộ nhớ activation.

Huấn luyện ở 512x512 thay vì 1024x1024. Điều này giảm bộ nhớ activation khoảng 75%.

Bạn có thể lo lắng rằng huấn luyện 512x512 tạo ra kết quả kém hơn độ phân giải gốc. Trong thực tế, LoRA được huấn luyện ở độ phân giải thấp hơn chuyển tốt sang inference độ phân giải cao hơn. Các yếu tố phong cách và khái niệm bạn đang huấn luyện vẫn xuất hiện ở tạo ảnh 1024x1024.

Nếu bạn cần huấn luyện độ phân giải cao hơn cho các trường hợp sử dụng cụ thể, 768x768 đôi khi có thể với các tối ưu hóa mạnh mẽ khác. Kiểm tra cẩn thận và giám sát bộ nhớ.

Sử dụng huấn luyện mixed precision

Mixed precision sử dụng FP16 hoặc BF16 cho hầu hết các hoạt động trong khi giữ các giá trị quan trọng ở FP32.

BF16 được khuyến nghị cho GPU Ampere và mới hơn. Nó xử lý gradient tốt hơn FP16 do phạm vi động lớn hơn.

Bật mixed precision trong cấu hình huấn luyện của bạn. Trong Kohya SS, chọn BF16 từ dropdown precision. Trong script tùy chỉnh, sử dụng context manager autocast của PyTorch.

Mixed precision giảm khoảng một nửa bộ nhớ cho trọng số mô hình và activation. Kết hợp với các tối ưu hóa khác, nó thiết yếu cho huấn luyện 24GB.

Bật attention tiết kiệm bộ nhớ

Các triển khai attention tiêu chuẩn phân bổ tensor trung gian lớn. Các biến thể tiết kiệm bộ nhớ như xFormers hoặc Flash Attention xử lý attention theo chunk.

xFormers hoạt động tốt cho huấn luyện và được hỗ trợ rộng rãi. Cài đặt riêng và bật trong cấu hình huấn luyện của bạn.

Flash Attention cung cấp hiệu suất tốt hơn trên các GPU gần đây. Kiểm tra xem framework huấn luyện của bạn có hỗ trợ không.

Attention tiết kiệm bộ nhớ có thể giảm sử dụng bộ nhớ attention 80% hoặc hơn. Đối với các mô hình lớn như Flux, điều này chuyển thành tiết kiệm vài GB.

Sử dụng optimizer 8-bit

Optimizer Adam tiêu chuẩn lưu trữ hai giá trị FP32 cho mỗi tham số. Adam 8-bit lượng tử hóa chúng thành INT8, cắt giảm bộ nhớ optimizer một nửa.

Cài đặt thư viện bitsandbytes và cấu hình Adam 8-bit trong thiết lập huấn luyện của bạn. Kohya SS hỗ trợ trực tiếp điều này. Script tùy chỉnh cần import optimizer 8-bit từ bitsandbytes.

Tác động chất lượng là tối thiểu. Adam 8-bit hội tụ tương tự như độ chính xác đầy đủ trong hầu hết các trường hợp.

Cài đặt nào hoạt động tốt nhất cho Kohya SS?

Kohya SS là công cụ huấn luyện Flux LoRA phổ biến nhất. Đây là các cấu hình cụ thể hoạt động.

Cấu hình được khuyến nghị

Sử dụng các cài đặt này làm điểm bắt đầu cho huấn luyện Flux LoRA RTX 4090.

Độ phân giải: 512,512 Kích thước batch: 1 Bước tích lũy gradient: 4 Mixed precision: bf16 Gradient checkpointing: Đã bật xFormers: Đã bật Optimizer: AdamW8bit Network rank: 16-32 Network alpha: Bằng rank hoặc một nửa Learning rate: 1e-4 Bước huấn luyện: 1000-2000 cho nhân vật, 2000-4000 cho phong cách

Cấu hình này sử dụng khoảng 20GB VRAM, để lại dư địa cho ổn định.

Điều chỉnh cho các loại huấn luyện khác nhau

LoRA nhân vật có thể sử dụng rank thấp hơn khoảng 16 và ít bước hơn khoảng 1000. Các đặc điểm nhân vật tương đối đơn giản để nắm bắt.

LoRA phong cách được hưởng lợi từ rank cao hơn khoảng 32-64 và nhiều bước huấn luyện hơn khoảng 3000-4000. Phong cách nghệ thuật có nhiều biến thể hơn để học.

LoRA khái niệm cho các đối tượng hoặc tư thế cụ thể khác nhau rộng rãi. Bắt đầu với cài đặt nhân vật và điều chỉnh dựa trên kết quả.

Rank cao hơn cần nhiều VRAM hơn. Nếu bạn đẩy rank lên 64 hoặc cao hơn, giám sát sử dụng bộ nhớ trong quá trình huấn luyện. Bạn có thể cần giảm độ phân giải xuống 448x448.

Cấu hình caption

Caption tốt cải thiện chất lượng huấn luyện đáng kể. Flux phản hồi tốt với mô tả ngôn ngữ tự nhiên.

Sử dụng BLIP hoặc các công cụ tương tự để tạo caption ban đầu, sau đó tinh chỉnh thủ công. Xóa các mô tả không chính xác và thêm từ trigger của bạn.

Caption dài hơn thường hoạt động tốt hơn cho Flux so với SD 1.5. Bao gồm các chi tiết liên quan về nội dung ảnh ngoài chủ đề.

Tránh caption lặp lại trong dataset của bạn. Thay đổi ngôn ngữ trong khi giữ từ trigger nhất quán.

Tạo ảnh mẫu

Bật tạo mẫu trong quá trình huấn luyện để giám sát tiến độ. Đặt tần suất mẫu mỗi 100-200 bước.

Tạo mẫu thêm overhead bộ nhớ. Nếu bạn gặp OOM trong sampling, tăng tích lũy gradient hoặc giảm tần suất mẫu.

Mẫu cho bạn biết khi huấn luyện diễn ra tốt và khi bạn đã overfit. Dừng huấn luyện khi mẫu trông tốt nhưng chưa bắt đầu xuống cấp.

Cách khắc phục lỗi OOM dai dẳng?

Nếu crash tiếp tục sau khi tối ưu hóa, điều tra các yếu tố bổ sung này.

Phân mảnh VRAM

Bộ cấp phát bộ nhớ của PyTorch có thể phân mảnh VRAM theo thời gian, gây ra lỗi ngay cả khi tổng bộ nhớ trống có vẻ đủ.

Chạy huấn luyện với biến môi trường PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128. Điều này thay đổi chiến lược cấp phát để giảm phân mảnh.

Bắt đầu mới mà không có các hoạt động GPU trước đó giúp ích. Khởi động lại hệ thống của bạn hoặc ít nhất kết thúc tất cả các tiến trình Python trước khi huấn luyện.

Giám sát sử dụng bộ nhớ

Quan sát sử dụng VRAM trong quá trình huấn luyện để xác định chính xác khi OOM xảy ra.

Chạy nvidia-smi -l 1 trong terminal riêng để xem sử dụng bộ nhớ cập nhật mỗi giây.

Ghi nhận sử dụng đỉnh ngay trước khi crash. Điều này cho bạn biết bạn cần giảm bao nhiêu.

Nếu crash xảy ra ngay lập tức, cấu hình tải mô hình của bạn sai. Nếu crash xảy ra sau vài bước, tích lũy activation hoặc phân mảnh là vấn đề.

Vấn đề dataset

Ảnh độ phân giải cực cao trong dataset của bạn có thể gây OOM khi tải cho huấn luyện.

Tiền xử lý dataset của bạn để đảm bảo tất cả ảnh ở hoặc dưới độ phân giải huấn luyện. Không có lợi ích khi tải ảnh 4K khi huấn luyện ở 512x512.

Xác minh rằng tỷ lệ khung hình là hợp lý. Ảnh rất rộng hoặc cao có thể yêu cầu nhiều bộ nhớ hơn trong quá trình xử lý ngay cả ở cùng tổng số pixel.

Các tiến trình khác sử dụng VRAM

Kiểm tra các ứng dụng khác đang tiêu thụ bộ nhớ GPU trước khi huấn luyện.

Đóng trình duyệt web, Discord và các ứng dụng được tăng tốc GPU khác. Ngay cả vài trăm MB được tiêu thụ ở nơi khác có thể đẩy bạn qua giới hạn.

Nhiều tiến trình Python có thể giữ VRAM từ các lần chạy thất bại trước đó. Khởi động lại trình thông dịch Python hoặc toàn bộ hệ thống của bạn để có slate sạch.

Bug framework huấn luyện

Đôi khi, bug trong Kohya SS hoặc các công cụ huấn luyện khác gây rò rỉ bộ nhớ.

Cập nhật lên phiên bản mới nhất của công cụ huấn luyện của bạn. Các bản sửa lỗi liên quan đến bộ nhớ phổ biến trong các bản cập nhật.

Kiểm tra GitHub issues của công cụ để tìm các báo cáo phù hợp với triệu chứng của bạn. Những người khác có thể đã tìm thấy giải pháp tạm thời hoặc bản sửa lỗi.

Đối với người dùng muốn huấn luyện LoRA mà không quản lý các hạn chế kỹ thuật này, Apatero.com cung cấp huấn luyện dựa trên cloud với phần cứng cấp chuyên nghiệp. Bạn xác định công việc huấn luyện và dataset của mình trong khi nền tảng xử lý quản lý bộ nhớ và tối ưu hóa tự động.

Các phương pháp huấn luyện thay thế là gì?

Nếu huấn luyện 4090 vẫn có vấn đề, hãy xem xét các phương án thay thế này.

Huấn luyện cloud

Các instance cloud với 48GB+ VRAM loại bỏ hoàn toàn các ràng buộc bộ nhớ.

RunPod, Vast.ai và Lambda Labs cung cấp instance GPU theo giờ. Instance A100 80GB huấn luyện Flux LoRA thoải mái ở độ phân giải đầy đủ.

Chi phí cho một lần chạy huấn luyện LoRA điển hình là $5-15 tùy thuộc vào nhà cung cấp và loại GPU. Điều này hợp lý cho nhu cầu huấn luyện thỉnh thoảng.

Tải lên dataset của bạn, chạy huấn luyện, tải xuống LoRA. Quy trình mất tổng cộng 1-2 giờ cho hầu hết các dự án.

Huấn luyện độ chính xác thấp hơn

Huấn luyện FP8 thử nghiệm giảm bộ nhớ hơn nữa so với BF16. Một số công cụ cộng đồng hỗ trợ điều này.

Huấn luyện FP8 ít được chứng minh về chất lượng hơn BF16. Kiểm tra kết quả cẩn thận trước khi áp dụng cho các dự án quan trọng.

Tiết kiệm bộ nhớ thêm 30-40% có thể làm cho các cấu hình trước đây không thể hoạt động.

Các biến thể LoRA nhỏ hơn

LoKr, LoHa và các thích ứng low-rank tương tự sử dụng ít bộ nhớ hơn LoRA tiêu chuẩn.

Các biến thể này hoạt động tốt cho một số mục tiêu huấn luyện nhưng có thể kém hiệu suất cho các phong cách hoặc khái niệm phức tạp.

Thử nghiệm với các phương án thay thế nếu cấu hình LoRA tiêu chuẩn vẫn có vấn đề mặc dù đã tối ưu hóa.

Câu hỏi thường gặp

Tại sao huấn luyện crash sau đúng 1 bước mỗi lần?

Điều này cho thấy một ngưỡng bộ nhớ cụ thể đang bị vượt qua. Forward pass vừa nhưng thêm gradient trong backward pass vượt quá VRAM. Bật gradient checkpointing và giảm độ phân giải cùng nhau thay vì từng bước.

Tôi có thể huấn luyện ở 1024x1024 trên RTX 4090 không?

Về lý thuyết có thể với tối ưu hóa cực độ bao gồm rank tối thiểu, checkpointing nặng và 8-bit cho mọi thứ. Thực tế, kết quả ở 512x512 đủ tốt để đấu tranh với bộ nhớ không đáng. Huấn luyện ở 512x512 và tạo ở 1024x1024.

Kích thước batch 1 có tạo ra LoRA kém hơn batch lớn hơn không?

Không đáng kể. Tích lũy gradient cung cấp động lực huấn luyện tương đương. Một số cho rằng batch rất nhỏ có nhiễu hơn một chút nhưng gradient tích lũy làm mượt điều này. Sự khác biệt chất lượng là tối thiểu so với các yếu tố khác.

Làm sao tôi biết LoRA của tôi đang huấn luyện đúng cách?

Ảnh mẫu trong quá trình huấn luyện cho thấy tiến độ. Bạn nên thấy từ trigger của mình ảnh hưởng đến các tạo ảnh sau 200-400 bước. Chuyển phong cách đầy đủ thường xuất hiện vào 800-1000 bước. Nếu mẫu không thay đổi hoặc ngay lập tức xuống cấp, điều chỉnh learning rate.

Tại sao sử dụng VRAM tăng dần trong quá trình huấn luyện?

Phân mảnh bộ nhớ hoặc rò rỉ gây tăng dần. Bộ cấp phát tạo các mảnh nhỏ không thể tái sử dụng. Đặt biến môi trường max_split_size_mb và đảm bảo không có rò rỉ bộ nhớ trong code tùy chỉnh.

Tôi nên sử dụng xFormers hay attention native PyTorch?

xFormers cung cấp hiệu quả bộ nhớ tốt hơn cho hầu hết các kịch bản huấn luyện. Attention native đôi khi hoạt động tốt hơn cho các kiến trúc cụ thể. Bắt đầu với xFormers và chỉ chuyển nếu bạn gặp vấn đề.

Tôi nên sử dụng network rank nào cho Flux LoRA?

Bắt đầu với 16 cho nhân vật và khái niệm đơn giản, 32 cho phong cách và chủ đề phức tạp. Rank cao hơn nắm bắt nhiều chi tiết hơn nhưng cần nhiều bộ nhớ và nhiều dữ liệu huấn luyện hơn. Kiểm tra rank thấp hơn trước vì chúng thường hoạt động tốt.

Tôi cần bao nhiêu ảnh huấn luyện?

Đối với nhân vật, 10-20 ảnh tốt hoạt động tốt. Đối với phong cách, 50-200 ảnh cung cấp phạm vi tốt hơn. Chất lượng quan trọng hơn số lượng. Ảnh có caption tốt và đa dạng đánh bại hàng trăm shot tương tự.

Tôi có thể tiếp tục huấn luyện sau crash OOM không?

Nếu bạn đã bật checkpointing trong cấu hình huấn luyện của mình, có. Kohya SS lưu tiến độ định kỳ. Tiếp tục từ checkpoint cuối cùng sau khi sửa cài đặt bộ nhớ.

Có cách nào để dự đoán liệu cấu hình của tôi sẽ OOM trước khi bắt đầu không?

Ước tính sơ bộ dựa trên kích thước mô hình, kích thước batch, độ phân giải và tối ưu hóa. Các công cụ như accelerate estimate-memory cung cấp ước tính. Nhưng bộ nhớ thực tế khác nhau tùy theo chi tiết triển khai, vì vậy luôn bắt đầu bảo thủ.

Kết luận và quy trình được khuyến nghị

Huấn luyện Flux LoRA trên RTX 4090 yêu cầu quản lý bộ nhớ cẩn thận nhưng tạo ra kết quả tuyệt vời khi được cấu hình đúng cách. Các cài đặt chính là gradient checkpointing, kích thước batch 1 với tích lũy, độ phân giải 512x512 và huấn luyện mixed precision.

Bắt đầu với cấu hình bảo thủ được cung cấp trong hướng dẫn này. Chạy thử nghiệm ngắn 100-200 bước để xác minh ổn định trước khi cam kết với các lần chạy huấn luyện đầy đủ.

Giám sát mẫu của bạn trong quá trình huấn luyện. Dừng khi LoRA nắm bắt khái niệm mục tiêu của bạn trước khi chất lượng bắt đầu xuống cấp từ overtraining.

Nếu bạn liên tục đấu tranh với giới hạn bộ nhớ hoặc muốn cài đặt chất lượng cao hơn những gì 24GB cho phép, huấn luyện cloud cung cấp một phương án thay thế thực tế. Các dịch vụ như Apatero.com làm cho điều này có thể truy cập mà không cần tự quản lý cơ sở hạ tầng cloud.

RTX 4090 của bạn là phần cứng có khả năng cho huấn luyện LoRA cục bộ. Với cấu hình thích hợp, bạn có thể huấn luyện Flux LoRA tùy chỉnh cho các nhân vật, phong cách và khái niệm cụ thể của mình trong khi giữ mọi thứ trên máy của riêng bạn.