[PyTorch] Debug linear layer when saving original input and using debug quantizer

[timmoon10]

Description

#1865 introduced a failure in the distributed debug tests. The root cause is because we only all-gather the row-wise data for DebugQuantizedTensor:

TransformerEngine/transformer_engine/pytorch/distributed.py

Line 1394 in f8933bb

rowwise_total = gather_along_first_dim(rowwise, process_group, False, final_quantizer)[0]

However, if the linear layer is caching its original input tensor and requantizing in the backward pass, the correct behavior is to only quantize the column-wise data. This PR is a hacky workaround that only applies the debug quantizer to the gathered input tensor.

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

• Modify linear backward to avoid all-gathering debug tensor with only column-wise data

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

[hxbai]

Hi Tim, why do we need the rowwise data here?

[timmoon10]

In principle we shouldn't need it, but I was running into issues where FP8 casts were failing without it. Actually, I don't think we need it once we skip the debug quantizer case.

[hxbai]

I think this may be due to that the condition is modified. For (Float8Quantizer, Float8CurrentScalingQuantizer) cases, it doesn't support only quantizing the colwise data. But if backward_input_needs_gather is False, (Float8Quantizer, Float8CurrentScalingQuantizer) tensors would also go into the else path, which will cause the error. This hurts the performance of blockwise FP8 and MXFP8. Do you need me to create a fix PR for this, or will you fix this together with the DebugTensor?

[pggPL]

I think this line is source of error

TransformerEngine/transformer_engine/pytorch/distributed.py

Line 1408 in f8933bb

out_obj.rowwise_gemm_tensor = out_obj.rowwise_gemm_tensor

The rowwise/columnwise in DebugQuantizedTensors are the tensors used in gemms -> they can be both the same Float8Tensor object for example. And update_usage() does nothing in debug tensors currently.

[timmoon10]

@pggPL I tried changing

TransformerEngine/transformer_engine/pytorch/distributed.py

Line 1408 in f8933bb

out_obj.rowwise_gemm_tensor = out_obj.rowwise_gemm_tensor

to

out_obj.columnwise_gemm_tensor = out_obj.rowwise_gemm_tensor 

However, the error reappeared when I applied the debug quantizer to the local input tensor.

For now, I think we should merge this as a quick bugfix and we can fix the edge cases for the debug tensor later.

[ksivaman]

LGTM

[timmoon10]

/te-ci pytorch L1