RuntimeError: Batching rule not implemented for aten::_make_dual; the fallback path doesn't work on out= or view ops

[JieaiMai]

RuntimeError: Batching rule not implemented for aten::_make_dual; the fallback path doesn't work on out= or view ops
UserWarning: .test(ckpt_path="best") is called with Trainer configured with multiple ModelCheckpoint callbacks. It will use the best checkpoint path from first checkpoint callback.
rank_zero_warn(
Restoring states from the checkpoint path at ./rfmcsp-conditional-perov/ekgeh2jv/checkpoints/epoch=1184-step=14220.ckpt
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [4]
Loaded model weights from checkpoint at ./rfmcsp-conditional-perov/ekgeh2jv/checkpoints/epoch=1184-step=14220.ckpt
Testing DataLoader 0: 0%| | 0/1 [00:00<?, ?it/s]Error executing job with overrides: ['data=perov', 'model=null_params']
Traceback (most recent call last):
File "/data0/mai/examples/flowmm/scripts_model/run.py", line 206, in main
run(cfg)
File "/data0/mai/examples/flowmm/scripts_model/run.py", line 193, in run
trainer.test(datamodule=datamodule, ckpt_path=ckpt_path)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 780, in test
return call._call_and_handle_interrupt(
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 38, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 829, in _test_impl
results = self._run(model, ckpt_path=self.ckpt_path)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1098, in _run
results = self._run_stage()
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1174, in _run_stage
return self._run_evaluate()
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1214, in _run_evaluate
eval_loop_results = self._evaluation_loop.run()
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
self.advance(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 152, in advance
dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
self.advance(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 137, in advance
output = self._evaluation_step(**kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 234, in _evaluation_step
output = self.trainer._call_strategy_hook(hook_name, *kwargs.values())
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1480, in _call_strategy_hook
output = fn(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 399, in test_step
return self.model.test_step(*args, **kwargs)
File "/data0/mai/examples/flowmm/src/flowmm/model/model_pl.py", line 1030, in test_step
return self.shared_eval_step(
File "/data0/mai/examples/flowmm/src/flowmm/model/model_pl.py", line 736, in shared_eval_step
loss_dict = self.loss_fn(batch)
File "/data0/mai/examples/flowmm/src/flowmm/model/model_pl.py", line 546, in loss_fn
return self.rfm_loss_fn(batch, *args, **kwargs)
File "/data0/mai/examples/flowmm/src/flowmm/model/model_pl.py", line 594, in rfm_loss_fn
x_t, u_t = manifold.cond_u(x0, x1, t)
File "/data0/mai/examples/flowmm/src/flowmm/rfm/vmap.py", line 235, in cond_u
return self.vmap(CondU, (x0, x1, t), data_in_dim=data_in_dim)
File "/data0/mai/examples/flowmm/src/flowmm/rfm/vmap.py", line 125, in vmap
return torch.func.vmap(wrapper, (0, 0, data_in_dim), randomness="different")(
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/apis.py", line 188, in wrapped
return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/vmap.py", line 266, in vmap_impl
return _flat_vmap(
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/vmap.py", line 38, in fn
return f(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/vmap.py", line 379, in _flat_vmap
batched_outputs = func(*batched_inputs, **kwargs)
File "/data0/mai/examples/flowmm/src/flowmm/rfm/vmap.py", line 118, in wrapper
return torch.func.functional_call(
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/functional_call.py", line 143, in functional_call
return nn.utils.stateless._functional_call(
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/nn/utils/stateless.py", line 264, in _functional_call
return module(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/data0/mai/examples/flowmm/src/flowmm/rfm/vmap.py", line 24, in forward
return method(*args, **kwargs)
File "/data0/mai/examples/flowmm/src/flowmm/rfm/manifolds/product.py", line 106, in cond_u
x_t, u_t = self._cond_u(manifold, x0p, x1p, t)
File "/data0/mai/examples/flowmm/src/flowmm/rfm/manifolds/product.py", line 95, in _cond_u
x_t, u_t = jvp(path, (t,), (torch.ones_like(t).to(t),))
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py", line 919, in jvp
return _jvp_with_argnums(func, primals, tangents, argnums=None, strict=strict, has_aux=has_aux)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/vmap.py", line 38, in fn
return f(*args, **kwargs)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py", line 962, in _jvp_with_argnums
flat_duals = tuple(fwAD.make_dual(p, t)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py", line 962, in
flat_duals = tuple(fwAD.make_dual(p, t)
File "/data0/jieai/miniforge3/envs/flowmm/lib/python3.9/site-packages/torch/autograd/forward_ad.py", line 121, in make_dual
return torch._VF._make_dual(tensor, tangent, level=level)
RuntimeError: Batching rule not implemented for aten::_make_dual; the fallback path doesn't work on out= or view ops.

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace
It seems that this error occurs during evaluation, but it works fine during training. How can I solve this problem?

[Dipannoy]

I am also finding similar error. Have you been able to solve the issue?

[bkmi]

Not sure. There seems to be an issue with the jvp.

I encourage you to setup a breakpoint before that line and investigate there.

File "/data0/mai/examples/flowmm/src/flowmm/rfm/manifolds/product.py", line 95, in _cond_u
x_t, u_t = jvp(path, (t,), (torch.ones_like(t).to(t),))

[Dipannoy]

Actually, I have done some debugging before calling jvp. The tensor 't' looks same both for training and testing when I print the following line

print(f"t: {t.shape}, {t.dtype}, contiguous: {t.is_contiguous()}, requires_grad: {t.requires_grad}")
The output is: Input tensors: torch.Size([1]) torch.float32 True False

But training is working fine, when testing starts it shows the error. I have also tried 'vjp' in place of jvp. But it comes up with other size mismatch problems.

[bkmi]

It could be an issue earlier with CUDA.

Try CUDA_LAUNCH_BLOCKING=1 and the hydra flag.

[Dipannoy]

Thanks Benjamin for your suggestions!

But trying the breakpoint and CUDA_LAUNCH_BLOCKING=1, I am not finding anything clear about the problem.

I was checking the source code file python3.9/site-packages/torch/autograd/forward_ad.py
I printed the tensor and tangent in the make_dual(tensor, tangent, *, level=None) function.

--------------------------------Tensor and Tangent--------------------------------------
BatchedTensor(lvl=1, bdim=0, value=
    tensor([[0.3264],
            [0.9967],
            [0.0706],
            ...,
            [0.8715],
            [0.6707],
            [0.7588]], device='cuda:0')
)
BatchedTensor(lvl=1, bdim=0, value=
    tensor([[1.],
            [1.],
            [1.],
            ...,
            [1.],
            [1.],
            [1.]], device='cuda:0')
)

Then using the breakpoint I get the following sequence.

> /home/chris/micromamba/envs/flowmm/lib/python3.9/site-packages/torch/autograd/forward_ad.py(125)make_dual()
-> return torch._VF._make_dual(tensor, tangent, level=level)
(Pdb) n
--Return--
> /home/chris/micromamba/envs/flowmm/lib/python3.9/site-packages/torch/autograd/forward_ad.py(125)make_dual()->GradTrackingT...da:0')
    )
)
-> return torch._VF._make_dual(tensor, tangent, level=level)
(Pdb) n
--Return--
> /home/chris/micromamba/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py(962)<genexpr>()->GradTrackingT...da:0')
    )
)
-> flat_duals = tuple(fwAD.make_dual(p, t)
(Pdb) n
--Call--
> /home/chris/micromamba/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py(962)<genexpr>()->GradTrackingT...da:0')
    )
)
-> flat_duals = tuple(fwAD.make_dual(p, t)
(Pdb) n
> /home/chris/micromamba/envs/flowmm/lib/python3.9/site-packages/torch/_functorch/eager_transforms.py(962)<genexpr>()->GradTrackingT...da:0')
    )
)
-> flat_duals = tuple(fwAD.make_dual(p, t)

What I have understood, the error occurs from
flat_duals = tuple(fwAD.make_dual(p, t) But I couldn't debug more as eager_transforms.py is written in private helper function.

[liuthe]

I encountered the same problem

[bkmi]

Perhaps it has to do with inference mode?

A quick google revealed this result:
pytorch/pytorch#138800

[liuthe]

Perhaps it has to do with inference mode?

A quick google revealed this result: pytorch/pytorch#138800

Yes! I am still stuck in the problem. This line "x_t, u_t = jvp(path, (t,), (torch.ones_like(t).to(t),))" works well during the training process, but outputs error during the testing process. The t variable looks like the same in both processes. Maybe it has something to do with the path method?

[sunyrain]

same problem……