Controller-based Privacy Engine for Better Transformer Compatibility

[evgri243]

// I'm not sure how to make changes of that scale or whether they are desired, but I'd like to highlight them for discussion. Probably, we should keep it out for a while until the implementation is feature complete and well-tested; meanwhile we can use the PR to discuss the changes.//

Summary

This PR introduces PrivacyEngineGradSampleController, an alternative implementation of Opacus's PrivacyEngine that attaches hooks directly to models without wrapping them in GradSampleModule. This solves compatibility issues with transformers and other models that have complex attribute access patterns.

Motivation

The current PrivacyEngine wraps models in a GradSampleModule, which creates several issues:

1. Type checking breaks: isinstance(model, BertModel) returns False after wrapping
2. State dict complexity: Wrapped models have _module. prefixes in state dicts
3. Attribute access issues: Complex __getattr__ behavior in transformers can break
4. Debugging difficulty: Model structure is hidden behind wrapper

These issues are particularly problematic with HuggingFace transformers and other libraries that perform introspection on model objects.

Solution

Instead of wrapping the model, we:

1. Attach hooks directly to model submodules via register_forward_hook() and register_full_backward_hook()
2. Manage hooks externally through a GradSampleController class
3. Add attributes directly to parameters using setattr() (e.g., param.grad_sample)

The model remains unchanged - no wrapper, no indirection, no type issues.

Implementation

Files Added

1. opacus/grad_sample_controller.py (~480 lines)

   • GradSampleController class that manages hook lifecycle
   • Captures activations in forward pass
   • Computes per-sample gradients in backward pass
   • Cleanup methods to remove hooks and attributes

2. opacus/privacy_engine_gsc.py (~530 lines)

   • PrivacyEngineGradSampleController class with same API as PrivacyEngine
   • Creates GradSampleController instead of wrapping model
   • All other functionality identical (accounting, clipping, noise)

3. opacus/tests/privacy_engine_gsc_test.py (~260 lines)

   • Comprehensive test suite
   • Tests initialization, hook attachment, grad computation
   • Tests state dict compatibility, checkpoints, cleanup

Key Differences from Current Approach

Feature	PrivacyEngine (Current)	PrivacyEngineHookBased (New)
Model wrapping	Yes (GradSampleModule)	No
Type preservation	❌ No	✅ Yes
State dict	Has _module. prefix	Clean, no prefix
Direct attribute access	No	Direct
Transformer compatibility	Can break	Better
Privacy guarantees	Same	Same
Performance	Baseline	Similar

Usage

from opacus.privacy_engine_hook_based import PrivacyEngineHookBased

model = BertModel(...)
optimizer = torch.optim.Adam(model.parameters())
dataloader = ...

privacy_engine = PrivacyEngineHookBased()

# Model is NOT wrapped!
model, optimizer, dataloader = privacy_engine.make_private(
    module=model,
    optimizer=optimizer,
    data_loader=dataloader,
    noise_multiplier=1.0,
    max_grad_norm=1.0,
)

# Train normally
for data, target in dataloader:
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# Clean up when done
privacy_engine.cleanup()

Validation

The code was used to do Zetta 7B transformer LoRA alignment using DPOTrainer from TRL library.

Correctness

1. Hook implementation: Mirrors GradSampleModule.add_hooks() exactly
2. Grad computation: Uses same create_or_accumulate_grad_sample() and promote_current_grad_sample() functions
3. DPOptimizer compatibility: Only requires param.grad_sample attribute, which we provide
4. Privacy accounting: Uses same accountant classes and mechanisms

Key Implementation Details

1. Grad samplers: Automatically imported from GradSampleModule.GRAD_SAMPLERS (registered via decorators)
2. Hook lifecycle: Proper enable/disable/remove/cleanup
3. Validation: Includes same buffer checking as GradSampleModule
4. Attribute cleanup: Removes all Opacus-added attributes on cleanup

API Compatibility

The new class maintains full API compatibility with PrivacyEngine:

• Same make_private() signature
• Same make_private_with_epsilon() signature
• Same save_checkpoint() and load_checkpoint() methods
• Same get_epsilon() method

