Skip to content

Change FP8 recipe defaults#112

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 10 commits into from
Mar 30, 2023
Merged

Change FP8 recipe defaults #112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
23b172d
Change FP8 recipe defaults
ksivaman Mar 21, 2023
112e54a
Merge branch 'main' into change_fp8_defaults
ksivaman Mar 22, 2023
7bb9898
Increase default amax history length
ksivaman Mar 22, 2023
e5198f3
Always check history size
ksivaman Mar 22, 2023
bedea6b
no amax history for onnx export
ksivaman Mar 23, 2023
ab2d1d7
Merge branch 'main' into change_fp8_defaults
ksivaman Mar 23, 2023
e894a6e
revert onnx export test changes
ksivaman Mar 24, 2023
a48963f
Merge branch 'main' into change_fp8_defaults
ksivaman Mar 28, 2023
de1eb75
Merge branch 'main' into change_fp8_defaults
ksivaman Mar 29, 2023
48bb4df
Fix indices in onnx test
ksivaman Mar 29, 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
17 changes: 10 additions & 7 deletions tests/pytorch/test_onnx_export.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -105,12 +105,15 @@ def to_numpy(tensor):
return tensor.cpu().numpy()


def set_layer_scale(module: torch.nn.Module, scale: float):
module.fp8_init()
def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
"""Initialize the FP8 quantization scales in module"""
NB_SCALES_PER_GEMM = 3 # One scale per: input, weights, and output GEMM tensors.
nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
module.fp8_init(num_gemms)
module.fp8_meta["scaling_fwd"].scale = torch.ones(
2, dtype=torch.float32, device="cuda") / scale
nb_total_scales, dtype=torch.float32, device="cuda") / scale
module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
2, dtype=torch.float32, device="cuda") * scale
nb_total_scales, dtype=torch.float32, device="cuda") * scale


def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
Expand DownExpand Up@@ -678,7 +681,7 @@ def forward(self, inp):
precision
).to(device='cuda')
if use_fp8:
set_layer_scale(model.linear, scale_factor)
set_layer_scale(model.linear, scale_factor, num_gemms=1)
do_export(model, inp, fname, use_fp8)

