Zero redundancy distributed learning with differential privacy

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.

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.

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.

• 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.