Getting proximal to work with augmented Lagrangian optimizers - #2298
Getting proximal to work with augmented Lagrangian optimizers#2298singh-jaydeep wants to merge 12 commits into
Conversation
Memory benchmark result| Test Name | %Δ | Master (MB) | PR (MB) | Δ (MB) | Time PR (s) | Time Master (s) |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
test_objective_jac_w7x | -0.11 % | 4.244e+03 | 4.239e+03 | -4.48 | 31.81 | 30.04 |
test_proximal_jac_w7x_with_eq_update | -3.02 % | 6.816e+03 | 6.610e+03 | -205.84 | 154.26 | 163.71 |
test_proximal_freeb_jac | -0.68 % | 1.354e+04 | 1.344e+04 | -91.46 | 76.94 | 83.25 |
test_proximal_freeb_jac_blocked | -0.74 % | 7.868e+03 | 7.810e+03 | -58.32 | 65.36 | 72.32 |
test_proximal_freeb_jac_batched | -0.96 % | 7.878e+03 | 7.803e+03 | -75.43 | 64.64 | 70.98 |
+ test_proximal_jac_ripple | -11.94 % | 3.829e+03 | 3.372e+03 | -457.24 | 45.82 | 55.19 |
+ test_proximal_jac_ripple_bounce1d | -16.46 % | 3.911e+03 | 3.267e+03 | -643.71 | 58.74 | 71.12 |
test_eq_solve | 0.31 % | 1.829e+03 | 1.835e+03 | 5.59 | 53.94 | 54.00 |
test_objective_quadratic_flux_jac | -0.38 % | 1.894e+03 | 1.887e+03 | -7.28 | 36.18 | 34.91 |For the memory plots, go to the summary of |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #2298 +/- ##
=======================================
Coverage 94.35% 94.36%
=======================================
Files 101 101
Lines 29053 29119 +66
=======================================
+ Hits 27414 27477 +63
- Misses 1639 1642 +3
🚀 New features to boost your workflow:
|
| # If self._state has not yet computed and stored the tangents, | ||
| # this function gets called. | ||
| def build_tangents(): | ||
| return _proximal_get_tangents( |
There was a problem hiding this comment.
I think you can safely move this function to ProximalState. Almost all the arguments are state attributes anyway. You just need to handle tuple(state.dimc_per_thing) or pass it too.
There was a problem hiding this comment.
I ended up passing in eq_idx as well. Let me know if this is what you were thinking.
The state.get_tangents method is now called directly by the ProximalProjection wrapper. Just as an aside: while refactoring this I realized the tangent cache wasn't set up correctly, and resulting in lots of cache misses. This is due to the difference in how jvp and vjp were written, with some wrappers (like LCP) routing grad through jvp and others (ProximalProjection, after my changes) going through vjp. Now the jvp and vjp methods go through the same cache. A consequence is the added flag cache_tangents on ProximalState. If you don't want to use the cache, as is likely the case in proximal-lsq-exact (you don't need to share tangents across multiple wrappers), you use the default False. For proximal-auglag methods we set True.
|
Also a quick question, with these changes is there a way to still use the old way? It would be nice if this only happened for 'proximal-auglag' or 'proximal-lsq-auglag'. |
|
@YigitElma I implemented your suggested changes; let me know if anything is unclear. I'll think about whether there is a simple method to let users choose the old way. |
YigitElma
left a comment
YigitElma
left a comment
There was a problem hiding this comment.
Looks much cleaner than befoer. Thanks for all the changes!
Does this memory improvement have a specific reason? I cannot think of any. Can there be a bug?
| # Does not include self._constraint | ||
| self._objective._set_things(self.things) | ||
|
|
||
| self._eq_idx = self.things.index(self._state.eq) |
There was a problem hiding this comment.
A bit of a nitpick, but I prefer keeping these in the build method, and removing this new definition of the _set_things all together. I personally think having many small helper methods actually makes the code harder to read (need to jump from function to function). But others prefer that I am fine.
There was a problem hiding this comment.
The main reason for this being separate was due to a sync issue in optimizer.optimize(). When the different wrappers are built, they may have different "things". The combine_args call later in that function calls set_things with the definitive set of things, allowing us to sync the layout vectors (like dimx_per_thing) at that time. Another option could be to just collect the list of things earlier in optimize(), and pass that to the wrapper. I can try that.
| ) | ||
|
|
||
| xs = jnp.split(x, np.cumsum(self._dimc_per_thing)) | ||
| xs = jnp.split(x, np.cumsum(self._dimc_per_thing)[:-1]) |
There was a problem hiding this comment.
Was this a bug before? Can you explain why we need this?
There was a problem hiding this comment.
Yeah you're right, we should not need this. At one point I was concerned about the last entry in cumsum, but it doesn't seem to be a problem.
Co-authored-by: Yigit Gunsur Elmacioglu <102380275+YigitElma@users.noreply.github.com>
|
My best guess for the memory change is that the tangent cache is just returning the proximal tangents on the second and third call in the test. This would be a test specific thing, and wouldn't correspond to a savings in practice. I will try explicitly clearing the cache in the test between calls and see if that gives identical results to master. |
If the optimizer is |
2 participants
One way of addressing #873. Right now, the only nonlinear constraints which proximal methods accept are force balance related ones.
ProximalProjectionalready handles the process of translating between the full set of state variables (including R_lmn, Z_lmn, L_lmn) and the reduced set of optimization variables, and ensures that an objective wrapped in proximal gets the proper tangents. For augmented lagrangian methods, the nonlinear constraints need those same tangents.This adds a class,
ProximalState, which is responsible for managing all properties of the equilibrium subproblem. Any proximal wrapped objective, e.g. the objective or a nonlinear constraint, can access the up-to-date equilibrium this way. The state also stores the equilibrium tangents for the present equilibrium, so multiple calls to grad start to just read the cache. It seems to work for the few cases I have tried.Since this is an issue others have thought about more, there may be better ways of getting this to work. Lmk and I can refactor.
Also some of the diff is due to changes from #2239, which will disappear once that is merged.
Resolves #873
Resolves #1720