Private Training Large-scale Models with Efficient DP-SGD
Authors: Liangyu Wang
Junxiao Wang
Jie Ren
Zihang Xiang
David E. Keyes
Di Wang
Venue: NeurIPS 2025 poster ()
Keywords: Differential Privacy
Large Language Models
DP-SGD
Primary area: applications

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.

Private Training Large-scale Models with Efficient
DP-SGD
Liangyu Wang KAUST
Junxiao Wang Guangzhou University
Jie Ren KAUST
Zihang Xiang KAUST
David E. Keyes KAUST
Di Wang" KAUST
Abstract
As large language models (LLMs) increasingly underpin technological advance- ments, 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 integra- tion via Differentially Private Stochastic Gradient Descent (DP-SGD) introduces substantial challenges, primarily due to the complexities of per-sample gradi- ent clipping. Current explicit methods, such as Opacus, necessitate extensive storage for per-sample gradients. significantly inflating memory requirements. Con- versely, implicit methods like GhostClip reduce storage needs by recalculating gradients multiple times, which leads to inefficiencies due to redundant compu- tations. 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 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 mem- ory demands and achieves 90% throughput compared to the Non-DP method on a four-A100 100 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- training of LLMs. FlashDP's code has been open-sourced in https:/gghh....ommuuasssu/d////h
1 Introduction
The transformer architecture Vaswani et al. 2017) has revolutionized fields like natural language processing (Gao et al.) 2024: Xie et al., (023), embodied AI (Song et al. 2023 Duan et al. 2022 Xu et al.) 2024), and Al- Al-generated content (AIGC) (Cao et al. [023] Wu et al. 2023). with Large Language Models (LLMs) demonstrating exceptional abilities in text generation, complex query responses, and various language tasks due training massive datasets. These models, exemplified by ChatGPT, are applied across diverse areas, including healthcare, where they enhance diagnosis and drug discovery by analyzing medical data (Toma et al.) [023] Ali et al. 2023 Sheikhalishahi et al. 2019; (allam 2023 Biswas 2023). However, the extensive capabilities of LLMs raise significant
privacy concerns, particularly they can inadvertently expose or generate sensitive information,
owing to their potential to memorize data from large training sets (Pang et al.) [0024 Nasr et al. 2023
Carlini al. 2023; Ippolito al. 2022 McCoy et al. 2023
Tirumala al. 2022 Zhang al.
2023 Ashkboos et al. 2023).
Cooresponding Author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Differential Privacy (DP) ensures privacy by adding noise during data processing, such that any single data point's influence on outcomes minimal (Dwork 2006). As the most commonly adopted methods for ensuring DP in deep learning models, Differentially Private Stochastic Gradient Descent (DP-SGD) based methods (Abadi et al. 2016) adapt traditional stochastic gradient descent by clipping gradients per sample and adding noise. A Although DP-SGD's application in LLMs is increasing, recent research (Li et al., 2022 Bu et al. 2023b Anil et al., 2022 Hoory et al., 2021) | primarily targets the fine-tuning phase, providing privacy only for fine-tuned data. While some studies (Lee & Kifer 22021 Li al. 2022; Bu et al. 2023b have applied DP-SGD to pre-training, they often exhibit limited scalability or reduced training efficiency. This is primarily due to the significant computational and memory overheads inherent to per-sample gradient processing in DP-SGD, which make end-to-end pre-training of large models particularly challenging.
Integrating DP into L.M training via DP-SGD/Adam poses significant chal- lenges, particularly due to per-sample gradient clipping. This crucial privacy technique involves adjusting each data sample's s gradients to limit their influ- ence on model updates. While criti- cal for maintaining strict privacy stan- dards, this approach requires com- puting and storing individual gradi- ents, significantly raising computa- tional and memory demands. Man- aging these gradients is especially tax- ing in LLMs, which are known for their large parameter spaces. Each gradient must be carefully clipped and aggregated before updating model parameters, straining computational resources, and prolonging training times. These scalability issues are par- ticularly acute in settings with limited hardware, creating significant barriers efficiently training privacy-aware LLMs (Li et al.
Rorwrrd
Bookward
Forword
Backword
9
Team
Oxtpat
Butter
Oxppet
Uas buffir)
Fiakes
[[G]]
Salar
Polale aud
Figure 1: Comparison of different training methods. (a) Stan- dard DP-SGD: Stores per-sample gradients G (red explicit cache), increasing memory usage (blue buffer). (b) FlashDP: Optimizes gradient processing by consolidating computa- tions into a single pass, reducing redundancy and memory use.
2022 Bu
et al. 2023b).
Current research on DP-SGD for training LLMs can be categorized into two classes: explicit methods
like Opacus (Yousefpour et al. 202T) stand out by directly storing per-sample gradients. This
approach, while straightforward. significantly increases the memory footprint (Appendix Table 4
which becomes prohibitive for state-of-the-ar LLMs characterized by billions of parameters (Touvron al. 2023; Achiam et al. (023). Such a substantial increase in memory requirements hampers scalability and renders these methods impractical for deployment in large-scale model training environments. The direct storage of gradients, essential for ensuring the privacy guarantees of DP, thus poses substantial barrier to the efficient implementation of DP in LLMs.
Conversely, implicit methods, exemplified by innovations such as GhostClip (Li et al. 2021). address the memory challenge by circumventing the need for persistent storage of per-sample gradients. These methods segment the DP-SGD process into multiple discrete computational tasks, ostensibly to mitigate memory demands. However, this strategy necessitates the frequent recalculation of per-sample gradients, which introduces a high degree of computational redundancy (Table 4). This redundancy not only undermines training efficiency but also extends the duration of the training process significantly. For LLMs, which require substantial computational resources and extended training times, the inefficiencies introduced by such redundant computations become critical bottleneck. These implicit methods, while innovative in reducing memory usage, thus struggle to deliver a practical solution for the privacy-preservee training of LLMs at scale.
To effectively tackle the challenges presented by existing methods of integrating DP into the training of LLMs, we introduce FlashDP, a novel, cache-friendly implicit algorithm designed to streamline the per-layer clipping DP-SGD process (Figure (a)). We opt for per-layer clipping in our research primarily due to its efficiency in managing both memory consumption and accuracy. especially vital the differentially private training of expansive language models (Bu et al. 2023a He et al. 2022).

---

It has been shown that this type of method not only sustains commendable accuracy compared to the standard DP-SGD but also mitigates memory overhead, a critical consideration when training large-scale models under privacy constraints. FlashDP uniquely implements a unified computational strategy that performs the gradient operations required for DP-SGD in a single pass (Figure [] (b)). This innovative approach not only eliminates the need for multiple recalculations of per- sample gradients but also consolidates the entire process into one cohesive computational task. To be specific, FlashDP's architecture, which consolidates the entire DP-SGD process into a single GPU kernel, eliminates redundant computations and optimizes data flow within the GPU. This integration results in streamlined workflow that efficiently manages memory and processing resources. Also, FlashDP reorganizes the GPU operations to maximize data throughput and minimize latency, effectively enhancing the overall efficiency of the training process. These architectural improvements significantly reduce the volume of memory transfers and computational redundancies, therehy optimizing both the speed and resource utilization during the training of LLMs with DP.
By re-designing the gradient computation workflow, FlashDP dramatically reduces the volume of memory transfers by 50% and decreases redundant computational tasks by 20% compared to previous implicit methods. This optimization is achieved through an advanced caching mechanism that efficiently manages gradient data and computation within GPU memory, minimizing the data movement across the system. As a result, FlashDP significantly alleviates the memory overhead traditionally associated with DP-SGD, enhancing the model's 's scalability and training speed.
The practical impact of these improvements is substantial. On a computational platform equipped with four NVIDIA A 110 GPUs, FlashDP achieves a remarkable 90% throughput compared to the non-DP method during the pre-training phase of the Llama-13B model, state-of-the-art [.LM known for its extensive data and computation demands. Crucially, this enhanced performance is attained without any degradation in the accuracy or dilution of the privacy guarantees compared to the original per-layer clipped DP-SGD. FlashDP thus not only meets but exceeds the operational requirements for effective and efficient privacy-preserving training of LLMs.
Our contributions can be summarized as follows:
Enhanced Throughput for LLM training with DP: We propose FlashDP, which effectively resolves the issue of low throughput in DP-SGD/Adam with per-layer clipping during the training of LLMs. By optimizing the computational workflow and integrating more efficient handling of per-sample gradients, FlashDP significantly enhances the processing speed without compromising the model's accuracy or privacy integrity.
Innovative GPU I/O Optimization: Our study pioneers the exploration of DP-SGD from the perspective of GPU input/output operations. FlashDP's architecture, which consolidates the entire DP-SGD process into a single GPU kernel, eliminates redundant computations and optimizes data flow within the GPU. This approach not only reduces the computational load but also minimizes the number of GPU memory accesses, setting a new standard for efficiency in DP implementations.
Experimental Validation of Efficiency and Scalability: In practical LLM models involv- ing Llama-13B. FlashDP matches the speed and memory usage of non-DP training methods and achieves a significant 90% throughput compared with Non-DP methods. This perfor- mance is achieved on computational platform equipped with four NVIDIA A100 GPUs. Importantly, it accomplishes this without any degradation in the precision or the privacy guarantees typically observed in the previous per-layer clipped DP-SGD implementations. This capability demonstrates FlashDP's effectiveness in scaling DP applications to larger and more complex LLMSs without the usual trade-offs.
2 Related Work
Improving Time and Memory Complexities of DP-SGD. The transition from standard stochastic gradient descent to DP-SGD introduces substantial modifications in memory and computational demands. In conventional settings, parameter updates are efficiently computed by aggregating gradi- ents across all samples within a batch. This approach is both memory-efficient and computationally straightforward. contrast, DP-SGD mandates that each sample's gradients be preserved, clipped, and subsequently aggregated to uphold privacy guarantees. Recent innovations in DP-SGD have
3

---

primarily concentrated on ameliorating its computational and memory inefficiencies. TF-Privacy vectorizes the loss to calculate per-sample gradients through backpropagation, which is efficient in terms of memory but slow in execution (Abadi et al., 2015). Opacus (Yousefpour et al. 2021 ) and (Rochette et al., 2019) enhance the training efficiency by employing the outer product method [Goodfellow] 2015), albeit at the cost of increased memory usage needed to store per-sample gradi- ents. This memory overhead is mitigated in FastGradClip (Lee & Kifer 2020 by distributing the space complexity across two stages of backpropagation, effectively doubling the time complexity. Additionally, ghost clipping techniques (Goodfellow) 2013), (Li et al., 202T). (Bu et al.) 2022) allow for clipping per-sample gradients without full instantiation, optimizing both time and space, particularly when feature dimensions are constrained. Furthermore, (Bu et al. 2023b) introduces a "book-kcepingg (BK) method that achieves high throughput and memory efficiency, but still leaves room for improvement in fully addressing the computational and memory bottlenecks inherent in large-scale DP training.
While these methodologies have made significant strides in mitigating the extensive computational and memory demands typically associated with managing per-sample gradients in DP-SGD, they have not addressed the optimization of DP training from the perspective of GPU architecture and memory access. Additionally, the approaches detailed thus far do not cater effectively to the training of today's ' LLMs. FlashDP aims to enhance the efficiency and feasibility of training LLMs under the constraints of differential privacy, ensuring both high performance and adherence to privacy standards.
DP for Large Language Models. The field of privacy-preserving LLMs is characterized by the use or exclusion of DP and extensions. (He al., 2022) evaluated the precision equivalence of per-layer clipping with flat clipping on LLMs. (Kerrigan et al. (020)) demonstrated that public pretraining could facilitate downstream DP fine-tuning, although they did not explore fine-tuning large pre-trained models using DP-SGD. (Qu et al., [02 explored the fine-tuning of BERT for language understanding tasks under local DP. (Bommasani et al. 202T) suggested the potential for cost-effective private learning through fine-tuning large pre-trained language models. (Anil et al.) 202I) and (Dupuy et al.] (022) extended these studies to BERT, pretraining and fine-tuning under global DP, respectively, with (Anil et al.] (021) addressing datasets comprising hundreds of millions of examples, and (Dupuy et al. (022) reporting on datasets of utterances with relatively high e values. Our research distinguishes itself by focusing on pre-training and fine-tuning large language models with high throughput and low memory usage.
3 Understanding the Limitations of Previous Methods
In this section, we introduce the previous non-DP, explicit, and implicate methods of DP-SGD from the GPU I/O perspective to see their weakness, which motivates framework. Due to the space limit, please refer to Appendix B for the background on DP, Transformers. GPU architecture, and CUDA programming. As discussed in Section B.2 the linear operation is crucial in the architecture of LLMss, particularly within Multi-Head Attention (MHA) and Feedforward Network (FFN) modules. Given its significance, we utilize the linear operation as exemplar elucidate the training workflow on GPUs, as shown in Figure See Appendix C for details.
In the standard non-private training workflow of a linear layer, the forward pass involves matrix multiplication Y = XWT between the activation tensor X E R BxTxx and the weight matrix W E DxP, resulting in the output Y E BxTxxD where B, T, P, and D denote the batch size, sequence length, input feature dimension, and output feature dimension, respectively. The backward pass calculates the output gradient Vy E R BxTxx D and the weight gradient Vw E RPx P via Vw = ΣB ( X. Figure 6 (a) illustrates this process, showing that the activation tensor X and weights W are stored in HBM for efficient access during computations, while intermediate operations utilize SRAM to enhance memory access time and throughput.
The explicit DP-SGD workflow, as depicted in Figure (b), categorizes the process into four stages to ensure privacy adherence by explicitly managing per-sample gradients. Stage involves computing per-sample gradients G En V using batched GEMM operations on SRAM to minimize latency, with subsequent storage of the gradients back to HBM. Stage 2 requires reloading these gradients to compute their norm ||G|| = VED Ep G2, then storing the results back in HBM. Stage 3 includes loading the gradients and their norms for the per-layer clipping operations, ensuring

---

that no gradient norm exceeds the predefined threshold C, with the clipped gradients G' written back to HBM. Stage 4 focuses on adding Gaussian noise the clipped gradients in SRAM for privacy preservation, followed by their aggregation for model updates, and storing the final noisy gradient Vw back HBM. This explicit handling of per-sample gradients not only increases memory usage but also complicates processing due to frequent memory swaps and disrupts efficient GPU utilization by breaking down kernel fusion strategies, becoming notably impractical for LLMs with their extensive parameter and gradient sizes, severely impacting training efficiency.
The implicit DP-SGD workflow, illustrated in Figure 6 (c), employs a method such as GhostClip to recalculate gradients in a fused manner, thus circumventing the need for explicit storage of per- sample gradients. Stage consolidates the first three stages of the explicit method into single fused computational step, where the activation tensor X and output gradient tensor Vy are loaded into SRAM. Per-sample gradient tensor recalculations, norm calculations, and the per-layer clipping are integrated into one operation, minimizing latency and avoiding repeated data transfers to HBM. Stage mirrors the explicit method's final stage, where the recalculated and clipped gradients G' undergo Gaussian noise addition in SRAM, followed by aggregation and storage in HBM for model updates. This approach reduces memory usage but increases computational load due to the redundancy of multiple gradient recalculations, which can significantly extend training times, rendering the method less practical due to the increased time complexity proportional to T.
To address the previous limitations, the subsequent section will introduce FlashDP, a novel strategy designed to address these inefficiencies by rethinking the execution pipeline of DP-SGD, Without delving into specifics here, FlashDP's architecture will streamline the integration of per-sample gradient computation and clipping, potentially reducing the operational bottlenecks observed in existing methods.
4 FlashDP Algorithm Design
IBB
Hack
X
Blcckviter Al--eddac
Bladimbe
Fir
Backe
duc.
Figure 2: Illustration of Fl depicts the core algorithm design of FlashDP. Its features are integrated with on-chip per-sample gradient norm calculations. The workflow incorporates block-wise all-reduce and synchronization to facilitate efficient norm aggregation. SRAM (orange) and HBM (green) are optimally utilized to manage memory efficiently, addressing the kernel fusion challenges and reducing computational redundancy inherent traditional DP-SGD implementations.
4.1 Algorithmic Enhancements FlashDP
FlashDP introduces a suite of algorithmic enhancements designed to reconcile the computational demands and memory constraints associated with DP-SGD. At the heart of these enhancements is the Block-wise All-Reduce algorithm, which integrates several critical operations into a unified kernel execution, thereby optimizing on-chip memory utilization and enhancing computational throughput.
Efficient Kernel Fusion through Block-wise All-Reduce. Central to FlashDP's strategy is our proposed Hierarchical Reduction Architecture (HRA), which encompasses more than just reduction operations. HRA is structured approach that manages the computation and synchronization of data across various stages, beginning with intra-block reduction of gradient norms within individual GPU blocks. This phase employs an HRA-based reduction strategy executed in shared memory,
5

---

culminating in single norm scaler block. Such design significantly reduces the data footprint necessary for subsequent inter-block communications, optimizing the efficiency of the all-reduce operation across the GPU grid.
Following the compact intra-block reduction, FlashDP coordinates global all-reduce operation across blocks, which computes global gradient norm crucial for consistent gradient clipping across the entire mini-batch. Efficiently handled in HBM thanks the minimized data size from earlier reductions, this step avoids the common memory bottlenecks typically associated with large-scale data operations in HBM, thus maintaining high computational throughput.
Algorithm 1 Algorithm: FlashDP with Block-wise All-Reduce on GPUs
Require:
Input activation tensor X E BBXxxP and output gradient tensor Vy RR BxTxx in GPU
HBM
Require:
Clipping threshold C, noise scale a
Require: Block dimensions b, t, d, and for batch size, sequence length, output features, and input features, respectively.
1: Split block for output gradient tensor Bvy E Rbxtxd, input activation tensor Bx E R 6xxxx based on GPU on-chip SRAM size M.
2: for each training backward iteration do
3: for each block input index ip = 1,2,....
in parallel do
4:
for each block output feature in 1, ........ 2 in parallel do for each block batch size it 2, ................................ in parallel do Load output gradient block B and input activation block Bx from HBM to SRAM.
5:
6:
7:
Compute per-sample gradients block Be - ET B Bx on-chip SRAM. Intra-block Reduce: Compute per-sample gradients Ea Ep BG 2 on-chip SRAM.
8:
square block [B
9:
Inter-block Reduce: Offload all per-sample gradients norm square blocks || Ball from SRAM to HBM, and perform block-wise all-reduce. Block-wise synchronization: Wait until all blocks finish the all-reduce operation to get all-reduced per-sample gradients norm square blocks ||
10:
11:
Upload || BB || from HBM to SRAM.
Compute clipped per-sample gradients block B' = BG/ max 1, y /a
12:
on-chip
SRAM.
13:
Add noise to clipped per-sample gradients block and aggregate to compute parameter gradient block B En B'g + +OO,, C²C on-chip SRAM.
14:
Offload parameter gradient block Bv from SRAM to HBM. end for
16:
end for
end for
18:
end for
Return entire parameter gradient Vw.
The strategic implementation of HRA not only facilitates these reductions but also orchestrates synchronized updates and data consistency across the GPU architecture. By managing data flow from the point of loading through to final computation and storage, HRA ensures that the most intensive computations are confined to the faster, on-chip memory. This methodical approach leverages the GPU's capabilities to facilitate high-performance differentially private training, minimizing memory and bandwidth overhead.
The practical implementation and operational dynamics of the FlashDP approach are thoroughly illustrated in Algorithm [ and visually depicted in Figure 2 FlashDP innovatively reduces the four distinct stages typically involved in explicit DP-SGD into a single streamlined stage. This consolidation is achieved without adding any extra computational steps, thereby enhancing the overall efficiency of the process. Here is detailed breakdown of this single streamlined stage:

---

Optimized Block Processing and Memory Management (Line 1-6). Initially, FlashDP partitions the input activation tensor X and the output gradient tensor Vy into blocks based on the SRAM capacity. This strategic partitioning is crucial for managing the limited on-chip memory more effectively and ensuring that data transfers between the HBM and SRAM are minimized.
Fused Computation of Gradients and Norms (Line 7-8). Within the GPU's SRAM, FlashDP simultaneously computes the per-sample gradients block and their norms square (intra-block reduce) for each block. This computation leverages the GPU's powerful batched GEMM operations, enabling it to handle large data sets efficiently.
Block-wise All-Reduce (Line 9-11). After computing the gradient norms, FlashDP performs a wlsck-wise All-Reduce operation in parallel to aggregate these norms across all blocks (inter-block reduce). This all-reduce operation is crucial for obtaining global view of gradient norms square, which is necessary for consistent gradient clipping across the entire batch. This step is executed efficiently within the SRAM, reducing the latency and memory bandwidth requirements typically associated with inter-GPU communications.
Per-layer Gradient Clipping and Noise Addition in SRAM (Line 12-13). Following the gradient and norm calculations, clipping is performed directly on the chip. Each gradient is scaled according to the computed norms and predefined clipping threshold C, ensuring compliance with DP standards. Immediately after clipping, Gaussian noise based on the noise scale a and the clipping threshold is added to each gradient block.
Efficient Parameter Aggregation (Line 14-19). The final step in the FlashDP algorithm involves aggregating the noisy, clipped gradients across all blocks and batches directly within SRAM. This aggregation is optimized to minimize memory accesses, ensuring that only the final gradient used for the model update is transferred back to HBM.
4.2 Adaptive Kernel Implementation
The implementation of the Fl algorithm leverages the robust and versatile capabilities of the PyTorch framework (Paszke et al. 2019), which is renowned for its intuitive handling of automatic differentiation and dynamic computational graphs. One of the critical features of our implementation involves customizing Py MyTorch's 's autograd functionality to accommodate the specific needs of differ- ential privacy during the training of deep neural networks. To this end, operators that necessitate trainable parameters are intricately defined by wrapping them within PyTorch's 's autograd function.
However, implementing the Block-wise All-Reduce algorithm has presented unique challenges, primarily due to the limitations of CUDA's programming model in facilitating block-wise synchro- nization. Block-wise synchronization is essential in our algorithm; without it, clip operations might be executed prematurely, while the inter-block reduce operation is still incomplete, leading to numerical inaccuracies in the computation of per-sample gradients' norm squares. There are two primary methods to implement synchronization: 1. cooperative groups (CG) [nd [and 2. adaptive kernel. We opted for the second method because the grid synchronization required by CG necessitates launching all blocks simultaneously. which is impractical for DP applications. Although CG supports global synchronization, its requirement launch all blocks concurrently imposes hardware constraints, making it unsuitable for large-scale DP training. We therefore opted for the adaptive kernel method.
To address this limitation, FlashDP's implementation employs an adaptive approach. Instead of relying on a monolithic kernel to perform the entire Block-wiss All-Reduce operation, the process is split across different kernels, which are executed iteratively over the batch dimension. This iterative approach allows for synchronization points between the execution of kernels, using the inherent block synchronization that occurs at kernel launch and completion.
The execution flow in FlashDP is as follows: (1) Intra-block Reduction: Each block computes the norms of its gradients and performs HRA-bbsed reduction within the shared memory. This step employs a shuffle-reduce mechanism, optimizing intra-block operations by minimizing memory footprint and synchronization overhead. This results in a single norm value per block. (2) Inter-block Reduction: Each block transfers the outcome of its intra-block reduction to the HBM. This transfer is facilitated through atomic operations for several reasons. Firstly, the result of the intra-block reduction comprises only a single element, and each block elects only one thread to perform the
7

---

atomic operation on this element. This approach minimizes potential bottlenecks, as the differing execution speeds across blocks prevent serious serialization issues. Secondly, atomic operations benefit from acceleration by the hardware instruction set, ensuring that these operations are executed swiftly and efficiently. (3) Inter-kernel Synchronization: After the completion of the inter-block reduction, FlashDP leverages the termination of the kernel as a natural synchronization point. At this juncture, all blocks have finished their individual reductions. (4) Iterative Kernel Launch: For each batch element, a new kernel is launched serially, maintaining synchronization across kernels. This approach involves broadcasting operations where source operands are dimensionally disparate, ensuring uniform data handling across computational units.
This implementation strategy, while divergent from the ideal single- kernel solution, allows FlashDP to function effectively within the current constraints of CUDA. It underscores FlashDP's s adaptability and represents a practical solution to the block synchronization challenge, ensuring accurate gradient norm calculations essential for maintaining the model's differential privacy.
Experiments
Our experimental suite is methodically designed to assess the robustness and efficiency of FlashDP across a range of training paradigms and hardware configurations. We explore FlashDP's performance in terms of memory efficiency and throughput under varying batch sizes, its adaptability to Automatic Mixed Precision (AMP) training (Appendix Section [2.2 its scalability when employing Distributed Data Parallel (DDP) and Pipeline Parallel (PP) techniques (Appendix Section F.3)
and utility
evaluation (Appendix Section F.4
Table 1: Differential Batch-size Analysis. The table displays multi-panel comparison of memory usage and throughput for four differential privacy methods-NonDE Opacus, GhostClip, BK, and FlashDP-across different batch sizes (1, 2, 4, and 8) when applied to GPT-2 models of varying sizes (small, medium, and large). Instances of 1-' in the table indicate scenarios where the corresponding method failed to execute due to memory constraints.
Model
B
NonDP
Upucus
Momoyy Usage (MB ko04)
GoosCCip
DX
FashDP
NonDP
Opscus
Tironghput (Jokenskec xlp4)
Clbssi up
BX
FIgEDP
GPT2-mmal
0.30
0.35
0.466,0)992
0.35(x1.05)
0,500x11 00)
2.31
0.91(00.32)
0.37 20)
L 54)
L83
GPT2-medium
1.26
1.55x1.21)
1.12/s0.89)
1 1.6)(x1.33)
0,42(x0.38
0.39
0.75
0.86c80.78)
GPT2-large
2.48
399x1161)
2.170x0.88)
2.730x1
2.48(v1 00)
0.27 0.25xx0.45)
0.40000.00
0.51(v0.89)
GPT2-5mall
038
(303129)
0.79(50.91)
0.87X1 (0)
3:22
0.02600 20)
L9110039)
2.30700.22
GPT2-medium
2
2.07
1.896
1.87(80.90)
2.44(11188
2.07(x] (0)
1.28
0.74 58)
0.59 46)
081(00633)
1.02)00.80)
GPT2-large
3.91
4 47951.23)
3.53(00.90
4.81((123))
3.91(11.00) 00)
0.68
(.3))0056)
0.380 56)
6.45000666
0.9%)
GPT2-amal
1.31
10131133)
1.44500.933
14350066)
T,2
GPI2-medium
4
3.58
4.26(x1.19)
4.00kx1.12 3.33(x0.93)
3.58tx1
1.42
0.900x0.63)
5.81100.7)
0.95px0.67
113
GPT2-lagge
6,0
6.15(x0:93)
6.600c 00)
6.00 00)
0.76
0,50000066
0.53/10 70)
0.64(v0.84)
GPT2-mmll
2366
34461.20)
2.75(00995)
2.66(1.000 L.0)
5,30
244(0099)
(.9200331)
240
27238172)
GPT 2-eeiim
8
6.60
6.24(x0.95)
6.60)x1000
6.600x1100)
1.52
0.99)x0.65)
103(80.68
1.19
GP12-Jarge
5.1 Experimental Setup
Our experiments utilize the Wikitext dataset (Merity 2016) and are conducted on NVIDIA A100
(80GB) GPUs using the PyTorch framework (Paszke et al. 2019). We assess the performance of FlashDP across various configurations by comparing it with established explicit methods Opacus (Yousefpour et al., 2021). and implicit method GhostClip (Li et al. (021) and BK (Bu et al.) 2023a).
all in the per-layer clipping mode, under different training paradigms The tested models include GPT-2 (Radford et al.] 2019) with sequence length of 1024 and Llama (Touvron et al. (023) models, both with a sequence length of 2048. More experimental settings and explanations can be found in
Appendix D
5.2 Results of Batch Size& Micro Batch Size
Efficient batch processing is crucial in LLM training due to its high computational and memory demands. By examining both batch and micro-bateh sizes, we assess FlashDP's s ability to manage
As (Bu et al. 2023a| He et al. 2022) demonstrated that for LLMs, compared to global clipping, lay-r clipping is more memory efficientt and üime-efficient while achieving comparable performance in terms of both privacy preservation and accuracy. Here, we only consider per-layer clipping baselines.
8

---

memory more effectively and maintain high throughput. This also tests the practicality of gradient ac- cumulation (GA), which allows larger effective batch sizes by splitting them into smaller, manageable micro-batches. The experiment results of different micro batch sizes can be seen in Appendix F.1
In Table D FlashDP was benchmarked against traditional DP-SGD methods like Opacus, GhostClip, and BK, as well as a non-DP (NonDP) configuration, demonstrating superior memory efficiency and throughput. FlashDP utilized approximately 38% less memory than Opacus and nearly matched the NonDP configuration while processing the GPT-2 large model at a batch size of 1. It achieved a throughput nearly double that of Opacus and only slightly lower than NonDP, showcasing its effective balance between privacy preservation and computational efficiency. Opacus exhibited the highest memory usage, which escalated with batch size, leading to failure at a batch size of 8. GhostClip. while more memory- -eficccee than Opacus, suffered from reduced throughput at higher batch sizes due to gradient re-computation. BK's performance was intermediate, lacking distinct advantages. Overall, FlashDP not only maintained lower memory usage and higher throughput than the DP methods across all batch sizes but also approached the efficiency of NonDP configurations.
5.3 Results of Distributed Training
Distributed Data Parallel (DDP) (Li et al. 2020 and Pipeline Parallel (PP) (Kim et al. 2020) are two advanced techniques crucial for scaling the training of LLMs efficiently across multiple GPUs or nodes.
10C000
00000
2000
Norte
00000
60000
FlancCe
60090
42000
3000
10000
(a) Memory Usage
(b) Throughput
Figure 3: Memory and Throughput for Llama Models Using Pipeline Parallel Training. (a) Memory usage for Llama-3B 3B, Llama-7B. and Llama- 13B models. (b) Throughput in tokens per second across these model sizes. A value of indicates out of memory.
Distributed Data Parallel (DDP). Figure 7 in Appendix illustrates the performance of different methods in a DDP setting across GPT-2 models of varying sizes. FlashDP showcases superior memory usage efficiency and higher throughput across all model sizes when compared to Opacus and BK. Notably, even as the model size increases, FlashDP maintains a competitive edge close to the NonDP benchmarks, highlighting its effective parameter distribution and gradient computation across multiple GPUs This is crucial in scenarios where training speed and model scalability are priorities.
Pipeline Parallel (PP). In the PP scenario depicted in Figure 5 FlashDP was tested with Llama models varying from 3 billion to 13 billion parameters. The results indicate that FlashDP not only scales efficiently with increasing model size but also demonstrates significant throughput improvements compared to Opacus and BK. Particularly, FlashDP's ability to handle the largest model (Llama- 133)) with minimal throughput degradation illustrates its robustness in managing extensive computational loads, characteristic of PP environments.
6 Conclusion
In this paper, we introduce FlashDP, a cache-friendly approach to pey- DP-SGD that improves memory efficiency and computational throughput for large language model (LLM) training. By optimizing GPU I/O through a unified Block-wise All-Reduce algorithm and a Hierarchical Reduction Architecture (HRA), FlashDP significantly reduces memory transactions and eliminates redundant computations. We also adopt an adaptive kernel design to overcome CUDA's synchronization limitations. Experiments show that FlashDP achieves memory usage close to non-private baselines and maintains 90% throughput during I.lam- training on four 100s, without compromising privacy accuracy. Fl may enable the deployment of preseriynr LLMs in sensitive

---

org/ org/ Software available from domains such as healthcare, education, and finance, where data protection is critical. At the same time, it highlights the need for responsible release practices to mitigate potential misuse under the guise of privacy.
Acknowledgments and Disclosure of Funding
Di Wang and I Liangyu Wang are supported in part by the funding BAS/1/1689-01-01 and funding from KAUST Center of Excellence for Generative AI, under award number 5940.
References
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, lan Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard, Yangging Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Murray. Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke. Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL http:://ttt:ss: http:://tensoffoo..ooo tensorflow.org.
Martin Abadi, Andy Chu, Ian Goodfellow, Brendan McMahan, Ilya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Proceedings of the 2016 ACM SIGSAC conference on computer and communications security, pp. 308-318, 2016.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et Gpt-4 technical report. arXiv preprint arXiv: arXiv:2303.08744 2023.
Stephen R Ali, Thomas Dobbs, Hayley A Hutchings, and Iain Whitaker. Using chatgpt to write patient clinic letters. The Lancet Digital Health, 5(4):e 179-e181, 2023.
Rohan Anil, Badih Ghazi, Vineet Gupta, Ravi Kumar, and Pasin Manurangsi. Large-scale differen- tially private bert. arXiv preprint Wiv::2001.00022. 2021.
Rohan Anil, Badih Ghazi, Vineet Gupta, Ravi Kumar, and Pasin Manurangsi. Large-scale differen- tially private bert. Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 6481-6491, 2022.
Salch Ashkboos, Ilia Markov, Elias Frantar, Tingxuan Zhong, Xincheng Wang, Jie Ren, Torsten Hoefler, and Dan Alistarh. Towards end-to-end 4-bit inference on generative large language models. arXiv preprint arXiv:2310.092599 2023.
Som S Biswas. Role of chat gpt public health. Annals of biomedical engineering, 51(5):868-869, 2023.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney Arx, Michael S Bernstein, Jcannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportuni- ties and risks of foundation models. arXiv preprint arXiv:2 arXiv:2/08.072. 108.07258, 2021.
Zhiqi Bu. Jialin Mao, and Shiyun Xu. Scalable and efficient training of large convolutional neural networks with differential privacy. Advances in Neural Information Processing Systems, 35: 38305-38318, 2022.
Zhiqi Bu, Ruixuan Liu, Yu-Xiang Wang, Sheng Zha, and George Karypis. On the accuracy and efficiency of group-wise clipping in differentially private optimization. arXiv preprint arXiv:2310. 19..2.55.. 2023a.
Zhiqi Bu, Yu-Xiang Wang, Sheng Zha, and George Karypis. Differentially private optimization on large model at small cost. In International Conference on Machine Learning, pp. 3192-3218. PMLR, 2023b.
10

---

Yihan Cao, Siyu Li, Yixin Liu, Zhiling Yan, Yutong Dai, Philip S Yu, and Lichao Sun. A comprehen sive survey ai-generated content (aige): A history of generative from gan chatgpt. arXiv preprint arXiv:2303.04222, 2023.
Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian Tramer, and Chiyuan Zhang. Quantifying memorization across neural language models. In International Conference on Learning Representations, 2023.
Jiafei Duan, Samson Yu, Hui Li Tan, Hongyuan Zhu, and Cheston Tan. A survey of embodied ai: From simulators to research tasks. IEEE Transactions on Emerging Topics in Computational
Intelligence, 6(2):230-22) 244, 2022.
Christophe Dupuy, Radhika Arava, Rahul Gupta, and Anna Rumshisky. An efficient dp-sgd mech- anism for large scale nlu models. In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4118-4122. IEEE, 2022.
Cynthia Dwork. Differential privacy. In International colloquium on automata, languages, and programming, pp. 1-12. Springer, 2006.
Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. In Theory of Cryptography: Third Theory of Cryptography Conference, TCC 2006, New York, NY, USA, March 4-7, 2006. Proceedings 3, pp. (66. Springer, 2006.
Mingqi Gao, Xinyu Hu, Jie Ruan, Xiao Pu, and Xiaojun Wan. Llm-based nlg evaluation: Current status and challenges. arXiv preprint FXXi::4422.01833, 2024.
lan Goodfellow. Efficient per-example gradient computations. arXiv preprint arXiv:1510.000 2015.
Jiyan He, Xuechen Li, Da Yu, Huishuai Zhang, Janardhan Kulkarni, Yin Tat Lee, Arturs Backurs, Nenghai Yu, and Jiang Bian. Exploring the limits of differentially private deep learning with group-wise clipping. arXiv preprint arXiv:2212.01539 2022.
Shlomo Hoory, Amir Feder, Avichai Tendler, Sofia Erell, Alon Peled-Cohen, Itay Laish, Hootan Nakhost, Uri Stemmer, Ayelet Benjamini, Avinatan Hassidim, al. Learning and evaluating a differentially private pre-trained language model. In Findings of the Association for Computational Linguistics: EMNLP 2021, pp. 1178-1189, 2021.
Daphne Ippolito. Florian Tramèr, Milad Nasr, Chiyuan Zhang, Matthew Jagielski, Katherine Lee, Christopher A Choquette-Chop and Nicholas Carlini. Preventing verbatim memorization in language models gives false sense of privacy. arXiv preprint arXiv:2210. /7546, 2022.
Gavin Kerrigan, Dylan Slack, and Jens Tuyls. Differentially private language models benefit from public pre- pre-training, arXiv preprint wrXiv:2009.05888 2020.
Chiheon Kim, Heungsub Lee, Myungryong Jeong, Woonhyuk Back, Boogeon Yoon, Ildoo Kim, Sungbin Lim, and Sungwoong Kim. torchgpipe: On-the-fly pipeline parallelism for training giant models. arXiv preprint arXiv:2004.0990 2020.
Jaewoo Lee and Daniel Kifer. Scaling up differentially private deep learning with fast per-example gradient clipping. arXiv preprint arXiv:2009.00 201.:2000003166, 2020.
Jaewoo Lee and Daniel Kifer. Scaling up differentially private deep learning with fast per-ex gradient clipping. Proceedings on Privacy Enhancing Technologies, 2021.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006. 15704, 2020.
Xuechen Li. Florian Tramer, Percy Liang, and Tatsunori Hashimoto, Large language models can be strong differentially private learners. arXiv preprint arXiv:2110.05677 2021.
Xuechen Li, Florian Tramer, Percy Liang, and Tatsunori Hashimoto. Large language models can be strong differentially private learners. In International Conference on Learning Representations, 2022.

---

Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. Fineweb-edu: the finest collection of educational content, 2024. URL htttps://uuggiigfacce co/ /
Haggginggaacee////nseeebeeee
R Thomas McCoy, Paul Smolensky, Tal Linzen, Jianfeng Gao, and Asli Celikyilmaz. How much do language models copy from their training data? evaluating linguistic novelty text generation using raven. Transactions of the Association for Compulational Linguistics, 11:652-670, 2023.
Stephen Merity. The wikitext long term dependency language modeling dataset. Salesforce Metamind, 9, 2016.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv: 2017.
Milad Nasr, Nicholas Carlini, Jonathan Hayase, Matthew Jagielski, A Feder Cooper, Daphne Ippolito, Christopher Choquette-Choo, Eric Wallace, Florian Tramèr, and Katherine Lee. Scalable extraction of training data from (production) language models. arXiv preprint arXiv:2311.17035,
2023.
Qi Pang. Jinhao Zhu, Helen Möllering. Wenting Zheng, and Thomas Schneider. Bolt: Privacy- preserving, accurate and efficient inference for transformers. In 2024 IEEE Symposium on Security and Privacy (SP). pp. 130-130. IEEE Computer Society, 2024.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, pegsoperformance deep learning library. Advances in neural information processing systems, 32, 2019.
Chen Qu, Weize Kong, Liu Yang, Mingyang Zhang, Michael Bendersky, and Marc Najork. Privacy- adaptive bert for natural language understanding. arXiv preprint arXiv:2104.07504, 190, 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Gaspar Rochette, Andre Manoel, and Eric W Tramel. Efficient per-example gradient computations in convolutional neural networks. arXiv preprint arXiv:1912.06015. 2019.
Malik Sallam. Chatgpt utility in healthcare education, research, and practice: systematic review on the promising perspectives and valid concerns. In Healthcare, volume 11, pp. 887. MDPI, 2023.
Seyedmostafa Sheikhalishahi, Riccardo Miotto, Joel T Dudley, Alberto Lavelli, Fabio Rinaldi, Venet Osmani, et al. Natural language processing of clinical notes on chronic diseases: systematic review. JMIR medical informaties, 12239, 2019.
Chan Hee Song, Jiaman Wu, Clayton Washington, Brian M Sadler, Wei-Lun Chao, and Yu Su. Lim-planner Few-shot grounded planning for embodied agents with large language models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 2998-3009, 2023.
Kushal Tirumala, Aram Markosyan, Luke Zettlemoyer, and Armen Aghajanyan. Memorization without overfitting: Analyzing the training dynamics of large language models. Advances in Neural Information Processing Systems, 35:38274-38290, 2022.
Augustin Toma, Patrick R Lawler, Jimmy Ba, Rahul G Krishnan, Barry B Rubin, and Bo Wang. Clin- ical camel: An open-source expert-level medical language model with dialogue-based knowledge encoding, arXiv preprint arXiv:2305.1203/ 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. /397/, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
12

---

Jiayang Wu, Wensheng Gan, Zefeng Chen, Shicheng Wan, and Hong Lin. Ai-generated content (aige): A survey. arXiv preprint 204x:22304.06632, 2023.
Yaqi Xie, Chen Yu, Tongyao Zhu, Jinbin Bai, Ze Gong, and Harold Soh. Translating natural language to planning goals with large-lngguage models. arXiv preprint arXiv:2302.05/28. 2023.
Zhiyuan Xu, Kun Wu, Junjie Wen, Jinming Li, Ning Liu, Zhengping Che, and Jian Tang. A survey on robotics with foundation models: toward embodied ai. arXiv preprint arXiv:2402.02385. 2024.
Ashkan Yousefpour, Igor Shilov, Alexandre Sablayrolles, Davide Testuggine, Karthik Prasad, Mani Malek, John Nguyen, Sayan Ghosh, Akash Bharadwaj. Jessica Zhao, et al. Opacus: User-friendly differential privacy library in pytorch. arXiv preprint 109i::2109./22
2021.
Chiyuan Zhang, Daphne Ippolito, Katherine Lee, Matthew Jagielski, Florian Tramèr, and Nicholas Carlini. Counterfactual memorization neural language models. Advances in Neural Information Processing Systems, 36:39321-39366 2023.
Explanation of FlashDP A
Note that FlashDP does not interfere with any sampling strategy, since these decisions are made during the data loading stage.
B Preliminaries
B.1 Differential Privacy
Definition 1. (Differential Privacy (Dwork et al. 2006 Given data universe X, two datasets X, X' S X are adjacent they differ by one data example. A randomized algorithm M is (c, 6)- differentially private for all adjacent datasets X, X' and all events S in the output space of M. we have Pr(M(X) E S) S ef Pr(M((') E 66
Differentially Private Stochastic Gradient Descent (DP-SGD) (Abadi al. 2016). DP-SGD is an adaptation of this principle for machine learning models, where privacy is preserved during the training process by modifying the gradient computation.
In the context of model parameterized by weights for loss L, the standard SGD update is modified in DP-SGD to include a mechanism for privacy preservation. Specifically, the gradient V VL(0, .x;) for each training example æ; is first computed, and then processed as follows to incorporate privacy:
1. Clipping: Each gradient clipped to a maximum norm C, defined as: o = g.min(1, ||o:Ila where 9; = VL(6,
C
2. Noise Addition:
Gaussian noise is added to the aggregated clipped gradients ensure differential privacy:
B
9= B M d. +
g
where is the batch size, and or is the noise scale, determined by the privacy budget, subsampling rate, and iteration number.
The model parameters then updated using the noisy, aggregated gradient: 0 0 n/ where is the learning rate. This approach to privacy-preservin training addresses the fundamental trade-off between accuracy and privacy by controlling the granularity of the updates through the parameters C and a.
In this work, we actually use Per-layer clipping and Differentially Private Adam (DP-Adam) instead of standard DP-SGD. The key distinction of the per-layer DP-SGD compared to standard DP-SGD lies in its approach to clipping gradients layer by layer and incorporating noise accordingly. (While there are various adaptations of per-layer DP-SGD, we focus on the simplest format that directly extends from the standard DP-SGD.) [Bu et al. 2023a have demonstrated that per-layer clipping not
13

---

only matches the accuracy of global clipping but also significantly enhances memory and throughput efficiency. While DP-Adam incorporates the same mechanisms for gradient clipping and noise addition as described for DP-SGD, it also leverages the adaptive learning rates characteristic of Adam. The detailed algorithms can found Algorithm
Algorithm 2 Common Gradient Processing in DP-SGD and DP-Adam
Require: (;): Loss function for parameter 0 and input 2;
Require: C: Clipping threshold
Require: Noise scale
Require: B: Batch size
for
B do
2: Compute gradient: 9i = V (.)
3: Clip gradient: g = 9: min (1,
C
||
4: end for
5: Aggregate clipped gradients and add Gaussian noise: g Ef
Algorithm 3 Per-Layer Gradient Processing in DP-SGD
Require: L(0()), æ;)): Loss function for layer parameters ((9) and input
Require: C(I); Clipping threshold for layer
Require: g(); Noise scale for layer I
Require: L: Total number of layers
1: for each layer = to L do
2: for to do
3: Compute gradient for layer I: 9/") VC(0())
(-0
4: Clip gradient for layer 1: g; Di min (1, 19
5: end for
6:
Aggregate clipped gradients for layer and add Gaussian noise: z()) (0,, (o(9
1
> 9:
B
7: end for
Algorithm 4 DP-SGD Specific Steps
Require: 0: Model parameters
Require: ry: Learning rate
1: for each training step do
2: Perform common gradient processing as in Algorithm
2
3: Update model parameters: / +- ng
4: end for
B.2 Transformers
The transformer architecture, proposed by Vaswani et al. (Vaswani et al. 2017). is predicated on self- attention mechanisms that process input tokens in parallel, significantly improving the performance and training efficiency of pepuence-to-sequene tasks. This architecture has become the backbone of LLMs.
In transformer model, the input tensor X of size only focus on text data as the input), where B is the batch size, T is the sequence length (number of tokens), and the embedding size of a token, undergoes series transformations through multi-head self-attention and feedforward neural network blocks. For each token in the sequence, the transformer computes weighted sum all tokens in the input, where the weights are determined through the self-attention mechanism.
P (since we are considering LLM, we
14

---

Algorithm 5 DP-Adam Specific Steps
Require: m, 0: Estimates of the first and second moments (initially 0)
1: for each training step do
2: Perform common gradient processing as in Algorithm 2
3: Update moment estimates: m <- 5.m + (1 81)9
4: t - 32v (1 32)g²
5: Compute adaptive learning rate: ń = (v
6: Update parameters: 0 t- 0 fpm
7: end for
Multi-Head Attention (MHA). The attention mechanism is primarily built upon linear transforma- tions where the query Q, key K, and value V matrices are obtained as follows:
Q = XWQ. K K K, V XWW
where We. Q. WK K. and Wv are the weight matrices that are subject training.
Feedforward Network (FFN), The FFN in the transformer consists of two linear transformations with ReLU activation in between:
FFN(x) ReLU(x W
Here, W, and W2 are the weight matrices, all of which are trainable parameters of the linear layers within the FFN.
Layer Normalization (LN). LN is applied post-attention and FFN in each layer of the transformer. It normalizes the output of each neuron to have a mean of zero and a variance of one, which are then scaled and shifted by the trainable parameter vectors 7 and f, respectively:
x --
LayerNorm(x) = o
+
V +
where j and o² are the mean and variance calculated over the last dimension of the input tensor x, small constant added for numerical stability, and denotes element-wise multiplication. The layer normalization parameters (scale) and /3 (shift) are learned to optimally scale and shift the normalized data.
The key trainable parameters in the transformer model are:
Weights of the WHA mechanism, including query Wo. key WK. and value Wy matrices, each of size P P.
2. Position wise FFN weights W, of size H and Wa of size x P, where H is the hidden layer size.
3. LN parameters y and B, which are vectors of size P.
It is important highlight that the bulk of the trainable parameters in the transformer model stems from MHA and FFN modules, both of which consist of linear transformations. These linear parameters responsible for the significantly contribute to its parameter count. In contrast, the trainable parameters LN represent relatively smaller portion of the model's total parameters. Therefore, we focus on the linear parameters gradient computation.
majority of transformations within the transformer and
DP-SGD for Training Transformers. The process of udapting DP-SGD transformers is formalized as follows: For each batch of input data X and corresponding loss function L, compute the par-sample gradients Gø for all trainable parameters (Wq, Wk. K, Wv, W1, 33}3
G VoL X) E R Bx|0|
where VoC X) denotes the computation gradients of the loss with respect to the parameters for the batch X.
15

---

B.3 GPU Architecture and CUDA Programming
High performance in deep learning, particularly operations like General Matrix Matrix Mul- tiplication (GEMM), is largely attributable to the parallel processing power of modern Graphics Processing Units (GPUs). The architectural design GPUs, with their numerous cores and hierarchi- cal memory systems, is optimized for the parallel execution of operations, making them ideal for the matrix-intensive computations required in neural network training.
GPU Architecture. At the heart of GPU's computational efficiency are its Streaming Multiprocessors (SMs), which are essentially multiprocessor units that execute a large number threads concurrently, Each SM is powerhouse of performance, containing set of processing cores and a block of on-chip memory, primarily Shared Random Access Memory (SRAM), which includes registers and shared memory. Shared memory, an ultra-fast SRAM, allows threads within the same block to exchange data without involving the slower global memory (HBM), thus acting as crucial facilitator for matrix blocking.
CUDA and GEMM. The quintessential challenge in optimizing GEMM lies in the meticulous orches- tration of data movement and computation, an endeavor where matrix blocking emerges as a pivotal strategy. Leveraging the robust architecture GPUs and the sophisticated abstractions provided by CUDA (Compute Unified Device Architecture), matrix blocking transforms the theoretical prowess of parallel computation into a practical performance paradigm.
Principles of Matrix Blocking, Matrix blocking, also known as matrix tiling, a technique inge- niously conceived to enhance data locality and parallelism. It systematically partitions extensive matrix operands into smaller, manageable sub-matrices or "blocks' that can be independently dis- patched the GPU's SMs. The judicious use of shared memory within SMs for these blocks reduces the frequency and volume global memory accesses, a common bottleneck due to its higher latency. Blocking pivotal in minimizing the communication overhead between the slow global memory and the fast but limited on-chip shared memory. This stratagem leverages the temporal and spatial locality by reusing data within the fast-access memory hierarchies, significantly reducing the volume of data shuttled to and from the global memory, thereby enhancing the computational throughput.
Mathematical Formalization of Blocking GEMM. Consider the GEMM operation defined as C = A B, where A E mxx, B E R=x,, and the resultant matrix C E mRxp, Blocking decomposes this operation into smaller, tractable computations over blocks such that:
N
Σ Aik x Baj+
(5)
where N the number blocks, and each Cij. Aiss and B&j represents a sub-matrix or block within C, A, and B. respectively. The indices j, and k denote the specific block within the partitioned matrices.
The dimensions of each block are chosen based on the GPU's shared memory constraints and the size of the SMs' thread blocks, enabling optimal utilization of resources. These dimensions are represented as Bn, x B, for A₁ and B, B, for Bæj, leading to a block Be in size of Bm Bp for Ctj. Hence, the computational paradigm shifts to:
Bo
Σ (Bax x BB,)),
(6)
where each multiplication within the summation is an independent block-level GEMM that can be exccuted in parallel.
C Details of Training Workflow
C.1 Non-private Training Workflow
In the standard training regime without privacy constraints, the linear forward operation takes an activation tensor X E BxTxP and a weight matrix W E RDxP, producing an output Y E B BxTxD according the matrix multiplication Y T.T, where B, T, P, and D indicate the batch
16

---

size, sequence length (token length), feature dimension of input activation tensor X, and feature dimension of output activation tensor Y. respectively.
During the backward pass, the gradient of the output with respect to the loss, denoted by Vy E RBxIxD, is computed be of the same dimensions as the output tensor Y. Subsequently, the gradient with respect to the weight matrix W, denoted by Vw E "xP", is obtained by summing the product of the transpose of the gradient tensor of each batch item and the corresponding input tensor, expressed as Vw EB B (similar for other notations).
X, where 2B represents the summation along the dimension
Figure 6 (a) illustrates the computational workflow for the forward and backward pass a linear operation within this conventional training framework. As shown in the figure, the activation tensor X and the weights W reside in HBM, which allows for rapid parallel access and is typically used for storing larger datasets and model parameters during GPU computations. The intermediate dot products and summations are handled using SRAM, shown in orange, which is faster than HBM and suitable for storing temporary, small blocks of data during computation. This setup minimizes memory access time and maximizes throughput.
Clarification on Gradient Formulations and Reviewer Feedback. To clarify the gradient compu- tation in our NonDP baseline and address a reviewer's concern, we now present two mathematically equivalent formulations:
Format 1 (Used in Our Implementation), The default in frameworks like PyTorch is use batched matrix operations. Let Vy E BBxxx be the gradient of the output and X E IR Rxxxx be the input activation. The weight gradient Vw E PxP is computed as:
Vw = (Vr) X
This corresponds to the batched GEMM routine invoked during loss. backwarc () and does not require computing or storing per-sample gradients.
Format 2 (Shown in Figure oa for Comparison Only). For structural alignment with the DP workflows, we also illustrate an equivalent formulation that computes per-sample gradients and then aggregates them:
B
G1 (V ((ll
VW
Σ (b))
(8)
This version is shown in Figure (a) only to highlight architectural differences across methods. We reiterate that this is not used in our actual NonDP implementation.
(a)
Summary. We emphasize that our implementation of the NonDP baseline strictly uses Format 7 (batched GEMM) and does not compute or store per-sample gradients. The inclusion of per-sample nodes in Figure a) is purely illustrative and will be clarified in the revised caption and main text.
C.2 Explicit DP-SGD Workflow
Figure 6 (b) terms the explicit method (e.g., Opacus, FastClip), demonstrates the traditional DP approach where per-sample gradients are stored explicitly, resulting in increased memory usage due to the retention of individual gradient information for noise addition and clipping. The explicit DP-SGD workflow is normally organized into four distinct stages to ensure adherence to privacy constraints:
Stage Per-sample Gradient Computation. this initial stage, the activation tensor X E R BXTxP and the output gradient tensor Vy BBXTTD are loaded in blocks from the HBM to the on-chip SRAM. The pam-l gradients tensor G e BxF is computed by performing the operation G = Ex V directly on the SRAM to minimize latency, effectively implementing batched GEMM operation, where each slice of G is per-sample gradient. After computation, the per-sample gradients are written back to the HBM for further processing.
Stage 2: Gradient Norm Computation. The computed per-sample gradients G are again loaded into SRAM in smaller blocks. The norm of per-sample gradient is then computed on-chip, |[G]] ED p Σn p G e B Then, this norm calculation is stored in HBM.
17

---

Stage 3: Gradient Clipping, This stage involves loading both the per-sample gradients G and its norm |G]| from the HBM into SRAM. The clipping operation performed by computing G' = G/ max
1,
(this division occurs in dimension B), ensuring that each gradient's norm |
does not exceed the clipping threshold C. The clipped gradients G' are then stored back in HBM.
Stage 4: Noise Addition and Aggregation. In the final stage, the clipped per-sample gradients is added to each, according to the specified noise scale a. This process ensures differential privacy by obfuscating the contributions of individual training examples. The noisy, aggregated gradient for the weight update, Vw En G' + W00, C²),, is computed and then written to HBM, ready for updating the model parameters.
G' are loaded into SRAM, and Gaussian noise (0,0 c²C
Limitations. Standard DP-SGD requires the explicit storage of per-sample gradients in HBM, which is crucial for computing the gradient norms needed for clipping. This requirement substantially increases the memory footprint. This method becomes impractical for LLMs, which have large model parameters and gradients due to extended sequence lengths. The extensive memory needed to store these gradients often exceeds the available HBM capacity, leading to frequent data swapping between memory and processing units, which severely slows down the training process. Crucially, the computation of gradient norms breaks down standard kernel fusion strategies, preventing the efficient integration of gradient computation and subsequent processing steps into a single operation, resulting in increased latency and inefficient GPU utilization.
C.3 Implicit DP-SGD Workflow
Figure (c) illustrates the implicit method (e.g., GhostClip, BK), which optimizes the DP-SGD process by recalculating gradients in a fused manner, thereby avoiding the explicit storage of per- sample gradients. This approach reduces memory demands but introduces computational redundancy due to multiple gradient recalculations. The implicit DP-SGD workflow is normally organized into two distinct stages:
Stage 1: Fused Computation (corresponds to Stage 1-3 of the explicit method). In the implicit method, stages 1 through 3 of the explicit method are executed in fused computational process. This involves loading the activation tensor X E R BxTxP and the output gradient tensor Vy E IR BxTx into SRAM. The per-sample gradients tensor G E Bx P is recalculated by integrating gradient computation, norm calculation, and clipping into a single pass. This minimizes latency and avoids repeated data transfers to HBM. During this fused operation, the per-sample gradient norms calculated ||G]] directly on the chip. Clipping is simultaneously performed by scaling the gradients:
G' = max (1, | where C' is the clipping threshold. These operations are performed without storing the intermediate states, reducing the memory footprint.
Stage 2: Noise Addition and Aggregation (corresponds to stage 4 of the explicit method), The clipped gradients G' are recalculated and loaded into SRAM where Gaussian noise N (22)) is added, adhering to the specified noise scale a. The final aggregate gradient then computed and written back to HBM for the model update.
Limitations of Implicit methods: Implicit methods attempt to mitigate the high memory usage by segmenting the gradient computation and clipping it into several smaller, manageable tasks. However, these methods involve multiple recalculations of the per-sample gradients, which is computationally expensive.
D Additional Experiments Settings and Explanations
D.1 Experiments Settings
Our evaluations mainly focus on memory usage (MB) and througbput (tokens/sec) to determine the efficiency. We also show the loss of the validation data to measure the utility of private pre-training. Unless specified otherwise, the settings for each experiment use GPT-2 models with sequence length of 1024, and Ilaaa models with a sequence length of 2048, employing the AdamW optimizer as the base
18

---

Batch Size & Micro Batch Size For the batch size experiment, we vary the batch sizes at 1, 2, 4, and 8, using GPT-2 models of small, medium, and large scales to test the method's s scalability and efficiency. Similarly, in the micro-batch size experiment, we set the micro-batch sizes at 1, 2, 4, and 8, with a gradient accumulation step of
Experiments on Testing Utility We conduct an experiment to evaluate the performance of the GPT2-small model trained from scratch using DP-SGD and FlashDP under differential privacy constraints, with epsilon values set at 0.2, 0.5, and 0.8. The model is trained on the Fineweb-edu (Lzzkkov et al.) 2024) dataset. Key hyperparameters include a total batch size of 524,288 tokens, micro batch size per device of 32, and a sequence length of 1024. We use maximum learning rate of 6 x 10 4 and a minimum learning rate of 6 x 10-5, with weight decay set at 0.1 and gradient clipping at 1.0. The model undergoes training with validation frequency every 250 steps and model saving every 5000 steps, using both DP-SGD and FlashDP, enabling differential privacy with delta set at 1 10-6 and a clipping threshold of 100. The training aims to compare utility across different privacy levels and analyze the trade-offs between privacy and utility. We use the validation loss as the evaluation metric in Table
Distributed Training DDP involves distributing the model's parameters across several devices, and each device computes gradients for a subset of the data independently. This method is beneficial for managing models that fit within the memory limits of a single GPU but need faster processing through parallel execution. On the other hand, Pipeline Parallel (PP) splits the model's s layers across different devices, allowing different parts of the model to be processed simultaneously. PP is particularly useful for very large models that exceed the memory capacity of individual GPUs, enabling concurrent processing of different stages of the model across the pipeline. The experiments with DDP and PP are designed evaluate the effectiveness of FlashDP in a distributed training context, assessing its performance in terms of memory usage and throughput across various model sizes and batch sizes. These experiments are critical to demonstrate that FlashDP can maintain its efficiency and scalability when applied to state-of-the-a LLMs, which require substantial computational resources and sophisticated training mechanisms to manage their size and complexity.
In this setup, we explore the scaling capabilities of FlashDP using DDP on four A100 GPUs (80GB each) by training GPT-2 models of small, medium, and large sizes with fixed sequence lengths of 1024 and varying batch sizes of 8, 4, and 2. Additionally, PP experiments are conducted on Llama models of sizes 3B, 7B, and 13B evaluate throughput and memory efficiency across different stages of the model pipeline. It is important to note that GhostClip and BK do not support the distributed modes we used.
D.2 Additional Explanations
GhostClip initially supports only global clipping: however, it can be easily adapted to per-layer clipping as outlined in Algorithm 6
E More Experimental Results
E.1 Results of Micro Batch Size
Table 2: Micro Batch Size Analysis. Comparing memory and throughput at varying micro batch sizes B (1, 2, 4, 8) and the same gradient accumulation steps (4) for GPT-2 sizes with differential privacy methods under consistent settings with Table [
Model
B
NonDP
Opaeis
Memory Usage (MB xio4)
GhoeeClip
BK
RashDP
NonDP
Opscus
Throughput (bokensfiee Klo4)
Chost
BK
FlshDP
GPI2-mmll
0.31
(991)1330)
(15Tcc
(.
0.5((11200)
1,3000011)
1,7550337
[
GPI2-medium
1.26
169x1134)
1.25(x0.9
1.8101144)
1,26
0.610x0.48)
0.45(x0.35)
0.86800.68
0.91(x0.72)
GPT2-large
2.48
3.64xx1.47)
2446x0.99)
3.216x1.29)
2.48tx1.00)
0.67
0.39xx043)
0.32 46)
0.47(x0.69)
0.53rx0.89
GPT2-small
1
0.87
1155x1.32)
1.0x.115)
1066x1.22 22)
0.87xx1.00
3.22
168650.52)
0.92150229)
L916x0.39)
2.32x0072)
GPT 2-F2dim
2
2.07
2380x1.39)
2.01(00.97)
2.62(1127))
(0)(x1.00)
1.38
0.88/x0.64)
0.65(10.7))
0.8(10064)
1.04v(1.75)
GPT2-large
2
3.91
6.7)x1.55)
3.83x0).9)
3.91(x1.00) 4.43xx1.1)
0.46x00622
0.43(0.58)
0.49(x0.66
0.5%v.1.80
1.53
2.10x1.37)
148(x0.97)
1.53tx1.00 1.73(x1.13)
2.49)x0.67)
L.50:30.4)
2.300x0.62)
2.50x0.70
GPT2 mediom
4
3.58
5.516x1.54)
3.46(x0.97)
1.04(x1.1)
3,58tx) 00)
1.48
0.97(x0.66)
0.86150.59 58)
0.99(s0.67)
1291x0.87)
GPT2-lage
4
6.60
6.601x1 00)
0.79
0.551x0.67)
0.65(20.82)
GPT2-small
8
2.86
4.000x140)
2.78(x0.97)
3.06)x1.0))
2.86 00)
3.57
2.600;0.67)
(.99:00.51)
2.441x0.63)
2.73xx0.11
GPT2 medium
6.60
6.37(x0.97)
7.161x1.08)
6.60(x1.00)
1.55
1.03xx0.66)
(05(x0.65)
1.19(x0.77)
GPT2-large
19

---

Algorithm 6 Per-Layer GhostClip
Require: C(0()) ,::): Loss function for layer parameters ((0) and input oy
Require: C(); Clipping threshold for layer
Require: o(); Noise scale for layer |
Require: L: Total number of layers
for each layer
L do
2: for i = to B do
Compute gradient norm for layer 1: ol || (VC( ,2))] by first computing per-sample gradient in-place then computing per-sample norm.
4: Clip gradient for layer I: 9, g: min
(()
1, lls 1|>
c())
by re-computing per-sample gradient
9i in in-place.
(1)
5: end for
6: Aggregate clipped gradients for layer and add Gaussian noise: g() 1 -E= 9: B
(0)
(0, ((0))))))
7: end for
Table further explores the impact of varying micro batch sizes, crucial factor for managing memory in constrained environments and optimizing the use gradient accumulation steps. FlashDP consistently displayed minimal memory footprint increases and maintained high throughput efficiency, even as micro batch sizes increased. For example, at a micro batch size for the GPT-2 medium model, FlashDP's memory usage 6.49 x 10 MB marraiainllly higher than its usage at smaller micro batch sizes and significantly lower than Opacus at the same size. This robust performance underscores FlashDP's effective management of memory, which is essential for scaling up the training of large models without excessive hardware requirements.
To be specific. 1) Opacus showed consistent increase in memory usage as micro batch sizes increased, which is indicative of its inefficient memory handling under fragmented gradient computations. 2) GhostClip, while better in memory usage compared Opacus, didn't scale well in throughput, which decreased noticeably with larger micro batches, reflecting the computational cost of gradient recaleulations. 3) BK displayed trends similar Opacus but generally used slightly less memory and provided slightly better throughput, suggesting a more optimized handling of gradient accumulation steps. 4) FlashDP maintained minimal increases in memory usage with increasing miero batch sizes and consistently provided the highest throughput, highlighting its effective integration of operations within the computational workflow. To summarize, as the micro batch size increases, FlashDP's memory usage increases only slightly and still maintains the highest throughput, demonstrating its efficient memory management techniques.
E.2
Results of AMP Training Scalability
(a) Memory Usage- float16
(b) Throughput- float16
(c) Throughhut bfloat16
Figure Memory and Throughput Analysis GPT-2 Models Using Automatic Mixed Precision (AMP) Training Across Float16 and BFloat16 Precision.: (a) Demonstrates the memory usage for GPT-2 small, medium, and large models with 6 precision. (b) shows throughput using Float16 precision, and (c) shows throughput with BFloat precision.
20

