Enabling FP8 all-gather for TE Float8Tensor when using Torch FSDP2

[youngeunkwon0405]

Description

This PR enables FP8 all-gather for TE Float8Tensor when using the Torch FSDP2 (a.k.a. per-parameter-sharding FSDP).
This feature will be automatically enabled when a user creates a module with the transformer_engine.pytorch.fp8_model_init.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refractor

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[youngeunkwon0405]

/te-ci pytorch L0 L1

[denera]

LGTM!

[timmoon10]

Having _to_copy in this list puts us in a weird position. @youngeunkwon0405 Where did you get this list of ops and could we figure out a way to remove _to_copy?

The trouble is that we are implicitly using torch.Tensor.to as a dequantize function, so always expect the _to_copy op to output a plain PyTorch tensor. The reason for this design was to work with Mcore's logic for maintaining FP32 master weights (see logic for DDP and distopt). With this PR, we now see many spurious errors whenever we dequantize an FP8 tensor with to/float/half/etc.

If the current impl of _to_copy leads to insurmountable problems with FSDP2, we'll probably need to remove the implicit dequantization and change Mcore so that it explicitly calls Float8Tensor.dequantize.

[youngeunkwon0405]

I was following Torch AO's implementation.
https://github.com/pytorch/ao/blob/main/torchao/float8/fsdp_utils.py#L86-L99

I checked that it was okay for _to_copy to not preserve the tensor class currently. But leaved the warning for the future reference.
This PR did not change the functional behavior of the Float8Tensor _to_copy it only adds a warning here.