Skip to content

Amax reduction interval#154

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 4 commits into from
Apr 18, 2023
+73 −8
Merged

Amax reduction interval #154

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
629b838
amax reduction internval
erhoo82 Apr 14, 2023
0db4d6a
Update transformer_engine/pytorch/fp8.py
erhoo82 Apr 17, 2023
dd4b8f6
check TP group initialized
erhoo82 Apr 17, 2023
9bacf57
Merge branch 'main' into slym/amax_reduce_interval
ksivaman Apr 18, 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
33 changes: 32 additions & 1 deletion transformer_engine/pytorch/fp8.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
fp8_meta: Dict[str, Any],
tp_group: dist_group_type,
tp_size: int,
forward: bool = True,
) -> None:
"""Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
if amax_buffer_key not in _global_fp8_buffer:
return None

# Reduce AMAX in DP-domain at an interval.
global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
if _dp_amax_reduce_interval is None:
_dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

tp_amax_reduce = False
if forward:
if _dp_amax_reduce_forward_idx == 0:
reduce_group = fp8_meta["fp8_group"]
else:
tp_amax_reduce = True
_dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
else:
if _dp_amax_reduce_backward_idx == 0:
reduce_group = fp8_meta["fp8_group"]
else:
tp_amax_reduce = True
_dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

if tp_amax_reduce:
if tp_size > 1:
reduce_group = tp_group
else:
return None

chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

wait_handle = reduce_tensor_across_group_op_max(
contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
fp8_meta["async_amax_reduction"],
)

Expand Down
48 changes: 41 additions & 7 deletions transformer_engine/pytorch/module.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
fp8: bool,
fp8_meta: Dict[str, Any],
tp_group: dist_group_type,
tp_size: int,
name: str = ""
) -> None:
"""Checks and prep for BWD."""
if fp8:
global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

if fp8 and fp8_meta["recipe"].reduce_amax:
if fp8_meta["first_module"]:
_amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
_amax_reduce_handle_bwd = global_amax_reduction(
fp8_meta,
tp_group,
tp_size,
forward=False
)
delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
self.fp8_meta["recipe"] = get_default_fp8_recipe()
self.fp8_meta_tensors_initialized = False
self.tp_group = None
self.tp_group_initialized = False
self.tp_size = 1
self.sequence_parallel = False
self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
reduce_func = partial(
global_amax_reduction,
self.fp8_meta,
self.tp_group,
self.tp_size,
forward=True
)
setup_amax_forward_global_reduce_func(reduce_func)

def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
fp8_meta: Dict[str, Any],
fuse_wgrad_accumulation: bool,
tp_group: Union[dist_group_type, None],
tp_size: int,
sequence_parallel: bool,
tensor_parallel: bool,
activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
ctx.inp_shape = inp.shape
ctx.parallel_mode = parallel_mode
ctx.tp_group = tp_group
ctx.tp_size = tp_size
ctx.return_layernorm_output = return_layernorm_output
ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
def backward(
ctx, *grad_outputs: Tuple[torch.Tensor, ...]
) -> Tuple[Union[torch.Tensor, None], ...]:
with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
with _prepare_backward(
ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
):
(
inputmat,
ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
None,
None,
None,
None,
)


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
self.fp8_meta,
self.fuse_wgrad_accumulation,
self.tp_group,
self.tp_size,
self.sequence_parallel,
self.tp_size > 1,
self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
fp8_meta: Dict[str, Any],
fuse_wgrad_accumulation: bool,
tp_group: Union[dist_group_type, None],
tp_size: int,
sequence_parallel: bool,
tensor_parallel: bool,
activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
ctx.inp_shape = inp.shape
ctx.parallel_mode = parallel_mode
ctx.tp_group = tp_group
ctx.tp_size = tp_size
ctx.requires_dgrad = inp.requires_grad

# Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
def backward(
ctx, grad_output: torch.Tensor
) -> Tuple[Union[torch.Tensor, None], ...]:
with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
with _prepare_backward(
ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
):
(
inputmat,
inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
None,
None,
None,
None,
)


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
self.fp8_meta,
self.fuse_wgrad_accumulation,
self.tp_group,
self.tp_size,
self.sequence_parallel,
self.tp_size > 1,
self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
fp8_meta: Dict[str, Any],
fuse_wgrad_accumulation: bool,
tp_group: Union[dist_group_type, None],
tp_size: int,
sequence_parallel: bool,
tensor_parallel: bool,
activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
ctx.tensor_parallel = tensor_parallel
ctx.inp_shape = inp.shape
ctx.tp_group = tp_group
ctx.tp_size = tp_size
ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
ctx.return_layernorm_output = return_layernorm_output
ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
def backward(
ctx, *grad_outputs: Tuple[torch.Tensor, ...]
) -> Tuple[Union[torch.Tensor, None], ...]:
with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
with _prepare_backward(
ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
):
(
inputmat,
ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
None,
None,
None,
None,
)


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
self.fp8_meta,
self.fuse_wgrad_accumulation,
self.tp_group,
self.tp_size,
self.sequence_parallel,
self.tp_size > 1,
self.activation_dtype,
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" + '
Amax reduction interval by erhoo82 · Pull Request #154 · NVIDIA/TransformerEngine · GitHub
Skip to content
Amax reduction interval#154

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 4 commits into

from

Apr 18, 2023

+73

−8

 Merged


Amax reduction interval
#154
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
629b838

amax reduction internval


erhoo82 Apr 14, 2023
0db4d6a

Update transformer_engine/pytorch/fp8.py


erhoo82 Apr 17, 2023
dd4b8f6

check TP group initialized


erhoo82 Apr 17, 2023
9bacf57

Merge branch 'main' into slym/amax_reduce_interval


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


33 changes: 32 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 forward: bool = True,
) -> None:
 """Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
 if amax_buffer_key not in _global_fp8_buffer:
 return None

 # Reduce AMAX in DP-domain at an interval.
 global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
 if _dp_amax_reduce_interval is None:
 _dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

 tp_amax_reduce = False
 if forward:
 if _dp_amax_reduce_forward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
 else:
 if _dp_amax_reduce_backward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

 if tp_amax_reduce:
 if tp_size > 1:
 reduce_group = tp_group
 else:
 return None

 chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
 contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

 wait_handle = reduce_tensor_across_group_op_max(
 contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
 fp8_meta["async_amax_reduction"],
 )

