Gỡ lỗi OOM bằng XProf

Lỗi Hết bộ nhớ (OOM) xảy ra khi hết dung lượng Bộ nhớ băng thông cao (HBM) của bộ tăng tốc (GPU hoặc TPU). Một số nguyên nhân phổ biến dẫn đến vấn đề OOM và kỹ thuật gỡ lỗi được trình bày chi tiết trong E1000 – Tài liệu OOM HBM thời gian biên dịch và Tài liệu JAX về việc phân bổ bộ nhớ GPU.

Trang này mô tả cách sử dụng công cụ Trình xem bộ nhớ của XProf để trực quan hoá mức sử dụng bộ nhớ của chương trình JAX, xác định các phiên bản sử dụng cao nhất và gỡ lỗi OOM. Quá trình này bao gồm các bước sau:

  1. Chạy chương trình bằng jax.profiler.trace để ghi lại hồ sơ.
  2. Khởi động XProf ở chế độ nền và sử dụng công cụ Trình xem bộ nhớ để xem thông tin chi tiết về mức sử dụng bộ nhớ.

Chương trình mẫu

Chương trình JAX sau đây dẫn đến lỗi OOM:

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

Trên một máy TPU, chương trình này sẽ gặp lỗi:

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(Trên máy có GPU, lỗi này sẽ có dạng: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

Chạy XProf

Cài đặt xprof (pip install xprof) và khởi động một phiên bản XProf, chỉ định thư mục nơi hồ sơ được lưu trữ:

xprof --logdir=/tmp/xprof/ --port=6006

Chuyển đến phiên bản (trên máy cục bộ, tại http://localhost:6006). Trong trình đơn thả xuống Tools (Công cụ), hãy chọn Memory Viewer (Trình xem bộ nhớ) và trong cửa sổ công cụ Memory Viewer, hãy chọn HBM trong trình đơn thả xuống Memory Types (Các loại bộ nhớ) (thường được chọn theo mặc định).

Trang XProf Memory Viewer cho chương trình ví dụ ở trên

XProf: Tài liệu về công cụ Trình xem bộ nhớ mô tả các thành phần của công cụ và thông tin được trình bày.

Tập trung vào phần HLO Ops at Peak Memory Allocation (Các thao tác HLO tại điểm phân bổ bộ nhớ cao nhất) cho thấy 3 biểu đồ vùng đệm tại điểm sử dụng bộ nhớ cao nhất. Vùng đệm bao gồm: * Đầ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 và SparseCore Temporaries: Bộ nhớ động cần thiết cho các phép tính trung gian (chẳng hạn như các hoạt động, độ dốc, v.v.)

Bạn có thể di chuột lên biểu đồ vùng đệm để biết thêm thông tin chi tiết về Op, chẳng hạn như kích thước, hình dạng, loại phân bổ, v.v. Điều này có thể giúp bạn xác định và đánh giá những thao tác có thể có các giá trị tạm thời cao hoặc kéo dài, mọi tensor đầu vào/trung gian/đầu ra lớn có khoảng đệm không hiệu quả, v.v. đang góp phần vào mức sử dụng bộ nhớ cao nhất và cần được điều chỉnh hoặc tối ưu hoá.

Tìm hiểu các kỹ thuật gỡ lỗi cụ thể trong E1000: Gỡ lỗi.