Zero redundancy distributed learning with differential privacy

## Abstract

Abstract Deep learning using large models have achieved great success in a wide range of domains. However, training these models on billions of parameters is very challenging in terms of the training speed, memory cost, and communication efficiency, especially under the privacy-preserving regime with differential privacy (DP). On the one hand, DP optimization has comparable efficiency to the standard non-private optimization on a single GPU, but on multiple GPUs, existing DP distributed learning (such as pipeline parallel) has suffered from significantly worse efficiency. On the other hand, the Zero Redundancy Optimizer (ZeRO) is a state-of-the-art solution to the standard distributed learning, exhibiting excellent training efficiency on large models, but to work compatibly with DP is technically complicated. In this work, we develop a new systematic solution, DP-ZeRO, (I) to scale up the trainable DP model size, e.g. to GPT-100B, (II) to obtain the same computation and communication efficiency as the standard ZeRO, and (III) to enable mixed-precision DP training. Our DP-ZeRO, like the standard ZeRO, has the potential to train models with arbitrary size and is evaluated on the world’s largest DP models in terms of the number of trainable parameters.

## 1 Introduction

Recent advances in differentially private (DP) deep learning have witnessed the power of large pre-trained models, achieving comparable accuracy to state-of-the-art (SOTA) non-private models across computer vision De et al. (); Bu et al. (); Mehta et al. (); Xie et al. (), natural language processing Yu et al. (); Li et al. (); Bu et al. (), and many other tasks. Similar to their non-DP counter-parts, it has been observed that larger DP models tend to have better performance. For example, the DP accuracy increases from 83% using RoBERTa-base (123M parameters) to 86% using RoBERTa-large (354M parameters) on GLUE datasets Li et al. (); Bu et al. (); Yu et al. (); the DP BLEU score increases from 61 using GPT2-small (124M parameters) to 64 using GPT2-large (800M parameters) on E2E dataset Li et al. (); Bu et al. (); a similar trend is also observed using ViT (Base/Large/Huge) up to 600M parameters to achieve state-of-the-art DP accuracy on ImageNet, around 81% at $\epsilon=8$ Mehta et al. ().

Driven by this success and the surge of computational power, it is high time to enable DP deep learning at the same scale of the standard non-DP one, e.g., GPT3-175B (Brown et al., ) and LLaMA-63B (Touvron et al., ; ) with billions of trainable parameters. Specifically, such a DP training system must have high time and memory efficiency, low communication cost, and the compatibility with general neural network architectures.

For small to moderately large models (e.g. with less than a billion parameters) that fit within the memory of a single GPU, a range of DP algorithms are feasible, producing the same result at different efficiency. Examples include TensorFlow-privacy Subramani et al. (), Opacus Yousefpour et al. (); Bu et al. (), ghost clipping (GhostClip) Goodfellow (); Li et al. (); Bu et al. (), and Book-Keeping (BK) Bu et al. (), among which the BK algorithm has allowed DP optimization to be almost as efficient as the standard one. To be specific, the time/space complexity of BK algorithm is $1.08\times/1.05\times$ of the standard optimization on ViT-Large (300M parameters, 147 layers) and $1.03\times/1.01\times$ on GPT2-large (800M parameters, 220 layers).

