[RFC] Modularizing FSDPv2 buffer for custom backend implementation

# [RFC] Modularizing FSDPv2 buffer for custom backend implementation

## Background

While working on the TorchTPU effort, one inefficiency we have noticed when analysing performance in FSDPv2 is the significant toil that comes from buffer concatenation as part of FSDP parameter grouping. We found that in the reduce scatter process for `foreach_reduce` in _fsdp_collective.py ([specific code line reference](https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py#L597-L632)), a couple operations are being called that trigger buffer allocations for TPUs.

We suspect that other backends might encounter this inefficiency. Therefore we propose that FSDPv2 have these sections modularized such that the option for different actions may be taken.

## The proposal and our specific change

The core change we would like to achieve for our implementation is changing the foreach_reduce_scatter to be done per parameter rather than concatenating them into one tensor as it is currently done. In TPUs, this avoids an inefficient buffer allocation and materialization. We have a rough prototype on how this might be achieved in [a small prototype](https://github.com/pytorch/pytorch/compare/main...pgmoka:pytorch:fsdp-optimal-collectives?expand=1). From these changes we see pretty significant gains on FSDPv2 overhead on TPU.

To do that, we need a point of entry into the code. We currently have `reduce_scatter_comm`, but we need to also control the fallback and earlier sections of code which cover buffer allocation for the code. Therefore, we need a wider section of code to be modularized and replaceable.

## Code path control gap

We would like the opinion of pytorch maintainers on what might be the best way to control code flow here. We have identified from looking at the code:
1. Using a flag such that users might call something like `torch.distributed.fsdp.fully_shard(module, *, mesh=None, reshard_after_forward=None, shard_placement_fn=None, mp_policy=MixedPrecisionPolicy(param_dtype=None, reduce_dtype=None, output_dtype=None, **cast_forward_inputs=True**), offload_policy=OffloadPolicy(), ignored_params=None, dp_mesh_dims=None, group_collectives=True)` which might trigger the alternative path we are proposing above
2. Modularize that section of code into a custom op(or series of custom ops) that the backend may override. This is expanding upon what has been done for `all_gather_copy_in_cuda` (a mechanism we are already using).
3. Potentially another pytorch way that might be preferable

We would like an opinion on how to carry on with this point from pytorch maintainers.

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @kapilsh @zhaojuanmao @mrshenli @rohan-varma @chauhang @mori360 @ppwwyyxx

Start with torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py at the referenced foreach_reduce_scatter code (lines 597-632), then review reduce_scatter_comm and the linked optimal-collectives prototype. The first step is to get maintainer guidance on the extension point; done requires an agreed design for backend-specific control of buffer allocation and per-parameter reduce-scatter behavior.

Written by the indexing model from the issue text.