Is this page helpful?

core.fusions.fused_embedding_wgrad#

Embedding weight gradient accumulated straight into main_grad.

PyTorch’s embedding backward builds a dense [num_embeddings, hidden] weight gradient (zero-fill, then the rows of the tokens that occur), and Megatron’s DDP hook then adds the whole tensor into main_grad. For a 129280 x 7168 vocabulary that is about 1 ms per embedding use, nearly all of it on rows of tokens that do not occur. accumulate_embedding_wgrad adds only the rows of the tokens that occur, with the same arithmetic, so main_grad gets the same values:

  • PyTorch’s embedding backward (embedding_dense_backward_cuda, sort-based for more than 3072 tokens) sorts the token ids together with their positions, sums each token’s gradient rows in position order into fp32 partial sums of up to 10 rows (NROWS_PER_THREAD), adds the partial sums in order starting from 0, and rounds once to the gradient dtype.

  • main_grad.add_(grad) adds that in fp32 and rounds to main_grad’s dtype.

Rows of tokens that do not occur are left untouched instead of having 0 added to them.

Module Contents#

Functions#

_max_fn

_min_fn

_sort_tokens_kernel

Sort (token, position) pairs, then give every sorted slot the first and last slot of its token’s run and list the runs longer than NROWS (long_runs[0] is their count).

_partial_sums_kernel

Sum NROWS-row chunks of every run in fp32; a run of at most NROWS rows is finished here.

_long_run_sums_kernel

Add the partial sums of a run longer than NROWS rows in order, round, add to main_grad.

_position_bits

can_accumulate_embedding_wgrad

Whether accumulate_embedding_wgrad supports these tensors.

accumulate_embedding_wgrad

main_grad += embedding_dense_backward(grad, ids), touching only the rows of ids.

Data#

MIN_TOKENS

MAX_TOKENS

_NROWS

_BLOCK_H

API#

core.fusions.fused_embedding_wgrad.MIN_TOKENS#

3073

core.fusions.fused_embedding_wgrad.MAX_TOKENS#

8192

core.fusions.fused_embedding_wgrad._NROWS#

10

core.fusions.fused_embedding_wgrad._BLOCK_H#

1024

core.fusions.fused_embedding_wgrad._max_fn(a, b)#
core.fusions.fused_embedding_wgrad._min_fn(a, b)#
core.fusions.fused_embedding_wgrad._sort_tokens_kernel(
ids_ptr,
sorted_ids_ptr,
positions_ptr,
run_start_ptr,
run_end_ptr,
long_runs_ptr,
n,
POS_BITS: triton.language.constexpr,
NROWS: triton.language.constexpr,
BLOCK: triton.language.constexpr,
)#

Sort (token, position) pairs, then give every sorted slot the first and last slot of its token’s run and list the runs longer than NROWS (long_runs[0] is their count).

core.fusions.fused_embedding_wgrad._partial_sums_kernel(
sorted_ids_ptr,
positions_ptr,
run_start_ptr,
run_end_ptr,
grad_ptr,
main_grad_ptr,
partial_ptr,
hidden,
grad_stride,
main_grad_stride,
NROWS: triton.language.constexpr,
BLOCK_H: triton.language.constexpr,
)#

Sum NROWS-row chunks of every run in fp32; a run of at most NROWS rows is finished here.

core.fusions.fused_embedding_wgrad._long_run_sums_kernel(
sorted_ids_ptr,
run_end_ptr,
long_runs_ptr,
grad_ptr,
main_grad_ptr,
partial_ptr,
hidden,
main_grad_stride,
NROWS: triton.language.constexpr,
BLOCK_H: triton.language.constexpr,
)#

Add the partial sums of a run longer than NROWS rows in order, round, add to main_grad.

core.fusions.fused_embedding_wgrad._position_bits(num_tokens: int) → int#
core.fusions.fused_embedding_wgrad.can_accumulate_embedding_wgrad(
ids: torch.Tensor,
grad: torch.Tensor,
main_grad: torch.Tensor,
) → bool#

Whether accumulate_embedding_wgrad supports these tensors.

It requires Triton, between MIN_TOKENS and MAX_TOKENS tokens, bf16 or fp32 gradient and main_grad on the GPU, and a token id and position that pack into 31 bits.

core.fusions.fused_embedding_wgrad.accumulate_embedding_wgrad(
ids: torch.Tensor,
grad: torch.Tensor,
main_grad: torch.Tensor,
) → None#

main_grad += embedding_dense_backward(grad, ids), touching only the rows of ids.

Parameters:

  • ids – token ids of the embedding lookup, num_tokens elements.

  • grad – gradient of the lookup’s output, [*ids.shape, hidden].

  • main_grad – [num_embeddings, hidden] buffer the weight gradient is added to.

Requires can_accumulate_embedding_wgrad(ids, grad, main_grad).