Added MCore FSDP support for TE

[sanandaraj5597]

Description

Added support for gradient accumulation fusion when using MCore FSDP with TE.

[shjwudp]

Thank you! This is very beneficial for improving FSDP's performance. I think it would be best to make the memory allocation for main_grad an explicit function. What do you think?

[sanandaraj5597]

@shjwudp I think I updated the code based on the recent code factor. Please check once.

[timmoon10]

Is weight.main_grad available during the backward pass or do you need to go through weight.get_main_grad()? The op fuser API tries to access main_grad in the backward pass:

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

Lines 980 to 989 in 3a298e6

	if ctx.weight_requires_grad and accumulate_into_main_grad:
	if not hasattr(self.weight, "main_grad"):
	raise RuntimeError(
	"BasicLinear op is configured with "
	"accumulate_into_main_grad=True, "
	"but weight parameter does not have main_grad attribute"
	)
	grad_weight = self.weight.main_grad.detach()
	else:
	accumulate_into_main_grad = False

[sanandaraj5597]

@timmoon10 that's a good point. When we use TE with MCore FSDP, the main grads buffers are lazily allocated during backward. So, if we are using the fuser with FSDP then we need to allocate the buffer and then use it.