Implement fused kernel for FP8 scale update

[timmoon10]

Running GPT 175B, I've found that the forward pass is often bottlenecked by GPU kernel launch overheads, especially in the forward pass. Profiling the Python code finds that ~20% of the Transformer layer forward pass is spent in amax_and_scale_update (compare to 9% launching GEMM kernels). This function is called in every forward pass of Linear, LayerNormLinear, and LayerNormMLP, and each call involves ~10 small GPU operations. nvfuser and torch.compile do fuse some of the operations, leading to some improvement in GPU runtime, but the extra CPU overhead results in somewhat worse performance in the CPU-bound case.

This is an experiment with using a hand-written kernel to reduce these overheads. Alternative approaches:

• Only update scaling factors once per training step instead of once per microbatch
• Batch the scale updates for all layers together

[timmoon10]

/te-ci pytorch

[yaox12]

From my profiling, _default_get_amax is not fused well either. Can we follow paddle's

TransformerEngine/transformer_engine/paddle/fp8.py

Line 238 in 753eed3

tex.amax_and_scale_update_inplace(_amax_history=fp8_meta[fp8_meta_tensor_key].amax_history,

to fused the whole amax_and_scale_update?

[ptrendx]

Right, instead of writing another kernel in the framework-aware portion, let's actually add the Paddle kernel to the common part and use in both places.

[timmoon10]

One thing that'll make this tricky is that the Paddle kernel is in-place. The PyTorch kernel includes the amax history roll, so making that in-place would make the kernel much more complicated than it's worth. Perhaps we should change the Paddle implementation to be out-of-place?

[timmoon10]

I've moved the fused kernel to the core C++ library and modified the PyTorch and Paddle extensions so they both use it. It turns out that making the kernel work in-place was not difficult.

This includes the Paddle bugfix from #633.

[timmoon10]

/te-ci

[timmoon10]

/te-ci

[timmoon10]

/te-ci

[timmoon10]

The PyTorch and JAX tests as passing in pipeline 12465152 and the Paddle tests are failing due to a problem with cuDNN v9 in the upstream Paddle container. The Paddle tests pass when I use an older Paddle container in pipeline 12474691. This is ready to merge.

[ksivaman]

/te-ci

[ksivaman]

Have we tested e2e numerics for this change against the previous versions? @timmoon10

[timmoon10]

The tests pass, but I haven't tried full-scale training runs.

[timmoon10]

@ksivaman The numerics are tested thoroughly in https://github.com/NVIDIA/TransformerEngine/blob/7fc00c0819b3249d1abbb935719a81202d097e9d/tests/pytorch/test_recipe.py, so I think full-scale convergence runs are excessively careful.

[timmoon10]

/te-ci

[timmoon10]

After internal discussion with @ptrendx and @ksivaman, there is no objection to merging without review.