Support for torch.func

Description & Motivation

Torch.func offers lots of flexibility. It's however not supported by lightning.
A few bugs I've noticed so far when trying to use it:

• When we init our model and then use a trainer, if we have used model.to("meta") to only keep the function version, it then fails as the trainers try to send it to device
• Params are not sent to the device, which then create device issues

Minimal example:

class Module(L.LightningModule):
    def __init__(self, cfg):
        super().__init__()
        models = [torch.nn.Linear(in_features, out_features) for i in range(num_models)]
        self.fmodel = copy.deepcopy(models[0])
        self.model_params, self.model_buffers = torch.func.stack_module_state(models)

   def training_step(self, batch, batch_idx):
       # It is possible to vmap directly over torch.func.functional_call,
      # but wrapping it in a function makes it clearer what is going on.
      def call_single_model(params, buffers, data):
          return torch.func.functional_call(self.fmodel, (params, buffers), (data,))
      output = torch.vmap(call_single_model, (0, 0, None))(self.model_params,  self.model_buffers, batch)
      return output.mean()

Pitch

No response

Alternatives

No response

Additional context

No response

cc @borda

Start with the minimal example using torch.func.stack_module_state, functional_call, vmap, and model.to("meta"), then trace how Trainer moves models and parameters to a device. Done means the example runs through training without device errors while preserving the torch.func setup for model parameters and buffers.