From fab7cc8887f127d9ea73cb7d0a386727865a1b21 Mon Sep 17 00:00:00 2001
From: Aymen Noor <mn05052002@gmail.com>
Date: Thu, 9 Apr 2026 18:50:19 -0400
Subject: [PATCH 1/2] git commit -m "Fix sample_rate calculation with
 WeightedRandomSampler

When DataLoader uses WeightedRandomSampler, sample_rate was incorrectly
computed from len(data_loader) instead of the true dataset size, causing
privacy budget to burn 100x-1000x faster than expected.

Fixed by capturing dataset size and batch_size before _prepare_data_loader()
replaces the sampler, ensuring privacy accounting uses actual dataset size.

Fixes #[813]"
---
 opacus/privacy_engine.py              | 16 ++++++++--
 opacus/tests/test_weighted_sampler.py | 46 +++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 opacus/tests/test_weighted_sampler.py

diff --git a/opacus/privacy_engine.py b/opacus/privacy_engine.py
index 9d8947c70..6074c0982 100644
--- a/opacus/privacy_engine.py
+++ b/opacus/privacy_engine.py
@@ -409,6 +409,18 @@ def make_private(
         )
         if poisson_sampling:
             module.forbid_grad_accumulation()
+            
+        if hasattr(data_loader.dataset, '__len__'):
+            true_dataset_size = len(data_loader.dataset)
+            print(f"DEBUG: true_dataset_size = {true_dataset_size}")
+        else:
+            raise ValueError(
+                "Dataset must have __len__ for privacy accounting. "
+                "IterableDataset is not supported."
+    )
+        original_batch_size = data_loader.batch_size
+        if original_batch_size is None:
+            raise ValueError("DataLoader must have a batch_size")
 
         data_loader = self._prepare_data_loader(
             data_loader,
@@ -418,8 +430,8 @@ def make_private(
             rand_on_empty=rand_on_empty,
         )
 
-        sample_rate = 1 / len(data_loader)
-        expected_batch_size = int(len(data_loader.dataset) * sample_rate)
+        sample_rate = original_batch_size / true_dataset_size
+        expected_batch_size = int(true_dataset_size * sample_rate)
 
         # expected_batch_size is the *per worker* batch size
         if distributed:
diff --git a/opacus/tests/test_weighted_sampler.py b/opacus/tests/test_weighted_sampler.py
new file mode 100644
index 000000000..d58b0be9b
--- /dev/null
+++ b/opacus/tests/test_weighted_sampler.py
@@ -0,0 +1,46 @@
+import torch
+from torch.utils.data import TensorDataset, DataLoader, WeightedRandomSampler
+from opacus import PrivacyEngine
+
+
+def test_weighted_sampler_privacy_accounting():
+    """
+    Test that WeightedRandomSampler doesn't break privacy accounting.
+    
+    Regression test for issue where sample_rate was computed from
+    sampler.num_samples instead of the true dataset size, causing
+    privacy budget to burn 100x-1000x faster than expected.
+    """
+    # Dataset with 100,000 samples
+    X = torch.randn(100_000, 10)
+    y = torch.randint(0, 2, (100_000,))
+    dataset = TensorDataset(X, y)
+    
+    # WeightedRandomSampler with only 128 samples per epoch
+    weights = torch.ones(100_000)
+    sampler = WeightedRandomSampler(weights, num_samples=128, replacement=True)
+    loader = DataLoader(dataset, batch_size=16, sampler=sampler)
+    
+    model = torch.nn.Linear(10, 2)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    
+    privacy_engine = PrivacyEngine()
+    model, optimizer, loader = privacy_engine.make_private_with_epsilon(
+        module=model,
+        optimizer=optimizer,
+        data_loader=loader,
+        epochs=1,
+        target_epsilon=8.0,
+        target_delta=1e-5,
+        max_grad_norm=1.0,
+    )
+    
+    # Verify privacy accounting uses true dataset size
+    expected_sample_rate = 16 / 100_000  # 0.00016
+    actual_sample_rate = optimizer.expected_batch_size / len(loader.dataset)
+    
+    assert optimizer.expected_batch_size == 16, \
+        f"expected_batch_size should be 16, got {optimizer.expected_batch_size}"
+    
+    assert abs(actual_sample_rate - expected_sample_rate) < 1e-6, \
+        f"sample_rate should be {expected_sample_rate}, got {actual_sample_rate}"
\ No newline at end of file

From 5a5a00d757d7c12d6c5a56fd9e2a03599fd99667 Mon Sep 17 00:00:00 2001
From: Aymen Noor <mn05052002@gmail.com>
Date: Thu, 9 Apr 2026 22:03:07 -0400
Subject: [PATCH 2/2] removed debug statement

---
 opacus/privacy_engine.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/opacus/privacy_engine.py b/opacus/privacy_engine.py
index 6074c0982..1436ac0dc 100644
--- a/opacus/privacy_engine.py
+++ b/opacus/privacy_engine.py
@@ -409,10 +409,9 @@ def make_private(
         )
         if poisson_sampling:
             module.forbid_grad_accumulation()
-            
+
         if hasattr(data_loader.dataset, '__len__'):
-            true_dataset_size = len(data_loader.dataset)
-            print(f"DEBUG: true_dataset_size = {true_dataset_size}")
+            true_dataset_size = len(data_loader.dataset)S
         else:
             raise ValueError(
                 "Dataset must have __len__ for privacy accounting. "