Title: Private Training Large-scale Models with Efficient DP-SGD
Venue: NeurIPS 2025
Accepted: yes
Keywords: Differential Privacy

Large Language Models

DP-SGD
Abstract: As large language models (LLMs) increasingly underpin technological advancements, the privacy of their training data emerges as a critical concern. Differential Privacy (DP) serves as a rigorous mechanism to protect this data, yet its integration via Differentially Private Stochastic Gradient Descent (DP-SGD) introduces substantial challenges, primarily due to the complexities of per-sample gradient clipping. Current explicit methods, such as Opacus, necessitate extensive storage for per-sample gradients, significantly inflating memory requirements. Conversely, implicit methods like GhostClip reduce storage needs by recalculating gradients multiple times, which leads to inefficiencies due to redundant computations. This paper introduces FlashDP, an innovative cache-friendly per-layer DP-SGD that consolidates necessary operations into a single task, calculating gradients only once in a fused manner. This approach not only diminishes memory movement by up to 50\% but also cuts down redundant computations by 20\%, compared to previous methods. Consequently, FlashDP does not increase memory demands and achieves a 90\% throughput compared to the Non-DP method on a four-A100 system during the pre-training of the Llama-13B model, while maintaining parity with standard per-layer clipped DP-SGD in terms of accuracy. These advancements establish FlashDP as a pivotal development for efficient and privacy-preserving training of LLMs. FlashDP's code has been open-sourced in https://github.com/kaustpradalab/flashdp.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces FlashDP, a system designed to enhance the efficiency of training LLMs with per-layer DP-SGD. FlashDP proposes a hardware-software co-design approach where key DP-SGD operations are consolidated into a single, cache-friendly computational pass, calculating gradients only once in a fused manner to optimize GPU operations. The empirical results demonstrate substantial improvements in throughput and memory usage for DP pre-training tasks compared to existing methods.\n\nSTRENGTHS AND WEAKNESSES:\n**Disclaimer**: My background is primarily in DP, and I have limited expertise in system-level optimization or GPU architecture. Therefore, this review focuses on the paper's contributions from a DP practitioner/researcher's perspective and may not fully capture or critically assess the technical novelty of the system-level innovations.\n\n**Strengths**\n- The paper presents FlashDP as a result of hardware-software co-design, introducing novel algorithmic ideas like Hierarchical Reduction Architecture (HRA) and Block-wise All-Reduce, which aim to efficiently calculate per-layer gradient norms with minimal data movement by using on-chip memory for local reductions before global aggregation. It further details how these ideas can be implemented effectively on current GPUs, addressing CUDA limitations through an adaptive kernel approach.\n- FlashDP demonstrates strong empirical performance (in terms of throughput and memory efficiency) across different experimental setups compared to baselines like Opacus, GhostClip, and BK for per-layer clipped DP-SGD.\n\n**Weaknesses**\n- The paper showcases FlashDP's effectiveness in DP pre-training for LLMs. While technically challenging, DP pre-training, especially when using largely public datasets, is often less of a priority and concern for DP practitioners compared to DP fine-tuning on private, sensitive data. Common DP fine-tuning paradigms such as DP-LoRA are not explored. Therefore, the focus of this paper does not align with the most pressing practical needs and common workflows within the DP research community.\n- FlashDP's design and optimizations focus on per-layer clipping. While the authors support this choice with existing literature, global clipping is generally more prevalent. Adapting FlashDP for global clipping appears challenging, as it would require all per-sample gradients for every layer to be available before the clipping operation can occur. \n- While the throughput gains are impressive and clearly demonstrated, the paper does not report on end-to-end wall-time for representative training tasks. Practitioners often rely on wall-time to gauge the actual time and cost savings. Without this, it's harder to fully assess the practical efficiency gains.\n- FlashDP's performance relies on deep integration with CUDA and employs an adaptive multi-kernel approach to manage GPU synchronization. While this addresses the limitations of existing system, it might pose challenges for portability to other hardware or even future CUDA versions. Furthermore, the inherent complexity of such low-level optimizations could make FlashDP harder to adopt, maintain, or modify for researchers lacking specialized GPU programming expertise.\n\n**Minor points**\n- Line 39-40: none of the references mentioned are about DP pre-training. Cite DP BERT instead\n- Some word choices are uncommon for research papers: Line 62 \"prolonging\" --> \"increasing\"; Line 181 \"exemplar\" --> \"example\"\n- I think a brief discussion on the sampling strategies, in particular the choice between shuffle-based sampling and Poisson subsampling, is important for a system that supports large-scale DP training \n- The apostrophes appear to be incorrectly rendered (e.g., 'FlashDPâA˘ Zs', 'PyTorchâA˘ Zs', 'modelâA˘ Zs')\n\nQUESTIONS:\n- The name \"FlashDP\" is reminiscent of \"FlashAttention\", and both appear to achieve significant performance gains through I/O-aware designs and kernel fusion. Could you elaborate on any conceptual similarities or inspirations drawn from FlashAttention in the design of FlashDP?\n- From a practical standpoint, how do you envision FlashDP benefiting DP practitioners in their common workflows? Do you foresee it being developed into a standalone, user-friendly library, or do you anticipate its core optimizations being integrated into existing, widely-used frameworks such as Opacus or DP-Transformers to enhance their performance?\n\nLIMITATIONS:\nThe authors didn't discuss the limitations in the main text, though they provide a brief discussions in the checklist (focus on per-layer clipping rather than global clipping, and reliance on fused CUDA kernels)", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the problem of efficiently training large-scale models with Differential Privacy (DP). The authors propose FlashDP, a cache-friendly per-layer DP-SGD that achieves computational advantages over existing methods. The authors show that FlashDP achieves a higher throughput compared to the Non-DP method.\n\nSTRENGTHS AND WEAKNESSES:\n* Achieves noteworthy computational advantages over existing methods;\n\n* Differential Privacy is arguably a field with growing adoption and better computational efficiency is a critical need;\n\n* Weakness: The method utilizes dedicated GPU kernels, which appears to be a general approach to enhance performance. It is compared to an implementation that does not have a custom kernel, which seems not a fair comparison.\n\nQUESTIONS:\n# Questions\n* The methodological improvements seem to come from computational considerations such as memory access, context switching, and redundant memory swaps. Given the authors' background in that literature, is this a common approach to improve code by lumping ops into a custom GPU kernel? If not, what is special about this method that makes it amenable; if it is common, what is the novelty of this paper?\n\n* Continuing the previous question, the paper claims that the custom kernel is more efficient and I assume that it was custom designed for the problem. a) To what other problems would this custom kernel be applicable? b) Would an automated tool be able to generate such a fused kernel for the same problem?\n\n* Quote from conclusion: \"At the same time, it highlights the need for responsible release practices to mitigate potential misuse under the guise of privacy.\" This sentence reads confusingly, but it seems important as the final sentence of the conclusion. What aspect in the paper \"highlights\" which \"release practices\"?\n\nLIMITATIONS:\nAs stated in the final question. Please explain what is meant with the line \"At the same time, it highlights the need for responsible release practices to mitigate potential misuse under the guise of privacy.\"", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an efficient implementation of DP-SGD training. Firstly, the paper adopts per-sample gradient clipping, so that we do not need to have cross-layer synchronization of gradient norms. Secondly, the paper fuses the operation of per-layer gradient clipping into several custom CUDA kernels to make the computation more efficient. Results show that the method improves the memory consumption and computation cost significantly compared to the vanilla DP-SGD, is compatible with several distributed training pipelines, and can be used to train models up to 13B.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n* The efficiency improvements are impressive.\n\n* The paper shows the feasibility of applying the method to train 13B models and using it in distributed training pipelines. This demonstrates the practical usability of the approach.\n\nWeakness:\n\n* Some results are missing/incomplete.\n\n* It seems like it will take non-trivial efforts to apply the approach to networks with new/non-standard network layers.\n\nSee \"Questions\" for details about the above points.\n\nQUESTIONS:\n* My main concern is that some important results are missing.\n  * Tables 1 and 2 have missing numbers.\n  * The paper claims that the method maintains \"parity with standard per-layer clipped DP-SGD in terms of accuracy.\" I agree that it is true by design, but I have some questions about the results.\n      * The only results I see are Table 3 in the appendix, where we see that DP-SGD and FlashDP have the same validation loss up to 4 decimal points. However, even if we train a model using the same codebase twice, due to the random seeds for data shuffling, etc. (if we do not fix the seed) and the randomness in GPU computation, we don't expect to see the same results. Could the authors explain why the results of DP-SGD and FlashDP match so well?\n      * Since some of the baselines (such as GhostClipping) use full model clipping (instead of per-layer clipping), it would be necessary to show if per-layer clipping leads to accuracy degradation in a full training process compared to full model clipping.\n\n* (Please correct me if my understanding is wrong.) Since the method requires custom CUDA kernels for the backpropagation process, we will need different implementations for different network layers. For example, Algorithm 1 lists how we implement it for a linear layer in transformers. If we have a different type of layer that has different input/output shapes (e.g., convolution layers in vision models), then we have to implement a different CUDA kernel. That means that the method is not very scalable in terms of implementation. In comparison, TF-Privacy does not need separate implementation for different network layers; while Opacus needs it, the implementation can be on the PyTorch level, which is easier than CUDA kernels. The paper needs to discuss the limitations.\n\n* Do the authors plan to release the code? I don't see it in the attachment. \n\n* Line 286: \"the grid synchronization required by CG necessitates launching all blocks simultaneously, which is impractical for DP applications.\" Could you explain it in more detail?\n\n* Table 1: Why does GhostClip have even less memory usage than Non-DP?\n\n* Typos:\n  * Line 279: PyTorch\n  * Line 331: FlashDP\n  * Line 360: FlashDP\n\nLIMITATIONS:\nI don't see discussions about limitations and the potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed memory- and compute-efficient LLM training method  (FLashDP) under the privacy constraints of differential privacy (DP). Specifically, authors have tacked the scalability challenges of DP-SGD (Differentially Private Stochastic Gradient Descent) primarily stemming from the per-sample gradient clipping, which demands huge GPU memory and pose a significant challenge for LLM's with billions parameters. \n\nThe authors validate the efficiency of FlashDP on GPT-2 (1B+) and LLaMA (up to 13B) models using 4 A100 GPUs. They show memory and throughput improvements over SoTA DP methods, and compare the validation loss to non-DP baselines.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths \n\n1. Training LLMs under differential privacy is a major challenge due to the fundamental tension between privacy and degraded utility. This work directly addresses that trade-off in the context of pre-training, which is far more challenging and demanding than fine-tuning. \n\n2. Unlike most prior work that focuses only on fine-tuning existing LLMs under DP, this paper targets the more ambitious and impactful goal of DP pre-training. This broaden the scope and applicability of DP to foundational model training, which is critical for ensuring privacy at scale.\n\n3. The proposed FlashDP is scalable to 13B models. \n\n4. The architectural advancement and engineering improvement in FlashDP  are clearly presented in the paper, and the limitations of prior work are well-motivated. The experimental results showing memory and throughput gains are well depicted. \n\n\n## Weaknesses\n\nA potential weakness is that the comparison of validation loss with non-DP baselines is not extended to downstream task evaluations. Particularly for longer context lengths (e.g., >1K), improvements in validation loss may not directly translate to downstream performance [2], potentially due to effects such as the dilution phenomenon discussed in [1]. \n\n[1] Fang et al., What is Wrong with Perplexity for Long-context Language Modeling?, ICLR 2025 \n\n\n[2] Hu et al., Can Perplexity Reflect Large Language Model's Ability in Long Text Understanding?, ICLR 2024\n\nQUESTIONS:\nSee the Weakness.\n\nLIMITATIONS:\nNot explicitly discussed in the paper. While authors have pointed to **Appendix A.3** for limitations (in NeurIPS checklist), I did not find the limitations of the author's proposed methods.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces FlashDP which allows efficient training of LLMs with differentially private SGD. FlashDP optimizes by implementing specific CUDA kernels which allows for efficient per-example gradient computation and reducing the HBM/SRAM IO overhead. Though FlashDP requires more than one CUDA kernel due to a required global synchronization of gradient norm square, the overall implementation is still more efficient than the previous approaches. FlashDP was compared against popular implementations such as Ocapus and showed similar memory and throughput as the non-DP approaches.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. The paper focused on an important question of implementing DP-SGD efficiently for LLMs. This is a crucial step for making LLMs more private and responsible as they are touching more of the private user data.\n2. The algorithm design is thoroughly described, and overall it is intuitive to understand even for readers without CUDA programming experience. \n3. The results indeed demonstrated the efficacy of FlashDP, showing better trade-off between memory and compute than the baseline approaches.\n\nWeaknesses\n1. Extensibility of the proposed method seems to be limited. The implementation was based on CUDA kernel, and thus might not apply to other hardwares e.g. TPUs or Apple’s M chips. Also, it seems like one has to implement CUDA kernels for different operations if extended to other architectures, making it harder to be extended.\n2. The main results in Table 1 only considered GPT-2 families which is a bit outdated. It would also be interesting to know the performance differences between different implementations on fine-tuning or post-training where DP is likely more needed than pre-training.\n3. Minor: multiple rendering of FlashDP’s as FlashDPâA˘ Zs, maybe some latex compiling issue.\n\nQUESTIONS:\n1. Is there a detailed profiling of the GPU I/O overhead reduced in FlashDP compared with existing approaches?\n2. How much faster would a single kernel implementation run, ignoring the inaccurate gradient clipping factor? Just curious about how much the block-wise synchronization overhead brings.\n\nLIMITATIONS:\nLimitations discussed. No potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-06"}