[PyTorch] Userbuffers support in operation-based API

[timmoon10]

Description

This PR adds basic support in the linear operation for using Userbuffers to overlap tensor-parallel communication with GEMMs. This is implemented as fused operations:

model = te.ops.Sequential(
    te.ops.BasicLinear(...),
    te.ops.Bias(...),
    te.ops.ReduceScatter(...),
)  # Fused into UserbuffersForwardLinear

I've tried to avoid touching the core UB infrastructure in transformer_engine/pytorch/module/base.py, so I've kept the messy API and hackily worked around some bugs. This feature should be considered experimental and unstable.

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

• Add fused operation for linear forward with Userbuffers
• Add fused operation for linear backward with Userbuffers

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

[timmoon10]

/te-ci pytorch

[denera]

LGTM, at least on the Userbuffers code side.

[timmoon10]

/te-ci pytorch

[timmoon10]

The previous UB GEMM+RS impl has a correctness bug. UB splits up the GEMM into multiple chunks so that it can do the GEMM to compute one output chunk at the same time it is doing an RS on another output chunk. However, each output chunk requires applying a different chunk of the bias. We previously used the same bias pointer for all chunks, while this PR computes the correct offsets in the bias pointer.

[timmoon10]

• The data kwarg allows us to easily initialize Float8Tensors that use the UB workspace buffer
• The data_transpose kwarg is added for completeness

[timmoon10]

/te-ci pytorch

[timmoon10]

/te-ci L1

[timmoon10]

/te-ci pytorch jax paddle L1

[timmoon10]

/te-ci pytorch L1