FlashDP: Memory-Efficient and High-Throughput DP-SGD Training for Large Language Models

Thank you for your thorough review and the positive remarks on our work.

For the weaknesses you mentioned, regarding the use of PyTorch, we opted for this framework due to its widespread adoption in both private and non-private training contexts. While our implementation is currently PyTorch-based, the core modifications we made to the CUDA code, which underlies both PyTorch and JAX, suggest that our method could indeed be extended to JAX. Adapting FlashDP to JAX would entail substantial work, including rewriting the model and submodule code for JAX and adjusting the underlying source code. We recognize the potential benefits of such an adaptation and plan to explore this in future work.

As for your question about the losses reported in Table 2, they are indeed the same. FlashDP is designed to optimize the efficiency of GPU utilization without altering the fundamental computational processes of DP-SGD.

Thanks to the authors for their response. While there is no further need to address the concerns in my initial review, reviewer 1eKM has raised some concerns about the presentation surrounding whether the current method vs. past work applies to per-layer vs flat gradients which I am in agreement with. In particular, I was under the impression FlashDP used flat clipping which made the comparisons to BK and GhostClip much more impressive. Now, as the other reviewers have pointed out, it is unclear how much of the improvement in the comparisons is because of the different end goals of flat clipping vs per-layer clipping, and how much is because of any novelties in FlashDP.

For now I am reducing my score in light of the concerns raised in reviewer 1eKM's review and the resulting discussion. I would be happy to continue discussing in the followups to reviewer 1eKM's review to keep the discussion in one place, and happy to revisit my score if the discussion there can alleviate the reviewers' shared concerns. I apologize for the misunderstanding in my initial review and for reducing the score late in the review process.

Thank you for your insightful comments and questions. We appreciate your acknowledgment of the importance of our work and its potential benefits. Below, we address your points to clarify the aspects of our method:

For the weaknesses you mentioned:

1. We explain the background on GPUs, CUDA, and chunked matrix multiplication in detail in Appendix Section A.3. Let's explain it here in layman's terms: the memory in a GPU can be roughly divided into two parts, the large but relatively slow High Bandwidth Memory (HBM), and the small but fast Static Random Access Memory (SRAM), which we usually refer to as the GPU memory. We can think of the per-sample gradient as an intermediate variable in the actual gradient computation, and Opacus stores this complete intermediate variable in the HBM, while we store this intermediate variable in chunks in SRAM when we compute the matrix multiplication in chunks. In this way, we both reduce the utilization of the HBM and increase the speed.

2. Our algorithm does not apply only to the MLP layer, but to the more basic Linear layer. almost all of the parameters in the transformer come from the Linear layer, including the attention as well as the MLP. this is explained in detail in the original Appendix Section A.2.

3. Technical Terminologies: The GEMM (General Matrix Multiply) and HBM (High Bandwidth Memory) are usually used in systems areas. Due to page constraints, we have explained these terms in the original Appendix Section A.3.

For the questions you mentioned:

1. For the convolution operation, although it is not discussed in this paper due to space constraints, it can be easily transformed into a GEMM operation by a simple transformation. [1] and [2]

[1] https://inria.hal.science/inria-00112631/document [2] https://arxiv.org/abs/1704.04428

2. We have implemented it as a package and provided our complete code and detailed documentation in an anonymous repository, accessible here: FlashDP4Review.

3. The observed difference in memory usage between GhostClip and NonDP implementations with PyTorch can primarily be attributed to the underlying memory management mechanisms. GhostClip implements a strategy where it explicitly deletes the accumulated gradient to conserve memory, which can be observed in their code implementation. This approach contrasts with NonDP (PyTorch), which employs a memory pool strategy to optimize allocation time during runtime. In the NonDP setup, used memory, such as the accumulated gradient, is placed into a memory pool. This means that once a computation requires more memory, PyTorch allocates it from the memory pool instead of performing new allocations. This approach speeds up memory allocation but can lead to slightly higher memory usage than what is strictly necessary for the computation at hand, as documented in the PyTorch documentation.

Thank you for your valuable feedback and detailed review. We appreciate the opportunity to clarify and improve our submission based on your insights.

For the major weaknesses you mentioned:

1. Code Availability for Reproducibility: We acknowledge the critical importance of reproducibility in research. To address this, we have provided our complete implementation and detailed documentation in an anonymous repository, accessible here: FlashDP4Review.

2. Gradient Clipping Methodology: In our implementation, we adopt per-layer clipping, which is increasingly recognized as a standard practice in differential privacy for large language models (LLMs). This approach has been shown to effectively match or even surpass the state-of-the-art in various language and vision tasks, as documented in the following studies: [1], [2], [3]. These references demonstrate that per-layer clipping can be seamlessly integrated into existing codebases with minimal modifications. Importantly, this method achieves compliance with differential privacy guarantees while maintaining the performance benefits of full gradient clipping but with significantly reduced memory overhead and computational complexity. [4] Recent papers also have shown that per-layer clipping could have comparable performance with the original full gradient clipping.

[1] https://arxiv.org/abs/2110.05679 [2] https://arxiv.org/abs/2210.00038 [3] https://arxiv.org/abs/2206.07136 [4] https://arxiv.org/abs/2212.01539

For the minor weaknesses you mentioned:

1. Exclusion of Methods in Figures: The absence of GhostClip and BK methods from certain figures was due to their incompatibility with the specific experimental setups tested. For example, in Figure 4, the GhostClip method does not support the Llama model structure, and in Figure 5, the BK method does not support pipeline parallelism. Please refer to the documentation private-transformers and fast-differential-privacy for more details. We have clarified these exclusions and ensured detailed justifications in Appendix Section D of the revised version.

2. Algorithm Execution Details: The "for loops" in Algorithm 1 (lines 3, 4, and 5) are designed to execute in parallel, not sequentially. We have revised the paper.

3. Model Configuration: We used a model configuration consistent with OpenLlama's 3B model, as detailed here: OpenLlama-3B.

Thank you for your response. Some of my concerns have been addressed. However, I would like to raise some remaining major concerns:

Regarding per-layer clipping:

we adopt per-layer clipping, which is increasingly recognized as a standard practice in differential privacy for large language models (LLMs). This approach has been shown to effectively match or even surpass the state-of-the-art in various language and vision tasks, as documented in the following studies: [1], [2], [3].

I have carefully re-examined papers [1,2] specifically looking for experimental comparisons or at least mentions of per-layer clipping, but I could not find any such references. Moreover, as I mentioned in my initial review, these algorithms seem to be specifically designed to address the challenges of global flat clipping ([1] performs backward pass twice, while [2] additionally stores gradients with respect to the outputs of the all linear layers, and all these overheads just to compute global per-example gradient norm). To substantiate the claim that "per-layer clipping is documented in the following studies," I would kindly request the authors to explicitly point out the specific pages and paragraphs in these papers where per-layer clipping is discussed.

 [4] Recent papers also have shown that per-layer clipping could have comparable performance with the original full gradient clipping.

While this work indeed addresses per-layer clipping and appears significantly more relevant to the authors' problem formulation than [1,2], I notice that the authors neither compare their approach with the implementation from [4] nor mention this work in their paper. I would appreciate if the authors could explain why this relevant research is not referenced in their study.

To summarize, I have a major concern regarding the validity and fairness of the experimental setup. Specifically, as I understand:

• The authors' main baselines are methods from [1,2], which were developed and designed to effectively solve a different problem (flat clipping instead of per-layer clipping).

• Meanwhile, other effective algorithms, such as [4], which are specifically designed for efficient implementation of per-layer clipping and would provide a more appropriate comparison, are not included in the paper.

Also, in my opinion, Differentially Private optimizers with flat clipping and with per-layer clipping represent distinct approaches. Therefore, the fact that the authors specifically focus on per-layer clipping should be explicitly stated in the paper to avoid potential misinterpretation.

Regarding absence of GhostClip and BK methods from certain figures

I would again kindly request that the authors explicitly point to the specific location in the documentation of these frameworks where the mentioned limitations are stated. For instance, I could not find any substring "llama" in the README of repository [5].

[1] https://arxiv.org/abs/2110.05679 

[2] https://arxiv.org/abs/2210.00038 

[3] https://arxiv.org/abs/2206.07136 

[4] https://arxiv.org/abs/2212.01539

[5] private-transformers

I'm also not the original reviewer, and I would like to move my thread to here as well for better context. After reading this thread and the revised text, I now understand that the method is doing per-layer gradient. I have related responses and concerns:

1. "This indicates that their clipping is confined to individual layers, aligning with the concept of per-layer clipping" is actually not true. They just calculate the per-layer norm one by one and the norm for the full gradient is calculated later on for the full-gradient clipping.
2. It is not clear to me that the method proposed in this paper is really benefitting from the per-layer gradient or the advanced memory-level usage, because per-layer gradient itself can save memory: similar to ghost-clipping, in their first backward when calculating the per-layer norm for a layer, they can at the same time calculate the clipped gradient for this layer and release the per-sample gradient for this layer. Or in the equivalent words, it seems not fair to compare the proposed algorithm with the full-gradient clipping methods, and given the current comparison in the paper, it is unclear how much memory saving is from the advanced memory-level implementation.

Upon further reflection, we acknowledge that our initial selection of baseline comparisons may not have been ideally matched, and including work [4] would indeed provide a more relevant comparison than references [1] and [2]. However, we have discovered that slight modifications to the FlashDP algorithm can also support flat clipping. Now we are working on modifying the FlashDP algorithm to support both per-layer and flat-clipping techniques. We are actively conducting experimental validations of these adjustments. We intend to update our manuscript with the revised algorithm and the corresponding experimental results if time permits.

We greatly appreciate the opportunity to engage in this dialogue with you. We are committed to refining FlashDP and making a meaningful contribution to the differential privacy community. If you have any further questions or require additional information about FlashDP and its developments, we are more than happy to continue this conversation and provide the necessary details.