if precision in (torch.bfloat16, ):
Expand DownExpand Up@@ -736,7 +739,7 @@ def test_export_layernorm_linear(
zero_centered_gamma=zero_centered_gamma,
).to(device='cuda')
if use_fp8:
set_layer_scale(model, scale_factor)
set_layer_scale(model, scale_factor, num_gemms=1)
do_export(model, inp, fname, use_fp8)
if not use_fp8:
validate_result(fname, inp, model, atol=1e-3)
Expand DownExpand Up@@ -792,7 +795,7 @@ def test_export_layernorm_mlp(
zero_centered_gamma=zero_centered_gamma,
).to(device='cuda')
if use_fp8:
set_layer_scale(model, scale_factor)
set_layer_scale(model, scale_factor, num_gemms=2)
do_export(model, inp, fname, use_fp8)
if not use_fp8:
validate_result(fname, inp, model, atol=1e-3)
Expand Down
8 changes: 4 additions & 4 deletions transformer_engine/common/recipe.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,10 @@ class DelayedScaling:
fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID
Controls the FP8 data format used during forward and backward
pass.
amax_history_len : int, default = 1
amax_history_len : int, default = 1024
The length of the amax history window used for
scaling factor computation.
amax_compute_algo : {'max', 'most_recent', Callable}, default = 'most_recent'
amax_compute_algo : {'max', 'most_recent', Callable}, default = 'max'
Algorithm used for choosing the `amax` value for the
scaling factor computation. There are 2 predefined
choices: `max` chooses the largest `amax` in the history
Expand DownExpand Up@@ -125,8 +125,8 @@ def scaling_factor_compute(amax: Tensor,
margin: int = 0
interval: int = 1
fp8_format: Format = Format.HYBRID
amax_history_len: int = 1
amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "most_recent"
amax_history_len: int = 1024
amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "max"
override_linear_precision: _OverrideLinearPrecision = _OverrideLinearPrecision()
scaling_factor_compute_algo: Optional[Callable] = None
reduce_amax: bool = True
Expand Down
25 changes: 19 additions & 6 deletions transformer_engine/pytorch/module.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -13,6 +13,7 @@

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init

Expand DownExpand Up@@ -187,6 +188,23 @@ def __init__(self) -> None:
def set_meta_tensor(self, fwd: bool) -> None:
"""Init scales and amaxes for fwd | bwd."""
fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

if self.fp8_meta_tensors_initialized:
# Handle changed amax history size.
curr_len = self.fp8_meta[fp8_meta_tensor_key].amax_history.shape[0]
need_len = self.fp8_meta["recipe"].amax_history_len
if need_len < curr_len:
self.fp8_meta[fp8_meta_tensor_key].amax_history = (
self.fp8_meta[fp8_meta_tensor_key]
.amax_history[: self.fp8_meta["recipe"].amax_history_len].clone()
)
elif need_len > curr_len:
extra_rows = need_len - curr_len
self.fp8_meta[fp8_meta_tensor_key].amax_history = F.pad(
self.fp8_meta[fp8_meta_tensor_key].amax_history, pad=(0, 0, 0, extra_rows)
)
return

# Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
# 2 (grad_output and grad_input) for bwd
num_fp8_tensors = (
Expand DownExpand Up@@ -222,12 +240,9 @@ def set_meta_tensor(self, fwd: bool) -> None:

def init_fp8_meta_tensors(self) -> None:
"""Init scales and amaxes."""
# Checkpoint loaded
if self.fp8_meta_tensors_initialized:
return

self.set_meta_tensor(True)
self.set_meta_tensor(False)
self.fp8_meta_tensors_initialized = True

def get_extra_state(self) -> torch.Tensor:
"""Save before checkpointing."""
Expand DownExpand Up@@ -280,7 +295,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
self.fp8_meta["scaling_fwd"].amax_history.copy_(amax_history_fwd)
self.fp8_meta["scaling_bwd"].scale.copy_(scale_bwd)
self.fp8_meta["scaling_bwd"].amax_history.copy_(amax_history_bwd)
self.fp8_meta_tensors_initialized = True

# Restore global FP8 buffer state.
set_global_fp8_buffer(state[4])
Expand DownExpand Up@@ -310,7 +324,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
self.fp8_meta["scaling_fwd"].amax_history.copy_(state["amax_history_fwd"])
self.fp8_meta["scaling_bwd"].scale.copy_(state["scale_bwd"])
self.fp8_meta["scaling_bwd"].amax_history.copy_(state["amax_history_bwd"])
self.fp8_meta_tensors_initialized = True

def set_activation_dtype(self, inp: torch.Tensor) -> None:
"""Get activation data type for AMP."""
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" + '
Change FP8 recipe defaults by ksivaman · Pull Request #112 · NVIDIA/TransformerEngine · GitHub
Skip to content
Change FP8 recipe defaults#112

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

from

Mar 30, 2023
 Merged


Change FP8 recipe defaults
#112
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
23b172d

Change FP8 recipe defaults


ksivaman Mar 21, 2023
112e54a

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 22, 2023
7bb9898

Increase default amax history length


ksivaman Mar 22, 2023
e5198f3

Always check history size


ksivaman Mar 22, 2023
bedea6b

no amax history for onnx export


ksivaman Mar 23, 2023
ab2d1d7

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 23, 2023
e894a6e

revert onnx export test changes


ksivaman Mar 24, 2023
a48963f

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 28, 2023
de1eb75

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 29, 2023
48bb4df

Fix indices in onnx test


ksivaman Mar 29, 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


17 changes: 10 additions & 7 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,12 +105,15 @@ def to_numpy(tensor):
 return tensor.cpu().numpy()


def set_layer_scale(module: torch.nn.Module, scale: float):
 module.fp8_init()
def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
 """Initialize the FP8 quantization scales in module"""
 NB_SCALES_PER_GEMM = 3 # One scale per: input, weights, and output GEMM tensors.
 nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
 module.fp8_init(num_gemms)
 module.fp8_meta["scaling_fwd"].scale = torch.ones(
2, dtype=torch.float32, device="cuda") / scale
nb_total_scales, dtype=torch.float32, device="cuda") / scale
 module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
2, dtype=torch.float32, device="cuda") * scale
nb_total_scales, dtype=torch.float32, device="cuda") * scale


def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
Expand DownExpand Up@@ -678,7 +681,7 @@ def forward(self, inp):
 precision
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model.linear, scale_factor)
 set_layer_scale(model.linear, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)

 if precision in (torch.bfloat16, ):
