Skip to content

TBE backward: consume precomputed index-preproc tensors - #6235

Open
gchalump wants to merge 2 commits into
pytorch:mainfrom
gchalump:export-D113624507
Open

gchalump wants to merge 2 commits into
pytorch:mainfrom
gchalump:export-D113624507

Conversation

@gchalump

Copy link
Copy Markdown
Contributor

Summary:
Everything except the forward for moving grad-independent index preprocessing
off the TBE backward critical path. Two halves, folded into one diff:

  • Backend driver (consume): adds an optional preproc_tensors (Tensor[]?)
    argument to the TBE backward driver so it can accept precomputed
    index-preprocessing outputs (transpose_embedding_input + find_long_segments)
    and skip recomputing them inline. When the arg is None (the default / common
    path), the driver runs the inline preprocessing exactly as before -- no
    behavior change. Applied across the backward host/meta/cu templates, the PT2
    CPU/CUDA wrappers, and split_embeddings_utils.

  • PT2 autograd (route): threads the optional preproc_tensors bundle through the
    PT2 training autograd backward into the backend backward op. When the forward
    emits the 12 index-preproc grad slots (grad_outputs[1..12]), backward packs
    them into preproc_tensors and forwards them to the driver's consume path;
    otherwise preproc_tensors stays nullopt and the driver recomputes the preproc
    inline exactly as before.

Non-breaking on its own: no forward output is emitted yet, so grad_outputs stays
size 1 until the forward-emit diff (top of stack) is applied. The backward
simply gains the ability to accept and route the extra grads when they later
appear. Builds on the add-op base (D114645411, below in stack).

Differential Revision: D113624507

)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/3107


Add the standalone `tbe_bwd_indices_preproc` CUDA op that runs the two grad-independent index-preprocessing steps (transpose_embedding_input, find_long_segments) and returns the 12-tensor bundle. Folded into embedding_backward_split_grad_template.cu. Adds a unit test asserting the op's bundle yields an allclose weight update versus the inline path.

Differential Revision: D114645411
Summary:
Everything except the forward for moving grad-independent index preprocessing
off the TBE backward critical path. Two halves, folded into one diff:

- Backend driver (consume): adds an optional `preproc_tensors` (Tensor[]?)
  argument to the TBE backward driver so it can accept precomputed
  index-preprocessing outputs (transpose_embedding_input + find_long_segments)
  and skip recomputing them inline. When the arg is None (the default / common
  path), the driver runs the inline preprocessing exactly as before -- no
  behavior change. Applied across the backward host/meta/cu templates, the PT2
  CPU/CUDA wrappers, and split_embeddings_utils.

- PT2 autograd (route): threads the optional preproc_tensors bundle through the
  PT2 training autograd backward into the backend backward op. When the forward
  emits the 12 index-preproc grad slots (grad_outputs[1..12]), backward packs
  them into preproc_tensors and forwards them to the driver's consume path;
  otherwise preproc_tensors stays nullopt and the driver recomputes the preproc
  inline exactly as before.

Non-breaking on its own: no forward output is emitted yet, so grad_outputs stays
size 1 until the forward-emit diff (top of stack) is applied. The backward
simply gains the ability to accept and route the extra grads when they later
appear. Builds on the add-op base (D114645411, below in stack).

Differential Revision: D113624507
@meta-cla meta-cla Bot added the cla signed label Aug 27, 2026
@meta-codesync

meta-codesync Bot commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

@gchalump has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113624507.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant