Repository: facebookresearch/fairscale

Issue:
Compatibility with pytest 7
I just ran into an error in the test suite with the very recently release pytest 7 in https://github.com/conda-forge/fairscale-feedstock/pull/9.

I was too slow to copy out the log before azure deleted the run, but it was something about using assert in a boolean context or similar. After pinning to `pytest <7` the tests pass again for fairscale 0.4.5. Would be good if this could be fixed. :)


Discussion:
Thanks for the report. It is interesting and nice to have for sure. @tmarkstrum @anj-s 

Resolving patch:
diff --git a/benchmarks/oss.py b/benchmarks/oss.py
--- a/benchmarks/oss.py
+++ b/benchmarks/oss.py
@@ -109,7 +109,8 @@ def validate_benchmark(measurements, final_loss, args, check_regression):
         assert max_memory < 1.05 * golden_data["reference_memory"], (
             f"Memory use regression detected: " f"{max_memory} vs. {1.05* golden_data['reference_memory']}"
         )
-        assert abs(cast(float, final_loss) - golden_data["reference_loss"]) < 1e-2, (
+        # any min_loss < than golden + epsilon is OK.
+        assert cast(float, final_loss) - golden_data["reference_loss"] < 1e-2, (
             f"Loss regression detected: " f"{final_loss} vs. {golden_data['reference_loss']}"
         )
         logging.info("[Regression Test] VALID")
@@ -176,6 +177,7 @@ def train(
 
     measurements = []
     final_loss: Optional[float] = -1.0
+    min_loss = 100.0
     need_profiling = args.profile
 
     for epoch in range(args.epochs):
@@ -264,14 +266,21 @@ def run_closure(closure, scaler, optimizer):
                 logging.info("... State dict collected")
 
         measurements.append(n_items / epoch_runtime)
+        min_loss = min(min_loss, final_loss)
         if dist.get_rank() == 0:
-            logging.info(f"Epoch {epoch} - processed {measurements[-1]:.2f} img per sec. Loss {final_loss:.3f}")
+            logging.info(
+                f"Epoch {epoch} - processed {measurements[-1]:.2f} img per sec. "
+                f"Loss {final_loss:.3f} min loss {min_loss:.3f}"
+            )
 
     training_stop = time.monotonic()
     img_per_sec = n_items / (training_stop - training_start) * args.epochs
     logging.info(f"[{dist.get_rank()}] : Training done. {img_per_sec:.2f} img per sec inc. checkpoint")
 
-    validate_benchmark(measurements, final_loss, args, check_regression)
+    # Use min_loss to check instead of final_loss since the final_loss is a bit random.
+    # If the training min_loss reaches certain number, we can be reasonably certain the
+    # training process was correct.
+    validate_benchmark(measurements, min_loss, args, check_regression)
 
     dist.destroy_process_group()  # type: ignore
 
diff --git a/fairscale/nn/data_parallel/fsdp_optim_utils.py b/fairscale/nn/data_parallel/fsdp_optim_utils.py
--- a/fairscale/nn/data_parallel/fsdp_optim_utils.py
+++ b/fairscale/nn/data_parallel/fsdp_optim_utils.py
@@ -119,7 +119,7 @@ def _unflatten_optim_state(
     if not combined_state:
         return {}, global_to_local_id
 
-    # copy non tensor state to all global entries
+    # copy non tensor state (like the "step" count) to all global entries
     unflat_state = {i: copy.deepcopy(non_tensor_state[0]) for i in range(sum(num_global_params))}
 
     if non_tensor_state[0].keys() == combined_state[0].keys():
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -367,13 +367,16 @@ def __init__(
         # In a unit test dummy enviromnent, the process_group_reduce_scatter can be None.
         if self.process_group_reduce_scatter is not None:
             reduce_scatter_group_size = self.process_group_reduce_scatter.size()
-            # Roll back to use the default process group for reduce scatter operation when the world size and reduce scatter process group size are differnt.
+            # Roll back to use the default process group for reduce scatter operation when
+            # the world size and reduce scatter process group size are differnt.
             if self.world_size != reduce_scatter_group_size:
                 self.process_group_reduce_scatter = self.process_group
                 logging.warn(
-                    "Rolled back to use the default process group for the reduce scatter operation because the reduce_scatter process group"
-                    f"size is {reduce_scatter_group_size}, which is different with the world size {self.world_size}. Please make sure the process_group"
-                    "parameter uses all the available ranks for the optimized performance."
+                    "Rolled back to use the default process group for the reduce scatter "
+                    "operation because the reduce_scatter process group "
+                    f"size is {reduce_scatter_group_size}, which is different with the "
+                    f"world size {self.world_size}. Please make sure the process_group "
+                    "parameter uses all the available ranks for the optimal performance."
                 )
         self.reshard_after_forward = self._orig_reshard_after_forward = reshard_after_forward
         self.disable_reshard_on_root = disable_reshard_on_root
@@ -2309,6 +2312,20 @@ def _fsdp_instances(self) -> List["FullyShardedDataParallel"]:
         return [m for m in self.modules() if isinstance(m, FullyShardedDataParallel)]
 
     def _remove_uncollectable_params_from_optim_state_dict(self, osd: Dict) -> Dict:
+        """Return a new state dict filtering out the ones like MoE layers, which has
+        ``no_broadcast_optim_state`` flag set.
+
+        We also make rooms for the optimizer state on rank 0.
+        """
+        # In PyTorch version 1.12, Adam's `step` state changed from an int to a singleton
+        # tensor. We convert it back here. Otherwise, the step counter will be treated
+        # like a singleton tensor and comparison with original state dict would fail.
+        for _, bufs in osd["state"].items():
+            if "step" in bufs.keys():
+                assert type(bufs["step"]) is int or ou.is_singleton_tensor(bufs["step"])
+                if ou.is_singleton_tensor(bufs["step"]):
+                    bufs["step"] = bufs["step"].item()
+        # Get uncollected_ids.
         uncollected_ids = [i for i, m in enumerate(self._fsdp_instances) if m.no_broadcast_optim_state]
         new_dct = {"state": {k: v for k, v in osd["state"].items() if k not in uncollected_ids}}
         if self.rank == 0:


Test patch:
diff --git a/fairscale/utils/testing.py b/fairscale/utils/testing.py
--- a/fairscale/utils/testing.py
+++ b/fairscale/utils/testing.py
@@ -136,6 +136,15 @@ def get_smi_ver() -> str:
     return tuple(int(n) for n in numbering)
 
 
+def make_cudnn_deterministic() -> None:
+    """Make cudnn (matmul) deterministic"""
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # TF32 also make things nondeterministic. Disable it.
+    torch.backends.cuda.matmul.allow_tf32 = False  # type: ignore
+    torch.backends.cudnn.allow_tf32 = False  # type: ignore
+
+
 def dist_init(rank: int, world_size: int, filename: str, filename_rpc: str = "") -> bool:
     """
     Initialize torch distributed, based on a temporary file shared across ranks, which makes it possible for unrelated
@@ -218,8 +227,7 @@ def test_runner(
 ) -> None:
     # At this point we're in a new process, torch options need to be set again
     if deterministic:
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
+        make_cudnn_deterministic()
         torch.manual_seed(1357)
 
     test_func(rank, *args, **kwargs)
@@ -270,8 +278,7 @@ def worker_process(
     )
 
     if torch.cuda.is_available() and not hasattr(torch.backends.cudnn, "flags"):
-        torch.backends.cudnn.benchmark = False
-        torch.backends.cudnn.deterministic = True
+        make_cudnn_deterministic()
 
     try:
         with context:
diff --git a/tests/experimental/nn/test_multiprocess_pipe.py b/tests/experimental/nn/test_multiprocess_pipe.py
--- a/tests/experimental/nn/test_multiprocess_pipe.py
+++ b/tests/experimental/nn/test_multiprocess_pipe.py
@@ -22,6 +22,7 @@
 
 from fairscale.experimental.nn.distributed_pipeline import DistributedLoss, DistributedPipeline, PipelineModulesGraph
 from fairscale.utils import torch_version
+from fairscale.utils.testing import skip_if_single_gpu
 
 pytestmark = pytest.mark.skipif(
     not torch.cuda.is_available() or torch_version() < (1, 9, 0),
@@ -103,6 +104,7 @@ def create(devices):
 
 
 @rpc_test()
+@skip_if_single_gpu
 def create_multiple_layers():
     model = [RemoteModuleParams(nn.Linear, (4, 4), {}), RemoteModuleParams(nn.ReLU, (), {})]
     pipe = create_sequence_pipeline(model, balance=[1, 1], chunks=1, devices=["worker0/cpu", "worker0/cpu"])
@@ -110,6 +112,7 @@ def create_multiple_layers():
 
 @rpc_test(world_size=2)
 @pytest.mark.parametrize("devices", DEVICES)
+@skip_if_single_gpu
 def create_multiple_workers(devices):
     model = [RemoteModuleParams(nn.Linear, (4, 4), {}), RemoteModuleParams(nn.ReLU, (), {})]
     pipe = create_sequence_pipeline(model, balance=[1, 1], chunks=1, devices=devices[:2])
@@ -117,6 +120,7 @@ def create_multiple_workers(devices):
 
 @rpc_test(world_size=2)
 @pytest.mark.parametrize("devices", DEVICES)
+@skip_if_single_gpu
 def parameter_rrefs(devices):
     model = [RemoteModuleParams(nn.Linear, (4, 4), {}), RemoteModuleParams(nn.ReLU, (), {})]
     pipe = create_sequence_pipeline(model, balance=[1, 1], chunks=1, devices=devices[:2])
@@ -149,6 +153,7 @@ def forward_chunks(devices):
 @rpc_test(world_size=2)
 @pytest.mark.parametrize("devices", DEVICES)
 @pytest.mark.parametrize("checkpoint", ["never", "always", "except_last"])
+@skip_if_single_gpu
 def forward_multi(devices, checkpoint):
     device = devices[0].split("/")[1]
     torch.random.manual_seed(3)
@@ -166,6 +171,7 @@ def forward_multi(devices, checkpoint):
 
 @rpc_test(world_size=2)
 @pytest.mark.parametrize("devices", DEVICES)
+@skip_if_single_gpu
 def backward(devices):
     device = devices[0].split("/")[1]
     torch.random.manual_seed(3)
@@ -183,6 +189,7 @@ def backward(devices):
 
 @rpc_test(world_size=2)
 @pytest.mark.parametrize("devices", DEVICES)
+@skip_if_single_gpu
 def update(devices):
     device = devices[0].split("/")[1]
     torch.random.manual_seed(3)
@@ -223,6 +230,7 @@ def extract_partitions(graph: PipelineModulesGraph, pipeline: DistributedPipelin
 
 @rpc_test(world_size=2)
 @pytest.mark.parametrize("devices", DEVICES)
+@skip_if_single_gpu
 def multi_input_multi_output_layers(devices):
     device = devices[0].split("/")[1]
     torch.random.manual_seed(3)
@@ -289,6 +297,7 @@ def forward(self, input):
 
 @rpc_test(world_size=2)
 @pytest.mark.parametrize("devices", DEVICES)
+@skip_if_single_gpu
 def auto_graph_extract(devices):
     from fairscale.experimental.nn.distributed_pipeline.trace import make_graph
 
diff --git a/tests/nn/data_parallel/test_fsdp.py b/tests/nn/data_parallel/test_fsdp.py
--- a/tests/nn/data_parallel/test_fsdp.py
+++ b/tests/nn/data_parallel/test_fsdp.py
@@ -801,8 +801,11 @@ def __init__(self, group, wrapper_config, checkpoint_act=False, delay_before_fre
 
         if wrapper_config is not None:
             # we create a process group of size 1 for the expert params
-            expert_group = torch.distributed.new_group([group.rank()])  # world size 1 means no shard
-            expert = FullyShardedDataParallel(expert, expert_group, **wrapper_config)
+            # we also need to pass that group as the reduce_scatter group.
+            expert_group = torch.distributed.new_group([group.rank()])
+            expert = FullyShardedDataParallel(
+                expert, process_group=expert_group, process_group_reduce_scatter=expert_group, **wrapper_config
+            )
 
             shared = FullyShardedDataParallel(shared, group, **wrapper_config)
 
diff --git a/tests/nn/data_parallel/test_fsdp_grad_acc.py b/tests/nn/data_parallel/test_fsdp_grad_acc.py
--- a/tests/nn/data_parallel/test_fsdp_grad_acc.py
+++ b/tests/nn/data_parallel/test_fsdp_grad_acc.py
@@ -13,7 +13,7 @@
 import torch
 
 from fairscale.nn.data_parallel import FullyShardedDataParallel
-from fairscale.utils.testing import DummyProcessGroup, objects_are_equal
+from fairscale.utils.testing import DummyProcessGroup, make_cudnn_deterministic, objects_are_equal
 
 from .test_fsdp import DistributedTest, NestedWrappedModule, rename_test, spawn_and_init
 
@@ -64,6 +64,7 @@ def _test_nested_wrapper(self, rank, group, config):
 
     @classmethod
     def _test_grad_acc(self, model, batch_dim, use_no_sync_context=True):
+        make_cudnn_deterministic()
         # Generate two input batches. We'll test that we get the same grads if
         # we train on them sequentially while accumulating grads (with no_sync
         # or without no_sync) vs. concatenating the batches and training in one go.
diff --git a/tests/nn/data_parallel/test_fsdp_multiple_wrapping.py b/tests/nn/data_parallel/test_fsdp_multiple_wrapping.py
--- a/tests/nn/data_parallel/test_fsdp_multiple_wrapping.py
+++ b/tests/nn/data_parallel/test_fsdp_multiple_wrapping.py
@@ -66,7 +66,7 @@ def forward(self, x):
 # We use strings for precision and flatten instead of bool to
 # make the pytest output more readable.
 @skip_if_no_cuda
-@pytest.mark.parametrize("world_size", [1, 2])
+@pytest.mark.parametrize("world_size", [1, 2] if torch.cuda.device_count() > 1 else [1])
 @pytest.mark.parametrize("precision", ["full", "mixed"])
 @pytest.mark.parametrize("flatten", ["flatten", "no_flatten"])
 def test(world_size, precision, flatten):
diff --git a/tests/nn/data_parallel/test_fsdp_optimizer_utils.py b/tests/nn/data_parallel/test_fsdp_optimizer_utils.py
--- a/tests/nn/data_parallel/test_fsdp_optimizer_utils.py
+++ b/tests/nn/data_parallel/test_fsdp_optimizer_utils.py
@@ -24,11 +24,13 @@
 )
 
 
-def first_tensor_numel(dct):
+def all_tensors_numel_except_for_step(dct):
+    """Compute the sum of numel from all tensors from a dict, except when the key is `step`."""
+    ret = 0
     for k, v in dct.items():
-        if torch.is_tensor(v):
-            return v.numel()
-    return 0
+        if k != "step" and torch.is_tensor(v):
+            ret += v.numel()
+    return ret
 
 
 def assert_equal(a, b):
@@ -123,8 +125,8 @@ def _test_consolidated_optimizer(self, config, rank, group, optim_fn=torch.optim
         assert_equal(len(sd["state"]), len(unwrapped_sd["state"]))
         assert_equal(len(sd["param_groups"][0]["params"]), len(unwrapped_sd["param_groups"][0]["params"]))
         assert_equal(
-            sum([first_tensor_numel(v) for k, v in sd["state"].items()]),
-            sum([first_tensor_numel(v) for k, v in unwrapped_sd["state"].items()]),
+            sum([all_tensors_numel_except_for_step(v) for k, v in sd["state"].items()]),
+            sum([all_tensors_numel_except_for_step(v) for k, v in unwrapped_sd["state"].items()]),
         )
 
         original_shard_sd = fsdp_optim.state_dict()
@@ -133,8 +135,8 @@ def _test_consolidated_optimizer(self, config, rank, group, optim_fn=torch.optim
         original_shard_sd = recursive_copy_to_device(original_shard_sd, non_blocking=False, device="cpu")
         # Before asserting that the dicts are equal, we check keys individually to allow nice tracebacks.
         assert_equal(
-            [first_tensor_numel(v) for k, v in shard_sd["state"].items()],
-            [first_tensor_numel(v) for k, v in original_shard_sd["state"].items()],
+            [all_tensors_numel_except_for_step(v) for k, v in shard_sd["state"].items()],
+            [all_tensors_numel_except_for_step(v) for k, v in original_shard_sd["state"].items()],
         )
         assert_equal(
             [v for k, v in shard_sd["param_groups"][0].items()],
diff --git a/tests/nn/data_parallel/test_sharded_ddp_features.py b/tests/nn/data_parallel/test_sharded_ddp_features.py
--- a/tests/nn/data_parallel/test_sharded_ddp_features.py
+++ b/tests/nn/data_parallel/test_sharded_ddp_features.py
@@ -207,6 +207,7 @@ def closure():
 @pytest.mark.parametrize("reduce_buffer_size", [0, 2 ** 20])
 @pytest.mark.parametrize("backend", ["gloo", "nccl"])
 @pytest.mark.parametrize("device", available_devices)
+@skip_if_single_gpu
 def test_inputs(reduce_buffer_size, backend, device):
     # Check that the ShardedDDP wrapper accepts tuple(tensors) as inputs
     world_size = 2
diff --git a/tests/nn/model_parallel/test_layers.py b/tests/nn/model_parallel/test_layers.py
--- a/tests/nn/model_parallel/test_layers.py
+++ b/tests/nn/model_parallel/test_layers.py
@@ -1,5 +1,3 @@
-# coding=utf-8
-
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 #
 # This source code is licensed under the BSD license found in the
@@ -136,7 +134,7 @@ def run_test_initialize_affine_weight(rank, model_parallel_size, filename, filen
             torch.distributed.get_rank(), error
         )
     )
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     # ------------
     # Row parallel
@@ -157,7 +155,7 @@ def run_test_initialize_affine_weight(rank, model_parallel_size, filename, filen
     print(
         "   row parallel max error (should be zero) on global rank {}: {}".format(torch.distributed.get_rank(), error)
     )
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     # Reset groups
     mpu.destroy_model_parallel()
@@ -217,18 +215,18 @@ def run_test_column_parallel_linear(rank, model_parallel_size, filename, filenam
     error = my_dLdA.sub(linear_layer.weight.grad).abs().max()
     torch.distributed.barrier()
     print("   error in dLdA on global rank {}: {}".format(torch.distributed.get_rank(), error))
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     my_dLdb = torch.split(dLdb, output_size_coeff, dim=0)[rank].contiguous().clone()
     error = my_dLdb.sub(linear_layer.bias.grad).abs().max()
     torch.distributed.barrier()
     print("   error in dLdb on global rank {}: {}".format(torch.distributed.get_rank(), error))
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     error = dLdX.sub(identity_layer.weight.grad).abs().max()
     torch.distributed.barrier()
     print("   error in dLdX on global rank {}: {}".format(torch.distributed.get_rank(), error))
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     # Reset groups
     mpu.destroy_model_parallel()
@@ -278,17 +276,17 @@ def run_test_row_parallel_linear(rank, model_parallel_size, filename, filename_r
     error = my_dLdA.sub(linear_layer.weight.grad).abs().max()
     torch.distributed.barrier()
     print("   error in dLdA on global rank {}: {}".format(torch.distributed.get_rank(), error))
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     error = dLdb.sub(linear_layer.bias.grad).abs().max()
     torch.distributed.barrier()
     print("   error in dLdb on global rank {}: {}".format(torch.distributed.get_rank(), error))
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     error = dLdX.sub(identity_layer.weight.grad).abs().max()
     torch.distributed.barrier()
     print("   error in dLdX on global rank {}: {}".format(torch.distributed.get_rank(), error))
-    assert error < 1.0e-6
+    assert error < 1.0e-6, error
 
     # Reset groups
     mpu.destroy_model_parallel()
diff --git a/tests/nn/moe/test_moe_layer.py b/tests/nn/moe/test_moe_layer.py
--- a/tests/nn/moe/test_moe_layer.py
+++ b/tests/nn/moe/test_moe_layer.py
@@ -13,6 +13,7 @@
 
 from fairscale.nn import MOELayer, Top2Gate
 from fairscale.utils import torch_version
+from fairscale.utils.testing import make_cudnn_deterministic
 
 pytestmark = pytest.mark.skipif(
     not (torch.cuda.is_available() and torch_version() >= (1, 8, 0)), reason="cuda and torch>=1.8.0 required"
@@ -68,6 +69,7 @@ def expert_params(device):
 @pg_test()
 @pytest.mark.parametrize("device", devices)
 def forward(device):
+    make_cudnn_deterministic()
     model_dim = 8
     num_experts = dist.get_world_size(dist.group.WORLD)
     input = torch.randn(4, 16, model_dim).to(device)
@@ -85,6 +87,7 @@ def forward(device):
 @pg_test()
 @pytest.mark.parametrize("device", devices)
 def forward_multi(device):
+    make_cudnn_deterministic()
     torch.set_printoptions(threshold=5000)
     num_local_experts = 4
     model_dim = 4
@@ -128,6 +131,7 @@ def forward(self, input):
 @pg_test()
 @pytest.mark.parametrize("device", devices)
 def forward_routing(device):
+    make_cudnn_deterministic()
     model_dim = 8
     num_experts = dist.get_world_size()
     input = torch.randn(4, 16, model_dim).to(device)
@@ -149,6 +153,7 @@ def forward_routing(device):
 @pg_test()
 @pytest.mark.parametrize("device", devices)
 def forward_routing_multi(device):
+    make_cudnn_deterministic()
     model_dim = 8
     num_local_experts = 4
     num_experts = dist.get_world_size(dist.group.WORLD) * num_local_experts
@@ -174,6 +179,7 @@ def forward_routing_multi(device):
 @pg_test()
 @pytest.mark.parametrize("device", devices)
 def backward(device):
+    make_cudnn_deterministic()
     loss = torch.nn.MSELoss()
     model_dim = 8
     num_experts = dist.get_world_size(dist.group.WORLD)
diff --git a/tests/nn/moe/test_top2gating.py b/tests/nn/moe/test_top2gating.py
--- a/tests/nn/moe/test_top2gating.py
+++ b/tests/nn/moe/test_top2gating.py
@@ -27,7 +27,7 @@ def do_test_forward(device):
     gate = Top2Gate(4, 6).to(device)
     capacity = 2 * 12 // 6
     l_aux, combine_weights, dispatch_mask = gate(input)
-    assert pytest.approx(l_aux.item(), 0.0283)
+    assert pytest.approx(l_aux.item(), rel=0.01) == 0.0267, l_aux
     assert combine_weights.shape == (12, 6, 4)
     assert dispatch_mask.shape == (12, 6, 4)
     assert torch.equal(combine_weights.bool(), dispatch_mask)
@@ -35,9 +35,9 @@ def do_test_forward(device):
     assert torch.all(combine_weights >= 0.0)
     assert torch.all(combine_weights <= 1.0)
     weights_sum = torch.sum(combine_weights).item()
-    assert round(weights_sum) == pytest.approx(weights_sum)
+    assert round(weights_sum) == pytest.approx(weights_sum), weights_sum
     # For this random seed, we get 12 slots filled.
-    assert weights_sum == pytest.approx(12.0)
+    assert weights_sum == pytest.approx(12.0), weights_sum
 
 
 def test_forward_cpu():
diff --git a/tests/optim/test_layerwise_gradient_scaler.py b/tests/optim/test_layerwise_gradient_scaler.py
--- a/tests/optim/test_layerwise_gradient_scaler.py
+++ b/tests/optim/test_layerwise_gradient_scaler.py
@@ -70,6 +70,7 @@ def load_data(model_type: str) -> Union[DataLoader, Tuple[Any, Any]]:
         torch.manual_seed(10)
         transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
 
+        # TODO: we should NOT do this download over and over again during test.
         train_ds = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
         train_ds_loader = torch.utils.data.DataLoader(train_ds, batch_size=128, shuffle=False, num_workers=2)
 
@@ -204,9 +205,11 @@ def train_vision_model(model: SimpleConvNet, per_layer_scaling=False):
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
 def test_vision_model() -> None:
-    # The os.environ below doesn't seem to be enough if the test is run on CI with many other tests together.
+    # The os.environ below doesn't seem to be enough if the test is run on CI with many other tests
+    # together.
     # see: https://app.circleci.com/pipelines/github/facebookresearch/fairscale/4086/workflows/72b1470a-55f8-4a45-afe5-04641b093bef/jobs/45179/tests#failed-test-0
     # Skipping for now.
+    # Also, TODO (Min): improving downloading code above before re-enable this.
     skip_a_test_if_in_CI()
     # Remove randomness from various sources while testing.
     torch.use_deterministic_algorithms(True)  # type: ignore
diff --git a/tests/optim/test_oss.py b/tests/optim/test_oss.py
--- a/tests/optim/test_oss.py
+++ b/tests/optim/test_oss.py
@@ -819,7 +819,7 @@ def run_grad_step(model, head, optimizer):
     dist.destroy_process_group()
 
 
-@skip_if_no_cuda
+@skip_if_single_gpu
 def test_state_dict_distributed():
     world_size = 2
     temp_file_name = tempfile.mkstemp()[1]
diff --git a/tests/optim/test_single_node_adascale.py b/tests/optim/test_single_node_adascale.py
--- a/tests/optim/test_single_node_adascale.py
+++ b/tests/optim/test_single_node_adascale.py
@@ -21,7 +21,7 @@
 
 from fairscale.optim import AdaScale
 from fairscale.utils.golden_testing_data import adascale_test_data
-from fairscale.utils.testing import skip_if_no_cuda
+from fairscale.utils.testing import make_cudnn_deterministic, skip_if_no_cuda
 from fairscale.utils.testing_memory import find_tensor_by_shape
 
 
@@ -63,6 +63,7 @@ def test_loss_accum_cpu():
 @pytest.mark.parametrize("test_case", adascale_test_data)
 def test_grad_accum(test_case, cpu):
     """Test the basic functionality on CPU/GPU with gradient accumulation without DDP"""
+    make_cudnn_deterministic()
     model = Linear(2, 2, bias=True)
     if not cpu:
         if torch.cuda.device_count() < 1: