Conversation
|
cc @frapac |
|
Docs are broken due to #520 |
|
Hi @klamike ! Can you expand a bit more on the context for this and potential use of it? |
|
@tmigot The idea is to define an API to talk about batches of models in a standardized way across JSO. Then each solver (up for discussion -- maybe we should do that here instead?) can define its own batch model types, which would mark the relevant assumptions (e.g. same jac/hess structure, same nvar/ncon, etc.) to enable exploiting that shared structure (importantly, not just in the evaluation of the model, but also in the solvers themselves). For example, if we assume the same jac/hess structure, we can use CUDSS uniform batch API for solving KKT systems instead of looping over single solves. This should bring substantial speedup over something like To be honest, this kind of change is very big for me to think about all the consequences of different design decisions at once. So, we decided it might be best to just implement something dumb first to get some visibility into tradeoffs. |
| isempty(updates) && error("Cannot create InplaceBatchNLPModel from empty collection.") | ||
| InplaceBatchNLPModel{M}(base_model, updates, Counters(), length(updates)) | ||
| end | ||
| # TODO: counters? |
There was a problem hiding this comment.
Note the counters are not properly wired up yet, both here and in foreach
| ] | ||
| @test hprods ≈ manual_hprods | ||
|
|
||
| if isa(bnlp, ForEachBatchNLPModel) # NOTE: excluding InplaceBatchNLPModel |
There was a problem hiding this comment.
Had to exclude the inplace version here since the closures in the operators reference the model state, which we are mutating without their knowledge
There was a problem hiding this comment.
It should be doable to put the update function inside the operator closure..
|
@klamike Can you rebase your branch with |
|
@amontoison I've rebased, tests should pass now. Curious what you think of this approach, as well as the downstream klamike/MadIPM.jl#1 and klamike/MadNLP.jl#1 (which so far do not depend on this PR) |
|
@klamike @andrewrosemberg
We don’t want to update all KKT systems if some of them have already converged. |
|
@klamike I worked a little bit on the batch support in export AbstractBatchNLPModel
export batch_obj, batch_obj!
export batch_grad, batch_grad!
export batch_cons, batch_cons!
export batch_jac_structure!, batch_jac_structure
export batch_hess_structure!, batch_hess_structure
export batch_jac_coord!, batch_jac_coord
export batch_hess_coord!, batch_hess_coord
export batch_jprod, batch_jprod!
export batch_jtprod, batch_jtprod!
export batch_hprod, batch_hprod!We don't need something else and I prefer that we extend it in the future if needed. What is still missing is a |
|
784ccb7 LGTM feel free to push it here! Thanks for spending some time on this. Agreed on all your points, except I'm not so sure we should force strided (for example in multi-GPUs/multi-node setting we may not want it). It can always be changed/added separately later. Also, it would be nice to have batch attributes similar to the operator availability ones you added recently for NLPModelMeta, like |
|
closing in favor of #540 |
4 participants
to be continued in #540...