Mã lỗi: E1001

Danh mục: Thời gian biên dịch: Vmem có phạm vi OOM

Lỗi này cho biết chương trình yêu cầu 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(...) ...

Phần phụ trợ XLA: TPU

Tổng quan

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

  • Phân bổ có phạm vi hướng dẫn: Bộ nhớ tạm thời trong Vmem trong khi thực thi một hướng dẫn HLO. Điều này bao gồm bộ đệm khoảng thời gian toán hạng (ví dụ: để đệm gấp đôi) và tràn thanh ghi.
  • Phân bổ có phạm vi chương trình: Phân bổ tồn tại ngoài phạm vi của một hướng dẫn HLO. Đâ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 hướng dẫn HLO.

Lỗi Vmem có phạm vi OOM trong thời gian biên dịch xảy ra khi các phân bổ có phạm vi hướng dẫn vượt quá giới hạn phân bổ cho hướng dẫn đó. Giới hạn này do

và

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 kernel tuỳ chỉnh vượt quá giới hạn phân bổ.

Gỡ lỗi

Phân tích kỹ thông báo lỗi để xác định xem lỗi có bắt nguồn từ kernel tuỳ chỉnh hay HLO tiêu chuẩn hay không. Lỗi do kernel 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" ...
  • Vmem có phạm vi OOM của kernel tuỳ chỉnh: Nếu lỗi trỏ đến một kernel tuỳ chỉnh → Chuyển đến Điều chỉnh lại kernel.
  • Vấn đề về Vmem không phải kernel: Nếu Vmem OOM xảy ra do một thao tác không phải kernel tuỳ chỉnh, thì có khả năng là lỗi trình biên dịch nội bộ. Vui lòng gửi báo cáo lỗi trên XLA kèm theo kết xuất HLO.

Điều chỉnh lại kernel

Nếu lỗi bắt nguồn từ một kernel 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 kernel:

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