Add Optimization Cookbook

[samanklesaria]

What does this PR do?

This PR adds a guide that shows some common techniques for working with Flax models during optimization. These include:

• Calculation of Exponential Moving Averages
• Optimizing only a low rank addition to certain weights (LORA)
• Using different learning rates for different parameters to implement the maximal update parameterization
• Using second order optimizers like LBFGS.
• Specifying sharding for optimization state that differs from that of parameter state
• Gradient accumulation

This document emphasizes a style as close to pure jax as possible: to that end, it shows how the flax version of each technique only requires minor deviation from the often more intuitive pure-jax version.

Warnings:

• This PR requires the fix Add _graph_node_set_key method for List class #5171 , which is why the diff looks so big here. This PR shouldn't be merged until after Add _graph_node_set_key method for List class #5171
• This PR includes both the notebook file used to construct the optimization examples as well as the RST file that contains what gets used in the documentation. When review is complete (and I don't need to actually test any of the python snippets anymore) I will delete the notebook file.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[cgarciae]

I feel we could simplify the intro by doing the following:

1. Define a single model at the begining, simply reuse it on all examples (its just a guide).
2. Inline the loss function.

model = nnx.Sequential(
  nnx.Linear(2,8, rngs=rngs),
  nnx.relu,
  nnx.Linear(8,8, rngs=rngs),
)

optimizer = nnx.Optimizer(
  model,
  tx=optax.adam(1e-3),
  wrt=nnx.Param)

...

@nnx.jit
def train_step(model, optimizer, ema, x, y):
  loss_fn = lambda m, x, y: jnp.sum((m(x) - y) ** 2)
  loss, grads = nnx.value_and_grad(loss_fn)(model, x, y)
  optimizer.update(model, grads)
  ema.update(model)
  return loss

[cgarciae]

I thin this is enough:

Suggested change

	param_tys = nnx.map_state(lambda p, v: list(p)[-1], state)
	param_tys = nnx.map_state(lambda p, v: p[-1], state)

[cgarciae]

We could also use jax.tree. map_with_path as in the JAX example.

[cgarciae]

instead of creating this fake model it would be a good opportunity to create the optimizer_sharding API on Variable before finishing this guide.

[cgarciae]

After fully reading the guide I'm getting the sense that having the JAX versions makes explanations a bit longer and slightly harder to understand (cause you have to mentally filter for the version you are interested in) and having the JAX version doesn't necessarily make understanding the NNX version easier.

[samanklesaria]

After fully reading the guide I'm getting the sense that having the JAX versions makes explanations a bit longer and slightly harder to understand (cause you have to mentally filter for the version you are interested in) and having the JAX version doesn't necessarily make understanding the NNX version easier.

Fair enough! I'll convert it to nnx-only.