Question about `inference_mode=False` in `pl.Trainer` and enabling gradients in validation and test step

[fedeotto]

Hello,

I’ve been following your work with great interest and have spent some time exploring and re-implementing your model to better understand and control its behavior.

While I believe my re-implementation is faithful to your original design, I’ve been having difficulty reproducing the same performance results at a comparable scale. I noticed that in your PyTorch Lightning setup, inference_mode=False is set in the trainer, and gradient computation is enabled during both the validation_step() and test_step() methods.

I haven’t yet applied this approach in my own implementation, so I wanted to ask about the rationale behind allowing gradient computation during inference. This isn’t something I’ve typically encountered, and I’m curious whether it addresses a known issue or could potentially explain some of the performance discrepancies I’m seeing.

Many thanks for your interesting work!

Best regards,
Federico

[FranklinHu1]

Hi Federico,

Thanks for your interest in our work!

One thing to clarify is that the results produced by the model and published in the paper were done without PyTorch lightning, instead using the training utilities implemented under the training package. Regarding the gradient computation being enabled in the validation and test steps: this is actually related to a subtle bug in the implementation of PyTorch's transformer class where disabling gradients can cause different behaviors in the transformer stack (see pytorch/pytorch#106630). The simple workaround is to enable gradient information during the validation and test steps as that is always correct so long as we do not backpropagate/train on those gradients.

I hope that helps, and would be interested in hearing more about your experiments with the model.

Cheers!

Frank