What is the best practise to use FSDP with opacus?

username: Pier297
cooked: <p>Hi all!</p>
<p>I was wondering how/if it’s possible to use opacus together with PyTorch FSDP (or deepspeed) to allow for fine-tuning of large LM that doesn’t fit on a single gpu.</p>
<p>Right now what I managed to do is basically have each gpu compute a sample gradient, clip it and then accumulate the gradients of the different processes so that then I can add the noise.<br>
The problem is that this is very slow and I was wondering if there was any better way to do this?</p>
<p>Thank you so much and have a good day:)<br>
Pier</p>

username: ffuuugor
cooked: <p>Good question - supporting FSDP could be useful for larger models, however there are certain challenges.</p>
<p>The biggest problem is - in order to properly clip the gradient you need to know the overall norm of the entire gradient, which in case of FSDP is spread across multiple machines. It’s not 100% clear from your post if you do this, but after you compute per sample gradient on each GPU you need to broadcast the norm to get the clipping coefficient.</p>
<p>This naturally adds one more point of synchronization between GPUs, which inevitably impacts performance. I don’t have a good intuition on how much it should impact the training speed - what exactly do you mean by “very slow” here?</p>
<p>One more thing - from your description it looks like you add noise post-syncronisation, you don’t need to do that, you can add noise on one of the workers, the <a href="https://github.com/pytorch/opacus/blob/main/opacus/distributed.py" rel="noopener nofollow ugc">same way we do for DDP</a></p>
<p>Hope this helps</p>

username: Pier297
cooked: <p>Hi and thank you a lot for your response,</p>
<p>the way I clip the sample gradient is by using the <a href="https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.no_sync" rel="noopener nofollow ugc">no_sync</a> context manager when calling <code>loss.backward()</code> so that on each gpu I only have the gradient of a single sample (I think), I then clip it and add to a local variable that basically accumulated the clipped gradients until the batch is over.</p>
<p>Then my real problem, but I guess it’s more of a FSDP question (have just asked about it <a href="https://github.com/pytorch/pytorch/issues/92820" rel="noopener nofollow ugc">here</a> in case feel like taking a look ahah), is how I can set back the noisy gradient inside the model parameters <code>.grad</code> attribute so that I can take the optimization step.</p>
<p>Many thanks for your help!</p>

username: ffuuugor
cooked: <p>Ok, based on what I’ve read about FSDP and <code>no_sync</code>, what you’re doing is broadly right (bear in mind though I’ve never used FSDP myself, so might be wrong about this).</p>
<p>The part that confuses me a little is this:</p>
<blockquote>
<p>so that on each gpu I only have the gradient of a single sample (I think), I then clip it and add to a local variable that basically accumulated the clipped gradients until the batch is over.</p>
</blockquote>
<p>Without getting too deep into FSDP context, it sounds like you’re computing per-sample gradients one-by-one? If that’s the case, that’s a good explanation why your code is slow - thing that makes opacus fast is vectorized per-sample gradient computation.</p>

username: Ali_Dadsetan
cooked: <p>I think <a href="https://github.com/awslabs/fast-differential-privacy" rel="noopener nofollow ugc">this repository</a> has implemented what you guys were discussing here. Also, this is the <a href="https://arxiv.org/abs/2311.11822" rel="noopener nofollow ugc">link to the paper</a>.</p>