pytorch
/
opacus
/
17145830368
/
3
80%
main: 80%
main: 80%
|
Ran
|
Files
70
|
Run time
2s
|
Badge
Embed ▾
README BADGES
|
coverage: 47.815% (-0.1%) from 47.924%
push
github
facebook-github-botTensor Parallelism -- Review Version (#776) Summary: Pull Request resolved: https://github.com/pytorch/opacus/pull/776 We add a class named `grad_sample_module_fast_gradient_clipping_tp` to support 1D tensor parallelism for Opacus. We demonstrate the effectiveness on a toy example (examples/tp_toy.py) and llama 2 model (examples/tp_llama.py). The key contribution is to understand when merging per-sample norms from different local devices, the operation should be `reduce_all` or just keep the norm from device 0. Specifically, under the following exceptions: 1. The parameter is not a DTensor. 2. The model is nn.embedding and under RowWiseParallel. 3. The parameter is nn.linear.bias under RowWiseParallel. We should not merge the per-sample norm from local devices, but maintain the one from device 0. Currently, we do not support 1. Vanilla Opacus mode (non-GC). 2. Sequential Parallelism (usually applied to LayerNormalization). 3. Modify/customize the placements of module outputs, which might lead to wrong behavior of norm calculation. 4. Trainable parameter being a plain tensor or replicated because of the constraint of the optimizer. Reviewed By: aparna-aketi Differential Revision: D79062242 fbshipit-source-id: ffaf071d1
1466 of 3066 relevant lines covered (47.81%)
0.48 hits per line
