Mã lỗi: 1000

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: 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.

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

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.
• TensorCore + SparseCore Temporaries: 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.).
• 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 IR HLO 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:

• TensorCore (TC) + SparseCore (SC) HBM Usage Exceeds Limit (Đã vượt quá giới hạn sử dụng HBM của TensorCore (TC) + SparseCore (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". → Chuyển đến Phần 1. Cân bằng việc sử dụng HBM TC và SC.
• Phân bổ lớn ngoài dự kiến: Nếu lỗi có nội dung "Ran out of memory in memory space HBM" (Hết bộ nhớ trong không gian bộ nhớ HBM), hãy kiểm tra nhật ký để biết danh sách các hoạt động 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 Mục 2. Phân bổ lớn ngoài dự kiến.
• Tổng số lượt phân bổ vượt quá giới hạn HBM: Nếu lỗi có nội dung "Ran out of memory in memory space HBM" (Hết bộ nhớ trong không gian bộ nhớ HBM) nhưng không có tensor lớn bất thường nào trong nhật ký → Chuyển đến Mục 3. Tổng số lượt phân bổ vượt quá giới hạn HBM.

Phần 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" 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 khả năng xử lý song song.
  • Kiểm tra chi phí bổ sung: SparseCore căn chỉnh các bảng nhúng thành 32B (8 số thực). Bảng có chiều rộng tính năng nhỏ (ví dụ: < 8 floats) incur significant padding overhead, wasting HBM.
  • Giảm mức sử dụng vùng nhớ heap: 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 vùng nhớ 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 rằng các bảng nhúng được phân đoạn theo cách thức phù 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 đến Mục 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.

Phần 2. Phân bổ lớn ngoài dự kiến

Nếu có một hoặc nhiều lượt phân bổ lớn bất thường 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 hoạt động phân bổ lớn để biết gợi ý về mã nguồn JAX của chúng.

• Xoá cấu phần phần mềm 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ả:
  • 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 là Replication – 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.

Phần 3. Tổng số lượt phân bổ vượt quá giới hạn HBM

Nếu chương trình hết dung lượng do tổng số các lần phân bổ vượt quá giới hạn HBM, thì việc trực quan hoá hồ sơ bộ nhớ thường sẽ giúp 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á dấu vết bộ nhớ.

A. 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 đệ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à gây ra mức lãng phí bộ nhớ là 0%.
• 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.

B. Điều chỉnh cấu hình

Bạn thường có thể giải quyết lỗi OOM bằng cách điều chỉnh cấu hình theo những cách 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ớ.
• Donate Input Buffers (Quyên góp 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à các yêu cầu về chất lượng cho phép.

C. 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 so 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á các phương pháp song song hoá dữ liệu, tensor hoặc quy trình 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 chuyển tải máy chủ JAX: Chuyển tải các tensor lớn sang bộ nhớ CPU của máy chủ. Ví dụ: chuyển tải hoạt động và chuyển tải trạng thái trình tối ưu hoá.

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

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 các phương án này khi không còn cách nào khác vì chúng có thể ảnh hưởng tiêu cực đến hiệu suất.

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ể 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:

Ngoài ra, hãy sử 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.

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 Trình xem bộ nhớ XProf để trực quan hoá 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.