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

Offen
  • Stimmt etwas nicht?
#24,766 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen

Maintainer antworten meist innerhalb von 1 Tag

Dieses Issue hat noch niemand übernommen.

Bewertung

Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Anfängerfreundlichkeit
45/100
Issue-Typ
Bug
Klarheit
Größtenteils klar
Aktivitätsstatus
Aktiv
Tech-Stack
python, pytorch
Bereich
machine-learning, performance

Rechercherichtung

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.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Beschreibung

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.

Vorherrschende Sprache
Python
Sterne
5.2k
Forks
1.2k
Ø Merge
1 T. 23 Std.
Gemergte PRs (30 T.)
700

Entwicklungsumgebung

Erste Schritte

  1. Lesen Sie das ganze Issue und danach den Beitragsleitfaden des Projekts.
  2. Schreiben Sie ins Issue, dass Sie es übernehmen — das erspart doppelte Arbeit.
  3. Forken Sie das Repository und arbeiten Sie in einem Branch.
  4. Öffnen Sie einen Pull Request, der die Issue-Nummer nennt.

Mehr aus pytorch/executorch

Alle Issues in pytorch/executorch

Ähnliche Issues

Weitere Issues zu Python

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.