Expand Down

48 changes: 41 additions & 7 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
 return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
 fp8: bool,
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 name: str = ""
) -> None:
 """Checks and prep for BWD."""
 if fp8:
 global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

 if fp8 and fp8_meta["recipe"].reduce_amax:
 if fp8_meta["first_module"]:
 _amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
 _amax_reduce_handle_bwd = global_amax_reduction(
 fp8_meta,
 tp_group,
 tp_size,
 forward=False
 )
 delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
 self.fp8_meta["recipe"] = get_default_fp8_recipe()
 self.fp8_meta_tensors_initialized = False
 self.tp_group = None
 self.tp_group_initialized = False
 self.tp_size = 1
 self.sequence_parallel = False
 self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

 if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
 set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
 reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
 reduce_func = partial(
 global_amax_reduction,
 self.fp8_meta,
 self.tp_group,
 self.tp_size,
 forward=True
 )
 setup_amax_forward_global_reduce_func(reduce_func)

 def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.return_layernorm_output = return_layernorm_output
 ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
 ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.requires_dgrad = inp.requires_grad

 # Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
 def backward(
 ctx, grad_output: torch.Tensor
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
 ):
 (
 inputmat,
 inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
 ctx.tensor_parallel = tensor_parallel
 ctx.inp_shape = inp.shape
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
 ctx.return_layernorm_output = return_layernorm_output
 ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
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('^' + ".*" + '
Amax reduction interval by erhoo82 · Pull Request #154 · NVIDIA/TransformerEngine · GitHub
Skip to content
Amax reduction interval#154

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 4 commits into

from

Apr 18, 2023

+73

−8

 Merged


Amax reduction interval
#154
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
629b838

amax reduction internval


erhoo82 Apr 14, 2023
0db4d6a

Update transformer_engine/pytorch/fp8.py


erhoo82 Apr 17, 2023
dd4b8f6

check TP group initialized


erhoo82 Apr 17, 2023
9bacf57

Merge branch 'main' into slym/amax_reduce_interval


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


33 changes: 32 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 forward: bool = True,
) -> None:
 """Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
 if amax_buffer_key not in _global_fp8_buffer:
 return None

 # Reduce AMAX in DP-domain at an interval.
 global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
 if _dp_amax_reduce_interval is None:
 _dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

 tp_amax_reduce = False
 if forward:
 if _dp_amax_reduce_forward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
 else:
 if _dp_amax_reduce_backward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

 if tp_amax_reduce:
 if tp_size > 1:
 reduce_group = tp_group
 else:
 return None

 chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
 contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

 wait_handle = reduce_tensor_across_group_op_max(
 contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
 fp8_meta["async_amax_reduction"],
 )

Expand Down

48 changes: 41 additions & 7 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
 return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
 fp8: bool,
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 name: str = ""
) -> None:
 """Checks and prep for BWD."""
 if fp8:
 global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

 if fp8 and fp8_meta["recipe"].reduce_amax:
 if fp8_meta["first_module"]:
 _amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
 _amax_reduce_handle_bwd = global_amax_reduction(
 fp8_meta,
 tp_group,
 tp_size,
 forward=False
 )
 delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
 self.fp8_meta["recipe"] = get_default_fp8_recipe()
 self.fp8_meta_tensors_initialized = False
 self.tp_group = None
 self.tp_group_initialized = False
 self.tp_size = 1
 self.sequence_parallel = False
 self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

 if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
 set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
 reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
 reduce_func = partial(
 global_amax_reduction,
 self.fp8_meta,
 self.tp_group,
 self.tp_size,
 forward=True
 )
 setup_amax_forward_global_reduce_func(reduce_func)

 def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.return_layernorm_output = return_layernorm_output
 ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
 ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.requires_dgrad = inp.requires_grad

 # Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
 def backward(
 ctx, grad_output: torch.Tensor
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
 ):
 (
 inputmat,
 inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
 ctx.tensor_parallel = tensor_parallel
 ctx.inp_shape = inp.shape
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
 ctx.return_layernorm_output = return_layernorm_output
 ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
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('^' + ".*" + '
Amax reduction interval by erhoo82 · Pull Request #154 · NVIDIA/TransformerEngine · GitHub
Skip to content
Amax reduction interval#154

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 4 commits into

from

Apr 18, 2023

+73

−8

 Merged


Amax reduction interval
#154
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
629b838

amax reduction internval


erhoo82 Apr 14, 2023
0db4d6a

Update transformer_engine/pytorch/fp8.py


erhoo82 Apr 17, 2023
dd4b8f6

check TP group initialized


erhoo82 Apr 17, 2023
9bacf57

Merge branch 'main' into slym/amax_reduce_interval


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


33 changes: 32 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 forward: bool = True,
) -> None:
 """Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
 if amax_buffer_key not in _global_fp8_buffer:
 return None

 # Reduce AMAX in DP-domain at an interval.
 global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
 if _dp_amax_reduce_interval is None:
 _dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

 tp_amax_reduce = False
 if forward:
 if _dp_amax_reduce_forward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
 else:
 if _dp_amax_reduce_backward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

 if tp_amax_reduce:
 if tp_size > 1:
 reduce_group = tp_group
 else:
 return None

 chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
 contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

 wait_handle = reduce_tensor_across_group_op_max(
 contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
 fp8_meta["async_amax_reduction"],
 )

Expand Down

48 changes: 41 additions & 7 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
 return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
 fp8: bool,
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 name: str = ""
) -> None:
 """Checks and prep for BWD."""
 if fp8:
 global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

 if fp8 and fp8_meta["recipe"].reduce_amax:
 if fp8_meta["first_module"]:
 _amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
 _amax_reduce_handle_bwd = global_amax_reduction(
 fp8_meta,
 tp_group,
 tp_size,
 forward=False
 )
 delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
 self.fp8_meta["recipe"] = get_default_fp8_recipe()
 self.fp8_meta_tensors_initialized = False
 self.tp_group = None
 self.tp_group_initialized = False
 self.tp_size = 1
 self.sequence_parallel = False
 self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

 if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
 set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
 reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
 reduce_func = partial(
 global_amax_reduction,
 self.fp8_meta,
 self.tp_group,
 self.tp_size,
 forward=True
 )
 setup_amax_forward_global_reduce_func(reduce_func)

 def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.return_layernorm_output = return_layernorm_output
 ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
 ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.requires_dgrad = inp.requires_grad

 # Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
 def backward(
 ctx, grad_output: torch.Tensor
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
 ):
 (
 inputmat,
 inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
 ctx.tensor_parallel = tensor_parallel
 ctx.inp_shape = inp.shape
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
 ctx.return_layernorm_output = return_layernorm_output
 ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
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" + '
Amax reduction interval by erhoo82 · Pull Request #154 · NVIDIA/TransformerEngine · GitHub
Skip to content
Amax reduction interval#154

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 4 commits into

from

Apr 18, 2023

+73

−8

 Merged


Amax reduction interval
#154
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
629b838

amax reduction internval


erhoo82 Apr 14, 2023
0db4d6a

Update transformer_engine/pytorch/fp8.py


erhoo82 Apr 17, 2023
dd4b8f6

check TP group initialized


erhoo82 Apr 17, 2023
9bacf57

Merge branch 'main' into slym/amax_reduce_interval


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


33 changes: 32 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 forward: bool = True,
) -> None:
 """Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
 if amax_buffer_key not in _global_fp8_buffer:
 return None

 # Reduce AMAX in DP-domain at an interval.
 global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
 if _dp_amax_reduce_interval is None:
 _dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

 tp_amax_reduce = False
 if forward:
 if _dp_amax_reduce_forward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
 else:
 if _dp_amax_reduce_backward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

 if tp_amax_reduce:
 if tp_size > 1:
 reduce_group = tp_group
 else:
 return None

 chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
 contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

 wait_handle = reduce_tensor_across_group_op_max(
 contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
 fp8_meta["async_amax_reduction"],
 )

Expand Down

48 changes: 41 additions & 7 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
 return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
 fp8: bool,
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 name: str = ""
) -> None:
 """Checks and prep for BWD."""
 if fp8:
 global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

 if fp8 and fp8_meta["recipe"].reduce_amax:
 if fp8_meta["first_module"]:
 _amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
 _amax_reduce_handle_bwd = global_amax_reduction(
 fp8_meta,
 tp_group,
 tp_size,
 forward=False
 )
 delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
 self.fp8_meta["recipe"] = get_default_fp8_recipe()
 self.fp8_meta_tensors_initialized = False
 self.tp_group = None
 self.tp_group_initialized = False
 self.tp_size = 1
 self.sequence_parallel = False
 self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

 if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
 set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
 reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
 reduce_func = partial(
 global_amax_reduction,
 self.fp8_meta,
 self.tp_group,
 self.tp_size,
 forward=True
 )
 setup_amax_forward_global_reduce_func(reduce_func)

 def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.return_layernorm_output = return_layernorm_output
 ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
 ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.requires_dgrad = inp.requires_grad

 # Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
 def backward(
 ctx, grad_output: torch.Tensor
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
 ):
 (
 inputmat,
 inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
 ctx.tensor_parallel = tensor_parallel
 ctx.inp_shape = inp.shape
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
 ctx.return_layernorm_output = return_layernorm_output
 ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
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('^' + ".*" + '
Amax reduction interval by erhoo82 · Pull Request #154 · NVIDIA/TransformerEngine · GitHub
Skip to content
Amax reduction interval#154

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 4 commits into

from

Apr 18, 2023

+73

−8

 Merged


Amax reduction interval
#154
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
629b838

amax reduction internval


erhoo82 Apr 14, 2023
0db4d6a

Update transformer_engine/pytorch/fp8.py


erhoo82 Apr 17, 2023
dd4b8f6

check TP group initialized


erhoo82 Apr 17, 2023
9bacf57

Merge branch 'main' into slym/amax_reduce_interval


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


33 changes: 32 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 forward: bool = True,
) -> None:
 """Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
 if amax_buffer_key not in _global_fp8_buffer:
 return None

 # Reduce AMAX in DP-domain at an interval.
 global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
 if _dp_amax_reduce_interval is None:
 _dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

 tp_amax_reduce = False
 if forward:
 if _dp_amax_reduce_forward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
 else:
 if _dp_amax_reduce_backward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

 if tp_amax_reduce:
 if tp_size > 1:
 reduce_group = tp_group
 else:
 return None

 chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
 contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

 wait_handle = reduce_tensor_across_group_op_max(
 contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
 fp8_meta["async_amax_reduction"],
 )

