Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "07e225932235ccb0db5271b0874d00f086f28423"
Unverified Commit 8a7ab3dd authored by jberchtold-nvidia's avatar jberchtold-nvidia Committed by GitHub
Browse files

[JAX] NVFP4 support in TE/JAX (#2254) 


Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: default avatarPhuong Nguyen <phuonguyen@nvidia.com>
parent e99be1b6
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 1619 additions and 466 deletions
examples/jax/encoder/common.py
View file @ 8a7ab3dd
......@@ -33,6 +33,13 @@ def is_mxfp8_supported():
return gpu_arch >= 100
@lru_cache
def is_nvfp4_supported():
"""Return if FP8 has hardware supported"""
gpu_arch = get_device_compute_capability(0)
return gpu_arch >= 100
def assert_params_sufficiently_sharded(params, mesh, tolerance=0.01, print_info=False):
"""Checks whether most params are sharded across sharding axis.
......@@ -98,7 +105,7 @@ def assert_params_sufficiently_sharded(params, mesh, tolerance=0.01, print_info=
)
def get_fp8_recipe_from_name_string(name: str):
def get_quantization_recipe_from_name_string(name: str):
"""Query recipe from a given name string"""
match name:
case "DelayedScaling":
......@@ -107,5 +114,7 @@ def get_fp8_recipe_from_name_string(name: str):
return recipe.MXFP8BlockScaling()
case "Float8CurrentScaling":
return recipe.Float8CurrentScaling()
case "NVFP4BlockScaling":
return recipe.NVFP4BlockScaling()
case _:
raise ValueError(f"Invalid fp8_recipe, got {name}")
raise ValueError(f"Invalid quantization_recipe, got {name}")
examples/jax/encoder/run_test_multiprocessing_encoder.sh
View file @ 8a7ab3dd
......@@ -10,9 +10,11 @@ TEST_CASES=(
"test_te_delayed_scaling_fp8"
"test_te_current_scaling_fp8"
"test_te_mxfp8"
"test_te_nvfp4"
"test_te_bf16_shardy"
"test_te_delayed_scaling_fp8_shardy"
"test_te_current_scaling_fp8_shardy"
"test_te_nvfp4_shardy"
)
: ${TE_PATH:=/opt/transformerengine}
......
examples/jax/encoder/test_model_parallel_encoder.py
View file @ 8a7ab3dd
......@@ -21,13 +21,13 @@ from jax.sharding import PartitionSpec, NamedSharding
from common import (
is_bf16_supported,
get_fp8_recipe_from_name_string,
get_quantization_recipe_from_name_string,
assert_params_sufficiently_sharded,
)
import transformer_engine.jax as te
import transformer_engine.jax.cpp_extensions as tex
import transformer_engine.jax.flax as te_flax
from transformer_engine.jax.quantize import is_fp8_available, ScalingMode
from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode
DEVICE_DP_AXIS = "data"
......@@ -36,6 +36,7 @@ NAMED_BROADCAST_AXIS = "my_broadcast_axis"
NAMED_TP_AXIS = "my_tp_axis"
PARAMS_KEY = "params"
PARAMS_AXES_KEY = PARAMS_KEY + "_axes"
SR_KEY = "sr_rng"
DROPOUT_KEY = "dropout"
INPUT_KEY = "input_rng"
......@@ -121,6 +122,8 @@ def train_epoch(state, train_ds, batch_size, rngs, var_collect, train_fn):
epoch_accuracy = []
for perm in perms:
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
batch_inputs = train_ds["sentence"][perm, ...]
batch_masks = train_ds["mask"][perm, ...]
batch_labels = train_ds["label"][perm, ...]
......@@ -135,11 +138,11 @@ def train_epoch(state, train_ds, batch_size, rngs, var_collect, train_fn):
return state, avg_loss, avg_accuracy, var_collect
def eval_step(state, inputs, masks, labels, var_collect):
def eval_step(state, inputs, masks, labels, var_collect, rngs):
"""Computes loss and accuracy for a single batch."""
def loss_fn(var_collect, disable_dropout=False):
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout)
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout, rngs=rngs)
one_hot = jax.nn.one_hot(labels.astype(jnp.int32), 2)
loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
return loss, logits
......@@ -150,7 +153,7 @@ def eval_step(state, inputs, masks, labels, var_collect):
return loss, accuracy
def eval_model(state, test_ds, batch_size, var_collect, eval_fn):
def eval_model(state, test_ds, batch_size, var_collect, eval_fn, rngs):
"""Evaluation loop."""
test_ds_size = len(test_ds["sentence"])
num_steps = test_ds_size // batch_size
......@@ -159,11 +162,13 @@ def eval_model(state, test_ds, batch_size, var_collect, eval_fn):
all_accuracy = []
for batch_start in range(0, valid_size, batch_size):
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
batch_end = batch_start + batch_size
batch_inputs = test_ds["sentence"][batch_start:batch_end]
batch_masks = test_ds["mask"][batch_start:batch_end]
batch_labels = test_ds["label"][batch_start:batch_end]
loss, accuracy = eval_fn(state, batch_inputs, batch_masks, batch_labels, var_collect)
loss, accuracy = eval_fn(state, batch_inputs, batch_masks, batch_labels, var_collect, rngs)
all_loss.append(loss)
all_accuracy.append(accuracy)
......@@ -223,7 +228,7 @@ def get_datasets(max_seq_len):
def check_fp8(state, var_collect, inputs, masks, labels):
"Check if model includes FP8."
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0)}
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0), SR_KEY: jax.random.PRNGKey(0)}
func_jaxpr = str(jax.make_jaxpr(train_step)(state, inputs, masks, labels, var_collect, rngs))
assert "f8_e5m2" in func_jaxpr or "f8_e4m3" in func_jaxpr
......@@ -257,7 +262,7 @@ def train_and_evaluate(args):
), "Test batch size needs to be multiple of 32 for MXFP8"
if args.use_fp8:
fp8_recipe = get_fp8_recipe_from_name_string(args.fp8_recipe)
fp8_recipe = get_quantization_recipe_from_name_string(args.fp8_recipe)
else:
fp8_recipe = None
......@@ -275,7 +280,8 @@ def train_and_evaluate(args):
rng = jax.random.PRNGKey(args.seed)
rng, params_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
init_rngs = {PARAMS_KEY: params_rng, DROPOUT_KEY: dropout_rng}
rng, sr_rng = jax.random.split(rng)
init_rngs = {PARAMS_KEY: params_rng, DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
input_shape = [args.batch_size, args.max_seq_len]
mask_shape = [args.batch_size, 1, args.max_seq_len, args.max_seq_len]
......@@ -355,7 +361,14 @@ def train_and_evaluate(args):
train_step, in_shardings=in_shardings, out_shardings=out_shardings
)
in_shardings = (state_sharding, inputs_sharding, masks_sharding, labels_sharding, None)
in_shardings = (
state_sharding,
inputs_sharding,
masks_sharding,
labels_sharding,
None,
None,
)
out_shardings = (None, None)
jit_eval_step = jax.jit(
eval_step, in_shardings=in_shardings, out_shardings=out_shardings
......@@ -367,22 +380,24 @@ def train_and_evaluate(args):
if args.dry_run:
labels = jnp.zeros(label_shape, dtype=jnp.bfloat16)
rngs = {DROPOUT_KEY: dropout_rng}
rngs = {DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng_state}
jit_train_step(state, inputs, masks, labels, var_collect, rngs)
print("PASSED")
return None
for epoch in range(1, args.epochs + 1):
# Split and reassign to 'rng' to ensure unique rng for each step
rng, input_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng}
rng, sr_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
state, train_loss, train_accuracy, var_collect = train_epoch(
state, train_ds, args.batch_size, rngs, var_collect, jit_train_step
)
test_loss, test_accuracy = eval_model(
state, test_ds, args.test_batch_size, var_collect, jit_eval_step
state, test_ds, args.test_batch_size, var_collect, jit_eval_step, rngs
)
print(
......@@ -402,16 +417,16 @@ def encoder_parser(args):
parser.add_argument(
"--batch-size",
type=int,
default=128,
default=256,
metavar="N",
help="input batch size for training (default: 128)",
help="input batch size for training (default: 256)",
)
parser.add_argument(
"--test-batch-size",
type=int,
default=128,
default=256,
metavar="N",
help="input batch size for testing (default: 128)",
help="input batch size for testing (default: 256)",
)
parser.add_argument(
"--max-seq-len",
......@@ -466,8 +481,9 @@ def encoder_parser(args):
class TestEncoder(unittest.TestCase):
"""Encoder unittests"""
is_fp8_supported, fp8_reason = is_fp8_available(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_fp8_available(ScalingMode.MXFP8_1D_SCALING)
is_fp8_supported, fp8_reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING)
is_nvfp4_supported, nvfp4_reason = is_scaling_mode_supported(ScalingMode.NVFP4_1D_SCALING)
def setUp(self):
"""Run 5 epochs for testing"""
......@@ -477,7 +493,7 @@ class TestEncoder(unittest.TestCase):
def test_te_bf16(self):
"""Test Transformer Engine with BF16"""
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_delayed_scaling_fp8(self):
......@@ -485,7 +501,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "DelayedScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.361 and actual[1] > 0.84
@unittest.skipIf(not is_mxfp8_supported, mxfp8_reason)
def test_te_mxfp8(self):
......@@ -493,14 +509,22 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "MXFP8BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_nvfp4_supported, nvfp4_reason)
def test_te_nvfp4(self):
"""Test Transformer Engine with NVFP4"""
self.args.use_fp8 = True
self.args.fp8_recipe = "NVFP4BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.40 and actual[1] > 0.82
@unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16")
def test_te_bf16_with_sp(self):
"""Test Transformer Engine with BF16 + SP"""
self.args.enable_sp = True
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_delayed_scaling_fp8_with_sp(self):
......@@ -509,7 +533,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "DelayedScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_mxfp8_supported, mxfp8_reason)
def test_te_mxfp8_with_sp(self):
......@@ -518,14 +542,23 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "MXFP8BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_nvfp4_supported, nvfp4_reason)
def test_te_nvfp4_with_sp(self):
"""Test Transformer Engine with NVFP4"""
self.args.enable_sp = True
self.args.use_fp8 = True
self.args.fp8_recipe = "NVFP4BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.40 and actual[1] > 0.82
@unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16")
def test_te_bf16_shardy(self):
"""Test Transformer Engine with BF16"""
self.args.enable_shardy = True
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_delayed_scaling_fp8_shardy(self):
......@@ -534,7 +567,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "DelayedScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_delayed_scaling_fp8_with_sp_shardy(self):
......@@ -544,24 +577,27 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "DelayedScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.361 and actual[1] > 0.84
@unittest.skipIf(not is_mxfp8_supported, mxfp8_reason)
@unittest.skipIf(
tex.gemm_uses_jax_dot(), "`jax.nn.scaled_matmul()` does not support the Shardy partitioner."
)
def test_te_mxfp8_shardy(self):
"""Test Transformer Engine with MXFP8"""
self.args.enable_shardy = True
self.args.use_fp8 = True
self.args.fp8_recipe = "MXFP8BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_nvfp4_supported, nvfp4_reason)
def test_te_nvfp4_shardy(self):
"""Test Transformer Engine with NVFP4"""
self.args.enable_shardy = True
self.args.use_fp8 = True
self.args.fp8_recipe = "NVFP4BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.40 and actual[1] > 0.82
@unittest.skipIf(not is_mxfp8_supported, mxfp8_reason)
@unittest.skipIf(
tex.gemm_uses_jax_dot(), "`jax.nn.scaled_matmul()` does not support the Shardy partitioner."
)
def test_te_mxfp8_with_sp_shardy(self):
"""Test Transformer Engine with MXFP8 + SP"""
self.args.enable_shardy = True
......@@ -569,7 +605,17 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "MXFP8BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.39 and actual[1] > 0.83
assert actual[0] < 0.36 and actual[1] > 0.84
@unittest.skipIf(not is_nvfp4_supported, nvfp4_reason)
def test_te_nvfp4_with_sp_shardy(self):
"""Test Transformer Engine with NVFP4"""
self.args.enable_shardy = True
self.args.enable_sp = True
self.args.use_fp8 = True
self.args.fp8_recipe = "NVFP4BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.40 and actual[1] > 0.82
if __name__ == "__main__":
......
examples/jax/encoder/test_multigpu_encoder.py
View file @ 8a7ab3dd
......@@ -19,17 +19,18 @@ from flax.training import train_state
from jax.experimental import mesh_utils
from jax.sharding import PartitionSpec, NamedSharding
from common import is_bf16_supported, get_fp8_recipe_from_name_string
from common import is_bf16_supported, get_quantization_recipe_from_name_string
import transformer_engine.jax as te
import transformer_engine.jax.cpp_extensions as tex
import transformer_engine.jax.flax as te_flax
from transformer_engine.jax.quantize import is_fp8_available, ScalingMode
from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode
DEVICE_DP_AXIS = "data"
PARAMS_KEY = "params"
PARAMS_AXES_KEY = PARAMS_KEY + "_axes"
DROPOUT_KEY = "dropout"
SR_KEY = "sr_rng"
INPUT_KEY = "input_rng"
......@@ -97,6 +98,8 @@ def train_epoch(state, train_ds, batch_size, rngs, var_collect, train_fn):
epoch_accuracy = []
for perm in perms:
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
batch_inputs = train_ds["sentence"][perm, ...]
batch_masks = train_ds["mask"][perm, ...]
batch_labels = train_ds["label"][perm, ...]
......@@ -111,11 +114,11 @@ def train_epoch(state, train_ds, batch_size, rngs, var_collect, train_fn):
return state, avg_loss, avg_accuracy, var_collect
def eval_step(state, inputs, masks, labels, var_collect):
def eval_step(state, inputs, masks, labels, var_collect, rngs):
"""Computes loss and accuracy for a single batch."""
def loss_fn(var_collect, disable_dropout=False):
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout)
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout, rngs=rngs)
one_hot = jax.nn.one_hot(labels.astype(jnp.int32), 2)
loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
return loss, logits
......@@ -126,7 +129,7 @@ def eval_step(state, inputs, masks, labels, var_collect):
return loss, accuracy
def eval_model(state, test_ds, batch_size, var_collect, eval_fn):
def eval_model(state, test_ds, batch_size, var_collect, eval_fn, rngs):
"""Evaluation loop."""
test_ds_size = len(test_ds["sentence"])
num_steps = test_ds_size // batch_size
......@@ -135,11 +138,13 @@ def eval_model(state, test_ds, batch_size, var_collect, eval_fn):
all_accuracy = []
for batch_start in range(0, valid_size, batch_size):
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
batch_end = batch_start + batch_size
batch_inputs = test_ds["sentence"][batch_start:batch_end]
batch_masks = test_ds["mask"][batch_start:batch_end]
batch_labels = test_ds["label"][batch_start:batch_end]
loss, accuracy = eval_fn(state, batch_inputs, batch_masks, batch_labels, var_collect)
loss, accuracy = eval_fn(state, batch_inputs, batch_masks, batch_labels, var_collect, rngs)
all_loss.append(loss)
all_accuracy.append(accuracy)
......@@ -199,7 +204,7 @@ def get_datasets(max_seq_len):
def check_fp8(state, var_collect, inputs, masks, labels):
"Check if model includes FP8."
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0)}
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0), SR_KEY: jax.random.PRNGKey(0)}
func_jaxpr = str(jax.make_jaxpr(train_step)(state, inputs, masks, labels, var_collect, rngs))
assert "f8_e5m2" in func_jaxpr or "f8_e4m3" in func_jaxpr
......@@ -254,7 +259,7 @@ def train_and_evaluate(args):
), "Test batch size needs to be multiple of 32 for MXFP8"
if args.use_fp8:
fp8_recipe = get_fp8_recipe_from_name_string(args.fp8_recipe)
fp8_recipe = get_quantization_recipe_from_name_string(args.fp8_recipe)
else:
fp8_recipe = None
......@@ -270,6 +275,7 @@ def train_and_evaluate(args):
rng = jax.random.PRNGKey(args.seed)
rng, params_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
rng, sr_rng = jax.random.split(rng)
init_rngs = {PARAMS_KEY: params_rng, DROPOUT_KEY: dropout_rng}
input_shape = [args.batch_size, args.max_seq_len]
......@@ -322,7 +328,14 @@ def train_and_evaluate(args):
train_step, in_shardings=in_shardings, out_shardings=out_shardings
)
in_shardings = (state_sharding, inputs_sharding, masks_sharding, labels_sharding, None)
in_shardings = (
state_sharding,
inputs_sharding,
masks_sharding,
labels_sharding,
None,
None,
)
out_shardings = (None, None)
jit_eval_step = jax.jit(
eval_step, in_shardings=in_shardings, out_shardings=out_shardings
......@@ -334,22 +347,24 @@ def train_and_evaluate(args):
if args.dry_run:
labels = jnp.zeros(label_shape, dtype=jnp.bfloat16)
rngs = {DROPOUT_KEY: dropout_rng}
rngs = {DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
jit_train_step(state, inputs, masks, labels, var_collect, rngs)
print("PASSED")
return None
for epoch in range(1, args.epochs + 1):
# Split and reassign to 'rng' to ensure unique rng for each step
rng, input_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng}
rng, sr_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
state, train_loss, train_accuracy, var_collect = train_epoch(
state, train_ds, args.batch_size, rngs, var_collect, jit_train_step
)
test_loss, test_accuracy = eval_model(
state, test_ds, args.test_batch_size, var_collect, jit_eval_step
state, test_ds, args.test_batch_size, var_collect, jit_eval_step, rngs
)
print(
......@@ -369,16 +384,16 @@ def encoder_parser(args):
parser.add_argument(
"--batch-size",
type=int,
default=256,
default=512,
metavar="N",
help="input batch size for training (default: 256)",
help="input batch size for training (default: 512)",
)
parser.add_argument(
"--test-batch-size",
type=int,
default=256,
default=512,
metavar="N",
help="input batch size for testing (default: 256)",
help="input batch size for testing (default: 512)",
)
parser.add_argument(
"--max-seq-len",
......@@ -430,8 +445,9 @@ def encoder_parser(args):
class TestEncoder(unittest.TestCase):
"""Encoder unittests"""
is_fp8_supported, fp8_reason = is_fp8_available(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_fp8_available(ScalingMode.MXFP8_1D_SCALING)
is_fp8_supported, fp8_reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING)
is_nvfp4_supported, nvfp4_reason = is_scaling_mode_supported(ScalingMode.NVFP4_1D_SCALING)
def setUp(self):
"""Run 5 epochs for testing"""
......@@ -441,7 +457,7 @@ class TestEncoder(unittest.TestCase):
def test_te_bf16(self):
"""Test Transformer Engine with BF16"""
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
assert actual[0] < 0.51 and actual[1] > 0.75
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_delayed_scaling_fp8(self):
......@@ -449,7 +465,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "DelayedScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
assert actual[0] < 0.51 and actual[1] > 0.75
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_current_scaling_fp8(self):
......@@ -457,7 +473,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "Float8CurrentScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
assert actual[0] < 0.51 and actual[1] > 0.749
@unittest.skipIf(not is_mxfp8_supported, mxfp8_reason)
def test_te_mxfp8(self):
......@@ -465,6 +481,14 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "MXFP8BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.51 and actual[1] > 0.75
@unittest.skipIf(not is_nvfp4_supported, nvfp4_reason)
def test_te_nvfp4(self):
"""Test Transformer Engine with NVFP4"""
self.args.use_fp8 = True
self.args.fp8_recipe = "NVFP4BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
@unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16")
......@@ -472,7 +496,7 @@ class TestEncoder(unittest.TestCase):
"""Test Transformer Engine with BF16"""
self.args.enable_shardy = True
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
assert actual[0] < 0.51 and actual[1] > 0.75
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_delayed_scaling_fp8_shardy(self):
......@@ -481,7 +505,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "DelayedScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
assert actual[0] < 0.51 and actual[1] > 0.75
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_current_scaling_fp8_shardy(self):
......@@ -490,18 +514,24 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "Float8CurrentScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
assert actual[0] < 0.51 and actual[1] > 0.749
@unittest.skipIf(not is_mxfp8_supported, mxfp8_reason)
@unittest.skipIf(
tex.gemm_uses_jax_dot(), "`jax.nn.scaled_matmul()` does not support the Shardy partitioner."
)
def test_te_mxfp8_shardy(self):
"""Test Transformer Engine with MXFP8"""
self.args.enable_shardy = True
self.args.use_fp8 = True
self.args.fp8_recipe = "MXFP8BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.51 and actual[1] > 0.75
@unittest.skipIf(not is_nvfp4_supported, nvfp4_reason)
def test_te_nvfp4_shardy(self):
"""Test Transformer Engine with NVFP4"""
self.args.enable_shardy = True
self.args.use_fp8 = True
self.args.fp8_recipe = "NVFP4BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.52 and actual[1] > 0.74
......
examples/jax/encoder/test_multiprocessing_encoder.py
View file @ 8a7ab3dd
......@@ -25,7 +25,8 @@ from common import (
is_bf16_supported,
is_fp8_supported,
is_mxfp8_supported,
get_fp8_recipe_from_name_string,
is_nvfp4_supported,
get_quantization_recipe_from_name_string,
)
import transformer_engine.jax as te
import transformer_engine.jax.cpp_extensions as tex
......@@ -39,6 +40,7 @@ NAMED_BROADCAST_AXIS = "my_broadcast_axis"
NAMED_TP_AXIS = "my_tp_axis"
PARAMS_KEY = "params"
PARAMS_AXES_KEY = PARAMS_KEY + "_axes"
SR_KEY = "sr_rng"
DROPOUT_KEY = "dropout"
INPUT_KEY = "input_rng"
......@@ -175,6 +177,8 @@ def train_epoch(
epoch_accuracy = []
for perm in perms:
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
batch_input = sentence[perm, ...]
batch_mask = mask[perm, ...]
batch_label = label[perm, ...]
......@@ -200,11 +204,11 @@ def train_epoch(
return state, avg_loss, avg_accuracy, var_collect
def eval_step(state, inputs, masks, labels, var_collect):
def eval_step(state, inputs, masks, labels, var_collect, rngs):
"""Computes loss and accuracy for a single batch."""
def loss_fn(var_collect, disable_dropout=False):
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout)
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout, rngs=rngs)
one_hot = jax.nn.one_hot(labels, 2)
loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
return loss, logits
......@@ -216,7 +220,16 @@ def eval_step(state, inputs, masks, labels, var_collect):
def eval_model(
state, test_ds, batch_size, var_collect, eval_fn, mesh, inputs_pspec, masks_pspec, labels_pspec
state,
test_ds,
batch_size,
var_collect,
eval_fn,
mesh,
inputs_pspec,
masks_pspec,
labels_pspec,
rngs,
):
"""Evaluation loop."""
global_input_shape, input_named_sharding, sentence = shard_array_wrapper(
......@@ -233,7 +246,8 @@ def eval_model(
all_accuracy = []
for batch_input, batch_mask, batch_label in zip(sentence, mask, label):
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
shard_input = jax.make_array_from_single_device_arrays(
global_input_shape, input_named_sharding, [batch_input]
)
......@@ -244,7 +258,7 @@ def eval_model(
global_label_shape, label_named_sharding, [batch_label]
)
loss, accuracy = eval_fn(state, shard_input, shard_mask, shard_label, var_collect)
loss, accuracy = eval_fn(state, shard_input, shard_mask, shard_label, var_collect, rngs)
all_loss.append(loss)
all_accuracy.append(accuracy)
......@@ -303,7 +317,7 @@ def get_datasets(max_seq_len):
def check_fp8(state, var_collect, inputs, masks, labels):
"Check if model includes FP8."
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0)}
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0), SR_KEY: jax.random.PRNGKey(0)}
func_jaxpr = str(jax.make_jaxpr(train_step)(state, inputs, masks, labels, var_collect, rngs))
assert "f8_e5m2" in func_jaxpr or "f8_e4m3" in func_jaxpr
......@@ -372,7 +386,7 @@ def train_and_evaluate(args):
), "Test batch size needs to be multiple of 32 for MXFP8"
if args.use_fp8:
fp8_recipe = get_fp8_recipe_from_name_string(args.fp8_recipe)
fp8_recipe = get_quantization_recipe_from_name_string(args.fp8_recipe)
else:
fp8_recipe = None
......@@ -390,7 +404,8 @@ def train_and_evaluate(args):
rng = jax.random.PRNGKey(args.seed)
rng, params_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
init_rngs = {PARAMS_KEY: params_rng, DROPOUT_KEY: dropout_rng}
rng, sr_rng = jax.random.split(rng)
init_rngs = {PARAMS_KEY: params_rng, DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
input_shape = [args.batch_size, args.max_seq_len]
mask_shape = [args.batch_size, 1, args.max_seq_len, args.max_seq_len]
......@@ -444,7 +459,14 @@ def train_and_evaluate(args):
train_step, in_shardings=in_shardings, out_shardings=out_shardings
)
in_shardings = (state_sharding, inputs_sharding, masks_sharding, labels_sharding, None)
in_shardings = (
state_sharding,
inputs_sharding,
masks_sharding,
labels_sharding,
None,
None,
)
out_shardings = (None, None)
jit_eval_step = jax.jit(
eval_step, in_shardings=in_shardings, out_shardings=out_shardings
......@@ -456,14 +478,16 @@ def train_and_evaluate(args):
if args.dry_run:
labels = jnp.zeros(label_shape, dtype=jnp.bfloat16)
rngs = {DROPOUT_KEY: dropout_rng}
rngs = {DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng_state}
jit_train_step(state, inputs, masks, labels, var_collect, rngs)
print("PASSED")
else:
for epoch in range(1, args.epochs + 1):
# Split and reassign to 'rng' to ensure unique rng for each step
rng, input_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng}
rng, sr_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
state, train_loss, train_accuracy, var_collect = train_epoch(
state,
......@@ -488,6 +512,7 @@ def train_and_evaluate(args):
inputs_pspec,
masks_pspec,
labels_sharding.spec,
rngs,
)
if args.process_id == 0:
print(
......@@ -508,16 +533,16 @@ def encoder_parser(args):
parser.add_argument(
"--batch-size",
type=int,
default=128,
default=256,
metavar="N",
help="input batch size for training (default: 128)",
help="input batch size for training (default: 256)",
)
parser.add_argument(
"--test-batch-size",
type=int,
default=128,
default=256,
metavar="N",
help="input batch size for testing (default: 128)",
help="input batch size for testing (default: 256)",
)
parser.add_argument(
"--max-seq-len",
......@@ -629,7 +654,7 @@ class TestEncoder(unittest.TestCase):
def test_te_current_scaling_fp8(self):
"""Test Transformer Engine with CurrentScaling FP8"""
result = self.exec(True, "Float8CurrentScaling")
assert result[0] < 0.43 and result[1] > 0.80
assert result[0] < 0.432 and result[1] > 0.80
@unittest.skipIf(
not is_mxfp8_supported(), "Device compute capability 10.0+ is required for MXFP8"
......@@ -639,6 +664,14 @@ class TestEncoder(unittest.TestCase):
result = self.exec(True, "MXFP8BlockScaling")
assert result[0] < 0.43 and result[1] > 0.80
@unittest.skipIf(
not is_nvfp4_supported(), "Device compute capability 10.0+ is required for NVFP4"
)
def test_te_nvfp4(self):
"""Test Transformer Engine with NVFP4"""
result = self.exec(True, "NVFP4BlockScaling")
assert result[0] < 0.451 and result[1] > 0.79
@unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16")
def test_te_bf16_shardy(self):
"""Test Transformer Engine with BF16"""
......@@ -659,19 +692,24 @@ class TestEncoder(unittest.TestCase):
def test_te_current_scaling_fp8_shardy(self):
"""Test Transformer Engine with CurrentScaling FP8"""
result = self.exec(True, "Float8CurrentScaling", enable_shardy=True)
assert result[0] < 0.43 and result[1] > 0.80
assert result[0] < 0.432 and result[1] > 0.80
@unittest.skipIf(
not is_mxfp8_supported(), "Device compute capability 10.0+ is required for MXFP8"
)
@unittest.skipIf(
tex.gemm_uses_jax_dot(), "`jax.nn.scaled_matmul()` does not support the Shardy partitioner."
)
def test_te_mxfp8_shardy(self):
"""Test Transformer Engine with MXFP8"""
result = self.exec(True, "MXFP8BlockScaling", enable_shardy=True)
assert result[0] < 0.43 and result[1] > 0.80
@unittest.skipIf(
not is_nvfp4_supported(), "Device compute capability 10.0+ is required for NVFP4"
)
def test_te_nvfp4_shardy(self):
"""Test Transformer Engine with NVFP4"""
result = self.exec(True, "NVFP4BlockScaling", enable_shardy=True)
assert result[0] < 0.451 and result[1] > 0.79
if __name__ == "__main__":
train_and_evaluate(encoder_parser(None))
examples/jax/encoder/test_single_gpu_encoder.py
View file @ 8a7ab3dd
......@@ -16,14 +16,15 @@ from datasets import load_dataset
from flax import linen as nn
from flax.training import train_state
from common import is_bf16_supported, get_fp8_recipe_from_name_string
from common import is_bf16_supported, get_quantization_recipe_from_name_string
import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax
from transformer_engine.jax.quantize import is_fp8_available, ScalingMode
from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode
PARAMS_KEY = "params"
DROPOUT_KEY = "dropout"
SR_KEY = "sr_rng"
INPUT_KEY = "input_rng"
......@@ -92,6 +93,8 @@ def train_epoch(state, train_ds, batch_size, rngs, var_collect):
epoch_accuracy = []
for perm in perms:
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
batch_inputs = train_ds["sentence"][perm, ...]
batch_masks = train_ds["mask"][perm, ...]
batch_labels = train_ds["label"][perm, ...]
......@@ -107,11 +110,11 @@ def train_epoch(state, train_ds, batch_size, rngs, var_collect):
@jax.jit
def eval_step(state, inputs, masks, labels, var_collect):
def eval_step(state, inputs, masks, labels, var_collect, rngs):
"""Computes loss and accuracy for a single batch."""
def loss_fn(var_collect, disable_dropout=False):
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout)
logits = state.apply_fn(var_collect, inputs, masks, disable_dropout, rngs=rngs)
one_hot = jax.nn.one_hot(labels.astype(jnp.int32), 2)
loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
return loss, logits
......@@ -122,7 +125,7 @@ def eval_step(state, inputs, masks, labels, var_collect):
return loss, accuracy
def eval_model(state, test_ds, batch_size, var_collect):
def eval_model(state, test_ds, batch_size, var_collect, rngs):
"""Evaluation loop."""
test_ds_size = len(test_ds["sentence"])
num_steps = test_ds_size // batch_size
......@@ -131,11 +134,15 @@ def eval_model(state, test_ds, batch_size, var_collect):
all_accuracy = []
for batch_start in range(0, valid_size, batch_size):
# Split and reassign to 'rngs' to ensure unique rng for each step
rngs = {key: jax.random.split(rngs[key])[1] for key in rngs}
batch_end = batch_start + batch_size
batch_inputs = test_ds["sentence"][batch_start:batch_end]
batch_masks = test_ds["mask"][batch_start:batch_end]
batch_labels = test_ds["label"][batch_start:batch_end]
loss, accuracy = eval_step(state, batch_inputs, batch_masks, batch_labels, var_collect)
loss, accuracy = eval_step(
state, batch_inputs, batch_masks, batch_labels, var_collect, rngs
)
all_loss.append(loss)
all_accuracy.append(accuracy)
......@@ -195,7 +202,7 @@ def get_datasets(max_seq_len):
def check_fp8(state, var_collect, inputs, masks, labels):
"Check if model includes FP8."
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0)}
rngs = {DROPOUT_KEY: jax.random.PRNGKey(0), SR_KEY: jax.random.PRNGKey(0)}
func_jaxpr = str(jax.make_jaxpr(train_step)(state, inputs, masks, labels, var_collect, rngs))
assert "f8_e5m2" in func_jaxpr or "f8_e4m3" in func_jaxpr
......@@ -208,14 +215,15 @@ def train_and_evaluate(args):
rng = jax.random.PRNGKey(args.seed)
rng, params_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
init_rngs = {PARAMS_KEY: params_rng, DROPOUT_KEY: dropout_rng}
rng, sr_rng = jax.random.split(rng)
init_rngs = {PARAMS_KEY: params_rng, DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
input_shape = [args.batch_size, args.max_seq_len]
mask_shape = [args.batch_size, 1, args.max_seq_len, args.max_seq_len]
label_shape = [args.batch_size]
if args.use_fp8:
fp8_recipe = get_fp8_recipe_from_name_string(args.fp8_recipe)
fp8_recipe = get_quantization_recipe_from_name_string(args.fp8_recipe)
else:
fp8_recipe = None
......@@ -238,21 +246,25 @@ def train_and_evaluate(args):
if args.dry_run:
labels = jnp.zeros(label_shape, dtype=jnp.bfloat16)
rngs = {DROPOUT_KEY: dropout_rng}
rngs = {DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
train_step(state, inputs, masks, labels, var_collect, rngs)
print("PASSED")
return None
for epoch in range(1, args.epochs + 1):
# Split and reassign to 'rng' to ensure unique rng for each step
rng, input_rng = jax.random.split(rng)
rng, dropout_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng}
rng, sr_rng = jax.random.split(rng)
rngs = {INPUT_KEY: input_rng, DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
state, train_loss, train_accuracy, var_collect = train_epoch(
state, train_ds, args.batch_size, rngs, var_collect
)
test_loss, test_accuracy = eval_model(state, test_ds, args.test_batch_size, var_collect)
test_loss, test_accuracy = eval_model(
state, test_ds, args.test_batch_size, var_collect, rngs
)
print(
f"Epoch: {epoch:>2} "
......@@ -329,8 +341,9 @@ def encoder_parser(args):
class TestEncoder(unittest.TestCase):
"""Encoder unittests"""
is_fp8_supported, fp8_reason = is_fp8_available(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_fp8_available(ScalingMode.MXFP8_1D_SCALING)
is_fp8_supported, fp8_reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING)
is_nvfp4_supported, nvfp4_reason = is_scaling_mode_supported(ScalingMode.NVFP4_1D_SCALING)
def setUp(self):
"""Run 3 epochs for testing"""
......@@ -340,7 +353,7 @@ class TestEncoder(unittest.TestCase):
def test_te_bf16(self):
"""Test Transformer Engine with BF16"""
actual = train_and_evaluate(self.args)
assert actual[0] < 0.45 and actual[1] > 0.79
assert actual[0] < 0.452 and actual[1] > 0.788
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_delayed_scaling_fp8(self):
......@@ -348,7 +361,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "DelayedScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.455 and actual[1] > 0.79
assert actual[0] < 0.457 and actual[1] > 0.784
@unittest.skipIf(not is_fp8_supported, fp8_reason)
def test_te_current_scaling_fp8(self):
......@@ -356,7 +369,7 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "Float8CurrentScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.455 and actual[1] > 0.79
assert actual[0] < 0.461 and actual[1] > 0.784
@unittest.skipIf(not is_mxfp8_supported, mxfp8_reason)
def test_te_mxfp8(self):
......@@ -364,7 +377,15 @@ class TestEncoder(unittest.TestCase):
self.args.use_fp8 = True
self.args.fp8_recipe = "MXFP8BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.455 and actual[1] > 0.79
assert actual[0] < 0.457 and actual[1] > 0.784
@unittest.skipIf(not is_nvfp4_supported, nvfp4_reason)
def test_te_nvfp4(self):
"""Test Transformer Engine with NVFP4"""
self.args.use_fp8 = True
self.args.fp8_recipe = "NVFP4BlockScaling"
actual = train_and_evaluate(self.args)
assert actual[0] < 0.476 and actual[1] > 0.775
if __name__ == "__main__":
......
examples/jax/mnist/test_single_gpu_mnist.py
View file @ 8a7ab3dd
......@@ -18,11 +18,11 @@ from flax.training import train_state
import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax
from transformer_engine.jax.quantize import is_fp8_available, ScalingMode
from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode
DIR = str(Path(__file__).resolve().parents[1])
sys.path.append(str(DIR))
from encoder.common import is_bf16_supported, get_fp8_recipe_from_name_string
from encoder.common import is_bf16_supported, get_quantization_recipe_from_name_string
IMAGE_H = 28
IMAGE_W = 28
......@@ -189,7 +189,7 @@ def train_and_evaluate(args):
label_shape = [args.batch_size]
if args.use_fp8:
fp8_recipe = get_fp8_recipe_from_name_string(args.fp8_recipe)
fp8_recipe = get_quantization_recipe_from_name_string(args.fp8_recipe)
else:
fp8_recipe = None
......@@ -308,8 +308,8 @@ def mnist_parser(args):
class TestMNIST(unittest.TestCase):
"""MNIST unittests"""
is_fp8_supported, fp8_reason = is_fp8_available(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_fp8_available(ScalingMode.MXFP8_1D_SCALING)
is_fp8_supported, fp8_reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING)
@classmethod
def setUpClass(cls):
......
tests/jax/test_custom_call_compute.py
View file @ 8a7ab3dd
This diff is collapsed.
tests/jax/test_distributed_layernorm_mlp.py
View file @ 8a7ab3dd
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
import re
from typing import Callable, Sequence, Union, Optional
import pytest
......@@ -17,7 +18,11 @@ from utils import (
)
from transformer_engine.common import recipe
from transformer_engine.jax.quantize import is_fp8_available, ScalingMode
from transformer_engine.jax.quantize import (
is_fp8_available,
ScalingMode,
get_quantize_config_with_recipe,
)
from transformer_engine.jax import fp8_autocast
from transformer_engine.jax.flax import LayerNormMLP
from transformer_engine.jax.layernorm_mlp import layernorm_mlp
......@@ -33,19 +38,20 @@ from transformer_engine.jax.sharding import (
W_JOINED_AXES,
)
from transformer_engine.jax.sharding import MeshResource
from transformer_engine.jax.quantize import QuantizerFactory
from transformer_engine.jax.quantize import (
QuantizerFactory,
get_supported_quantization_recipes,
is_scaling_mode_supported,
)
from transformer_engine.jax.cpp_extensions.misc import get_min_device_compute_capability
is_fp8_supported, reason = is_fp8_available()
is_mxfp8_supported, reason = is_fp8_available(ScalingMode.MXFP8_1D_SCALING)
is_fp8_supported, reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING)
is_nvfp4_supported, reason = is_scaling_mode_supported(ScalingMode.NVFP4_1D_SCALING)
SUPPORTED_RECIPES = []
if is_fp8_supported:
SUPPORTED_RECIPES.append(pytest.param(recipe.DelayedScaling(), id="DelayedScaling"))
SUPPORTED_RECIPES.append(pytest.param(recipe.Float8CurrentScaling(), id="CurrentScaling"))
if is_mxfp8_supported:
SUPPORTED_RECIPES.append(pytest.param(recipe.MXFP8BlockScaling(), id="MXFP8BlockScaling"))
SUPPORTED_RECIPES = get_supported_quantization_recipes()
SUPPORTED_RECIPES = [pytest.param(r, id=r.__class__.__name__) for r in SUPPORTED_RECIPES]
DTYPES = [jnp.bfloat16, jnp.float16]
INPUT_SHAPE = [[4, 128, 256]] # [batch, seqlen, hidden_in]
......@@ -141,6 +147,7 @@ class TestDistributedLayernormMLP:
layernorm_type: str = "rmsnorm",
activation_type: Sequence[Union[str, Callable]] = ("gelu",),
multi_gpus: bool = False,
quantization_recipe: recipe.Recipe = None,
) -> jnp.ndarray:
if multi_gpus:
......@@ -154,7 +161,9 @@ class TestDistributedLayernormMLP:
dot_1_input_axes = dot_2_input_axes = None
kernel_1_axes = kernel_2_axes = None
quantizer_sets = QuantizerFactory.create_set(n_quantizer_sets=2)
quantizer_sets = QuantizerFactory.create_set(
n_quantizer_sets=2, fp8_recipe=quantization_recipe
)
# out = ((x * kernel_1) + bias_1) * kernel_2 + bias_2
return jnp.mean(
......@@ -182,7 +191,7 @@ class TestDistributedLayernormMLP:
use_bias,
input_shape,
dtype,
fp8_recipe,
quantization_recipe,
use_shardy,
with_jax_gemm,
):
......@@ -202,7 +211,9 @@ class TestDistributedLayernormMLP:
# Single GPU
with fp8_autocast(
enabled=fp8_recipe is not None, fp8_recipe=fp8_recipe, mesh_resource=MeshResource()
enabled=quantization_recipe is not None,
fp8_recipe=quantization_recipe,
mesh_resource=MeshResource(),
):
single_jitter = jax.jit(
value_and_grad_func,
......@@ -214,7 +225,9 @@ class TestDistributedLayernormMLP:
devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape)
mesh = Mesh(devices, mesh_axes)
with mesh, fp8_autocast(
enabled=fp8_recipe is not None, fp8_recipe=fp8_recipe, mesh_resource=mesh_resource
enabled=quantization_recipe is not None,
fp8_recipe=quantization_recipe,
mesh_resource=mesh_resource,
):
k1_sharding = NamedSharding(mesh, PartitionSpec("fsdp", None, "tpsp"))
k2_sharding = NamedSharding(mesh, PartitionSpec("tpsp", "fsdp"))
......@@ -254,10 +267,16 @@ class TestDistributedLayernormMLP:
multi_fwd, multi_grads = multi_jitter(*multi_inputs, *static_inputs, True)
fwd_test_type = dtype if fp8_recipe is None else jnp.float8_e4m3fn
bwd_test_type = dtype if fp8_recipe is None else jnp.float8_e5m2
fwd_test_type = bwd_test_type = dtype
if quantization_recipe is not None:
quantize_config = get_quantize_config_with_recipe(quantization_recipe)
fwd_test_type = quantize_config.FWD_DTYPE
bwd_test_type = quantize_config.BWD_DTYPE
assert_allclose(multi_fwd, single_fwd, dtype=fwd_test_type)
if fwd_test_type == jnp.float16 and use_bias:
assert_allclose(multi_fwd, single_fwd, atol=0.04, rtol=1.5)
else:
assert_allclose(multi_fwd, single_fwd, dtype=fwd_test_type)
for i in range(len(inputs)):
if multi_grads[i] is not None:
......@@ -278,13 +297,12 @@ class TestDistributedLayernormMLP:
err_msg=f"multi_grads[{i}] is not close",
)
@pytest.mark.skipif(not is_fp8_supported, reason=reason)
@pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs())
@pytest_parametrize_wrapper("input_shape", INPUT_SHAPE)
@pytest_parametrize_wrapper("activation_type", [("gelu",), ("gelu", "linear")])
@pytest_parametrize_wrapper("dtype", DTYPES)
@pytest_parametrize_wrapper("use_bias", [True, False])
@pytest_parametrize_wrapper("fp8_recipe", [None] + SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("quantization_recipe", [None] + SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("with_jax_gemm", [False, True])
def test_layernorm_mlp_grad(
self,
......@@ -293,27 +311,28 @@ class TestDistributedLayernormMLP:
use_bias,
input_shape,
dtype,
fp8_recipe,
quantization_recipe,
with_jax_gemm,
):
if dtype == jnp.float16 and quantization_recipe is not None and quantization_recipe.nvfp4():
pytest.skip("NVFP4 GEMM + Float16 output is unsupported!")
self._test_layernorm_mlp_grad(
mesh_config,
activation_type,
use_bias,
input_shape,
dtype,
fp8_recipe,
quantization_recipe,
use_shardy=False,
with_jax_gemm=with_jax_gemm,
)
@pytest.mark.skipif(not is_fp8_supported, reason=reason)
@pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs())
@pytest_parametrize_wrapper("input_shape", INPUT_SHAPE)
@pytest_parametrize_wrapper("activation_type", [("gelu",), ("gelu", "linear")])
@pytest_parametrize_wrapper("dtype", DTYPES)
@pytest_parametrize_wrapper("use_bias", [True, False])
@pytest_parametrize_wrapper("fp8_recipe", [None] + SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("quantization_recipe", [None] + SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("with_jax_gemm", [False, True])
def test_layernorm_mlp_grad_shardy(
self,
......@@ -322,18 +341,18 @@ class TestDistributedLayernormMLP:
use_bias,
input_shape,
dtype,
fp8_recipe,
quantization_recipe,
with_jax_gemm,
):
if with_jax_gemm and isinstance(fp8_recipe, recipe.MXFP8BlockScaling):
pytest.skip("`jax.nn.scaled_matmul()` does not support the Shardy partitioner.")
if dtype == jnp.float16 and quantization_recipe is not None and quantization_recipe.nvfp4():
pytest.skip("NVFP4 GEMM + Float16 output is unsupported!")
self._test_layernorm_mlp_grad(
mesh_config,
activation_type,
use_bias,
input_shape,
dtype,
fp8_recipe=fp8_recipe,
quantization_recipe=quantization_recipe,
use_shardy=True,
with_jax_gemm=with_jax_gemm,
)
......@@ -346,7 +365,7 @@ class TestDistributedLayernormMLP:
input_shape,
dtype,
use_fp8,
fp8_recipe,
quantization_recipe,
use_shardy,
with_jax_gemm,
):
......@@ -355,14 +374,16 @@ class TestDistributedLayernormMLP:
layernorm_type = "rmsnorm"
rng = jax.random.PRNGKey(0)
subkeys = jax.random.split(rng, 2)
subkeys = jax.random.split(rng, 3)
x = jax.random.normal(subkeys[0], (batch, seqlen, hidden_in), dtype)
init_rngs = {"params": subkeys[1]}
init_rngs = {"params": subkeys[1], "sr_rng": subkeys[2]}
with use_jax_gemm(enabled=with_jax_gemm):
# Single GPUs
with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe, mesh_resource=MeshResource()):
with fp8_autocast(
enabled=use_fp8, fp8_recipe=quantization_recipe, mesh_resource=MeshResource()
):
ln_mlp_single = LayerNormMLP(
layernorm_type=layernorm_type,
intermediate_dim=INTERMEDIATE,
......@@ -371,7 +392,7 @@ class TestDistributedLayernormMLP:
)
params_single = ln_mlp_single.init(init_rngs, x, deterministic=True)
mlp_out_single, ln_out_single = ln_mlp_single.apply(
params_single, x, deterministic=True
params_single, x, deterministic=True, rngs={"sr_rng": subkeys[2]}
)
# Multi GPUs
......@@ -379,7 +400,7 @@ class TestDistributedLayernormMLP:
devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape)
mesh = Mesh(devices, mesh_axes)
with mesh, fp8_autocast(
enabled=use_fp8, fp8_recipe=fp8_recipe, mesh_resource=mesh_resource
enabled=use_fp8, fp8_recipe=quantization_recipe, mesh_resource=mesh_resource
):
ln_mlp_sharded = LayerNormMLP(
layernorm_type=layernorm_type,
......@@ -399,7 +420,7 @@ class TestDistributedLayernormMLP:
)
params_sharded = ln_mlp_sharded.init(init_rngs, x, deterministic=True)
mlp_out_sharded, ln_out_sharded = ln_mlp_sharded.apply(
params_sharded, x, deterministic=True
params_sharded, x, deterministic=True, rngs={"sr_rng": subkeys[2]}
)
# Make sure params values are the same
......@@ -411,8 +432,8 @@ class TestDistributedLayernormMLP:
rtol = None
l40_tolerance_update = (
get_min_device_compute_capability() == 89
and fp8_recipe == recipe.DelayedScaling()
and use_fp8
and quantization_recipe.delayed()
and dtype == jnp.float16
and activation_type == ("gelu",)
)
......@@ -430,8 +451,8 @@ class TestDistributedLayernormMLP:
# within tolerance to the float32 ground truth.
jax_triton_gemm_precision_tolerance_update = (
with_jax_gemm
and fp8_recipe is not None
and (fp8_recipe.delayed() or fp8_recipe.float8_current_scaling())
and quantization_recipe is not None
and (quantization_recipe.delayed() or quantization_recipe.float8_current_scaling())
and dtype in (jnp.bfloat16, jnp.float16)
and activation_type == ("gelu", "linear"),
)
......@@ -457,22 +478,30 @@ class TestDistributedLayernormMLP:
input_shape,
dtype,
use_fp8=False,
fp8_recipe=None,
quantization_recipe=None,
use_shardy=False,
with_jax_gemm=with_jax_gemm,
)
@pytest.mark.skipif(not is_fp8_supported, reason=reason)
@pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs())
@pytest_parametrize_wrapper("activation_type", [("gelu",), ("gelu", "linear")])
@pytest_parametrize_wrapper("use_bias", [True, False])
@pytest_parametrize_wrapper("input_shape", INPUT_SHAPE)
@pytest_parametrize_wrapper("dtype", DTYPES)
@pytest_parametrize_wrapper("fp8_recipe", SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("quantization_recipe", SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("with_jax_gemm", [False, True])
def test_layernorm_mlp_layer_fp8(
self, mesh_config, activation_type, use_bias, input_shape, dtype, fp8_recipe, with_jax_gemm
self,
mesh_config,
activation_type,
use_bias,
input_shape,
dtype,
quantization_recipe,
with_jax_gemm,
):
if dtype == jnp.float16 and quantization_recipe is not None and quantization_recipe.nvfp4():
pytest.skip("NVFP4 GEMM + Float16 output is unsupported!")
self._test_layernorm_mlp(
mesh_config,
activation_type,
......@@ -480,7 +509,7 @@ class TestDistributedLayernormMLP:
input_shape,
dtype,
use_fp8=True,
fp8_recipe=fp8_recipe,
quantization_recipe=quantization_recipe,
use_shardy=False,
with_jax_gemm=with_jax_gemm,
)
......@@ -501,24 +530,30 @@ class TestDistributedLayernormMLP:
input_shape,
dtype,
use_fp8=False,
fp8_recipe=None,
quantization_recipe=None,
use_shardy=True,
with_jax_gemm=with_jax_gemm,
)
@pytest.mark.skipif(not is_fp8_supported, reason=reason)
@pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs())
@pytest_parametrize_wrapper("activation_type", [("gelu",), ("gelu", "linear")])
@pytest_parametrize_wrapper("use_bias", [True, False])
@pytest_parametrize_wrapper("input_shape", INPUT_SHAPE)
@pytest_parametrize_wrapper("dtype", DTYPES)
@pytest_parametrize_wrapper("fp8_recipe", SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("quantization_recipe", SUPPORTED_RECIPES)
@pytest_parametrize_wrapper("with_jax_gemm", [False, True])
def test_layernorm_mlp_layer_fp8_shardy(
self, mesh_config, activation_type, use_bias, input_shape, dtype, fp8_recipe, with_jax_gemm
self,
mesh_config,
activation_type,
use_bias,
input_shape,
dtype,
quantization_recipe,
with_jax_gemm,
):
if with_jax_gemm and isinstance(fp8_recipe, recipe.MXFP8BlockScaling):
pytest.skip("`jax.nn.scaled_matmul()` does not support the Shardy partitioner.")
if dtype == jnp.float16 and quantization_recipe is not None and quantization_recipe.nvfp4():
pytest.skip("NVFP4 GEMM + Float16 output is unsupported!")
self._test_layernorm_mlp(
mesh_config,
activation_type,
......@@ -526,7 +561,7 @@ class TestDistributedLayernormMLP:
input_shape,
dtype,
use_fp8=True,
fp8_recipe=fp8_recipe,
quantization_recipe=quantization_recipe,
use_shardy=True,
with_jax_gemm=with_jax_gemm,
)
tests/jax/test_helper.py
View file @ 8a7ab3dd
......@@ -10,20 +10,27 @@ import jax.numpy as jnp
import numpy as np
from utils import assert_allclose
from transformer_engine.common.recipe import DelayedScaling, MXFP8BlockScaling, Float8CurrentScaling
from transformer_engine.common.recipe import (
DelayedScaling,
MXFP8BlockScaling,
Float8CurrentScaling,
NVFP4BlockScaling,
)
from transformer_engine.common.recipe import Format as FP8Format
from transformer_engine.jax import fp8_autocast, get_delayed_scaling
from transformer_engine.jax import fp8_autocast
from transformer_engine.jax.quantize import (
get_quantize_config,
is_fp8_available,
is_scaling_mode_supported,
ScalingMode,
update_collections,
TensorSource,
)
from transformer_engine.jax.quantize.helper import _format2dtypes
from transformer_engine.jax.sharding import MeshResource, global_mesh_resource
is_fp8_supported, reason = is_fp8_available()
is_mxfp8_supported, mxfp8_reason = is_fp8_available(ScalingMode.MXFP8_1D_SCALING)
is_fp8_supported, reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING)
is_mxfp8_supported, mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING)
is_nvfp4_supported, nvfp4_reason = is_scaling_mode_supported(ScalingMode.NVFP4_1D_SCALING)
class TestHelper(unittest.TestCase):
......@@ -52,14 +59,16 @@ class TestFP8Functions(unittest.TestCase):
def _check_default_state(self):
self.assertFalse(get_quantize_config().is_fp8_enabled())
def _compare_delay_scaling(self, ref, test):
self.assertTrue(ref.margin == test.margin)
self.assertTrue(ref.fp8_format == test.fp8_format)
self.assertTrue(ref.amax_history_len == test.amax_history_len)
self.assertTrue(ref.amax_compute_algo == test.amax_compute_algo)
def _compare_delay_scaling(self, test):
self.assertEqual(get_quantize_config().MARGIN, test.margin)
self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp8_format)[0])
self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp8_format)[1])
self.assertEqual(get_quantize_config().AMAX_HISTORY_LEN, test.amax_history_len)
self.assertEqual(get_quantize_config().AMAX_COMPUTE_ALGO.value, test.amax_compute_algo)
def _compare_current_scaling(self, test):
self.assertEqual(get_quantize_config().FP8_FORMAT, test.fp8_format)
self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp8_format)[0])
self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp8_format)[1])
for tensor_source in TensorSource:
self.assertEqual(
get_quantize_config().get_scaling_mode(tensor_source),
......@@ -67,13 +76,26 @@ class TestFP8Functions(unittest.TestCase):
)
def _compare_mxfp8_scaling(self, test):
self.assertEqual(get_quantize_config().MARGIN, test.margin)
self.assertEqual(get_quantize_config().FP8_FORMAT, test.fp8_format)
self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp8_format)[0])
self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp8_format)[1])
for tensor_source in TensorSource:
self.assertEqual(
get_quantize_config().get_scaling_mode(tensor_source), ScalingMode.MXFP8_1D_SCALING
)
def _compare_nvfp4_scaling(self, test):
self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp4_format)[0])
self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp4_format)[1])
for tensor_source in TensorSource:
target_scaling_mode = (
ScalingMode.NVFP4_2D_SCALING
if tensor_source == TensorSource.KERNEL
else ScalingMode.NVFP4_1D_SCALING
)
self.assertEqual(
get_quantize_config().get_scaling_mode(tensor_source), target_scaling_mode
)
@unittest.skipIf(not is_fp8_supported, reason=reason)
def test_fp8_autocast_delayed_scaling(self):
self._check_default_state()
......@@ -86,14 +108,14 @@ class TestFP8Functions(unittest.TestCase):
ds = DelayedScaling(margin=5.0, fp8_format=FP8Format.E4M3, amax_history_len=1)
with fp8_autocast(enabled=True, fp8_recipe=ds, mesh_resource=MeshResource()):
self.assertTrue(get_quantize_config().is_fp8_enabled())
self._compare_delay_scaling(get_delayed_scaling(), ds)
self._compare_delay_scaling(ds)
self._check_default_state()
ds = DelayedScaling(margin=3.0, fp8_format=FP8Format.HYBRID, amax_history_len=1)
with fp8_autocast(enabled=True, fp8_recipe=ds, mesh_resource=MeshResource()):
self.assertTrue(get_quantize_config().is_fp8_enabled())
self._compare_delay_scaling(get_delayed_scaling(), ds)
self._compare_delay_scaling(ds)
self._check_default_state()
......@@ -133,16 +155,27 @@ class TestFP8Functions(unittest.TestCase):
self._check_default_state()
bs = MXFP8BlockScaling(margin=5.0, fp8_format=FP8Format.E4M3)
bs = MXFP8BlockScaling()
with fp8_autocast(enabled=True, fp8_recipe=bs, mesh_resource=MeshResource()):
self.assertTrue(get_quantize_config().is_fp8_enabled())
self._compare_mxfp8_scaling(bs)
self._check_default_state()
bs = MXFP8BlockScaling(margin=3.0, fp8_format=FP8Format.HYBRID)
@unittest.skipIf(not is_nvfp4_supported, reason=nvfp4_reason)
def test_fp8_autocast_nvfp4_block_scaling(self):
self._check_default_state()
with fp8_autocast(
enabled=False, fp8_recipe=NVFP4BlockScaling(), mesh_resource=MeshResource()
):
self._check_default_state()
self._check_default_state()
bs = NVFP4BlockScaling()
with fp8_autocast(enabled=True, fp8_recipe=bs, mesh_resource=MeshResource()):
self.assertTrue(get_quantize_config().is_fp8_enabled())
self._compare_mxfp8_scaling(bs)
self._compare_nvfp4_scaling(bs)
self._check_default_state()
tests/jax/utils.py
View file @ 8a7ab3dd
......@@ -1544,6 +1544,12 @@ def dtype_tols(
rtol = eps_relaxed
if atol is None:
atol = max(ulp, eps_relaxed)
# Manually set tols for nvfp4
if dtype == jnp.float4_e2m1fn:
atol = 0.05
rtol = 0.1
return {"rtol": rtol, "atol": atol}
......
transformer_engine/jax/__init__.py
View file @ 8a7ab3dd
......@@ -34,7 +34,7 @@ load_framework_extension("jax")
from . import flax
from . import quantize
from .quantize import fp8_autocast, update_collections, get_delayed_scaling
from .quantize import fp8_autocast, update_collections
from .quantize import NVTE_FP8_COLLECTION_NAME
from .sharding import MeshResource
......@@ -47,7 +47,6 @@ __all__ = [
"NVTE_FP8_COLLECTION_NAME",
"fp8_autocast",
"update_collections",
"get_delayed_scaling",
"MeshResource",
"flax",
"quantize",
......
transformer_engine/jax/cpp_extensions/__init__.py
View file @ 8a7ab3dd
......@@ -3,6 +3,7 @@
# See LICENSE for license information.
"""Python interface for c++ extensions"""
from .activation import *
from .amax import *
from .attention import *
from .normalization import *
from .quantization import *
......
transformer_engine/jax/cpp_extensions/activation.py
View file @ 8a7ab3dd
......@@ -1314,7 +1314,10 @@ def act_lu(
)
return out
if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
if (
quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING
or quantizer.scaling_mode.is_nvfp4_scaling
):
# Current scaling does not support fused operations. Perform dact in higher precision then quantize after.
out = act_lu(
x=x,
......@@ -1488,7 +1491,10 @@ def quantize_dact_dbias(
if war_output is not None:
return war_output
if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
if (
quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING
or quantizer.scaling_mode.is_nvfp4_scaling
):
# Current scaling does not support fused operations. Perform dact in higher precision then quantize after.
out = dact_lu(
dz=dz,
......
transformer_engine/jax/cpp_extensions/amax.py 0 → 100644
View file @ 8a7ab3dd
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""JAX/TE custom ops for amax calculation"""
from enum import Enum
import jax
import jax.numpy as jnp
from jax import dtypes, ffi
from jax.experimental.custom_partitioning import SdyShardingRule
from jax.sharding import PartitionSpec
from .base import BasePrimitive, register_primitive
from .misc import (
get_padded_spec,
NamedSharding,
)
from ..sharding import (
global_mesh_resource,
lax_paral_op,
)
from ..quantize import (
get_wgrad_sign_vector,
get_sign_from_vector,
)
__all__ = ["AmaxScope", "calculate_amax", "calculate_post_rht_amax"]
class AmaxScope(Enum):
"""
Amax Scope Enum
"""
LOCAL = 1
TPSP = 2
FSDP = 3
def all_reduce_amax_along_TPSP_and_FSDP(self, amax, data_spec, transpose_batch_sequence, mesh):
"""Reduce the amax based on its scope"""
gmesh = global_mesh_resource()
sequence_dim = 0 if transpose_batch_sequence else 1
# Run AR across TPSP only when tensor-sequence is detected in the input spec
if self is AmaxScope.TPSP and data_spec[sequence_dim] == gmesh.tpsp_resource:
return lax_paral_op(amax, jax.lax.pmax, gmesh.tpsp_resource, mesh)
# Run AR across FSDP
if self is AmaxScope.FSDP:
return lax_paral_op(amax, jax.lax.pmax, gmesh.fsdp_resource, mesh)
return amax
class AmaxCalculationPrimitive(BasePrimitive):
"""
Amax Calculation Primitive with custom_partitioning
"""
name = "jax_local_amax"
multiple_results = False
impl_static_args = (
1,
2,
) # amax_scope, transpose_batch_sequence
inner_primitive = None
outer_primitive = None
@staticmethod
def abstract(
x_aval,
*,
amax_scope,
transpose_batch_sequence,
):
"""
amax calcuation abstract
"""
del amax_scope, transpose_batch_sequence
dtype = dtypes.canonicalize_dtype(x_aval.dtype)
assert dtype in [jnp.float32, jnp.float16, jnp.bfloat16]
out_aval = jax.core.ShapedArray(shape=(1,), dtype=jnp.float32)
return out_aval
@staticmethod
def impl(
x,
amax_scope,
transpose_batch_sequence,
):
"""
amax calcuation implementation
"""
del amax_scope, transpose_batch_sequence
amax = jnp.amax(jnp.abs(x), keepdims=True).astype(jnp.float32).reshape((1,))
return amax
@staticmethod
def infer_sharding_from_operands(
amax_scope,
transpose_batch_sequence,
mesh,
arg_infos,
result_infos,
):
"""
amax calcuation infer_sharding_from_operands
"""
del (amax_scope, transpose_batch_sequence, arg_infos, result_infos) # Unused.
amax_sharding = NamedSharding(
mesh,
PartitionSpec(None),
desc="AmaxCalculationPrimitive.out_sharding",
)
return amax_sharding
@staticmethod
def partition(
amax_scope,
transpose_batch_sequence,
mesh,
arg_infos,
result_infos,
):
"""
amax calcuation partition
"""
del result_infos
x_spec = get_padded_spec(arg_infos[0])
amax_sharding = NamedSharding(
mesh,
PartitionSpec(None),
desc="AmaxCalculation.amax_sharding",
)
def sharded_impl(x):
amax = AmaxCalculationPrimitive.impl(
x,
amax_scope=amax_scope,
transpose_batch_sequence=transpose_batch_sequence,
)
amax = amax_scope.all_reduce_amax_along_TPSP_and_FSDP(
amax, x_spec, transpose_batch_sequence, mesh
)
return amax
arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos)
return mesh, sharded_impl, amax_sharding, arg_shardings
@staticmethod
def shardy_sharding_rule(amax_scope, transpose_batch_sequence, mesh, value_types, result_types):
"""
amax calcuation shardy_sharding_rule
"""
del amax_scope, transpose_batch_sequence, mesh, result_types
prefix = "AmaxCal"
input_spec = tuple(f"{prefix}_{i}" for i in range(len(value_types[0].shape)))
output_spec = (f"{prefix}_amax",)
return SdyShardingRule((input_spec,), (output_spec,))
register_primitive(AmaxCalculationPrimitive, outer_only=True)
class RHTAmaxCalculationPrimitive(BasePrimitive):
"""
Amax Calculation Primitive with custom_partitioning for calculating regular and post-Random Hadamard Transform (RHT) amax using TE's fused kernels.
"""
name = "te_rht_amax_ffi"
multiple_results = True
impl_static_args = (
1, # amax_scope
2, # transpose_batch_sequence
3, # rht_matrix_random_sign_mask_t
4, # produce_regular_amax
5, # flatten_axis
)
inner_primitive = None
outer_primitive = None
@staticmethod
def abstract(
x_aval,
*,
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
):
"""
amax calcuation abstract
"""
del (
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
)
dtype = dtypes.canonicalize_dtype(x_aval.dtype)
assert dtype in [jnp.bfloat16], f"RHT requires input to be bfloat16, but got {dtype}"
amax_aval = jax.core.ShapedArray(shape=(1,), dtype=jnp.float32)
post_rht_amax_aval = jax.core.ShapedArray(shape=(1,), dtype=jnp.float32)
return amax_aval, post_rht_amax_aval
@staticmethod
def lowering(
ctx,
x,
*,
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
):
"""
te_dbias_quantize_p lowering rules
"""
del amax_scope, transpose_batch_sequence
(x_aval,) = ctx.avals_in
assert x_aval.dtype in [jnp.float32, jnp.float16, jnp.bfloat16]
flatten_axis = flatten_axis if flatten_axis >= 0 else flatten_axis + len(x_aval.shape)
assert 0 < flatten_axis < len(x_aval.shape), "Flatten axis out of bounds!"
return ffi.ffi_lowering(
RHTAmaxCalculationPrimitive.name,
)(
ctx,
x,
rht_matrix_random_sign_mask_t=rht_matrix_random_sign_mask_t,
produce_regular_amax=produce_regular_amax,
flatten_axis=flatten_axis,
)
@staticmethod
def impl(
x,
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
):
"""
amax calcuation implementation
"""
assert RHTAmaxCalculationPrimitive.inner_primitive is not None
(
amax,
post_rht_amax,
) = RHTAmaxCalculationPrimitive.inner_primitive.bind(
x,
amax_scope=amax_scope,
transpose_batch_sequence=transpose_batch_sequence,
rht_matrix_random_sign_mask_t=rht_matrix_random_sign_mask_t,
produce_regular_amax=produce_regular_amax,
flatten_axis=flatten_axis,
)
return amax, post_rht_amax
@staticmethod
def infer_sharding_from_operands(
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
mesh,
arg_infos,
result_infos,
):
"""
amax calcuation infer_sharding_from_operands
"""
del (
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
arg_infos,
result_infos,
) # Unused.
amax_sharding = NamedSharding(
mesh,
PartitionSpec(None),
desc="RHTAmaxCalculationPrimitive.out_sharding",
)
return amax_sharding, amax_sharding
@staticmethod
def partition(
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
mesh,
arg_infos,
result_infos,
):
"""
amax calcuation partition
"""
del result_infos
x_spec = get_padded_spec(arg_infos[0])
amax_sharding = NamedSharding(
mesh,
PartitionSpec(None),
desc="RHTAmaxCalculationPrimitive.amax_sharding",
)
out_shardings = (amax_sharding, amax_sharding)
def sharded_impl(x):
amax, post_rht_amax = RHTAmaxCalculationPrimitive.impl(
x,
amax_scope=amax_scope,
transpose_batch_sequence=transpose_batch_sequence,
rht_matrix_random_sign_mask_t=rht_matrix_random_sign_mask_t,
produce_regular_amax=produce_regular_amax,
flatten_axis=flatten_axis,
)
amax = amax_scope.all_reduce_amax_along_TPSP_and_FSDP(
amax, x_spec, transpose_batch_sequence, mesh
)
post_rht_amax = amax_scope.all_reduce_amax_along_TPSP_and_FSDP(
post_rht_amax, x_spec, transpose_batch_sequence, mesh
)
return amax, post_rht_amax
arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos)
return mesh, sharded_impl, out_shardings, arg_shardings
@staticmethod
def shardy_sharding_rule(
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
mesh,
value_types,
result_types,
):
"""
amax calcuation shardy_sharding_rule
"""
del (
amax_scope,
transpose_batch_sequence,
rht_matrix_random_sign_mask_t,
produce_regular_amax,
flatten_axis,
mesh,
result_types,
)
prefix = "RHTAmaxCal"
input_spec = tuple(f"{prefix}_{i}" for i in range(len(value_types[0].shape)))
output_amax_spec = (f"{prefix}_amax",)
output_post_rht_amax_spec = (f"{prefix}_post_rht_amax",)
return SdyShardingRule((input_spec,), (output_amax_spec, output_post_rht_amax_spec))
register_primitive(RHTAmaxCalculationPrimitive)
def calculate_amax(x: jnp.ndarray, amax_scope: AmaxScope, transpose_batch_sequence: bool):
"""
Compute the maximum absolute value (amax) of the input tensor.
"""
assert AmaxCalculationPrimitive.outer_primitive is not None
return AmaxCalculationPrimitive.outer_primitive.bind(
x,
amax_scope=amax_scope,
transpose_batch_sequence=transpose_batch_sequence,
)
def calculate_post_rht_amax(
x: jnp.ndarray,
amax_scope: AmaxScope,
transpose_batch_sequence: bool,
produce_regular_amax: bool,
flatten_axis: int,
):
"""Compute the post-Random Hadamard Transform (RHT) amax of the input tensor, and optionally the regular amax.
Args:
x: Input tensor.
amax_scope: The scope for amax reduction (local, TPSP, or FSDP).
transpose_batch_sequence: Whether the input tensor has its batch and sequence dimensions transposed.
produce_regular_amax: Whether to compute and return the regular amax alongside the post-RHT amax.
flatten_axis: The axis at which to flatten the input tensor before applying RHT.
Returns:
A tuple containing:
- The regular amax if `produce_regular_amax` is True, otherwise None.
- The post-RHT amax.
"""
amax, post_rht_amax = RHTAmaxCalculationPrimitive.outer_primitive.bind(
x,
amax_scope=amax_scope,
transpose_batch_sequence=transpose_batch_sequence,
rht_matrix_random_sign_mask_t=get_sign_from_vector(get_wgrad_sign_vector()),
produce_regular_amax=produce_regular_amax,
flatten_axis=flatten_axis,
)
if produce_regular_amax:
return amax, post_rht_amax
return None, post_rht_amax
transformer_engine/jax/cpp_extensions/gemm.py
View file @ 8a7ab3dd
......@@ -32,6 +32,7 @@ from ..quantize import (
AbstractBaseTensor,
NoScaleTensor,
ScaledTensor,
ScaledTensor1x,
ScaledTensor2x,
GroupedScaledTensor1x,
ScalingMode,
......@@ -43,6 +44,7 @@ from ..quantize import (
noop_quantizer_set,
is_fp8_gemm_with_all_layouts_supported,
apply_padding_to_scale_inv,
should_use_rht,
)
from .misc import get_padded_spec, is_all_reduce_in_float32
from ..sharding import (
......@@ -138,6 +140,7 @@ def _quantize_gemm_operands(lhs, rhs, lhs_quantizer, rhs_quantizer, contracting_
need_lhs_colwise = lhs_is_transposed and (
lhs_quantizer.scaling_mode.is_1d_block_scaling()
or not is_fp8_gemm_with_all_layouts_supported()
or lhs_quantizer.scaling_mode.is_nvfp4_scaling
)
flatten_axis = max(lhs_cdims) + 1 if lhs_is_transposed else min(lhs_cdims)
lhs_q = lhs_quantizer.quantize(
......@@ -153,6 +156,7 @@ def _quantize_gemm_operands(lhs, rhs, lhs_quantizer, rhs_quantizer, contracting_
need_rhs_colwise = not rhs_is_transposed and (
rhs_quantizer.scaling_mode.is_1d_block_scaling()
or not is_fp8_gemm_with_all_layouts_supported()
or rhs_quantizer.scaling_mode.is_nvfp4_scaling
)
flatten_axis = min(rhs_cdims) if rhs_is_transposed else max(rhs_cdims) + 1
rhs_q = rhs_quantizer.quantize(
......@@ -165,9 +169,27 @@ def _quantize_gemm_operands(lhs, rhs, lhs_quantizer, rhs_quantizer, contracting_
assert not isinstance(lhs_q, ScaledTensor2x)
assert not isinstance(rhs_q, ScaledTensor2x)
def uses_rht(q: AbstractBaseTensor) -> bool:
return isinstance(q, ScaledTensor1x) and should_use_rht(
q.scaling_mode, is_colwise=q.is_colwise
)
# TODO(jberchtold): Move RHT usage check to a bool flag on the ScaledTensor class
assert uses_rht(lhs_q) == uses_rht(rhs_q), (
"With NVFP4_1D_SCALING, if one operand is colwise quantized, the other must be colwise"
" quantized as well. This is to ensure the RHT is applied to both and will cancel out in"
" the GEMM."
)
return lhs_q, rhs_q
def _get_nvfp4_tensor_scale_inv(amax):
DATA_DTYPE_MAX = jnp.finfo(jnp.float4_e2m1fn.dtype).max.astype(jnp.float32)
SCALE_DTYPE_MAX = jnp.finfo(jnp.float8_e4m3fn.dtype).max.astype(jnp.float32)
return amax / (DATA_DTYPE_MAX * SCALE_DTYPE_MAX)
def collective_gemm_bootstrap(
num_total_devices,
num_devices_per_process,
......@@ -345,7 +367,7 @@ class GemmPrimitive(BasePrimitive):
name = "te_gemm_ffi"
multiple_results = True
impl_static_args = 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
impl_static_args = (8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18)
inner_primitive = None
outer_primitive = None
......@@ -357,6 +379,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype,
contracting_dims,
scaling_mode,
......@@ -404,7 +428,9 @@ class GemmPrimitive(BasePrimitive):
lhs_is_transposed, rhs_is_transposed = _get_gemm_layout(operand_ndims, contracting_dims)
if scaling_mode != ScalingMode.NO_SCALING:
assert _compatible_fp8_gemm_dtypes(lhs.dtype, rhs.dtype), (
assert scaling_mode.is_nvfp4_scaling or _compatible_fp8_gemm_dtypes(
lhs.dtype, rhs.dtype
), (
"cuBLAS GEMM quantized operands have incompatible data types: "
f"{lhs.dtype} x {rhs.dtype}."
)
......@@ -484,6 +510,8 @@ class GemmPrimitive(BasePrimitive):
f"expected {pre_gelu_dtype} but found {gelu_input.dtype}."
)
pre_gelu_out = jax.core.ShapedArray(shape=pre_gelu_shape, dtype=pre_gelu_dtype)
assert alpha.size == 1 and alpha.dtype == jnp.float32
assert beta.size == 1 and beta.dtype == jnp.float32
# Declare cuBLAS workspace
workspace_size = get_cublas_workspace_size_bytes()
......@@ -510,6 +538,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype,
contracting_dims,
scaling_mode,
......@@ -530,7 +560,7 @@ class GemmPrimitive(BasePrimitive):
(lhs_aval.ndim, rhs_aval.ndim), (lhs_cdims, rhs_cdims)
)
args = (lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input)
args = (lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input, alpha, beta)
kwargs = {
"scaling_mode": int(scaling_mode.value),
"lhs_axis_boundary": max(lhs_cdims) + 1 if lhs_transposed else min(lhs_cdims),
......@@ -563,6 +593,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype,
contracting_dims,
scaling_mode,
......@@ -626,6 +658,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype=out_dtype,
contracting_dims=contracting_dims,
scaling_mode=scaling_mode,
......@@ -675,6 +709,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype,
contracting_dims,
scaling_mode,
......@@ -694,6 +730,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype,
contracting_dims,
scaling_mode,
......@@ -1001,6 +1039,9 @@ class GemmPrimitive(BasePrimitive):
gelu_input_specs = (None,)
arg_shardings += (NamedSharding(mesh, PartitionSpec(*gelu_input_specs)),)
# Alpha, beta
arg_shardings += (none_sharding, none_sharding)
# Assemble output shardings
out_shardings = [NamedSharding(mesh, PartitionSpec(*out_specs))]
......@@ -1014,7 +1055,7 @@ class GemmPrimitive(BasePrimitive):
pre_gelu_specs = (None,)
out_shardings.append(NamedSharding(mesh, PartitionSpec(*pre_gelu_specs)))
def _sharded_impl(lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input):
def _sharded_impl(lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input, alpha, beta):
# We should not fuse bias in the output reduction case
sharded_fuse_bias = fuse_bias and reduce_spec is None
outputs = GemmPrimitive.impl(
......@@ -1024,6 +1065,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype=out_dtype,
contracting_dims=contracting_dims,
scaling_mode=scaling_mode,
......@@ -1114,8 +1157,10 @@ class GemmPrimitive(BasePrimitive):
rhs_non_cspec = tuple(rhs_specs[i] for i in range(operand_ndims[1]) if i not in rhs_cdims)
out_spec = (*lhs_non_cspec, *rhs_non_cspec)
bias_spec = rhs_non_cspec if fuse_bias else ("…4",)
dbias_spec = bias_spec if grad else ("…5")
gelu_spec = out_spec if fuse_gelu else ("…6",)
gelu_spec = out_spec if fuse_gelu else ("…5",)
alpha_spec = ("_6",)
beta_spec = ("_7",)
dbias_spec = bias_spec if grad else ("…8")
return SdyShardingRule(
operand_mappings=(
......@@ -1125,6 +1170,8 @@ class GemmPrimitive(BasePrimitive):
rhs_scale_specs,
bias_spec,
gelu_spec,
alpha_spec,
beta_spec,
),
result_mappings=(
out_spec,
......@@ -1178,6 +1225,7 @@ def _te_gemm(
# Quantize operands (if necessary)
lhs_q, rhs_q = _quantize_gemm_operands(lhs, rhs, lhs_quantizer, rhs_quantizer, contracting_dims)
lhs_amax = rhs_amax = None
# Extract GEMM custom op inputs from quantized operands
if isinstance(lhs_q, ScaledTensor):
assert isinstance(rhs_q, ScaledTensor) or rhs_quantizer is not None, (
......@@ -1192,6 +1240,7 @@ def _te_gemm(
lhs_scale_inv = lhs_q.scale_inv
if lhs_q.data_layout == "T":
lhs_cdims = transpose_dims(lhs_q.ndim, lhs_cdims, flatten_axis=lhs_q.flatten_axis)
lhs_amax = lhs_q.amax
if isinstance(rhs_q, ScaledTensor):
assert isinstance(lhs_q, ScaledTensor) or lhs_quantizer is not None, (
......@@ -1201,7 +1250,11 @@ def _te_gemm(
if isinstance(rhs_q, ScaledTensor2x):
# Choose the quantization of the contracting dimension(s)
rhs_q = rhs_q.get_rowwise_tensor() if rhs_is_transposed else rhs_q.get_colwise_tensor()
assert rhs_q.scaling_mode == lhs_q.scaling_mode, (
assert (
rhs_q.scaling_mode == lhs_q.scaling_mode
or rhs_q.scaling_mode.is_nvfp4_scaling
and lhs_q.scaling_mode.is_nvfp4_scaling
), (
"cuBLAS GEMM quantized operands have mismatched scaling types, "
f"LHS:{lhs_q.scaling_mode} x RHS:{rhs_q.scaling_mode}."
)
......@@ -1209,6 +1262,15 @@ def _te_gemm(
rhs_scale_inv = rhs_q.scale_inv
if rhs_q.data_layout == "T":
rhs_cdims = transpose_dims(rhs_q.ndim, rhs_cdims, flatten_axis=rhs_q.flatten_axis)
rhs_amax = rhs_q.amax
alpha = jnp.ones((1,), jnp.float32)
beta = jnp.zeros((1,), jnp.float32)
if scaling_mode.is_nvfp4_scaling:
assert lhs_amax is not None and rhs_amax is not None
lhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(lhs_amax)
rhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(rhs_amax)
alpha = lhs_tensor_scale_inv * rhs_tensor_scale_inv
# Dummy empties for bias and gelu
out_dtype = lhs_q.dq_dtype if isinstance(lhs_q, ScaledTensor) else lhs_data.dtype
......@@ -1224,6 +1286,8 @@ def _te_gemm(
rhs_scale_inv,
bias,
gelu_input,
alpha,
beta,
out_dtype=out_dtype,
contracting_dims=(lhs_cdims, rhs_cdims),
scaling_mode=scaling_mode,
......@@ -1514,15 +1578,17 @@ def _jax_gemm_tensor_scaling_fp8(lhs, rhs, dim_nums, precision):
@partial(jax.jit, static_argnums=(2,))
def _jax_gemm_mxfp8_1d(
def _jax_scaled_matmul(
lhs: ScaledTensor, rhs: ScaledTensor, dim_nums: Tuple[Tuple[Sequence[int], Sequence[int]]]
):
"""
JAX GEMM for MXFP8 via scaled_matmul
"""
assert (
rhs.scaling_mode == ScalingMode.MXFP8_1D_SCALING
), "rhs does not have MXFP8 1D scaling mode"
assert rhs.scaling_mode in (
ScalingMode.MXFP8_1D_SCALING,
ScalingMode.NVFP4_1D_SCALING,
ScalingMode.NVFP4_2D_SCALING,
), f"rhs does not have MXFP8 or NVFP4 scaling mode, got rhs.scaling_mode={rhs.scaling_mode}"
(lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dim_nums
......@@ -1537,21 +1603,48 @@ def _jax_gemm_mxfp8_1d(
f" {rhs.is_colwise}"
)
if lhs.scaling_mode == ScalingMode.MXFP8_1D_SCALING:
out_dtype = lhs.dq_dtype
assert (
lhs.data_layout == "N" and rhs.data_layout == "N"
), f"Got lhs.data_layout={lhs.data_layout}, rhs.data_layout={rhs.data_layout}"
else:
if lhs.data_layout == "T":
lhs_contract = transpose_dims(
lhs.data.ndim, lhs_contract, flatten_axis=lhs.flatten_axis
)
if rhs.data_layout == "T":
rhs_contract = transpose_dims(
rhs.data.ndim, rhs_contract, flatten_axis=rhs.flatten_axis
)
out_dtype = jnp.float32
# Reshape + Transpose (if needed)
# [..., M, K] -> [1, reduce(..., M), K]
# [..., K, M] -> [1, reduce(..., M), K]
lhs_3d = _shape_normalization(lhs.data, (lhs_contract, lhs_batch))
rhs_3d = _shape_normalization(rhs.data, (rhs_contract, rhs_batch))
lhs_scale_3d = _shape_normalization(lhs.scale_inv, (lhs_contract, lhs_batch))
rhs_scale_3d = _shape_normalization(rhs.scale_inv, (rhs_contract, rhs_batch))
lhs_3d = _shape_normalization(lhs.data, (lhs_contract, lhs_batch), lhs.data_layout == "T")
rhs_3d = _shape_normalization(rhs.data, (rhs_contract, rhs_batch), rhs.data_layout == "T")
lhs_scale_3d = _shape_normalization(
lhs.scale_inv, (lhs_contract, lhs_batch), lhs.data_layout == "T"
)
rhs_scale_3d = _shape_normalization(
rhs.scale_inv, (rhs_contract, rhs_batch), rhs.data_layout == "T"
)
# JAX scaled_matmul only supports NT now (TN-gemm)
# * Expected shape:
# * lhs_data (B, M, K) * rhs_data (B, N, K)
# * lhs_scale (B, M, K_block) * rhs_scale (B, N, K_block)
out_3d = jax.nn.scaled_matmul(
lhs_3d, rhs_3d, lhs_scale_3d, rhs_scale_3d, preferred_element_type=lhs.dq_dtype
lhs_3d, rhs_3d, lhs_scale_3d, rhs_scale_3d, preferred_element_type=out_dtype
)
if lhs.scaling_mode.is_nvfp4_scaling:
assert lhs.amax is not None and rhs.amax is not None
lhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(lhs.amax)
rhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(rhs.amax)
alpha = lhs_tensor_scale_inv * rhs_tensor_scale_inv
out_3d = (out_3d * alpha).astype(lhs.dq_dtype)
# Reshape [1, reduce(..., M), N] -> [..., M, N]
lhs_remain_shape = tuple(
lhs.data.shape[dim] for dim in range(len(lhs.data.shape)) if dim not in lhs_contract
......@@ -1560,6 +1653,7 @@ def _jax_gemm_mxfp8_1d(
rhs.data.shape[dim] for dim in range(len(rhs.data.shape)) if dim not in rhs_contract
)
out = out_3d.reshape(*lhs_remain_shape, *rhs_remain_shape)
return out
......@@ -1575,7 +1669,7 @@ def _jax_gemm(
"""
dim_nums = (contracting_dims, ((), ()))
def _jax_gemm_fp8_impl(lhs, rhs):
def _jax_gemm_impl(lhs, rhs):
if lhs.scaling_mode.is_tensor_scaling():
assert (
rhs.scaling_mode == lhs.scaling_mode
......@@ -1587,15 +1681,15 @@ def _jax_gemm(
)
return _jax_gemm_tensor_scaling_fp8(lhs, rhs, dim_nums, precision)
if lhs.scaling_mode == ScalingMode.MXFP8_1D_SCALING:
return _jax_gemm_mxfp8_1d(lhs, rhs, dim_nums)
if lhs.scaling_mode.is_1d_block_scaling:
return _jax_scaled_matmul(lhs, rhs, dim_nums)
raise NotImplementedError(f"Unsupported ScalingMode: {lhs.scaling_mode}")
lhs_q, rhs_q = _quantize_gemm_operands(lhs, rhs, lhs_quantizer, rhs_quantizer, contracting_dims)
if isinstance(lhs_q, ScaledTensor) and isinstance(rhs_q, ScaledTensor):
return _jax_gemm_fp8_impl(lhs_q, rhs_q)
return _jax_gemm_impl(lhs_q, rhs_q)
if (
isinstance(lhs, jnp.ndarray)
......
transformer_engine/jax/cpp_extensions/misc.py
View file @ 8a7ab3dd
......@@ -6,8 +6,6 @@
import os
import functools
from typing import Tuple
from importlib.metadata import version as get_pkg_version
from packaging.version import Version as PkgVersion
import numpy as np
......@@ -75,7 +73,8 @@ def jax_dtype_to_te_dtype(jax_dtype):
jnp.int64.dtype: TEDType.kInt64,
jnp.float8_e4m3fn.dtype: TEDType.kFloat8E4M3,
jnp.float8_e5m2.dtype: TEDType.kFloat8E5M2,
jnp.uint8.dtype: TEDType.kByte,
jnp.float8_e8m0fnu.dtype: TEDType.kFloat8E8M0,
jnp.float4_e2m1fn.dtype: TEDType.kFloat4E2M1,
}
if jax_dtype not in converter:
......@@ -151,16 +150,6 @@ def get_cudnn_version() -> Tuple[int, int, int]:
return (major, minor, patch)
@functools.lru_cache(maxsize=None)
def jax_version_meet_requirement(version: str):
"""
Helper function checking if required JAX version is available
"""
jax_version = PkgVersion(get_pkg_version("jax"))
jax_version_required = PkgVersion(version)
return jax_version >= jax_version_required
def get_xla_flag(flag: str, default=None, cast=str):
"""
Returns the value of a flag/option in XLA_FLAGS environment variable if present or returns the default value.
......
transformer_engine/jax/cpp_extensions/normalization.py
View file @ 8a7ab3dd
......@@ -28,7 +28,10 @@ from .misc import (
get_cudnn_version,
)
from .quantization import _quantize_dbias_impl, AmaxScope
from ..sharding import all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp_tpsp
from ..sharding import (
all_reduce_max_along_all_axes_except_PP,
all_reduce_sum_along_dp_fsdp_tpsp,
)
from ..quantize import ScaledTensor, ScaledTensorFactory, NoScaleTensor
from ..quantize import (
Quantizer,
......@@ -1031,7 +1034,10 @@ def layernorm_fwd(
)
return out, mu, rsigma
if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
if (
quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING
or quantizer.scaling_mode.is_nvfp4_scaling
):
# Current scaling does not support fused operations. Perform norm in higher precision then quantize after.
out, mu, rsigma = layernorm_fwd(
x=x,
......@@ -1276,7 +1282,10 @@ def rmsnorm_fwd(
)
return out, rsigma
if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING:
if (
quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING
or quantizer.scaling_mode.is_nvfp4_scaling
):
# Current scaling does not support fused operations. Perform norm in higher precision then quantize after.
out, rsigma = rmsnorm_fwd(
x=x,
......
transformer_engine/jax/csrc/extensions.h
View file @ 8a7ab3dd
......@@ -85,7 +85,7 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedQuantizeHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(DequantizeHandler);
pybind11::tuple GetDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size,
DType in_dtype, DType out_dtype,
DType in_dtype, DType out_dtype, DType scale_dtype,
JAXX_Scaling_Mode scaling_mode,
QuantizeLayout q_layout);
......@@ -138,6 +138,10 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(CollectiveGemmInitHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmD2HGroupSizesHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmHandler);
// Amax
XLA_FFI_DECLARE_HANDLER_SYMBOL(RHTAmaxCalculationInitializeHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(RHTAmaxCalculationHandler);
// Cudnn helpers
XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnHandleInitHandler);
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment