From 07183be7a0da7de31fc27ce3f26d42de3f8e5823 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 31 Aug 2026 15:54:23 -0700 Subject: [PATCH 1/2] Add SiTU-GLU support to JAX Signed-off-by: Jeremy Berchtold --- tests/jax/test_custom_call_compute.py | 99 +++++++++++-------- tests/jax/test_distributed_layernorm_mlp.py | 18 +++- .../include/transformer_engine/activation.h | 3 +- transformer_engine/jax/activation.py | 25 +++-- .../jax/cpp_extensions/activation.py | 58 +++++++++-- transformer_engine/jax/csrc/extensions.h | 13 ++- .../jax/csrc/extensions/activation.cpp | 11 +++ .../jax/csrc/extensions/pybind.cpp | 1 + transformer_engine/jax/flax/module.py | 20 +++- transformer_engine/jax/flax/transformer.py | 5 +- transformer_engine/jax/layernorm_mlp.py | 4 + 11 files changed, 188 insertions(+), 69 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index d79c05f8ae..23580319e1 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -199,6 +199,7 @@ def assert_dequantized_grouped_scaled_tensor( ("squared_relu",), ("squared_relu", "linear"), ("clamped_silu", "clamped_linear"), + ("situ", "situ_linear"), ] ACTIVATION_TYPES = { @@ -210,6 +211,17 @@ def assert_dequantized_grouped_scaled_tensor( } +def make_activation_params(activation_type): + """Create non-default parameters for configurable activation tests.""" + if activation_type == ("clamped_silu", "clamped_linear"): + return tex.activation.ActivationParams.create( + activation_type, limit=0.75, alpha=1.702, glu_linear_offset=0.5 + ) + if activation_type == ("situ", "situ_linear"): + return tex.activation.ActivationParams.create(activation_type, beta1=2.0, beta2=8.0) + return None + + class TestActivation: def ref_act(self, x, activation_type, act_params): return _jax_act_lu(x, activation_type, act_params=act_params).data @@ -228,6 +240,26 @@ def primitive_func(self, inputs, activation_type, quantizer, act_params): ) return jnp.mean(out) + @pytest.mark.parametrize("betas", [(4.0, 25.0), (2.0, 8.0)]) + def test_jax_situglu_reference(self, betas): + beta1, beta2 = betas + x = jnp.arange(24, dtype=jnp.float32).reshape(4, 2, 3) / 3.0 - 4.0 + params = tex.activation.ActivationParams.create( + ("situ", "situ_linear"), beta1=beta1, beta2=beta2 + ) + + output = _jax_act_lu(x, ("situ", "situ_linear"), act_params=params).data + gate, up = x[:, 0, :], x[:, 1, :] + expected = ( + beta1 * jnp.tanh(gate / beta1) * jax.nn.sigmoid(gate) + ) * (beta2 * jnp.tanh(up / beta2)) + + assert_allclose(output, expected, dtype=x.dtype) + assert hash(params) == hash(params) + ffi_params = params.to_ffi_lowering_dict()["situglu"] + assert float(ffi_params["beta1"]) == pytest.approx(beta1) + assert float(ffi_params["beta2"]) == pytest.approx(beta2) + @pytest_parametrize_wrapper("shape", ALL_ACTIVATION_SHAPES) @pytest_parametrize_wrapper( "activation_type", @@ -244,16 +276,7 @@ def test_act_grad(self, shape, activation_type): value_n_grad_primitive_func = jit( value_and_grad(self.primitive_func, (0,)), static_argnums=(1, 3) ) - act_args = ( - {"limit": 0.75, "alpha": 1.702, "glu_linear_offset": 0.5} - if activation_type == ("clamped_silu", "clamped_linear") - else {} - ) - act_params = ( - tex.activation.ActivationParams.create(activation_type=activation_type, **act_args) - if activation_type == ("clamped_silu", "clamped_linear") - else None - ) + act_params = make_activation_params(activation_type) prim_out, (prim_grad,) = value_n_grad_primitive_func(x, activation_type, None, act_params) ref_out, (ref_grad,) = self.value_n_grad_ref_func(x, activation_type, act_params) assert_allclose(prim_out, ref_out, dtype=x.dtype) @@ -284,17 +307,7 @@ def test_act_grad_with_tensor_scaling_fp8( q_dtype=output_type, q_layout=QuantizeLayout.ROWWISE, ) - act_args = ( - {"limit": 0.75, "alpha": 1.702} - if activation_type == ("clamped_silu", "clamped_linear") - else {} - ) - - act_params = ( - tex.activation.ActivationParams.create(activation_type=activation_type, **act_args) - if activation_type == ("clamped_silu", "clamped_linear") - else None - ) + act_params = make_activation_params(activation_type) prim_out, (prim_grad,) = value_n_grad_primitive_func( x, activation_type, quantizer, act_params ) @@ -327,16 +340,7 @@ def test_act_forward_with_tensor_scaling_fp8( q_dtype=output_type, q_layout=q_layout, ) - act_args = ( - {"limit": 0.75, "alpha": 1.702} - if activation_type == ("clamped_silu", "clamped_linear") - else {} - ) - act_params = ( - tex.activation.ActivationParams.create(activation_type=activation_type, **act_args) - if activation_type == ("clamped_silu", "clamped_linear") - else None - ) + act_params = make_activation_params(activation_type) te_output = tex.act_lu(x, activation_type, te_quantizer, act_params) jax_output = _jax_act_lu(x, activation_type, jax_quantizer, act_params) assert_bitwise_scaled_tensors(te_output, jax_output) @@ -358,20 +362,18 @@ def test_act_forward_with_block_scaling_fp8( quantizer = QuantizerFactory.create( scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=output_type, q_layout=q_layout ) - act_args = ( - {"limit": 0.75, "alpha": 1.702} - if activation_type == ("clamped_silu", "clamped_linear") - else {} - ) - act_params = ( - tex.activation.ActivationParams.create(activation_type=activation_type, **act_args) - if activation_type == ("clamped_silu", "clamped_linear") - else None - ) + act_params = make_activation_params(activation_type) output = tex.act_lu(x, activation_type, quantizer, act_params) ref_out = self.ref_act(x, activation_type, act_params) assert_dequantized_scaled_tensor(output, ref_out) + @pytest.mark.parametrize("name", ["beta1", "beta2"]) + @pytest.mark.parametrize("value", [0.0, -1.0, jnp.inf, -jnp.inf, jnp.nan]) + def test_situglu_invalid_params(self, name, value): + kwargs = {name: value} + with pytest.raises(ValueError, match=rf"{name} must be finite and positive"): + tex.activation.ActivationParams.create(("situ", "situ_linear"), **kwargs) + NORM_OUTPUT_DTYPES = { "L0": [jnp.float8_e4m3fn], @@ -1631,7 +1633,9 @@ def ref_func(x, w, gamma, beta): @pytest.mark.skipif(not is_fp8_supported, reason=fp8_unsupported_reason) @pytest.mark.parametrize("m,n,k", [(64, 128, 128)]) - @pytest.mark.parametrize("activation_type", [("gelu",), ("gelu", "linear")]) + @pytest.mark.parametrize( + "activation_type", [("gelu",), ("gelu", "linear"), ("situ", "situ_linear")] + ) @pytest_parametrize_wrapper("recipe", supported_recipes) @pytest.mark.parametrize("norm_type", ["layernorm", "rmsnorm"]) @pytest_parametrize_wrapper("use_bias", [True, False]) @@ -1670,6 +1674,12 @@ def test_layernorm_mlp_grad( x=QuantizeMeta(), kernel=QuantizeMeta(), grad=QuantizeMeta() ), ) + activation_params = ( + {"beta1": 2.0, "beta2": 8.0} + if activation_type == ("situ", "situ_linear") + else None + ) + ref_activation_params = make_activation_params(activation_type) if norm_type == "layernorm": beta = jax.random.normal(subkeys[3], (k,)).astype(jnp.bfloat16) @@ -1688,6 +1698,7 @@ def prim_func(x, gamma, kernel_1, kernel_2, bias_1, bias_2): zero_centered_gamma=zero_centered_gamma, epsilon=eps, activation_type=activation_type, + activation_params=activation_params, quantizer_sets=quantizer_sets, ) ) @@ -1701,7 +1712,9 @@ def _ref_func_impl(x, gamma, kernel_1, kernel_2, bias_1, bias_2): bias_1_shape = (1,) * (linear_1_out.ndim - bias_1.ndim) + bias_1.shape linear_1_out += jnp.reshape(bias_1, bias_1_shape) - x = _jax_act_lu(linear_1_out, activation_type).data + x = _jax_act_lu( + linear_1_out, activation_type, act_params=ref_activation_params + ).data linear_2_out = jax.lax.dot_general(x, kernel_2, (((1,), (0,)), ((), ()))) if use_bias: bias_2_shape = (1,) * (linear_2_out.ndim - bias_2.ndim) + bias_2.shape diff --git a/tests/jax/test_distributed_layernorm_mlp.py b/tests/jax/test_distributed_layernorm_mlp.py index abf579d48e..854575cd62 100644 --- a/tests/jax/test_distributed_layernorm_mlp.py +++ b/tests/jax/test_distributed_layernorm_mlp.py @@ -337,6 +337,12 @@ def _test_layernorm_mlp( ): batch, seqlen, hidden_in = input_shape layernorm_type = "rmsnorm" + if activation_type == ("situ", "situ_linear"): + activation_params = {"beta1": 2.0, "beta2": 8.0} + elif activation_type == ("clamped_silu", "clamped_linear"): + activation_params = {"limit": 0.75, "alpha": 1.702, "glu_linear_offset": 0.5} + else: + activation_params = None rng = jax.random.PRNGKey(0) subkeys = jax.random.split(rng, 3) @@ -353,6 +359,7 @@ def _test_layernorm_mlp( layernorm_type=layernorm_type, intermediate_dim=INTERMEDIATE, activations=activation_type, + activation_params=activation_params, use_bias=use_bias, return_layernorm_output=True, ) @@ -372,6 +379,7 @@ def _test_layernorm_mlp( layernorm_type=layernorm_type, intermediate_dim=INTERMEDIATE, activations=activation_type, + activation_params=activation_params, scale_axes=LN_SCALE_AXES, ln_bias_axes=LN_BIAS_AXES, kernel_axes_1=KERNEL_1_AXES, @@ -431,7 +439,15 @@ def _test_layernorm_mlp( @pytest_parametrize_wrapper("input_shape", INPUT_SHAPE) @pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs()) - @pytest_parametrize_wrapper("activation_type", [("gelu",), ("silu", "linear")]) + @pytest_parametrize_wrapper( + "activation_type", + [ + ("gelu",), + ("silu", "linear"), + ("clamped_silu", "clamped_linear"), + ("situ", "situ_linear"), + ], + ) @pytest_parametrize_wrapper("dtype", DTYPES) @pytest_parametrize_wrapper("use_bias", [True, False]) @pytest_parametrize_wrapper("with_jax_gemm", [False, True]) diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 1be68a4daf..072cf1125c 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -40,7 +40,8 @@ enum class NVTE_Activation_Type { QGEGLU, SRELU, SREGLU, - CLAMPED_SWIGLU + CLAMPED_SWIGLU, + SITUGLU }; /*! \brief Computes the GeLU activation of the input. diff --git a/transformer_engine/jax/activation.py b/transformer_engine/jax/activation.py index b2b90a10c9..c4daeb6ba5 100644 --- a/transformer_engine/jax/activation.py +++ b/transformer_engine/jax/activation.py @@ -17,23 +17,30 @@ from .quantize.quantizer import Quantizer +ActivationParams = tex.activation.ActivationParams + +__all__ = ["activation", "ActivationParams"] + + def activation( x: jnp.ndarray, activation_type: Sequence[Union[str, Callable]], quantizer: Optional[Quantizer] = None, - act_params: Optional[tex.activation.ActivationParams] = None, + act_params: Optional[ActivationParams] = None, ) -> jnp.ndarray: """Apply activation functions to input tensor with optional quantization. This function applies a sequence of activation functions to the input tensor. It supports string-based activation types (e.g., 'relu', 'gelu', ('gelu', 'linear')). + SiTU-GLU is selected with ``('situ', 'situ_linear')``; ``situ_linear`` computes + the soft-capped up branch and is not an identity function. Args: x: Input tensor to apply activations to activation_type: Sequence of activation functions quantizer: Optional quantizer for quantizing the output - act_params: Optional activation parameters. Currently used - just for ClampedSwiGLU. + act_params: Optional parameters for configurable activations such as + ClampedSwiGLU and SiTU-GLU. Returns: Activated output tensor @@ -54,8 +61,8 @@ def _activation(x, activation_type, quantizer, act_params): x: Input tensor activation_type: Sequence of activation functions quantizer: Optional quantizer - act_params: Optional activation parameters. Currently used - just for ClampedSwiGLU. + act_params: Optional parameters for configurable activations such as + ClampedSwiGLU and SiTU-GLU. Returns: Activated tensor @@ -71,8 +78,8 @@ def _activation_fwd_rule(x, activation_type, quantizer, act_params): x: Input tensor activation_type: Sequence of activation functions quantizer: Optional quantizer - act_params: Optional activation parameters. Currently used - just for ClampedSwiGLU. + act_params: Optional parameters for configurable activations such as + ClampedSwiGLU and SiTU-GLU. Returns: Tuple of (output, context) for backward pass @@ -88,8 +95,8 @@ def _activation_bwd_rule(activation_type, act_params, ctx, g): Args: activation_type: Sequence of activation functions - act_params: Optional activation parameters. Currently used - just for ClampedSwiGLU. + act_params: Optional parameters for configurable activations such as + ClampedSwiGLU and SiTU-GLU. ctx: Context from forward pass g: Gradient from upstream diff --git a/transformer_engine/jax/cpp_extensions/activation.py b/transformer_engine/jax/cpp_extensions/activation.py index 5058192c3f..498428b23f 100644 --- a/transformer_engine/jax/cpp_extensions/activation.py +++ b/transformer_engine/jax/cpp_extensions/activation.py @@ -2,8 +2,9 @@ # # See LICENSE for license information. """JAX/TE custom ops for activation""" -from typing import Sequence, Union, Callable, Optional, Tuple +import math import operator +from typing import Sequence, Union, Callable, Optional, Tuple from functools import reduce, partial from dataclasses import dataclass @@ -54,6 +55,7 @@ ("squared_relu",): NVTE_Activation_Type.SRELU, ("squared_relu", "linear"): NVTE_Activation_Type.SREGLU, ("clamped_silu", "clamped_linear"): NVTE_Activation_Type.CLAMPED_SWIGLU, + ("situ", "situ_linear"): NVTE_Activation_Type.SITUGLU, } @@ -88,13 +90,39 @@ def to_ffi_lowering_dict(self): } +@dataclass(frozen=True) +class SiTUGLUParams: + """Soft-cap parameters for the SiTU-GLU activation function.""" + + beta1: float = 4.0 + beta2: float = 25.0 + + def __post_init__(self): + """Validate and canonicalize parameters for stable JIT caching.""" + for name in ("beta1", "beta2"): + value = float(getattr(self, name)) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{name} must be finite and positive, got {value}") + object.__setattr__(self, name, value) + + def __hash__(self): + """Return a stable hash for use as a JAX static argument.""" + return hash((self.beta1, self.beta2)) + + def to_ffi_lowering_dict(self): + """Convert parameters to the form consumed by the XLA FFI binding.""" + return { + "beta1": np.float32(self.beta1), + "beta2": np.float32(self.beta2), + } + + @dataclass(frozen=True) class ActivationParams: - """Parameters for various activation functions. - Currently only Clamped SwiGLU activation has parameters. - """ + """Parameters for activation functions with configurable behavior.""" clamped_swiglu: ClampedSwigluParams = ClampedSwigluParams() + situglu: SiTUGLUParams = SiTUGLUParams() @staticmethod def create(activation_type, **kwargs): @@ -104,13 +132,20 @@ def create(activation_type, **kwargs): "clamped_silu", "clamped_linear", } + SITU_ACTIVATION_TYPES = { + ("situ", "situ_linear"), + "situ", + "situ_linear", + } if activation_type in CLAMPED_ACTIVATION_TYPES: - return ActivationParams(ClampedSwigluParams(**kwargs)) + return ActivationParams(clamped_swiglu=ClampedSwigluParams(**kwargs)) + if activation_type in SITU_ACTIVATION_TYPES: + return ActivationParams(situglu=SiTUGLUParams(**kwargs)) return ActivationParams() # Default params for activations without parameters def __hash__(self): """Custom hash function to ensure dataclass is hashable for jax jit to work""" - return hash((self.clamped_swiglu,)) + return hash((self.clamped_swiglu, self.situglu)) def to_ffi_lowering_dict(self): """Convert the activation parameters to a dictionary format for FFI lowering. @@ -118,7 +153,10 @@ def to_ffi_lowering_dict(self): dict: A dictionary representation of the activation parameters consumable by XLA FFI bindings for activation functions. """ - return {"clamped_swiglu": self.clamped_swiglu.to_ffi_lowering_dict()} + return { + "clamped_swiglu": self.clamped_swiglu.to_ffi_lowering_dict(), + "situglu": self.situglu.to_ffi_lowering_dict(), + } def _convert_to_activation_function(fn_or_string, act_params: ActivationParams): @@ -129,6 +167,9 @@ def _convert_to_activation_function(fn_or_string, act_params: ActivationParams): limit = act_params.clamped_swiglu.limit offset = act_params.clamped_swiglu.glu_linear_offset return lambda x: jnp.clip(x, min=-limit, max=limit) + offset + if fn_or_string == "situ_linear": + beta2 = act_params.situglu.beta2 + return lambda x: beta2 * jnp.tanh(x / beta2) if fn_or_string == "quick_gelu": return lambda x: jax.nn.sigmoid(1.702 * x) * x if fn_or_string == "squared_relu": @@ -137,6 +178,9 @@ def _convert_to_activation_function(fn_or_string, act_params: ActivationParams): limit = act_params.clamped_swiglu.limit alpha = act_params.clamped_swiglu.alpha return lambda x: jax.nn.sigmoid(alpha * jnp.minimum(x, limit)) * jnp.minimum(x, limit) + if fn_or_string == "situ": + beta1 = act_params.situglu.beta1 + return lambda x: beta1 * jnp.tanh(x / beta1) * jax.nn.sigmoid(x) if isinstance(fn_or_string, str): return getattr(jax.nn, fn_or_string) if callable(fn_or_string): diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 580219baf2..518ba18b83 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -42,8 +42,14 @@ struct ClampedSwigluConfig { float glu_linear_offset; }; +struct SiTUGLUConfig { + float beta1; + float beta2; +}; + struct ActivationConfig { ClampedSwigluConfig clamped_swiglu; + SiTUGLUConfig situglu; }; struct GemmConfig { @@ -239,9 +245,14 @@ XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(transformer_engine::jax::ClampedSwigluConf ::xla::ffi::StructMember("alpha"), ::xla::ffi::StructMember("glu_linear_offset")); +XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(transformer_engine::jax::SiTUGLUConfig, + ::xla::ffi::StructMember("beta1"), + ::xla::ffi::StructMember("beta2")); + XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( transformer_engine::jax::ActivationConfig, - ::xla::ffi::StructMember("clamped_swiglu")); + ::xla::ffi::StructMember("clamped_swiglu"), + ::xla::ffi::StructMember("situglu")); XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( transformer_engine::jax::GemmConfig, diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index 6325a700d1..7ec0644377 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -24,6 +24,8 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal auto swiglu_limit = act_params.clamped_swiglu.limit; auto swiglu_alpha = act_params.clamped_swiglu.alpha; auto swiglu_glu_linear_offset = act_params.clamped_swiglu.glu_linear_offset; + auto situ_beta1 = act_params.situglu.beta1; + auto situ_beta2 = act_params.situglu.beta2; auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type()); auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type()); @@ -141,6 +143,9 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal nvte_clamped_swiglu_v2(input_tensor.data(), output_tensor.data(), swiglu_limit, swiglu_alpha, swiglu_glu_linear_offset, stream); break; + case NVTE_Activation_Type::SITUGLU: + nvte_situglu(input_tensor.data(), output_tensor.data(), situ_beta1, situ_beta2, stream); + break; default: NVTE_ERROR("Unsupported ActivationEnum"); break; @@ -273,6 +278,8 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, auto swiglu_limit = act_params.clamped_swiglu.limit; auto swiglu_alpha = act_params.clamped_swiglu.alpha; auto swiglu_glu_linear_offset = act_params.clamped_swiglu.glu_linear_offset; + auto situ_beta1 = act_params.situglu.beta1; + auto situ_beta2 = act_params.situglu.beta2; auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type()); auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type()); @@ -451,6 +458,10 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, nvte_clamped_dswiglu_v2(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), swiglu_limit, swiglu_alpha, swiglu_glu_linear_offset, stream); break; + case NVTE_Activation_Type::SITUGLU: + nvte_dsituglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), situ_beta1, + situ_beta2, stream); + break; default: NVTE_ERROR("Unsupported ActivationEnum"); break; diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 3927e2686e..1f8da6771b 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -225,6 +225,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("SRELU", NVTE_Activation_Type::SRELU) .value("SREGLU", NVTE_Activation_Type::SREGLU) .value("CLAMPED_SWIGLU", NVTE_Activation_Type::CLAMPED_SWIGLU) + .value("SITUGLU", NVTE_Activation_Type::SITUGLU) .export_values(); pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 17c9a242f0..22c3dd0e33 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -26,7 +26,7 @@ from ..layernorm import layernorm from ..layernorm_dense import layernorm_dense from ..layernorm_mlp import layernorm_mlp -from ..activation import activation +from ..activation import activation, ActivationParams from ..softmax import softmax, SoftmaxFusionType from ..sharding import with_sharding_constraint_by_logical_axes from ..attention import AttnSoftmaxType @@ -948,8 +948,9 @@ class LayerNormMLP(TransformerEngineBase): Each activation has its own transformation layer. activation_params: dict, default = None The parameters needed(if any) by the activation functions specified in :attr:`activations`. - At the moment only ('clamped_silu', 'clamped_linear') which is clamped_swiglu used in GPT OSS - need additional parameters. + ClampedSwiGLU and SiTU-GLU require additional parameters. SiTU-GLU is selected with + ``('situ', 'situ_linear')``; ``situ_linear`` is the soft-capped up branch rather than + an identity function. intermediate_dropout_rng_name: str, default = 'dropout' The key in given RNGs via flax.linen.Module.apply that for generating Dropout masks. intermediate_dropout_rate: float, default = 0.0 @@ -1087,6 +1088,7 @@ def __call__(self, inputs: Array, deterministic: bool = False) -> Array: ("quick_gelu", "linear"), ("squared_relu", "linear"), ("clamped_silu", "clamped_linear"), + ("situ", "situ_linear"), ] act_pool = [("gelu",), ("silu",), ("relu",), ("quick_gelu",), ("squared_relu",)] normalized_acts = [] @@ -1096,7 +1098,7 @@ def __call__(self, inputs: Array, deterministic: bool = False) -> Array: normalized_acts.append(act.lower()) normalized_acts = tuple( reversed(normalized_acts) - if (normalized_acts[0] == "linear" or normalized_acts[0] == "clamped_linear") + if normalized_acts[0] in ("linear", "clamped_linear", "situ_linear") else normalized_acts ) @@ -1293,7 +1295,15 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args): x = checkpoint_name(x, self.ffn1_ckpt_name) if is_act_implemented: - z = activation(x, normalized_acts) + z = activation( + x, + normalized_acts, + act_params=( + ActivationParams.create(normalized_acts, **self.activation_params) + if self.activation_params + else None + ), + ) else: activations = [] x = jnp.split(x, num_activations, axis=-2) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 4b497826cc..42b8464ecf 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -1926,8 +1926,9 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods The sequence of activation functions to apply after the first linear transformation. Each activation has its own transformation layer. mlp_activation_params: dict = None - This is only used when ``('clamped_silu', 'clamped_linear')`` is in :attr:`mlp_activations`. At the moment - ``ClampedSwiglu`` is the only activation that requires parameters. + Parameters for configurable MLP activations. SiTU-GLU uses + ``mlp_activations=('situ', 'situ_linear')`` with ``beta1`` and ``beta2`` values; + ``situ_linear`` denotes its soft-capped up branch. use_bias: bool, default = False Indicate whether to enable bias shifting for QKVO projections, FC1 and FC2. If set to ``False``, the layer will not learn additive biases. diff --git a/transformer_engine/jax/layernorm_mlp.py b/transformer_engine/jax/layernorm_mlp.py index 4c324c208e..2f930aa7f1 100644 --- a/transformer_engine/jax/layernorm_mlp.py +++ b/transformer_engine/jax/layernorm_mlp.py @@ -87,6 +87,10 @@ def layernorm_mlp( ffn1_ckpt_name: Name for checkpointing the first feed-forward network ffn2_ckpt_name: Name for checkpointing the second feed-forward network activation_type: Activation function(s) to apply after the first dense layer transformation + activation_params: Optional parameters for configurable activations. SiTU-GLU uses + ``{"beta1": 4.0, "beta2": 25.0}`` with activation type + ``("situ", "situ_linear")``. ``situ_linear`` is the soft-capped up branch, + not an identity function. collective_op_sets: Tuple of two collective gemm config sets for the two dense layer transformations quantizer_sets: Tuple of two quantizer sets for the two dense layer transformations From 0b8eb80413de58f63b75000eaa7ac88c40b44ae9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:22:56 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_custom_call_compute.py | 14 +++++--------- .../jax/csrc/extensions/activation.cpp | 4 ++-- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 23580319e1..31ab6cd922 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -250,9 +250,9 @@ def test_jax_situglu_reference(self, betas): output = _jax_act_lu(x, ("situ", "situ_linear"), act_params=params).data gate, up = x[:, 0, :], x[:, 1, :] - expected = ( - beta1 * jnp.tanh(gate / beta1) * jax.nn.sigmoid(gate) - ) * (beta2 * jnp.tanh(up / beta2)) + expected = (beta1 * jnp.tanh(gate / beta1) * jax.nn.sigmoid(gate)) * ( + beta2 * jnp.tanh(up / beta2) + ) assert_allclose(output, expected, dtype=x.dtype) assert hash(params) == hash(params) @@ -1675,9 +1675,7 @@ def test_layernorm_mlp_grad( ), ) activation_params = ( - {"beta1": 2.0, "beta2": 8.0} - if activation_type == ("situ", "situ_linear") - else None + {"beta1": 2.0, "beta2": 8.0} if activation_type == ("situ", "situ_linear") else None ) ref_activation_params = make_activation_params(activation_type) @@ -1712,9 +1710,7 @@ def _ref_func_impl(x, gamma, kernel_1, kernel_2, bias_1, bias_2): bias_1_shape = (1,) * (linear_1_out.ndim - bias_1.ndim) + bias_1.shape linear_1_out += jnp.reshape(bias_1, bias_1_shape) - x = _jax_act_lu( - linear_1_out, activation_type, act_params=ref_activation_params - ).data + x = _jax_act_lu(linear_1_out, activation_type, act_params=ref_activation_params).data linear_2_out = jax.lax.dot_general(x, kernel_2, (((1,), (0,)), ((), ()))) if use_bias: bias_2_shape = (1,) * (linear_2_out.ndim - bias_2.ndim) + bias_2.shape diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index 7ec0644377..ac0c101fcd 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -459,8 +459,8 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, swiglu_limit, swiglu_alpha, swiglu_glu_linear_offset, stream); break; case NVTE_Activation_Type::SITUGLU: - nvte_dsituglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), situ_beta1, - situ_beta2, stream); + nvte_dsituglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), + situ_beta1, situ_beta2, stream); break; default: NVTE_ERROR("Unsupported ActivationEnum");