Mã lỗi: E0100

Danh mục: Thời gian chạy: Lỗi phân bổ vùng đệm

Lỗi này cho biết trình phân bổ bộ nhớ của thời gian chạy XLA:TPU không tìm thấy một khối bộ nhớ phù hợp trên HBM của trình tăng tốc cho việc phân bổ được yêu cầu.

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

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

XLA backends: TPU

Tổng quan

Lỗi này sẽ xảy ra khi:

  • Lỗi phân bổ vùng đệm do người dùng thực hiện thông qua jax.device_put

hoặc

  • Không phân bổ được đầu ra của chương trình do người dùng lên lịch.

Những lỗi này thường xảy ra do một số lý do:

  • Hết bộ nhớ (OOM): Người dùng đang cố gắng phân bổ một khối bộ nhớ lớn hơn tổng dung lượng bộ nhớ trống có trên HBM của TPU.
  • Phân mảnh bộ nhớ: Quá trình phân bổ không thành công vì không có một khối trống liền kề nào trong không gian bộ nhớ đủ lớn để đáp ứng kích thước được yêu cầu. Tổng dung lượng bộ nhớ trống đủ để phân bổ, nhưng được phân tán trên không gian bộ nhớ thành các khối nhỏ, không liền kề.

Thời gian chạy TPU có một số cơ chế để thử lại các lỗi phân bổ, bao gồm:

  • Nếu có các hoạt động huỷ phân bổ được xếp hàng đợi, thì thời gian chạy sẽ thử lại các hoạt động phân bổ không thành công.
  • Trên các OOM do phân mảnh, thời gian chạy có thể tự động kích hoạt quá trình chống phân mảnh và thử lại.
  • Thời gian chạy TPU ưu tiên việc phân bổ vùng đệm hơn là giữ cho các chương trình được tải. Nếu quá trình phân bổ vùng đệm không thành công do không đủ HBM, hệ thống sẽ loại bỏ các chương trình TPU đã tải cho đến khi có đủ bộ nhớ cho vùng đệm.

Vì vậy, lỗi gặp phải sau các biện pháp giảm thiểu nêu trên thường yêu cầu người dùng hành động.

Gỡ lỗi

  • Giảm mức sử dụng bộ nhớ của mô hình
    • Giảm kích thước lô: Việc giảm kích thước lô sẽ trực tiếp làm giảm mức sử dụng bộ nhớ.
    • Phân đoạn tham số: Đối với các mô hình rất lớn, hãy sử dụng các kỹ thuật như tính song song của mô hình hoặc phân đoạn để phân phối các tham số trên HBM của 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 hoạt động trên các chuỗi (chẳng hạn như mô hình ngôn ngữ), việc giảm độ dài chuỗi đầu vào có thể giảm đáng kể mức sử dụng bộ nhớ.
    • Đóng góp vùng đệm: Sử dụng các tính năng của khung (chẳng hạn như: jax.jit(..., donate_argnums=...)) để báo hiệu cho XLA rằng một số vùng đệm đầu vào nhất định có thể bị ghi đè và dùng lại cho đầu ra. Đọc bài viết Quyên góp qua Buffer để biết thêm thông tin.
    • Tối ưu hoá chiến lược điểm kiểm tra: Thay vì lưu toàn bộ trạng thái mô hình cùng một lúc, hãy cân nhắc chỉ lưu trọng số mô hình hoặc sử dụng chiến lược kiểm tra phân đoạn.
  • Bố cục bộ nhớ và khoảng đệm địa chỉ
    • Bộ nhớ TPU được phân bổ theo khối và khoảng đệm có thể làm tăng kích thước thực tế của các tensor.
  • Đảm bảo không rò rỉ bộ nhớ
    • Đảm bảo các tham chiếu đến đối tượng jax.Array không được giữ lâu hơn dự kiến. Việc giữ 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.

Xem thêm Mã lỗi: E1000 để biết các chiến lược khác mà bạn có thể sử dụng để giảm lượng bộ nhớ mà mỗi chương trình sử dụng.

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.