Add TE/JAX high-level modules, unittests and examples

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))

examples/jax/encoder/test_single_gpu_bf16_training.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with BF16 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ optimizer = optax.sgd(0.001, 0.9)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

examples/jax/encoder/test_single_gpu_fp8_training.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+""" Encoder with FP8 Training on single GPU"""
+import jax
+import jax.numpy as jnp
+import optax
+from cuda import cudart
+from flax.core.frozen_dict import FrozenDict
+from flax.training import train_state
+
+import transformer_engine.jax as te
+from transformer_engine.jax.fp8 import FP8Helper
+from transformer_engine.common.recipe import Format as FP8Format
+from transformer_engine.common.recipe import DelayedScaling
+
+PARAMS_KEY = 'params'
+
+BATCH = 32
+SEQLEN = 512
+HIDDEN = 1024
+
+
+def gpu_has_fp8():
+ """GPU arch has to support FP8"""
+ cudaSuccess = cudart.cudaError_t.cudaSuccess
+ ret, gpu_id = cudart.cudaGetDevice()
+ assert ret == cudaSuccess
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor
+ _, major = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ flag = cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor
+ _, minor = cudart.cudaDeviceGetAttribute(flag, gpu_id)
+ sm_arch = major * 10 + minor
+ return sm_arch >= 89
+
+
+def network():
+ """NLP Encoder"""
+ encoder = te.TransformerLayer(hidden_size=HIDDEN,
+ mlp_hidden_size=4 * HIDDEN,
+ hidden_dropout=0.0,
+ attention_dropout=0.0,
+ layernorm_type='rmsnorm',
+ mlp_activations=('gelu', 'linear'),
+ layer_type=te.TransformerLayerType.ENCODER,
+ transpose_batch_sequence=True,
+ dtype=jnp.bfloat16)
+ return encoder
+
+
+def synthesis_data(data_rng):
+ """Dataset generator"""
+ return jax.random.normal(data_rng, [SEQLEN, BATCH, HIDDEN], jnp.bfloat16)
+
+
+def train_step(batch, state, others):
+ """Training function."""
+
+ def loss_fn(collections):
+ logits = state.apply_fn(collections, batch)
+ loss = jnp.mean(logits)
+ return loss
+
+ grad_fn = jax.value_and_grad(loss_fn)
+ loss, grads = grad_fn(FrozenDict({PARAMS_KEY: state.params, **others}))
+ grads, params_grads = grads.pop(PARAMS_KEY)
+ state = state.apply_gradients(grads=params_grads)
+ others = FP8Helper.update_fp8_metas(grads)
+ return loss, state, others
+
+
+def test_encoder():
+ """Encoder example"""
+ if gpu_has_fp8() is False:
+ print("GPU doesn't support FP8")
+ return
+
+ rng = jax.random.PRNGKey(0)
+ rng, init_rng, data_rng = jax.random.split(rng, 3)
+ inputs = synthesis_data(data_rng)
+ optimizer = optax.sgd(0.001, 0.9)
+
+ with te.fp8_autocast(enabled=True, fp8_recipe=DelayedScaling(fp8_format=FP8Format.HYBRID)):
+ encoder = network()
+ variables = jax.jit(encoder.init)(init_rng, inputs)
+ variables, params = variables.pop(PARAMS_KEY)
+ state = train_state.TrainState.create(apply_fn=encoder.apply, params=params, tx=optimizer)
+ jitted_train_step = jax.jit(train_step)
+ assert "fp8" in str(jax.make_jaxpr(jitted_train_step)(inputs, state, variables))
+
+ for i in range(5):
+ rng, data_rng = jax.random.split(rng)
+ inputs = synthesis_data(data_rng)
+ loss, state, variables = jitted_train_step(inputs, state, variables)
+ print(f"Step {i} - Loss: {loss}")
+
+
+if __name__ == "__main__":
+ test_encoder()

qa/L0_jax_unittest/test.sh

@@ -6,3 +6,4 @@ set -xe
 
 : ${TE_PATH:=/opt/transformerengine}
 pytest -Wignore -v $TE_PATH/tests/jax
+pytest -Wignore -v $TE_PATH/examples/jax

tests/jax/test_custom_call_compute.py

@@ -60,15 +60,15 @@ def test_compile_bf16(self):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  # x = input, matrix 2d
  # y = input, matrix 2d (weight)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -84,13 +84,13 @@ def test_compile_fp8(self, compute_type):
 
  def func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  value_n_grad_func = value_and_grad(func, (0, 1))
  value_n_grad_func_compiled = jit(value_n_grad_func).lower(a, b).compile()
@@ -104,13 +104,13 @@ def test_forward_bf16(self, m, n, k):
  b = jax.random.normal(subkeys[1], (k, n), jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None))
