The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
timmoon10 merged 7 commits into from
May 16, 2023
Merged

The Implementation of Praxis's Modules #158

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b75b0b9
Adding JAX/Praxis modules and dependencies.
mingxu1067 May 10, 2023
23410e2
Adding UTs to JAX/Praxis modules.
mingxu1067 May 10, 2023
08d09ab
Remove praxis as a dependency due to not strictly needed
mingxu1067 May 10, 2023
11a8e66
Repalce is_fp8_supported to is_fp8_available
mingxu1067 May 10, 2023
4823c01
Make Praxis as an optional dependency.
mingxu1067 May 15, 2023
c72007e
Adding a workaround to FP8 training on Praxis.
mingxu1067 May 15, 2023
d3b693e
Merge branch 'main' into mingh/praxis_modules_for_review
mingxu1067 May 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion README.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
import jax
import jax.numpy as jnp
import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax
from transformer_engine.common import recipe

BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
model = te.flax.DenseGeneral(features=HIDDEN)
model = te_flax.DenseGeneral(features=HIDDEN)

def loss_fn(params, other_vars, inp):
out = model.apply({'params':params, **other_vars}, inp)
Expand Down
15 changes: 7 additions & 8 deletions examples/jax/encoder/test_model_parallel_encoder.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
def __call__(self, x, mask, disable_dropout=False):
x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

