Repository: pytorch
Issue: DISABLED test_distributed_checkpoint_state_dict_type1_cuda (__main__.TestDistributedCheckpointCUDA)
Labels: oncall: distributed,module: flaky-tests,skipped

Platforms: linux

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/flakytest?name=test_distributed_checkpoint_state_dict_type1_cuda&suite=TestDistributedCheckpointCUDA&limit=100) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/35670700293).

Over the past 3 hours, it has been determined flaky in 3 workflow(s) with 4 failures and 3 successes.

**Debugging instructions (after clicking on the recent samples link):**
DO NOT ASSUME THINGS ARE OKAY IF THE CI IS GREEN. We now shield flaky tests from developers so CI will thus be green but it will be harder to parse the logs.
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Click on the Test step of the job so that it is expanded. Otherwise, the grepping will not work.
3. Grep for `test_distributed_checkpoint_state_dict_type1_cuda`
4. There should be several instances run (as flaky tests are rerun in CI) from which you can study the logs.



<details><summary>Sample error message</summary>

```
Traceback (most recent call last):
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 597, in wrapper
    self._join_processes(fn)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 837, in _join_processes
    self._check_return_codes(elapsed_time)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 886, in _check_return_codes
    raise RuntimeError(error)
RuntimeError: Process 1 exited with error code 10 and exception:
Traceback (most recent call last):
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 726, in run_test
    getattr(self, test_name)()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 599, in wrapper
    fn()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 3128, in wrapper
    method(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 3128, in wrapper
    method(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_device_type.py", line 485, in instantiated_test
    raise rte
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_device_type.py", line 465, in instantiated_test
    result = test(self, **param_kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 199, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/distributed/checkpoint_utils.py", line 44, in wrapper
    func(self, *args, **kwargs)
  File "/var/lib/jenkins/workspace/test/distributed/fsdp/test_distributed_checkpoint.py", line 67, in test_distributed_checkpoint
    state_dict = model.state_dict()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2204, in state_dict
    module.state_dict(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2204, in state_dict
    module.state_dict(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2204, in state_dict
    module.state_dict(
  [Previous line repeated 1 more time]
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2210, in state_dict
    hook_result = hook(self, destination, prefix, local_metadata)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_state_dict_utils.py", line 714, in _post_state_dict_hook
    processed_state_dict = _post_state_dict_hook_fn[fsdp_state._state_dict_type](
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_state_dict_utils.py", line 432, in _local_post_state_dict_hook
    sharded_tensor = init_from_local_shards(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/_shard/sharded_tensor/__init__.py", line 407, in init_from_local_shards
    return ShardedTensor._init_from_local_shards(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/_shard/sharded_tensor/api.py", line 753, in _init_from_local_shards
    dist.all_gather_object(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 81, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 3037, in all_gather_object
    input_tensor.resize_(max_object_size)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate more than 1EB memory.
Exception raised from allocate at /var/lib/jenkins/workspace/c10/cuda/CUDACachingAllocator.cpp:3623 (most recent call first):
C++ CapturedTraceback:
#4 std::_Function_handler<std::shared_ptr<c10::LazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > const> (), c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from Logging.cpp:0
#5 c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) from ??:0
#6 c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::allocate(unsigned long) from :0
#7 at::native::resize_bytes_cuda(c10::StorageImpl*, unsigned long) from ??:0
#8 at::native::resize_cuda_(at::Tensor const&, c10::ArrayRef<long>, std::optional<c10::MemoryFormat>) from ??:0
#9 at::_ops::resize_::redispatch(c10::DispatchKeySet, at::Tensor const&, c10::ArrayRef<c10::SymInt>, std::optional<c10::MemoryFormat>) from ??:0
#10 torch::ADInplaceOrView::resize_(c10::DispatchKeySet, at::Tensor const&, c10::ArrayRef<c10::SymInt>, std::optional<c10::MemoryFormat>) from VariableTypeManual.cpp:0
#11 at::_ops::resize_::redispatch(c10::DispatchKeySet, at::Tensor const&, c10::ArrayRef<c10::SymInt>, std::optional<c10::MemoryFormat>) from ??:0
#12 torch::autograd::VariableType::(anonymous namespace)::resize_(c10::DispatchKeySet, at::Tensor const&, c10::ArrayRef<c10::SymInt>, std::optional<c10::MemoryFormat>) from VariableTypeManual.cpp:0
#13 at::_ops::resize_::call(at::Tensor const&, c10::ArrayRef<c10::SymInt>, std::optional<c10::MemoryFormat>) from ??:0
#14 torch::autograd::THPVariable_resize_(_object*, _object*, _object*) from python_variable_methods.cpp:0
#15 method_vectorcall_VARARGS_KEYWORDS from :0
#16 _PyEval_EvalFrameDefault from ??:0
#17 _PyFunction_Vectorcall from ??:0
#18 PyObject_Call from ??:0
#19 _PyEval_EvalFrameDefault from ??:0
#20 _PyFunction_Vectorcall from ??:0
#21 _PyEval_EvalFrameDefault from ??:0
#22 method_vectorcall from :0
#23 PyObject_Call from ??:0
#24 _PyEval_EvalFrameDefault from ??:0
#25 _PyFunction_Vectorcall from ??:0
#26 _PyEval_EvalFrameDefault from ??:0
#27 _PyFunction_Vectorcall from ??:0
#28 _PyEval_EvalFrameDefault from ??:0
#29 _PyFunction_Vectorcall from ??:0
#30 _PyEval_EvalFrameDefault from ??:0
#31 _PyFunction_Vectorcall from ??:0
#32 _PyEval_EvalFrameDefault from ??:0
#33 method_vectorcall from :0
#34 _PyEval_EvalFrameDefault from ??:0
#35 method_vectorcall from :0
#36 _PyEval_EvalFrameDefault from ??:0
#37 method_vectorcall from :0
#38 _PyEval_EvalFrameDefault from ??:0
#39 method_vectorcall from :0
#40 _PyEval_EvalFrameDefault from ??:0
#41 method_vectorcall from :0
#42 _PyEval_EvalFrameDefault from ??:0
#43 _PyFunction_Vectorcall from ??:0
#44 PyObject_Call from ??:0
#45 _PyEval_EvalFrameDefault from ??:0
#46 _PyFunction_Vectorcall from ??:0
#47 PyObject_Call from ??:0
#48 _PyEval_EvalFrameDefault from ??:0
#49 _PyFunction_Vectorcall from ??:0
#50 PyObject_Call from ??:0
#51 _PyEval_EvalFrameDefault from ??:0
#52 method_vectorcall from :0
#53 _PyEval_EvalFrameDefault from ??:0
#54 method_vectorcall from :0
#55 _PyEval_EvalFrameDefault from ??:0
#56 method_vectorcall from :0
#57 _PyEval_EvalFrameDefault from ??:0
#58 method_vectorcall from :0
#59 _PyEval_EvalFrameDefault from ??:0
#60 _PyFunction_Vectorcall from ??:0
#61 _PyEval_EvalFrameDefault from ??:0
#62 method_vectorcall from :0
#63 PyObject_Call from ??:0
#64 _PyEval_EvalFrameDefault from ??:0
#65 _PyFunction_Vectorcall from ??:0
#66 _PyEval_EvalFrameDefault from ??:0
#67 _PyFunction_Vectorcall from ??:0
#68 _PyEval_EvalFrameDefault from ??:0
#69 _PyFunction_Vectorcall from ??:0
#70 _PyEval_EvalFrameDefault from ??:0
#71 _PyFunction_Vectorcall from ??:0
#72 _PyEval_EvalFrameDefault from ??:0
#73 _PyEval_Vector from :0
#74 PyEval_EvalCode from ??:0
#75 run_eval_code_obj from :0
#76 run_mod from :0
#77 PyRun_StringFlags.localalias from :0
#78 PyRun_SimpleStringFlags.localalias from :0
#79 Py_RunMain.localalias from :0
#80 Py_BytesMain from ??:0
#81 __libc_start_main from ??:0
#82 _start from ??:0


To execute this test, run the following from the base repo dir:
    python test/distributed/fsdp/test_distributed_checkpoint.py TestDistributedCheckpointCUDA.test_distributed_checkpoint_state_dict_type1_cuda

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0
```

</details>


Test file path: `distributed/fsdp/test_distributed_checkpoint.py`

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @clee2000 @wdvr