---

Automatic Mixed Precision (AMP) (Micikevicius et al. 2017) training involves utilizing lower precision formats like f and 6 within a traiiing session to reduce computational demands and memory usage. This strategy is particularly valuable for large language models (LLMss,, which typically require substantial computational resources. By employing AMP, training processes can be accelerated, and larger models or batches can be managed more efficiently without proportional increases in hardware capacity. The integration of differential privacy with AMP, especially in techniques like FlashDP, is critical for exploring the practical limits of DP-SGD. This experiment assesses how FlashDP adapts to AMP settings compared to other methods, and evaluates the impact memory efficiency and processing speed, which crucial for the scalability of private training in constrained environments.
In our experiments, we analyze GPT-2 models of varying sizes using batch sizes of 8, 4. and 2 across 6 and bfloat precision formats to measure memory usage and throughput, examining FlashDDP's performance relative to NonDP, Opacus, and BK methods. It is important to note that GhostClip does not support AMP, and Opacus does not support the bfloat16 precision format.
Memory Usage Analysis. As depicted in Figure (a), the memory usage across GPT-2 models of different sizes indicates that FlashDP, when utilizing AMP in both f and 6 formats, maintains lower memory consumption compared to Opacus and BK, and closely approximates the NonDP configuration. This showcases FlashDP's effective use of AMP to minimize memory overhead, facilitating the training of large models under stringent privacy constraints.
Throughput Analysis with Floatl6 and BF In terms of throughput, Figure (b) and 5(c) present a comprehensive look at the advantages of using f and 6 precision formats under AMP. FlashDP consistently outperforms Opacus and BK in throughput metrics across both precision types. This is especially notable in larger model configurations, where the differences in throughput become more pronounced, highlighting FlashDP's capability to handle extensive computational loads efficiently. As demonstrated in Figure 5(b), FlashDP exhibits significant throughput ad advantages over the other DP methods. This performance indicative of the efficient computational optimizations that FlashDP leverages within the AMP framework. As shown in Figure while 6 typically offers slightly lower computational throughput than floatl6 due to numerical properties, FlashDP's implementation still ensures that it outperforms other differential privacy methods. This underscores FlashDP's robust performance across varying precision settings.
E.3 Results Distributed Training
Distributed Data Parallel (DDP) (Li al., 2020 and Pipeline Parallel (PP) (Kim al. 2020) are two advanced techniques crucial for scaling the training of LLMs efficiently across mu multiple GPUs or nodes.
300007
Nesbe
Nocbe
Opeooo
70000
500C
4000
30002
200C
1000
0100-22 21
1800-11 10
Lane
Lamra A0
Likin Lob
(a) Memory Usage
(b) Throughput
Figure 5: Memory and Throughput for Llama Models Using Pipeline Parallel Training. (a) Memory usage for Llama-3B, Llama-7B, and Llama-13B models. (b) Throughput in tokens per second across these model sizes. value of indicates of memory.
Distributed Data Parallel (DDP). Figure in Appendix illustrates the performance of different methods in a DDP setting across GPT-2 models of varying sizes. FlashDP showcases superior memory usage efficiency and higher throughput across all model sizes when compared to Opacus and BK. Notably, even as the model size increases, FlashDP maintains competitive edge close to the
21

---

NonDP benchmarks, highlighting its effective parameter distribution and gradient computation across multiple GPUs. This is crucial scenarios where training speed and model scalability are priorities. Pipeline Parallel (PP). In the PP scenario depicted in Figure 5 FlashDP was tested with Llama models varying from 3 billion to 13 billion parameters. The results indicate that FlashDP not only scales efficiently with increasing model size but also demonstrates significant throughput improvements compared to Opacus and BK. Particularly, FlashDP's ability to handle the largest model (Llama-133) with minimal throughput degradation illustrates its robustness in managing extensive computational loads, characteristic of PP environments.
E.4
Results of Utility
Table 3: FlashDP Pretrain Precision validation on GPT2-small with different privacy 6.
Validation loss
Method
<= 0.2
= 0.5
< 0.8
DP-SGD
4,8082
4,8063
4,8061
FlashDP
4.8082
4.8063
4.8061
In our study, FlashDP is meticulously optimized for DP-SGD, focusing on enhancing GPU I/O and system-level efficiencies without altering the fundamental algorithmic components of per-layer DP-SGD. We conducted experiments on utility with GPT-2 small to support this, whose results are shown in Table From the table, we can easily see that FlashDP demonstrates an identical validation loss that of DP-SGD across all privacy levels.
F Additional Tables and More Figures
Prooad
Nodkoadd
Ruthoge
Fovend
Maule
Opgel
Dape
Orggit
Indive
tepldh
Explicit Method (c.g Opacus, FustC lip)
(c) Implicit Method GhostClip, BK)
(4) Non-DP
(d) FlashDP
Figure 6: Comparison of different training methods. (a) Non-DP: Basic training without DP. (b) Explicit Method (e.g., Opacus, FastClip): Stores per-sample gradients G (red explicit cache), increas- ing memory usage. (c) Implicit Method (e.g., GhostClip. BK): Reduces memory by recalculating gradients in fused manners (blue dotted box) but implicitly calculating the per-sample gradient twice, causing computational redundancy. (d) FlashDP: Optimizes gradient processing by consolidating computations into a single pass, reducing redundancy and memory use.
Table 4: Comparison of Backward Propagation Methods.
Per-sample Gradient
Method
Implicit Fusion
Cache
Recalculation
Non-DP
Explicit-DD DP
Implicit-Dp
X
FlashDP
X
22

---

50000
NenDP
NanDP
Mammory Usage 0182 001-11
Throaggo www.eraaaaaa.o.oo..ooo.oo Pawrann
25000
20000
20000
15000
10000
10000
agge
aarge
(a) Memory Usage
(b) Throughput
Figure 7: Memory and Throughput for GPT Models Using Distributed Data Parallel Training. (a) Memory usage for GPT-samll, GPT-medium, and GPT-llrge models. (b) Throughput in tokens per second across these model sizes. value indicates out memory.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly describe the paper's key contributions, including the design of FlashDP for efficient per-layer DP-SGD training. GPU I/O op- timization via fused kernels, and scalability to large models like Llama- 13B. These are substantiated by both theoretical and experimental results (see Sections and 4āAş6).
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The conclusion and discussions throughout the paper acknowledge the current scope FlashDP, including its focus on per-layer clipping rather than global clipping, and reliance fused CUDA kernels, which may require hardware-specific tuning (see Section 6 and Appendix A.3).
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
23

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
The paper is primarily systems- and engineering-orrentte and does not include formal theoretical results requiring assumptions or proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper provides detailed descriptions of model sizes, batch configurations. datasets, and hardware used in experiments (see Section 5.1 and Appendix C), allowing reproduction of key results even without code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may
24

---

be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Due to ongoing code cleanup and anonymization, code is not yet released. However, we commit open-sourcing FlashDP with reproduction scripts upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so aÄIJNoâÄi is an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https wtt.:/ddeee/.oSSSemiiieieeeeeeee
/ //nips.cc/pppp
wttsssuwmssssioiiii.. ssionPolicy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question:
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
25

---

The training settings, model variants, datasets (Wikitext), and hardware (A100 GPUs) are clearly described in Section 5.1 and Appendix C. with baseline comparison details provided in Table 1.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Error bars are not reported because it would be too computationally expensive.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section 5.1 specifies that experiments were conducted on NVIDIA A100 GPUs (80GB), and various resource constraints (e.g., out-of -fmooor for Opacus) are reported in Table and Figure 3.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
26

---

9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research follows the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The conclusion section discusses the positive impact of enabling scalable privacy-preee pinaecypreserving LLMs. Potential negative societal impacts (e.g., misuse of DP-trained models) are acknowledged.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymenn of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not release pretrained models datasets with high risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
27

---

Released models that have high risk for misuse or dual-use use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All datasets and tools used are publicly available and appropriately cited (e.g..
Wikitext, PyTorch, Opacus). See Section 5.1 and References.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (c.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pataasess.paperswtt paperswi thcode. cow/coo////ooooo /atasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new datasets models are released in this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This research does not involve human participants crowdsourcing.
28

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No IRB approval is required as the study does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: While our work targets LLM training, it does not an LLM method- ological component or for content generation. No LLMs were used to produce the paper
content.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy lattps://neurre /ttp://rrrp.cccooooooerrre////// for what should or should not be described.
29


Reviews:
[{"review_id": "quHQ4qF0Fx", "reviewer": "Reviewer_AcaE", "summary": "This 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.", "questions": "1. 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.", "limitations": "Limitations discussed. No potential negative societal impact.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths\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.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognizing the importance and clarity of our work, and for the constructive feedback. Below we provide detailed responses to the concerns raised.\n\n1. Hardware extensibility\n\nWe agree that generalizing FlashDP beyond CUDA-based GPUs (e.g., to TPUs or Apple’s M-series chips) is an important and challenging task. However, this limitation is not unique to FlashDP, but rather inherent to all low-level GPU operator optimization efforts. For instance, when FlashAttention, a widely recognized operator fusion method, was first released, it only supported NVIDIA Ampere GPUs (e.g., A100), and did not work on previous-generation GPUs, let alone other hardware platforms. This did not diminish its value to the community.\n\nSimilarly, FlashDP is a CUDA-native design targeting NVIDIA architectures where per-layer DP-SGD presents clear performance bottlenecks. Extending FlashDP to non-CUDA hardware is certainly possible, but would require new backend-specific system engineering efforts. We believe that presenting FlashDP at a top venue like NeurIPS will help us attract community interest and contributions toward making such extensions a reality.\n\n2. Model architecture and relevance\n\nWhile the primary results in Table 1 are reported on GPT-2 variants, our paper also includes experiments on the LLaMA family (including LLaMA-13B), which shares a decoder-only transformer architecture with GPT models, in Figures 3 and 5. Given that the vast majority of today's open-source and production LLMs adopt similar architectures (e.g., LLaMA, GPT-NeoX, Mistral), we believe these two model families sufficiently represent mainstream pretraining setups.\n\nFurthermore, as FlashDP is implemented at the operator level, it remains fully applicable to fine-tuning and post-training settings as well. Since these stages typically reuse the same model structure and training codebase as pretraining, FlashDP can be directly integrated to speed up and reduce memory usage during private fine-tuning tasks.\n\n3. Micro issue: LaTeX character encoding\n\nThank you for pointing out the formatting issue with rendered text like “FlashDPâA˘ Zs.” These were caused by LaTeX encoding artifacts in PDF generation. We will fix all such issues in the camera-ready version.\n\n4. Operator-level DP profiling results\n\nWe appreciate your suggestion to analyze the performance of the fused DP-specific operator chain (i.e., linear backward + clipping + noise addition). To that end, we conducted an isolated operator-level benchmark comparing three setups:\n- NonDP: standard Linear.backward()\n- DP-SGD: naïve implementation of Linea per-sample backward + per-sampling clipping + noise\n- FlashDP: fused kernel covering all of the DP-SGD operators in a single pass\n\nWe vary the matrix dimension M and set P = T = D = M. The achieved TFLOPS are reported below:\n| M    | NonDP (TFLOPS) | DP-SGD (TFLOPS) | FlashDP (TFLOPS) |\n| ---- | -------------- | --------------- | ---------------- |\n| 2048 | 75             | 41              | 70               |\n| 2560 | 76             | 47              | 74               |\n| 3072 | 78             | 52              | 77               |\n| 3584 | 80             | 55              | 79               |\n| 4096 | 83             | 59              | 80               |\n| 4608 | 84             | 61              | 81               |\n| 5120 | 84             | 62              | 81               |\n| 5632 | 84             | 65              | 82               |\n| 6144 | 85             | 66              | 83               |\n| 6656 | 85             | 66              | 83               |\n| 7168 | 87             | 69              | 84               |\nThese results show that naïve DP-SGD significantly reduces GPU utilization, while FlashDP restores throughput nearly to NonDP levels. This validates the core motivation of FlashDP: fusing gradient computation, norm calculation, clipping, and noise addition into an efficient single kernel, avoiding costly memory access and redundant computation."}, {"review_id": "UI3jw2C0rU", "reviewer": "Reviewer_nQGS", "summary": "This 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.", "questions": "See the Weakness.", "limitations": "Not 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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## 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", "was_revised": true, "final_justification": "My recommendation is primarily based on the large-scale experimental evaluation, as authors evaluated  FlashDP on GPT-2 (1B+) and LLaMA (up to 13B) models, demonstrating clear memory and throughput gains over SOTA DP methods, and also comparing validation loss against non-DP baselines. While the scale and empirical results are compelling, **I did not verify the correctness of the method itself, as my expertise lies in cryptographically secure privacy-preserving machine learning (PPML) rather than differential privacy**.", "rebuttal": "We sincerely thank the reviewer for their encouraging feedback and strong support for our work. We greatly appreciate your recognition of FlashDP's ambition to support DP pre-training, and your positive assessment of its scalability, architectural clarity, and empirical rigor.\n\nYou raised an excellent point regarding the limitations of using validation loss or perplexity (PPL) as the sole measure of model utility, especially in the long-context regime. We agree that PPL improvements may not always correlate with downstream task performance, particularly due to phenomena like context dilution, as discussed in the papers you cited ([1], [2]). These observations are very insightful and valuable.\n\nThat said, the primary purpose of our utility evaluation was to confirm that FlashDP maintains full accuracy parity with standard DP-SGD, which our results do show. This alignment is expected, since FlashDP modifies only the execution efficiency of gradient clipping and noise addition, but does not alter the algorithmic behavior or statistical properties of DP-SGD. At the operator level, FlashDP and DP-SGD behave identically in both computation and output.\n\nNonetheless, we fully agree that exploring downstream task evaluations, as well as alternative utility metrics beyond PPL, are crucial for a more complete understanding of model quality under DP constraints. We will cite and briefly discuss the two references you provided in the camera-ready version, and we see this as a natural direction for future work.\n\nFor limitation, while FlashDP currently focuses on CUDA-optimized implementation, a key direction for future work is expanding support to other hardware backends such as AMD GPUs or custom AI accelerators. This would help further democratize differentially private LLM training. We believe this is primarily an engineering effort and can benefit from community collaboration as adoption grows. We sincerely apologize for the oversight—although we referred to Appendix A.3 in the checklist, we inadvertently omitted the actual discussion of limitations in the appendix. We will correct this in the camera-ready version."}, {"review_id": "MsGqL9TGXN", "reviewer": "Reviewer_Sx9F", "summary": "This 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.", "questions": "* 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", "limitations": "I don't see discussions about limitations and the potential negative societal impact of their work.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "The rebuttal has addressed my concerns. Therefore, I have increased my score from 3 to 4. The paper would be clearer to readers with the added clarifications on the generalizability of the implemented CUDA kernels, as discussed.", "rebuttal": "We sincerely thank the reviewer for the valuable feedback and questions. Below we provide our point-by-point response.\n\n- Q1: \"Tables 1 and 2 have missing numbers.\"\n\nSome entries in the tables are missing. As clearly indicated in the table header, these are due to GPU memory limitations. In such cases, the baseline methods fail to run due to out-of-memory (OOM) errors. Rather than omitting these settings altogether—which we felt might obscure important limitations—we chose to include them with missing results to indicate the limitations of certain methods under constrained GPU memory.\n\n- Q2: \"Could the authors explain why the results of DP-SGD and FlashDP match so well?\"\n\nFlashDP is designed to improve the GPU efficiency of DP-SGD operators without modifying their mathematical semantics. Therefore, it does not introduce any accuracy degradation by design. In our experiments, to ensure consistent and fair comparison, we used the same random seed across both DP-SGD and FlashDP. This configuration ensures that the observed matching in validation loss is expected and repeatable. We will make this setting explicit in the camera-ready version.\n\n- Q3: \"Since some of the baselines (such as GhostClipping) use full model clipping (instead of per-layer clipping)...\"\n\nWhile the reviewer raises an important point regarding the potential impact of clipping strategy on model utility, prior work has shown that per-layer clipping can achieve comparable accuracy to global clipping, especially in large-scale models. Specifically, as cited in our submission (Lines 164–165), He et al. [1] demonstrate that per-layer clipping matches the utility of global clipping in various DP training setups. This suggests that our choice of per-layer clipping is both practically effective and theoretically grounded.\n\n[1] Jiyan He, Xuechen Li, Da Yu, Huishuai Zhang, Janardhan Kulkarni, Yin Tat Lee, Arturs Backurs, Nenghai Yu, and Jiang Bian. Exploring the limits of differentially private deep learning with group-wise clipping. arXiv preprint arXiv:2212.01539, 2022.\n\nTherefore, the use of per-layer clipping in our experiments is unlikely to introduce utility degradation relative to global clipping, and it aligns with trends in scalable DP training practice. We will clarify this more explicitly in the camera-ready version to avoid confusion.\n\n- Q4: \"Since the method requires custom CUDA kernels for the backpropagation process, we will need different implementations for different network layers...\"\n\nFirst, in today’s large models—especially transformer-based LLMs—the vast majority of parameters reside in linear layers. Since DP-SGD operates on per-parameter gradients, it is most effective to optimize the gradient processing for layers with the highest parameter count. In contrast, optimizing smaller layers (e.g., LayerNorm) yields limited practical gains.\nSecond, the reviewer mentioned concerns about input/output shape variation. Our algorithm is independent of specific input shapes; as shown in Algorithm 1, the logic does not rely on any shape-specific assumptions.\nLastly, while convolutional layers are indeed common in vision models, they can be losslessly converted into linear operations via im2col, which is already a standard practice in libraries such as cuDNN. In fact, our FlashDP implementation already supports convolutional layers internally, though we did not include related experiments as this work focuses on LLMs.\n\n---\nFurther explanation: convolution workflow via im2col / col2im:\n\n- Forward: Convolution as GEMM via im2col\n\nGiven:\n- Input tensor: $X \\in \\mathbb{R}^{C_{\\text{in}} \\times H \\times W}$\n- Weight tensor: $W \\in \\mathbb{R}^{C_{\\text{out}} \\times C_{\\text{in}} \\times K_H \\times K_W}$\n- Output tensor: $Y \\in \\mathbb{R}^{C_{\\text{out}} \\times H_{\\text{out}} \\times W_{\\text{out}}}$\n\n1. Flatten input patches via im2col (no time cost):\n$$\nX_{\\text{col}} \\in \\mathbb{R}^{(K_H \\cdot K_W \\cdot C_{\\text{in}}) \\times (H_{\\text{out}} \\cdot W_{\\text{out}})}\n$$\n\n2. Flatten weights (no time cost):\n$$\nW_{\\text{row}} \\in \\mathbb{R}^{C_{\\text{out}} \\times (K_H \\cdot K_W \\cdot C_{\\text{in}})}\n$$\n\n3. Matrix multiplication (linear forward):\n$$\nY_{\\text{col}} = W_{\\text{row}} \\cdot X_{\\text{col}} \\in \\mathbb{R}^{C_{\\text{out}} \\times (H_{\\text{out}} \\cdot W_{\\text{out}})}\n$$\n\n4. Reshape $Y_{\\text{col}}$ to obtain output $Y$\n\n- Backward\n\nGiven:\n- Loss gradient w.r.t. output: $\\frac{\\partial \\mathcal{L}}{\\partial Y} = G_Y \\in \\mathbb{R}^{C_{\\text{out}} \\times H_{\\text{out}} \\times W_{\\text{out}}}$\n- Flatten: $G_Y^{\\text{col}} \\in \\mathbb{R}^{C_{\\text{out}} \\times (H_{\\text{out}} \\cdot W_{\\text{out}})}$\n\n1. Gradient w.r.t. weight (the backward pass of the linear layer for weight gradients, where FlashDP is applied):\n$$\n\\frac{\\partial \\mathcal{L}}{\\partial W_{\\text{row}}} = G_Y^{\\text{col}} \\cdot X_{\\text{col}}^\\top\n$$\n\n2. Gradient w.r.t. input (linear backward for input gradient computation):\n$$\nG_{X_{\\text{col}}} = W_{\\text{row}}^\\top \\cdot G_Y^{\\text{col}}\n$$\n\n3. Recover input gradient via col2im (no time cost):\n$$\n\\frac{\\partial \\mathcal{L}}{\\partial X} = \\text{col2im}\\left(G_{X_{\\text{col}}}, \\text{input shape}, K_H, K_W, \\text{stride}, \\text{padding}\\right)\n$$\nHere, $\\text{col2im}(\\cdot)$ denotes the inverse transformation of $\\text{im2col}$: it maps the unfolded gradient matrix $G_{X_{\\text{col}}}$ back to the original input tensor shape by aggregating overlapping gradients from sliding windows into their corresponding spatial positions.\n\nAs shown in the equations above, convolution operations can be equivalently expressed as matrix multiplications preceded by im2col and followed by col2im. These two transformations are deterministic, involve no learnable parameters, and do not perform any computation beyond memory reshaping and indexing. In practice, their overhead is negligible—especially on modern accelerators where such data layout operations are heavily optimized. Therefore, from an implementation perspective, supporting convolutions does not require new CUDA kernels; our existing linear-layer kernels can be reused directly after applying these standard transformations. This also highlights the modularity of our design: once DP operators are optimized for linear layers, they naturally extend to convolutional layers through standardized tensor reshaping procedures.\n\n---\n\n- Q5: \"Do the authors plan to release the code?\"\n\nWe confirm that we will release the full FlashDP codebase with the camera-ready version.\n\n- Q6: \"Line 286: 'the grid synchronization required by CG necessitates launching all blocks simultaneously, which is impractical for DP applications.'...\"\n\nCG indeed allows global synchronization across blocks, but they require all blocks to be launched simultaneously, which imposes strict hardware constraints on grid size and shared memory usage. In the setting of DP training, where each sample must undergo individual gradient clipping and noise addition, high parallelism is essential to maintain training throughput. We will elaborate on this in Section 4.2 of the revised version.\n\n- Q7: \"Table 1: Why does GhostClip have even less memory usage than Non-DP?\"\n\nThis is due to an implementation detail in the official GhostClip code. Specifically, in line 503 of privacy_engine.py, the authors explicitly delete all intermediate variables except .summed_grad with the comment: “Aggressive memory saving — delete everything except .summed_grad to save memory!” While this technique indeed reduces memory usage, it significantly harms runtime performance due to recomputation. We chose not to adopt this strategy, as FlashDP aims to balance both memory efficiency and high throughput.\n\n- Q8: \"Typos...\"\n\nWe acknowledge the typographical and compiling issues (e.g., “FlashDPâA˘ Zs”) and will correct them in the camera-ready version."}, {"review_id": "JnlqZbtDVd", "reviewer": "Reviewer_JGX8", "summary": "This 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.", "questions": "# 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\"?", "limitations": "As 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.\"", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "* 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.", "was_revised": true, "final_justification": "Private comment to the AC has been made before. For many rounds in the reviews, my main argument was the computational scope of the paper instead of the machine-learning angle. It seems that two versions of FlashAttention have been published in this conference before. Therefore, I have increased the score.", "rebuttal": "We appreciate your thoughtful review and detailed questions. Below, we address each of your points in turn.\n\n- Q1: \"Is this a common approach to improve code by lumping ops into a custom GPU kernel? What is the novelty of this paper? To what other problems would this custom kernel be applicable? Would an automated tool be able to generate such a fused kernel for the same problem?\"\n\nOperator fusion is indeed a widely used optimization technique for improving GPU utilization, especially for workloads that involve multiple consecutive tensor operations. However, in practice, operator fusion remains largely domain-specific and often needs manual tailoring to the computational patterns of a specific application.\n\nIn the case of FlashDP, we specifically target the gradient computation workflow under DP constraints, which introduces a unique and challenging fusion scenario:\n1. Per-sample gradient computation involves a matrix multiplication, followed immediately by a norm-reduction step that requires aggregation across different axes.\n2. This disrupts typical fusion patterns because the intermediate output of the matmul needs to be preserved and reused across threads or even blocks for the subsequent reduction.\n3. If these intermediate results are stored in HBM, both memory usage and latency increase significantly.\n\nTo address this, FlashDP preserves the blocking pattern used in matrix multiplication and introduces a two-stage hierarchical reduction (intra-block and inter-block), enabling norm computation without writing the intermediate matmul output to HBM. This design is tightly coupled with GPU architecture, particularly SRAM usage and block scheduling.\n\nCurrent automated operator fusion tools (such as torch.compile) are not aware of DP-specific semantic constraints, particularly the privacy-critical separation of per-sample gradients and the cross-sample norm reduction logic. As a result, they cannot discover the optimal fusion strategy we implemented in FlashDP. Without this DP constraint, standard workloads would allow straightforward fusion and scheduling, which automated compilers can often handle effectively.\n\nIn short, FlashDP highlights a novel fusion challenge unique to DP training, and provides a solution that is not yet supported by automated tooling.\n\nAlso, FlashDP is not limited to DP-SGD. It is a general backend for any norm-bound-clipping-based training algorithm. As long as per-sample norm computation is needed, the idea from FlashDP can provide the same memory and compute benefits.\n\n- Q2: \"Quote from conclusion: 'At the same time, it highlights the need for responsible release practices to mitigate potential misuse under the guise of privacy.'...\"\n\nThank you for pointing this out. This sentence was added in response to the NeurIPS Paper Checklist requirement to consider potential risks of misuse and to discuss responsible mitigation practices. Specifically, our concern is that as tools like FlashDP make large-scale DP training more efficient and accessible, it becomes easier for practitioners to claim differential privacy without rigorous accounting or transparency (e.g., not reporting ε/δ, noise scale, or clipping strategy).\n\nThus, this sentence is intended to highlight that performance gains should not come at the cost of accountability, and any deployment of DP-trained models should still follow sound privacy analysis and transparent reporting.\n\nWe agree that this statement could be made more precise and will revise it in the camera-ready version for clarity and context.\n\n- Q3: Small observations / formatting issues\n\nThank you for noting the minor formatting and citation issues. Specifically:\n1. \"Line 279...\": We will fix the LaTeX rendering bugs such as “FlashDPâA˘ Zs” and “PyTorchâA˘ Zs” which were caused by incorrect character encoding.\n2. \"Use \\citet{} instead of \\citep{}...\": For citations like “Opacus (Yousefpour et al., 2021)” and “(He et al., 2022)”, we agree that \\citet{} is preferable when citing as part of the sentence. We will fix these instances.\n3. \"Implicate -> implicit...\": We will correct “implicate” to “implicit” in the phrase “explicit and implicate methods of DP-SGD”.\n\nAll of these corrections will be addressed in the camera-ready version."}, {"review_id": "50Er4RrFUN", "reviewer": "Reviewer_bYKW", "summary": "This 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.", "questions": "- 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?", "limitations": "The 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)", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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')", "was_revised": true, "final_justification": "I'm not fully convinced that the paper will be of sufficient interest to the DP community, because its focus diverges from the common practices (DP-LoRA, global clipping, etc). \n\nThat said, I agree that the paper presents a solid piece of work, and wouldn’t object if other reviewers would like to champion acceptance. I also appreciate the authors’ detailed response, hence have increased my score.", "rebuttal": "We thank the reviewer for the thoughtful and DP-focused feedback. Below, we respond to each point raised:\n\n- Q1: \"The paper showcases FlashDP's effectiveness in DP pre-training for LLMs...\"\n\nWe agree with the reviewer that DP fine-tuning is of critical importance, particularly for protecting user-specific or private downstream data. However, we also argue that DP pre-training is equally important. As already cited in our submission, several recent works [1, 2, 3] have emphasized the necessity and feasibility of applying differential privacy to the pre-training stage of large models, especially in the context of LLMs. Also, pre-training on large, publicly available corpora poses risks related to copyright and data licensing, and differential privacy offers a principled mechanism to protect against unintended memorization or redistribution of such content. Recent work has also explored DP pre-training in this context.\n\n[1] Xuechen Li, Florian Tramer, Percy Liang, and Tatsunori Hashimoto. Large language models can be strong differentially private learners. In International Conference on Learning Representations, 2022.\n\n[2] Zhiqi Bu, Jialin Mao, and Shiyun Xu. Scalable and efficient training of large convolutional neural networks with differential privacy. Advances in Neural Information Processing Systems, 35: 38305–38318, 2022.\n\n[3] Zhiqi Bu, Yu-Xiang Wang, Sheng Zha, and George Karypis. Differentially private optimization on large model at small cost. In International Conference on Machine Learning, pp. 3192–3218. PMLR, 2023b.\n\nFurthermore, our work does not exclude the use of FlashDP in fine-tuning settings. In fact, FlashDP can be directly applied to full-parameter fine-tuning, since the training flow is nearly identical to pre-training.\n\nAs for low-rank fine-tuning methods like DP-LoRA: FlashDP's optimization strategy—focusing on fused per-layer weight gradient computation for linear layers—can theoretically be adapted to LoRA. However, LoRA’s trainable weights are typically very small, and the DP-induced overhead on such layers is negligible in practice. In other words, optimizing DP-related operators in DP-LoRA yields much smaller performance gains compared to optimizing them in full-parameter training, since the privacy-related computation is already lightweight in DP-LoRA. This is, in fact, one of the main reasons for the popularity of DP-LoRA: it sidesteps the performance and memory concerns of DP-SGD altogether. Thus, we intentionally focus FlashDP on full-parameter updates where performance constraints are significant.\n\n- Q2: \"FlashDP's design and optimizations focus on per-layer clipping...\"\n\nWe acknowledge that global clipping is commonly used in practice. However, per-layer clipping has also gained significant traction recently, particularly in LLM-scale DP training, due to its superior memory efficiency and empirical performance. FlashDP is designed to serve this need effectively. Moreover, as cited in our submission (Lines 164–165), recent work [4] has shown that per-layer clipping can achieve comparable utility to global clipping. This further supports the practicality of adopting per-layer clipping as a viable alternative in large-scale DP training.\n\n[4] Jiyan He, Xuechen Li, Da Yu, Huishuai Zhang, Janardhan Kulkarni, Yin Tat Lee, Arturs Backurs, Nenghai Yu, and Jiang Bian. Exploring the limits of differentially private deep learning with group-wise clipping. arXiv preprint arXiv:2212.01539, 2022.\n\nThat said, we agree there is value in supporting both paradigms. In fact, we are actively developing FlashDP-v2, a new version that extends our optimization techniques to support global clipping. This extension introduces considerable system-level complexity (e.g., asynchronous kernel launching, CPU-GPU overlap, synchronization management). These system-level complexities make it challenging to present FlashDP-v2 and FlashDP within a single paper, so we have decided to describe them in two separate works.\n\nIn our view, it is reasonable and modular for different backends to support different clipping modes. Practitioners can simply use FlashDP for per-layer clipping, and FlashDP-v2 (forthcoming) for global clipping—both are independent DP-compatible operator modules. We will note this as future work in the camera-ready version.\n\n- Q3: \"While the throughput gains are impressive and clearly demonstrated...\"\n\nWe appreciate this point. While we do not explicitly report wall-time, our paper reports throughput in tokens per second (tokens/s). Since each training step processes a fixed number of tokens (batch size × sequence length), wall-time can be easily computed as:\n\nwall-time = total_tokens / throughput\n\nThus, throughput and wall-time are effectively interchangeable metrics in our setting. We will clarify this equivalence explicitly in the camera-ready version.\n\n- Q4: \"FlashDP's performance relies on deep integration with CUDA...\"\n\nFirst, FlashDP is already implemented as a user-facing PyTorch interface. Users do not need to write or modify any CUDA code to adopt it. All optimization logic is encapsulated within the framework, making it simple to integrate.\n\nSecond, while FlashDP is currently optimized for CUDA-based GPUs, its design principles—cache-aware fused computation and hierarchical reduction—apply broadly to many modern parallel architectures, as multi-level cache hierarchies are common across platforms. We believe the FlashDP design can be adapted to other hardware in the future, and we plan to explore this direction.\n\n- Other Points:\n\n1. \"Line 39-40...\": Line 39–40 citation correction: Thank you for noting that none of the cited references pertain to DP pre-training. We will replace or remove them and instead cite more relevant works in the camera-ready version.\n\n2. \"Some word choices are uncommon for research papers...\": Thank you for pointing this out. We agree that the word choices “prolonging” (Line 62) and “exemplar” (Line 181) are uncommon in academic writing. We will revise them to “increasing” and “example,” respectively, in the camera-ready version.\n\n3. \"I think a brief discussion on the sampling strategies...\": We thank the reviewer for highlighting this point. We will add a brief discussion on sampling methods (e.g., Poisson vs. uniform) in the revised paper. Note that FlashDP does not interfere with any sampling strategy, as these decisions are made at the data loading stage. FlashDP operates after sampling, during the backward pass.\n\n4. \"The name 'FlashDP' is reminiscent of 'FlashAttention'...\": Inspired by the naming convention of FlashAttention, which similarly optimizes GPU kernel fusion and memory access, we adopted the name “FlashDP” to reflect our method’s core: efficient, cache-friendly operator fusion for DP training.\n\n5. \"From a practical standpoint...\": FlashDP is designed as a standalone, modular library. We plan to release it publicly with documentation and PyTorch integration. Furthermore, we would be happy to contribute upstream to widely used DP libraries such as Opacus and DP-Transformers, and plan to submit PRs to facilitate adoption.\n\n6. For limitation, while FlashDP currently focuses on CUDA-optimized implementation, a key direction for future work is expanding support to other hardware backends such as AMD GPUs or custom AI accelerators. This would help further democratize differentially private LLM training. We believe this is primarily an engineering effort and can benefit from community collaboration as adoption grows. We sincerely apologize for the oversight—although we referred to Appendix A.3 in the checklist, we inadvertently omitted the actual discussion of limitations in the appendix. We will correct this in the camera-ready version."}]

Author rebuttal:

Meta-review: There was a consensus among the reviewers that this work addresses an important and challenging problem: the efficient private training of large-scale models. Reviewers were uniformly impressed by the significant empirical results, which convincingly demonstrate that the proposed FlashDP method achieves substantial gains in throughput and memory efficiency, nearing non-private performance levels on models as large as 13B parameters. During the extensive author-reviewer discussion, several limitations were clarified. The primary ones are the lack of generalizability of the implemented CUDA kernels, and its specific focus on per-layer clipping, as opposed to the more common global clipping paradigm. Furthermore, the evaluation centers on pre-training efficiency rather than fine-tuning, which is a common focus for DP practitioners . However, the authors engaged constructively with this feedback and have committed to addressing these points by adding a dedicated limitations section that discusses the hardware dependency and clipping strategy, and by incorporating a discussion of their utility evaluation metrics.
Decision: Accept (poster)