Expand DownExpand Up@@ -736,7 +739,7 @@ def test_export_layernorm_linear(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand DownExpand Up@@ -792,7 +795,7 @@ def test_export_layernorm_mlp(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=2)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/recipe.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,10 @@ class DelayedScaling:
 fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID
 Controls the FP8 data format used during forward and backward
 pass.
 amax_history_len : int, default = 1
 amax_history_len : int, default = 1024
 The length of the amax history window used for
 scaling factor computation.
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'most_recent'
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'max'
 Algorithm used for choosing the `amax` value for the
 scaling factor computation. There are 2 predefined
 choices: `max` chooses the largest `amax` in the history
Expand DownExpand Up@@ -125,8 +125,8 @@ def scaling_factor_compute(amax: Tensor,
 margin: int = 0
 interval: int = 1
 fp8_format: Format = Format.HYBRID
 amax_history_len: int = 1
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "most_recent"
 amax_history_len: int = 1024
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "max"
 override_linear_precision: _OverrideLinearPrecision = _OverrideLinearPrecision()
 scaling_factor_compute_algo: Optional[Callable] = None
 reduce_amax: bool = True
Expand Down

25 changes: 19 additions & 6 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,6 +13,7 @@

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init

Expand DownExpand Up@@ -187,6 +188,23 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

 if self.fp8_meta_tensors_initialized:
 # Handle changed amax history size.
 curr_len = self.fp8_meta[fp8_meta_tensor_key].amax_history.shape[0]
 need_len = self.fp8_meta["recipe"].amax_history_len
 if need_len < curr_len:
 self.fp8_meta[fp8_meta_tensor_key].amax_history = (
 self.fp8_meta[fp8_meta_tensor_key]
 .amax_history[: self.fp8_meta["recipe"].amax_history_len].clone()
 )
 elif need_len > curr_len:
 extra_rows = need_len - curr_len
 self.fp8_meta[fp8_meta_tensor_key].amax_history = F.pad(
 self.fp8_meta[fp8_meta_tensor_key].amax_history, pad=(0, 0, 0, extra_rows)
 )
 return

 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
Expand DownExpand Up@@ -222,12 +240,9 @@ def set_meta_tensor(self, fwd: bool) -> None:

 def init_fp8_meta_tensors(self) -> None:
 """Init scales and amaxes."""
 # Checkpoint loaded
 if self.fp8_meta_tensors_initialized:
 return

 self.set_meta_tensor(True)
 self.set_meta_tensor(False)
 self.fp8_meta_tensors_initialized = True

 def get_extra_state(self) -> torch.Tensor:
 """Save before checkpointing."""
Expand DownExpand Up@@ -280,7 +295,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(amax_history_fwd)
 self.fp8_meta["scaling_bwd"].scale.copy_(scale_bwd)
 self.fp8_meta["scaling_bwd"].amax_history.copy_(amax_history_bwd)
 self.fp8_meta_tensors_initialized = True

 # Restore global FP8 buffer state.
 set_global_fp8_buffer(state[4])
Expand DownExpand Up@@ -310,7 +324,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(state["amax_history_fwd"])
 self.fp8_meta["scaling_bwd"].scale.copy_(state["scale_bwd"])
 self.fp8_meta["scaling_bwd"].amax_history.copy_(state["amax_history_bwd"])
 self.fp8_meta_tensors_initialized = True

 def set_activation_dtype(self, inp: torch.Tensor) -> None:
 """Get activation data type for AMP."""
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('^' + ".*" + '
Change FP8 recipe defaults by ksivaman · Pull Request #112 · NVIDIA/TransformerEngine · GitHub
Skip to content
Change FP8 recipe defaults#112

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

from

Mar 30, 2023
 Merged


Change FP8 recipe defaults
#112
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
23b172d

Change FP8 recipe defaults


ksivaman Mar 21, 2023
112e54a

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 22, 2023
7bb9898

Increase default amax history length


ksivaman Mar 22, 2023
e5198f3

Always check history size


ksivaman Mar 22, 2023
bedea6b

no amax history for onnx export


ksivaman Mar 23, 2023
ab2d1d7

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 23, 2023
e894a6e

revert onnx export test changes


ksivaman Mar 24, 2023
a48963f

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 28, 2023
de1eb75

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 29, 2023
48bb4df

Fix indices in onnx test


ksivaman Mar 29, 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


17 changes: 10 additions & 7 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,12 +105,15 @@ def to_numpy(tensor):
 return tensor.cpu().numpy()


def set_layer_scale(module: torch.nn.Module, scale: float):
 module.fp8_init()
def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
 """Initialize the FP8 quantization scales in module"""
 NB_SCALES_PER_GEMM = 3 # One scale per: input, weights, and output GEMM tensors.
 nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
 module.fp8_init(num_gemms)
 module.fp8_meta["scaling_fwd"].scale = torch.ones(
2, dtype=torch.float32, device="cuda") / scale
nb_total_scales, dtype=torch.float32, device="cuda") / scale
 module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
2, dtype=torch.float32, device="cuda") * scale
nb_total_scales, dtype=torch.float32, device="cuda") * scale


def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
Expand DownExpand Up@@ -678,7 +681,7 @@ def forward(self, inp):
 precision
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model.linear, scale_factor)
 set_layer_scale(model.linear, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)

 if precision in (torch.bfloat16, ):
Expand DownExpand Up@@ -736,7 +739,7 @@ def test_export_layernorm_linear(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand DownExpand Up@@ -792,7 +795,7 @@ def test_export_layernorm_mlp(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=2)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/recipe.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,10 @@ class DelayedScaling:
 fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID
 Controls the FP8 data format used during forward and backward
 pass.
 amax_history_len : int, default = 1
 amax_history_len : int, default = 1024
 The length of the amax history window used for
 scaling factor computation.
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'most_recent'
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'max'
 Algorithm used for choosing the `amax` value for the
 scaling factor computation. There are 2 predefined
 choices: `max` chooses the largest `amax` in the history
Expand DownExpand Up@@ -125,8 +125,8 @@ def scaling_factor_compute(amax: Tensor,
 margin: int = 0
 interval: int = 1
 fp8_format: Format = Format.HYBRID
 amax_history_len: int = 1
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "most_recent"
 amax_history_len: int = 1024
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "max"
 override_linear_precision: _OverrideLinearPrecision = _OverrideLinearPrecision()
 scaling_factor_compute_algo: Optional[Callable] = None
 reduce_amax: bool = True
Expand Down

25 changes: 19 additions & 6 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,6 +13,7 @@

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init

Expand DownExpand Up@@ -187,6 +188,23 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

 if self.fp8_meta_tensors_initialized:
 # Handle changed amax history size.
 curr_len = self.fp8_meta[fp8_meta_tensor_key].amax_history.shape[0]
 need_len = self.fp8_meta["recipe"].amax_history_len
 if need_len < curr_len:
 self.fp8_meta[fp8_meta_tensor_key].amax_history = (
 self.fp8_meta[fp8_meta_tensor_key]
 .amax_history[: self.fp8_meta["recipe"].amax_history_len].clone()
 )
 elif need_len > curr_len:
 extra_rows = need_len - curr_len
 self.fp8_meta[fp8_meta_tensor_key].amax_history = F.pad(
 self.fp8_meta[fp8_meta_tensor_key].amax_history, pad=(0, 0, 0, extra_rows)
 )
 return

 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
Expand DownExpand Up@@ -222,12 +240,9 @@ def set_meta_tensor(self, fwd: bool) -> None:

 def init_fp8_meta_tensors(self) -> None:
 """Init scales and amaxes."""
 # Checkpoint loaded
 if self.fp8_meta_tensors_initialized:
 return

 self.set_meta_tensor(True)
 self.set_meta_tensor(False)
 self.fp8_meta_tensors_initialized = True

 def get_extra_state(self) -> torch.Tensor:
 """Save before checkpointing."""
Expand DownExpand Up@@ -280,7 +295,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(amax_history_fwd)
 self.fp8_meta["scaling_bwd"].scale.copy_(scale_bwd)
 self.fp8_meta["scaling_bwd"].amax_history.copy_(amax_history_bwd)
 self.fp8_meta_tensors_initialized = True

 # Restore global FP8 buffer state.
 set_global_fp8_buffer(state[4])
Expand DownExpand Up@@ -310,7 +324,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(state["amax_history_fwd"])
 self.fp8_meta["scaling_bwd"].scale.copy_(state["scale_bwd"])
 self.fp8_meta["scaling_bwd"].amax_history.copy_(state["amax_history_bwd"])
 self.fp8_meta_tensors_initialized = True

 def set_activation_dtype(self, inp: torch.Tensor) -> None:
 """Get activation data type for AMP."""
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('^' + ".*" + '
Change FP8 recipe defaults by ksivaman · Pull Request #112 · NVIDIA/TransformerEngine · GitHub
Skip to content
Change FP8 recipe defaults#112

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

from

Mar 30, 2023
 Merged


Change FP8 recipe defaults
#112
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
23b172d

Change FP8 recipe defaults


ksivaman Mar 21, 2023
112e54a

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 22, 2023
7bb9898

Increase default amax history length


ksivaman Mar 22, 2023
e5198f3

Always check history size


ksivaman Mar 22, 2023
bedea6b

no amax history for onnx export


ksivaman Mar 23, 2023
ab2d1d7

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 23, 2023
e894a6e

revert onnx export test changes


ksivaman Mar 24, 2023
a48963f

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 28, 2023
de1eb75

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 29, 2023
48bb4df

Fix indices in onnx test


ksivaman Mar 29, 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


17 changes: 10 additions & 7 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,12 +105,15 @@ def to_numpy(tensor):
 return tensor.cpu().numpy()


def set_layer_scale(module: torch.nn.Module, scale: float):
 module.fp8_init()
def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
 """Initialize the FP8 quantization scales in module"""
 NB_SCALES_PER_GEMM = 3 # One scale per: input, weights, and output GEMM tensors.
 nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
 module.fp8_init(num_gemms)
 module.fp8_meta["scaling_fwd"].scale = torch.ones(
2, dtype=torch.float32, device="cuda") / scale
nb_total_scales, dtype=torch.float32, device="cuda") / scale
 module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
2, dtype=torch.float32, device="cuda") * scale
nb_total_scales, dtype=torch.float32, device="cuda") * scale


def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
Expand DownExpand Up@@ -678,7 +681,7 @@ def forward(self, inp):
 precision
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model.linear, scale_factor)
 set_layer_scale(model.linear, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)

 if precision in (torch.bfloat16, ):
