Skip to content
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

Jump to bottom

Introduce LayerNorm optimization from latest Apex #277

Merged
merged 4 commits into from
Nov 2, 2023
Merged

Introduce LayerNorm optimization from latest Apex #277

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
34c9b34
Introduce LayerNorm optimization from https://github.com/NVIDIA/apex/…
Quentin-Anthony Nov 1, 2023
7c59960
Fix args call
Quentin-Anthony Nov 1, 2023
8434496
Ad-hoc apex version check
Quentin-Anthony Nov 1, 2023
046319f
Remove unnecessary TransformerConfig arg
Quentin-Anthony Nov 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions megatron/arguments.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -626,6 +626,9 @@ def _add_network_size_args(parser):
group.add_argument('--apply-layernorm-1p', action='store_true',
help='Adjust LayerNorm weights such that they are centered '
'around zero. This improves numerical stability.')
group.add_argument('--disable-mem-efficient-ln', action='store_false',
help='Disable the memory-efficient fused LayerNorm optimization '
'introduced in https://github.com/NVIDIA/apex/pull/1715')
group.add_argument('--apply-residual-connection-post-layernorm',
action='store_true',
help='If set, use original BERT residula connection '
Expand Down
12 changes: 10 additions & 2 deletions megatron/model/fused_layer_norm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from torch.nn import init
import importlib
from torch.nn import functional as F
import inspect

from megatron.core.utils import make_viewless_tensor

Expand All @@ -31,10 +32,12 @@ class MixedFusedLayerNorm(torch.nn.Module):
def __init__(self, normalized_shape, eps=1e-5,
no_persist_layer_norm=True,
sequence_parallel=False,
apply_layernorm_1p=False):
apply_layernorm_1p=False,
mem_efficient_ln=True):
super(MixedFusedLayerNorm, self).__init__()

self.apply_layernorm_1p = apply_layernorm_1p
self.mem_efficient_ln = mem_efficient_ln

global fused_layer_norm_cuda
fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
Expand Down Expand Up @@ -83,7 +86,12 @@ def forward(self, input):
return F.layer_norm(input, self.normalized_shape, weight, self.bias, self.eps)

if self.no_persist_layer_norm:
return FusedLayerNormAffineFunction.apply(input, weight, self.bias, self.normalized_shape, self.eps)
# Apex does not have versions yet (https://github.com/NVIDIA/apex/pull/1648), so we need to inspect
# the function manually on whether the extra arg introduced in https://github.com/NVIDIA/apex/pull/1715 exists yet
if 'memory_efficient' in inspect.getfullargspec(FusedLayerNormAffineFunction.forward).args:
return FusedLayerNormAffineFunction.apply(input, weight, self.bias, self.normalized_shape, self.eps, self.mem_efficient_ln)
else:
return FusedLayerNormAffineFunction.apply(input, weight, self.bias, self.normalized_shape, self.eps)
else:
output = FastLayerNormFN.apply(input, weight, self.bias, self.eps)

Expand Down
12 changes: 8 additions & 4 deletions megatron/model/transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -919,7 +919,8 @@ def __init__(self, config,
eps=config.layernorm_epsilon,
no_persist_layer_norm=args.no_persist_layer_norm,
sequence_parallel=config.sequence_parallel,
apply_layernorm_1p=args.apply_layernorm_1p)
apply_layernorm_1p=args.apply_layernorm_1p,
mem_efficient_ln=not args.disable_mem_efficient_ln)
else:
self.input_layernorm = LayerNorm(
config.hidden_size,
Expand All @@ -944,7 +945,8 @@ def __init__(self, config,
eps=config.layernorm_epsilon,
no_persist_layer_norm=not config.persist_layer_norm,
sequence_parallel=config.sequence_parallel,
apply_layernorm_1p=args.apply_layernorm_1p)
apply_layernorm_1p=args.apply_layernorm_1p,
mem_efficient_ln=not args.disable_mem_efficient_ln)
else:
self.post_attention_layernorm = LayerNorm(
config.hidden_size,
Expand All @@ -967,7 +969,8 @@ def __init__(self, config,
eps=config.layernorm_epsilon,
no_persist_layer_norm=not config.persist_layer_norm,
sequence_parallel=config.sequence_parallel,
apply_layernorm_1p=args.apply_layernorm_1p)
apply_layernorm_1p=args.apply_layernorm_1p,
mem_efficient_ln=not args.disable_mem_efficient_ln)
else:
self.post_inter_attention_layernorm = MixedFusedRMSNorm(config.hidden_size, config.layernorm_epsilon)

Expand Down Expand Up @@ -1726,7 +1729,8 @@ def build_layer(layer_number, n_e):
eps=config.layernorm_epsilon,
no_persist_layer_norm=args.no_persist_layer_norm,
sequence_parallel=config.sequence_parallel,
apply_layernorm_1p=args.apply_layernorm_1p)
apply_layernorm_1p=args.apply_layernorm_1p,
mem_efficient_ln=not args.disable_mem_efficient_ln)
else:
self.final_layernorm = LayerNorm(
config.hidden_size,
Expand Down