diff --git a/CHANGELOG.md b/CHANGELOG.md index dfeee4bf11..1a46a132cc 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py index 5badddc61f..64c55f2587 100644 --- a/desc/optimize/aug_lagrangian_ls.py +++ b/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 @@ def lagjac(z, y, mu, *args): 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 @@ def lagjac(z, y, mu, *args): 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 @@ def lagjac(z, y, mu, *args): # 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 @@ def lagjac(z, y, mu, *args): 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 @@ def lagjac(z, y, mu, *args): 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 @@ def lagjac(z, y, mu, *args): 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, diff --git a/desc/optimize/fmin_scalar.py b/desc/optimize/fmin_scalar.py index ee40126a93..4c5a85f2b6 100644 --- a/desc/optimize/fmin_scalar.py +++ b/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, diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py index 125247f1e6..a73d14079c 100644 --- a/desc/optimize/least_squares.py +++ b/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,6 +456,9 @@ 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, @@ -458,7 +466,7 @@ def lsqtr( # noqa: C901 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, diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py index ad5b4a4b3c..41845ad392 100644 --- a/desc/optimize/utils.py +++ b/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.""" diff --git a/tests/benchmarks/memory_benchmark_cpu.py b/tests/benchmarks/memory_benchmark_cpu.py index 6308278961..3c49fdfb01 100644 --- a/tests/benchmarks/memory_benchmark_cpu.py +++ b/tests/benchmarks/memory_benchmark_cpu.py @@ -1,6 +1,7 @@ """Benchmark memory usage of various functions.""" import gc +import os import pickle import subprocess import sys @@ -8,25 +9,20 @@ 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 = {} diff --git a/tests/benchmarks/memory_funcs.py b/tests/benchmarks/memory_funcs.py index f17e1dada1..0f1250bd1e 100644 --- a/tests/benchmarks/memory_funcs.py +++ b/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,