Skip to content

Feat : Add DP-SGD Transformer example using Flax NNX API | Issue #120 #126

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
debanganghosh08 wants to merge 23 commits into
base: main
Choose a base branch
from
Open

Feat : Add DP-SGD Transformer example using Flax NNX API | Issue #120 #126

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
968cee0
Implemented Deplayed Preconditioners with alternating-phase protocol …
debanganghosh08 Jan 21, 2026
f014b51
Implement User-Level Sampling (ULS) for Transformers with per-user av…
debanganghosh08 Jan 20, 2026
4a08a05
Refactor: Use UserSelectionStrategy and clipped_grad(keep_batch_dim=F…
debanganghosh08 Jan 22, 2026
ee22708
style: fix remaining pylint R0917 and whitespace violations
debanganghosh08 Jan 24, 2026
8cbc13b
style: fix line length in user_level_transformer and sync nnx example
debanganghosh08 Jan 24, 2026
eceff37
style: fix pytype import errors and finalize production standards
debanganghosh08 Jan 24, 2026
ebe4b07
Refactor transformer DP-SGD example and update CI workflow
debanganghosh08 Jan 26, 2026
e2bd4bd
Add requirements file for transformer examples to fix CI dependencies
debanganghosh08 Jan 26, 2026
5198d8f
chore: align dependencies with new modular CI in pyproject.toml
debanganghosh08 Jan 27, 2026
ea34efa
refactor: align Transformer examples with execution_plan and producti…
debanganghosh08 Jan 30, 2026
08c76fd
fix: resolve TOML corruption and standardize formatting across all ex…
debanganghosh08 Jan 31, 2026
83ea418
chore: sync infrastructure files from ULS branch
debanganghosh08 Feb 1, 2026
b148409
refactor: revert to Adam optimizer and ensure clean TOML infrastructure
debanganghosh08 Feb 1, 2026
cc99d51
chore: final production sync and docstring fix
debanganghosh08 Feb 1, 2026
ba34809
fix: restore infrastructure integrity and synchronize transformer logic
debanganghosh08 Feb 1, 2026
0374a12
Refactor: Final production hardening for NNX Transformer, secure RNG …
debanganghosh08 Mar 15, 2026
bfc9e8e
Fix: Align test dependencies in pyproject.toml and resolve pyink form…
debanganghosh08 Mar 15, 2026
e0dda1e
Fix: Add Windows environment marker to flax in both test and examples…
debanganghosh08 Mar 15, 2026
f347fb5
Fix: Add Windows environment marker
debanganghosh08 Mar 15, 2026
6c1bc3b
Refactor: Final production hardening, privacy-safe 2D data units, and…
debanganghosh08 Mar 23, 2026
c505741
Fix: Resolve BandMFExecutionPlanConfig constructor error and optimize…
debanganghosh08 Mar 23, 2026
1ca4683
Refactor: Implement batch padding for JIT optimization, migrate to pl…
debanganghosh08 Apr 5, 2026
2206902
Chore: Synchronize with upstream/main and verify JIT optimization logic
debanganghosh08 Apr 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions examples/distributed_noise_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,10 +162,10 @@ def run(pytree_like_model_params):
t0 = time.time()
compiled_run = run.lower(model_params).compile()
t1 = time.time()
print(f'[BandMF] Compilation time: {t1-t0:.3f} seconds')
print(f'[BandMF] Compilation time: {t1 - t0:.3f} seconds')
state, noisy_grad = jax.block_until_ready(compiled_run(model_params))
t2 = time.time()
print(f'[BandMF] Per-step run time: {(t2-t1)/steps:.3f} seconds')
print(f'[BandMF] Per-step run time: {(t2 - t1) / steps:.3f} seconds')

return state, noisy_grad

Expand Down
Loading