Reduce module size or the number of steps to avoid over-time tests (#739)

Summary:

As titled.

Facebook
The buck test has the limit of 10 mins. To avoid overtime failure, we slightly reduce the number of parameters or the number of repetitions. This won't reduce the credibility of the tests.

Differential Revision: D70707205

Author: HuanyuZhang

Diff for: opacus/tests/grad_sample_module_fast_gradient_clipping_test.py

@@ -54,8 +54,8 @@ def __init__(self):
         super(SampleModule, self).__init__()
         self.fc1 = nn.Linear(2, 2)
         self.fc3 = nn.Linear(2, 1024)
-        self.fc4 = nn.Linear(1024, 1024)
-        self.fc5 = nn.Linear(1024, 1)
+        self.fc4 = nn.Linear(1024, 10)
+        self.fc5 = nn.Linear(10, 1)
         self.layer_norm = nn.LayerNorm(2)
 
     def forward(self, x):
@@ -119,7 +119,7 @@ def setUp_data_sequantial(self, size, length, dim):
 
     @given(
         size=st.sampled_from([10]),
-        length=st.sampled_from([1]),
+        length=st.sampled_from([5]),
         dim=st.sampled_from([2]),
     )
     @settings(deadline=1000000)
@@ -195,12 +195,12 @@ def test_norm_calculation_fast_gradient_clipping(self, size, length, dim):
         diff = flat_norms_normal - flat_norms_gc
 
         logging.info(f"Diff = {diff}"),
-        msg = "Fail: Gradients from vanilla DP-SGD and from fast gradient clipping are different"
+        msg = "Fail: Per-sample gradient norms from vanilla DP-SGD and from fast gradient clipping are different"
         assert torch.allclose(flat_norms_normal, flat_norms_gc, atol=1e-3), msg
 
     @given(
         size=st.sampled_from([10]),
-        length=st.sampled_from([1, 5]),
+        length=st.sampled_from([5]),
         dim=st.sampled_from([2]),
     )
     @settings(deadline=1000000)

Diff for: opacus/tests/privacy_engine_test.py

@@ -268,7 +268,7 @@ def _compare_to_vanilla(
         do_clip=st.booleans(),
         do_noise=st.booleans(),
         use_closure=st.booleans(),
-        max_steps=st.sampled_from([1, 4]),
+        max_steps=st.sampled_from([1, 3]),
     )
     @settings(suppress_health_check=list(HealthCheck), deadline=None)
     def test_compare_to_vanilla(
@@ -660,7 +660,7 @@ def test_checkpoints(
 
     @given(
         noise_multiplier=st.floats(0.5, 5.0),
-        max_steps=st.integers(8, 10),
+        max_steps=st.integers(3, 5),
         secure_mode=st.just(False),  # TODO: enable after fixing torchcsprng build
     )
     @settings(suppress_health_check=list(HealthCheck), deadline=None)