Migration is trivial: Just change import statement.

Testing

Comprehensive test suite covers:

• ✅ Initialization
• ✅ Hook attachment
• ✅ Per-sample gradient computation
• ✅ Optimizer step functionality
• ✅ State dict preservation (no _module. prefix)
• ✅ Direct attribute access
• ✅ Checkpoint save/load
• ✅ Cleanup (hooks and attributes removed)

Benefits

1. Better transformer compatibility: No wrapper means no __getattr__ issues
2. Simpler state management: Direct model access, no delegation
3. Cleaner checkpoints: No _module. prefix to handle
4. Type checking works: isinstance(model, MyModel) returns True
5. Easier debugging: Model structure unchanged

Trade-offs

1. Explicit cleanup needed: Must call privacy_engine.cleanup() to remove hooks
2. Parameter attributes: Adds attributes directly to parameters (cleaned up on cleanup())
3. Less battle-tested: New implementation, though logic is identical to existing code

Backward Compatibility

• ✅ Does not modify existing PrivacyEngine
• ✅ Can be used alongside existing code
• ✅ Same privacy guarantees
• ✅ Compatible with same DPOptimizer classes
• ✅ Works with existing accountants

Future Work

• [ ] Support for ghost clipping mode (currently only supports hooks/functorch)
• [ ] Support for ghost clipping model + PrivacyEngineAdapriveClipping
• [ ] FSDP support

Checklist

• [x] Implementation complete
• [x] Tests written and passing (locally, pending CI)
• [x] Documentation written
• [x] Examples provided
• [x] Code follows Opacus style (Meta copyright, type hints, docstrings)
• [x] No breaking changes to existing code
• [x] CI tests pass

[meta-codesync[bot]]

@facebook-github-bot has imported this pull request. If you are a Meta employee, you can view this in D85086663. (Because this pull request was imported automatically, there will not be any future comments.)

[iden-kalemaj]

@evgri243 thank you for this heavy-lifting change. I will take some time to digest and also discuss internally with the team. In the meanwhile I have some questions for you:

1. Can you provide some examples of failure modes for the current PrivacyEngine approach with HuggingFace transformers. In your tutorial you use BertForSequenceClassification, but this should work fine with the current approach? I'd like to better understand the extent to which the current PrivacyEngine fails to handle certain models, i.e., which model types have you encountered that do not work with PrivacyEngine, and can we handle the failure modes with smaller changes within the current approach?

2. A main consideration is to avoid code duplication and maintaining several approaches in parallel. With your method, can we avoid some duplication by having PrivacyEngineGradSampleController inherit from PrivacyEngine?

3. What would the changes needed for extending your methodology to ghost clipping and FSDP?

As an intermediate step, we could place your approach in the "research" folder, which we do not actively maintain, but are happy to support you in maintaining it. This would allow some time for the method to be digested before moving it into the main opacus folder.

[evgri243]

Thanks for consideration. It is still work in progress, but I'd love to know your opinion. Let me answer your questions in words, then I'll come with examples if needed:

1. Yeah. The example is more Opacus oriented. We don't experience issues with the models as we mostly use LoRA that limits us to linear layers only anyway. The problem is the trainer. We have a custom loop to replace a standard trainers, but trainers from try library (DPO, KTO) are much harder to substitute due to complicated data preparation and loss implementations. Those trainers are really into isinstance(model, PeftModel) checks, calling direct methods model.from_pretrained(...) on checkpointing, or accessing properties directly model.loss_function. We implemented originally getattr but checkpoint restoration started failing as the model started saving as _module.<tensor_name>, but recovering through forwarding with getattr without it. We overriden state_dict as well, but it all started falling apart. And then we realized that we can avoid wrapping in the first place... and here we are.
2. We can look into privacy engine to wrap it deeper into the existing one. Yeah, duplication is major and existence of two parallel modes is even worse, but keeping it completely to ourselves was unfair as the wrap-less mode just works.
3. We are on ghost clipping now, but thanks to the same KTO and DPO Trainers integrating the loss wrapper is another headache which we so far fail to implement. We have standard ghost clipping wired, but haven't had a proper chance to test it. FSDP is the final goal and we are slowly going there...

