`None` grad at first iteration

[ripaul]

Hi, I've been trying to use IVON for a MDE task, but I get the following error which leaves me kinda clueless:

Traceback (most recent call last):
  File "/p/project1/.../src/mde.py", line 205, in <module>
    main(args)
  File "/p/project1/.../src/mde.py", line 182, in main
    train_loss = train(model, optim, train_dataloader)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/p/project1/.../src/mde.py", line 87, in train
    return _train_epoch_ivon(model, optim, dataloader, n_samples)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/p/project1/.../src/mde.py", line 69, in _train_epoch_ivon
    with optim.sampled_params(train=True):
  File "/p/software/jurecadc/stages/2024/software/Python/3.11.3-GCCcore-12.3.0/lib/python3.11/contextlib.py", line 144, in __exit__
    next(self.gen)
  File "/p/project1/.../venv/lib/python3.11/site-packages/ivon/_ivon.py", line 132, in sampled_params
    self._restore_param_average(train, param_avg, noise)
  File "/p/project1/.../venv/lib/python3.11/site-packages/ivon/_ivon.py", line 149, in _restore_param_average
    param_grads.append(p.grad.flatten())
                       ^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'flatten'

the corresponding calling code snippet is this one:

 64 def _train_epoch_ivon(model, optim, dataloader, n_samples=1):
 65     train_loss = 0
 66 
 67     for x, y in tqdm(dataloader):
 68         for _ in range(n_samples):
 69             with optim.sampled_params(train=True):
 70                 optim.zero_grad()
 71                 x = x.to(device)
 72                 y = y.to(device)
 73                                                   
 74                 ypred = predict(model, x, y.shape)              
 75                 loss = F.l1_loss(normalize(ypred), normalize(y))
 76
 77                 train_loss += loss.item()
 78
 79                 loss.backward()
 80
 81         optim.step()
 82                                        
 83     return train_loss / len(dataloader)

Is there any obvious mistake I'm overlooking? n_samples is set to 1.

Btw, this is running on Python 3.11, PyTorch 2.1.2 and ivon-opt 0.1.3.

Thanks in advance for any help! :)

[ripaul]

I did some follow up on this and found that the issue were unused parameters in the model, which would never get a gradient update. I guess one could add some kind of safeguard to _restore_param_average to account for this? To me it happened working with a "third-party" implementation from huggingface, so I assume this issue may hit others as well.

[ysngshn]

Hi Richard! Thanks for your interest in our work!

I think the issue is that the current IVON implementation assumes at each step all parameters have gradients. So if only parts of the parameters are involved in the autodiff process then some parameters will receive no gradient (param.grad = None) and IVON can't compute with None 😅

We are getting multiple feedback about this problem and are thinking of a solution to remove this constraint, hopefully soon!

In the meantime maybe you can try the dirty workaround of resetting grads to zero instead of None:

optim.zero_grad(set_to_none=False)

Of course, this changes the behavior and IVON will track some unnecessary zero gradients. But it might get your code running for now.

[ripaul]

hey, thanks! I didn't know about the set_to_none option, but this seems like a good workaround for the moment. I had solved the issue by first passing a dummy input through the model, then calling backward and finally only passing the parameters to IVON which had a populated grad.