Fix memory overhead of linear layer when all gather from sequence parallel

[yuzhongw-nvidia]

Description

gather_along_first_dim (all gather from sequence parallel region) in Linear and LayerNormLinear introduces memory overhead:

1. Some of the tensors are not released immediately, e.g., grad_output, inputmat_total, and ln_out_total.
2. When executing _post_process_fp8_blockwise_gather, it will call Float8BlockwiseQTensorBase._transpose_columnwise_data, whose old implementation is as follows. It will make the old _columnwise_data become unreachable and unable to be released.

self._columnwise_data = tex.fp8_transpose(
    self._columnwise_data, self._fp8_dtype, out=None
)

These two bugs are harmful especially for long context training, because the tensor gathered from SP region will be very large and the overhead is unacceptable. For example, in the MoE layer of DSV3, there will be 6 tensors that were not released in time; and in the dense layer of DSV3, there will be 18 tensors that were not released in time.

Suppose seq_len=65536, there will be at least 18 * 65536 * 7168 = 7.875 GB memory allocation overhead, and this overhead will increase proportionally with the sequence length. Furthermore, this usually leads to much more overhead in device memory usage. Still take 64K context length as an example, the device memory usage decreases from 59187 MB to 40777 MB for DeepSeek-V3-TP16PP8EP32VPP4CP1-full-recompute with the fix. With this huge improvement, we could further turn into selective recomputation to greatly improve the E2E performance (267 TFLOPS -> 307 TFLOPS).

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]

Thanks! It does seem a problem if the tensor GC will happen after a long period of time, while enforcing GC directly brings too much overhead. CC @timmoon10 to also take a look at this PR.

[timmoon10]

Shouldn't Python refcounting deallocate the old data automatically? If not, then there's a larger bug in tex.fp8_transpose that we should fix. The code will become unmanageable if we have to apply this kind of trick everywhere we call a C++ extension.

I see that we return at::Tensor instead of Pybind11 objects. I wonder if it is not properly handling the refcount.

[yuzhongw-nvidia]

It is a very good question that also confuses me. I have a similar question: why we need to use clear_tensor_data to release tensors that unused, but not let Python GC to deallocate them? I guess the root cause may be similar and my current fix is a similar trick with clear_tensor_data.

I'm not very familiar with the code and deeper implementation. Could you please share this problem with some TE experts to help solve it? Or do you think we could merge this PR first and try to find out the root cause later, because it is a little bit emergent for the runnability and perf of DSV3 / MLA long context training.

[timmoon10]

We call clear_tensor_data to manually deallocate memory before Python GC. A good example is here in the LayerNormLinear backward:

TransformerEngine/transformer_engine/pytorch/module/layernorm_linear.py

Line 896 in 5f2b831

clear_tensor_data(ln_out_total)

The forward GEMM input tensor is stored within the autograd ctx, so GC will not deallocate until after the backward has finished. However, we don't need this buffer after the wgrad GEMM and ideally it would be reused for the LayerNorm grad.

[yuzhongw-nvidia]

As I described in this thread, I think it is not the business of C++ extension, but the business of Float8BlockQuantizer.make_empty (in this case, the unreleased tensor is also created by make_empty, and the C++ extension is only creating the new tensor). In other words, I believe the issue is: tensors created by Float8BlockQuantizer.make_empty can only be released by manually deleting _columnwise_data.data and _rowwise_data.data. (I believe that's why we need to delete the _columnwise_data.data and _rowwise_data.data in clear_tensor_data, but not to delete _columnwise_data and _rowwise_data.)

[timmoon10]

Will this have an effect? The only benefit from deallocating the input tensor is that the buffer might be reused for the tensor-parallel reduce-scatter on the output. However, if we've done a tensor-parallel all-gather on the input, we won't do a reduce-scatter.

[yuzhongw-nvidia]

Yes, it will have an effect. I have just confirmed this by an experiment, and if I remove this line, the inputmat_total will be in the memory for a very long time.

I think that's the most serious problem: the data of FP8 tensors created by make_empty may not be released by Python GC as expected. Do you have any idea about it?

[timmoon10]

/te-ci pytorch L1

[timmoon10]

/te-ci pytorch L1

[timmoon10]

/te-ci pytorch L1