I thank the authors for their detailed response. Their specific references to the relevant sections in articles [1,2] have confirmed that we are indeed addressing the same material in our discussion.

The key question from which all my current major concerns stem is the following: "Are the implementations of FlashDP, GhostClip[1], and BK[2] mathematically equivalent?" To avoid potential misunderstandings, I would like to present my current understanding of the situation systematically and in detail.

1. There are two different clipping strategies in the field of DP optimization: 1. Flat clipping and 2. Group-wise clipping. The difference between these two strategies is described in [3], Section 2, page 3, paragraphs 3, 4. Specifically, in the case of Flat clipping, there exists a single group of parameters $\theta$ and threshold $C$, while for Group-wise clipping, the parameters are divided into $K$ disjoint sets {$\theta_k$}$_{k=1}^K$, each with its own clipping threshold $C_k$.
2. These two strategies are not mathematically equivalent when $K > 1$. By this, I mean that for any value $C>0$ for Flat clipping and any values $C_k>0, k=1,\dots,K,$ there exists a vector $g$ such that Flat clipping and Group-wise clipping will return different results $g_{\text{flat}} \neq g_{\text{group}}.$
3. Per-layer clipping is a special case of Group-wise clipping, where $K$ equals the number of layers $L$ and each set $\theta_k$ is some instance of torch.nn.Parameter $W_l$ (or its equivalent in another deep learning framework).
4. FlashDP supports only per-layer clipping.
5. GhostClip (Section 4 in [1]) and BK (Algorithm 1 in [2]) implement Flat clipping. I draw this conclusion based on the following considerations: For GhostClip, from Section 4.1 paragraph 2 in [1] where the authors describe how to calculate $\|\nabla \mathcal{L}_i\|_2$, i.e., the full gradient norm, using per-layer gradient norms $\{\|\nabla_W^{(1)}\mathcal{L}_i\|_2,\dots, \|\nabla_W^{(L)}\mathcal{L}_i\|_2\}$ and also from Algorithm 2 in Appendix C in [2] (page 15); For BK, from lines 3-7 of Algorithm 1 in [2], where they compute the value of $C_i$, which requires knowledge of $\|\frac{\partial \mathcal{L}_i}{\partial W}\|_F^2$, i.e., the full gradient norm.

If the sequence of statements above is correct, then it follows that the FlashDP algorithm is not mathematically equivalent to GhostClip and BK.

If this is the case, I return to the concerns I mentioned above. Specifically, I think that:

1. since FlashDP and its baselines implement different mathematical functions, the experimental comparisons are not fair and do not allow for drawing conclusions about the efficiency of FlashDP.
2. the effectiveness of FlashDP should be verified through comparisons with other implementations of per-layer clipping, such as the per-layer clipping from [3] (see Figure 1).

[1] https://arxiv.org/abs/2110.05679

[2] https://arxiv.org/abs/2210.00038

[3] https://arxiv.org/abs/2212.01539

Thank you for your constructive feedback. We appreciate the opportunity to clarify and further explain our approach and findings.

For the weaknesses you mentioned:

Reducing memory storage or accelerating DP-SGD has been a hot topic, with contributions like Ghost Clipping and DP-JAX, some of which have appeared at security conferences, such as [1], now integrated into Opacus. Our system-level approach is relatively new to the security community, making this research impactful for a security-focused conference.

[1] Scaling up Differentially Private Deep Learning with Fast Per-Example Gradient Clipping (PoPETS 2021)

We would like to highlight that we are the first to integrate several CUDA-related techniques, such as intra- and inter-block shuffle reduction and atomic add, into DP-SGD. These innovations significantly reduce the number of HBM loads and stores and improve the efficiency of DP-SGD, which has historically been bottlenecked by memory movement rather than computation. While these techniques have long been established, adapting them to DP-SGD presents crucial challenges, which we have effectively addressed in our work. **By our approach, for the first time, we make it possible to pre-train a Llama-13B model using DP-SGD with 4 A100 GPU. **

For your questions:

1. Regarding Figure 1: The buffer in Figure 1 is related to GPU High Bandwidth Memory (HBM). By reducing the frequency of HBM accesses, we significantly decrease the latency and overhead associated with memory operations, thus improving throughput.

2. Clarification of subscript T: In our manuscript, 'T' denotes the sequence length, as described in Section 3.

3. Matching the performance of DP-SGD: To address the inherent variability introduced by differential privacy, we ensure consistency in our experiments by using the same seeds to generate noise across comparative runs.

4. In our implementation, we employ subsampling without replacement rather than Poisson subsampling. We opt for fixed-size mini batches because they align well with available GPU memory configurations, which enhances learning throughput. This approach is well-supported by existing literature that investigates the associated privacy budgets, such as the studies found in [1] and [2]. Using subsampling without replacement is a standard practice in scenarios like ours where consistent batch sizes contribute to more efficient computation.

[1] https://arxiv.org/pdf/2408.10456 [2] https://arxiv.org/pdf/1808.00087