[JAX] Expose sigmoid GLU activation type

[jberchtold-nvidia]

Expose the activation type added in this PR here into TE/JAX #2656

[CodersAcademy006]

@jberchtold-nvidia Hi @jberchtold-nvidia,

It looks like GLU support was added internally for JAX in #2656 via the
("sigmoid", "linear") mapping in ActivationEnum, but there isn't a
direct ("glu",) mapping exposed yet.

Would it make sense to expose GLU explicitly by adding:

("glu",): NVTE_Activation_Type.GLU

to transformer_engine/jax/cpp_extensions/activation.py, and adding
corresponding test coverage in the JAX activation tests?

Happy to open a PR if this approach aligns with the intended API.

[jberchtold-nvidia]

Hi @CodersAcademy006, thanks for noticing this! I hadn't realized the latest version of that PR that merged did add internal JAX support for this activation type.

Would it make sense to expose GLU explicitly by adding:
("glu",): NVTE_Activation_Type.GLU

On the TE/JAX side, we use the convention of a tuple to indicate a gated activation as that is what JAX-based frameworks like MaxText typically use, e.g. here: https://github.com/AI-Hypercomputer/maxtext/blob/d9d167ec14d1a5a2ee9f1c8f0904236d3cb107cb/src/maxtext/configs/models/llama3.1-70b.yml#L23

For this GLU activation in TE, as you noticed, this is exposed to JAX here

("sigmoid", "linear"): NVTE_Activation_Type.GLU,

So this activation type is already supported in JAX, but we lack unit tests.

All that should need to be done should be a one-line change to tests/jax/test_custom_call_compute.py
to extend ALL_ACTIVATION_TYPES to include ("sigmoid", "linear").

If you'd like to make the PR, lmk and feel free to tag me for review and for running CI.
If you'd like me to make the PR, I can do that as well.

Thanks for noticing this!