To enable the DP distributed learning of these not-too-large models, one can directly use DDP (distributed data parallelism) (), where each mini-batch of data is partitioned to smaller micro-batches and each GPU computes one micro-batch with a full copy of the DP model. A line of researches (Yousefpour et al., ; De et al., ; Kurakin et al., ) have reported that DDP with DP usually either incurs huge memory cost due to caching the per-sample gradients, or suffers from $2-9\times$ slower training speed than non-DP optimization De et al. (); Bu et al. (). While the efficiency issues can be addressed through a better DP algorithm, such as BK, the feasibility issue remains insurmountable because DDP cannot train models that exceed the capacity of one GPU. Notably, the efficiency of BK algorithm is enhanced by two key techniques: mixed ghost norm (computing per-sample gradient norms almost for free) and book-keeping trick (only using one round of full back-propagation, not two rounds as in Li et al. (); Bu et al. ()), which are detailed in [Appendix A](#A1) and will also be leveraged in our DP-ZeRO solution.

As the model size further increases beyond a reasonable bound for one GPU (e.g. 32GB memory, which roughly translates to 2B model training with Adam), the model must be partitioned in addition to the data, e.g. using pipeline parallelism and model parallelism, so that each GPU only holds a partial shard of the model (see [Figure 2](#S2.F2)). In He et al. (), DP is combined with pipeline parallelism to fine-tune about 0.1% of GPT3-175B. Yet, the pipeline parallelism can be inefficient due to a non-DP-related issue, known as the pipeline bubble, where GPUs are idle while waiting for data to process.

**Table 1: Summary of DP distributed learning.**
| Distributed<br>solution | Parallelism | Model<br>sharding | Standard<br>version | DP<br>version | Remark |
| --- | --- | --- | --- | --- | --- |
| DDP | Data | No |  | Yousefpour et al. () | unable to fit large model and DP is memory costly |
| DDP | Data | No | Frostig et al. () | De et al. () | unable to fit large model and DP is slow |
| GPipe | Pipeline | Yes | Huang et al. () | He et al. () | pipeline bubble wastes GPU time |
| ZeRO | Data(&Model) | Yes | Rajbhandari et al. () | Ours | speed & memory & communication efficient |

Generally speaking, more advanced distributed methods such as Zero Redundancy Optimizer Rajbhandari et al. () (ZeRO) and mixed-precision training have not be paired with DP due to the lack of algorithmic advances. In this work, we develop DP-ZeRO, equipping state-of-the-art distributed learning solution with DP (see comparison in [Table 1](#S1.T1)), without altering the mathematics of DP optimization. We summarize our contributions as follows.

- 1.
We propose the zero redundancy distributed learning with differential privacy (DP-ZeRO), demonstrating the same level of communication efficiency, computation efficiency (speed and memory), and scalability (e.g. to GPT3 level and hundreds of GPUs) as the standard ZeRO.
- 2.
We enable the mixed-precision training with DP by addressing the issue of loss scaling. This advance allows us to reduce the memory cost by roughly 50% and allow significantly faster communication that was previously not enjoyed by DP distributed learning.
- 3.
We enable DP deep learning with more than 1B trainable parameters for the first time. E.g. we are the first to train the full GPT2-XL, ViT-Gigantic, ViT-10B and GPT-100B with DP.
- 4.
We will open-source our codebase that automatically applies DP-ZeRO for general tasks (e.g. classification and language understanding), general network architectures (e.g. ResNet, ViT, GPT), and general distributed solutions (including DeepSpeed and FSDP), with one line of code change.

Figure: Figure 1: Total/trainable parameters of existing DP models (blue) and ours by DP-ZeRO (red).
Refer to caption: /html/2311.11822/assets/x1.png

## 2 Preliminary

### 2.1 Differential privacy

DP provides a formal privacy guarantee, making it difficult to extract any information from training data. The privacy guarantee is characterized by $(\epsilon,\delta)$-DP in [Definition 2.1](#S2.Thmothertheorem1), with smaller $(\epsilon,\delta)$ representing lower privacy risk.

###### Definition 2.1 ( Dwork et al. ( 2006 ) ) .

A randomized algorithm $M$ is $(\varepsilon,\delta)$-DP if, for any two neighboring datasets $S,S^{\prime}$ that differ by one sample and for any event $E$, we have $\mathbb{P}[M(S)\in E]\leqslant\mathrm{e}^{\varepsilon}\mathbb{P}\left[M\left(S^{\prime}\right)\in E\right]+\delta.$

In DP deep learning, the gradients are made private by post-processing through per-sample gradient clipping and random noising:

$$ $\displaystyle\text{private gradient: }\mathbf{G}_{[m]}:=\sum_{i}C_{i}(R_{m})\bm{g}_{[m],i}+\sigma_{\text{DP}}\|[R_{1},\cdots,R_{M}]\|\cdot\mathcal{N}(0,\mathbf{I}),$ (1) $$

Here the gradient of all trainable parameters is partitioned into $M$ groups, i.e. $\bm{g}_{[m],i}$ is the $i$-th per-sample gradient of the $m$-th group’s parameters, where $m\in\{1\cdots M\}$ is the group index. $C_{i}$ is the per-sample gradient clipping factor so that $\|C_{i}\bm{g}_{[m],i}\|\leq R_{m}$ and $R_{m}$ is the clipping threshold. That is, DP optimization is enabled when the standard optimizers such as stochastic gradient descent (SGD) and Adam (Kingma & Ba, ) update the trainable parameters with the private gradient, instead of the standard gradient $\sum_{i}\bm{g}_{i}$.

###### Definition 2.1 ( Dwork et al. ( 2006 ) ) .

##### Mathematical gradient partition

In ([1](#S2.E1)), the trainable parameters and their gradients are mathematically partitioned into $M$ groups, e.g. in all-layer clipping, all parameters form one group ($M=1$) Abadi et al. (); in layer-wise clipping McMahan et al. (); He et al. (), each layer’s parameters form a group ($M=$ number of layers).
Empirical evidence and theoretical analysis show that different partitions have the same training speed, though a finer partition (e.g. layer-wise) has lighter memory footprint(^1^11We note that DP optimization under different group-wise clippings can have the same computation and communication efficiency (under the BK algorithm), with or without ZeRO.).

##### Per-sample gradient clipping

In ([1](#S2.E1)), a number of clipping functions $C_{i}=C(\|\bm{g}_{i}\|;R)$ are available. Most works Abadi et al. (); Li et al. (); Yu et al. () use the vanilla clipping $C_{i}=\min(R/\|\bm{g}_{i}\|,1)$. Recently, Bu et al. (); Yang et al. () advocate the automatic clipping $C_{i}=1/(\|\bm{g}_{i}\|+0.01)$ which is hyperparameter-free and comparably accurate.
Note if $C_{i}\equiv 1$, then the clipped gradient reduces to the standard gradient. The main overhead of DP optimization is the computation of per-sample gradient norms. On a single GPU, the mixed ghost clipping Bu et al. () has reduced the time complexity to $<10\%$ on large models like GPT2.

##### Privacy accounting

In ([1](#S2.E1)), adding Gaussian noise to the clipped gradient protects the privacy that is quantifiable by the privacy accounting theory Abadi et al. (); Bu et al. (); Dong et al. (); Zhu et al. (); Gopi et al. (); Koskela et al. (). The privacy guarantee is increasing in the noise level $\sigma_{\text{DP}}$, independent of $R_{m}$, learning rate, clipping function and model architectures, with $\sigma_{\text{DP}}=0$ leading to $\epsilon=\infty$ (non-private).

### 2.2 Zero Redundancy Optimizer (ZeRO)

#### 2.2.1 Parallel computing

Parallel computing is necessary to train large-scale models and is critical to the optimization efficiency. For models that fit in a single GPU, data parallelism (DataP) can be used to speed up the training by partition the mini-batch of samples into multiple micro-batches. Then, each GPU (holding a full copy of parameters) executes the forward and backward propagation of one micro-batch, from which the parameter gradients are generated and averaged across GPUs to update the trainable parameters. However, for models that do not fit in a single GPU, the model parameters need to be sharded by alternative solutions such as ZeRO Rajbhandari et al. (), model parallelism (ModelP) and pipeline parallelism (PipeP).

ModelP partitions a model vertically, e.g. using 3 GPUs to store the parameters of one layer. As a consequence, ModelP does not scale efficiently beyond a single node due to fine-grained computation and expensive communication between layers. Implementation-wise, ModelP frameworks usually require heavy code integration that may not be generalizable in model architectures. In contrast, PipeP partitions a model horizontally across layers, e.g. storing 3 layers in each GPU. Each GPU deals with all micro-batches sequentially, though PipeP can be inefficient due to the pipeline bubble, which is overcome by ZeRO Rajbhandari et al. (). ZeRO is an advanced data parallel method that eliminates memory redundancies during the training, and improves the training speed and communication volume proportionally to the number of GPUs. Unlike basic DataP, ZeRO partitions a model’s states across GPUs and gather/reduce in a just-in-time manner, thus sustaining the high efficiency of very large model training. Notice that ZeRO can work compatibly with ModelP and optionally offload the model states to CPUs Ren et al. (); Rajbhandari et al. ().

#### 2.2.2 Model state partition

A major part of the training memory is consumed by the model states(^2^22Another important part of memory consumption is the batch-size-related variables such as the activation tensors, which is instantiated during the forward propagation and independent to DP (which modifies the gradient during the back-propagation).). ZeRO has three stages (ZeRO1/2/3) that partition these model states by different levels, with lower level being faster but more memory costly. For instance, in Table 2 of (Rajbhandari et al., ), ZeRO1/2/3 at most train 7.6/14.4/128B models on 64 V100 GPUs.

We take an example of mixed-precision Adam optimizer to train a model with $\Psi_{\text{model}}$ parameters, which maintains a master copy (fp32) of optimizer states
, and the half-precision parameters and gradients.

##### Optimizer state partition

The optimizer states are the (master) parameters, variance and momentum, each taking $4\Psi_{\text{model}}$ memory. ZeRO1 only applies the optimizer state partition and updates the parameters locally, reducing the memory cost of model states from $16\Psi_{\text{model}}$ for basic DataP to $(4+\frac{12}{N_{d}})\Psi_{\text{model}}$ at each of $N_{d}$ GPUs.

Figure: Figure 2: Mathematical (left two, same GPU allocation, different accuracy) and hardware (right two, different GPU allocation, same accuracy) gradient partition. Orange and blue are gradient groups $\bm{g}_{[1]}$ and $\bm{g}_{[2]}$ in DP optimization ([1](#S2.E1)).
Refer to caption: /html/2311.11822/assets/figures/partition.png

##### Hardware gradient partition

In addition to ZeRO1, during the back-propagation, ZeRO2 (and ZeRO3) further partitions the $2\Psi_{\text{model}}$ gradients into different GPUs, reducing the memory cost to $(2+\frac{14}{N_{d}})\Psi_{\text{model}}$ at each GPU. [Figure 2](#S2.F2) illustrates the difference between the hardware partition and the mathematical partition in [Section 2.1](#S2.SS1).

##### Parameter partition

In addition to ZeRO2, ZeRO3 also partitions the $2\Psi_{\text{model}}$ fp16 parameters, further reducing the memory cost to $\frac{16}{N_{d}}\Psi_{\text{model}}$ at each GPU.

### 2.3 Mixed-precision training

Mixed-precision training (Micikevicius et al., ) performs the forward and backward propagation on the half precision (fp16 or bf16) parameters, activations, and gradients, while performing the model update in full precision (fp32). Compared to the full-precision training, it is capable of saving the memory by $\approx 50\%$ and accelerating the computation by $\approx 20\%$. We note that fp16 has better precision but a limited range: its representable numbers are among $10^{-8}\sim 10^{5}$, and vice versa for bf16. Therefore, loss scaling is necessary when using fp16 to prevent small gradients from being rounded to zero, thereby preserving the model’s accuracy (see [Appendix C](#A3) for details). In contrast, bf16 usually does not need loss scaling since it has the same range as fp32. However, bf16 is not as widely supported as fp16, e.g. only available on NVIDIA Ampere GPUs or above ().

## 3 Differentially private ZeRO

### 3.1 Algorithm

Our DP-ZeRO algorithm introduces the per-sample gradient clipping and noising to the standard ZeRO (Rajbhandari et al., ), while maintaining the efficiency. At high level, an iteration of ZeRO consists of the following steps:

$$ $\left(\hbox{\pagecolor{Plum!20} all-gather}\to\hbox{\pagecolor{LimeGreen!20}forward}\to\hbox{\pagecolor{Plum!20}all-gather}\to\hbox{\pagecolor{LimeGreen!20}backward}\to\hbox{\pagecolor{Plum!20}reduce}\right)^{\times L\text{ layers}}\to\hbox{\pagecolor{LimeGreen!20}update(SGD/Adam/...)}$ $$

The operations in purple are global and require communication among GPUs, whereas the operations in green are locally computed within each GPU. In particular, DP optimization is only different from the standard optimization in the back-propagation, which can be decomposed into

$$ $\hbox{\pagecolor{LimeGreen!20}backward}=(\text{output gradient}\to\text{clipping factor}\to\text{parameter gradient}\to\text{noising})$ $$

To give more details, we consider training a neural network of linear layers using $N_{d}$ GPUs. We emphasize that the following procedure is sufficiently generic to cover other layer types, such as convolution, embedding, normalization, and so on, which are all supported by DP-ZeRO. The full algorithm is depicted in [Figure 3](#S3.F3), where we denote the $j$-th micro-batched variables like $\bm{a}_{l}^{(j)}$, for $1\leq j\leq N_{d}$.

Figure: Figure 3: Algorithm of DP-ZeRO with mixed-precision training.
Refer to caption: /html/2311.11822/assets/figures/flowchart.png

The forward propagation of DP optimization is the same as that of the standard optimization:

$$ $\displaystyle\bm{s}_{l}=\bm{a}_{l}\mathbf{W}_{l}+\mathbf{b}_{l},\bm{a}_{l+1}=\phi_{l}(\bm{s}_{l}).$ (2) $$

At the $l$-th layer, $\bm{a}_{l}\in\mathbb{R}^{BT_{l}d_{l}}$ is the layer’s input, also known as the activation, $\bm{s}_{l}\in\mathbb{R}^{BT_{l}p_{l}}$ is the layer’s output, $\mathbf{W}_{l}\in\mathbb{R}^{d_{l}p_{l}}$ is the weight, $\mathbf{b}_{l}\in\mathbb{R}^{p_{l}}$ is the bias, and $\phi_{l}$ is any inter-layer operation such as ReLU, tanh or pooling. We denote $B$ as the physical micro-batch size(^3^33The micro-batch size $B$ is the number of samples processed by each GPU, which determines the time and memory efficiency, but not the performance. The logical batch size that determines the performance is $B\times N_{d}\times\text{gradient_accumulation_steps}$.) and $T_{l}$ as the hidden feature dimension (e.g. sentence length or number of pixels). During the forward propagation, the activations $\{\bm{a}_{l}\}$ are computed and stored in the computation graph, and the loss $L=\sum_{i}L_{i}$ is derived, where $L_{i}$ is the per-sample losses. During the back-propagation, the output gradient is first computed based on the previous layer.

$$ $\frac{\partial L}{\partial\bm{s}_{l}}=\frac{\partial L}{\partial\bm{s}_{l+1}}\frac{\partial\bm{s}_{l+1}}{\partial\bm{a}_{l+1}}\circ\frac{\partial\bm{a}_{l+1}}{\partial\bm{s}_{l}}=\frac{\partial L}{\partial\bm{s}_{l+1}}\mathbf{W}_{l+1}\circ\phi_{l}^{\prime}(\bm{s}_{l}),$ $$

in which $\circ$ is element-wise multiplication. Specifically, the use of parameter $\mathbf{W}_{l+1}$ necessitates the all-gather operation when the model is partitioned into multiple GPUs, which is not needed in single GPU training.
Next, the activation $\bm{a}_{l}$ is used together with $\frac{\partial L}{\partial\bm{s}_{l}}$ to compute the parameter gradient:

$$ DP gradient: $\displaystyle\frac{\partial\sum_{i}C_{i}L_{i}}{\partial\mathbf{W}_{l}}+\sigma\mathcal{N}(0,\mathbf{I})=\bm{a}_{l}^{\top}\text{diag}(C_{1},\cdots,C_{B})\frac{\partial L}{\partial\bm{s}_{l}}+\sigma\mathcal{N}(0,\mathbf{I}).$ $$

Note the standard gradient can be viewed as $C_{i}=1,\sigma=0$. Here the per-sample gradient norm (or the clipping factor $C_{i}$) can be computed at small cost, as we have discussed in [Section 1](#S1).

### 3.2 Time efficiency of DP-ZeRO

The time efficiency of DP-ZeRO consists of two parts: the local computation (including forward and backward propagation) and the global communication (including intra-node and inter-node communication). Given that the only difference between DP-ZeRO and ZeRO is the back-propagation, we claim that DP-ZeRO could enjoy high efficiency on-par with the standard ZeRO when (I) DP back-propagation exhibits a time efficiency comparable to the standard, similar to the single GPU training, and/or (II) the time efficiency of the parts other than back-propagation is not
insignificant. We give the time of each part of DP-ZeRO in ([3](#S3.E3)) to illustrate our claim.

$$ $\displaystyle\frac{\text{DP-ZeRO speed}}{\text{standard ZeRO speed}}=\frac{{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\text{back-propagation}}+{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\text{forward propagation}}+{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\text{communication}}}{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\text{DP back-propagation}}+{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\text{forward propagation}}+{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\text{communication}}}$ (3) $$

To be explicit, we summarize the time complexity in [Table 2](#S3.T2) and refer to [Appendix B](#A2) for details.

**Table 2: Time complexity of one iteration in distributed learning(^5^55Here 0.666 is figurative and dependent on settings. Notice that $\Psi_{\text{train}}\ll\Psi_{\text{model}}$ when most parameters are frozen.). We denote $\Psi_{\text{train}}$ to be the number of trainable parameters ($\Psi_{\text{train}}=\Psi_{\text{model}}$ in full parameter training), and define $B,T$ below ([2](#S3.E2)).**
|  | forward propagation | back-propagation | communication |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | activation | attention | output grad | param grad | DP clip and noise |  |
| complexity | $2BT\Psi_{\text{model}}$ | $O(BT^{2})$ | $2BT\Psi_{\text{model}}$ | $2BT\Psi_{\text{train}}$ | $0.666BT\Psi_{\text{train}}$ | $O(\Psi_{\text{model}})$ |

In what follows, we analyze the absolute and relative speed (to standard ZeRO) of DP-ZeRO under important settings.

#### 3.2.1 Number of computation devices

When scaling from one GPU (zero communication) to one node (multiple GPUs) and to multiple nodes, the communication efficiency decreases sub-linearly ($\uparrow$communication). On a single node, multiple GPUs can communicate using the high-speed intra-node connections such as NVLink/NVSwitch (Foley & Danskin, ; Ishii et al., ). On multiple nodes, which are necessary for large models, the inter-node connections are $3\sim 24\times$ slower than the intra-node connections (Li et al., ; Zhang et al., ). In short, DP-ZeRO can be as fast as ZeRO by ([3](#S3.E3)) when multiple nodes are employed.

#### 3.2.2 Memory-efficient distributed learning

The communication volume is specific to different distributed algorithms, most of which trade the communication or speed for memory, in order to feasibly train very large models. For example, ZeRO3 (but not ZeRO1/2) needs to all-gather the sharded parameters at each iteration, hence suffering from 50% extra communication volume ($\uparrow$communication). Another example is the activation check-pointing (also known as gradient check-pointint Chen et al. ()), where a second forward propagation re-computes the expensive activations during back-propagation, though at a 33% slower speed ($\uparrow$forward propagation). These techniques improve the relative speed of DP-ZeRO but worsens the absolute speed.

#### 3.2.3 Parameter efficient fine-tuning

Parameter efficient fine-tuning (PEFT), such as LoRA (Hu et al., ), Adapter (Houlsby et al., ), and BiTFiT (Zaken et al., ), optimizes a small fraction (e.g. $\Psi_{\text{train}}=0.1\%\Psi_{\text{model}}$) of model parameters and thus boosts the efficiency of back-propagation and communication ($\downarrow$DP back-propagation $\downarrow$back-propagation $\downarrow$communication). Consequently, (I) the communication volume of the gradient can be reduced possibly by $1000\times$; (II) the local computation can accelerate by 50% (Hu et al., ; Bu et al., ), which can be seen by treating $\Psi_{\text{train}}$ in [Table 2](#S3.T2) as almost zero; (III) the memory cost is saved on the non-trainable layers, which translates to larger batch size and faster computation. Hence, both relative and absolute speed of DP-ZeRO improve using PEFT.

### 3.3 Memory efficiency of DP-ZeRO

We claim that DP-ZeRO is as memory efficient as the standard ZeRO, similar to the single GPU training, when we use (I) the mixed ghost norm trick Bu et al. (; ), instead of GhostClip (Goodfellow, ; Li et al., ) or per-sample gradient instantiation (Yousefpour et al., ); (II) the layer-wise clipping style instead of the all-layer clipping, so that the book-keeping (Bu et al., ) does not store all output gradients; (III) a large number of GPUs so that the micro-batch size $B$ (i.e. per-GPU batch size) is small: specially, when
$B=1$ in the gradient accumulation, the per-sample gradient is free. We empirically verify our claim in [Figures 5](#S4.F5) and [7](#S4.F7).

### 3.4 Mixed-precision training with DP

Figure: Figure 4: Accuracy of mixed-precision training with loss scaling. ViT-large on CIFAR100.
Refer to caption: /html/2311.11822/assets/x3.png

We now analyze the intricacy in mixed-precision training with DP, which is not unique to DP-ZeRO but present in the general DP optimization. We emphasize that the per-sample gradient clipping already plays the role of scaling, and hence DP mixed-precision training must not use loss scaling, as illustrated in [Table 3](#S3.T3). Specifically, in standard mixed-precision training, there are two steps of scaling: (I) scaling up the loss ${L}_{i}$ by $10^{3}\sim 10^{9}$ (and consequently the output gradient $\frac{\partial{L}_{i}}{\partial\bm{s}}$ as well as $\frac{\partial{L}_{i}}{\partial\mathbf{W}}$) before the back-propagation, to prevent the underflow where fp16 gradient is too small to be distinguished from 0, and (II) scaling down the parameter gradient $\frac{\partial{L}_{i}}{\partial\mathbf{W}}$, by the same factor, after the back-propagation to recover the correct magnitude of gradient. However, in DP mixed-precision training, scaling up the loss may cause overflow, while scaling down the gradient incorrectly over-shrinks the gradient and worsens the performance. See [Figure 4](#S3.F4) for a real example. We explain this intricacy step-by-step in [Appendix C](#A3).

**Table 3: Illustration of overflow and underflow issues during mixed-precision training (ghost norm).**
| loss scale=$10^{3}$ | activation $\bm{a}_{l}$ | output grad $\frac{\partial L}{\partial\bm{s}_{l}}$ | per-sample grad norm | clipping factor | param grad | param grad |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| (scaled) | $\text{vec}(\bm{a}_{l}\bm{a}_{l}^{\top})$ | $\text{vec}(\frac{\partial L}{\partial\bm{s}_{l}}\frac{\partial L}{\partial\bm{s}_{l}}^{\top})$ | (not scaled down) | (if scaled down) |  |  |  |
| standard w/o scaling | $10^{-3}\sim 10^{2}$ | ${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}10^{-8}}\sim 10^{1}$ | N/A | N/A | 1 | $10^{-7}\sim 10^{1}$ | $10^{-7}\sim 10^{1}$ |
| standard w/ scaling | $10^{-3}\sim 10^{2}$ | $10^{-5}\sim 10^{4}$ | N/A | N/A | 1 | $10^{-4}\sim 10^{4}$ | $10^{-7}\sim 10^{1}$ |
| DP w/o scaling | $10^{-3}\sim 10^{2}$ | ${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}10^{-8}}\sim 10^{1}$ | $10^{2}\sim 10^{3}$ | $10^{-6}\sim 10^{0}$ | $10^{-3}\sim 10^{2}$ | $10^{-7}\sim 10^{1}$ | $10^{-7}\sim 10^{1}$ |
| DP w/ scaling | $10^{-3}\sim 10^{2}$ | $10^{-5}\sim 10^{4}$ | $10^{2}\sim 10^{3}$ | $10^{0}\sim{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}10^{6}}$ | $10^{-6}\sim 10^{-1}$ | $10^{-7}\sim 10^{1}$ | ${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}10^{-10}\sim 10^{-2}}$ |

## 4 Empirical performance of DP-ZeRO

We evaluate DP-ZeRO on five aspects: model architectures, efficiency, scalability, compatibility with various distributed learning and DP techniques. We use DP-ZeRO to train ResNet (He et al., ), ViT (Dosovitskiy et al., ; Zhai et al., ) and GPT (Radford et al., ; Brown et al., ), which are workhorses in computer vision and language tasks. We measure the time and memory efficiency of DP-ZeRO under settings such as PEFT and multiple precision formats (fp32 or fp16/bf16). We evaluate the scalability of DP-ZeRO in terms of number of GPUs and number of model parameters. Our experiments scale from single node (8 GPUs) to multiple nodes, up to 256 GPUs, and train models up to 100B trainable parameters. Moreover, DP-ZeRO is compatible with mainstream implementations of ZeRO(^6^66DP-ZeRO is implemented on DeepSpeed (supporting ZeRO1/2/3), FSDP () (supporting ZeRO3), MiCS (Zhang et al., ) (supporting ZeRO2/3), and any distributed optimizers supported on them.) and with different clipping styles, clipping functions, privacy accountants, and so on.
We leave the experimental details in [Appendix D](#A4). By default, we use AdamW, mixed-precision training, layer-wise clipping style, $B=4$, and A100 GPU with 40GB memory, unless otherwise stated.

### 4.1 Generality of DP-ZeRO

DP-ZeRO is generally applicable to different neural network architectures, clipping styles, and precision formats. We test DP-ZeRO1 on single node and observe that different clipping styles are equally fast, but layer-wise clipping is more memory efficient than all-layer clipping.
Comparing to the standard ZeRO, our DP-ZeRO enjoys almost the same speed and memory efficiency, while the gap will be further closed as we move to more advanced stages of ZeRO.

Figure: Figure 5: Efficiency on ViT-Gigantic (left, 1.8B), ResNet152 (middle) and GPT2-XL (right, 1.5B).
Refer to caption: /html/2311.11822/assets/x4.png

### 4.2 Lighter training of DP-ZeRO

Figure: Figure 6: Efficiency of DP-ZeRO with lighter training. Upper: $>2\times$ speedup with lighter training; ViT-5B, $B=1$. Lower: DP-ZeRO benefits more than ZeRO from PEFT.
Refer to caption: /html/2311.11822/assets/x10.png

DP-ZeRO can employ low-memory optimizers and train on fewer parameters, therefore vastly reducing the memory and communication cost. On a single node, we demonstrate that DP-ZeRO actually benefits (more than standard ZeRO and single-GPU training) from lighter training, following from our discussion in [Section 3.2](#S3.SS2).

##### Low-memory optimizers

Low-memory optimizers can boost the training efficiency at the cost of accuracy degradation.
For example, SGD requires only 1/3 optimizer states of Adam and significantly saves the memory; 1-bit Adam (Tang et al., ) and signSGD compress the gradient and reduce the communication volume up to $32\times$.

##### Fewer trainable parameters

In the fine-tuning phase, as we analyzed in [Section 3.2.3](#S3.SS2.SSS3), PEFT improves both the local computation and the communication volume. Hence, DP-ZeRO allows PEFT on ViT and GPT to be $\approx 2\times$ faster than full fine-tuning, whereas the single GPU acceleration is $\leq 1.5\times$.

###### Remark 4.1 .

We leverage DP-ZeRO3 with SGD to train ViT-10B (full parameters) and ViT-22B (PEFT; 1M trainable parameters) on one node. See [Appendix D](#A4) for configurations.

###### Remark 4.1 .

### 4.3 Three stages of DP-ZeRO

DP-ZeRO supports all stages of ZeRO under different implementations including DeepSpeed (default) and FSDP.

In [Figure 7](#S4.F7), the efficiency of DP-ZeRO catches up with the standard ZeRO when we move up the stages. For instance of ViT-Gigantic, the throughput increases from 83% by DP-ZeRO1 to 95-97% by DP-ZeRO3. Following ([3](#S3.E3)), we can attribute the relatively fast training of DP-ZeRO to the increase cost of communication, especially in DP-ZeRO3. Additionally, we observe that the throughput of DP-ZeRO1/2 improves to over 95% on 4 nodes, as predicted by [Section 3.2.1](#S3.SS2.SSS1). Notice that
we save DP-ZeRO3 of GPT to [Section 4.4](#S4.SS4) on super-large scale.

Figure: Figure 7: Efficiency of DP-ZeRO on ViT-Gigantic and GPT2-XL under different implementations.
Refer to caption: /html/2311.11822/assets/x13.png

### 4.4 Scalability of DP-ZeRO

We evaluate the scalability of DP-ZeRO3 in terms of large sequence length (2048), large model size ($7\sim 100$B), and large number of GPUs (up to 256). We use A100 with 80GB memory, as well as the activation check-pointing and ModelP.

In [Figure 8](#S4.F8) (left), we observe that for a fixed model with 26B trainable parameters, DP-ZeRO is super-linearly scalable to the number of GPUs, achieving $>95\%$ speed of the standard ZeRO. Here super-linearity is a property of ZeRO (see Figure 3 in Rajbhandari et al. ()) which allows more GPUs to shard the model states (and reduce the per-GPU memory cost) more aggressively, and to train faster since the micro-batch size is larger. Furthermore, in [Figure 8](#S4.F8) (right), for a fixed number of GPUs, DP-ZeRO is linearly scalable to the model size, achieving the same speed as the standard ZeRO. In short, DP-ZeRO is almost equal to ZeRO in terms of training efficiency in super-large scale.

Figure: Figure 8: Scalability of DP-ZeRO3 on 26B model (left) and 256 GPUs (right). ‘max B’ means we fit the maximum micro-batch in each GPU and ‘fixed B’ means $B=2$.
Refer to caption: /html/2311.11822/assets/x15.png

###### Remark 4.2 .

In comparison to DP-ZeRO, DataP (DP or standard) at most fits 5B models (Rajbhandari et al., ) in 80GB memory, regardless of the number of GPUs. We cannot compare to DP-PipeP in He et al. () because the codebase and experiment details (e.g. number of trainable parameters and sequence length) are not publicly available. Nevertheless, since DP-ZeRO resembles the efficiency of standard ZeRO, it suffices to demonstrate the usefulness of DP-ZeRO by comparing ZeRO to PipeP.

###### Remark 4.2 .

## 5 Discussion

In this work, we develope DP-ZeRO that enables the optimization of models up to $100$B trainable parameters, thus allowing DP distributed learning to be as efficient and scalable as the standard one. We believe this is a significant milestone to pave the path towards DP foundation models, especially for its open-source nature (link to be released).

We emphasize that, since DP only modifies the back-propagation, our DP-ZeRO is orthogonal to any large-scale training techniques that are not tied to back-propagation: for example, activation check-pointing, CPU offloading, weight/activation quantization (Dettmers et al., ; Xiao et al., ), tensor parallelism (Narayanan et al., ) and other techniques yet to come.
Note that to make fair comparisons in this paper, DP and non-DP optimization are implemented without fusing the operations such as tensor multiplications. Therefore, the efficiency of DP deep learning can be further and ultimately improved, from an engineering perspective, by implementing the DP back-propagation with operator fusion via CUDA kernels and C++ coding.

## 6 Acknowledgement

We thank Dhananjay Ram, Zhenghui Jin and Cole Hawkins for technical support on the AWS cluster, Xinwei Zhang for validating part of our code, Yu-Xiang Wang for the insightful initial discussion, and the hosts of FSDP (especially Min Xu) and DeepSpeed Github repositories for debugging part of our experiments.

## Appendix A The Book-Keeping (BK) algorithm

### A.1 Efficient computation of per-sample gradient norms

The mixed ghost norm Bu et al. () is the state-of-the-art technique to compute the per-sample gradient norm of the weight, almost for free. It hybridizes two basic techniques – the per-sample gradient instantiation and the ghost norm – to compute the Frobenius norm of weight gradient,

$$ $\displaystyle\left\|\bm{a}_{l,i}^{\top}\frac{\partial L}{\partial\bm{s}_{l,i}}\right\|^{2}\overset{\text{per-sample grad}}{=\joinrel=\joinrel=\joinrel=}\|\frac{\partial L_{i}}{\partial\mathbf{W}_{l}}\|_{\text{Fro}}^{2}\overset{\text{ghost norm}}{=\joinrel=\joinrel=\joinrel=}\text{vec}\left(\frac{\partial L}{\partial\bm{s}_{l,i}}\frac{\partial L}{\partial\bm{s}_{l,i}}^{\top}\right)\cdot\text{vec}(\bm{a}_{l,i}\bm{a}_{l,i}^{\top})$ (4) $$

where ”vec” flattens the tensor to an one-dimensional vector. In words, [Equation 4](#A1.E4) gives two equations that are equivalent mathematically, but significantly different in efficiency:

- •
$\|X\|_{\text{Fro}}^{2}=\left\|A^{\top}B\right\|^{2}$ firstly computes $A^{\top}B$ and then its norm.
- •
$\|X\|_{\text{Fro}}^{2}=\text{vec}\left(AA^{\top}\right)\cdot\text{vec}(BB^{\top})$ firstly computes $AA^{\top},BB^{\top}\in\mathbb{R}^{TT}$ and then their dot product.

In summary, the mixed ghost norm always applies the cheaper of two techniques at each layer of a neural network.

Finally, we note that the per-sample gradient norm of the bias is computed differently. This is because

$$ $\frac{\partial L_{i}}{\partial b_{l}}=\mathbf{1}^{\top}\frac{\partial L}{\partial\bm{s}_{l,i}}$ $$

is not actually a product of tensors like $X=A^{\top}B$. In fact, the multiplication with $\mathbf{1}$ turns out to be a summation along the first dimension, and it suffices to use per-sample gradient instantiation for the bias.

### A.2 Book-keeping the output gradient

The BK algorithm uses two rounds of back-propagation (though each round only takes half the complexity, hence the total complexity of DP back-propagation matches the non-DP back-propagation). Therefore, output gradients $\frac{\partial L}{\partial\bm{s}_{l,i}}$ are kept to avoid repeated computation. Notice that the output gradient are relatively cheap to book-keep (see Figure 4 and Figure 5 in Bu et al. ()).

## Appendix B Component-wise time complexity of DP-ZeRO

In this section, we explain the time complexity of each part in [Table 2](#S3.T2), and demonstrate how the complexity can be different under different settings.

Forward propagation: The matrix multiplication during forward propagation results in $2BT\Psi_{\text{model}}$ complexity (see Bu et al. ()). Notice that, if the activation check-pointing is used, essentially two rounds of forward propagation take place in one iteration. Hence the time complexity doubles and becomes $4BT\Psi_{\text{model}}$.

Back-propagation: This contains two sub-processes: the output gradients are computed at all layers, taking $2BT\Psi_{\text{model}}$ complexity; the parameter gradients are computed only at trainable layers (a few if doing PEFT), taking $2BT\Psi_{\text{train}}$ complexity. Clearly, in full parameter training, the total is $4BT\Psi_{\text{model}}$, and in PEFT, about $2BT\Psi_{\text{model}}$.

Attention: The time complexity of attention is $O(BT^{2})$ in Vaswani et al. (), where $T$ is the sequence length (a.k.a. token length). When $T$ is large, e.g. training with long context like $T=8192$, this cost is prohibitively high. In this regard, a line of researches have proposed linear complexity attention, including but not limited to Wang et al. (); Katharopoulos et al. (); Shen et al. ().

Communication: For algorithms that don’t shard the model, such as data parallelism and ZeRO1/2, the communication is only used to send gradients and optimizer states. Hence the communication volume is proportional to the number of trainable parameters $O(\Psi_{\text{train}})$. Otherwise, for algorithms such as ZeRO3 and tensor parallelism, the communication volume is proportional to the number of total parameters $O(\Psi_{\text{model}})$, because the forward propagation needs to gather the parameters from many GPUs. This makes a big difference in PEFT when $\Psi_{\text{train}}\ll\Psi_{\text{model}}$.

## Appendix C Loss scaling in mixed-precision training

We write the per-sample gradient with loss scaling $S$ as

$$ $\frac{\partial C_{i}L_{i}}{\partial\mathbf{W}_{l}}=C_{i}\frac{1}{S}\cdot\left(\bm{a}_{l,i}^{\top}\big{(}S\cdot\frac{\partial L}{\partial\bm{s}_{l,i}}\big{)}\right)$ $$

This covers the standard gradient ($C_{i}=1$) and DP gradient (e.g. $C_{i}=1/||\bm{g}_{i}||$, computed by the mixed ghost norm in [Appendix A](#A1)), in which $S$ enlarges the output gradient to avoid underflow, and $\frac{1}{S}$ shrinks the parameter gradient to the correct magnitude.

Recall that a standard mixed-precision training (with loss scaling) uses steps $1\to 2\to 3\to 5\to 6$ (^7^77See [https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#lossscaling](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#lossscaling).), or $1\to 3\to 6$ without loss scaling.

- 1.
Forward propagation (fp16 weights and activations) and get the loss.
- 2.
Scaling up the loss by a factor $S$.
- 3.
Backward propagation on the scaled loss (fp16 parameters and their gradients).
- 4.
Per-sample gradient clipping (sensitivity $=1$) and noising for DP.
- 5.
Scaling down the parameter gradient by a factor of $1/S$.
- 6.
Update the parameters with their gradients.

If we follow the same procedure under the DP regime, say using a hook function to be called after back-propagation creates the gradients like in Opacus (Yousefpour et al., ), Private-Transformers (Li et al., ), FastDP (Bu et al., ), then the per-sample clipping factor is scaled up $S$ times so as to normalize the gradient. Hence per-sample gradient clipping has already played the role of scaling down. If we scale down the gradient for a second time, the gradient is incorrectly over-shrunk. This is the case in Yu et al. () and in the alternative implementation of (Li et al., , Appendix T) (see also [Figure 4](#S3.F4)). To be sure, this approach is still DP, but the performance does not match fp32 DP training correctly, and usually degrades too much to be useful.

One walk-around is to prevent per-sample gradient clipping to scale down the gradients and let step 5 do its job, i.e. $1\to 2\to 3\to 4^{*}\to 5\to 6$. We note that (Li et al., , Appendix T) follows this path (though no experiment results or codes are available at the time of writing) by modifying step 4: clipping threshold (sensitivity)$=S$ instead of $1$, so that the clipped gradient is $S$ times larger than the DP f32 training, to be scaled down by step 5. However, this introduces additional design decisions and does not prevent overflow when using fp16 (due to step 2, see [Table 3](#S3.T3)).

Another walk-around is to delete step 5 and let per-sample gradient clipping do its job, i.e. $1\to 2\to 3\to 4\to 6$. However, this approach is harder to implement because in the standard process step 2 and 5 are simultaneously enabled or disabled. Also we cannot prevent overflow when using fp16 as we still use step 2.

Therefore, we propose to not use loss scaling (or equivalently we set $S=1$ statically for all steps) during DP mixed-precision training, i.e. $1\to 3\to 4\to 6$. Although, by not using step 2, we cannot prevent underflow when using fp16, this is much less a problem compared to overflow: underflow (treating small values as 0) makes the training less accurate but does not fail the training like overflow (treating large values as NAN). Lastly, the underflow issue is perfectly mitigated by bf16, which we recommend for DP mixed-precision training whenever possible .

**Table 4: Mixed-precision training with DP or not.**
|  | steps | fp16 issue | note | reference |
| --- | --- | --- | --- | --- |
| standard | 136 | underflow |  | Micikevicius et al. () |
| standard | 12356 | none |  | Micikevicius et al. () |
| DP | 123456 | overflow | incorrect due to over-shrinking | Li et al. () |
| DP | $1234^{*}56$ | overflow | different clipping threshold | Li et al. () |
| DP | 1346 | underflow | perfect with bf16 | ours |
| DP | 12346 | overflow | hard to implement | ours |

## Appendix D Experiment settings

Datasets: To evaluate the efficiency, it suffices to declare the data’s dimension (e.g. micro-batch size and feature dimension) without specifying the dataset (though sometimes specifying the dataset means declaring the dimension, e.g. MNIST usually means 28*28 pixels). This is the norm in system papers such as Rajbhandari et al. (; ); . In this work, vision models are trained with 224*224 pixels at ImageNet scale; GPT models are trained with sequence length 100, except in [Figure 8](#S4.F8) where sequence length is 2048.

[Figure 4](#S3.F4) and [Table 3](#S3.T3): We train ViT-large (300M parameters) and CIFAR100, 5 epochs, learning rate 5e-4, logical batch size 1000.

[Figure 6](#S4.F6): To fit as large a model as possible, we set $B=1$ and use SGD. We set 48 attention heads, 21 layers, MLP=4*width (also known as embedding dimension), and modify width for all models. For instance, ViT-10B uses width=$768*22$, ViT-22B uses width=$768*34$.

[Figure 8](#S4.F8): We train AdamW with layer-wise clipping. DP distributed learning is based on MiCS (ZeRO3) using bf16 mixed-precision training. Most of GPT configuration is the same as Touvron et al. () (Table 2) in terms of embedding dimension, attention heads and number of layers. However, GPT-100B uses the configuration from Brown et al. () (Table 2.1) but a smaller width.

## Appendix E Codebase design

### E.1 With forward & backward hooks

Hooks(^8^88See [https://pytorch.org/tutorials/beginner/former_torchies/nnft_tutorial.html#forward-and-backward-function-hooks](https://pytorch.org/tutorials/beginner/former_torchies/nnft_tutorial.html#forward-and-backward-function-hooks).) are important functions to enrich the deep learning optimization. To be specific, there are

- 1.
forward modular hook (nn.register_forward_hook),
- 2.
backward modular hook (nn.register_backward_hook),
- 3.
backward tensor hook (tensor.register_hook).

DP libraries including Opacus Yousefpour et al. (), Private-transformers Li et al. (), Private-Vision Bu et al. (), FastDP Bu et al. (; ), FastGradClip Lee & Kifer () and so on, use modular hooks to modify the standard optimization. However, ZeRO libraries including DeepSpeed and FSDP use tensor hooks. This difference in the types of hooks and many other differences (e.g. both ZeRO libraries and DP libraries modify the optimizer’s step function) cause non-trivial problems when combining DP with ZeRO. For example, to keep DP optimization as efficient as the standard, it is necessary to not waste time on computing the non-private gradient. However, if we skip such computation, then ZeRO’s tensor hook will not be triggered and the corresponding distributed-learning-related operations cannot carry on. For another example, because DP and ZeRO add different types of hooks, the number of hooks is larger than either optimization and they slows down the training: consider an 100-layer network, each layer with weight and bias (2 tensors), then DP-ZeRO in this subsection needs 100 modular hooks and 200 tensor hooks, adding to a total of 300 hooks. In addition, the Book-Keeping algorithm (in FastDP) in its original form cannot be implemented together with ZeRO3, because all model states are partitioned including the output gradients which are meant to be book-kept. To work around this requires rewriting the distributed solution’s communication mechanism, and if successful, still requiring additional communication cost during the second back-propagation. Similar problems are present for Opacus and FastGradClip, which instantiates per-sample gradients that will be partitioned in ZeRO2/3 and requires additional communication cost when gathered to create the privatized gradient.

As a consequence, the hooks are fully supported on DP-ZeRO1 and partially supported on DP-ZeRO2/3 under the layer-wise clipping.

### E.2 Without hooks

Instead of registering hooks on top of the original (non-DP) back-propagation, we can directly modify the back-propagation following [Appendix A](#A1): e.g., given the activation and output gradient,

$$ $\frac{\partial C_{i}L_{i}}{\partial\mathbf{W}_{l}}=\bm{a}_{(l),i}^{\top}\frac{\partial L}{\partial\bm{s}_{(l),i}}\Big{/}\sqrt{\text{vec}\left(\frac{\partial L}{\partial\bm{s}_{(l),i}}\frac{\partial L}{\partial\bm{s}_{(l),i}}^{\top}\right)\cdot\text{vec}(\bm{a}_{(l),i}\bm{a}_{(l),i}^{\top})}$ $$

This approach requires rewriting the back-propagation for each layer type (linear, embedding, convolution, normalization, …) and can be done at different levels (Pytorch, C++, CUDA kernel).

### E.3 User interface

DP-ZeRO can be enabled by one piece of code: after the model is instantiated,

The codebase is designed not to modify the optimizer, hence DP-ZeRO can work with arbitrary optimizer. Because of this design, our DP-ZeRO will not distinguish micro-batches. This is different from the gradient accumulation in Opacus (version ==0.x) and Private-Vision, where only the last micro-batch is processed by ”optimizer.step()” but all other micro-batches are processed by ”optimizer.virtual_step()”. In other words, the noise $\sigma_{\text{DP}}N(0,I)$ is added on the last micro-batch, after the micro-batches are accumulated. But DP-ZeRO adds the noise on each micro-batch equally. Note that the noise level per micro-batch is $\sigma_{\text{DP}}/\sqrt{N_{d}}$ if a random seed is set across $N_{d}$ GPUs, or $\sigma_{\text{DP}}/N_{d}$ otherwise.

## Peer Reviews

confidence: 3
contribution: 1
initial_rating: 1
presentation: 1
questions: See weaknesses.
soundness: 1
strengths: - The paper tackles an important problem of private training of large models.
summary: This work presents a new differentially private version of the Zero Redundancy (ZeRO) optimizer called DP-ZeRO. ZeRO is an optimization technique to improve the efficiency of training very large models. The paper's main contributions are
- DP-ZeRO has the same level of communication efficiency, computation efficiency and scalability as standard ZeRO.
- Enables training 1B+ parameter models with DP for the first time.
- Codebase to run training that will be open sourced.
- Discussion around loss scaling and clipping.

The paper goes on to describe the ZeRO optimizer as well as the overhead of DP-SGD which requires per example clipping. Then, the paper describes that DP-ZeRO is basically the same as ZeRO except for the backpropagation step as it requires per sample gradient norms to scale the per example gradient. However, the authors do not clarify what is required to overcome this challenge and only say that it has a small cost. The authors also explain that many other techniques like PEFT, low memory optimizers and activation rematerialization are compatible with DP-ZeRO. Then there is an experimental section describing some experiments where the authors show that there is a minimal overhead of running DP-ZeRO compared to ZeRO and linear scaling.
weaknesses: - It's not clear what novel algorithm is presented in the paper and what it actually does. Can the authors please clarify?
- Usually, for large models that require model parallelism, either very small per core batch sizes are used (such as batch size 1) and/or PEFT is used which means that the cost of per example gradients and per example gradient norms is usually is small. 
- The larger practical challenge of private training are the use of very large batch sizes which is not mentioned in the paper.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 1
questions: - I don't quite understand the discussion and what is illustrated in the Figure 2. Why is there different accuracy between the left two but not between the right two? Is this illustration somehow representative of the proposed partitioning scheme? And if so, which part of the illustration?
- I'm also bit confused on how the clipping is actually performed. You mention in Section 2.1 the per-group clipping and in Section 3.1. you state that the clipping is done locally on each GPU. Does this mean that the parameter groups are now split amongst the GPUs? It would help to follow the discussion if authors could summarize the procedure. I guess Figure 3 is meant to summarize it, but I'm struggling to see the different groups from this figure.
- I guess you apply the clipping per group, so maybe in the discussion of "Per-sample gradient clipping" you could use $\mathbf{q}_{[m],i}$ in the denominators?
- In "Optimizer state partition" you discuss the memory costs. Can you please clarify where does the $16\Psi$ for the basic DataP come from?
- When you say (e.g. in the beginning of page 5) "$2\Psi_{model}$ fp16 parameters", do you actually mean "fp16 parameters that take $2\Psi_{model}$ bytes of memory"?
- At the end of Page 5 the clipping bound does not appear in the noise scaling. Are you somehow assuming that the maximum norm of the gradient is $1$? If so please be explicit about it.
- Are the numerator and denominator flipped in eq. (3)?

typos
- Page 4: "... states. and ..."
soundness: 2
strengths: Fine-tuning large pre-trained models has proved to be an appealing avenue for mitigating the utility hit caused by DP training. The proposed pipeline extends the modern state-of-the-art DP fine tuning to very large models, which have been difficult (or even impossible) to train previously due to additional computation costs by per-example gradient computation. Therefore the proposed method is an promising direction towards DP training that benefits from the vast information condenced in the large models.

Empirical evaluation demonstrates a computational efficiency comparable to the non-DP ZeRO for many different architectures. Furthermore, authors also demonstrate high scalability of the method w.r.t. the number of available GPUs.
summary: This paper presents a new efficient pipeline for distributed DP optimization or large deep learning models. Authors combine the existing ZeRO distributed pipeline with mixed precision and DP-SGD. This allows the DP deep learning to be extended to modern large models with (hundreds of) billions of parameters. Authors demonstrate empirically, that the throughput between the non-DP and DP pipelines is rather comparable, suggesting that the additional computation cost of DP can be rather moderate.
weaknesses: While I can believe the argument, that existing distributed DP methods scale somewhat poorly, I think authors should still compare the throughput of the proposed method to some existing DDP baselines. For example, how would the throughput compare against DataP? Also, in Remark 4.2, authors state that the "the usefulness of DP-ZeRO can be demonstrated by comparing ZeRO to PipeP.", but I don't see this comparison nowhere.

The clarity of the paper could be improved somewhat significantly. I think some summary of the full pipeline would help the readibility a lot. Currently the components of the proposed pipeline are spread across the paper, which makes assessing e.g. the privacy properties a bit difficult. Also, I don't forward/backward propagation discussion in Section 3.1. I guess this is not novel analysis of this work, and I don't at its current form understand its purpose. I will list some further questions regarding the clarity in the next section.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - Q1: Why are you not fully utilizing the memory of the GPU even though this should improve bring speedups? Experiments regarding Figures 4, 6 and 7 do not use max B but some fixed B. 

- Q2: Contribution 4: *"We will open-source our codebase as the first DP distributed learning library"*: You name this as a contribution but the anonymous git just has code for fastDP or is it somehow integrated into each other? This is not clear to me from the README.

- Q3: How does the scaling of your method compare to naively scaling with DDP as e.g., illustrated in Beltran et al. [2] Figure 7 for smaller models? Is your method then also benefical or does it only worth for very large models?

- Q4: Line 155: *"These techniques improve the relative speed of DP-ZeRO but worsens the absolute speed."*: It is unclear to me what this means. Could you clarify please?

- Q5: Section 3.2.3: You write that "local computation can accelerate by 50%" as can be seen in Table 2. When treating training_params as almost zero I arrive at  PEFT: $O(BT^2)+4BT \Phi_\text{model}$ which is more than half of full training: $O(BT^2)+6.666BT \Phi_\text{model}$. What is my misunderstanding here because I cannot clearly see it?

- Q6: Line 360: *"a large number of GPUs so that the micro-batch size B (i.e. per-GPU batch size) is small"*: Recent work [3, 4, 5] suggests that large (logical) batch sizes are benefical under DP (in comparison in non-DP they are not), how does this influence the relative cost of DP here.

- Q7: Lines 461-464: Is it possible that there could be a higher speedup with PEFT when fully utilizing the GPU memory? In Figure 6 only 25% of the GPU memory are used for PEFT, but more memory should result in more speed-up.


[2] Beltran, S. R., Tobaben, M., Loppi, N., & Honkela, A. (2024). Towards Efficient and Scalable Training of Differentially Private Deep Learning. arXiv:2406.17298.

[3] De, S., Berrada, L., Hayes, J., Smith, S. L., & Balle, B. (2022). Unlocking high-accuracy differentially private image classification through scale.arXiv:2204.13650.

[4] Räisä, O., Jälkö, J., & Honkela, A. Subsampling is not Magic: Why Large Batch Sizes Work for Differentially Private Stochastic Optimisation. ICML 2024.

[5] Ponomareva, N., Hazimeh, H., Kurakin, A., Xu, Z., Denison, C., McMahan, H. B., ... & Thakurta, A. G. (2023). How to dp-fy ml: A practical guide to machine learning with differential privacy. JAIR, 77, 1113-1201.
soundness: 2
strengths: - S1: Introduces a method to train very large models under DP (e.g., GPT-100B or ViT-10B) which is very timely even with PEFT as models grow a lot.
- S2: There are analysis both empirically as well as in terms of time complexity.
- S3: The authors make the method work (e.g., address problems with loss scaling and integrate it with many clipping methods).
- S4: The authors promise to release a framework, although the linked anonymous git is not very structured (See Q2).
summary: The authors propose a DP version of Zero Redundancy Optimizer (ZeRO) that enables the training of very large models under DP. They compare the non-DP and variants in different settings and observe a good level of communication and computation efficiency as well as scalability. Additionally, a core contribution is the enabling of mixed-precision training through appropriate handling of the loss. These advances enable the authors to train very large models under DP for the first time (e.g., GPT-100B or ViT-10B). The final contribution is the promised release of the codebase as a framework to allows to use DP-ZeRO for many applications.
weaknesses: While I believe that the paper has good contributions there are some things that need to be addressed.

- W1: Contribution 3: *"We enable DP deep learning with more than 1B trainable parameters for the first time."*: This claim is very strong and I am not sure if it is correct as Ghost Clipping can train all parameters of GPT-2-large which is smaller than a billion but not too far from it with a smaller GPU (24 GB vs. 40 GB) easily (physical batch size of 10) [see Figure 4a of Li et al. (2022) [1]]. Another recent work trains a ResNet of almost one billion parameters with out-of-the-box Opacus on 40 GB VRAM [see Figure 3b of Beltran et al. (2024) [2]]. I would suggest to the authors to make their claim more precise as their contributions are good, but this claim goes too far and might not be true. You also mention also in Line 072 that 2 Billion is possible without your contribution. Something like "We enable DP deep learning with **far more** than 1B trainable parameters for the first time." might be a better choice.

- W2: Figure 1: This suggests that there are no other works that achieved this large trainable parameters. Unfortunately, I was unable to find citations to these works so it is hard to say how comprehensive the list is. I would recommend adding appropriate citations (perhaps in the Appendix) and perhaps also some numbers from recent papers like [2] that trained ViT-Huge.

- W3: Experimental design regarding utilizing the GPUs: It is unclear to me why the authors chose to not utilize the full VRAM of the GPUs in their experiments (Figure 4, 6 and 7) when evaluating the methods. Their experiment in Figure 8 compares between maximum B and fixed B=2 and there seems to be a gap. Similarly, related work [2] has also noted that maximum possible GPU memory usage improves the speed. (See Question 1 as well).

- W4: Presentation of Figures: The captions do not have comprehensive information about what is shown (e.g., Figure 4 and 6) and at least Figure 1 is never mentioned in text. I think it would be easier to grasp Figures 2, 6 and 8 using subfigures instead of this very compressed format that impairs the understanding. 

- W5: The usage of $\Psi$ is a bit confusing at first. I recommend adding the unit bytes to it (if I understood it correctly). Also it is confusing that the total $\Psi_\text{model} = 16$ if there are 3 components (line 197)?

Minor:
- Line 068: *"DDP cannot train models that exceed the capacity of one GPU"*: Would be good to make this clearer as this refers to cases where the model weights do not fit anymore.

- Line 128: *"In DP deep learning, the gradients are made private by **post-processing** through per-sample gradient clipping and random noising"*. Post-processing is not the correct term here as post-processing in DP usually refers to the property that something being DP cannot be reversed. I recommend rephrasing this.

- Line 155: *"With $\sigma_\text{DP} = 0$ leading to $\epsilon = 0$ (not private)"*: It would be good to emphasize that this is only in regard to the noise, but clipping is still different than non-DP.

- Line 192: The sentence seems broken *"optimizer states. and"*.

- Line 341: *"These techniques improve the relative speed of DP-ZeRO but worsens the absolute speed."*: It is unclear what this means. Please rephrase. 

- Line 456: This implies that SGD leads to worse accuracy than Adam which is not universially true. I would recommend reprashing this.

[1] Li, X., Tramer, F., Liang, P., & Hashimoto, T. Large Language Models Can Be Strong Differentially Private Learners. In ICLR 2022.

[2] Beltran, S. R., Tobaben, M., Loppi, N., & Honkela, A. (2024). Towards Efficient and Scalable Training of Differentially Private Deep Learning. arXiv:2406.17298.

---

confidence: 2
contribution: 1
initial_rating: 3
presentation: 1
questions: 1. In lines 12-13 you mentioned the communication efficiency challnage of training deep networks. As you do not specify at this point that distributed settings are considered, it is not clear what communication has in regular (fully centralized) deep learning.
2. In line 8 you use the symbol $\epsilon$. It is not clear at this stage that this symbols the privacy budget rather than something else.
3. In line 64 you mentioned that DDP with DP usually incurs huge memory cost due to caching the per-sample gradients, isn't it because of caching a full copy of the DP model? 
4. In lines 72-73: "...which roughly translates to 2B model training with Adam" what is 'B' stands for?
5. There is no referencing in the paper text for Figure 1.
6. As ZeRO is not reviewed on its own, is it not clear what 'model states' means in line 187.
7. In line 192: "a master copy (fp32) of optimizer states. and the half-precision...". What does 'fp' means? "states. and" is probably a typo.
8. In lines 196-198, where does the 4, 16, and (4+...) come from? not clear or intuitive.
9. In line 222: "...on the half precision (fp16 or bf16) parameters" what is 'bf'?
10. In equation (3) lines 308-309, it seems that that the numerator and denominator of the fraction are reversed.
11. In lines 484-485: "Here super-linearity is a property of ZeRO (see Figure 3 in Rajbhandari et al. (2020))..." quite not easy-to-follow; forcing the reader to jiggle many times to Rajbhandari et al. (2020) paper to understand this paper.
soundness: 1
strengths: Significance:
- The paper studies DP in the challenging distributed learning case, where the model size increases beyond a reasonable bound for one GPU.
- The model can be partitioned into shards via, e.g., pipeline parallelism and model parallelism. Yet, the pipeline parallelism can be inefficient
due to a non-DP-related issue – the pipeline bubble, where GPUs are idle while waiting for data to process; and the authors are focusing on the model parallelism alternative, by adopting the ZeRO optimizer.
summary: This paper incorporates DP into the established ZeRO optimizer, used in distributed learning (via, e.g., multiple GPUs) which especially relaxes the burden of training large models. In such cases, where the model size increases beyond a reasonable bound for one GPU, the model must be partitioned (in addition to the data, as in conventional distributed learning), so that each GPU only holds a partial shard of the model. The authors specifically apply DP with the mixed-precision training, as adopted in ZeRO; and are thus able to experimentally show that the incorporation of DP into ZeRO does not harm the latter's performance.
weaknesses: Originality:
-  The claimed contribution lacks novelty. To the best of my understanding, the proposed solution is to incorporate the Gaussian mechanism into the ZeRO optimizer (in which mixed-precision is used). Yet, the idea of incorporating the Gaussian mechanism into pipeline parallelism (instead of model parallelism) was already learned in He et al. (2022)

Clarity:
- The ZeRO optimizer, which is a main tool in you method, it not properly covered, for instance in a dedicated section of its own. This makes the paper difficult to follow as a stand-alone work. 
- The threat model is not mentioned: what is the quantity that should be preserved as private, and why, and from whom; all are not clarified. This can by implied from the context and from the motivation for studying the work of He et al. (2022), but is not included in the presented paper, and therefore not well-clarified.