Mã lỗi: E1000

Danh mục: Thời gian biên dịch: HBM OOM

Lỗi này cho biết chương trình yêu cầu nhiều Bộ nhớ băng thông cao (HBM) hơn mức có sẵn trên thiết bị TPU.

Thông báo lỗi mẫu:

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

XLA backends: TPU

Tổng quan

XLA thực hiện các bước kiểm tra để đảm bảo tổng kích thước của tất cả các hoạt động phân bổ tĩnh cần thiết phù hợp với HBM của thiết bị.

Trình biên dịch quản lý dung lượng HBM cố định của TPU cho một số loại phân bổ:

• Đầu vào và đầu ra của chương trình: Các lô huấn luyện, trạng thái của trình tối ưu hoá, v.v.
• TPU tạm thời: Bộ nhớ động cần thiết cho các phép tính trung gian (ví dụ: lượt kích hoạt, độ dốc, v.v.).
• Tệp nhị phân đã biên dịch: Mã máy cho cả TensorCore (TC) và SparseCore (SC).
• Chi phí hệ thống: Không gian dành riêng cho Thời gian chạy XLA (ví dụ: bộ đệm trong nguồn cấp dữ liệu trên các thế hệ TPU cũ).
• Hằng số: Các giá trị hằng số được nhúng trong HLO IR sẽ được phân bổ trên HBM.
• Cấu trúc bên trong của trình biên dịch: Phân bổ ở cấp chương trình và theo HLO (ví dụ: thông tin định tuyến cho các nút trong lưới).

Lỗi này xảy ra khi trình biên dịch XLA không thể điều chỉnh tất cả các hoạt động phân bổ ở trên vào HBM của thiết bị.

Gỡ lỗi

Phân tích cẩn thận thông báo lỗi và nhật ký để xác định danh mục HBM OOM nào dưới đây mô tả đúng nhất lỗi của bạn:

• "TC Hbm usage: X, SC Hbm usage Y": Nếu lỗi phân tích rõ ràng mức sử dụng HBM, thì mức sử dụng TensorCore (TC) + SparseCore (SC) tổng hợp sẽ vượt quá giới hạn HBM. → Chuyển đến Tình huống 1. Cân bằng việc sử dụng HBM TC và SC.
• "Ran out of memory in memory space HBM" (Hết bộ nhớ trong không gian bộ nhớ HBM): Kiểm tra nhật ký để biết danh sách các mục phân bổ lớn nhất trên HBM.
  • Trong trường hợp có một hoặc nhiều tensor lớn bất thường (ví dụ: > 50% giới hạn HBM) → Chuyển đến Trường hợp 2. Hết bộ nhớ do phân bổ quá nhiều một cách bất ngờ.
  • Nếu không có tensor lớn bất thường trong nhật ký → Chuyển đến Tình huống 3. Hết bộ nhớ do tổng hợp các hoạt động phân bổ.

Tình huống 1. Cân bằng việc sử dụng HBM TC và SC

Nếu lỗi phân tích rõ ràng mức sử dụng, ví dụ: "TC Hbm usage: X, SC Hbm usage Y", điều này có nghĩa là mức sử dụng TensorCore (TC) + SparseCore (SC) tổng hợp vượt quá giới hạn HBM. So sánh hai giá trị để xác định điểm tắc nghẽn:

• Mức sử dụng SparseCore cao
  • Tối ưu hoá mức sử dụng ngăn xếp HBM: Mức tiêu thụ bộ nhớ ngăn xếp HBM tăng theo feature_width, max_unique_nz_per_row và logical_replica_count. Bạn có thể giảm mức sử dụng ngăn xếp cao nhất bằng cách điều chỉnh cờ --xla_sc_num_serialized_tables_to_optimize_hbm để tuần tự hoá quá trình xử lý các bảng. Điều này làm giảm tính song song.
  • Kiểm tra chi phí chung của khoảng đệm: SparseCore căn chỉnh các bảng nhúng thành 32B (8 số thực). Các bảng có chiều rộng đặc trưng nhỏ (ví dụ: < 8 số thực) sẽ phải chịu chi phí bổ sung đáng kể, làm lãng phí HBM.
  • Giảm mức sử dụng heap: Các giá trị cao cho maximum_parallel_iterations sẽ làm tăng lượng dữ liệu đầu vào được tìm nạp trước vào heap HBM. Việc giảm giá trị này có thể giải phóng đáng kể bộ nhớ.
  • Xác minh việc phân đoạn: Đảm bảo các bảng nhúng được phân đoạn theo cách thích hợp trên tất cả các chip. Xem phần Cách giới hạn được chuyển đổi thành bảng.
  • Hãy xem SC: Performance and memory bottlenecks (SC: Hiệu suất và các điểm nghẽn về bộ nhớ) để biết thêm ý tưởng.
