[PyTorch] Fix FP8 activation recompute

[ksivaman]

Description

The amax reduction of all backward tensors happens in the first module (one of the base modules) in a given fp8_autocast. The ctx.reduce_and_update_bwd_fp8_tensors flag is saved by querying the FP8GlobalStateManager.is_first_fp8_module() which only returns True for the first module in the fp8_autocast. However, this introduces a bug during activation recompute since the recompute phase runs outside the fp8 context, and the first module flags are never set. This results in the amaxes for gradients not getting reduced.

Fixes #1190

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

The activation_recompute_forward maintains a queue structure to pass values of the IS_FIRST_FP8_MODULE flag from the forward phase to the recompute phase. During the recompute phase, it is reset back to not disturb any nested autocasts.

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

[ksivaman]

/te-ci pytorch

[denera]

LGTM!

[ptrendx]

@ksivaman Could you put some more information about the bug and the fix in the description?

[ksivaman]

/te-ci pytorch