PyTorch Lightning Support?

I’m trying to utilise opacus with the PyTorch Lightning framework which we use as a wrapper around a lot of our models. I can see that there was an effort to integrate this partially into PyTorch Lightning late last year but this seems to have stalled due to lack of bandwidth.

I’ve created a simple MVP but there seems to be a compatibility problem with even this simple model; it throws AttributeError: 'Parameter' object has no attribute 'grad_sample' as soon as it hits the optimization step.

What’s the likely underlying cause of this? I can see on the opacus GitHub that similar errors have been encountered before where it’s been caused by unsupported layers but as the gist shows, this model is incredibly simple so I don’t think it’s any of the layers.

This is with:


Hi James! Would you have time to file a bug and share a colab so I can take a look? Integrating with Lightning is indeed on our plate :slight_smile:

Sure thing; didn’t want to go straight to a bug report as it didn’t feel right

1 Like