[PyTorch] Remove special handling for FP8 params in FP8 recipe infrastructure

[timmoon10]

Description

#1142 exposed a very subtle bug that caused non-deterministic test failures in test_fusible_ops_with_userbuffers.py.

Bug description

test_fusible_ops_with_userbuffers.py runs multiple test cases at a time because launching a parallel job is expensive, so it constructs and destroys multiple TE models with FP8 parameters. Python IDs may be reused after an object is deallocated, so the Python ID for FP8 tensors is sometimes reused. However, Float8Tensor.post_optimizer_step_fwd_amax_reduction uses Python IDs to check whether to perform amax reductions and FP8 scale updates. I observed that this was causing FP8 scale updates at weird times, which corrupted UB buffers, which caused hangs.

🫠

In short, the problem is from this weird callback in Float8Tensor:

TransformerEngine/transformer_engine/pytorch/tensor/float8_tensor.py

Line 77 in 2643ba1

def post_optimizer_step_fwd_amax_reduction(param: Float8Tensor) -> None:

This hack was added in #575 so that we would properly update FP8 scales for FP8 params after the optimizer step. However, we've made improvements since then:

• Avoid amax roll for non-run modules #825: There is no harm if we perform FP8 scale updates more often than necessary.
• [PyTorch] Refactor FP8 workspaces in linear modules #820: FP8 weights keep a private copy of its FP8 scale-inv, so we can perform FP8 scale updates whenever we want.

Thus, there's no need to do an FP8 scale update for the weights immediately after the optimizer step. We just need to do it sometime before the next optimizer step and there should be no change in numerics. In fact, these FP8 scales are already participating in the forward pass amax reduction and scale update, so avoiding those operations reduces runtime overheads. Also, this just makes Float8Tensor more sane and less tightly coupled with the FP8 recipe infrastructure.

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

• Remove special handling for FP8 params from FP8 recipe infrastructure

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

[timmoon10]

/te-ci pytorch L1 L3

[ksivaman]

/te-ci pytorch