Repository: pytorch
Issue: torch.compile() breaks when using DeepSpeed ZeRO Level 3 sharding
Labels: needs reproduction,triaged,oncall: pt2,module: dynamo,dynamo-triage-june2024

### 🐛 Describe the bug

torch.compile() breaks when using DeepSpeed ZeRO Level 3 sharding. I am fine-tuning Llama 2 using the Transformers codebase, and added a `torch.compile()` decorator over the MLP layer in this line of code:

https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L250

I am using Ubuntu 22.04, Torch 2.1.1, Transformers 4.35.2 and DeepSpeed 0.12.4. Here is my DeepSpeed config file, which works fine when `torch.compile()` is not used:

```
{
  "bf16": {
    "enabled": true
  },

  "zero_optimization": {
    "stage": 3,

    "zero_quantized_weights": false,
    "zero_hpz_partition_size": 1,
    "zero_quantized_gradients": false,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 7e8,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 7e8,
    "stage3_max_reuse_distance": 7e8,
    "stage3_gather_16bit_weights_on_model_save": true
  },

  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 1,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false,
  "flops_profiler": {
    "enabled": false,
    "profile_step": 1,
    "module_depth": -1,
    "top_modules": 1,
    "detailed": true,
    "output_file": null
  }
}

```

### Error logs

7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
7e72bd4e-02:     return fn(*args, **kwargs)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 451, in checkpoint
7e72bd4e-02:     return CheckpointFunction.apply(function, preserve, *args)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/autograd/function.py", line 539, in apply
7e72bd4e-02:     return super().apply(*args, **kwargs)  # type: ignore[misc]
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 230, in forward
7e72bd4e-02:     outputs = run_function(*args)
7e72bd4e-02:   File "/scratch/brr/brr/llm/models/llama_flash_attention.py", line 718, in custom_forward
7e72bd4e-02:     return module(*inputs, output_attentions, None)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
7e72bd4e-02:     return self._call_impl(*args, **kwargs)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
7e72bd4e-02:     result = forward_call(*args, **kwargs)
7e72bd4e-02:   File "/scratch/brr/brr/llm/models/llama_flash_attention.py", line 491, in forward
7e72bd4e-02:     hidden_states = self.mlp(hidden_states)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
7e72bd4e-02:     return self._call_impl(*args, **kwargs)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
7e72bd4e-02:     result = forward_call(*args, **kwargs)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
7e72bd4e-02:     return fn(*args, **kwargs)
7e72bd4e-02:   File "/scratch/brr/brr/llm/models/llama_flash_attention.py", line 264, in forward
7e72bd4e-02:     return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
7e72bd4e-02:     return self._call_impl(*args, **kwargs)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1557, in _call_impl
7e72bd4e-02:     args_result = hook(self, args)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 14, in wrapped_fn
7e72bd4e-02:     get_accelerator().range_push(func.__qualname__)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in <resume in wrapped_fn>
7e72bd4e-02:     ret_val = func(*args, **kwargs)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 392, in _pre_forward_module_hook
7e72bd4e-02:     self.pre_sub_module_forward_function(module)
7e72bd4e-02:   File "/scratch/miniconda3/envs/brr/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 493, in pre_sub_module_forward_function
7e72bd4e-02:     def pre_sub_module_forward_function(self, sub_module):
7e72bd4e-02: TypeError: 'int' object is not iterable

### Minified repro

Tried running this, but it did not appear to output minified code, just the same error message I posted above.

### Versions

Collecting environment information...
PyTorch version: 2.1.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 545.23.08
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.3
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.3
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.3
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.3
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             96
On-line CPU(s) list:                0-95
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz
CPU family:                         6
Model:                              106
Thread(s) per core:                 2
Core(s) per socket:                 24
Socket(s):                          2
Stepping:                           6
BogoMIPS:                           4000.04
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr\
 sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology cpuid tsc_known_freq pni pclm\
ulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor\
 lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vp\
id ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb a\
vx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves wbnoinvd arat avx512vbmi umip pku ospke avx512_vbmi2 gfni va\
es vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid fsrm md_clear arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          3 MiB (96 instances)
L1i cache:                          3 MiB (96 instances)
L2 cache:                           192 MiB (48 instances)
L3 cache:                           32 MiB (2 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-47
NUMA node1 CPU(s):                  48-95
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] numpyro==0.9.2
[pip3] torch==2.1.1
[pip3] torchaudio==2.1.1
[pip3] torchvision==0.16.1
[pip3] triton==2.1.0
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpyro                   0.9.2                    pypi_0    pypi
[conda] torch                     2.1.1                    pypi_0    pypi
[conda] torchaudio                2.1.1                    pypi_0    pypi
[conda] torchvision               0.16.1                   pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi

cc @ezyang @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @msaroufim @wconstab @bdhirsh @zou3519 @aakhundov