Mã lỗi: 1001

Danh mục: Thời gian biên dịch: Lỗi hết bộ nhớ (OOM) Vmem có phạm vi

Lỗi này cho biết chương trình cần nhiều Bộ nhớ vectơ có phạm vi (Vmem) hơn mức đã phân bổ.

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

RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...

XLA Backends: TPU

Tổng quan

TPU có Bộ nhớ vectơ (VMEM) là bộ nhớ tạm cục bộ chỉ được TensorCore (TC) sử dụng. Trình biên dịch quản lý Vmem cho nhiều loại hoạt động phân bổ:

  • Phân bổ theo phạm vi hướng dẫn: Bộ nhớ tạm thời trong Vmem trong khi thực thi một chỉ dẫn HLO. Điều này bao gồm vùng đệm khoảng toán hạng (ví dụ: để đệm đôi) và tràn thanh ghi.
  • Phân bổ theo phạm vi chương trình: Phân bổ nằm ngoài phạm vi của một chỉ dẫn HLO duy nhất. Đây thường là các kết quả tạm thời và trung gian của HLO, là đầu vào và/hoặc đầu ra của các chỉ dẫn HLO.

Lỗi hết bộ nhớ ảo (OOM) Vmem có phạm vi thời gian biên dịch xảy ra khi các hoạt động phân bổ có phạm vi hướng dẫn vượt quá hạn mức phân bổ cho hướng dẫn đó. Giới hạn này được kiểm soát

Những lỗi này thường do lỗi trình biên dịch nội bộ hoặc do một nhân tuỳ chỉnh vượt quá hạn mức phân bổ.

Gỡ lỗi

Phân tích kỹ thông báo lỗi để xác định xem lỗi bắt nguồn từ một nhân tuỳ chỉnh hay HLO tiêu chuẩn. Lỗi do một nhân tuỳ chỉnh phải có chữ ký sau:

Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
  • OOM Vmem có phạm vi tuỳ chỉnh của nhân: Nếu lỗi này chỉ đến một nhân tuỳ chỉnh → Chuyển đến phần Điều chỉnh lại nhân.
  • Vấn đề về Vmem không phải là hạt nhân: Nếu Vmem OOM xảy ra do một thao tác không phải là hạt nhân tuỳ chỉnh, thì đó có thể là một lỗi trình biên dịch nội bộ. Vui lòng báo cáo lỗi trên XLA bằng một kết xuất HLO.

Điều chỉnh lại Kernel

Nếu lỗi bắt nguồn từ một nhân tuỳ chỉnh, hãy sử dụng các kỹ thuật sau để giảm yêu cầu về bộ nhớ của nhân:

  • Điều chỉnh kích thước khối: Giảm kích thước khối (kích thước ô) trong cấu hình hạt nhân để giảm mức sử dụng Vmem có phạm vi.
  • Đặt giới hạn Vmem theo phạm vi trên mỗi nhân: Yêu cầu rõ ràng lượng bộ nhớ cần thiết cho nhân cụ thể đó bằng cách sử dụng tham số vmem_limit_bytes
  • Sửa đổi tính năng tô màu bộ nhớ: Tô màu/hạn chế rõ ràng các đầu vào/đầu ra của nhân đối với VMEM bằng cách sử dụng pallas.tpu.with_memory_space_constraint. Nhưng hãy cẩn thận để không tô màu quá nhiều đầu vào/đầu ra cho Vmem, vì điều đó có thể gây ra lỗi OOM VMEM tổng thể.
  • Nếu khó điều chỉnh lại theo từng nhân cụ thể hoặc vấn đề ảnh hưởng đến nhiều nhân, bạn có thể điều chỉnh giới hạn Vmem chung bằng cách sử dụng cờ --xla_tpu_scoped_vmem_limit_kib.