I thought about "research" or "contrib". My major problem is to make it package able, but I guess it is not a major implementation issue to add yet another package.

[iden-kalemaj]

Thank you for these explanations. I like your solution and have also experienced some annoyances with accessing attributes of the model post-wrapping. I also understand the use case better now.

I believe we can minimize code duplication which would make it more reasonable to introduce this into Opacus.

1. having GradSampleController extend GradSampleModule to avoid duplicate code. It's okay if this requires some factorization on the side of GradSampleModule.
2. having PrivacyEngineGradSampleController extend PrivacyEngine to avoid the duplicated functions. The more we can minimize new code in PrivacyEngineGradSampleController the better. Likewise we can factorize things in PrivacyEngine as needed.

Do these make sense?

Regarding ghost clipping and FSDP. You mention that you mostly use LoRA. Ghost clipping does not give any memory advantage with LoRA fine-tunining since the effective linear layer width is small, so just wanted to give a heads up that ghost clipping might not be needed for your use case.

We did not implement FSDP with vanilla (non-ghost) clipping since this required more significant effort, though we did put some work into this and if you're interested in using extending FSDP + vanilla, then we'd welcome PRs here.

[evgri243]

@iden-kalemaj give me a few days to give it a try. There is a catch though: GradSampleModule is nn.Module, tracked and controlled by torch. GradSampleController is a simple class, untacking it may turn an issue.

[iden-kalemaj]

That's a good point... let me know how it's looking once you work on it.

[evgri243]

Together (let's be honest with Claude) we actually refactored something reasonable to merge the designs of both Modules/Controllers and PrivacyEngines.

It is still WIP, but you may take a look at the direction at least.

[evgri243]

I will try to add ghost_fsdp support and adapt from our sources a properly working Transformers DPTrainer based on the controller.

[evgri243]

Hello!

I hope to attend to it completely during the holidays and return back with better compatibility.

[iden-kalemaj]

Hi @evgri243, apologies for the delayed response. I like how the integration has improved with the new version.

I have some small recommendations about naming:

• In privacy_engine, rename the parameter return_controller to something more intuitive like wrap_model. Let's try to hide the details of the controller from the user.
• Rename GradSampleHooksMixin to HooksHandler or something similar?

A few questions:

• Can we avoid duplication of additional functions in GradSampleModuleFastGradientClipping?
• In fast_gradient_utils, what is the purpose of functions like add and rmul?

[evgri243]

@iden-kalemaj I've majorly refactored the PR, starting practically from scratch. I tried to keep changes as limited as possible this time to limit the change surface and make it easier to read. It is still work in progress and we should test it on our tasks early next year to make sure everything is correct.

[evgri243]

@david-stan may you have a look in it as well as a co-author.

[evgri243]

Ok. Calling it "as limited as possible" is an overstatement, but I've drastically changed the design and did my best to avoid unnecessary changes

[iden-kalemaj]

@evgri243 and @david-stan, thank you for these changes, and for the hard work on improving the functionality. I like the idea of splitting the hooks handling from the model wrapping at the very base abstract class.

I am somewhat worried about people expecting that the return of the privacy_engine.make_private would be the model as opposed to the hooks object, when using wrap_model=False, but given that it is not the default we can assume people have read the documentation before using this mode.

In the examples, it would be great to have another simpler example of training a hugging face model that would not have been supported with model wrapping. This can be in a separate PR, however.

Finally, there are changes in the code that are not related to the new non-wrapping functionality. I assume these are meant to clean up code, but they make reviewing harder. Can these be placed into separate PRs? I left comments for some of them, but there might be other changes.

[evgri243]

I addressing your cutting all unnecessary changes surgically. Should be out first half next week.

[evgri243]

@iden-kalemaj it should be the shortest I can get it. Some other important functionality: https://github.com/meta-pytorch/opacus/pull/805 and https://github.com/meta-pytorch/opacus/pull/806