Release v2.2 bugfixes megatron ifu

[sudhu2k]

Description

This PR fixes some bugs in the release v2.2 branch that I found when working on the Megatron IFU r0.15.0.

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 refactoring

Changes

1. Fixes assertion of Float8Tensor to include Float8Currentscaling tensor.
2. cherry picks a commit which checks if the function is of type module, before accessing fsdp_wrapped, fsdp_group. (It was trying to access them even when it isn't of type TEModule.

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

[ipanfilo]

I think it is better to create separate branch rather than using release_v2.2_rocm.
Is there corresponding PR for dev branch?
Also, please update copyright on modified files

[sudhu2k]

Dev branch doesn't need these fixes, since they already have them

TransformerEngine/transformer_engine/pytorch/module/base.py

Line 1393 in 51f74fa

if isinstance(tensor, QuantizedTensor):

TransformerEngine/transformer_engine/pytorch/distributed.py

Lines 706 to 710 in 51f74fa

	if isinstance(function, TransformerEngineBaseModule):
	# If this TE module is FSDP-wrapped, clear its FSDP group information because there's no need
	# to scatter/gather activations that we will recompute anyway.
	setattr(function, "fsdp_wrapped", False)
	setattr(function, "fsdp_group", None)

This is only for release_v2.2 branch since Megatron's CI uses this specific commit.
Let me know if you still think we should have a separate branch. I've updated copyright on the files which I made changes to, not the cherry-picks.

[ipanfilo]

It is not a question of branch per se but rather versioning. TE 2.2 was released a while ago. So changes on top of that should be versioned like 2.2.post0 or something like this.

[sudhu2k]

Yeah that makes sense! I'll bump the version to 2.2.post0.