Is the delayed scaling computation overlapped with another computation?

[avizon-aws]

In order for the delayed scaling recipe to be efficient, the scales for the next iteration are computed before hand. I think the part for the computation of the scales is here:

TransformerEngine/transformer_engine/pytorch/fp8.py

Line 314 in f8eb799

def reduce_tensor_across_group_op_max(tensor: torch.Tensor, group: dist_group_type) -> None:

But I dont understand if this computation is done in parallel with another computation to save time and for efficiency. Could anyone throw some light on this?Also where is the amax calculation for the input and weights done?

[Advait251206]

Yes, that's generally the idea behind the delayed scaling recipe. Instead of computing the scaling factor needed for the current FP8 operation, Transformer Engine records the amax values during the current iteration and uses them to determine the scaling factors for the next iteration. This avoids introducing additional synchronization points into the critical execution path.

Regarding reduce_tensor_across_group_op_max(), this function is responsible for reducing (taking the maximum of) the recorded amax values across the relevant distributed process group when FP8 synchronization is enabled. The actual overlap depends on the execution context:

• Since PyTorch CUDA operations are launched asynchronously, these reduction operations can often be queued on CUDA streams and overlap with other GPU work, provided there are no immediate synchronization dependencies.
• The implementation itself does not explicitly guarantee overlap; whether communication and computation execute concurrently depends on factors such as CUDA stream usage, NCCL scheduling, and whether later operations need to wait for the reduction to complete.
• The delayed scaling design helps hide much of this overhead because the updated scales are only required in the following iteration, giving the runtime an opportunity to overlap communication with ongoing computation.

As for the amax calculation itself, it is not performed in fp8.py. Instead, the amax values are generated inside the FP8 kernels (GEMMs and other FP8-enabled operations) implemented by Transformer Engine. During execution, these kernels track the maximum absolute value of the tensors they process (inputs, weights, gradients, outputs, etc.) and write those values into the FP8 metadata (FP8TensorMeta). The Python layer in fp8.py is then responsible for managing the amax history, performing any required distributed reductions, and updating the scaling factors according to the selected FP8 recipe (such as delayed scaling).

So, in summary:

• amax collection: Performed inside the FP8 CUDA kernels.
• amax reduction: Managed by functions such as reduce_tensor_across_group_op_max() when synchronization across ranks is required.
• Scale update: Computed from the recorded amax history and applied to the next iteration as part of the delayed scaling recipe.

This separation is what allows delayed scaling to minimize synchronization overhead while maintaining stable FP8 scaling across iterations.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206