Repository: pytorch
Issue: `NotImplementedError` when using `torch.distributed.launch`
Labels: oncall: distributed,module: functorch

First of all, really thanks for the nice work and huge contribution.

I tried to get the pre-sample grad from the instruction here https://pytorch.org/functorch/1.13/notebooks/per_sample_grads.html, and it worked for me on CPU-mode, single-GPU mode. However, I encounter the Error: 
```NotImplementedError: Cannot access storage of TensorWrapper```
 when I switch to the multi-GPU mode by calling 
```python -m torch.distributed.launch --nproc_per_node=2 train.py```
The Traceback shows the Error comes from the forward function when it is called from `grad_and_value()`. Note again, the same code works fine for single-GPU mode.

The whole code is as follows:
```python
model.eval()
fmodel, params, buffers = make_functional_with_buffers(model)
for batch_idx, (input, target) in enumerate(loader_plot):
    input, target = input.cuda(), target.cuda()
    if mixup_fn is not None:
        input, target = mixup_fn(input, target)
    if args.channels_last:
        input = input.contiguous(memory_format=torch.channels_last)
        
    

    def compute_loss_stateless_model(params, buffers, sample, target):
        batch = sample.unsqueeze(0)
        targets = target.unsqueeze(0)

        predictions = fmodel(params, buffers, batch) 
        loss = loss_fn(predictions, targets)
        return loss

    ft_compute_grad = grad_and_value(compute_loss_stateless_model)


    ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, None, 0, 0))
    with torch.no_grad():
        per_sample_opt = ft_compute_sample_grad(params, buffers, input, target) # Error comes from here
        losses = per_sample_opt[1].cpu().numpy()
        sample_grads = per_sample_opt[0]
        device = sample_grads[0].device
        global_grad_norm = torch.zeros(len(losses), device=device)
        for layer_grads in sample_grads:
            dims = list(range(1, layer_grads.dim()))
            global_grad_norm.add_(layer_grads.pow(2).sum(dim=dims))

        global_grad_norm = torch.sqrt(global_grad_norm).cpu().numpy()


    
    losses_plot = losses if losses_plot is None else np.append(losses_plot, losses)
    gradnorm_plot = global_grad_norm if gradnorm_plot is None else np.append(gradnorm_plot, global_grad_norm)
    
    torch.cuda.synchronize()
output_dir = get_outdir(args.output if args.output else './output/train', 'grad_figures')        
np.savez('{}/{}_{}.npz'.format(output_dir,epoch,args.local_rank),
         loss = losses_plot, norm = gradnorm_plot)
```


The Error information is as follows:
```
Traceback (most recent call last):
  File "train.py", line 1139, in <module>
    main()
  File "train.py", line 856, in main
    train_metrics = train_one_epoch(
  File "train.py", line 947, in train_one_epoch
    per_sample_opt = ft_compute_sample_grad(params, buffers, input, target)
  File "/usr/local/lib/python3.8/dist-packages/functorch/_src/vmap.py", line 362, in wrapped
    return _flat_vmap(
  File "/usr/local/lib/python3.8/dist-packages/functorch/_src/vmap.py", line 35, in fn
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/functorch/_src/vmap.py", line 489, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/functorch/_src/vmap.py", line 35, in fn
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/functorch/_src/eager_transforms.py", line 1111, in wrapper
    output = func(*args, **kwargs)
  File "train_grad2.py", line 938, in compute_loss_stateless_model
    predictions = fmodel(params, buffers, batch)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/functorch/_src/make_functional.py", line 282, in forward
    return self.stateless_model(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/distributed.py", line 1060, in forward
    self.require_forward_param_sync = False
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/profiler.py", line 493, in __exit__
    torch.ops.profiler._record_function_exit(self.handle)
  File "/usr/local/lib/python3.8/dist-packages/torch/_ops.py", line 442, in __call__
    return self._op(*args, **kwargs or {})
NotImplementedError: Cannot access storage of TensorWrapper
```

Thanks in advance if you could help me to find the solution. Best

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @zou3519 @Chillee @samdow @soumith