Expand Down

48 changes: 41 additions & 7 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
 return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
 fp8: bool,
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 name: str = ""
) -> None:
 """Checks and prep for BWD."""
 if fp8:
 global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

 if fp8 and fp8_meta["recipe"].reduce_amax:
 if fp8_meta["first_module"]:
 _amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
 _amax_reduce_handle_bwd = global_amax_reduction(
 fp8_meta,
 tp_group,
 tp_size,
 forward=False
 )
 delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
 self.fp8_meta["recipe"] = get_default_fp8_recipe()
 self.fp8_meta_tensors_initialized = False
 self.tp_group = None
 self.tp_group_initialized = False
 self.tp_size = 1
 self.sequence_parallel = False
 self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

 if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
 set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
 reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
 reduce_func = partial(
 global_amax_reduction,
 self.fp8_meta,
 self.tp_group,
 self.tp_size,
 forward=True
 )
 setup_amax_forward_global_reduce_func(reduce_func)

 def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.return_layernorm_output = return_layernorm_output
 ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
 ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.requires_dgrad = inp.requires_grad

 # Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
 def backward(
 ctx, grad_output: torch.Tensor
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
 ):
 (
 inputmat,
 inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
 ctx.tensor_parallel = tensor_parallel
 ctx.inp_shape = inp.shape
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
 ctx.return_layernorm_output = return_layernorm_output
 ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
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('^' + ".*" + '
Amax reduction interval by erhoo82 · Pull Request #154 · NVIDIA/TransformerEngine · GitHub
Skip to content
Amax reduction interval#154

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 4 commits into

from

Apr 18, 2023

+73

−8

 Merged


Amax reduction interval
#154
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
629b838

amax reduction internval


erhoo82 Apr 14, 2023
0db4d6a

Update transformer_engine/pytorch/fp8.py


erhoo82 Apr 17, 2023
dd4b8f6

check TP group initialized


erhoo82 Apr 17, 2023
9bacf57

Merge branch 'main' into slym/amax_reduce_interval


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


33 changes: 32 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 forward: bool = True,
) -> None:
 """Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
 if amax_buffer_key not in _global_fp8_buffer:
 return None

 # Reduce AMAX in DP-domain at an interval.
 global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
 if _dp_amax_reduce_interval is None:
 _dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

 tp_amax_reduce = False
 if forward:
 if _dp_amax_reduce_forward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
 else:
 if _dp_amax_reduce_backward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

 if tp_amax_reduce:
 if tp_size > 1:
 reduce_group = tp_group
 else:
 return None

 chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
 contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

 wait_handle = reduce_tensor_across_group_op_max(
 contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
 fp8_meta["async_amax_reduction"],
 )

Expand Down

48 changes: 41 additions & 7 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
 return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
 fp8: bool,
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 name: str = ""
) -> None:
 """Checks and prep for BWD."""
 if fp8:
 global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

 if fp8 and fp8_meta["recipe"].reduce_amax:
 if fp8_meta["first_module"]:
 _amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
 _amax_reduce_handle_bwd = global_amax_reduction(
 fp8_meta,
 tp_group,
 tp_size,
 forward=False
 )
 delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
 self.fp8_meta["recipe"] = get_default_fp8_recipe()
 self.fp8_meta_tensors_initialized = False
 self.tp_group = None
 self.tp_group_initialized = False
 self.tp_size = 1
 self.sequence_parallel = False
 self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

 if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
 set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
 reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
 reduce_func = partial(
 global_amax_reduction,
 self.fp8_meta,
 self.tp_group,
 self.tp_size,
 forward=True
 )
 setup_amax_forward_global_reduce_func(reduce_func)

 def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.return_layernorm_output = return_layernorm_output
 ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
 ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.requires_dgrad = inp.requires_grad

 # Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
 def backward(
 ctx, grad_output: torch.Tensor
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
 ):
 (
 inputmat,
 inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
 ctx.tensor_parallel = tensor_parallel
 ctx.inp_shape = inp.shape
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
 ctx.return_layernorm_output = return_layernorm_output
 ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
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); }
})();
})();
Amax reduction interval by erhoo82 · Pull Request #154 · NVIDIA/TransformerEngine · GitHub
Skip to content
Amax reduction interval#154

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 4 commits into

from

Apr 18, 2023

+73

−8

 Merged


Amax reduction interval
#154
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


4 commits


Select commit
Hold shift + click to select a range
629b838

amax reduction internval


erhoo82 Apr 14, 2023
0db4d6a

Update transformer_engine/pytorch/fp8.py


erhoo82 Apr 17, 2023
dd4b8f6

check TP group initialized


erhoo82 Apr 17, 2023
9bacf57

Merge branch 'main' into slym/amax_reduce_interval


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


33 changes: 32 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
# See LICENSE for license information.

"""FP8 utilities for TransformerEngine"""
import os
from contextlib import contextmanager
from collections import deque
from typing import Callable, List, Optional, Dict, Any, Tuple, Union
Expand DownExpand Up@@ -30,6 +31,9 @@
_amax_reduce_handle_fwd = None
_is_fp8_available = None
_reason_for_no_fp8 = ""
_dp_amax_reduce_interval = None
_dp_amax_reduce_forward_idx = 0
_dp_amax_reduce_backward_idx = 0


def _check_fp8_support() -> Tuple[bool, str]:
Expand DownExpand Up@@ -545,6 +549,8 @@ def reduce_tensor_across_group_op_max(

def global_amax_reduction(
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 forward: bool = True,
) -> None:
 """Concatenate, reduce, and split amaxes in the global buffer."""
Expand All@@ -555,12 +561,37 @@ def global_amax_reduction(
 if amax_buffer_key not in _global_fp8_buffer:
 return None

 # Reduce AMAX in DP-domain at an interval.
 global _dp_amax_reduce_interval, _dp_amax_reduce_forward_idx, _dp_amax_reduce_backward_idx
 if _dp_amax_reduce_interval is None:
 _dp_amax_reduce_interval = int(os.getenv("NVTE_DP_AMAX_REDUCE_INTERVAL", "1"))

 tp_amax_reduce = False
 if forward:
 if _dp_amax_reduce_forward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_forward_idx = (_dp_amax_reduce_forward_idx + 1) % _dp_amax_reduce_interval
 else:
 if _dp_amax_reduce_backward_idx == 0:
 reduce_group = fp8_meta["fp8_group"]
 else:
 tp_amax_reduce = True
 _dp_amax_reduce_backward_idx = (_dp_amax_reduce_backward_idx + 1) % _dp_amax_reduce_interval

 if tp_amax_reduce:
 if tp_size > 1:
 reduce_group = tp_group
 else:
 return None

 chunk_sizes = [x.numel() for x in _global_fp8_buffer[amax_buffer_key]]
 contiguous_amax = torch.cat(_global_fp8_buffer[amax_buffer_key])

 wait_handle = reduce_tensor_across_group_op_max(
 contiguous_amax,
fp8_meta["fp8_group"],
reduce_group,
 fp8_meta["async_amax_reduction"],
 )

Expand Down

48 changes: 41 additions & 7 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,7 +105,13 @@ def get_workspace() -> torch.Tensor:
 return _cublas_workspace

@contextmanager
def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> None:
def _prepare_backward(
 fp8: bool,
 fp8_meta: Dict[str, Any],
 tp_group: dist_group_type,
 tp_size: int,
 name: str = ""
) -> None:
 """Checks and prep for BWD."""
 if fp8:
 global _amax_reduce_handle_bwd
Expand All@@ -132,7 +138,12 @@ def _prepare_backward(fp8: bool, fp8_meta: Dict[str, Any], name: str = "") -> N

 if fp8 and fp8_meta["recipe"].reduce_amax:
 if fp8_meta["first_module"]:
 _amax_reduce_handle_bwd = global_amax_reduction(fp8_meta, forward=False)
 _amax_reduce_handle_bwd = global_amax_reduction(
 fp8_meta,
 tp_group,
 tp_size,
 forward=False
 )
 delete_key_from_amax_buffer(forward=False)


Expand DownExpand Up@@ -186,7 +197,6 @@ def __init__(self) -> None:
 self.fp8_meta["recipe"] = get_default_fp8_recipe()
 self.fp8_meta_tensors_initialized = False
 self.tp_group = None
 self.tp_group_initialized = False
 self.tp_size = 1
 self.sequence_parallel = False
 self.fp8_weight_shapes = []
Expand DownExpand Up@@ -541,7 +551,13 @@ def prepare_forward(

 if self.fp8 and self.training and self.fp8_meta["recipe"].reduce_amax:
 set_fp8_context_id(self.fp8_meta["autocast_id_fwd"])
 reduce_func = partial(global_amax_reduction, self.fp8_meta, forward=True)
 reduce_func = partial(
 global_amax_reduction,
 self.fp8_meta,
 self.tp_group,
 self.tp_size,
 forward=True
 )
 setup_amax_forward_global_reduce_func(reduce_func)

 def set_nccl_overlap_warning_if_tp(self) -> None:
Expand DownExpand Up@@ -692,6 +708,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -867,6 +884,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.return_layernorm_output = return_layernorm_output
 ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
 ctx.zero_centered_gamma = zero_centered_gamma
Expand All@@ -890,7 +908,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormLinear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormLinear"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -1065,6 +1085,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1381,6 +1402,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -1427,6 +1449,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -1563,6 +1586,7 @@ def forward(
 ctx.inp_shape = inp.shape
 ctx.parallel_mode = parallel_mode
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.requires_dgrad = inp.requires_grad

 # Row Parallel Linear
Expand All@@ -1579,7 +1603,9 @@ def forward(
 def backward(
 ctx, grad_output: torch.Tensor
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_Linear"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_Linear"
 ):
 (
 inputmat,
 inputmat_t,
Expand DownExpand Up@@ -1730,6 +1756,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -1995,6 +2022,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand DownExpand Up@@ -2039,6 +2067,7 @@ def forward(
 fp8_meta: Dict[str, Any],
 fuse_wgrad_accumulation: bool,
 tp_group: Union[dist_group_type, None],
 tp_size: int,
 sequence_parallel: bool,
 tensor_parallel: bool,
 activation_dtype: torch.dtype,
Expand DownExpand Up@@ -2282,6 +2311,7 @@ def forward(
 ctx.tensor_parallel = tensor_parallel
 ctx.inp_shape = inp.shape
 ctx.tp_group = tp_group
 ctx.tp_size = tp_size
 ctx.bias_gelu_nvfusion = bias_gelu_nvfusion
 ctx.return_layernorm_output = return_layernorm_output
 ctx.set_parallel_mode = set_parallel_mode
Expand All@@ -2307,7 +2337,9 @@ def forward(
 def backward(
 ctx, *grad_outputs: Tuple[torch.Tensor, ...]
 ) -> Tuple[Union[torch.Tensor, None], ...]:
 with _prepare_backward(ctx.fp8, ctx.fp8_meta, name="_LayerNormMLP"):
 with _prepare_backward(
 ctx.fp8, ctx.fp8_meta, ctx.tp_group, ctx.tp_size, name="_LayerNormMLP"
 ):
 (
 inputmat,
 ln_weight,
Expand DownExpand Up@@ -2610,6 +2642,7 @@ def backward(
 None,
 None,
 None,
 None,
 )


Expand DownExpand Up@@ -2904,6 +2937,7 @@ def forward(
 self.fp8_meta,
 self.fuse_wgrad_accumulation,
 self.tp_group,
 self.tp_size,
 self.sequence_parallel,
 self.tp_size > 1,
 self.activation_dtype,
Expand Down