Remove large arrays before big memory operations

CHANGELOG.md

@@ -1,6 +1,14 @@
 Changelog
 =========
 
+Performance Improvements
+
+- Improves memory management to reduce the base memory used during optimization while using `lsq-exact`, `lsq-auglag` and `fmin-auglag` optimizers.
+
+
+v0.17.2
+-------
+
 New Features
 
 - Adds ``eq_fixed`` argument to ``BoundaryError`` to remove the equilibrium from the optimization. This can be used instead of adding a ``FixParameter(eq)`` constraint.

desc/optimize/aug_lagrangian_ls.py

@@ -25,6 +25,7 @@
     inequality_to_bounds,
     print_header_nonlinear,
     print_iteration_nonlinear,
+    scale_columns,
     solve_triangular_regularized,
 )
 
@@ -284,12 +285,10 @@
     diag_h = g * dv * scale
 
     g_h = g * d
-    # TODO: place this function under JIT to use in-place operation (#1669)
-    # we don't need unscaled J anymore, so we overwrite
-    # it with J_h = J * d to avoid carrying so many J-sized matrices
-    # in memory, which can be large
-    J *= d
-    J_h = J
+    # we don't need unscaled J anymore, so we overwrite it with J_h = J * d to avoid
+    # carrying so many J-sized matrices in memory, which can be large. The buffer is
+    # donated so the scaling doesn't allocate a second copy of J.
+    J_h = scale_columns(J, d)
     del J
     g_norm = jnp.linalg.norm(
         (g * v * scale if scaled_termination else g * v), ord=jnp.inf
@@ -523,9 +522,18 @@
             L = L_new
             cost = cost_new
             Lcost = Lcost_new
+            # Delete old arrays before computing new one
+            # otherwise the peak is bigger by J-sized arrays
+            del J_h, J_a
+            if tr_method == "svd":
+                del U, s, Vt
+            elif tr_method == "cho":
+                del B_h
+            elif tr_method == "qr":
+                del R
             J = lagjac(z, y, mu, *args)
             njev += 1
-            g = jnp.dot(J.T, L)
+            g = jnp.dot(L, J)
 
             if jac_scale:
                 scale, scale_inv = compute_jac_scale(J, scale_inv)
@@ -548,9 +556,10 @@
                 # if we update lagrangian params, need to recompute L and J
                 L = lagfun(f, c, y, mu)
                 Lcost = 0.5 * jnp.dot(L, L)
+                del J
                 J = lagjac(z, y, mu, *args)
                 njev += 1
-                g = jnp.dot(J.T, L)
+                g = jnp.dot(L, J)
 
                 if jac_scale:
                     scale, scale_inv = compute_jac_scale(J, scale_inv)
@@ -567,11 +576,7 @@
             d = v**0.5 * scale
             diag_h = g * dv * scale
             g_h = g * d
-            # we don't need unscaled J anymore, so we overwrite
-            # it with J_h = J * d to avoid carrying so many J-sized matrices
-            # in memory, which can be large
-            J *= d
-            J_h = J
+            J_h = scale_columns(J, d)
             del J
 
             if g_norm < gtol and constr_violation < ctol:
@@ -603,6 +608,9 @@
         success, message = False, STATUS_MESSAGES["maxiter"]
     x, s = z2xs(z)
     active_mask = find_active_constraints(z, zbounds[0], zbounds[1], rtol=xtol)
+    # after overwriting J_h with J*d, we have to revert back and store the
+    # unscaled version
+    J_h = scale_columns(J_h, 1 / d)
     result = OptimizeResult(
         x=x,
         s=s,
@@ -613,7 +621,7 @@
         fun=f,
         grad=g,
         v=v,
-        jac=J_h * 1 / d,  # after overwriting J_h, we have to revert back,
+        jac=J_h,
         optimality=g_norm,
         nfev=nfev,
         njev=njev,

desc/optimize/fmin_scalar.py

@@ -24,6 +24,7 @@
     compute_hess_scale,
     print_header_nonlinear,
     print_iteration_nonlinear,
+    scale_columns,
 )
 
 
@@ -213,7 +214,7 @@ def fmintr(  # noqa: C901
         else:
             hess.initialize(N, "hess")
         bfgs = True
-        H = hess.get_matrix()
+        H = jnp.asarray(hess.get_matrix())
     errorif(
         not (isinstance(hess, BFGS) or callable(hess)),
         ValueError,
@@ -239,15 +240,12 @@ def fmintr(  # noqa: C901
     diag_h = g * dv * scale
 
     g_h = g * d
-    # we don't need unscaled H anymore this iteration, so we overwrite
-    # it with H_h = d * H * d[:, None] to avoid carrying so many H-sized matrices
-    # in memory, which can be large
-    # TODO: place this function under JIT (#1669)
-    # doing operation H = d * H * d[:, None]
-    H *= d[:, None]
-    H *= d
-    H_h = H
+    # we don't need unscaled H anymore, so we overwrite it with H_h = d[:, None] * H * d
+    # to avoid carrying so many H-sized matrices in memory, which can be large. The
+    # buffer is donated so the scaling doesn't allocate a second copy of H.
+    H_h = scale_columns(H, d)
     del H
+    H_h = scale_columns(H_h, d[:, None])
 
     g_norm = jnp.linalg.norm(
         (g * v * scale if scaled_termination else g * v), ord=jnp.inf
@@ -422,9 +420,12 @@ def fmintr(  # noqa: C901
             g_old = g
             g = grad(x, *args)
             ngev += 1
+            # Delete old arrays before computing new one
+            # otherwise the peak is bigger by H-sized arrays
+            del H_h, H_a
             if bfgs:
                 hess.update(x - x_old, g - g_old)
-                H = hess.get_matrix()
+                H = jnp.asarray(hess.get_matrix())
             else:
                 H = hess(x, *args)
                 nhev += 1
@@ -439,14 +440,9 @@ def fmintr(  # noqa: C901
 
             g_h = g * d
 
-            # we don't need unscaled H anymore this iteration, so we overwrite
-            # it with H_h = d * H * d[:, None] to avoid carrying so many H-sized
-            # matrices in memory, which can be large
-            # doing operation H = d * H * d[:, None]
-            H *= d[:, None]
-            H *= d
-            H_h = H
+            H_h = scale_columns(H, d)
             del H
+            H_h = scale_columns(H_h, d[:, None])
 
             x_norm = jnp.linalg.norm(
                 ((x * scale_inv) if scaled_termination else x), ord=2
@@ -476,13 +472,17 @@ def fmintr(  # noqa: C901
     if (iteration == maxiter) and success is None:
         success, message = False, STATUS_MESSAGES["maxiter"]
     active_mask = find_active_constraints(x, lb, ub, rtol=xtol)
+    # after overwriting H_h with the scaled version, we have to revert back and
+    # store the unscaled one
+    H_h = scale_columns(H_h, 1 / d)
+    H_h = scale_columns(H_h, 1 / d[:, None])
     result = OptimizeResult(
         x=x,
         success=success,
         fun=f,
         grad=g,
         v=v,
-        hess=H_h / d[:, None] / d,  # unscale the hessian
+        hess=H_h,
         optimality=g_norm,
         nfev=nfev,
         ngev=ngev,

desc/optimize/least_squares.py

@@ -24,6 +24,7 @@
     compute_jac_scale,
     print_header_nonlinear,
     print_iteration_nonlinear,
+    scale_columns,
     solve_triangular_regularized,
 )
 
@@ -183,7 +184,8 @@ def lsqtr(  # noqa: C901
     # jax.pure_callback and gets stuck due to async dispatch
     J = jac(x, *args).block_until_ready()
     njev += 1
-    g = jnp.dot(J.T, f)
+    # g is J.T@f and this is equal to f@J for 1D f
+    g = jnp.dot(f, J)
 
     maxiter = setdefault(maxiter, n * 100)
     max_nfev = options.pop("max_nfev", 5 * maxiter + 1)
@@ -203,13 +205,11 @@ def lsqtr(  # noqa: C901
     diag_h = g * dv * scale
 
     g_h = g * d
-    # TODO: place this function under JIT to use in-place operation (#1669)
-    J *= d
 
-    # we don't need unscaled J anymore, so we overwrite
-    # it with J_h = J * d to avoid carrying so many J-sized matrices
-    # in memory, which can be large
-    J_h = J
+    # we don't need unscaled J anymore, so we overwrite it with J_h = J * d to avoid
+    # carrying so many J-sized matrices in memory, which can be large. The buffer is
+    # donated so the scaling doesn't allocate a second copy of J.
+    J_h = scale_columns(J, d)
     del J
     g_norm = jnp.linalg.norm(
         (g * v * scale if scaled_termination else g * v), ord=jnp.inf
@@ -405,9 +405,18 @@ def lsqtr(  # noqa: C901
             allx.append(x)
             f = f_new
             cost = cost_new
+            # Delete old arrays before computing new one
+            # otherwise the peak is bigger by J-sized arrays
+            del J_h, J_a
+            if tr_method == "svd":
+                del U, s, Vt
+            elif tr_method == "cho":
+                del B_h
+            elif tr_method == "qr":
+                del R
             J = jac(x, *args)
             njev += 1
-            g = jnp.dot(J.T, f)
+            g = jnp.dot(f, J)
 
             if jac_scale:
                 scale, scale_inv = compute_jac_scale(J, scale_inv)
@@ -418,11 +427,7 @@ def lsqtr(  # noqa: C901
             diag_h = g * dv * scale
 
             g_h = g * d
-            J *= d
-            # we don't need unscaled J anymore this iteration, so we overwrite
-            # it with J_h = J * d to avoid carrying so many J-sized matrices
-            # in memory, which can be large
-            J_h = J
+            J_h = scale_columns(J, d)
             del J
             x_norm = jnp.linalg.norm(
                 ((x * scale_inv) if scaled_termination else x), ord=2
@@ -451,14 +456,17 @@ def lsqtr(  # noqa: C901
     if (iteration == maxiter) and success is None:
         success, message = False, STATUS_MESSAGES["maxiter"]
     active_mask = find_active_constraints(x, lb, ub, rtol=xtol)
+    # after overwriting J_h with J*d, we have to revert back and store the
+    # unscaled version
+    J_h = scale_columns(J_h, 1 / d)
     result = OptimizeResult(
         x=x,
         success=success,
         cost=cost,
         fun=f,
         grad=g,
         v=v,
-        jac=J_h * 1 / d,  # after overwriting J_h, we have to revert back
+        jac=J_h,
         optimality=g_norm,
         nfev=nfev,
         njev=njev,

desc/optimize/utils.py

@@ -487,6 +487,16 @@ def compute_jac_scale(A, prev_scale_inv=None):
     return 1 / scale_inv, scale_inv
 
 
+@functools.partial(jit, donate_argnums=0)
+def scale_columns(A, d):
+    """Compute `A * d` reusing `A`'s buffer instead of allocating a second copy.
+
+    `A` is invalid after this call, so callers must rebind, ie
+    `A = scale_columns(A, d)`.
+    """
+    return A * d
+
+
 @jit
 def compute_hess_scale(H, prev_scale_inv=None):
     """Compute scaling factors based on diagonal of Hessian matrix."""

tests/benchmarks/memory_benchmark_cpu.py

@@ -1,32 +1,28 @@
 """Benchmark memory usage of various functions."""
 
 import gc
+import os
 import pickle
 import subprocess
 import sys
 import threading
 import time
 
 import numpy as np
-import psutil
 
 
 def monitor_ram(proc, interval, ram_usage, timestamps):
-    """Sample system RAM until *proc* finishes."""
-    while proc.poll() is None:  # check if child still running
-        info = psutil.virtual_memory()
-        used_mb = (info.total - info.available) / 1024 / 1024
-        ram_usage.append(used_mb)
-        timestamps.append(time.time())
-        time.sleep(interval)
-
-    # keep watching for an extra second
-    end = time.time() + 1.0
-    while time.time() < end:
-        info = psutil.virtual_memory()
-        ram_usage.append((info.total - info.available) / 1024 / 1024)
-        timestamps.append(time.time())
-        time.sleep(interval)
+    """Sample the child's resident set size until *proc* finishes."""
+    page_mb = os.sysconf("SC_PAGE_SIZE") / 1024 / 1024
+    with open(f"/proc/{proc.pid}/statm", "rb") as statm:
+        while proc.poll() is None:  # check if child still running
+            try:
+                statm.seek(0)
+                ram_usage.append(int(statm.read(64).split()[1]) * page_mb)
+            except OSError:  # child exited between the poll and the read
+                break
+            timestamps.append(time.time())
+            time.sleep(interval)
 
 
 def monitor_vram(proc, interval, vram_usage, timestamps):
@@ -72,7 +68,7 @@ def monitor_vram(proc, interval, vram_usage, timestamps):
 
 if __name__ == "__main__":
     mode = "CPU"  # "CPU" or "GPU"
-    interval = 0.01  # seconds between samples
+    interval = 0.001  # seconds between samples
 
     data = {}
 

tests/benchmarks/memory_funcs.py

@@ -225,8 +225,8 @@ def test_eq_solve():
     eq = desc.examples.get("precise_QA")
     eq.change_resolution(L=res, M=res, L_grid=2 * res, M_grid=2 * res)
     # this test is mostly for intermediate operations, so having a chunk size
-    # of 100 will be fine to see their effect
-    obj = ObjectiveFunction(ForceBalance(eq), jac_chunk_size=100, deriv_mode="batched")
+    # of 30 will be fine to see their effect
+    obj = ObjectiveFunction(ForceBalance(eq), jac_chunk_size=30, deriv_mode="batched")
     obj.build(verbose=0)
     eq.solve(
         objective=obj,