Fine-tuning HuggingFace wav2vec 2.0 with `torch.compile`
Labels: high priority, oncall: distributed, module: crash, triaged, module: nccl, ezyang's list, bug, oncall: pt2, module: dynamic shapes

### 🐛 Describe the bug

I followed the example to fine-tune HuggingFace's wav2vec 2.0 for speech recognition, using `torch.compile`, aiming to get faster training. However, I ran into an issue as outlined in the error logs.

I suspect that HuggingFace's wav2vec 2.0 is not yet supported in PyTorch 2.0 and needs some modification to ensure compatibility when running `torch.compile`. It's mostly related to creating mask tensors for SpecAugment.

This issue seems to also be related: [fairseq Hubert with torch.compile](https://github.com/pytorch/pytorch/issues/97511). Same issue also raised in [HuggingFace](https://github.com/huggingface/transformers/issues/22849).

### Error logs

```
***** Running training *****
  Num examples = 3,478
  Num Epochs = 15
  Instantaneous batch size per device = 4
  Total train batch size (w. parallel, distributed & accumulation) = 4
  Gradient Accumulation steps = 1
  Total optimization steps = 13,050
  Number of trainable parameters = 311,270,569
  0%|                                                                                                                                                                    | 0/13050 [00:00<?, ?it/sTraceback (most recent call last):
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 670, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.fake_example_inputs())
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/debug_utils.py", line 1055, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/__init__.py", line 1390, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_inductor/compile_fx.py", line 455, in compile_fx
    return aot_autograd(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/backends/common.py", line 48, in compiler_fn
    cg = aot_module_simplified(gm, example_inputs, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 2822, in aot_module_simplified
    compiled_fn = create_aot_dispatcher_function(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 2515, in create_aot_dispatcher_function
    compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1715, in aot_wrapper_dedupe
    return compiler_fn(flat_fn, leaf_flat_args, aot_config)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 2104, in aot_dispatch_autograd
    fx_g = make_fx(joint_forward_backward, aot_config.decompositions)(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 714, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 443, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 652, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 459, in wrapped
    out = f(*tensors)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1158, in traced_joint
    return functionalized_f_helper(primals, tangents)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1110, in functionalized_f_helper
    f_outs = flat_fn_no_input_mutations(fn, f_primals, f_tangents, meta, keep_input_mutations)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1078, in flat_fn_no_input_mutations
    outs = flat_fn_with_synthetic_bases_expanded(fn, primals, primals_after_cloning, maybe_tangents, meta, keep_input_mutations)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1050, in flat_fn_with_synthetic_bases_expanded
    outs = forward_or_joint(fn, primals_before_cloning, primals, maybe_tangents, meta, keep_input_mutations)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1019, in forward_or_joint
    backward_out = torch.autograd.grad(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/autograd/__init__.py", line 269, in grad
    return handle_torch_function(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/overrides.py", line 1534, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_inductor/overrides.py", line 38, in __torch_function__
    return func(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/autograd/__init__.py", line 303, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 487, in __torch_dispatch__
    return self.inner_torch_dispatch(func, types, args, kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 512, in inner_torch_dispatch
    out = proxy_call(self, func, args, kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py", line 345, in proxy_call
    out = func(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_ops.py", line 287, in __call__
    return self._op(*args, **kwargs or {})
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 987, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 1162, in dispatch
    op_impl_out = op_impl(self, func, *args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 453, in index_tensor
    check_no_bool_index_tensors(func, *args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 432, in check_no_bool_index_tensors
    raise DynamicOutputShapeException(func)
torch._subclasses.fake_tensor.DynamicOutputShapeException: aten.index.Tensor

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/wilson_bookbotkids_com/transformers/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py", line 775, in <module>
    main()
  File "/home/wilson_bookbotkids_com/transformers/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py", line 723, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 1664, in train
    return inner_training_loop(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 1940, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 2735, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 2767, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 82, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1684, in forward
    outputs = self.wav2vec2(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1316, in forward
    hidden_states = self._mask_hidden_states(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1249, in _mask_hidden_states
    if not getattr(self.config, "apply_spec_augment", True):
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1259, in <graph break in _mask_hidden_states>
    mask_time_indices = _compute_mask_indices(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py", line 1266, in <graph break in _mask_hidden_states>
    mask_time_indices = torch.tensor(mask_time_indices, device=hidden_states.device, dtype=torch.bool)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 337, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
    transformations(instructions, code_options)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 311, in transform
    tracer.run()
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
    super().run()
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
    and self.step()
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
    getattr(self, inst.opname)(inst)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1792, in RETURN_VALUE
    self.output.compile_subgraph(
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 517, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 588, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 675, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e) from e
torch._dynamo.exc.BackendCompilerFailed: debug_wrapper raised DynamicOutputShapeException: aten.index.Tensor

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True
```

### Minified repro

Install HuggingFace Transformers from source:
```
git clone https://github.com/huggingface/transformers.git
pip install transformers/
```

Run training script
```sh
python transformers/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py \
    --dataset_name="common_voice" \
    --model_name_or_path="facebook/wav2vec2-large-xlsr-53" \
    --dataset_config_name="tr" \
    --output_dir="./wav2vec2-common_voice-tr-demo-dist" \
    --preprocessing_num_workers="16" \
    --overwrite_output_dir \
    --num_train_epochs="15" \
    --per_device_train_batch_size="4" \
    --gradient_accumulation_steps="1" \
    --learning_rate="3e-4" \
    --warmup_steps="500" \
    --evaluation_strategy="steps" \
    --text_column_name="sentence" \
    --save_steps="400" \
    --eval_steps="100" \
    --logging_steps="1" \
    --layerdrop="0.0" \
    --save_total_limit="3" \
    --freeze_feature_encoder \
    --gradient_checkpointing \
    --chars_to_ignore , ? . ! - \; \: \" “ % ‘ ” � \
    --fp16 \
    --group_by_length \
    --do_train --do_eval \
    --torch_compile True
```

### Versions

<details>
<summary>Versions</summary>

Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (Debian 8.3.0-6) 8.3.0
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.28

Python version: 3.9.0 | packaged by conda-forge | (default, Nov 26 2020, 07:57:39)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-4.19.0-23-cloud-amd64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: 12.0.76
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 525.60.13
cuDNN version: Could not collect
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
Byte Order:          Little Endian
Address sizes:       46 bits physical, 48 bits virtual
CPU(s):              12
On-line CPU(s) list: 0-11
Thread(s) per core:  2
Core(s) per socket:  6
Socket(s):           1
NUMA node(s):        1
Vendor ID:           GenuineIntel
CPU family:          6
Model:               85
Model name:          Intel(R) Xeon(R) CPU @ 2.20GHz
Stepping:            7
CPU MHz:             2200.152
BogoMIPS:            4400.30
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            1024K
L3 cache:            39424K
NUMA node0 CPU(s):   0-11
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 rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat avx512_vnni md_clear arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.0.1+cu117
[pip3] torchaudio==2.0.2+cu117
[pip3] triton==2.0.0
[conda] numpy                     1.24.3                   pypi_0    pypi
[conda] torch                     2.0.1+cu117              pypi_0    pypi
[conda] torchaudio                2.0.2+cu117              pypi_0    pypi
[conda] triton                    2.0.0                    pypi_0    pypi

</details>

cc @ezyang @gchanan @zou3519 @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @msaroufim @wconstab @ngimel @bdhirsh @anijain2305 @soumith