pytorch / opacus / 13383064725

80%

push

github

Adds fast gradient clipping support for the Embedding layer. (#694)

Summary:
The algorithm used is described in the 'A Unified Fast Gradient Clipping Framework for DP-SGD' paper: https://proceedings.neurips.cc/paper_files/paper/2023/file/a45d344b28179c8da7646bc38ff50ad8-Paper-Conference.pdf.

## Types of changes

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Docs change / refactoring / dependency upgrade

## Motivation and Context / Related issue

Previously, Ghost clipping was not supported in Opacus for embedding layer. With default DP-SGD implementation, the training OOMs out on large embedding layers over large physical batch size (useful for privacy). The regular DP-SGD needs O(Bnd), where B=physical batch size, n=vocab size, d=embedding dimension. To give an example on memory needed: we've seen embeddings with [vocab size=1000000, dim=5] (and higher) in real-world differential privacy applications. With a physical batch size of 16,000, memory needed: 16000 × 1000000 × 5 x 4 = 298.02 GB.

With this change, we need significantly smaller memory: O(Br) where B is physical batch size, and r is number of unique  indices in the embedding sequence. We could successfully run DP-SGD over above example, using < 8GiB.

This is a good add to Opacus, enabling larger embedding layers over larger physical batch sizes to be trained with DP-SGD.

## How Has This Been Tested (if it applies)

Unit tests. Runs over large embedding layer training over realworld DP application.

## Checklist

- [x] The documentation is up-to-date with the changes I made.
- [x] I have read the **CONTRIBUTING** document and completed the CLA (see **CONTRIBUTING**).
- [x] All tests passed, and additional code has been covered with new tests.

Pull Request resolved: https://github.com/pytorch/opacus/pull/694

Reviewed B... (continued)

157 of 157 new or added lines in 5 files covered. (100.0%)

5301 of 6181 relevant lines covered (85.76%)

1.91 hits per line