[functorch] [vmap] [SymInt][fake tensor]
Labels: triaged, module: functorch

Goal:

1. Create an input and target tensor
2. Create a list of models for ensemble
3. Use vmap to create an ensemble
4. Obtain the vmap function and pass it to aot_function, to obtain the forward and backward graphs
5. Although the weights and parameters can fit into memory, the activations might not fit in memory and hence, I set `config.use_fake_tensor = True`

Repro:

```
from typing import List
import torch
from functorch.compile import aot_function, print_compile, config, aot_module
from functorch import make_functional_with_buffers, vmap, combine_state_for_ensemble
from functorch._src.named_members_polyfill import _named_parameters, _named_buffers
from torchvision.models import resnet18
config.use_fake_tensor = True


def fake_compiler(fx_g, inps):
    print(fx_g.code)
    output_node = [node for node in fx_g.graph.nodes if node.op == 'output'][0]
    output_data = [node.meta['val'] for node in output_node.args[0]]
    def new_f(*args):
        return output_data
    return new_f

inp = torch.randn(32, 3, 224, 224, dtype=torch.float32).cuda()
targets = torch.zeros(32, dtype=int).cuda()
b_models:List[torch.nn.Module] = [resnet18().cuda() for _ in range(5)]

func_model, params, buffers = combine_state_for_ensemble(b_models)
for p in params:
    p.requires_grad = True

def compute_loss(weights, buffers, batch, targets):
    output = func_model(weights, buffers, batch)
    loss = torch.nn.functional.nll_loss(output,targets)
    return loss

parallel_func = vmap(compute_loss, in_dims=(0,0,None, None))

aot_func = aot_function(parallel_func, fake_compiler)
out = aot_func(params, buffers, inp, targets)
out.mean().backward()
```
Error Stack Trace:

```
Traceback (most recent call last):
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/testing/test_patch.py", line 36, in <module>
    out = aot_func(params, buffers, inp, targets)
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/pytorch/functorch/_src/aot_autograd.py", line 670, in returned_function
    compiled_fn = create_aot_dispatcher_function(
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/torchdynamo/torchdynamo/utils.py", line 76, in time_wrapper
    r = func(*args, **kwargs)
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/pytorch/functorch/_src/aot_autograd.py", line 523, in create_aot_dispatcher_function
    aot_dispatch_autograd(flat_fn, fake_flat_tensor_args, aot_config)
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/pytorch/functorch/_src/aot_autograd.py", line 350, in aot_dispatch_autograd
    out = flat_fn(*flat_args)
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/pytorch/functorch/_src/aot_autograd.py", line 650, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/pytorch/functorch/_src/vmap.py", line 362, in wrapped
    return _flat_vmap(
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/pytorch/functorch/_src/vmap.py", line 35, in fn
    return f(*args, **kwargs)
  File "/net/rcstorenfs02/ifs/rc_labs/idreos_lab/users/spurandare/develop/pytorch/functorch/_src/vmap.py", line 486, in _flat_vmap
    vmap_level = _vmap_increment_nesting(batch_size, randomness)
TypeError: _vmap_increment_nesting(): incompatible function arguments. The following argument types are supported:
    1. (arg0: int, arg1: str) -> int

Invoked with: t0.size(0), 'error'
```
Debugging to check the values passed to `_vmap_increment_nesting()`
Debug Trace:
```
-> vmap_level = _vmap_increment_nesting(batch_size, randomness)
(Pdb) print(batch_size)
t0.size(0)
(Pdb) print(randomness)
error
(Pdb) print(type(batch_size))
<class 'torch.SymIntNode'>
(Pdb) exit()
```

cc @zou3519 @Chillee @samdow @soumith @kshitij12345 @janeyx99