te_Encoder = partial(te.flax.TransformerLayer,
te_Encoder = partial(te_flax.TransformerLayer,
hidden_size=256,
mlp_hidden_size=1024,
num_attention_heads=8,
hidden_dropout=0.1,
attention_dropout=0.1,
dropout_rng_name=DROPOUT_KEY,
layer_type=te.flax.TransformerLayerType.ENCODER,
layer_type=te_flax.TransformerLayerType.ENCODER,
enable_relative_embedding=False,
dtype=jnp.bfloat16)
x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

x = x.reshape(x.shape[0], -1)

x = te.flax.DenseGeneral(features=256,
x = te_flax.DenseGeneral(features=256,
kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
bias_axes=(NAMED_TP_AXIS,),
sharding_type=te.ShardingType.DP_TP_COL,
dtype=jnp.bfloat16)(x)

x = te.flax.DenseGeneral(features=256,
x = te_flax.DenseGeneral(features=256,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
else:
tensor[i] = vocab[word]

seq_len = len(tokens)
if seq_len > max_seq_len:
seq_len = max_seq_len
seq_len = min(len(tokens), max_seq_len)
mask_2d = mask_3d[j]
mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down
15 changes: 7 additions & 8 deletions examples/jax/encoder/test_multigpu_encoder.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
def __call__(self, x, mask, disable_dropout=False):
x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

te_Encoder = partial(te.flax.TransformerLayer,
te_Encoder = partial(te_flax.TransformerLayer,
hidden_size=256,
mlp_hidden_size=1024,
num_attention_heads=8,
hidden_dropout=0.1,
attention_dropout=0.1,
dropout_rng_name=DROPOUT_KEY,
layer_type=te.flax.TransformerLayerType.ENCODER,
layer_type=te_flax.TransformerLayerType.ENCODER,
enable_relative_embedding=False,
dtype=jnp.bfloat16)
x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

x = x.reshape(x.shape[0], -1)

x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
dtype=jnp.bfloat16)(x)

x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
dtype=jnp.bfloat16)(x)

x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
else:
tensor[i] = vocab[word]

seq_len = len(tokens)
if seq_len > max_seq_len:
seq_len = max_seq_len
seq_len = min(len(tokens), max_seq_len)
mask_2d = mask_3d[j]
mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

sharding_rules = te.flax.extend_logical_axis_rules(tuple())
sharding_rules = te_flax.extend_logical_axis_rules(tuple())
params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down
15 changes: 7 additions & 8 deletions examples/jax/encoder/test_multiprocessing_encoder.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
def __call__(self, x, mask, disable_dropout=False):
x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

te_Encoder = partial(te.flax.TransformerLayer,
te_Encoder = partial(te_flax.TransformerLayer,
hidden_size=256,
mlp_hidden_size=1024,
num_attention_heads=8,
hidden_dropout=0.1,
attention_dropout=0.1,
dropout_rng_name=DROPOUT_KEY,
layer_type=te.flax.TransformerLayerType.ENCODER,
layer_type=te_flax.TransformerLayerType.ENCODER,
enable_relative_embedding=False,
dtype=jnp.bfloat16)
x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

x = x.reshape(x.shape[0], -1)

x = te.flax.DenseGeneral(features=256,
x = te_flax.DenseGeneral(features=256,
kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
bias_axes=(NAMED_TP_AXIS,),
sharding_type=te.ShardingType.DP_TP_COL,
dtype=jnp.bfloat16)(x)

x = te.flax.DenseGeneral(features=256,
x = te_flax.DenseGeneral(features=256,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
else:
tensor[i] = vocab[word]

seq_len = len(tokens)
if seq_len > max_seq_len:
seq_len = max_seq_len
seq_len = min(len(tokens), max_seq_len)
mask_2d = mask_3d[j]
mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down
13 changes: 6 additions & 7 deletions examples/jax/encoder/test_single_gpu_encoder.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
def __call__(self, x, mask, disable_dropout=False):
x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

te_Encoder = partial(te.flax.TransformerLayer,
te_Encoder = partial(te_flax.TransformerLayer,
hidden_size=256,
mlp_hidden_size=1024,
num_attention_heads=8,
hidden_dropout=0.1,
attention_dropout=0.1,
dropout_rng_name=DROPOUT_KEY,
layer_type=te.flax.TransformerLayerType.ENCODER,
layer_type=te_flax.TransformerLayerType.ENCODER,
enable_relative_embedding=False,
dtype=jnp.bfloat16)
x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

x = x.reshape(x.shape[0], -1)

x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
else:
tensor[i] = vocab[word]

seq_len = len(tokens)
if seq_len > max_seq_len:
seq_len = max_seq_len
seq_len = min(len(tokens), max_seq_len)
mask_2d = mask_3d[j]
mask_2d[:seq_len, :seq_len] = 0

Expand Down
3 changes: 2 additions & 1 deletion examples/jax/mnist/test_single_gpu_mnist.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
@nn.compact
def __call__(self, x, disable_dropout=False):
if self.use_te:
nn_Dense = te.flax.DenseGeneral
nn_Dense = te_flax.DenseGeneral
else:
nn_Dense = nn.Dense

Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

timmoon10
merged 7 commits into

from

May 16, 2023
 Merged


The Implementation of Praxis's Modules
#158
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
b75b0b9

Adding JAX/Praxis modules and dependencies.


mingxu1067 May 10, 2023
23410e2

Adding UTs to JAX/Praxis modules.


mingxu1067 May 10, 2023
08d09ab

Remove praxis as a dependency due to not strictly needed


mingxu1067 May 10, 2023
11a8e66

Repalce is_fp8_supported to is_fp8_available


mingxu1067 May 10, 2023
4823c01

Make Praxis as an optional dependency.


mingxu1067 May 15, 2023
c72007e

Adding a workaround to FP8 training on Praxis.


mingxu1067 May 15, 2023
d3b693e

Merge branch 'main' into mingh/praxis_modules_for_review


mingxu1067 May 15, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
 import jax
 import jax.numpy as jnp
 import transformer_engine.jax as te
 import transformer_engine.jax.flax as te_flax
 from transformer_engine.common import recipe

 BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.flax.DenseGeneral(features=HIDDEN)
 model = te_flax.DenseGeneral(features=HIDDEN)

 def loss_fn(params, other_vars, inp):
 out = model.apply({'params':params, **other_vars}, inp)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
 masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 sharding_rules = te.flax.extend_logical_axis_rules(tuple())
 sharding_rules = te_flax.extend_logical_axis_rules(tuple())
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multiprocessing_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

13 changes: 6 additions & 7 deletions
examples/jax/encoder/test_single_gpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand Down

3 changes: 2 additions & 1 deletion
examples/jax/mnist/test_single_gpu_mnist.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.flax.DenseGeneral
 nn_Dense = te_flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

timmoon10
merged 7 commits into

from

May 16, 2023
 Merged


The Implementation of Praxis's Modules
#158
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
b75b0b9

Adding JAX/Praxis modules and dependencies.


mingxu1067 May 10, 2023
23410e2

Adding UTs to JAX/Praxis modules.


mingxu1067 May 10, 2023
08d09ab

Remove praxis as a dependency due to not strictly needed


mingxu1067 May 10, 2023
11a8e66

Repalce is_fp8_supported to is_fp8_available


mingxu1067 May 10, 2023
4823c01

Make Praxis as an optional dependency.


mingxu1067 May 15, 2023
c72007e

Adding a workaround to FP8 training on Praxis.


mingxu1067 May 15, 2023
d3b693e

Merge branch 'main' into mingh/praxis_modules_for_review


mingxu1067 May 15, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
 import jax
 import jax.numpy as jnp
 import transformer_engine.jax as te
 import transformer_engine.jax.flax as te_flax
 from transformer_engine.common import recipe

 BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.flax.DenseGeneral(features=HIDDEN)
 model = te_flax.DenseGeneral(features=HIDDEN)

 def loss_fn(params, other_vars, inp):
 out = model.apply({'params':params, **other_vars}, inp)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
 masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 sharding_rules = te.flax.extend_logical_axis_rules(tuple())
 sharding_rules = te_flax.extend_logical_axis_rules(tuple())
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multiprocessing_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

13 changes: 6 additions & 7 deletions
examples/jax/encoder/test_single_gpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand Down

3 changes: 2 additions & 1 deletion
examples/jax/mnist/test_single_gpu_mnist.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.flax.DenseGeneral
 nn_Dense = te_flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

timmoon10
merged 7 commits into

from

May 16, 2023
 Merged


The Implementation of Praxis's Modules
#158
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
b75b0b9

Adding JAX/Praxis modules and dependencies.


mingxu1067 May 10, 2023
23410e2

Adding UTs to JAX/Praxis modules.


mingxu1067 May 10, 2023
08d09ab

Remove praxis as a dependency due to not strictly needed


mingxu1067 May 10, 2023
11a8e66

Repalce is_fp8_supported to is_fp8_available


mingxu1067 May 10, 2023
4823c01

Make Praxis as an optional dependency.


mingxu1067 May 15, 2023
c72007e

Adding a workaround to FP8 training on Praxis.


mingxu1067 May 15, 2023
d3b693e

Merge branch 'main' into mingh/praxis_modules_for_review


mingxu1067 May 15, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
 import jax
 import jax.numpy as jnp
 import transformer_engine.jax as te
 import transformer_engine.jax.flax as te_flax
 from transformer_engine.common import recipe

 BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.flax.DenseGeneral(features=HIDDEN)
 model = te_flax.DenseGeneral(features=HIDDEN)

 def loss_fn(params, other_vars, inp):
 out = model.apply({'params':params, **other_vars}, inp)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
 masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 sharding_rules = te.flax.extend_logical_axis_rules(tuple())
 sharding_rules = te_flax.extend_logical_axis_rules(tuple())
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multiprocessing_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

13 changes: 6 additions & 7 deletions
examples/jax/encoder/test_single_gpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand Down

3 changes: 2 additions & 1 deletion
examples/jax/mnist/test_single_gpu_mnist.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.flax.DenseGeneral
 nn_Dense = te_flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

timmoon10
merged 7 commits into

from

May 16, 2023
 Merged


The Implementation of Praxis's Modules
#158
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
b75b0b9

Adding JAX/Praxis modules and dependencies.


mingxu1067 May 10, 2023
23410e2

Adding UTs to JAX/Praxis modules.


mingxu1067 May 10, 2023
08d09ab

Remove praxis as a dependency due to not strictly needed


mingxu1067 May 10, 2023
11a8e66

Repalce is_fp8_supported to is_fp8_available


mingxu1067 May 10, 2023
4823c01

Make Praxis as an optional dependency.


mingxu1067 May 15, 2023
c72007e

Adding a workaround to FP8 training on Praxis.


mingxu1067 May 15, 2023
d3b693e

Merge branch 'main' into mingh/praxis_modules_for_review


mingxu1067 May 15, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
 import jax
 import jax.numpy as jnp
 import transformer_engine.jax as te
 import transformer_engine.jax.flax as te_flax
 from transformer_engine.common import recipe

 BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.flax.DenseGeneral(features=HIDDEN)
 model = te_flax.DenseGeneral(features=HIDDEN)

 def loss_fn(params, other_vars, inp):
 out = model.apply({'params':params, **other_vars}, inp)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
 masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 sharding_rules = te.flax.extend_logical_axis_rules(tuple())
 sharding_rules = te_flax.extend_logical_axis_rules(tuple())
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multiprocessing_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

13 changes: 6 additions & 7 deletions
examples/jax/encoder/test_single_gpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand Down

3 changes: 2 additions & 1 deletion
examples/jax/mnist/test_single_gpu_mnist.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.flax.DenseGeneral
 nn_Dense = te_flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

timmoon10
merged 7 commits into

from

May 16, 2023
 Merged


The Implementation of Praxis's Modules
#158
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
b75b0b9

Adding JAX/Praxis modules and dependencies.


mingxu1067 May 10, 2023
23410e2

Adding UTs to JAX/Praxis modules.


mingxu1067 May 10, 2023
08d09ab

Remove praxis as a dependency due to not strictly needed


mingxu1067 May 10, 2023
11a8e66

Repalce is_fp8_supported to is_fp8_available


mingxu1067 May 10, 2023
4823c01

Make Praxis as an optional dependency.


mingxu1067 May 15, 2023
c72007e

Adding a workaround to FP8 training on Praxis.


mingxu1067 May 15, 2023
d3b693e

Merge branch 'main' into mingh/praxis_modules_for_review


mingxu1067 May 15, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
 import jax
 import jax.numpy as jnp
 import transformer_engine.jax as te
 import transformer_engine.jax.flax as te_flax
 from transformer_engine.common import recipe

 BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.flax.DenseGeneral(features=HIDDEN)
 model = te_flax.DenseGeneral(features=HIDDEN)

 def loss_fn(params, other_vars, inp):
 out = model.apply({'params':params, **other_vars}, inp)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
 masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 sharding_rules = te.flax.extend_logical_axis_rules(tuple())
 sharding_rules = te_flax.extend_logical_axis_rules(tuple())
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multiprocessing_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

13 changes: 6 additions & 7 deletions
examples/jax/encoder/test_single_gpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand Down

3 changes: 2 additions & 1 deletion
examples/jax/mnist/test_single_gpu_mnist.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.flax.DenseGeneral
 nn_Dense = te_flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

timmoon10
merged 7 commits into

from

May 16, 2023
 Merged


The Implementation of Praxis's Modules
#158
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
b75b0b9

Adding JAX/Praxis modules and dependencies.


mingxu1067 May 10, 2023
23410e2

Adding UTs to JAX/Praxis modules.


mingxu1067 May 10, 2023
08d09ab

Remove praxis as a dependency due to not strictly needed


mingxu1067 May 10, 2023
11a8e66

Repalce is_fp8_supported to is_fp8_available


mingxu1067 May 10, 2023
4823c01

Make Praxis as an optional dependency.


mingxu1067 May 15, 2023
c72007e

Adding a workaround to FP8 training on Praxis.


mingxu1067 May 15, 2023
d3b693e

Merge branch 'main' into mingh/praxis_modules_for_review


mingxu1067 May 15, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
 import jax
 import jax.numpy as jnp
 import transformer_engine.jax as te
 import transformer_engine.jax.flax as te_flax
 from transformer_engine.common import recipe

 BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.flax.DenseGeneral(features=HIDDEN)
 model = te_flax.DenseGeneral(features=HIDDEN)

 def loss_fn(params, other_vars, inp):
 out = model.apply({'params':params, **other_vars}, inp)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
 masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 sharding_rules = te.flax.extend_logical_axis_rules(tuple())
 sharding_rules = te_flax.extend_logical_axis_rules(tuple())
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multiprocessing_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

13 changes: 6 additions & 7 deletions
examples/jax/encoder/test_single_gpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand Down

3 changes: 2 additions & 1 deletion
examples/jax/mnist/test_single_gpu_mnist.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.flax.DenseGeneral
 nn_Dense = te_flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
The Implementation of Praxis's Modules#158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

timmoon10
merged 7 commits into

from

May 16, 2023
 Merged


The Implementation of Praxis's Modules
#158
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
b75b0b9

Adding JAX/Praxis modules and dependencies.


mingxu1067 May 10, 2023
23410e2

Adding UTs to JAX/Praxis modules.


mingxu1067 May 10, 2023
08d09ab

Remove praxis as a dependency due to not strictly needed


mingxu1067 May 10, 2023
11a8e66

Repalce is_fp8_supported to is_fp8_available


mingxu1067 May 10, 2023
4823c01

Make Praxis as an optional dependency.


mingxu1067 May 15, 2023
c72007e

Adding a workaround to FP8 training on Praxis.


mingxu1067 May 15, 2023
d3b693e

Merge branch 'main' into mingh/praxis_modules_for_review


mingxu1067 May 15, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


3 changes: 2 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -77,6 +77,7 @@ Flax
 import jax
 import jax.numpy as jnp
 import transformer_engine.jax as te
 import transformer_engine.jax.flax as te_flax
 from transformer_engine.common import recipe

 BATCH = 32
Expand All@@ -93,7 +94,7 @@ Flax

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.flax.DenseGeneral(features=HIDDEN)
 model = te_flax.DenseGeneral(features=HIDDEN)

 def loss_fn(params, other_vars, inp):
 out = model.apply({'params':params, **other_vars}, inp)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
DEVICE_TP_AXIS = 'model'
Expand All@@ -39,27 +40,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -174,9 +175,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -275,7 +274,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,6 +20,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

DEVICE_DP_AXIS = 'data'
PARAMS_KEY = 'params'
Expand All@@ -36,24 +37,24 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 x = te_flax.DenseGeneral(features=256, sharding_type=te.ShardingType.DP,
 dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
Expand DownExpand Up@@ -165,9 +166,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -257,7 +256,7 @@ def train_and_evaluate(args):
 masks = jnp.zeros(mask_shape, dtype=jnp.uint8)
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 sharding_rules = te.flax.extend_logical_axis_rules(tuple())
 sharding_rules = te_flax.extend_logical_axis_rules(tuple())
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

15 changes: 7 additions & 8 deletions
examples/jax/encoder/test_multiprocessing_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,6 +22,7 @@
from jax.experimental.pjit import pjit

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
DEVICE_DP_AXIS = 'data'
Expand All@@ -42,27 +43,27 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_BROADCAST_AXIS, NAMED_TP_AXIS),
 bias_axes=(NAMED_TP_AXIS,),
 sharding_type=te.ShardingType.DP_TP_COL,
 dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256,
 x = te_flax.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
Expand DownExpand Up@@ -248,9 +249,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand DownExpand Up@@ -356,7 +355,7 @@ def train_and_evaluate(args):
 abs_var_collect = jax.eval_shape(encoder.init, init_rngs, inputs, masks)

 customized_rules = ((NAMED_BROADCAST_AXIS, None), (NAMED_TP_AXIS, DEVICE_TP_AXIS))
 sharding_rules = te.flax.extend_logical_axis_rules(tuple()) + customized_rules
 sharding_rules = te_flax.extend_logical_axis_rules(tuple()) + customized_rules
 params_pspec = get_params_pspec(sharding_rules, abs_var_collect)
 inputs_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None)
 masks_pspec = jax.sharding.PartitionSpec(DEVICE_DP_AXIS, None, None, None)
Expand Down

13 changes: 6 additions & 7 deletions
examples/jax/encoder/test_single_gpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -17,6 +17,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

PARAMS_KEY = 'params'
DROPOUT_KEY = 'dropout'
Expand All@@ -31,23 +32,23 @@ class Net(nn.Module):
 def __call__(self, x, mask, disable_dropout=False):
 x = nn.Embed(num_embeddings=self.num_embed, features=256, dtype=jnp.bfloat16)(x)

 te_Encoder = partial(te.flax.TransformerLayer,
 te_Encoder = partial(te_flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
 hidden_dropout=0.1,
 attention_dropout=0.1,
 dropout_rng_name=DROPOUT_KEY,
 layer_type=te.flax.TransformerLayerType.ENCODER,
 layer_type=te_flax.TransformerLayerType.ENCODER,
 enable_relative_embedding=False,
 dtype=jnp.bfloat16)
 x = te_Encoder()(x, attention_mask=mask, deterministic=disable_dropout)

 x = x.reshape(x.shape[0], -1)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = te.flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)
 x = te_flax.DenseGeneral(features=256, dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand DownExpand Up@@ -160,9 +161,7 @@ def data_preprocess(dataset, vocab, word_id, max_seq_len):
 else:
 tensor[i] = vocab[word]

 seq_len = len(tokens)
 if seq_len > max_seq_len:
 seq_len = max_seq_len
 seq_len = min(len(tokens), max_seq_len)
 mask_2d = mask_3d[j]
 mask_2d[:seq_len, :seq_len] = 0

Expand Down

3 changes: 2 additions & 1 deletion
examples/jax/mnist/test_single_gpu_mnist.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,7 @@
from flax.training import train_state

import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax

IMAGE_H = 28
IMAGE_W = 28
Expand All@@ -32,7 +33,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.flax.DenseGeneral
 nn_Dense = te_flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down
Loading