Expand DownExpand Up@@ -736,7 +739,7 @@ def test_export_layernorm_linear(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand DownExpand Up@@ -792,7 +795,7 @@ def test_export_layernorm_mlp(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=2)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/recipe.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,10 @@ class DelayedScaling:
 fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID
 Controls the FP8 data format used during forward and backward
 pass.
 amax_history_len : int, default = 1
 amax_history_len : int, default = 1024
 The length of the amax history window used for
 scaling factor computation.
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'most_recent'
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'max'
 Algorithm used for choosing the `amax` value for the
 scaling factor computation. There are 2 predefined
 choices: `max` chooses the largest `amax` in the history
Expand DownExpand Up@@ -125,8 +125,8 @@ def scaling_factor_compute(amax: Tensor,
 margin: int = 0
 interval: int = 1
 fp8_format: Format = Format.HYBRID
 amax_history_len: int = 1
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "most_recent"
 amax_history_len: int = 1024
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "max"
 override_linear_precision: _OverrideLinearPrecision = _OverrideLinearPrecision()
 scaling_factor_compute_algo: Optional[Callable] = None
 reduce_amax: bool = True
Expand Down

25 changes: 19 additions & 6 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,6 +13,7 @@

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init

Expand DownExpand Up@@ -187,6 +188,23 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

 if self.fp8_meta_tensors_initialized:
 # Handle changed amax history size.
 curr_len = self.fp8_meta[fp8_meta_tensor_key].amax_history.shape[0]
 need_len = self.fp8_meta["recipe"].amax_history_len
 if need_len < curr_len:
 self.fp8_meta[fp8_meta_tensor_key].amax_history = (
 self.fp8_meta[fp8_meta_tensor_key]
 .amax_history[: self.fp8_meta["recipe"].amax_history_len].clone()
 )
 elif need_len > curr_len:
 extra_rows = need_len - curr_len
 self.fp8_meta[fp8_meta_tensor_key].amax_history = F.pad(
 self.fp8_meta[fp8_meta_tensor_key].amax_history, pad=(0, 0, 0, extra_rows)
 )
 return

 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
Expand DownExpand Up@@ -222,12 +240,9 @@ def set_meta_tensor(self, fwd: bool) -> None:

 def init_fp8_meta_tensors(self) -> None:
 """Init scales and amaxes."""
 # Checkpoint loaded
 if self.fp8_meta_tensors_initialized:
 return

 self.set_meta_tensor(True)
 self.set_meta_tensor(False)
 self.fp8_meta_tensors_initialized = True

 def get_extra_state(self) -> torch.Tensor:
 """Save before checkpointing."""
Expand DownExpand Up@@ -280,7 +295,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(amax_history_fwd)
 self.fp8_meta["scaling_bwd"].scale.copy_(scale_bwd)
 self.fp8_meta["scaling_bwd"].amax_history.copy_(amax_history_bwd)
 self.fp8_meta_tensors_initialized = True

 # Restore global FP8 buffer state.
 set_global_fp8_buffer(state[4])
Expand DownExpand Up@@ -310,7 +324,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(state["amax_history_fwd"])
 self.fp8_meta["scaling_bwd"].scale.copy_(state["scale_bwd"])
 self.fp8_meta["scaling_bwd"].amax_history.copy_(state["amax_history_bwd"])
 self.fp8_meta_tensors_initialized = True

 def set_activation_dtype(self, inp: torch.Tensor) -> None:
 """Get activation data type for AMP."""
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" + '
Change FP8 recipe defaults by ksivaman · Pull Request #112 · NVIDIA/TransformerEngine · GitHub
Skip to content
Change FP8 recipe defaults#112

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

from

Mar 30, 2023
 Merged


Change FP8 recipe defaults
#112
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
23b172d

Change FP8 recipe defaults


ksivaman Mar 21, 2023
112e54a

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 22, 2023
7bb9898

Increase default amax history length


ksivaman Mar 22, 2023
e5198f3

Always check history size


ksivaman Mar 22, 2023
bedea6b

no amax history for onnx export


ksivaman Mar 23, 2023
ab2d1d7

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 23, 2023
e894a6e

revert onnx export test changes


ksivaman Mar 24, 2023
a48963f

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 28, 2023
de1eb75

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 29, 2023
48bb4df

Fix indices in onnx test


ksivaman Mar 29, 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


17 changes: 10 additions & 7 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,12 +105,15 @@ def to_numpy(tensor):
 return tensor.cpu().numpy()


def set_layer_scale(module: torch.nn.Module, scale: float):
 module.fp8_init()
def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
 """Initialize the FP8 quantization scales in module"""
 NB_SCALES_PER_GEMM = 3 # One scale per: input, weights, and output GEMM tensors.
 nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
 module.fp8_init(num_gemms)
 module.fp8_meta["scaling_fwd"].scale = torch.ones(
2, dtype=torch.float32, device="cuda") / scale
nb_total_scales, dtype=torch.float32, device="cuda") / scale
 module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
2, dtype=torch.float32, device="cuda") * scale
nb_total_scales, dtype=torch.float32, device="cuda") * scale


def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
Expand DownExpand Up@@ -678,7 +681,7 @@ def forward(self, inp):
 precision
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model.linear, scale_factor)
 set_layer_scale(model.linear, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)

 if precision in (torch.bfloat16, ):
