Skip to content

[JAX] Adjust Module Structure.#169

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
ksivaman merged 11 commits into from
Apr 28, 2023
Merged

[JAX] Adjust Module Structure. #169

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
e3419c7
Adjust Module Structure.
mingxu1067 Apr 24, 2023
085f70b
Make changes be compatible to previous versions.
mingxu1067 Apr 25, 2023
9a3afe6
Adapt jax/examples to the new module structure.
mingxu1067 Apr 25, 2023
df99161
Update jax/docs and Add deprecated warning.
mingxu1067 Apr 25, 2023
2618ff3
Update README
mingxu1067 Apr 26, 2023
40361fc
Adding deprecated_wrapper
mingxu1067 Apr 26, 2023
8faaa5e
Adding deprecated warning to flax modules which imported via transfor…
mingxu1067 Apr 26, 2023
370205d
Fix CI errors and update docs.
mingxu1067 Apr 26, 2023
c69656c
Removing unnecessary deprecated warning in docs.
mingxu1067 Apr 27, 2023
5c1f169
Implementing __iter__ to DeprecatedEnum.
mingxu1067 Apr 27, 2023
ddeb9fd
Merge branch 'main' into mingh/module_structure_adjust
ksivaman Apr 27, 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
5 changes: 4 additions & 1 deletion README.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

import jax
Expand All@@ -90,7 +93,7 @@ JAX

# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
model = te.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
25 changes: 12 additions & 13 deletions docs/api/jax.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
:members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
:members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
:members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
:members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
:members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
:members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
:members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules
24 changes: 12 additions & 12 deletions examples/jax/encoder/test_model_parallel_encoder.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
te_Encoder = partial(te.flax.TransformerLayer,
hidden_size=256,
mlp_hidden_size=1024,
num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

x = te.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.DenseGeneral(features=256,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)
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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
return x
Expand Down
8 changes: 5 additions & 3 deletions examples/jax/encoder/test_multigpu_encoder.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
te_Encoder = partial(te.flax.TransformerLayer,
hidden_size=256,
mlp_hidden_size=1024,
num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
return x
Expand Down
6 changes: 3 additions & 3 deletions examples/jax/encoder/test_single_gpu_encoder.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
te_Encoder = partial(te.flax.TransformerLayer,
hidden_size=256,
mlp_hidden_size=1024,
num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

x = te.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 Down
2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
@nn.compact
def __call__(self, x, disable_dropout=False):
if self.use_te:
nn_Dense = te.DenseGeneral
nn_Dense = te.flax.DenseGeneral
else:
nn_Dense = nn.Dense

Expand Down
2 changes: 1 addition & 1 deletion tests/jax/test_layer.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down
2 changes: 1 addition & 1 deletion tests/jax/test_sharding.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down
53 changes: 53 additions & 0 deletions transformer_engine/common/utils.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
"""DeprecatedEnum"""

def __init__(self, enum_cls, msg):
self.enum_cls = enum_cls
self.msg = msg

def __iter__(self):
return iter(list(self.enum_cls.__members__.values()))

def __getattr__(self, name):
if name in self.enum_cls.__members__:
warnings.warn(self.msg, DeprecationWarning)
return self.enum_cls.__members__[name]
raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
"""Deprecate wrapper"""
if inspect.isclass(obj):
if issubclass(obj, Enum):
return DeprecatedEnum(obj, msg)

class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
"""DeprecatedCls"""

def __init__(self, *args, **kwargs):
warnings.warn(msg, DeprecationWarning)
super().__init__(*args, **kwargs)

return DeprecatedCls

if inspect.isfunction(obj):

def deprecated(*args, **kwargs):
warnings.warn(msg, DeprecationWarning)
return obj(*args, **kwargs)

return deprecated

raise NotImplementedError(
f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")
41 changes: 36 additions & 5 deletions transformer_engine/jax/__init__.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
flax.extend_logical_axis_rules,
"extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
"DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
"LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
flax.LayerNormDenseGeneral,
"LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
"LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
flax.TransformerEngineBase,
"TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
flax.RelativePositionBiases,
"RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
flax.TransformerLayerType,
"TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]
9 changes: 9 additions & 0 deletions transformer_engine/jax/flax/__init__.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
46 changes: 22 additions & 24 deletions transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
if original_init is None:
if not zero_centered_gamma:
return nn.initializers.ones
return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
bias_axes, dtype):
scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
sharding_type: ShardingType = ShardingType.SINGLE

def __post_init__(self):
if self.scale_init is None:
if not self.zero_centered_gamma:
self.scale_init = nn.initializers.ones
else:
self.scale_init = nn.initializers.zeros
self.scale_init = _obtain_default_layernorm_scale_init_if_need(
self.scale_init, self.zero_centered_gamma)
super().__post_init__()

@nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
def __post_init__(self):
if self.kernel_init is None:
self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
if self.scale_init is None:
if not self.zero_centered_gamma:
self.scale_init = nn.initializers.ones
else:
self.scale_init = nn.initializers.zeros
self.scale_init = _obtain_default_layernorm_scale_init_if_need(
self.scale_init, self.zero_centered_gamma)
super().__post_init__()

@nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
def __post_init__(self):
if self.kernel_init is None:
self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
if self.scale_init is None:
if not self.zero_centered_gamma:
self.scale_init = nn.initializers.ones
else:
self.scale_init = nn.initializers.zeros
self.scale_init = _obtain_default_layernorm_scale_init_if_need(
self.scale_init, self.zero_centered_gamma)
super().__post_init__()

@nn.compact
Expand Down
6 changes: 3 additions & 3 deletions transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
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;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
[JAX] Adjust Module Structure. by mingxu1067 · Pull Request #169 · NVIDIA/TransformerEngine · GitHub
Skip to content
[JAX] Adjust Module Structure.#169

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

ksivaman
merged 11 commits into

from

Apr 28, 2023
 Merged


[JAX] Adjust Module Structure.
#169
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


11 commits


Select commit
Hold shift + click to select a range
e3419c7

Adjust Module Structure.


mingxu1067 Apr 24, 2023
085f70b

Make changes be compatible to previous versions.


mingxu1067 Apr 25, 2023
9a3afe6

Adapt jax/examples to the new module structure.


mingxu1067 Apr 25, 2023
df99161

Update jax/docs and Add deprecated warning.


mingxu1067 Apr 25, 2023
2618ff3

Update README


mingxu1067 Apr 26, 2023
40361fc

Adding deprecated_wrapper


mingxu1067 Apr 26, 2023
8faaa5e

Adding deprecated warning to flax modules which imported via transfor…


mingxu1067 Apr 26, 2023
370205d

Fix CI errors and update docs.


mingxu1067 Apr 26, 2023
c69656c

Removing unnecessary deprecated warning in docs.


mingxu1067 Apr 27, 2023
5c1f169

Implementing __iter__ to DeprecatedEnum.


mingxu1067 Apr 27, 2023
ddeb9fd

Merge branch 'main' into mingh/module_structure_adjust


ksivaman Apr 27, 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


5 changes: 4 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

 import jax
Expand All@@ -90,7 +93,7 @@ JAX

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.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

25 changes: 12 additions & 13 deletions
docs/api/jax.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
 :members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules

24 changes: 12 additions & 12 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

 x = te.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.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
 dtype=jnp.bfloat16)(x)
 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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

8 changes: 5 additions & 3 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

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

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

 x = te.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 Down

2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.DenseGeneral
 nn_Dense = te.flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_layer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_sharding.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down

53 changes: 53 additions & 0 deletions
transformer_engine/common/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
 """DeprecatedEnum"""

 def __init__(self, enum_cls, msg):
 self.enum_cls = enum_cls
 self.msg = msg

 def __iter__(self):
 return iter(list(self.enum_cls.__members__.values()))

 def __getattr__(self, name):
 if name in self.enum_cls.__members__:
 warnings.warn(self.msg, DeprecationWarning)
 return self.enum_cls.__members__[name]
 raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
 """Deprecate wrapper"""
 if inspect.isclass(obj):
 if issubclass(obj, Enum):
 return DeprecatedEnum(obj, msg)

 class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
 """DeprecatedCls"""

 def __init__(self, *args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 super().__init__(*args, **kwargs)

 return DeprecatedCls

 if inspect.isfunction(obj):

 def deprecated(*args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 return obj(*args, **kwargs)

 return deprecated

 raise NotImplementedError(
 f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")

41 changes: 36 additions & 5 deletions
transformer_engine/jax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
 flax.extend_logical_axis_rules,
 "extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
 "DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
 "LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
 flax.LayerNormDenseGeneral,
 "LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
 "LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
 flax.TransformerEngineBase,
 "TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
 flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
 flax.RelativePositionBiases,
 "RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
 flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
 flax.TransformerLayerType,
 "TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
 'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
 'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
 'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
 'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]

9 changes: 9 additions & 0 deletions
transformer_engine/jax/flax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType

46 changes: 22 additions & 24 deletions
transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
 return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
 if original_init is None:
 if not zero_centered_gamma:
 return nn.initializers.ones
 return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
 bias_axes, dtype):
 scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
 sharding_type: ShardingType = ShardingType.SINGLE

 def __post_init__(self):
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand Down

6 changes: 3 additions & 3 deletions
transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down

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

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

ksivaman
merged 11 commits into

from

Apr 28, 2023
 Merged


[JAX] Adjust Module Structure.
#169
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


11 commits


Select commit
Hold shift + click to select a range
e3419c7

Adjust Module Structure.


mingxu1067 Apr 24, 2023
085f70b

Make changes be compatible to previous versions.


mingxu1067 Apr 25, 2023
9a3afe6

Adapt jax/examples to the new module structure.


mingxu1067 Apr 25, 2023
df99161

Update jax/docs and Add deprecated warning.


mingxu1067 Apr 25, 2023
2618ff3

Update README


mingxu1067 Apr 26, 2023
40361fc

Adding deprecated_wrapper


mingxu1067 Apr 26, 2023
8faaa5e

Adding deprecated warning to flax modules which imported via transfor…


mingxu1067 Apr 26, 2023
370205d

Fix CI errors and update docs.


mingxu1067 Apr 26, 2023
c69656c

Removing unnecessary deprecated warning in docs.


mingxu1067 Apr 27, 2023
5c1f169

Implementing __iter__ to DeprecatedEnum.


mingxu1067 Apr 27, 2023
ddeb9fd

Merge branch 'main' into mingh/module_structure_adjust


ksivaman Apr 27, 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


5 changes: 4 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

 import jax
Expand All@@ -90,7 +93,7 @@ JAX

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.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

25 changes: 12 additions & 13 deletions
docs/api/jax.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
 :members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules

24 changes: 12 additions & 12 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

 x = te.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.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
 dtype=jnp.bfloat16)(x)
 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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

8 changes: 5 additions & 3 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

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

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

 x = te.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 Down

2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.DenseGeneral
 nn_Dense = te.flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_layer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_sharding.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down

53 changes: 53 additions & 0 deletions
transformer_engine/common/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
 """DeprecatedEnum"""

 def __init__(self, enum_cls, msg):
 self.enum_cls = enum_cls
 self.msg = msg

 def __iter__(self):
 return iter(list(self.enum_cls.__members__.values()))

 def __getattr__(self, name):
 if name in self.enum_cls.__members__:
 warnings.warn(self.msg, DeprecationWarning)
 return self.enum_cls.__members__[name]
 raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
 """Deprecate wrapper"""
 if inspect.isclass(obj):
 if issubclass(obj, Enum):
 return DeprecatedEnum(obj, msg)

 class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
 """DeprecatedCls"""

 def __init__(self, *args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 super().__init__(*args, **kwargs)

 return DeprecatedCls

 if inspect.isfunction(obj):

 def deprecated(*args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 return obj(*args, **kwargs)

 return deprecated

 raise NotImplementedError(
 f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")

41 changes: 36 additions & 5 deletions
transformer_engine/jax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
 flax.extend_logical_axis_rules,
 "extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
 "DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
 "LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
 flax.LayerNormDenseGeneral,
 "LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
 "LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
 flax.TransformerEngineBase,
 "TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
 flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
 flax.RelativePositionBiases,
 "RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
 flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
 flax.TransformerLayerType,
 "TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
 'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
 'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
 'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
 'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]

9 changes: 9 additions & 0 deletions
transformer_engine/jax/flax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType

46 changes: 22 additions & 24 deletions
transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
 return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
 if original_init is None:
 if not zero_centered_gamma:
 return nn.initializers.ones
 return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
 bias_axes, dtype):
 scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
 sharding_type: ShardingType = ShardingType.SINGLE

 def __post_init__(self):
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand Down

6 changes: 3 additions & 3 deletions
transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down

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

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

ksivaman
merged 11 commits into

from

Apr 28, 2023
 Merged


[JAX] Adjust Module Structure.
#169
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


11 commits


Select commit
Hold shift + click to select a range
e3419c7

Adjust Module Structure.


mingxu1067 Apr 24, 2023
085f70b

Make changes be compatible to previous versions.


mingxu1067 Apr 25, 2023
9a3afe6

Adapt jax/examples to the new module structure.


mingxu1067 Apr 25, 2023
df99161

Update jax/docs and Add deprecated warning.


mingxu1067 Apr 25, 2023
2618ff3

Update README


mingxu1067 Apr 26, 2023
40361fc

Adding deprecated_wrapper


mingxu1067 Apr 26, 2023
8faaa5e

Adding deprecated warning to flax modules which imported via transfor…


mingxu1067 Apr 26, 2023
370205d

Fix CI errors and update docs.


mingxu1067 Apr 26, 2023
c69656c

Removing unnecessary deprecated warning in docs.


mingxu1067 Apr 27, 2023
5c1f169

Implementing __iter__ to DeprecatedEnum.


mingxu1067 Apr 27, 2023
ddeb9fd

Merge branch 'main' into mingh/module_structure_adjust


ksivaman Apr 27, 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


5 changes: 4 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

 import jax
Expand All@@ -90,7 +93,7 @@ JAX

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.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

25 changes: 12 additions & 13 deletions
docs/api/jax.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
 :members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules

24 changes: 12 additions & 12 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

 x = te.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.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
 dtype=jnp.bfloat16)(x)
 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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

8 changes: 5 additions & 3 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

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

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

 x = te.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 Down

2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.DenseGeneral
 nn_Dense = te.flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_layer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_sharding.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down

53 changes: 53 additions & 0 deletions
transformer_engine/common/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
 """DeprecatedEnum"""

 def __init__(self, enum_cls, msg):
 self.enum_cls = enum_cls
 self.msg = msg

 def __iter__(self):
 return iter(list(self.enum_cls.__members__.values()))

 def __getattr__(self, name):
 if name in self.enum_cls.__members__:
 warnings.warn(self.msg, DeprecationWarning)
 return self.enum_cls.__members__[name]
 raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
 """Deprecate wrapper"""
 if inspect.isclass(obj):
 if issubclass(obj, Enum):
 return DeprecatedEnum(obj, msg)

 class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
 """DeprecatedCls"""

 def __init__(self, *args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 super().__init__(*args, **kwargs)

 return DeprecatedCls

 if inspect.isfunction(obj):

 def deprecated(*args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 return obj(*args, **kwargs)

 return deprecated

 raise NotImplementedError(
 f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")

41 changes: 36 additions & 5 deletions
transformer_engine/jax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
 flax.extend_logical_axis_rules,
 "extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
 "DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
 "LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
 flax.LayerNormDenseGeneral,
 "LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
 "LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
 flax.TransformerEngineBase,
 "TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
 flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
 flax.RelativePositionBiases,
 "RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
 flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
 flax.TransformerLayerType,
 "TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
 'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
 'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
 'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
 'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]

9 changes: 9 additions & 0 deletions
transformer_engine/jax/flax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType

46 changes: 22 additions & 24 deletions
transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
 return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
 if original_init is None:
 if not zero_centered_gamma:
 return nn.initializers.ones
 return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
 bias_axes, dtype):
 scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
 sharding_type: ShardingType = ShardingType.SINGLE

 def __post_init__(self):
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand Down

6 changes: 3 additions & 3 deletions
transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down

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

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

ksivaman
merged 11 commits into

from

Apr 28, 2023
 Merged


[JAX] Adjust Module Structure.
#169
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


11 commits


Select commit
Hold shift + click to select a range
e3419c7

Adjust Module Structure.


mingxu1067 Apr 24, 2023
085f70b

Make changes be compatible to previous versions.


mingxu1067 Apr 25, 2023
9a3afe6

Adapt jax/examples to the new module structure.


mingxu1067 Apr 25, 2023
df99161

Update jax/docs and Add deprecated warning.


mingxu1067 Apr 25, 2023
2618ff3

Update README


mingxu1067 Apr 26, 2023
40361fc

Adding deprecated_wrapper


mingxu1067 Apr 26, 2023
8faaa5e

Adding deprecated warning to flax modules which imported via transfor…


mingxu1067 Apr 26, 2023
370205d

Fix CI errors and update docs.


mingxu1067 Apr 26, 2023
c69656c

Removing unnecessary deprecated warning in docs.


mingxu1067 Apr 27, 2023
5c1f169

Implementing __iter__ to DeprecatedEnum.


mingxu1067 Apr 27, 2023
ddeb9fd

Merge branch 'main' into mingh/module_structure_adjust


ksivaman Apr 27, 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


5 changes: 4 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

 import jax
Expand All@@ -90,7 +93,7 @@ JAX

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.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

25 changes: 12 additions & 13 deletions
docs/api/jax.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
 :members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules

24 changes: 12 additions & 12 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

 x = te.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.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
 dtype=jnp.bfloat16)(x)
 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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

8 changes: 5 additions & 3 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

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

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

 x = te.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 Down

2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.DenseGeneral
 nn_Dense = te.flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_layer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_sharding.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down

53 changes: 53 additions & 0 deletions
transformer_engine/common/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
 """DeprecatedEnum"""

 def __init__(self, enum_cls, msg):
 self.enum_cls = enum_cls
 self.msg = msg

 def __iter__(self):
 return iter(list(self.enum_cls.__members__.values()))

 def __getattr__(self, name):
 if name in self.enum_cls.__members__:
 warnings.warn(self.msg, DeprecationWarning)
 return self.enum_cls.__members__[name]
 raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
 """Deprecate wrapper"""
 if inspect.isclass(obj):
 if issubclass(obj, Enum):
 return DeprecatedEnum(obj, msg)

 class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
 """DeprecatedCls"""

 def __init__(self, *args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 super().__init__(*args, **kwargs)

 return DeprecatedCls

 if inspect.isfunction(obj):

 def deprecated(*args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 return obj(*args, **kwargs)

 return deprecated

 raise NotImplementedError(
 f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")

41 changes: 36 additions & 5 deletions
transformer_engine/jax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
 flax.extend_logical_axis_rules,
 "extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
 "DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
 "LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
 flax.LayerNormDenseGeneral,
 "LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
 "LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
 flax.TransformerEngineBase,
 "TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
 flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
 flax.RelativePositionBiases,
 "RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
 flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
 flax.TransformerLayerType,
 "TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
 'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
 'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
 'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
 'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]

9 changes: 9 additions & 0 deletions
transformer_engine/jax/flax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType

46 changes: 22 additions & 24 deletions
transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
 return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
 if original_init is None:
 if not zero_centered_gamma:
 return nn.initializers.ones
 return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
 bias_axes, dtype):
 scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
 sharding_type: ShardingType = ShardingType.SINGLE

 def __post_init__(self):
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand Down

6 changes: 3 additions & 3 deletions
transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down

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

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

ksivaman
merged 11 commits into

from

Apr 28, 2023
 Merged


[JAX] Adjust Module Structure.
#169
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


11 commits


Select commit
Hold shift + click to select a range
e3419c7

Adjust Module Structure.


mingxu1067 Apr 24, 2023
085f70b

Make changes be compatible to previous versions.


mingxu1067 Apr 25, 2023
9a3afe6

Adapt jax/examples to the new module structure.


mingxu1067 Apr 25, 2023
df99161

Update jax/docs and Add deprecated warning.


mingxu1067 Apr 25, 2023
2618ff3

Update README


mingxu1067 Apr 26, 2023
40361fc

Adding deprecated_wrapper


mingxu1067 Apr 26, 2023
8faaa5e

Adding deprecated warning to flax modules which imported via transfor…


mingxu1067 Apr 26, 2023
370205d

Fix CI errors and update docs.


mingxu1067 Apr 26, 2023
c69656c

Removing unnecessary deprecated warning in docs.


mingxu1067 Apr 27, 2023
5c1f169

Implementing __iter__ to DeprecatedEnum.


mingxu1067 Apr 27, 2023
ddeb9fd

Merge branch 'main' into mingh/module_structure_adjust


ksivaman Apr 27, 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


5 changes: 4 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

 import jax
Expand All@@ -90,7 +93,7 @@ JAX

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.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

25 changes: 12 additions & 13 deletions
docs/api/jax.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
 :members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules

24 changes: 12 additions & 12 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

 x = te.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.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
 dtype=jnp.bfloat16)(x)
 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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

8 changes: 5 additions & 3 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

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

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

 x = te.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 Down

2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.DenseGeneral
 nn_Dense = te.flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_layer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_sharding.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down

53 changes: 53 additions & 0 deletions
transformer_engine/common/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
 """DeprecatedEnum"""

 def __init__(self, enum_cls, msg):
 self.enum_cls = enum_cls
 self.msg = msg

 def __iter__(self):
 return iter(list(self.enum_cls.__members__.values()))

 def __getattr__(self, name):
 if name in self.enum_cls.__members__:
 warnings.warn(self.msg, DeprecationWarning)
 return self.enum_cls.__members__[name]
 raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
 """Deprecate wrapper"""
 if inspect.isclass(obj):
 if issubclass(obj, Enum):
 return DeprecatedEnum(obj, msg)

 class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
 """DeprecatedCls"""

 def __init__(self, *args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 super().__init__(*args, **kwargs)

 return DeprecatedCls

 if inspect.isfunction(obj):

 def deprecated(*args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 return obj(*args, **kwargs)

 return deprecated

 raise NotImplementedError(
 f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")

41 changes: 36 additions & 5 deletions
transformer_engine/jax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
 flax.extend_logical_axis_rules,
 "extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
 "DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
 "LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
 flax.LayerNormDenseGeneral,
 "LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
 "LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
 flax.TransformerEngineBase,
 "TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
 flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
 flax.RelativePositionBiases,
 "RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
 flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
 flax.TransformerLayerType,
 "TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
 'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
 'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
 'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
 'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]

9 changes: 9 additions & 0 deletions
transformer_engine/jax/flax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType

46 changes: 22 additions & 24 deletions
transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
 return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
 if original_init is None:
 if not zero_centered_gamma:
 return nn.initializers.ones
 return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
 bias_axes, dtype):
 scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
 sharding_type: ShardingType = ShardingType.SINGLE

 def __post_init__(self):
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand Down

6 changes: 3 additions & 3 deletions
transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down

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

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

ksivaman
merged 11 commits into

from

Apr 28, 2023
 Merged


[JAX] Adjust Module Structure.
#169
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


11 commits


Select commit
Hold shift + click to select a range
e3419c7

Adjust Module Structure.


mingxu1067 Apr 24, 2023
085f70b

Make changes be compatible to previous versions.


mingxu1067 Apr 25, 2023
9a3afe6

Adapt jax/examples to the new module structure.


mingxu1067 Apr 25, 2023
df99161

Update jax/docs and Add deprecated warning.


mingxu1067 Apr 25, 2023
2618ff3

Update README


mingxu1067 Apr 26, 2023
40361fc

Adding deprecated_wrapper


mingxu1067 Apr 26, 2023
8faaa5e

Adding deprecated warning to flax modules which imported via transfor…


mingxu1067 Apr 26, 2023
370205d

Fix CI errors and update docs.


mingxu1067 Apr 26, 2023
c69656c

Removing unnecessary deprecated warning in docs.


mingxu1067 Apr 27, 2023
5c1f169

Implementing __iter__ to DeprecatedEnum.


mingxu1067 Apr 27, 2023
ddeb9fd

Merge branch 'main' into mingh/module_structure_adjust


ksivaman Apr 27, 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


5 changes: 4 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

 import jax
Expand All@@ -90,7 +93,7 @@ JAX

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.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

25 changes: 12 additions & 13 deletions
docs/api/jax.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
 :members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules

24 changes: 12 additions & 12 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

 x = te.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.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
 dtype=jnp.bfloat16)(x)
 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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

8 changes: 5 additions & 3 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

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

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

 x = te.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 Down

2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.DenseGeneral
 nn_Dense = te.flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_layer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_sharding.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down

53 changes: 53 additions & 0 deletions
transformer_engine/common/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
 """DeprecatedEnum"""

 def __init__(self, enum_cls, msg):
 self.enum_cls = enum_cls
 self.msg = msg

 def __iter__(self):
 return iter(list(self.enum_cls.__members__.values()))

 def __getattr__(self, name):
 if name in self.enum_cls.__members__:
 warnings.warn(self.msg, DeprecationWarning)
 return self.enum_cls.__members__[name]
 raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
 """Deprecate wrapper"""
 if inspect.isclass(obj):
 if issubclass(obj, Enum):
 return DeprecatedEnum(obj, msg)

 class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
 """DeprecatedCls"""

 def __init__(self, *args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 super().__init__(*args, **kwargs)

 return DeprecatedCls

 if inspect.isfunction(obj):

 def deprecated(*args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 return obj(*args, **kwargs)

 return deprecated

 raise NotImplementedError(
 f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")

41 changes: 36 additions & 5 deletions
transformer_engine/jax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
 flax.extend_logical_axis_rules,
 "extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
 "DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
 "LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
 flax.LayerNormDenseGeneral,
 "LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
 "LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
 flax.TransformerEngineBase,
 "TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
 flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
 flax.RelativePositionBiases,
 "RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
 flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
 flax.TransformerLayerType,
 "TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
 'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
 'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
 'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
 'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]

9 changes: 9 additions & 0 deletions
transformer_engine/jax/flax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType

46 changes: 22 additions & 24 deletions
transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
 return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
 if original_init is None:
 if not zero_centered_gamma:
 return nn.initializers.ones
 return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
 bias_axes, dtype):
 scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
 sharding_type: ShardingType = ShardingType.SINGLE

 def __post_init__(self):
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand Down

6 changes: 3 additions & 3 deletions
transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down

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

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

ksivaman
merged 11 commits into

from

Apr 28, 2023
 Merged


[JAX] Adjust Module Structure.
#169
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


11 commits


Select commit
Hold shift + click to select a range
e3419c7

Adjust Module Structure.


mingxu1067 Apr 24, 2023
085f70b

Make changes be compatible to previous versions.


mingxu1067 Apr 25, 2023
9a3afe6

Adapt jax/examples to the new module structure.


mingxu1067 Apr 25, 2023
df99161

Update jax/docs and Add deprecated warning.


mingxu1067 Apr 25, 2023
2618ff3

Update README


mingxu1067 Apr 26, 2023
40361fc

Adding deprecated_wrapper


mingxu1067 Apr 26, 2023
8faaa5e

Adding deprecated warning to flax modules which imported via transfor…


mingxu1067 Apr 26, 2023
370205d

Fix CI errors and update docs.


mingxu1067 Apr 26, 2023
c69656c

Removing unnecessary deprecated warning in docs.


mingxu1067 Apr 27, 2023
5c1f169

Implementing __iter__ to DeprecatedEnum.


mingxu1067 Apr 27, 2023
ddeb9fd

Merge branch 'main' into mingh/module_structure_adjust


ksivaman Apr 27, 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


5 changes: 4 additions & 1 deletion
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -69,6 +69,9 @@ pyTorch
JAX
^^^

Flax
~~~~

.. code-block:: python

 import jax
Expand All@@ -90,7 +93,7 @@ JAX

 # Enable autocasting for the forward pass
 with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
 model = te.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

25 changes: 12 additions & 13 deletions
docs/api/jax.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -9,34 +9,33 @@ Jax
.. autoapiclass:: transformer_engine.jax.MajorShardingType
.. autoapiclass:: transformer_engine.jax.ShardingType
.. autoapiclass:: transformer_engine.jax.TransformerLayerType
.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)


.. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas


.. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.MultiHeadAttention(head_dim, num_heads, **kwargs)
 :members: __call__

.. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
.. autoapiclass:: transformer_engine.jax.flax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs)
 :members: __call__


.. autoapifunction:: transformer_engine.jax.extend_logical_axis_rules
.. autoapifunction:: transformer_engine.jax.fp8_autocast
.. autoapifunction:: transformer_engine.jax.update_collections
.. autoapifunction:: transformer_engine.jax.update_fp8_metas
.. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules

24 changes: 12 additions & 12 deletions
examples/jax/encoder/test_model_parallel_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -59,7 +59,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -73,17 +73,17 @@ def __call__(self, x, mask, disable_dropout=False):

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

 x = te.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.DenseGeneral(features=256,
 kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
 bias_axes=(NAMED_BROADCAST_AXIS,),
 sharding_type=te.ShardingType.DP_TP_ROW,
 dtype=jnp.bfloat16)(x)
 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,
kernel_axes=(NAMED_TP_AXIS, NAMED_BROADCAST_AXIS),
bias_axes=(NAMED_BROADCAST_AXIS,),
sharding_type=te.ShardingType.DP_TP_ROW,
dtype=jnp.bfloat16)(x)

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

8 changes: 5 additions & 3 deletions
examples/jax/encoder/test_multigpu_encoder.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -56,7 +56,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -70,9 +70,11 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

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

 x = nn.Dense(features=2, dtype=jnp.bfloat16)(x)
 return x
Expand Down

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

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -46,7 +46,7 @@ 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.TransformerLayer,
 te_Encoder = partial(te.flax.TransformerLayer,
 hidden_size=256,
 mlp_hidden_size=1024,
 num_attention_heads=8,
Expand All@@ -60,9 +60,9 @@ def __call__(self, x, mask, disable_dropout=False):

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

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

 x = te.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 Down

2 changes: 1 addition & 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@@ -47,7 +47,7 @@ class Net(nn.Module):
 @nn.compact
 def __call__(self, x, disable_dropout=False):
 if self.use_te:
 nn_Dense = te.DenseGeneral
 nn_Dense = te.flax.DenseGeneral
 else:
 nn_Dense = nn.Dense

Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_layer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -10,7 +10,7 @@
import pytest

from transformer_engine.common.recipe import Format
from transformer_engine.jax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType
from transformer_engine.jax.fp8 import FP8Helper
from utils import assert_allclose, is_fp8_supported
from utils import DecoderLayer as RefDecoderLayer
Expand Down

2 changes: 1 addition & 1 deletion
tests/jax/test_sharding.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,7 @@
import pytest
from jax.experimental import maps

from transformer_engine.jax import extend_logical_axis_rules
from transformer_engine.jax.flax import extend_logical_axis_rules
from transformer_engine.jax.sharding import get_dot_sharding_meta
from transformer_engine.jax.sharding import get_elementwise_sharding_meta
from transformer_engine.jax.sharding import get_fp8_meta_sharding_meta
Expand Down

53 changes: 53 additions & 0 deletions
transformer_engine/common/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""The utilities for Transformer Engine"""
import inspect
import warnings
from enum import Enum

warnings.simplefilter('default')


class DeprecatedEnum: # pylint: disable=too-few-public-methods
 """DeprecatedEnum"""

 def __init__(self, enum_cls, msg):
 self.enum_cls = enum_cls
 self.msg = msg

 def __iter__(self):
 return iter(list(self.enum_cls.__members__.values()))

 def __getattr__(self, name):
 if name in self.enum_cls.__members__:
 warnings.warn(self.msg, DeprecationWarning)
 return self.enum_cls.__members__[name]
 raise AttributeError(f"{self.enum_cls} does not contain {name}")


def deprecate_wrapper(obj, msg):
 """Deprecate wrapper"""
 if inspect.isclass(obj):
 if issubclass(obj, Enum):
 return DeprecatedEnum(obj, msg)

 class DeprecatedCls(obj): # pylint: disable=too-few-public-methods
 """DeprecatedCls"""

 def __init__(self, *args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 super().__init__(*args, **kwargs)

 return DeprecatedCls

 if inspect.isfunction(obj):

 def deprecated(*args, **kwargs):
 warnings.warn(msg, DeprecationWarning)
 return obj(*args, **kwargs)

 return deprecated

 raise NotImplementedError(
 f"deprecate_cls_wrapper only support Class and Function, but got {type(obj)}.")

41 changes: 36 additions & 5 deletions
transformer_engine/jax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2,10 +2,41 @@
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""

from . import flax
from .fp8 import fp8_autocast, update_collections, update_fp8_metas, get_delayed_scaling
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType
from .sharding import MajorShardingType, ShardingResource, ShardingType
from ..common.utils import deprecate_wrapper

extend_logical_axis_rules = deprecate_wrapper(
 flax.extend_logical_axis_rules,
 "extend_logical_axis_rules is moving to transformer_engine.jax.flax module")
DenseGeneral = deprecate_wrapper(flax.DenseGeneral,
 "DenseGeneral is moving to transformer_engine.jax.flax module")
LayerNorm = deprecate_wrapper(flax.LayerNorm,
 "LayerNorm is moving to transformer_engine.jax.flax module")
LayerNormDenseGeneral = deprecate_wrapper(
 flax.LayerNormDenseGeneral,
 "LayerNormDenseGeneral is moving to transformer_engine.jax.flax module")
LayerNormMLP = deprecate_wrapper(flax.LayerNormMLP,
 "LayerNormMLP is moving to transformer_engine.jax.flax module")
TransformerEngineBase = deprecate_wrapper(
 flax.TransformerEngineBase,
 "TransformerEngineBase is moving to transformer_engine.jax.flax module")
MultiHeadAttention = deprecate_wrapper(
 flax.MultiHeadAttention, "MultiHeadAttention is moving to transformer_engine.jax.flax module")
RelativePositionBiases = deprecate_wrapper(
 flax.RelativePositionBiases,
 "RelativePositionBiases is moving to transformer_engine.jax.flax module")
TransformerLayer = deprecate_wrapper(
 flax.TransformerLayer, "TransformerLayer is moving to transformer_engine.jax.flax module")
TransformerLayerType = deprecate_wrapper(
 flax.TransformerLayerType,
 "TransformerLayerType is moving to transformer_engine.jax.flax module")

__all__ = [
 'fp8_autocast', 'update_collections', 'update_fp8_metas', 'get_delayed_scaling',
 'MajorShardingType', 'ShardingResource', 'ShardingType', 'flax', 'DenseGeneral', 'LayerNorm',
 'LayerNormDenseGeneral', 'LayerNormMLP', 'TransformerEngineBase', 'MultiHeadAttention',
 'RelativePositionBiases', 'TransformerLayer', 'TransformerLayerType'
]

9 changes: 9 additions & 0 deletions
transformer_engine/jax/flax/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for JAX"""
from .module import DenseGeneral, LayerNorm
from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
from .transformer import extend_logical_axis_rules
from .transformer import MultiHeadAttention, RelativePositionBiases
from .transformer import TransformerLayer, TransformerLayerType

46 changes: 22 additions & 24 deletions
transformer_engine/jax/module.py → transformer_engine/jax/flax/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,15 +16,15 @@
from jax import nn as jax_nn
from jax import random as jax_random

from .dot import fp8_dot
from .fp8 import FP8GemmPackage, FP8Helper
from .layernorm import canonicalize_layernorm_type
from .layernorm import layernorm, layernorm_fp8_dot
from .mlp import fp8_ln_mlp, geglu
from .sharding import infer_sharding_type
from .softmax import is_softmax_kernel_available
from .sharding import MajorShardingType, ShardingType
from .softmax import softmax, SoftmaxType
from ..dot import fp8_dot
from ..fp8 import FP8GemmPackage, FP8Helper
from ..layernorm import canonicalize_layernorm_type
from ..layernorm import layernorm, layernorm_fp8_dot
from ..mlp import fp8_ln_mlp, geglu
from ..sharding import infer_sharding_type
from ..softmax import is_softmax_kernel_available
from ..sharding import MajorShardingType, ShardingType
from ..softmax import softmax, SoftmaxType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand All@@ -46,6 +46,13 @@ def _canonicalize_tuple(x):
 return (x,)


def _obtain_default_layernorm_scale_init_if_need(original_init, zero_centered_gamma):
 if original_init is None:
 if not zero_centered_gamma:
 return nn.initializers.ones
 return nn.initializers.zeros


def _create_layernorm_parameters(layernorm_type, shape, scale_init, scale_axes, bias_init,
 bias_axes, dtype):
 scale = nn_partitioning.param_with_axes('scale',
Expand DownExpand Up@@ -250,11 +257,8 @@ class LayerNorm(nn.Module):
 sharding_type: ShardingType = ShardingType.SINGLE

 def __post_init__(self):
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -551,11 +555,8 @@ class LayerNormDenseGeneral(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand DownExpand Up@@ -785,11 +786,8 @@ class LayerNormMLP(TransformerEngineBase):
 def __post_init__(self):
 if self.kernel_init is None:
 self.kernel_init = nn.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')
 if self.scale_init is None:
 if not self.zero_centered_gamma:
 self.scale_init = nn.initializers.ones
 else:
 self.scale_init = nn.initializers.zeros
 self.scale_init = _obtain_default_layernorm_scale_init_if_need(
 self.scale_init, self.zero_centered_gamma)
 super().__post_init__()

 @nn.compact
Expand Down

6 changes: 3 additions & 3 deletions
transformer_engine/jax/transformer.py → transformer_engine/jax/flax/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -18,9 +18,9 @@

from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP
from .module import LayerNorm, Softmax
from .softmax import SoftmaxType
from .sharding import infer_major_sharding_type, infer_sharding_type
from .sharding import global_shard_resource, ShardingType
from ..softmax import SoftmaxType
from ..sharding import infer_major_sharding_type, infer_sharding_type
from ..sharding import global_shard_resource, ShardingType

PRNGKey = Any
Shape = Tuple[int, ...]
Expand Down