[PyTorch][FP8 CS] Remove the unnecessary torch reciprocal op in fp8 current scaling code path

[zhongbozhu]

Description

at::reciprocal was part of FP8 current scaling recipe to make sure TPGEMM overlap works with FP8 current scaling. That's because the scales are not copied into user buffer for communication, but the old TPGEMM overlap codebase relies on the create_tensor API to construct the output tensor which requires a scaling factor tensor in it.

Now that dependency has been removed, we no longer need this op.

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

[zhongbozhu]

/te-ci pytorch L1

[yaox12]

/te-ci pytorch L1

[BestJuly]

LGTM.

[timmoon10]

Since we're no longer respecting the rowwise_data arg, we should error out if it is ever used:

NVTE_CHECK(!rowwise_data.has_value(), "Float8CurrentScalingQuantizer does not support creating tensor with pre-initialized buffer");

#1952 is a systematic attempt to remove rowwise_data, but it's quite large. This PR will be much easier to merge.

[zhongbozhu]

@timmoon10 Seems like this change will break test_fusible_ops.py, do you know about the reason?

[timmoon10]

It seems like numerical errors in the tex functions for activation forward. I wonder if the scale_inv isn't being properly updated, so it was relying on the reciprocal in create_tensor. I expect LayerNormMLP also experiences these numerical problems, but its test tolerances are much looser.

The root cause is because the activation backward pass doesn't support FP8 current scaling:

TransformerEngine/transformer_engine/pytorch/csrc/extensions/activation.cpp

Line 76 in e7251f9

py::object dactivation_helper(const at::Tensor& grad, const at::Tensor& input,

[zhongbozhu]

It seems like numerical errors in the tex functions for activation forward. I wonder if the scale_inv isn't being properly updated, so it was relying on the reciprocal in create_tensor. I expect LayerNormMLP also experiences these numerical problems, but its test tolerances are much looser.

The root cause is because the activation backward pass doesn't support FP8 current scaling:

TransformerEngine/transformer_engine/pytorch/csrc/extensions/activation.cpp

Line 76 in e7251f9

py::object dactivation_helper(const at::Tensor& grad, const at::Tensor& input,

Thanks for the root causing! I don't think layernorm_mlp will have this issue because the place where we call the dact function, the quantizer we gave is None here => so it's actually a None quantizer, calculating dact for the high precision, and then we call this line to do the quantization again here. So yes, the fusion is not supported, but layernorm_mlp has already taken that into account, in an implicit way of course.

[zhongbozhu]

@timmoon10 has the numerical issue in fused_op side been resolved? If it's already done, I can rebase and run the CI again.

[timmoon10]

#1952 is a comprehensive fix that removes the at::reciprocal and properly handles the activation backward with FP8 current scaling.

[zhongbozhu]

Update:

Will close this PR once #1952 is merged, which is a superset of this fix.