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

opacus/tests/grad_sample_module_fast_gradient_clipping_test.py

@@ -120,7 +120,7 @@ def setUp_data_sequantial(self, size, length, dim):
 
     @given(
         size=st.sampled_from([10]),
-        length=st.sampled_from([1, 10]),
+        length=st.sampled_from([5]),
         dim=st.sampled_from([2]),
     )
     @settings(deadline=1000000)
@@ -192,12 +192,12 @@ def test_norm_calculation_fast_gradient_clipping(self, size, length, dim):
         diff = flat_norms_normal - flat_norms_gc
 
         logging.info(f"Max difference between (vanilla) Opacus and FGC = {max(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, 10]),
+        length=st.sampled_from([5]),
         dim=st.sampled_from([2]),
     )
     @settings(deadline=1000000)

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)