Expand DownExpand Up@@ -736,7 +739,7 @@ def test_export_layernorm_linear(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand DownExpand Up@@ -792,7 +795,7 @@ def test_export_layernorm_mlp(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=2)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/recipe.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,10 @@ class DelayedScaling:
 fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID
 Controls the FP8 data format used during forward and backward
 pass.
 amax_history_len : int, default = 1
 amax_history_len : int, default = 1024
 The length of the amax history window used for
 scaling factor computation.
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'most_recent'
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'max'
 Algorithm used for choosing the `amax` value for the
 scaling factor computation. There are 2 predefined
 choices: `max` chooses the largest `amax` in the history
Expand DownExpand Up@@ -125,8 +125,8 @@ def scaling_factor_compute(amax: Tensor,
 margin: int = 0
 interval: int = 1
 fp8_format: Format = Format.HYBRID
 amax_history_len: int = 1
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "most_recent"
 amax_history_len: int = 1024
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "max"
 override_linear_precision: _OverrideLinearPrecision = _OverrideLinearPrecision()
 scaling_factor_compute_algo: Optional[Callable] = None
 reduce_amax: bool = True
Expand Down

25 changes: 19 additions & 6 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,6 +13,7 @@

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init

Expand DownExpand Up@@ -187,6 +188,23 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

 if self.fp8_meta_tensors_initialized:
 # Handle changed amax history size.
 curr_len = self.fp8_meta[fp8_meta_tensor_key].amax_history.shape[0]
 need_len = self.fp8_meta["recipe"].amax_history_len
 if need_len < curr_len:
 self.fp8_meta[fp8_meta_tensor_key].amax_history = (
 self.fp8_meta[fp8_meta_tensor_key]
 .amax_history[: self.fp8_meta["recipe"].amax_history_len].clone()
 )
 elif need_len > curr_len:
 extra_rows = need_len - curr_len
 self.fp8_meta[fp8_meta_tensor_key].amax_history = F.pad(
 self.fp8_meta[fp8_meta_tensor_key].amax_history, pad=(0, 0, 0, extra_rows)
 )
 return

 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
Expand DownExpand Up@@ -222,12 +240,9 @@ def set_meta_tensor(self, fwd: bool) -> None:

 def init_fp8_meta_tensors(self) -> None:
 """Init scales and amaxes."""
 # Checkpoint loaded
 if self.fp8_meta_tensors_initialized:
 return

 self.set_meta_tensor(True)
 self.set_meta_tensor(False)
 self.fp8_meta_tensors_initialized = True

 def get_extra_state(self) -> torch.Tensor:
 """Save before checkpointing."""
Expand DownExpand Up@@ -280,7 +295,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(amax_history_fwd)
 self.fp8_meta["scaling_bwd"].scale.copy_(scale_bwd)
 self.fp8_meta["scaling_bwd"].amax_history.copy_(amax_history_bwd)
 self.fp8_meta_tensors_initialized = True

 # Restore global FP8 buffer state.
 set_global_fp8_buffer(state[4])
Expand DownExpand Up@@ -310,7 +324,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(state["amax_history_fwd"])
 self.fp8_meta["scaling_bwd"].scale.copy_(state["scale_bwd"])
 self.fp8_meta["scaling_bwd"].amax_history.copy_(state["amax_history_bwd"])
 self.fp8_meta_tensors_initialized = True

 def set_activation_dtype(self, inp: torch.Tensor) -> None:
 """Get activation data type for AMP."""
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('^' + ".*" + '
Change FP8 recipe defaults by ksivaman · Pull Request #112 · NVIDIA/TransformerEngine · GitHub
Skip to content
Change FP8 recipe defaults#112

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

from

Mar 30, 2023
 Merged


Change FP8 recipe defaults
#112
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
23b172d

Change FP8 recipe defaults


ksivaman Mar 21, 2023
112e54a

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 22, 2023
7bb9898

Increase default amax history length


ksivaman Mar 22, 2023
e5198f3

Always check history size


ksivaman Mar 22, 2023
bedea6b

no amax history for onnx export


ksivaman Mar 23, 2023
ab2d1d7

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 23, 2023
e894a6e

revert onnx export test changes


ksivaman Mar 24, 2023
a48963f

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 28, 2023
de1eb75

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 29, 2023
48bb4df

Fix indices in onnx test


ksivaman Mar 29, 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


17 changes: 10 additions & 7 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,12 +105,15 @@ def to_numpy(tensor):
 return tensor.cpu().numpy()


def set_layer_scale(module: torch.nn.Module, scale: float):
 module.fp8_init()
def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
 """Initialize the FP8 quantization scales in module"""
 NB_SCALES_PER_GEMM = 3 # One scale per: input, weights, and output GEMM tensors.
 nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
 module.fp8_init(num_gemms)
 module.fp8_meta["scaling_fwd"].scale = torch.ones(
2, dtype=torch.float32, device="cuda") / scale
nb_total_scales, dtype=torch.float32, device="cuda") / scale
 module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
2, dtype=torch.float32, device="cuda") * scale
nb_total_scales, dtype=torch.float32, device="cuda") * scale


def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
Expand DownExpand Up@@ -678,7 +681,7 @@ def forward(self, inp):
 precision
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model.linear, scale_factor)
 set_layer_scale(model.linear, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)

 if precision in (torch.bfloat16, ):
Expand DownExpand Up@@ -736,7 +739,7 @@ def test_export_layernorm_linear(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand DownExpand Up@@ -792,7 +795,7 @@ def test_export_layernorm_mlp(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=2)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/recipe.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,10 @@ class DelayedScaling:
 fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID
 Controls the FP8 data format used during forward and backward
 pass.
 amax_history_len : int, default = 1
 amax_history_len : int, default = 1024
 The length of the amax history window used for
 scaling factor computation.
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'most_recent'
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'max'
 Algorithm used for choosing the `amax` value for the
 scaling factor computation. There are 2 predefined
 choices: `max` chooses the largest `amax` in the history
Expand DownExpand Up@@ -125,8 +125,8 @@ def scaling_factor_compute(amax: Tensor,
 margin: int = 0
 interval: int = 1
 fp8_format: Format = Format.HYBRID
 amax_history_len: int = 1
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "most_recent"
 amax_history_len: int = 1024
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "max"
 override_linear_precision: _OverrideLinearPrecision = _OverrideLinearPrecision()
 scaling_factor_compute_algo: Optional[Callable] = None
 reduce_amax: bool = True
Expand Down

25 changes: 19 additions & 6 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,6 +13,7 @@

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init

Expand DownExpand Up@@ -187,6 +188,23 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

 if self.fp8_meta_tensors_initialized:
 # Handle changed amax history size.
 curr_len = self.fp8_meta[fp8_meta_tensor_key].amax_history.shape[0]
 need_len = self.fp8_meta["recipe"].amax_history_len
 if need_len < curr_len:
 self.fp8_meta[fp8_meta_tensor_key].amax_history = (
 self.fp8_meta[fp8_meta_tensor_key]
 .amax_history[: self.fp8_meta["recipe"].amax_history_len].clone()
 )
 elif need_len > curr_len:
 extra_rows = need_len - curr_len
 self.fp8_meta[fp8_meta_tensor_key].amax_history = F.pad(
 self.fp8_meta[fp8_meta_tensor_key].amax_history, pad=(0, 0, 0, extra_rows)
 )
 return

 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
Expand DownExpand Up@@ -222,12 +240,9 @@ def set_meta_tensor(self, fwd: bool) -> None:

 def init_fp8_meta_tensors(self) -> None:
 """Init scales and amaxes."""
 # Checkpoint loaded
 if self.fp8_meta_tensors_initialized:
 return

 self.set_meta_tensor(True)
 self.set_meta_tensor(False)
 self.fp8_meta_tensors_initialized = True

 def get_extra_state(self) -> torch.Tensor:
 """Save before checkpointing."""
Expand DownExpand Up@@ -280,7 +295,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(amax_history_fwd)
 self.fp8_meta["scaling_bwd"].scale.copy_(scale_bwd)
 self.fp8_meta["scaling_bwd"].amax_history.copy_(amax_history_bwd)
 self.fp8_meta_tensors_initialized = True

 # Restore global FP8 buffer state.
 set_global_fp8_buffer(state[4])
Expand DownExpand Up@@ -310,7 +324,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(state["amax_history_fwd"])
 self.fp8_meta["scaling_bwd"].scale.copy_(state["scale_bwd"])
 self.fp8_meta["scaling_bwd"].amax_history.copy_(state["amax_history_bwd"])
 self.fp8_meta_tensors_initialized = True

 def set_activation_dtype(self, inp: torch.Tensor) -> None:
 """Get activation data type for AMP."""
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); }
})();
})();
Change FP8 recipe defaults by ksivaman · Pull Request #112 · NVIDIA/TransformerEngine · GitHub
Skip to content
Change FP8 recipe defaults#112

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

from

Mar 30, 2023
 Merged


Change FP8 recipe defaults
#112
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
23b172d

Change FP8 recipe defaults


ksivaman Mar 21, 2023
112e54a

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 22, 2023
7bb9898

Increase default amax history length


ksivaman Mar 22, 2023
e5198f3

Always check history size


ksivaman Mar 22, 2023
bedea6b

no amax history for onnx export


ksivaman Mar 23, 2023
ab2d1d7

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 23, 2023
e894a6e

revert onnx export test changes


ksivaman Mar 24, 2023
a48963f

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 28, 2023
de1eb75

Merge branch 'main' into change_fp8_defaults


ksivaman Mar 29, 2023
48bb4df

Fix indices in onnx test


ksivaman Mar 29, 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


17 changes: 10 additions & 7 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -105,12 +105,15 @@ def to_numpy(tensor):
 return tensor.cpu().numpy()


