PyTorch 1.10.0, `distributed/optim/test_zero_redundancy_optimizer` fails on A100

[casparvl]

🐛 Bug

The following two tests from distributed/optim/test_zero_redundancy_optimizer fail

ERROR: test_zero_model_parallel_with_bucket_view (__main__.TestZeroRedundancyOptimizerDistributed)
ERROR: test_zero_model_parallel_without_bucket_view (__main__.TestZeroRedundancyOptimizerDistributed)

With the following error:

ERROR: test_zero_model_parallel_without_bucket_view (distributed.optim.test_zero_redundancy_optimizer.TestZeroRedundancyOptimizerDistributed)
Check that ZeRO works with model parallelism where layers are sharded
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/tmp/casparl/eb_tmp/eb-goi82_nt/tmp69ke_p3z/lib/python3.9/site-packages/torch/testing/_internal/common_distributed.py", line 418, in wrapper
    self._join_processes(fn)
  File "/tmp/casparl/eb_tmp/eb-goi82_nt/tmp69ke_p3z/lib/python3.9/site-packages/torch/testing/_internal/common_distributed.py", line 637, in _join_processes
    self._check_return_codes(elapsed_time)
  File "/tmp/casparl/eb_tmp/eb-goi82_nt/tmp69ke_p3z/lib/python3.9/site-packages/torch/testing/_internal/common_distributed.py", line 682, in _check_return_codes
    raise RuntimeError(error)
RuntimeError: Process 0 exited with error code 10 and exception:
Traceback (most recent call last):
  File "/tmp/casparl/eb_tmp/eb-goi82_nt/tmp69ke_p3z/lib/python3.9/site-packages/torch/testing/_internal/common_distributed.py", line 536, in run_test
    getattr(self, test_name)()
  File "/tmp/casparl/eb_tmp/eb-goi82_nt/tmp69ke_p3z/lib/python3.9/site-packages/torch/testing/_internal/common_distributed.py", line 420, in wrapper
    fn()
  File "/tmp/casparl/eb_tmp/eb-goi82_nt/tmp69ke_p3z/lib/python3.9/site-packages/torch/testing/_internal/common_distributed.py", line 111, in wrapper
    return func(*args, **kwargs)
  File "/tmp/casparl/eb_build/PyTorch/1.10.0/foss-2021a-CUDA-11.3.1/pytorch/test/distributed/optim/test_zero_redundancy_optimizer.py", line 913, in test_zero_model_parallel_without_bucket_view
    self._test_zero_model_parallel(parameters_as_bucket_view=False)
  File "/tmp/casparl/eb_build/PyTorch/1.10.0/foss-2021a-CUDA-11.3.1/pytorch/test/distributed/optim/test_zero_redundancy_optimizer.py", line 885, in _test_zero_model_parallel
    assert torch.allclose(
AssertionError: Losses differ between local optim and ZeroRedundancyOptimizer)

To Reproduce

Steps to reproduce the behavior:

1. Build PyTorch from source (see my build command below)
2. Run python -m unittest distributed.optim.test_zero_redundancy_optimizer.TestZeroRedundancyOptimizerDistributed.test_zero_model_parallel_without_bucket_view -v or python -m unittest distributed.optim.test_zero_redundancy_optimizer.TestZeroRedundancyOptimizerDistributed.test_zero_model_parallel_with_bucket_view -v

Expected behavior

I expected the tests to pass.

Environment

Please copy and paste the output from our
environment collection script
(or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

• PyTorch Version (e.g., 1.0): PyTorch 1.10.0
• OS (e.g., Linux): CentOS 8
• How you installed PyTorch (conda, pip, source): source
• Build command you used (if compiling from source):

USE_CUPTI_SO=1 PYTORCH_BUILD_VERSION=1.10.0 PYTORCH_BUILD_NUMBER=1 MAX_JOBS=72 BLAS=Eigen WITH_BLAS=open USE_FFMPEG=1 BUILD_CUSTOM_PROTOBUF=0 USE_IBVERBS=1 USE_CUDA=1 CUDNN_LIB_DIR=/tmp/sw_stack_gpu/software/cuDNN/8.2.1.32-CUDA-11.3.1/lib64 CUDNN_INCLUDE_DIR=/tmp/sw_stack_gpu/software/cuDNN/8.2.1.32-CUDA-11.3.1/include USE_SYSTEM_NCCL=1 NCCL_INCLUDE_DIR=/tmp/sw_stack_gpu/software/NCCL/2.10.3-GCCcore-10.3.0-CUDA-11.3.1/include USE_METAL=0   /tmp/sw_stack_gpu/software/Python/3.9.5-GCCcore-10.3.0/bin/python setup.py build

• Python version: 3.9.5
• CUDA/cuDNN version: CUDA: 11.3.1, cuDNN: 8.2.1.32
• GPU models and configuration: N/A
• Any other relevant information: Ran test on a node with 4* A100 GPUs

Additional context

I modified the test to print the local_loss and ddp_loss, which are compared through an assert on torch.allclose(...), to check if they were really different. They were not very different at all:

Local_loss: 27.58694839477539, ddp_loss: 27.58181381225586

But that is too much for the default relative tolerance of a torch.allclose(...)

The test passes if I set

export NVIDIA_TF32_OVERRIDE=0

before running it. Thus, the cause of the small deviation seems to be the use of TensorFloat32.

I see three options for a fix:

1. Detect if TF32 is supported/used, and if so, increase the tolerance.
2. Increase the tolerance for this test in general, so that it pass regardless of whether TF32 is used
3. Disable the use of TF32 before running this test

I personally think option (3) is not very good, because it doesn't represent how a user would run PyTorch. I think (2) is the easiest to implement, and I would still consider it pretty safe (the tolerance doesn't need to be all that large for this case to pass succesfully). I'll create a PR based on (2) and we can discuss if that's ok in the PR.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang

[casparvl]

Two more tests that seem to show this behaviour are:

test_ops.TestCommonCUDA.test_reference_testing_addbmm_cuda_complex64
test_ops.TestCommonCUDA.test_reference_testing_addbmm_cuda_float32

Although I didn't get these failures the first time I compiled PyTorch & ran the full test suite, so I'm not sure if there is any randomness to it. Now, when running these individual tests, it seems pretty consistent though. Differences are again small, (see Traceback)

Traceback

======================================================================
FAIL: test_reference_testing_addbmm_cuda_complex64 (test_ops.TestCommonCUDA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1396, in wrapper
    method(*args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1396, in wrapper
    method(*args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_device_type.py", line 368, in instantiated_test
    result = test(self, **param_kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_device_type.py", line 734, in test_wrapper
    return test(*args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_device_type.py", line 917, in only_fn
    return fn(self, *args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1023, in wrapper
    fn(*args, **kwargs)
  File "/tmp/casparl/eb_build/PyTorch/1.10.0/foss-2021a-CUDA-11.3.1/pytorch/test/test_ops.py", line 202, in test_reference_testing
    self.compare_with_reference(op, op.ref, sample_input)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1621, in compare_with_reference
    self.assertEqual(actual, expected, exact_device=False)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1899, in assertEqual
    self.assertEqual(
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1877, in assertEqual
    super().assertTrue(result, msg=self._get_assert_msg(msg, debug_msg=debug_msg))
AssertionError: False is not true : Tensors failed to compare as equal!With rtol=0.0012 and atol=1e-05, found 2 element(s) (out of 50) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.5579074661595502 ((54.0982666015625+381.3050537109375j) vs. (53.867881774902344+381.81317138671875j)), which occurred at index (4, 2).

======================================================================
FAIL: test_reference_testing_addbmm_cuda_float32 (test_ops.TestCommonCUDA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1396, in wrapper
    method(*args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1396, in wrapper
    method(*args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_device_type.py", line 368, in instantiated_test
    result = test(self, **param_kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_device_type.py", line 734, in test_wrapper
    return test(*args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_device_type.py", line 917, in only_fn
    return fn(self, *args, **kwargs)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1023, in wrapper
    fn(*args, **kwargs)
  File "/tmp/casparl/eb_build/PyTorch/1.10.0/foss-2021a-CUDA-11.3.1/pytorch/test/test_ops.py", line 202, in test_reference_testing
    self.compare_with_reference(op, op.ref, sample_input)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1621, in compare_with_reference
    self.assertEqual(actual, expected, exact_device=False)
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1899, in assertEqual
    self.assertEqual(
  File "/tmp/casparl/eb_tmp/eb-k860rj68/tmpfhjobzgk/lib/python3.9/site-packages/torch/testing/_internal/common_utils.py", line 1877, in assertEqual
    super().assertTrue(result, msg=self._get_assert_msg(msg, debug_msg=debug_msg))
AssertionError: False is not true : Tensors failed to compare as equal!With rtol=1.3e-05 and atol=1.3e-05, found 49 element(s) (out of 50) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.107940673828125 (-417.7663269042969 vs. -417.874267578125), which occurred at index (1, 7).

And the problem is consistently resolved by setting NVIDIA_TF32_OVERRIDE=0. I guess I have to increase the tolerance here as well.

Update: it does seem inconsistent, and (I guess) dependent on the state of the RNG. If I run the full test_ops set of tests,

python -m unittest test_ops -v

The index that showed the maximum difference is (3,3) instead of (1,7) for the test_reference_testing_addbmm_cuda_float32 test.

[zhaojuanmao]

@casparvl feel free to send out PR, the option 2 sounds reasonable, thanks for contributing!

[casparvl]

@zhaojuanmao it's already there actually :) #67765

[pritamdamania87]

@malfet Was wondering if anyone of our CI jobs run on A100 machines? It would be hard to maintain a test suite sanity if we don't run on A100s and we discover several broken tests like these manually while running on A100s.