• Mức sử dụng TensorCore cao
  • Chuyển sang Tình huống 2.
• Cân bằng
  • Nếu không có tài khoản nào vượt quá hạn mức riêng lẻ nhưng tổng số lại quá cao, thì bạn đã đạt đến dung lượng của chip. Bạn phải cố gắng giảm mức sử dụng của cả hai thành phần này. Làm theo các đề xuất trong cả 3 phần.

Trường hợp 2. Hết bộ nhớ do phân bổ quá nhiều một cách bất ngờ

Nếu bạn thấy thông báo lỗi "Ran out of memory in memory space HBM" (Hết bộ nhớ trong không gian bộ nhớ HBM) và một hoặc nhiều lượt phân bổ lớn ngoài dự kiến xuất hiện trong nhật ký (> 50% giới hạn HBM), thì đó hầu như không phải là vấn đề về dung lượng phần cứng. Đây thường là lỗi cấu hình. Kiểm tra nhãn XLA (nếu có) của các mục phân bổ lớn để biết gợi ý về mã nguồn JAX của các mục đó.

• Xoá các cấu phần gỡ lỗi
  • Việc sử dụng jax.debug.print() trong các lần chạy quy mô lớn có thể buộc trình biên dịch hiện thực hoá toàn bộ tensor trong HBM để chuyển tensor đó sang CPU, làm gián đoạn quá trình hợp nhất và tăng mức sử dụng bộ nhớ tối đa. Xoá mọi jax.debug.print() còn sót lại.
• Khắc phục các hình dạng lưới hoặc phân đoạn không hiệu quả
  • Các hình dạng lưới không chính xác hoặc thiếu chú thích phân đoạn có thể khiến trình biên dịch mặc định thành sao chép – buộc trình biên dịch cố gắng điều chỉnh các tensor thực sự lớn trên một chip duy nhất.
  • Kiểm tra hình dạng của các phân bổ lớn và xác minh rằng XLA đã chỉ định và truyền tải đúng cách việc phân đoạn.

Trường hợp 3. Hết bộ nhớ do phân bổ tổng hợp

Nếu bạn thấy thông báo lỗi "Ran out of memory in memory space HBM" (Hết bộ nhớ trong không gian bộ nhớ HBM) và không có tensor lớn bất thường nào xuất hiện trong nhật ký, thì chương trình sẽ hết dung lượng do tổng số các mục phân bổ vượt quá giới hạn HBM. Trong trường hợp này, việc hình dung hồ sơ bộ nhớ thường rất hữu ích để xác định các vùng đệm cụ thể góp phần vào mức sử dụng cao nhất. Hãy xem phần Gỡ lỗi lỗi OOM bằng XProf để biết hướng dẫn từng bước về cách xác định những yếu tố đóng góp bộ nhớ cao nhất.

Sau khi bạn xác định được một số thành phần đóng góp hàng đầu, hãy làm theo các bước sau để tối ưu hoá mức sử dụng bộ nhớ.

Tình huống 3.A Điều chỉnh cấu hình

Bạn thường có thể giải quyết lỗi OOM bằng những điều chỉnh cấu hình sau:

• Giảm kích thước lô: Bộ nhớ cần thiết cho các kích hoạt và độ dốc trung gian tỷ lệ thuận với kích thước lô. Việc giảm kích thước lô thường có thể giúp giảm mức sử dụng bộ nhớ.
• Đóng góp các vùng đệm đầu vào: Khi sử dụng jax.jit, hãy chỉ định donate_argnums cho các tham số mô hình của bạn. Điều này cho phép XLA ghi đè bộ nhớ đầu vào bằng đầu ra.
• Bật độ chính xác hỗn hợp (bfloat16): Sử dụng bfloat16 hoặc lượng tử hoá (int8, v.v.) cho các tensor lớn nhất trong chương trình nếu kiến trúc mô hình và yêu cầu về chất lượng cho phép. Xin lưu ý rằng thay đổi này có thể ảnh hưởng đến hành vi của mô hình và bạn nên cân nhắc kỹ lưỡng.

Tình huống 3.B Tối ưu hoá cấu trúc và phân đoạn

Nếu các thay đổi về cấu hình là không đủ, thì có thể cấu trúc liên kết mô hình quá lớn đối với chế độ thiết lập phần cứng hiện tại.

• Sử dụng các thế hệ TPU mới hơn: Các TPU mới hơn thường cung cấp nhiều HBM hơn trên mỗi chip; hãy chuyển sang các thế hệ TPU mới hơn nếu có.
• Chạy trên một cấu trúc liên kết chip lớn hơn: Nếu trọng số mô hình quá lớn đối với cấu trúc liên kết hiện có, bạn có thể thử phân mảnh các trọng số đó trên nhiều chip hơn.
• Triển khai các kỹ thuật phân đoạn nâng cao:
  • Khám phá thêm các phương pháp song song hoá dữ liệu, tensor hoặc pipeline nâng cao hơn.
  • Chỉ định gợi ý phân đoạn cho các giá trị và đầu ra trung gian.
• Sử dụng tính năng giảm tải máy chủ JAX: Các kỹ thuật giảm tải máy chủ cho phép người dùng giảm tải các Tensor lớn sang bộ nhớ CPU của máy chủ (ví dụ: giảm tải kích hoạt và giảm tải trạng thái optimizer).

Tình huống 3.C Kiểm tra khoảng đệm và việc căn chỉnh tensor

Các hình dạng tensor không hiệu quả là một nguyên nhân phổ biến, âm thầm gây ra lỗi hết bộ nhớ trên TPU. Để đạt được hiệu suất cao nhất trên TPU, XLA sẽ đệm các phương diện của Tensor – thường là bội số của 128 cho phương diện nhỏ nhất và 8 cho phương diện nhỏ thứ hai. Khoảng đệm này ảnh hưởng đến cả mảng đầu vào và các tensor trung gian (tạm thời HLO), có khả năng làm tăng đáng kể mức sử dụng bộ nhớ, đặc biệt là với kích thước nhỏ. Xem phần Bố cục mảng.

• Kiểm tra các hình dạng của vùng đệm lớn: (Trên TPU phiên bản 5 có bố cục mặc định)
  • Khi di chuột lên một vùng đệm trong Xprof Memory Viewer (Trình xem bộ nhớ Xprof), thẻ thông tin chi tiết về vùng đệm sẽ xuất hiện. Thẻ này chứa thông tin chi tiết về vùng đệm, bao gồm cả thông tin về khoảng đệm.
  • Ví dụ: Hình dạng (129, 1024) có thể được thêm khoảng đệm thành (256, 1024), dẫn đến lãng phí gần 50% bộ nhớ.
  • Chỉnh sửa: Hình dạng của (128, 1024) không yêu cầu khoảng đệm và không gây lãng phí bộ nhớ.
• Căn chỉnh các phương diện: Đảm bảo tất cả các phương diện tensor lớn (kích thước lô, phương diện nhúng, kích thước ẩn) đều là bội số của 128. Xin lưu ý rằng thay đổi này có thể ảnh hưởng đến hành vi của mô hình và bạn nên cân nhắc kỹ lưỡng.

Trường hợp 3.D Điều chỉnh bộ nhớ khoá ảnh hưởng đến cờ XLA

Bạn có thể điều chỉnh các cờ bộ nhớ chính để đánh đổi hiệu suất nhằm giảm mức sử dụng bộ nhớ. Tuy nhiên, bạn chỉ nên sử dụng chiến lược này khi không còn cách nào khác vì chiến lược này có thể ảnh hưởng tiêu cực đến hiệu suất.

Tình huống 3.E Tune XLA rematerialization pass/manual checkpointing

Nếu mô hình gần như vừa với bộ nhớ, bạn có thể dùng trình trang trí jax.checkpoint với jax.grad để kiểm soát theo cách thủ công những thành phần trung gian được lưu trên đường truyền chuyển tiếp so với được tính toán lại trên đường truyền ngược, trao đổi các chu kỳ tính toán cho HBM.

Ngoài ra, bạn có thể buộc truyền XLA::Rematerialization để ưu tiên việc tiết kiệm bộ nhớ, có thể phải trả giá bằng việc biên dịch chậm hơn:

Xin lưu ý rằng bạn chỉ nên sử dụng biện pháp thay đổi cờ XLA khi không còn cách nào khác, vì biện pháp này có thể ảnh hưởng tiêu cực đến hiệu suất.

Tình huống 3.F Sử dụng các công cụ lập hồ sơ nâng cao

Gỡ lỗi OOM bằng XProf cung cấp một hướng dẫn về cách sử dụng XProf Memory Viewer để hình dung chế độ xem của trình biên dịch về mức sử dụng HBM.

Công cụ này cho phép bạn xem mức phân bổ bộ nhớ cao nhất và thời gian tồn tại của vùng đệm. Đây là thông tin quan trọng để hiểu chính xác những gì tiêu thụ HBM tại thời điểm sử dụng cao nhất. Để biết thông tin thiết lập phân tích tài nguyên chung, hãy xem bài viết Bắt đầu sử dụng Xprof và Phân tích tài nguyên bằng TensorBoard.