def set_layer_scale(module: torch.nn.Module, scale: float):
 module.fp8_init()
def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
 """Initialize the FP8 quantization scales in module"""
 NB_SCALES_PER_GEMM = 3 # One scale per: input, weights, and output GEMM tensors.
 nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
 module.fp8_init(num_gemms)
 module.fp8_meta["scaling_fwd"].scale = torch.ones(
2, dtype=torch.float32, device="cuda") / scale
nb_total_scales, dtype=torch.float32, device="cuda") / scale
 module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
2, dtype=torch.float32, device="cuda") * scale
nb_total_scales, dtype=torch.float32, device="cuda") * scale


def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
Expand DownExpand Up@@ -678,7 +681,7 @@ def forward(self, inp):
 precision
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model.linear, scale_factor)
 set_layer_scale(model.linear, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)

 if precision in (torch.bfloat16, ):
Expand DownExpand Up@@ -736,7 +739,7 @@ def test_export_layernorm_linear(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=1)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand DownExpand Up@@ -792,7 +795,7 @@ def test_export_layernorm_mlp(
 zero_centered_gamma=zero_centered_gamma,
 ).to(device='cuda')
 if use_fp8:
 set_layer_scale(model, scale_factor)
 set_layer_scale(model, scale_factor, num_gemms=2)
 do_export(model, inp, fname, use_fp8)
 if not use_fp8:
 validate_result(fname, inp, model, atol=1e-3)
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/recipe.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,10 @@ class DelayedScaling:
 fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID
 Controls the FP8 data format used during forward and backward
 pass.
 amax_history_len : int, default = 1
 amax_history_len : int, default = 1024
 The length of the amax history window used for
 scaling factor computation.
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'most_recent'
 amax_compute_algo : {'max', 'most_recent', Callable}, default = 'max'
 Algorithm used for choosing the `amax` value for the
 scaling factor computation. There are 2 predefined
 choices: `max` chooses the largest `amax` in the history
Expand DownExpand Up@@ -125,8 +125,8 @@ def scaling_factor_compute(amax: Tensor,
 margin: int = 0
 interval: int = 1
 fp8_format: Format = Format.HYBRID
 amax_history_len: int = 1
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "most_recent"
 amax_history_len: int = 1024
 amax_compute_algo: Union[Literal["max", "most_recent"], Callable] = "max"
 override_linear_precision: _OverrideLinearPrecision = _OverrideLinearPrecision()
 scaling_factor_compute_algo: Optional[Callable] = None
 reduce_amax: bool = True
Expand Down

25 changes: 19 additions & 6 deletions
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -13,6 +13,7 @@

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init

Expand DownExpand Up@@ -187,6 +188,23 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"

 if self.fp8_meta_tensors_initialized:
 # Handle changed amax history size.
 curr_len = self.fp8_meta[fp8_meta_tensor_key].amax_history.shape[0]
 need_len = self.fp8_meta["recipe"].amax_history_len
 if need_len < curr_len:
 self.fp8_meta[fp8_meta_tensor_key].amax_history = (
 self.fp8_meta[fp8_meta_tensor_key]
 .amax_history[: self.fp8_meta["recipe"].amax_history_len].clone()
 )
 elif need_len > curr_len:
 extra_rows = need_len - curr_len
 self.fp8_meta[fp8_meta_tensor_key].amax_history = F.pad(
 self.fp8_meta[fp8_meta_tensor_key].amax_history, pad=(0, 0, 0, extra_rows)
 )
 return

 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
Expand DownExpand Up@@ -222,12 +240,9 @@ def set_meta_tensor(self, fwd: bool) -> None:

 def init_fp8_meta_tensors(self) -> None:
 """Init scales and amaxes."""
 # Checkpoint loaded
 if self.fp8_meta_tensors_initialized:
 return

 self.set_meta_tensor(True)
 self.set_meta_tensor(False)
 self.fp8_meta_tensors_initialized = True

 def get_extra_state(self) -> torch.Tensor:
 """Save before checkpointing."""
Expand DownExpand Up@@ -280,7 +295,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(amax_history_fwd)
 self.fp8_meta["scaling_bwd"].scale.copy_(scale_bwd)
 self.fp8_meta["scaling_bwd"].amax_history.copy_(amax_history_bwd)
 self.fp8_meta_tensors_initialized = True

 # Restore global FP8 buffer state.
 set_global_fp8_buffer(state[4])
Expand DownExpand Up@@ -310,7 +324,6 @@ def set_extra_state(self, state: torch.Tensor) -> None:
 self.fp8_meta["scaling_fwd"].amax_history.copy_(state["amax_history_fwd"])
 self.fp8_meta["scaling_bwd"].scale.copy_(state["scale_bwd"])
 self.fp8_meta["scaling_bwd"].amax_history.copy_(state["amax_history_bwd"])
 self.fp8_meta_tensors_initialized = True

 def set_activation_dtype(self, inp: torch.Tensor) -> None:
 """Get activation data type for AMP."""
Expand Down