BUG: memory grows with every forward and adjoint evaluation when a mixed space has a Real component

Aperta
#5,534 5 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
45/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python

Direzione di ricerca

Start by tracing the example's mixed-space Function creation and the reported MixedDat.dataset / MixedDataSet cache in PyOP2; the issue does not name source files or tests. Reproduce the serial example and inspect cache growth across forward and adjoint evaluations. Done means memory stays roughly flat for the mixed space with a Real component while the plain Poisson comparison remains flat.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

bug LLM used pyop2

In adjoint inversions, memory grows with the number of forward and adjoint evaluations whenever the solve is on a mixed space with a Real component. We first saw this on an HPE Cray EX machine with Firedrake 2026.10.0, in a 256-rank inversion that was killed by the OOM killer after 10 optimisation iterations. The example below shows it in serial on a MacBook (Apple silicon) with Firedrake main at ef9e4b1e2.

The cause seems to be the MixedDataSet cache in PyOP2. MixedDat.dataset builds a MixedDataSet from the datasets of its components, and caches it on the Set of the first component with the tuple of datasets as the key. A Real component is a Global with its own GlobalDataSet, so every new Function on the mixed space gives a new key. The new entry stays in the cache for as long as the mesh lives. pyadjoint makes new Functions at every replay and adjoint step, so the cache grows with every evaluation.

import gc

import psutil
from firedrake import *
from firedrake.adjoint import *

continue_annotation()
mesh = UnitSquareMesh(128, 128)
V = FunctionSpace(mesh, "CG", 1)
Z = V * FunctionSpace(mesh, "R", 0)
x, y = SpatialCoordinate(mesh)
f = Function(V).interpolate(cos(pi * x) * cos(pi * y))

# Poisson problem with an extra uniform source c, chosen so that the mean of u is zero.
z = Function(Z)
u, c = TrialFunctions(Z)
v, d = TestFunctions(Z)
solve(inner(grad(u), grad(v)) * dx + c * v * dx + u * d * dx == f * v * dx, z,
      DirichletBC(Z.sub(0), 0, "on_boundary"))
J = assemble(z.subfunctions[0] ** 2 * dx)
rf = ReducedFunctional(J, Control(f))
pause_annotation()

for i in range(1, 201):
    rf(f)
    rf.derivative()
    if i % 50 == 0:
        gc.collect()
        print(f"after {i} evaluations: {psutil.Process().memory_info().rss / 2**20:.0f} MB")

Output:

after 50 evaluations: 503 MB
after 100 evaluations: 566 MB
after 150 evaluations: 655 MB
after 200 evaluations: 741 MB

The same script without the Real unknown (a plain Dirichlet Poisson problem on V) stays flat at 332 to 339 MB.

I used Claude Code to reproduce this and to write the example.

Lingua principale
Python
Stelle
684
Fork
204
Merge medio
4g 14h
PR unite (30g)
46

Preparare l'ambiente

  • Nessun Dockerfile né file Docker Compose
  • Ha un modello di pull request
  • Nessuna guida per i contributori

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di firedrakeproject/firedrake

Tutte le issue di firedrakeproject/firedrake

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.