Why not include a callback that stops the training when the Differential Privacy budget has been spent? See [20853](https://github.com/Lightning-AI/pytorch-lightning/pull/20853).