Mã lỗi: 0101

Danh mục: Thời gian chạy: Lỗi phân bổ chương trình

Lỗi này cho biết thời gian chạy XLA trên thiết bị TPU không tải được tệp thực thi chương trình XLA đã biên dịch vào HBM của TPU.

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

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

Các chương trình phụ trợ XLA: TPU

Tổng quan

Lỗi này thường xảy ra do một trong những nguyên nhân sau:

• Kích thước chương trình vượt quá HBM có sẵn: Chương trình XLA đã biên dịch, bao gồm các chỉ dẫn, dữ liệu tĩnh và mọi hằng số được nhúng, lớn hơn tổng lượng HBM miễn phí hiện có trên(các) lõi TPU cụ thể nơi chương trình đang được tải.
• Phân mảnh HBM: Mặc dù tổng HBM trống trên thiết bị có thể đủ về tổng thể, nhưng không có sẵn trong một khối liền kề duy nhất đủ lớn để chứa toàn bộ chương trình.

Điều quan trọng là bạn phải hiểu cách thời gian chạy TPU ưu tiên bộ nhớ. Việc phân bổ vùng đệm được ưu tiên hơn các chương trình đã tải. Nếu không phân bổ được vùng đệm, thời gian chạy sẽ loại bỏ các chương trình đã tải khỏi HBM để giải phóng dung lượng. Điều này có thể dẫn đến tình huống là một chương trình đã tải thành công trước đó hiện không thành công với lỗi OOM, vì HBM hiện đang chứa nhiều vùng đệm dữ liệu hơn.

Gỡ lỗi

• Giảm mức sử dụng bộ nhớ của vùng đệm: Việc giải phóng bộ nhớ mà các vùng đệm dữ liệu sử dụng sẽ giúp chương trình có thêm dung lượng:
  • Giảm kích thước lô: Đây là một trong những cách hiệu quả nhất để giảm lượng bộ nhớ được dùng cho các lượt kích hoạt.
  • Phân đoạn tham số: Đối với các mô hình rất lớn, hãy sử dụng tính song song của mô hình hoặc các kỹ thuật phân đoạn (chẳng hạn như FSDP hoặc Megascale) để phân phối các tham số và hoạt động tính toán của mô hình trên nhiều lõi hoặc máy chủ TPU.
  • Rút ngắn độ dài chuỗi/ngữ cảnh: Đối với các mô hình xử lý dữ liệu tuần tự (ví dụ: NLP), việc giảm độ dài chuỗi có thể giảm đáng kể mức sử dụng bộ nhớ.
  • Đóng góp bộ nhớ đệm: Sử dụng các tính năng của khung (ví dụ: jax.jit(..., donate_argnums=...)) để cho phép XLA sử dụng lại bộ nhớ của vùng đệm đầu vào để lưu trữ đầu ra, giảm mức sử dụng bộ nhớ tối đa.
• Giảm yêu cầu về bộ nhớ của chương trình đối với các tệp tạm thời:
  • Giảm mức sử dụng bộ nhớ của chương trình cho các tệp tạm thời bằng cách sử dụng cờ tpu_shared_memory_percent. Xin lưu ý rằng điều này có thể ảnh hưởng tiêu cực đến hiệu suất.
• Tối ưu hoá chiến lược thực thi/Giảm tải phân phát:
  • Quản lý việc tải chương trình: Nếu bạn đang biên dịch JIT nhiều hàm, hãy lưu ý rằng mỗi hàm có thể dẫn đến việc tải một chương trình. Hãy cố gắng cấu trúc khối lượng công việc để giảm thiểu số lượng chương trình được tải đồng thời.
• Đảm bảo không rò rỉ bộ nhớ:
  • Đảm bảo các tham chiếu đến các đối tượng jax.Array không được giữ lâu hơn dự kiến. Việc giữ lại các đối tượng jax.Array có thể ngăn việc huỷ cấp phát tự động ngay cả sau khi quá trình biên dịch chương trình hoàn tất.

Công cụ

• Bật cờ tpu_log_allocations_on_oom để trình phân bổ sẽ kết xuất một báo cáo chi tiết về tất cả các hoạt động phân bổ hiện tại khi xảy ra lỗi hết bộ nhớ. Điều này có thể rất hữu ích cho việc gỡ lỗi.
• Lập hồ sơ cho chương trình của bạn: Sử dụng trình phân tích bộ nhớ JAX hoặc trình phân tích tài nguyên TensorFlow để xem chi tiết mức sử dụng bộ nhớ của chương trình theo thời gian. Điều này có thể giúp xác định các mức tăng đột biến ngoài dự kiến về mức tiêu thụ bộ nhớ.