Repository: pytorch
Issue: vmap and rnn/lstm "accessing '.data' under vmap transform is not allowed"
Labels: triaged,module: functorch

### 🐛 Describe the bug

Quick disclaimer before I start: This is literally the first time I ever posted a bug report on GitHub; I've done my best to be clear and complete and follow the guidelines, but plz forgive me if anything here is incomplete/unclear :X

I have a use case where it'd be very ( very ) nice to be able to run a collection of models ( really a bunch of differently-weighted versions of the same model ) in a vectorized way rather than looping through each of them individually. The models here all have the same architecture, and are all ( at any given time ) doing inference on the same batch of data - they differ only in their weights. The error here only occurs with the LSTM component of the model. The rest of the model is just a bunch of linear layers, which I have tested independently, and are not subject to this error - only the LSTM does this. Below is minimal code to reproduce the problem ( written following guidance provided [here](https://pytorch.org/tutorials/intermediate/ensembling.html) ).
```python
import torch as pt

BSIZE     = 32
SEQLEN    = 16
LSTMSIZE  = 64
INPUTSIZE = 8
NUMMODELS = 4

def BuildLSTM( inputSize=INPUTSIZE, lstmSize=LSTMSIZE ):
	return pt.nn.LSTM( input_size=INPUTSIZE, hidden_size=LSTMSIZE, batch_first=True )

BASELSTM = BuildLSTM().to( 'meta' )
def fLSTM( params, buffers, inputBatch ): 
	return pt.func.functional_call( BASELSTM, (params,buffers), (inputBatch,) )

dummyBatch     = pt.rand( (BSIZE, SEQLEN, INPUTSIZE) ) #(32,16,8)
LSTMstack      = [ BuildLSTM() for _ in range( NUMMODELS ) ] #len = 4
PARAMS,BUFFERS = pt.func.stack_module_state( LSTMstack )

pt.vmap( fLSTM, in_dims=(0,0,None) )( PARAMS, BUFFERS, dummyBatch )
```

The exact error produced is:
```
Traceback (most recent call last):
  
  File "vmaptest.py", line 20, in <module>
    pt.vmap( fLSTM, in_dims=(0,0,None) )( PARAMS, BUFFERS, dummyBatch )
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/_functorch/vmap.py", line 434, in wrapped
    return _flat_vmap(
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/_functorch/vmap.py", line 619, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  
  File "vmaptest.py", line 18, in fLSTM
    return pt.func.functional_call( BASELSTM, (params,buffers), (inputBatch,) )
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/_functorch/functional_call.py", line 143, in functional_call
    return nn.utils.stateless._functional_call(
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/nn/utils/stateless.py", line 262, in _functional_call
    return module(*args, **kwargs)
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/nn/modules/rnn.py", line 760, in forward
    self._init_flat_weights()
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/nn/modules/rnn.py", line 139, in _init_flat_weights
    self.flatten_parameters()
  
  File "/home/USERNAME/miniconda3/envs/ENVNAME/lib/python3.8/site-packages/torch/nn/modules/rnn.py", line 167, in flatten_parameters
    if (not isinstance(fw.data, Tensor) or not (fw.data.dtype == dtype) or not fw.data.is_cuda or not torch.backends.cudnn.is_acceptable(fw.data)):
    
RuntimeError: accessing `data` under vmap transform is not allowed
```

Superficially, this seems quite similar to: https://github.com/pytorch/pytorch/issues/103161, minus the jacrev stuff. I'll add a couple additional observations in a reply to this issue to avoid making this post excessively long. Environment details below.

### Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04.1) 11.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.8.16 | packaged by conda-forge | (default, Feb  1 2023, 16:01:55)  [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.19.0-46-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080
Nvidia driver version: 525.125.06
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.2
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.6.0
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:                   43 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          12
On-line CPU(s) list:             0-11
Vendor ID:                       AuthenticAMD
Model name:                      AMD Ryzen 5 2600 Six-Core Processor
CPU family:                      23
Model:                           8
Thread(s) per core:              2
Core(s) per socket:              6
Socket(s):                       1
Stepping:                        2
BogoMIPS:                        6799.27
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb hw_pstate ssbd ibpb vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt sha_ni xsaveopt xsavec xgetbv1 xsaves clzero irperf xsaveerptr arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif overflow_recov succor smca sev sev_es
Virtualization:                  AMD-V
L1d cache:                       192 KiB (6 instances)
L1i cache:                       384 KiB (6 instances)
L2 cache:                        3 MiB (6 instances)
L3 cache:                        16 MiB (2 instances)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-11
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Mitigation; untrained return thunk; SMT vulnerable
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] pytorch-lightning==2.0.2
[pip3] torch==2.0.1
[pip3] torchmetrics==0.11.4
[conda] blas                 2.16                              mkl    conda-forge
[conda] libblas              3.8.0                          16_mkl    conda-forge
[conda] libcblas             3.8.0                          16_mkl    conda-forge
[conda] libjpeg-turbo        2.0.0                      h9bf148f_0    pytorch-nightly
[conda] liblapack            3.8.0                          16_mkl    conda-forge
[conda] liblapacke           3.8.0                          16_mkl    conda-forge
[conda] mkl                  2020.2                            256  
[conda] numpy                1.22.4                 py38h99721a1_0    conda-forge
[conda] pytorch              2.0.1     py3.8_cuda11.8_cudnn8.7.0_0    pytorch
[conda] pytorch-cuda         11.8                       h7e8668a_5    pytorch
[conda] pytorch-lightning    2.0.2                    pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex        1.0                              cuda    pytorch
[conda] torchmetrics         0.11.4                   pyhd8ed1ab_0    conda-forge
[conda] torchtriton          2.0.0                            py38    pytorch

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99