+ primitive_out = fp8_dot(fp8_gemm_pkg, *_format2dtypes(None))
  ref_out = jnp.dot(a, b)
 
  assert_allclose(primitive_out, ref_out)
@@ -128,20 +128,20 @@ def test_forward_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  # calculate amax
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  # calculate scale by amax
  fp8_meta = FP8Helper._update_fp8_metas_impl(fp8_meta)
 
  fp8_gemm_pkg = FP8GemmPackage(1, a, [b], *fp8_meta)
- primitive_out = fp8_dot(fp8_gemm_pkg, 0, *compute_type)
+ primitive_out = fp8_dot(fp8_gemm_pkg, *compute_type)
  ref_out = jnp.dot(a, b)
 
  ref_out = ref_out.astype(jnp.float32)
@@ -158,13 +158,13 @@ def test_grad_bf16(self, m, n, k):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.mean(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None)))
+ return jnp.mean(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None)))
 
  def ref_func(x, y):
  return jnp.mean(jnp.dot(x, y))
@@ -193,15 +193,15 @@ def test_grad_fp8_randint(self, m, n, k, compute_type):
  b = jax.random.randint(subkeys[1], (k, n), min_val, max_val).astype(jnp.bfloat16)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_meta = [fp8_max, fp8_metas_amax, fp8_metas_scale, fp8_metas_scale_inv]
 
  def primitive_func(x, y, metas):
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], *metas)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *compute_type))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *compute_type))
 
  def ref_func(x, y):
  return jnp.sum(jnp.dot(x, y))
@@ -232,13 +232,13 @@ def test_contracting_dims_bf16(self):
 
  def primitive_func(x, y):
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM, 1), jnp.float32)
  fp8_gemm_pkg = FP8GemmPackage(1, x, [y], fp8_max, fp8_metas_amax, fp8_metas_scale,
  fp8_metas_scale_inv)
- return jnp.sum(fp8_dot(fp8_gemm_pkg, 0, *_format2dtypes(None), ((2, 3), (0, 1))))
+ return jnp.sum(fp8_dot(fp8_gemm_pkg, *_format2dtypes(None), ((2, 3), (0, 1))))
 
  def ref_func(x, y):
  return jnp.sum(lax.dot_general(x, y, dimension_numbers=(((2, 3), (0, 1)), ((), ()))))
@@ -266,7 +266,7 @@ def test_grad_fp8_mlp_randint(self, m, n, k):
  s = jax.random.uniform(subkeys[3], (k,), jnp.bfloat16, 5, 8)
 
  fp8_max = FP8Helper.generate_fp8_max_array(FP8Helper.NUM_META_PER_GEMM * 2)
- fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_SIZE),
+ fp8_metas_amax = jnp.zeros((FP8Helper.NUM_META_PER_GEMM * 2, FP8Helper.AMAX_HISTORY_LEN),
  jnp.float32)
  fp8_metas_scale = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
  fp8_metas_scale_inv = jnp.ones((FP8Helper.NUM_META_PER_GEMM * 2, 1), jnp.float32)
@@ -283,7 +283,6 @@ def primitive_func(x, ln_s, y, z, metas):
  ln_s,
  None,
  "rmsnorm",
- 0,
  *compute_type,
  activations=activations))
 
@@ -305,7 +304,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  amax: jnp.ndarray,
  scale: jnp.ndarray,
  scale_inv: jnp.ndarray,
- amax_history_idx: int,
  fwd_dtype,
  bwd_dtype,
  epsilon=1e-6,
@@ -323,7 +321,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[:FP8Helper.NUM_META_PER_GEMM],
  scale_inv[:FP8Helper.NUM_META_PER_GEMM])
  linear_1_out = fp8_dot(fp8_gemm_1_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -341,7 +338,6 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
  scale[FP8Helper.NUM_META_PER_GEMM:],
  scale_inv[FP8Helper.NUM_META_PER_GEMM:])
  output = fp8_dot(fp8_gemm_2_pkg,
- amax_history_idx,
  fwd_dtype,
  bwd_dtype,
  contracting_dims,
@@ -350,7 +346,7 @@ def fp8_ln_mlp_py(inputs: jnp.ndarray,
 
  def ref_func(x, ln_s, y, z, metas):
  return jnp.mean(
- fp8_ln_mlp_py(x, ln_s, y, z, *metas, 0, *compute_type, activations=activations))
+ fp8_ln_mlp_py(x, ln_s, y, z, *metas, *compute_type, activations=activations))
 
  value_n_grad_primitive_func = jit(value_and_grad(primitive_func, (0, 1, 2, 3)))
  value_n_grad_ref_func = jit(value_and_grad(ref_func, (0, 1, 2, 3)))