JetBrains Research #
Research is crucial for progress and innovation, which is why at JetBrains we are passionate about both scientific and market research
It is a well-known problem by now that training LLMs on sensitive data raises serious privacy concerns. In a recent blog post, 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, a new library we’ve developed and now open-sourced. DPTrainer smoothly integrates Opacus and Hugging Face 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 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, 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 and a recent post 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 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 s. 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 FaceTrainer
and Transformers Reinforcement Learning (TRL)’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 , 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; 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:
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:
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 viaDPOptimizer
.Gradient clipping. Clips each sample’s gradient individually before aggregation, not the batch gradient. Supports flat, adaptive (AdaClip), and per-layer strategies viaclipping
andper_sample_max_grad_norm
.Gradient computation. Wraps the model in Opacus’sGradSampleModule
for per-sample gradients, which is required for DP-SGD correctness.Optimizer creation. Interceptscreate_optimizer
to wrap the Hugging Face-created optimizer withDPOptimizer
.Data . Overridesget_train_data
to return aDPData
with Poisson sub-sampling, which is what enables privacy amplification by sampling.Noise calibration. Given atarget_epsilon
and your training configuration,DPTrainer
computes the correctnoise_multiplier
automatically – no manual binary search.Privacy accounting. ADPCallback
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
: chooseclipping
"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.