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

Aberta
  • Algo errado?
#24,766 0 comentários 0 reações 0 responsáveis Ver no GitHub

Mantenedores costumam responder em até 1 dia

Ninguém assumiu esta issue ainda.

Avaliação

Dificuldade
4/5
Tempo estimado
3-5 dias
Facilidade para iniciantes
45/100
Tipo de issue
Bug
Clareza
Razoavelmente clara
Status de atividade
Ativa
Stack de tecnologia
python, pytorch
Domínio
machine-learning, performance

Direção de pesquisa

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.

Escrita pelo modelo de indexação a partir do texto da issue.

Descrição

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.

Linguagem predominante
Python
Estrelas
5.2k
Forks
1.2k
Merge médio
1d 23h
PRs com merge (30d)
700

Preparar o ambiente

Primeiros passos

  1. Leia a issue inteira e depois o guia de contribuição do projeto.
  2. Comente na issue dizendo que vai assumir — evita que duas pessoas façam o mesmo trabalho.
  3. Faça um fork do repositório e trabalhe em uma branch.
  4. Abra um pull request que referencie o número da issue.

Mais de pytorch/executorch

Todas as issues de pytorch/executorch

Issues semelhantes

Mais issues de Python

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.