Differential Privacy for Hugging Face Trainers – Without Rewriting Your Training Loop

JetBrains Research has open-sourced DPTrainer, a library that integrates Opacus with Hugging Face's Trainer to enable differential privacy (DP-SGD) training without rewriting training loops or modifying trainer source code. The tool addresses privacy concerns in LLM training by preventing membership inference attacks, allowing JetBrains to use sensitive IDE-generated data for model training.

JetBrains Research Research is crucial for progress and innovation, which is why at JetBrains we are passionate about both scientific and market research Research /research/category/research/ Differential Privacy for Hugging Face Trainers – Without Rewriting Your Training Loop It is a well-known problem by now that training LLMs on sensitive data raises serious privacy concerns. In a recent blog post https://blog.jetbrains.com/research/2026/06/membership-inference/ , we talked about membership inference attacks and our research on mitigating them. At JetBrains Research, we are deeply concerned about user privacy and continually developing new methods and tools to improve privacy protection. In this post, we present DPTrainer https://jb.gg/dptrainer , a new library we’ve developed and now open-sourced. DPTrainer smoothly integrates Opacus and Hugging Face Trainer https://huggingface.co/docs/transformers/en/main classes/trainer so that you can train privacy-preserving models without rewriting training loops or modifying trainer source code. The importance of differential privacy It’s been widely observed https://arxiv.org/pdf/2010.14701 that the quality of a model scales along three axes: size, compute, and data. Larger models offer more capacity but suffer from less efficient training and costlier inference. More compute used during training naturally incurs higher costs and takes more time. The data axis, on the other hand, is mostly constrained by the ability to acquire it in sufficient quality and quantity. Differential privacy is our solution to the data-gathering hurdle. Basically, differential privacy is a mathematical framework that protects individual data points used for training. The core guarantee: a model trained with differential privacy behaves almost identically whether or not any single example was included in the training set. For LLMs, which are known to memorize training data and can reproduce it in response to adversarial prompting, this is the strongest known defense against leakage. Even sophisticated Membership Inference Attacks https://blog.jetbrains.com/research/2026/06/membership-inference/ , given access to model weights, confidence scores, and the base model architecture, cannot determine whether a specific example protected by this method was included in the training set. In practice, differential privacy is applied to neural network training through what is known as the differentially private stochastic gradient descent DP-SGD . Rather than computing a single gradient over the entire batch, the DP-SGD computes one gradient per sample, clips it to bound outliners, aggregates the gradients in the batch and than injects noise making the footprint of any single example indistinguishable. By guaranteeing the privacy of our training method, we can exploit previously unavailable channels and use data generated every day through our IDEs see our data collection policy https://www.jetbrains.com/legal/docs/terms/product data collection/ detailed-data-collection and a recent post https://blog.jetbrains.com/blog/2025/09/30/detailed-data-sharing-for-better-ai/ on data sharing for AI . This gives us high data quantity due to the size of our user base, as well as high data quality, as the data is generated in the process of writing code, not just extracted from the final product. Such advantages guarantee that our upcoming models will hit above their weight pun intended . The gap it closes Opacus https://opacus.ai/ is the go-to library for DP-SGD in PyTorch. It provides everything you need: per-sample gradient computation, a DPOptimizer , privacy accountants, and Poisson-sampled data loaders. The catch is that it’s designed around a manual PyTorch training loop, which is inconvenient and not well integrated into the Hugging Face platform. Hugging Face Trainer and Transformers Reinforcement Learning TRL https://huggingface.co/docs/trl/index ’s alignment trainers e.g. SFTTrainer , DPOTrainer are the top high-level training APIs for transformers. They handle distributed training, checkpointing, evaluation, callbacks, and many other things you don’t want to reimplement. However, they have zero awareness of differential privacy. Wiring Opacus into a Trainer -based workflow requires touching model wrapping, optimizer creation, data loading, loss computation, checkpointing, and callback management. These interact in subtle ways, and getting any one wrong can break your privacy guarantee, and do it silently. To fix this, our researchers Evgeny Grigorenko and David Stanojevic created DPTrainer https://jb.gg/dptrainer ; and Mihajlo Linic now maintains it. DPTrainer handles these issues with care. A genuine drop-in replacement A key concept in differential privacy is the privacy budget . This concept represents the maximum theoretical risk of information leakage we are willing to accept. In other words, it is the maximum amount that any single datapoint could shift the output distribution. An important property of the privacy budget is that its expenditure is cumulative, forcing a trade-off between privacy and performance as higher privacy necessitates higher injection of noise into the gradient. DPTrainer extends transformers.Trainer , and incorporates the privacy budget with an added PrivacyArguments dataclass. Every standard training argument, callback, checkpoint, and evaluation workflow works unchanged, as can be seen in the following code: python from dptrainer import DPTrainer, PrivacyArguments privacy args = PrivacyArguments target epsilon=8.0, per sample max grad norm=1.0, trainer = DPTrainer model=model, args=training args, train dataset=train dataset, privacy args=privacy args, data collator=data collator, trainer.train Set a target epsilon to match your privacy budget, and DPTrainer will handle the rest. The internal accountant keeps track of the budget expenditure, and the remaining budget is saved during checkpointing so the run can be easily resumed. Privatizing TRL and other specialized trainers The real power comes from privatize trainer . Many workflows use Trainer subclasses: e.g. DPOTrainer for preference learning, SFTTrainer for instruction tuning, and Seq2SeqTrainer for generation. These all add task-specific loss functions and generation logic on top of the base class. Rewriting those to inherit from DPTrainer would be invasive and fragile. privatize trainer patches any Trainer -based class at runtime, injecting DPTrainer into its inheritance chain without touching the class’s own logic: python from trl import DPOTrainer from dptrainer import PrivacyArguments, privatize trainer privatize trainer DPOTrainer one line trainer = DPOTrainer model=model, args=training args, train dataset=train dataset, processing class=tokenizer, privacy args=PrivacyArguments target epsilon=8.0, per sample max grad norm=1.0 , trainer.train The patched trainer keeps all its original behavior e.g. reward computation, DPO loss, generation , while gaining DP-SGD. What DPTrainer handles, so you don’t have to DPTrainer automatically manages the following: Noise addition. Adds calibrated Gaussian noise to the aggregated gradients via DPOptimizer . Gradient clipping. Clips each sample’s gradient individually before aggregation, not the batch gradient. Supports flat, adaptive AdaClip , and per-layer strategies via clipping and per sample max grad norm . Gradient computation. Wraps the model in Opacus’s GradSampleModule for per-sample gradients, which is required for DP-SGD correctness. Optimizer creation. Intercepts create optimizer to wrap the Hugging Face-created optimizer with DPOptimizer . Data loading. Overrides get train dataloader to return a DPDataLoader with Poisson sub-sampling, which is what enables privacy amplification by sampling. Noise calibration. Given a target epsilon and your training configuration, DPTrainer computes the correct noise multiplier automatically – no manual binary search. Privacy accounting. A DPCallback hooks into the optimizer step and tracks the running privacy budget after every update. Checkpointing. Saves and restores accountant state alongside model weights, so your privacy budget tracking remains correct after resuming. Early stopping. A privacy-budget-aware stopping mechanism halts training automatically when the entire budget is exhausted. Flexible configuration PrivacyArguments exposes the knobs you’d expect: : set one or the other – they’re mutually exclusive. target epsilon / noise multiplier : choose clipping "flat" standard , "adaptive" AdaClip , or "per layer" .: toggle Poisson sub-sampling for privacy amplification. poisson sampling : grad sample mode "hooks ” default .: privacy accountant type RDP by default . accountant : log budget expenditure at training steps, eval, both, or not at all. epsilon log mode Try our DPTrainer Differential privacy is increasingly a compliance requirement, not just a research nicety. Regulations around training on personal data, and the growing awareness of membership inference attacks against LLMs, mean that teams need practical, auditable differential privacy training. The hard part has never been the math; it’s been the engineering. And DPTrainer removes that barrier. If you’re training transformers on sensitive data and using any part of the Hugging Face ecosystem, this is worth a try. Subscribe to JetBrains Research blog updates