DPTrainer

DPTrainer bridges Opacus and Hugging Face's Trainer ecosystem — including the alignment trainers in TRL (DPOTrainer, SFTTrainer, etc.) — so you can add DP-SGD to any Trainer-based workflow without rewriting training loops or modifying trainer source code.

The Problem

Opacus provides the building blocks for DP-SGD (per-sample gradients, DPOptimizer, privacy accountants, Poisson-sampled data loaders), but it is designed around a manual PyTorch training loop. Hugging Face Trainer is the dominant high-level training API for Transformers — and TRL extends it with alignment-specific trainers — but neither has any awareness of differential privacy.

Plugging Opacus into a Trainer-based pipeline requires coordinated changes across model wrapping, optimizer creation, data loading, loss computation, checkpointing, and callback management. These concerns interact in non-obvious ways, and getting any one of them wrong silently breaks the privacy guarantee.

What DPTrainer Handles

Installation

uv pip install --index-url https://europe-west4-python.pkg.dev/jetbrains-ml4se-fed/jbr-fed-python/simple DPTrainer

Or with pip:

pip install --index-url https://europe-west4-python.pkg.dev/jetbrains-ml4se-fed/jbr-fed-python/simple DPTrainer

Quick Start — DPTrainer as a Drop-In Replacement

DPTrainer extends transformers.Trainer, so all standard training arguments, callbacks, checkpointing, and evaluation workflows work unchanged. Pass a PrivacyArguments instance to control the privacy budget:

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, DataCollatorForLanguageModeling
from dptrainer import DPTrainer, PrivacyArguments

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

privacy_args = PrivacyArguments(
    target_epsilon=8.0,
    per_sample_max_grad_norm=1.0,
)

training_args = TrainingArguments(
    output_dir="./output",
    num_train_epochs=3,
    per_device_train_batch_size=32,
)

trainer = DPTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    privacy_args=privacy_args,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)

trainer.train()

Privatizing Any Trainer (TRL, Seq2Seq, etc.)

DPTrainer works as a direct replacement for Trainer, but many workflows use specialized subclasses — TRL's DPOTrainer, SFTTrainer, Hugging Face's Seq2SeqTrainer, and others — that add task-specific logic on top of Trainer. Rewriting them to inherit from DPTrainer would be invasive and fragile.

The privatize_trainer utility solves this by patching any Trainer-based class at runtime, swapping Trainer for DPTrainer in its inheritance chain. The patched class keeps all of its original behavior (custom loss, generation logic, reward computation) while gaining differential privacy:

from trl import DPOTrainer
from transformers import TrainingArguments
from dptrainer import PrivacyArguments
from dptrainer import privatize_trainer

privacy_args = PrivacyArguments(
    target_epsilon=8.0,
    per_sample_max_grad_norm=1.0,
)

# One-line patch — DPOTrainer now trains with DP-SGD
privatize_trainer(DPOTrainer)

# Use DPOTrainer exactly as before — no other code changes needed
trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    processing_class=tokenizer,
    privacy_args=privacy_args,
)
trainer.train()

Note: privatize_trainer patches the trainer class, not a specific instance. It injects DPTrainer into the class hierarchy so that all future instances of that class train with DP-SGD. Use it once at import time or before constructing trainer instances.

Limitation: privatize_trainer works only on subclasses of transformers.Trainer (e.g., Seq2SeqTrainer, DPOTrainer, SFTTrainer). It cannot be called on Trainer itself — doing so raises a ValueError. If you want DP training with the base Trainer, use DPTrainer directly as shown in the Quick Start section above.

When ghost clipping is enabled, privatize_trainer automatically inspects the patched class's MRO and warns if any intermediate class overrides compute_loss or training_step in a way that could bypass DPTrainer's loss wrapping.

PrivacyArguments Reference

Docs

For detailed documentation, please visit the docs.

License

Licensed under the MIT License. See LICENSE.