Meta optimization failing with functorch transforms

[thurbridi]

I'm trying to implement the code from this paper. Following TorchOpts tutorials, I'm using `adam(use_accelerated_op=True) as the meta optimizer. However, the following runtime error appears:

File "/home/casper/Development/modac/src/modac_torch2.py", line 385, in update_all_params
  new_eta, new_eta_opt_state, new_theta, new_theta_opt_state = outer_update(
                                                               ^^^^^^^^^^^^^
File "/home/casper/Development/modac/src/modac_torch2.py", line 325, in outer_update
  eta_grads, (new_theta, new_theta_opt_state) = torch.func.grad(
                                                ^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1380, in wrapper
  results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 39, in fn
  return f(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1245, in wrapper
  output = func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/Development/modac/src/modac_torch2.py", line 300, in outer_loss
  theta, theta_opt_state = inner_update(
                           ^^^^^^^^^^^^^
File "/home/casper/Development/modac/src/modac_torch2.py", line 258, in inner_update
  updates, new_theta_opt_state = theta_opt_update(
                                 ^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torchopt/base.py", line 196, in update_fn
  updates, new_s = fn(updates, s, params=params, inplace=inplace)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torchopt/combine.py", line 92, in update_fn
  flat_updates, state = inner.update(flat_updates, state, params=flat_params, inplace=inplace)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torchopt/base.py", line 196, in update_fn
  updates, new_s = fn(updates, s, params=params, inplace=inplace)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torchopt/transform/scale_by_adam.py", line 328, in update_fn
  out = tree_map_flat(
        ^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torchopt/transform/utils.py", line 65, in tree_map_flat
  return flat_arg.__class__(map(fn, flat_arg, *flat_args))  # type: ignore[call-arg]
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torchopt/transform/scale_by_adam.py", line 326, in op_fn
  return op(mu, nu, update, count)  # type: ignore[arg-type]
         ^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torchopt/accelerated_op/adam_op.py", line 168, in __call__
  new_mu = self.MuOp.apply(updates, mu, self.b1)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/casper/miniconda3/envs/torch2/lib/python3.11/site-packages/torch/autograd/function.py", line 509, in apply
  raise RuntimeError(
RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. For more details, please see https://pytorch.org/docs/master/notes/extending.func.html

[XuehaiPan]

Hi @thurbridi, the accelerated_op for the Adam optimizer only supports second-order gradients (see our MAML examples (bi-level optimization)). Have you ever tried to set use_accelerated_op=False? This will use the native PyTorch ops which support higher orders of autograd computation.

[thurbridi]

Thanks for the quick response, XuehaiPan!

I have, but then I run into the NaN problem, so my last alternative would be the zero_nan_hook, I suppose? I'll try that.