Added keep_fp8_weight_transpose_cache checks while updating transpose in fwd pass

[sudhu2k]

Description

While keep_fp8_weight_transpose_cache is False, the expectation is to not create any transpose in the forward pass and compute transpose in the backward pass. ad76b62#diff-ba97b0d1ae75d17a678bc38b4fa69ffec1e0ea007657a28d65565ee2cff35b95
The above commit introduced check to ensure transpose is created when input requires grad is True. But we don't want to create transpose when keep_fp8_weight_transpose_cache is False.
Without this check, it leads to creating transpose while transpose data ptr isn't initialized.

RuntimeError: /workspace/TransformerEngine/transformer_engine/common/transpose/transpose.hip:206 in function transpose: Assertion failed: output.data.dptr != nullptr. Output is not allocated.

Fixes # (13552)
https://ontrack-internal.amd.com/browse/SWDEV-553639

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

Added keep_fp8_weight_transpose_cache checks while updating transpose

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

[sudhu2k]

Hi @ipanfilo, @wangye805, @wenchenvincent
I've added the unit test for the fix. Also modified it to test layernorm linear and layernorm mlp along with linear.
The issue happened, because the backward call clears transpose if keep_fp8_weight_transpose_cache is set to false and in the next iteration of the forward pass, the code tried to create transpose on a cleared memory of transpose.
https://github.com/ROCm/TransformerEngine/blob/66340ddbd8917ddf5f471de61301fc979027179b/transformer_engine/pytorch/module/linear.py#L563C40-L564

[wangye805]

Please also add a section of documentation for keep_fp8_weight_transpose_cache, either in the source codes before it's definition or in our README

[Micky774]

Several meaningful test failures on CI, all stemming from

        if keep_fp8_weight_transpose_cache:
>           assert not transpose_is_empty_or_none, "Expected _transpose to be a valid, non-empty tensor when transpose cache is enabled."
E           AssertionError: Expected _transpose to be a valid, non-empty tensor when transpose cache is enabled.

[ipanfilo]

Keep it

[ipanfilo]

Several meaningful test failures on CI, all stemming from

        if keep_fp8_weight_transpose_cache:
>           assert not transpose_is_empty_or_none, "Expected _transpose to be a valid, non-empty tensor when transpose cache is enabled."
E           AssertionError: Expected _transpose to be a valid, non-empty tensor when transpose cache is enabled.

I might be OK to assert that _transpose is not valid when keep flag is False but there might be other reasons for _transpose not to be valid when the flag is True (default behaviour)

[wangye805]

nit: remove this to reduce unnecessary changes