FSDP grad fusion support

[sanandaraj5597]

This PR adds support gradient fusion for MCore FSDP.

[timmoon10]

I don't think this makes sense. If you configure a TE module with fuse_wgrad_accumulation=True (e.g. here), the correct behavior is to fuse wgrad accumulation. If Mcore FSDP doesn't support it, then it should be Mcore's responsibility to not set that arg.

[timmoon10]

The root problem is that Mcore DDP and FSDP have different behaviors and require different contracts with TE:

• DDP uses persistent main_grad buffers and it expects TE to accumulate into it. To adhere to this contract, Mcore zeros out the main_grad before the first microbatch step.
• FSDP uses temporary main_grad buffers and it expects TE to overwrite it.

I don't like this PR's approach of switching between these two cases based on whether Mcore is using DDP or FSDP, since that's not actually the important thing. It also needlessly blocks some possible optimizations (DDP might want to overwrite main_grads in the first microbatch, FSDP might want to accumulate into main_grads if a weight is shared).

There are a few possible redesigns:

1. Deprecate the fuse_wgrad_accumulation kwarg in favor of something like output_wgrad_to_main_grad. Then check param flags to decide whether to overwrite or accumulate into the main_grad:

    grad_weight: torch.Tensor
    accumulate: bool = False
    if output_wgrad_to_main_grad:
        if getattr(weight, "get_main_grad", None) is not None:
            grad_weight = weight.get_main_grad()
        else:
            grad_weight = weight.main_grad
        accumulate = getattr(weight, "_overwrite_main_grad", True)
    else:
        grad_weight = torch.empty(...)

    gemm(..., out=grad_weight, accumulate=accumulate)

Ensuring backward compatibility will be tricky.

2. Have separate kwargs for fuse_wgrad_accumulation and overwrite_wgrad_main_grad. This means that the two cases are separate code paths and backward compatibility is easier to maintain. However, it also means we can't change behavior between steps.
3. Keep the fuse_wgrad_accumulation kwarg and purely control behavior with param flags. This is basically the approach used in this PR, although it could be improved by using better names rather than just checking weight.__fsdp_param__. One problem is that fuse_wgrad_accumulation will no longer be an accurate name.

[timmoon10]

We should include this behavior in the documentation:

TransformerEngine/transformer_engine/pytorch/module/linear.py

Line 1059 in dd9433e

fuse_wgrad_accumulation : bool, default = 'False'

TransformerEngine/transformer_engine/pytorch/module/layernorm_linear.py

Line 1123 in dd9433e

fuse_wgrad_accumulation : bool, default = 'False'

TransformerEngine/transformer_engine/pytorch/module/layernorm_mlp.py

Line 1482 in dd9433e

fuse_wgrad_accumulation : bool, default = 'False'

TransformerEngine/transformer_engine/pytorch/module/grouped_linear.py

Line 517 in dd9433e

fuse_wgrad_accumulation : bool, default = 'False'

TransformerEngine/transformer_engine/pytorch/ops/linear.py

Line 52 in dd9433e

accumulate_into_main_grad: bool, default = `False`

TransformerEngine/transformer_engine/pytorch/ops/basic/basic_linear.py

Line 77 in dd9433e

accumulate_into_main_grad: bool, default = `False`

[timmoon10]

/te-ci pytorch

[timmoon10]

LGTM, pending CI.

It seems that DCO didn't like some commits, but they look fine to me. Maybe there's something misconfigured with your GitHub account's emails or maybe DCO is just buggy? In any case, I'm happy leaving this PR as-is and overriding DCO.