Skip to content

GitLab

Unverified Commit 1f4c3979 authored by Kirthi Shankar Sivamani's avatar Kirthi Shankar Sivamani Committed by GitHub
Browse files

Remove TF support (#467) 


Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
parent 2574a1ca
Expand all Show whitespace changes
Inline Side-by-side
Showing with 0 additions and 4724 deletions
transformer_engine/tensorflow/__init__.py deleted 100644 → 0
View file @ 2574a1ca
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Transformer Engine bindings for Tensorflow"""
from transformer_engine.common.recipe import DelayedScaling
from transformer_engine.common.recipe import Format
from .constants import TE_DType
from .fp8 import fp8_autocast
from .module import Dense
from .module import LayerNorm
from .module import LayerNormDense
from .module import LayerNormMLP
from .module import get_stream_id
from .transformer import MultiHeadAttention
from .transformer import TransformerLayer
transformer_engine/tensorflow/constants.py deleted 100644 → 0
View file @ 2574a1ca
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Enums for e2e transformer"""
import tensorflow as tf
import transformer_engine_tensorflow as tex
"""
This is a map: tf.dtype -> int
Used for passing dtypes into cuda
extension. Has one to one mapping
with enum in transformer_engine.h
"""
TE_DType = {
tf.int8: tex.DType.kByte,
tf.int32: tex.DType.kInt32,
tf.float32: tex.DType.kFloat32,
tf.half: tex.DType.kFloat16,
tf.bfloat16: tex.DType.kBFloat16,
}
AttnMaskTypes = ("causal", "padding")
AttnTypes = ("self", "cross")
LayerTypes = ("encoder", "decoder")
transformer_engine/tensorflow/csrc/extensions.cu deleted 100644 → 0
View file @ 2574a1ca
This diff is collapsed.
transformer_engine/tensorflow/csrc/get_stream_op.cpp deleted 100644 → 0
View file @ 2574a1ca
/*************************************************************************
* Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
*
* See LICENSE for license information.
************************************************************************/
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/xla/stream_executor/gpu/gpu_stream.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
class GetStreamOp : public OpKernel {
public:
explicit GetStreamOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* ctx) override {
Tensor* output = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output("stream_id", {1}, &output));
auto vec = output->vec<int64_t>();
se::Stream* stream = ctx->op_device_context()->stream();
auto gpu_stream = se::gpu::AsGpuStreamValue(stream);
vec(0) = static_cast<int64_t>(reinterpret_cast<uintptr_t>(gpu_stream));
}
};
REGISTER_OP("GetStream")
.Output("stream_id: int64")
.SetShapeFn(shape_inference::UnknownShape);
REGISTER_OP_NO_GRADIENT("GetStream");
REGISTER_KERNEL_BUILDER(
Name("GetStream").Device(DEVICE_GPU).HostMemory("stream_id"), GetStreamOp);
} // namespace tensorflow
transformer_engine/tensorflow/fp8.py deleted 100644 → 0
View file @ 2574a1ca
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""FP8 utilies for TransformerEngine"""
from contextlib import contextmanager
from typing import Generator, Optional, Dict, Any
import tensorflow as tf
import transformer_engine_tensorflow as tex
from transformer_engine.common.recipe import DelayedScaling, Format
_FP8_ENABLED = False
_FP8_RECIPE = None
_FP8_DISTRIBUTED_GROUP = None
_IS_FIRST_FP8_MODULE = False
_FP8_AUTOCAST_COUNTER = 0
_FP8_CURRENT_CONTEXT_ID = 0
_FP8_AUTOCAST_DEPTH = 0
_global_fp8_buffer = {}
_amax_forward_global_reduce_func = lambda: None
_buffer_delete_key_fwd = None
_buffer_delete_key_bwd = None
def get_meta_tensor_key(forward: bool = True) -> str:
"""Returns scaling key in `fp8_meta`."""
if forward:
return "scaling_fwd"
return "scaling_bwd"
def get_autocast_key(forward: bool = True) -> str:
"""Returns module position key in `fp8_meta`."""
if forward:
return "autocast_id_fwd"
return "autocast_id_bwd"
def get_amax_buffer_key(fp8_meta: Dict[str, Any], forward: bool = True) -> str:
"""Return a key in `_global_fp8_buffer` for the AMAX storage."""
if forward:
return f"FWD_AMAX_{fp8_meta['autocast_id_fwd']}"
return f"BWD_AMAX_{fp8_meta['autocast_id_bwd']}"
def set_amax_buffer_key_deletion(
fp8_meta: Dict[str, Any], forward: bool = True
) -> None:
"""Delete this amax key from global buffer during autocast end."""
if get_autocast_key(forward=forward) not in fp8_meta:
return
global _buffer_delete_key_fwd, _buffer_delete_key_bwd
if forward:
_buffer_delete_key_fwd = get_amax_buffer_key(fp8_meta, forward=forward)
else:
_buffer_delete_key_bwd = get_amax_buffer_key(fp8_meta, forward=forward)
def get_default_fp8_recipe():
"""FP8 recipe if not provided by user
Margin = 0, interval = 1, E4M3
"""
return DelayedScaling()
@contextmanager
def fp8_autocast(
enabled: bool = False,
fp8_recipe: Optional[DelayedScaling] = None,
) -> Generator[None, None, None]:
"""
Context manager for FP8 usage.
.. code-block:: python
with fp8_autocast(enabled=True):
out = model(inp)
.. note::
Support for FP8 in the Dense layer of Transformer Engine is currently
limited to tensors with shapes where both dimensions are divisible by 16.
In terms of the input to the full Transformer network, this typically
requires padding sequence length to be multiple of 16.
Parameters
----------
enabled: bool, default = `False`
whether or not to enable fp8
fp8_recipe: recipe.DelayedScaling, default = `None`
recipe used for FP8 training.
"""
global _FP8_ENABLED, _FP8_RECIPE, _FP8_DISTRIBUTED_GROUP, _FP8_AUTOCAST_DEPTH
global _IS_FIRST_FP8_MODULE, _FP8_AUTOCAST_COUNTER
global _global_fp8_buffer, _buffer_delete_key_fwd
fp8_state = (_FP8_ENABLED, _FP8_RECIPE, _FP8_DISTRIBUTED_GROUP)
try:
_FP8_ENABLED = enabled
_FP8_RECIPE = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe
if _FP8_AUTOCAST_DEPTH == 0:
_IS_FIRST_FP8_MODULE = True
_FP8_AUTOCAST_COUNTER += 1
_FP8_AUTOCAST_DEPTH += 1
yield
finally:
_FP8_ENABLED, _FP8_RECIPE, _FP8_DISTRIBUTED_GROUP = fp8_state
_IS_FIRST_FP8_MODULE = False
_FP8_AUTOCAST_DEPTH -= 1
if _FP8_AUTOCAST_DEPTH == 0:
if callable(_amax_forward_global_reduce_func):
_amax_forward_global_reduce_func()
delete_key_from_amax_buffer(forward=True)
def get_fp8_context_id() -> int:
"""Returns an ID for the current FP8 context."""
return _FP8_CURRENT_CONTEXT_ID
def set_fp8_context_id(ctx_id: int) -> None:
"""Sets the current FP8 context."""
global _FP8_CURRENT_CONTEXT_ID
_FP8_CURRENT_CONTEXT_ID = ctx_id
def new_fp8_context_id() -> int:
"""Returns global autocast counter as a proxy to be used
as the autocast ID for FP8 modules.
"""
return _FP8_AUTOCAST_COUNTER
def is_fp8_enabled():
"""Is FP8 enabled"""
return _FP8_ENABLED
def is_first_fp8_module():
"""Returns `True` only the first time when called multiple
times from within the same `fp8_autocast` context.
"""
global _IS_FIRST_FP8_MODULE
tmp = _IS_FIRST_FP8_MODULE
_IS_FIRST_FP8_MODULE = False
return tmp
def get_fp8_recipe():
"""Return the fp8 recipe"""
return _FP8_RECIPE
def _default_sf_compute(amax, scale, fp8_max, margin):
"""Default function to convert amax to scaling factor."""
sf = (fp8_max / amax) / (2 ** margin)
sf = tf.where(amax > 0.0, sf, scale)
sf = tf.where(tf.math.is_finite(amax), sf, scale)
return sf
def _roll_and_zero_out(amax_history):
"""Update amax history and set next amax to zero."""
amax_history = tf.roll(amax_history, -1, 0)
zeros = tf.zeros(shape=amax_history[0].shape)
updated = tf.tensor_scatter_nd_update(amax_history, [[0]], [zeros])
return updated
@tf.function(jit_compile=True)
def _reduce_max_and_default_sf_compute(amax_history, scale, fp8_max, margin):
"""Get amax using max algorithm and compute scaling factor."""
amax = tf.reduce_max(amax_history, axis=0)
sf = _default_sf_compute(amax, scale, fp8_max, margin)
updated = _roll_and_zero_out(amax_history)
return updated, sf
@tf.function(jit_compile=True)
def _most_recent_and_default_sf_compute(amax_history, scale, fp8_max, margin):
"""Get amax using most-recent algorithm and compute scaling factor."""
amax = amax_history[0]
sf = _default_sf_compute(amax, scale, fp8_max, margin)
updated = _roll_and_zero_out(amax_history)
return updated, sf
def fused_amax_and_scale_update(
amax_history: tf.Variable,
scale: tf.Variable,
scale_inv: tf.Variable,
fp8_max: float,
margin: int,
amax_compute_algo: str,
):
"""Amax to scale conversion."""
if amax_compute_algo == "max":
updated, sf = _reduce_max_and_default_sf_compute(
amax_history, scale, fp8_max, margin
)
else:
assert amax_compute_algo == "most_recent"
updated, sf = _most_recent_and_default_sf_compute(
amax_history, scale, fp8_max, margin
)
amax_history.assign(updated)
scale.assign(sf)
scale_inv.assign(1.0 / sf)
def amax_and_scale_update(
fp8_meta: Dict[str, Any],
fwd_update: bool,
) -> None:
"""Updates fp8 amaxes/scales for fwd | bwd."""
amax_compute = fp8_meta["recipe"].amax_compute_algo
sf_compute = fp8_meta["recipe"].scaling_factor_compute_algo
fp8_meta_tensor_key = "scaling_fwd" if fwd_update else "scaling_bwd"
fp8_max_key = "fp8_max_fwd" if fwd_update else "fp8_max_bwd"
if not callable(amax_compute) and sf_compute is None:
fused_amax_and_scale_update(
fp8_meta[fp8_meta_tensor_key]["amax_history"],
fp8_meta[fp8_meta_tensor_key]["scale"],
fp8_meta[fp8_meta_tensor_key]["scale_inv"],
fp8_meta[fp8_max_key],
fp8_meta["recipe"].margin,
fp8_meta["recipe"].amax_compute_algo,
)
else:
raise ValueError(
"We only support the fp8 recipe with 'max' or 'most_recent' "
"amax_compute_algo and default scaling_factor_compute_algo at this "
"moment."
)
def get_fp8_te_dtype(fp8_recipe: DelayedScaling, fprop_tensor: bool = True):
"""Get fp8 data type according to recipe and tensor"""
if fp8_recipe.fp8_format == Format.E4M3 or (
fp8_recipe.fp8_format == Format.HYBRID and fprop_tensor
):
return tex.DType.kFloat8E4M3
return tex.DType.kFloat8E5M2
def delete_key_from_amax_buffer(forward: bool = True) -> None:
"""Delete the key from global amax buffer."""
global _global_fp8_buffer, _buffer_delete_key_fwd, _buffer_delete_key_bwd
if forward:
if (
_buffer_delete_key_fwd is not None
and _buffer_delete_key_fwd in _global_fp8_buffer
):
del _global_fp8_buffer[_buffer_delete_key_fwd]
else:
if (
_buffer_delete_key_bwd is not None
and _buffer_delete_key_bwd in _global_fp8_buffer
):
del _global_fp8_buffer[_buffer_delete_key_bwd]
transformer_engine/tensorflow/jit.py deleted 100644 → 0
View file @ 2574a1ca
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""XLA functions and JIT utilities"""
from typing import Callable
import tensorflow as tf
@tf.function(jit_compile=True)
def _bgrad_dgelu_fused(grad_output, inp):
"""Bgrad-Dgelu fused"""
x = inp
tanh_out = tf.math.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
# sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
ff = 0.5 * x * (
(1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)
) + 0.5 * (1 + tanh_out)
dgelu = ff * grad_output
bgrad = tf.math.reduce_sum(dgelu, axis=0)
return bgrad, dgelu
def bgrad_dgelu_fused(grad_output, inp):
"""Bgrad-Dgelu fused"""
return _bgrad_dgelu_fused(grad_output, inp)
def bias_dropout_add(
x: tf.Tensor,
bias: tf.Variable,
residual: tf.Tensor,
prob: float,
training: bool,
) -> tf.Tensor:
"""dropout(inp + bias) + residual"""
# TODO(kaixih): Use stateless_dropout and specify the seed mainly for
# debugging purpose. Should allow random seed.
out = (
tf.nn.experimental.stateless_dropout(
x + bias,
rate=prob,
seed=[1, 0],
)
if training
else x + bias
)
out = residual + out
return out
def get_bias_dropout_add(training: bool) -> Callable:
"""bias_dropout_add based on training or not"""
def _bias_dropout_add(x, bias, residual, prob):
return bias_dropout_add(x, bias, residual, prob, training)
return _bias_dropout_add
@tf.function(jit_compile=True)
def bias_dropout_add_fused_train_(
x: tf.Tensor,
bias: tf.Variable,
residual: tf.Tensor,
prob: float,
) -> tf.Tensor:
"""Jit fused bias_dropout_add for training"""
return bias_dropout_add(x, bias, residual, prob, True)
def bias_dropout_add_fused_train(
x: tf.Tensor,
bias: tf.Variable,
residual: tf.Tensor,
prob: float,
) -> tf.Tensor:
"""Jit fused bias_dropout_add for training"""
return bias_dropout_add_fused_train_(x, bias, residual, prob)
@tf.function(jit_compile=True)
def bias_dropout_add_fused_inference_(
x: tf.Tensor,
bias: tf.Variable,
residual: tf.Tensor,
prob: float,
) -> tf.Tensor:
"""Jit fused bias_dropout_add for inference"""
return bias_dropout_add(x, bias, residual, prob, False)
def bias_dropout_add_fused_inference(
x: tf.Tensor,
bias: tf.Variable,
residual: tf.Tensor,
prob: float,
) -> tf.Tensor:
"""Jit fused bias_dropout_add for inference"""
return bias_dropout_add_fused_inference_(x, bias, residual, prob)
transformer_engine/tensorflow/module.py deleted 100644 → 0
View file @ 2574a1ca
This diff is collapsed.
transformer_engine/tensorflow/softmax.py deleted 100644 → 0
View file @ 2574a1ca
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Fused scaled masked softmax functions"""
from typing import Callable
import os
import transformer_engine_tensorflow as tex
import tensorflow as tf
from .module import get_stream_id
THREADS_PER_WARP = 32
THREADS_PER_BLOCK = 128
_default_causal_mask = {}
def _get_default_causal_mask(sq: int) -> tf.Tensor:
"""Return the causal upper triangular mask for softmax input"""
if sq not in _default_causal_mask:
# In TF, the mask specifies 1 to keep and 0 to mask. In "causal" mask
# mode, we compute the softmax of the lower triangular.
mask_operator = tf.linalg.LinearOperatorLowerTriangular(
tf.ones((sq, sq), dtype=tf.bool))
mask = mask_operator.to_dense()
_default_causal_mask[sq] = mask
return _default_causal_mask[sq]
class FusedScaleMaskSoftmax(tf.keras.Model):
"""
fused operation: scaling + mask + softmax
Arguments:
attn_mask_type: attention mask type (pad or causal)
mask_func: mask function to be applied.
softmax_in_fp32: if true, softmax in performed at fp32 precision.
scale: scaling factor used in input tensor scaling.
"""
def __init__(
self,
attn_mask_type: str,
mask_func: Callable,
softmax_in_fp32: bool,
scale: float,
) -> None:
super().__init__()
self.attn_mask_type = attn_mask_type
self.scaled_masked_softmax_fusion = bool(
int(os.getenv("NVTE_MASKED_SOFTMAX_FUSION", "1"))
)
self.mask_func = mask_func
self.softmax_in_fp32 = softmax_in_fp32
self.scale = scale
self.stream = get_stream_id()
assert (
self.scale is None or softmax_in_fp32
), "softmax should be in fp32 when scaled"
def __call__(self, inp: tf.Tensor, mask: tf.Tensor) -> tf.Tensor:
"""FusedScaleMaskSoftmax fprop"""
# [b, np, sq, sk]
assert len(inp.shape) == 4
self.input_in_fp16 = inp.dtype == tf.float16
self.input_in_bf16 = inp.dtype == tf.bfloat16
self.input_in_float16 = self.input_in_fp16 or self.input_in_bf16
if self.is_kernel_available(*inp.shape):
return self.forward_fused_softmax(inp, mask)
return self.forward_tf_softmax(inp, mask)
def is_kernel_available(self, b: int, np: int, sq: int, sk: int) -> bool:
"""Check FusedScaleMaskSoftmax kernel availability based on size"""
attn_batches = b * np
if (
self.scaled_masked_softmax_fusion # user want to fuse
and self.input_in_float16 # input must be fp16
and 16 < sk <= 4096 # sk must be 16 ~ 2048
and sq % 4 == 0 # sq must be divisor of 4
and attn_batches % 4 == 0 # np * b must be divisor of 4
):
if 0 <= sk <= 4096:
batch_per_block = self.get_batch_per_block(int(sk))
if self.attn_mask_type == "causal":
if attn_batches % batch_per_block == 0:
return True
else:
if sq % batch_per_block == 0:
return True
return False
@tf.custom_gradient
def scaled_masked_softmax(self, x: tf.Tensor, mask: tf.Tensor,
scale: float):
"""Scaled masked softmax."""
y = tex.scaled_masked_softmax_forward(x, mask, scale, self.stream)
def grad_fn(upstream):
dx = tex.scaled_masked_softmax_backward(upstream, y, scale,
self.stream)
return dx, None, None
return y, grad_fn
@tf.custom_gradient
def scaled_softmax(self, x: tf.Tensor, scale: float):
"""Scaled softmax."""
y = tex.scaled_softmax_forward(x, scale, self.stream)
def grad_fn(upstream):
dx = tex.scaled_softmax_backward(upstream, y, scale, self.stream)
return dx, None
return y, grad_fn
@tf.custom_gradient
def scaled_upper_triang_masked_softmax(self, x: tf.Tensor, scale: float):
"""Scaled upper triangular masked softmax."""
y = tex.scaled_upper_triang_masked_softmax_forward(x, scale,
self.stream)
def grad_fn(upstream):
dx = tex.scaled_upper_triang_masked_softmax_backward(
upstream, y, scale, self.stream
)
return dx, None
return y, grad_fn
def forward_fused_softmax(
self,
inp: tf.Tensor,
mask: tf.Tensor,
) -> tf.Tensor:
"""Fused masked softmax kernel"""
sq, sk = inp.shape[2], inp.shape[3]
scale = self.scale if self.scale is not None else 1.0
if self.attn_mask_type == "causal":
assert sq == sk, "causal mask is only for self attention"
# input is 3D tensor (attn_batches, sq, sk)
inp = tf.reshape(inp, (-1, sq, sk))
probs = self.scaled_upper_triang_masked_softmax(inp, scale)
return tf.reshape(probs, inp.shape)
# input is 4D tensor (b, np, sq, sk)
if mask is not None:
# The mask defined in TE kernels are different from TF. In TE, the
# mask specifies 1 to mask out and 0 to keep.
mask = tf.math.logical_not(mask)
ndims = len(mask.shape)
assert ndims <= 4, "mask ndims should be <= 4"
if len(mask.shape) < 4:
# Broadcasting the first dims of mask to match the input ndims.
broadcast_shape = [1] * (4 - ndims) + mask.shape[:]
mask = tf.reshape(mask, broadcast_shape)
return self.scaled_masked_softmax(inp, mask, scale)
return self.scaled_softmax(inp, scale)
def forward_tf_softmax(self, inp: tf.Tensor, mask: tf.Tensor) -> tf.Tensor:
"""Framework softmax"""
if self.input_in_float16 and self.softmax_in_fp32:
inp = tf.cast(inp, tf.float32)
if self.scale is not None:
inp = inp * self.scale
if self.attn_mask_type == "causal":
mask = _get_default_causal_mask(inp.shape[2])
mask_output = self.mask_func(inp, mask) if mask is not None else inp
probs = tf.nn.softmax(mask_output, axis=-1)
if self.input_in_float16 and self.softmax_in_fp32:
if self.input_in_fp16:
probs = tf.cast(probs, tf.half)
else:
probs = tf.cast(probs, tf.bfloat16)
return probs
@staticmethod
def get_batch_per_block(key_seq_len: int) -> int:
"""Softmax utility"""
pow2 = 1 << (key_seq_len - 1).bit_length()
warp_size = pow2 if pow2 < THREADS_PER_WARP else THREADS_PER_WARP
batches_per_warp = 2 if pow2 <= 128 else 1
warps_per_block = THREADS_PER_BLOCK / warp_size
batches_per_block = warps_per_block * batches_per_warp
return batches_per_block
transformer_engine/tensorflow/transformer.py deleted 100644 → 0
View file @ 2574a1ca
This diff is collapsed.
transformer_engine/tensorflow/utils.py deleted 100644 → 0
View file @ 2574a1ca
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
"""Utility functions for Transformer Engine modules"""
import tensorflow as tf
def attention_mask_func(
attention_scores: tf.Tensor, attention_mask: tf.Tensor
) -> tf.Tensor:
"""Get attention mask"""
return tf.where(attention_mask, attention_scores, -10000.0)
def ensure_divisibility(numerator: int, denominator: int) -> None:
"""Ensure that numerator is divisible by the denominator."""
assert (
numerator % denominator == 0
), f"{numerator} is not divisible by {denominator}"
def divide(numerator: int, denominator: int) -> int:
"""Ensure that numerator is divisible by the denominator and return
the division value."""
ensure_divisibility(numerator, denominator)
return numerator // denominator
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment