XNNPACK packs a reused weight once per use, so memory grows with every reuse

Đang mở
  • Có gì không ổn?
#24,766 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
45/100
Loại issue
Lỗi
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
machine-learning, performance

Hướng nghiên cứu

Start by tracing how XNNPACK's weight cache is configured in the partitioner and runtime backend options; the issue names no source files or tests. Check how cache settings reach program loading and how reused weights are keyed. Done means repeated uses share a packed weight and the repro shows memory growth near one weight copy, with the existing runtime option still understood.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

What happens

When one layer runs several times in a program, XNNPACK keeps a separate packed copy of its weights for each use. This is common in robot policies: torch.export unrolls a denoising loop, so one Linear becomes N delegate calls that share one weight. Each call packs that weight again, because the XNNPACK weight cache is off by default.

Repro
import os, sys, tempfile, torch
from executorch.exir import to_edge_transform_and_lower
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.runtime import Runtime

n, cache = 10, len(sys.argv) > 1

class Loop(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(4096, 4096)  # 64 MiB of float32 weights
    def forward(self, x):
        for _ in range(n):
            x = torch.tanh(self.lin(x))
        return x

ep = torch.export.export(Loop().eval(), (torch.randn(1, 4096),))
prog = to_edge_transform_and_lower(ep, partitioner=[XnnpackPartitioner()]).to_executorch()
path = os.path.join(tempfile.mkdtemp(), "loop.pte")
with open(path, "wb") as f:
    prog.write_to_file(f)

def rss_mib():
    return int([l for l in open("/proc/self/status") if l.startswith("VmRSS")][0].split()[1]) // 1024

if cache:
    Runtime.get().backend_registry.set_option("XnnpackBackend", {"weight_cache_enabled": True})
before = rss_mib()
Runtime.get().load_program(path).load_method("forward").execute([torch.randn(1, 4096)])
print(f"cache={cache} file_mib={os.path.getsize(path) // 2**20} rss_growth_mib={rss_mib() - before}")

Output (executorch 1.6.0, torch 2.15.0, Linux x86_64):

cache=False file_mib=64 rss_growth_mib=642
cache=True  file_mib=64 rss_growth_mib=65

On a real policy (LeRobot Diffusion Policy, 10 DDIM steps, 1.05 GB of weights, 16 threads), host memory is 10.4 GB without the cache and 1.9 GB with it, and latency goes from 358 ms to 239 ms.

Ask

Turn the weight cache on by default, or let the partitioner record in the program that the cache is needed, so a reused weight is packed once. Today the workaround is to set weight_cache_enabled through the runtime backend options before loading.

Related: #24667.

Ngôn ngữ chính
Python
Star
5.2k
Fork
1.2k
Merge trung bình
1 ngày 23 giờ
Pull request đã merge (30 ngày)
700

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của pytorch/executorch

Tất cả issue của pytorch/executorch

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.