From 401eee2552581a72dcbff297bd898c38cfd7f7b4 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Fri, 8 Dec 2023 01:05:29 +0000 Subject: [PATCH 1/9] Implemented deferred initialization via `device='meta'` option for te.Linear and added new PyTorch example to demonstrate its use with FullyShardedDataParallel execution. Signed-off-by: Alp Dener --- examples/pytorch/fsdp/README.md | 7 + examples/pytorch/fsdp/fsdp.py | 134 ++++++++++++++++++++ transformer_engine/pytorch/module/linear.py | 88 +++++++++---- 3 files changed, 204 insertions(+), 25 deletions(-) create mode 100644 examples/pytorch/fsdp/README.md create mode 100644 examples/pytorch/fsdp/fsdp.py diff --git a/examples/pytorch/fsdp/README.md b/examples/pytorch/fsdp/README.md new file mode 100644 index 0000000000..78cce0ba03 --- /dev/null +++ b/examples/pytorch/fsdp/README.md @@ -0,0 +1,7 @@ +# Basic Example for Using PyTorch Fully Sharded Data Parallel mode with Transformer Engine + +```bash +torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py +torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py --no-fp8 +torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py --defer-init +``` diff --git a/examples/pytorch/fsdp/fsdp.py b/examples/pytorch/fsdp/fsdp.py new file mode 100644 index 0000000000..f257c6acff --- /dev/null +++ b/examples/pytorch/fsdp/fsdp.py @@ -0,0 +1,134 @@ +# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import os +import argparse +import torch + +def lowercase(s): + return str(s).lower() + +def torch_dtype(d): + typemap = { + 'fp32' : torch.float32, + 'float32' : torch.float32, + 'fp16' : torch.float16, + 'float16' : torch.float16, + 'bf16' : torch.bfloat16, + 'bfloat16' : torch.bfloat16 + } + if lowercase(d) not in typemap.keys(): + raise TypeError + return typemap[lowercase(d)] + +parser = argparse.ArgumentParser(description="Run Transformer Engine modules with the " + + "torch.distributed.fsdp.FullyShardedDataParallel strategy.") +parser.add_argument("--no-fp8", action="store_true", default=False, + help="Disables the te.fp8_autocast() context.") +parser.add_argument('-i', "--num-iters", type=int, default=3, + help="Number of fake training iterations.") +parser.add_argument('-b', "--batch-size", type=int, default=32, + help="Input batch size.") +parser.add_argument('-s', "--seq-length", type=int, default=2048, + help="Input sequence length.") +parser.add_argument('-n', "--num-heads", type=int, default=64, + help="Number of attention heads.") +parser.add_argument('-d', "--head-dim", type=int, default=512, + help="Dimension of each attention head (number of KV channels).") +parser.add_argument('-l', "--num-layers", type=int, default=1, + help="Number of modules chained together with nn.Sequential.") +parser.add_argument("--seed", type=int, default=1234, + help="PyTorch RNG seed.") +parser.add_argument("--defer-init", action="store_true", + help="Use 'meta' device to defer module initialization until FSDP sharding.") +parser.add_argument("--dtype", type=torch_dtype, default=torch.bfloat16, + help="Data type for input tensor and Transformer Engine module parameters.") +args = parser.parse_args() + +import torch.distributed as dist +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel, MixedPrecision +from torch.distributed.fsdp.wrap import always_wrap_policy, transformer_auto_wrap_policy + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import Format, DelayedScaling + +local_rank = int(os.environ["LOCAL_RANK"]) +world_size = int(os.environ["WORLD_SIZE"]) + +dist.init_process_group(backend="nccl") +torch.cuda.set_device(local_rank) +if local_rank == 0: + print(f"[GPU-0] WORLD_SIZE = {world_size}\n\n", end='') + +dtype = args.dtype +seq_len = args.seq_length +batch_size = args.batch_size +num_heads = args.num_heads +head_dim = args.head_dim +hidden_size = num_heads * head_dim +ffn_hidden_size = 4 * hidden_size +torch.manual_seed(args.seed) + +layer_type = te.Linear +layer_args = (hidden_size, ffn_hidden_size) +layer_kwargs = { + 'bias': True, + 'params_dtype': dtype, + 'device': 'meta' if args.defer_init else 'cuda' +} + +if args.num_layers > 1: + if layer_type is te.Linear: + # input and output features need to have the same size in order to + # chain these layers back to back in nn.Sequential + ffn_hidden_size = hidden_size + layer_args = (hidden_size, hidden_size) + te_layer_list = [ layer_type(*layer_args, **layer_kwargs) for i in range(args.num_layers) ] + te_model = nn.Sequential(*te_layer_list) +else: + te_model = layer_type(*layer_args, **layer_kwargs) + +pre_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 +print(f"[GPU-{local_rank}] Pre-FSDP memory use = {pre_mem_use}MiB\n", end='') + +te_model = FullyShardedDataParallel(te_model, + use_orig_params=True, + mixed_precision=MixedPrecision( + param_dtype=dtype, + reduce_dtype=torch.float32, + ), + sync_module_states=True, + auto_wrap_policy=always_wrap_policy) + +post_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 +print(f"[GPU-{local_rank}] Post-FSDP memory use = {post_mem_use}MiB\n", end='') + +fp8_format = Format.HYBRID +fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") + +start = torch.cuda.Event(enable_timing=True) +end = torch.cuda.Event(enable_timing=True) + +optim = torch.optim.Adam(te_model.parameters(), lr=0.0001) +torch.cuda.synchronize() +start.record() +for i in range(args.num_iters): + x = torch.rand(seq_len//world_size, batch_size, hidden_size).to(dtype=dtype).cuda() + with te.fp8_autocast(enabled=not args.no_fp8, fp8_recipe=fp8_recipe): + y = te_model(x) + loss = y.sum() + loss.backward() + optim.step() + del x + if local_rank == 0: + print(f"[GPU-0] Iter. {i+1}\n", end='') +end.record() +torch.cuda.synchronize() + +train_time = start.elapsed_time(end)/1000. +max_memory_alloc = int(torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}") * 1e-6) +print(f"\n[GPU-{local_rank}] Training Time: {train_time}s\n" + + f"[GPU-{local_rank}] Avg. Iter. Time: {train_time /args.num_iters}s\n" + + f"[GPU-{local_rank}] Peak memory use = {max_memory_alloc}MiB\n\n", end='') diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 2a28d67292..c3c282df79 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -5,6 +5,7 @@ """Linear API""" import warnings from typing import Union, Optional, Callable, Tuple, List, Dict, Any +from dataclasses import dataclass import torch @@ -39,6 +40,7 @@ from ..cpp_extensions import ( fp8_gemm, gemm, + fp8_transpose, fp8_cast_transpose_fused, cast_to_fp8, ) @@ -82,7 +84,7 @@ def forward( ub_split_ag: bool, ub_atomic_gemm_rs: bool, ub_atomic_gemm_ag: bool, - ub_name: str, + ub_name: str ) -> torch.Tensor: # Make sure input dimensions are compatible in_features = weight.shape[-1] @@ -625,6 +627,10 @@ def __init__( if any([ub_atomic_gemm_rs, ub_atomic_gemm_ag]): assert ub_name is not None, "Userbuffer name [string] is not set." self.ub_name = ub_name + self.get_rng_state_tracker = get_rng_state_tracker + if device == 'meta': + assert self.parameters_split is None, ("Cannot split module parameters " + "on 'meta' device.") if ub_split_rs or ub_split_ag or ub_atomic_gemm_rs: assert ( @@ -655,8 +661,7 @@ def __init__( elif self.parallel_mode == "row": self.in_features = divide(self.in_features, self.tp_size) - if init_method is None: - init_method = get_default_init_method() + self.param_init_fn = get_default_init_method() if init_method is None else init_method self.sequence_parallel = (self.tp_size > 1) and sequence_parallel @@ -664,34 +669,16 @@ def __init__( self.out_features, self.in_features, device=device, dtype=params_dtype) - # TODO(ksivaman): This functionality works with FP8 outside TE. - initialize_affine_weight_gpu( - temp_weight, - init_method, - get_rng_state_tracker, - partition_dim=1 if self.parallel_mode == "row" else 0, - stride=1, - ) - - if self.primary_weights_in_fp8: - self.init_fp8_metadata() - self.fp8_meta["update_amax_and_scale_fwd"] = True - - self.weight_tensor = Float8Tensor.to_float8( - temp_weight, - fp8_meta=self.fp8_meta, - fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, - ) - else: - self.weight_tensor = temp_weight + self.weight_tensor = self.init_weight_tensor(temp_weight) if self.use_bias: self.bias_tensor = torch.empty(self.out_features, device=device, dtype=params_dtype) else: self.bias_tensor = torch.Tensor().to(dtype=params_dtype, device=device) - with torch.no_grad(): - self.bias_tensor.zero_() + if self.bias_tensor.device != torch.device('meta'): + with torch.no_grad(): + self.bias_tensor.zero_() # Configure parameter splits self.weight_names = [] @@ -788,6 +775,8 @@ def __init__( del self.weight_tensor del self.bias_tensor + self.reset_parameters(defer_init=(device == 'meta')) + self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features))) # For RPL, bias has to be added after TP collectives @@ -797,6 +786,55 @@ def __init__( else: self.gemm_bias_unfused_add = False + # Clean up weight and bias buffers + if self.parameters_split is None: + del self.weight_tensor + if self.use_bias: + del self.bias_tensor + + def init_weight_tensor(self, weights): + if weights.device == torch.device('meta'): + return weights + + # TODO(ksivaman): This functionality works with FP8 outside TE. + initialize_affine_weight_gpu( + weights, + self.param_init_fn, + self.get_rng_state_tracker, + partition_dim=1 if self.parallel_mode == "row" else 0, + stride=1, + set_tp_attributes=False, + ) + + if self.primary_weights_in_fp8: + self.init_fp8_metadata() + self.fp8_meta["update_amax_and_scale_fwd"] = True + weights = Float8Tensor.to_float8( + data=weights, + fp8_meta=self.fp8_meta, + fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, + ) + + return weights + + def reset_parameters(self, defer_init=False): + if defer_init: + return + + # materialize all parameters (move to real device, initialize values) + for wname, bname in zip(self.weight_names, self.bias_names): + if self._parameters[wname].device == torch.device('meta'): + self._parameters[wname].to(device='cuda') + self._parameters[wname] = torch.nn.Parameter( + self.init_weight_tensor(self._parameters[wname])) + + if self.use_bias: + if self._parameters[bname].device == torch.device('meta'): + self._parameters[bname].to(device='cuda') + + with torch.no_grad(): + self._parameters[bname].zero_() + def get_fp8_weights_scratchpad( self, is_first_microbatch: Union[bool, None], From 3107ac3fad1073506d336f5f95c7a0d61faf1aad Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Tue, 9 Jan 2024 19:52:06 +0000 Subject: [PATCH 2/9] correcting Float8Tensor initialization and fixing linting errors Signed-off-by: Alp Dener --- transformer_engine/pytorch/module/linear.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index c3c282df79..4a2dfccf10 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -5,7 +5,6 @@ """Linear API""" import warnings from typing import Union, Optional, Callable, Tuple, List, Dict, Any -from dataclasses import dataclass import torch @@ -40,7 +39,6 @@ from ..cpp_extensions import ( fp8_gemm, gemm, - fp8_transpose, fp8_cast_transpose_fused, cast_to_fp8, ) @@ -792,7 +790,10 @@ def __init__( if self.use_bias: del self.bias_tensor - def init_weight_tensor(self, weights): + def init_weight_tensor(self, weights: torch.Tensor) -> torch.Tensor: + """ + Initialize values for the given weight tensor if tensor lives on a real device. + """ if weights.device == torch.device('meta'): return weights @@ -809,7 +810,7 @@ def init_weight_tensor(self, weights): if self.primary_weights_in_fp8: self.init_fp8_metadata() self.fp8_meta["update_amax_and_scale_fwd"] = True - weights = Float8Tensor.to_float8( + weights = Float8Tensor( data=weights, fp8_meta=self.fp8_meta, fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, @@ -817,7 +818,12 @@ def init_weight_tensor(self, weights): return weights - def reset_parameters(self, defer_init=False): + def reset_parameters(self, defer_init: bool = False) -> None: + """ + Reset all module parameters to initial values. Unless deferred initialization + is specified, all parameters on a 'meta' device are also materialized on a real cuda + device before the values are reset to initial. + """ if defer_init: return From 003b6df44ef20748a7fdf3feda692753952691b9 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Tue, 9 Jan 2024 21:40:02 +0000 Subject: [PATCH 3/9] removed duplicate code from upstream rebase, local tests passing Signed-off-by: Alp Dener --- examples/pytorch/fsdp/README.md | 25 ++++++++++++++++++--- transformer_engine/pytorch/module/linear.py | 10 ++------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/examples/pytorch/fsdp/README.md b/examples/pytorch/fsdp/README.md index 78cce0ba03..5814a16002 100644 --- a/examples/pytorch/fsdp/README.md +++ b/examples/pytorch/fsdp/README.md @@ -1,7 +1,26 @@ # Basic Example for Using PyTorch Fully Sharded Data Parallel mode with Transformer Engine ```bash -torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py -torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py --no-fp8 -torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py --defer-init +# FSDP without deferred initialization: +# Duplicate modules initialized on each device. Load on device memory reduced only after +# torch.distributed.fsdp.FullyShardedDataParallel mode shards model parameters. +$ torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py +# Sample output on 8xL40S: +# [GPU-0] WORLD_SIZE = 8 +# [GPU-0] Pre-FSDP memory use = 8590.196736MiB +# ... +# [GPU-0] Post-FSDP memory use = 1073.774592MiB +# ... + +# FSDP with deferred initialization: +# Modules initialized with empty paramaters via `device='meta'` option. Zero load on device +# memory until torch.distributed.fsdp.FullyShardedDataParallel mode triggers a reset on +# on already sharded model parameters. +$ torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py --defer-init +# Sample output on 8xL40S: +# [GPU-0] WORLD_SIZE = 8 +# [GPU-0] Pre-FSDP memory use = 0.0MiB +# ... +# [GPU-0] Post-FSDP memory use = 1073.774592MiB +# ... ``` diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 4a2dfccf10..95a95b5064 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -627,8 +627,8 @@ def __init__( self.ub_name = ub_name self.get_rng_state_tracker = get_rng_state_tracker if device == 'meta': - assert self.parameters_split is None, ("Cannot split module parameters " - "on 'meta' device.") + assert parameters_split is None, ("Cannot split module parameters " + "on 'meta' device.") if ub_split_rs or ub_split_ag or ub_atomic_gemm_rs: assert ( @@ -784,12 +784,6 @@ def __init__( else: self.gemm_bias_unfused_add = False - # Clean up weight and bias buffers - if self.parameters_split is None: - del self.weight_tensor - if self.use_bias: - del self.bias_tensor - def init_weight_tensor(self, weights: torch.Tensor) -> torch.Tensor: """ Initialize values for the given weight tensor if tensor lives on a real device. From 6233cd30944821c72c3d9a1cdd2e07c4ba0c7634 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Wed, 10 Jan 2024 00:06:15 +0000 Subject: [PATCH 4/9] improved comments/documentation for FSDP example Signed-off-by: Alp Dener --- examples/pytorch/fsdp/README.md | 3 + examples/pytorch/fsdp/fsdp.py | 241 ++++++++++++++++++-------------- 2 files changed, 136 insertions(+), 108 deletions(-) diff --git a/examples/pytorch/fsdp/README.md b/examples/pytorch/fsdp/README.md index 5814a16002..6cd45fa9cd 100644 --- a/examples/pytorch/fsdp/README.md +++ b/examples/pytorch/fsdp/README.md @@ -24,3 +24,6 @@ $ torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsd # [GPU-0] Post-FSDP memory use = 1073.774592MiB # ... ``` + +**NOTE:** This example has `fp8_autocast()` enabled by default. To run on GPUs without Fp8 support +(e.g.: A100), add the `--no-fp8` option to the commands shown above. diff --git a/examples/pytorch/fsdp/fsdp.py b/examples/pytorch/fsdp/fsdp.py index f257c6acff..e470e6fe5c 100644 --- a/examples/pytorch/fsdp/fsdp.py +++ b/examples/pytorch/fsdp/fsdp.py @@ -4,7 +4,15 @@ import os import argparse + import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel, MixedPrecision +from torch.distributed.fsdp.wrap import always_wrap_policy + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import Format, DelayedScaling def lowercase(s): return str(s).lower() @@ -22,113 +30,130 @@ def torch_dtype(d): raise TypeError return typemap[lowercase(d)] -parser = argparse.ArgumentParser(description="Run Transformer Engine modules with the " + +def parse_fsdp_args(): + parser = argparse.ArgumentParser(description="Run Transformer Engine modules with the " + "torch.distributed.fsdp.FullyShardedDataParallel strategy.") -parser.add_argument("--no-fp8", action="store_true", default=False, - help="Disables the te.fp8_autocast() context.") -parser.add_argument('-i', "--num-iters", type=int, default=3, - help="Number of fake training iterations.") -parser.add_argument('-b', "--batch-size", type=int, default=32, - help="Input batch size.") -parser.add_argument('-s', "--seq-length", type=int, default=2048, - help="Input sequence length.") -parser.add_argument('-n', "--num-heads", type=int, default=64, - help="Number of attention heads.") -parser.add_argument('-d', "--head-dim", type=int, default=512, - help="Dimension of each attention head (number of KV channels).") -parser.add_argument('-l', "--num-layers", type=int, default=1, - help="Number of modules chained together with nn.Sequential.") -parser.add_argument("--seed", type=int, default=1234, - help="PyTorch RNG seed.") -parser.add_argument("--defer-init", action="store_true", - help="Use 'meta' device to defer module initialization until FSDP sharding.") -parser.add_argument("--dtype", type=torch_dtype, default=torch.bfloat16, - help="Data type for input tensor and Transformer Engine module parameters.") -args = parser.parse_args() - -import torch.distributed as dist -from torch import nn -from torch.distributed.fsdp import FullyShardedDataParallel, MixedPrecision -from torch.distributed.fsdp.wrap import always_wrap_policy, transformer_auto_wrap_policy - -import transformer_engine.pytorch as te -from transformer_engine.common.recipe import Format, DelayedScaling - -local_rank = int(os.environ["LOCAL_RANK"]) -world_size = int(os.environ["WORLD_SIZE"]) - -dist.init_process_group(backend="nccl") -torch.cuda.set_device(local_rank) -if local_rank == 0: - print(f"[GPU-0] WORLD_SIZE = {world_size}\n\n", end='') - -dtype = args.dtype -seq_len = args.seq_length -batch_size = args.batch_size -num_heads = args.num_heads -head_dim = args.head_dim -hidden_size = num_heads * head_dim -ffn_hidden_size = 4 * hidden_size -torch.manual_seed(args.seed) - -layer_type = te.Linear -layer_args = (hidden_size, ffn_hidden_size) -layer_kwargs = { - 'bias': True, - 'params_dtype': dtype, - 'device': 'meta' if args.defer_init else 'cuda' -} - -if args.num_layers > 1: - if layer_type is te.Linear: - # input and output features need to have the same size in order to - # chain these layers back to back in nn.Sequential - ffn_hidden_size = hidden_size - layer_args = (hidden_size, hidden_size) - te_layer_list = [ layer_type(*layer_args, **layer_kwargs) for i in range(args.num_layers) ] - te_model = nn.Sequential(*te_layer_list) -else: - te_model = layer_type(*layer_args, **layer_kwargs) - -pre_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 -print(f"[GPU-{local_rank}] Pre-FSDP memory use = {pre_mem_use}MiB\n", end='') - -te_model = FullyShardedDataParallel(te_model, - use_orig_params=True, - mixed_precision=MixedPrecision( - param_dtype=dtype, - reduce_dtype=torch.float32, - ), - sync_module_states=True, - auto_wrap_policy=always_wrap_policy) - -post_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 -print(f"[GPU-{local_rank}] Post-FSDP memory use = {post_mem_use}MiB\n", end='') - -fp8_format = Format.HYBRID -fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") - -start = torch.cuda.Event(enable_timing=True) -end = torch.cuda.Event(enable_timing=True) - -optim = torch.optim.Adam(te_model.parameters(), lr=0.0001) -torch.cuda.synchronize() -start.record() -for i in range(args.num_iters): - x = torch.rand(seq_len//world_size, batch_size, hidden_size).to(dtype=dtype).cuda() - with te.fp8_autocast(enabled=not args.no_fp8, fp8_recipe=fp8_recipe): - y = te_model(x) - loss = y.sum() - loss.backward() - optim.step() - del x + parser.add_argument("--no-fp8", action="store_true", default=False, + help="Disables the te.fp8_autocast() context.") + parser.add_argument('-i', "--num-iters", type=int, default=3, + help="Number of fake training iterations.") + parser.add_argument('-b', "--batch-size", type=int, default=32, + help="Input batch size.") + parser.add_argument('-s', "--seq-length", type=int, default=1048, + help="Input sequence length.") + parser.add_argument('-n', "--num-heads", type=int, default=64, + help="Number of attention heads.") + parser.add_argument('-d', "--head-dim", type=int, default=512, + help="Dimension of each attention head (number of KV channels).") + parser.add_argument('-l', "--num-layers", type=int, default=1, + help="Number of modules chained together with nn.Sequential.") + parser.add_argument("--seed", type=int, default=1234, + help="PyTorch RNG seed.") + parser.add_argument("--defer-init", action="store_true", + help="Use 'meta' device to defer module initialization until FSDP sharding.") + parser.add_argument("--dtype", type=torch_dtype, default=torch.bfloat16, + help="Data type for input tensor and Transformer Engine module parameters.") + return parser.parse_args() + +def train(args): + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + # Initialize torch.distributed global process group + dist.init_process_group(backend="nccl") + torch.cuda.set_device(local_rank) if local_rank == 0: - print(f"[GPU-0] Iter. {i+1}\n", end='') -end.record() -torch.cuda.synchronize() - -train_time = start.elapsed_time(end)/1000. -max_memory_alloc = int(torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}") * 1e-6) -print(f"\n[GPU-{local_rank}] Training Time: {train_time}s\n" + - f"[GPU-{local_rank}] Avg. Iter. Time: {train_time /args.num_iters}s\n" + - f"[GPU-{local_rank}] Peak memory use = {max_memory_alloc}MiB\n\n", end='') + print(f"[GPU-0] WORLD_SIZE = {world_size}\n\n", end='') + + dtype = args.dtype + seq_len = args.seq_length + batch_size = args.batch_size + num_heads = args.num_heads + head_dim = args.head_dim + hidden_size = num_heads * head_dim + ffn_hidden_size = 4 * hidden_size + torch.manual_seed(args.seed) + + # Construct a simple homogeneous model (only one layer type) with NO PARALLELISM + # NOTE: Change the layer type, args and kwargs variables below to test different models. + layer_type = te.Linear + layer_args = (hidden_size, ffn_hidden_size) + layer_kwargs = { + 'bias': True, + 'params_dtype': dtype, + 'device': 'meta' if args.defer_init else 'cuda' + } + + if args.num_layers > 1: + # Multi-layer model via torch.nn.Sequential + if layer_type is te.Linear: + # When layers are linear, input and output feature sizes need to be identical + ffn_hidden_size = hidden_size + layer_args = (hidden_size, hidden_size) + te_layer_list = [ layer_type(*layer_args, **layer_kwargs) for i in range(args.num_layers) ] + te_model = nn.Sequential(*te_layer_list) + else: + # Single layer model + te_model = layer_type(*layer_args, **layer_kwargs) + + # Print out allocated device memory before the model parameters are sharded by FSDP + pre_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 + print(f"[GPU-{local_rank}] Pre-FSDP memory use = {pre_mem_use}MiB\n", end='') + + # Wrap the model with FSDP + # NOTE: The TE model itself has no inherent parallelism. FSDP shards model parameters and + # controls all communication. + all_gpus = dist.new_group(backend='nccl') + te_model = FullyShardedDataParallel(te_model, + process_group=all_gpus, + use_orig_params=True, + mixed_precision=MixedPrecision( + param_dtype=dtype, + reduce_dtype=torch.float32, + ), + sync_module_states=True, + auto_wrap_policy=always_wrap_policy) + + # Print out allocated device memory after the model parameters are sharded + post_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 + print(f"[GPU-{local_rank}] Post-FSDP memory use = {post_mem_use}MiB\n", end='') + + # Fp8 setup for TE + fp8_format = Format.HYBRID + fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") + + # Optimizer must be created after the model is wrapped in FSDP and the parameters are sharded + optim = torch.optim.Adam(te_model.parameters(), lr=0.0001) + + # Start and time dummy "training" iterations + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + for i in range(args.num_iters): + # Generate a random input batch + x = torch.rand(seq_len, batch_size, hidden_size).to(dtype=dtype).cuda() + # fp8_autocast needs to be given the FSDP process group for amax reductions + with te.fp8_autocast(enabled=not args.no_fp8, fp8_recipe=fp8_recipe, fp8_group=all_gpus): + y = te_model(x) + loss = y.sum() + # calculate gradient and take training step outside the fp8_autocast context + loss.backward() + optim.step() + del x + if local_rank == 0: + print(f"[GPU-0] Iter. {i+1}\n", end='') + end.record() + torch.cuda.synchronize() + + # Print out "training" time and peak memory use stats + train_time = start.elapsed_time(end)/1000. + max_memory_alloc = int(torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}") * 1e-6) + print(f"\n[GPU-{local_rank}] Training Time: {train_time}s\n" + + f"[GPU-{local_rank}] Avg. Iter. Time: {train_time /args.num_iters}s\n" + + f"[GPU-{local_rank}] Peak memory use = {max_memory_alloc}MiB\n\n", end='') + + +if __name__ == "__main__": + args = parse_fsdp_args() + train(args) From febcbefc310a7d1976958d3f80aa6fc3948911b9 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 11 Jan 2024 21:59:39 +0000 Subject: [PATCH 5/9] converted reset_parameters() into a base module function Signed-off-by: Alp Dener --- examples/pytorch/fsdp/README.md | 4 ++ transformer_engine/pytorch/module/_common.py | 56 +++++++++++++++- transformer_engine/pytorch/module/base.py | 34 ++++++++++ transformer_engine/pytorch/module/linear.py | 68 ++------------------ 4 files changed, 97 insertions(+), 65 deletions(-) diff --git a/examples/pytorch/fsdp/README.md b/examples/pytorch/fsdp/README.md index 6cd45fa9cd..0a357f1b26 100644 --- a/examples/pytorch/fsdp/README.md +++ b/examples/pytorch/fsdp/README.md @@ -1,3 +1,7 @@ +# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + # Basic Example for Using PyTorch Fully Sharded Data Parallel mode with Transformer Engine ```bash diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index edc3da120d..6a5f4cc37f 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -4,12 +4,16 @@ """Internal function used by multiple modules.""" -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union, Callable +from dataclasses import dataclass import torch from .. import cpp_extensions as tex -from ..fp8 import get_fp8_te_dtype +from ..fp8 import get_fp8_te_dtype, FP8GlobalStateManager +from ..float8_tensor import Float8Tensor +from ..distributed import initialize_affine_weight_gpu +from ..utils import get_default_init_method def _get_normalization_func(normalization: str, fp8_output: bool, @@ -187,3 +191,51 @@ def _noop_cat( # Perform no-op concat return _NoopCatFunc.apply(split_ranges, full_tensor, *tensors) + + +@dataclass +class _ParameterInitMeta: + parent: torch.nn.Module + init_fn: Optional[Callable] = get_default_init_method() + get_rng_state_tracker: Optional[Callable] = None + partition_dim: Optional[int] = 0, + stride: Optional[int] = 1, + fp8_meta_index: Optional[int] = None + + def __post_init__(self): + assert self.parent is not None + if self.init_fn is None: + self.init_fn = get_default_init_method() + + def init_as_weight(self, param: torch.Tensor, set_tp_attributes: bool = False) -> torch.Tensor: + if param.device == torch.device('meta'): + return param + + initialize_affine_weight_gpu( + param, + self.init_fn, + self.get_rng_state_tracker, + partition_dim=self.partition_dim, + stride=self.stride, + set_tp_attributes=set_tp_attributes + ) + + if FP8GlobalStateManager.with_fp8_parameters(): + self.parent.init_fp8_metadata() + self.parent.fp8_meta["update_amax_and_scale_fwd"] = True + param = Float8Tensor( + param, + fp8_meta=self.parent.fp8_meta, + fp8_meta_index=self.fp8_meta_index + ) + + return param + + def init_as_bias(self, param: torch.Tensor) -> torch.Tensor: + if param.device == torch.device('meta'): + return param + + with torch.no_grad(): + param.zero_() + + return param diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index cf9634b2cc..948a512a95 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -16,6 +16,7 @@ import torch.nn.functional as F import transformer_engine_extensions as tex +from ._common import _ParameterInitMeta from ..export import is_in_onnx_export_mode from ..fp8 import ( get_default_fp8_recipe, @@ -234,6 +235,7 @@ def __init__(self) -> None: self.fp8_meta["async_amax_reduction"] = bool( int(os.getenv("NVTE_ASYNC_AMAX_REDUCTION", "0")) ) + self.param_init_meta = {} def set_meta_tensor(self, fwd: bool) -> None: """Init scales and amaxes for fwd | bwd.""" @@ -746,6 +748,38 @@ def get_fp8_weights_empty_tensors( ) return fp8_weight_tensors + def register_parameter(self, name, param, **kwargs): + """ + Thin wrapper around PyTorch parameter registration to stash additional parameter + metedata used in deferred initialization. + """ + super().register_parameter(name, param) + self.param_init_meta[name] = _ParameterInitMeta(self, **kwargs) + + def reset_parameters(self, defer_init: Optional[bool] = False) -> None: + """ + Reset all module parameters to initial values. Unless deferred initialization + is specified, all parameters on a 'meta' device are also materialized on a real cuda + device before the values are reset to initial. + """ + if defer_init: + return + + for name, param in self.named_parameters(recurse=False): + # Ensure parameter is on a real device + if param.device == torch.device('meta'): + param.to(device='cuda') + + if 'weight' in name: + # Initialize weight values on device + param = self.param_init_meta[name].init_as_weight(param) + elif 'bias' in name: + # Reset biases to zero + param = self.param_init_meta[name].init_as_bias(param) + + # Redo parameter wrap in case we broke it above + param = torch.nn.Parameter(param) + @abstractmethod def forward(self): """Needs override.""" diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 95a95b5064..90bc3bb8af 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -23,7 +23,6 @@ from ..fp8 import get_fp8_te_dtype, FP8GlobalStateManager from ..utils import ( divide, - get_default_init_method, cast_if_needed, assert_dim_for_fp8_exec, clear_tensor_data, @@ -32,7 +31,6 @@ set_tensor_model_parallel_attributes, get_distributed_world_size, allreduce, - initialize_affine_weight_gpu, reduce_scatter_along_first_dim, gather_along_first_dim, ) @@ -659,25 +657,17 @@ def __init__( elif self.parallel_mode == "row": self.in_features = divide(self.in_features, self.tp_size) - self.param_init_fn = get_default_init_method() if init_method is None else init_method - self.sequence_parallel = (self.tp_size > 1) and sequence_parallel - temp_weight = torch.empty( + self.weight_tensor = torch.empty( self.out_features, self.in_features, device=device, dtype=params_dtype) - self.weight_tensor = self.init_weight_tensor(temp_weight) - if self.use_bias: self.bias_tensor = torch.empty(self.out_features, device=device, dtype=params_dtype) else: self.bias_tensor = torch.Tensor().to(dtype=params_dtype, device=device) - if self.bias_tensor.device != torch.device('meta'): - with torch.no_grad(): - self.bias_tensor.zero_() - # Configure parameter splits self.weight_names = [] self.bias_names = [] @@ -742,7 +732,10 @@ def __init__( if is_subview: weight = weight[split_start:split_end] weight = torch.nn.Parameter(weight) - self.register_parameter(self.weight_names[i], weight) + self.register_parameter(self.weight_names[i], weight, + init_fn=init_method, + get_rng_state_tracker=get_rng_state_tracker, + fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT) # Construct bias parameter if needed if self.use_bias: @@ -784,57 +777,6 @@ def __init__( else: self.gemm_bias_unfused_add = False - def init_weight_tensor(self, weights: torch.Tensor) -> torch.Tensor: - """ - Initialize values for the given weight tensor if tensor lives on a real device. - """ - if weights.device == torch.device('meta'): - return weights - - # TODO(ksivaman): This functionality works with FP8 outside TE. - initialize_affine_weight_gpu( - weights, - self.param_init_fn, - self.get_rng_state_tracker, - partition_dim=1 if self.parallel_mode == "row" else 0, - stride=1, - set_tp_attributes=False, - ) - - if self.primary_weights_in_fp8: - self.init_fp8_metadata() - self.fp8_meta["update_amax_and_scale_fwd"] = True - weights = Float8Tensor( - data=weights, - fp8_meta=self.fp8_meta, - fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, - ) - - return weights - - def reset_parameters(self, defer_init: bool = False) -> None: - """ - Reset all module parameters to initial values. Unless deferred initialization - is specified, all parameters on a 'meta' device are also materialized on a real cuda - device before the values are reset to initial. - """ - if defer_init: - return - - # materialize all parameters (move to real device, initialize values) - for wname, bname in zip(self.weight_names, self.bias_names): - if self._parameters[wname].device == torch.device('meta'): - self._parameters[wname].to(device='cuda') - self._parameters[wname] = torch.nn.Parameter( - self.init_weight_tensor(self._parameters[wname])) - - if self.use_bias: - if self._parameters[bname].device == torch.device('meta'): - self._parameters[bname].to(device='cuda') - - with torch.no_grad(): - self._parameters[bname].zero_() - def get_fp8_weights_scratchpad( self, is_first_microbatch: Union[bool, None], From 8d848952b7b7a6d46a7dba26fa10fb8c062e8126 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Fri, 12 Jan 2024 18:15:47 +0000 Subject: [PATCH 6/9] fixed Float8Tensor creation with deferred init, all tests passing locally Signed-off-by: Alp Dener --- transformer_engine/pytorch/module/_common.py | 2 +- transformer_engine/pytorch/module/base.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 6a5f4cc37f..a76e9a422a 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -223,7 +223,7 @@ def init_as_weight(self, param: torch.Tensor, set_tp_attributes: bool = False) - if FP8GlobalStateManager.with_fp8_parameters(): self.parent.init_fp8_metadata() self.parent.fp8_meta["update_amax_and_scale_fwd"] = True - param = Float8Tensor( + param = Float8Tensor.to_float8( param, fp8_meta=self.parent.fp8_meta, fp8_meta_index=self.fp8_meta_index diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 948a512a95..efeccd3a77 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -768,7 +768,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: for name, param in self.named_parameters(recurse=False): # Ensure parameter is on a real device if param.device == torch.device('meta'): - param.to(device='cuda') + param = param.to(device='cuda') if 'weight' in name: # Initialize weight values on device @@ -778,7 +778,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: param = self.param_init_meta[name].init_as_bias(param) # Redo parameter wrap in case we broke it above - param = torch.nn.Parameter(param) + setattr(self, name, torch.nn.Parameter(param)) @abstractmethod def forward(self): From 1cc1be65950dd5669f283a91123ba9c05625ff7c Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Fri, 12 Jan 2024 22:43:33 +0000 Subject: [PATCH 7/9] extended deferred initialization to all TE modules Signed-off-by: Alp Dener --- examples/pytorch/fsdp/README.md | 32 ++++- examples/pytorch/fsdp/fsdp.py | 110 ++++++++++++------ transformer_engine/pytorch/module/_common.py | 42 +------ transformer_engine/pytorch/module/base.py | 24 +++- .../pytorch/module/layernorm.py | 16 ++- .../pytorch/module/layernorm_linear.py | 52 ++++----- .../pytorch/module/layernorm_mlp.py | 85 ++++++-------- transformer_engine/pytorch/module/linear.py | 4 + transformer_engine/pytorch/module/rmsnorm.py | 14 ++- transformer_engine/pytorch/utils.py | 15 +++ 10 files changed, 229 insertions(+), 165 deletions(-) diff --git a/examples/pytorch/fsdp/README.md b/examples/pytorch/fsdp/README.md index 0a357f1b26..d492ea4a57 100644 --- a/examples/pytorch/fsdp/README.md +++ b/examples/pytorch/fsdp/README.md @@ -11,10 +11,30 @@ $ torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py # Sample output on 8xL40S: # [GPU-0] WORLD_SIZE = 8 -# [GPU-0] Pre-FSDP memory use = 8590.196736MiB -# ... -# [GPU-0] Post-FSDP memory use = 1073.774592MiB -# ... +# [GPU-0] TransformerEngine Model: +# TransformerLayer( +# (self_attention): MultiheadAttention( +# (layernorm_qkv): LayerNormLinear() +# (core_attention): DotProductAttention( +# (flash_attention): FlashAttention() +# (fused_attention): FusedAttention() +# (unfused_attention): UnfusedDotProductAttention( +# (scale_mask_softmax): FusedScaleMaskSoftmax() +# (attention_dropout): Dropout(p=0.1, inplace=False) +# ) +# ) +# (proj): Linear() +# ) +# (layernorm_mlp): LayerNormMLP() +# ) +# [GPU-0] Pre-FSDP memory use = 83.935232MiB +# [GPU-0] Post-FSDP memory use = 10.491904MiB +# [GPU-0] Iter. 1 +# [GPU-0] Iter. 2 +# [GPU-0] Iter. 3 +# [GPU-0] Training Time: 6.647654296875s +# [GPU-0] Avg. Iter. Time: 2.2158847656250003s +# [GPU-0] Peak memory use = 3000MiB # FSDP with deferred initialization: # Modules initialized with empty paramaters via `device='meta'` option. Zero load on device @@ -23,9 +43,9 @@ $ torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsd $ torchrun --standalone --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) fsdp.py --defer-init # Sample output on 8xL40S: # [GPU-0] WORLD_SIZE = 8 -# [GPU-0] Pre-FSDP memory use = 0.0MiB # ... -# [GPU-0] Post-FSDP memory use = 1073.774592MiB +# [GPU-0] Pre-FSDP memory use = 0.0MiB +# [GPU-0] Post-FSDP memory use = 10.491904MiB # ... ``` diff --git a/examples/pytorch/fsdp/fsdp.py b/examples/pytorch/fsdp/fsdp.py index e470e6fe5c..5d30be6c97 100644 --- a/examples/pytorch/fsdp/fsdp.py +++ b/examples/pytorch/fsdp/fsdp.py @@ -4,12 +4,13 @@ import os import argparse +from functools import partial import torch import torch.distributed as dist from torch import nn from torch.distributed.fsdp import FullyShardedDataParallel, MixedPrecision -from torch.distributed.fsdp.wrap import always_wrap_policy +from torch.distributed.fsdp.wrap import always_wrap_policy, transformer_auto_wrap_policy import transformer_engine.pytorch as te from transformer_engine.common.recipe import Format, DelayedScaling @@ -30,27 +31,68 @@ def torch_dtype(d): raise TypeError return typemap[lowercase(d)] +te_layer_map = { + 'linear': te.Linear, + 'layernorm': te.LayerNorm, + 'rmsnorm': te.RMSNorm, + 'layernormlinear': te.LayerNormLinear, + 'layernormmlp': te.LayerNormMLP, + 'multiheadattention': te.MultiheadAttention, + 'transformerlayer': te.TransformerLayer +} +def te_layer(l): + if lowercase(l) not in te_layer_map.keys(): + raise TypeError + return te_layer_map[lowercase(l)] + +def get_layer_args(args): + hidden_size = args.num_heads * args.head_dim + layer_args = (hidden_size, ) + layer_kwargs = { + 'params_dtype': args.dtype, + 'device': 'meta' if args.defer_init else 'cuda' + } + if args.layer_type in [te.Linear, te.LayerNormLinear, te.LayerNormMLP]: + ffn_hidden_size = 3 * hidden_size if args.num_layers == 1 else hidden_size + layer_args += (ffn_hidden_size, ) + layer_kwargs['bias'] = True + if args.layer_type == te.LayerNormMLP: + layer_kwargs['seq_length'] = args.seq_length + elif args.layer_type == te.MultiheadAttention: + layer_args += (args.num_heads, ) + layer_kwargs['fuse_qkv_params'] = True + elif args.layer_type == te.TransformerLayer: + layer_args += (3 * hidden_size, args.num_heads) + layer_kwargs['fuse_qkv_params'] = True + layer_kwargs['seq_length'] = args.seq_length + return layer_args, layer_kwargs + def parse_fsdp_args(): parser = argparse.ArgumentParser(description="Run Transformer Engine modules with the " + "torch.distributed.fsdp.FullyShardedDataParallel strategy.") + parser.add_argument("-t", "--layer-type", type=te_layer, default=te.TransformerLayer, + choices=list(te_layer_map.values()), + help="TE module type used to construct the test model.") parser.add_argument("--no-fp8", action="store_true", default=False, help="Disables the te.fp8_autocast() context.") parser.add_argument('-i', "--num-iters", type=int, default=3, - help="Number of fake training iterations.") + help="Number of dummy 'training' iterations.") parser.add_argument('-b', "--batch-size", type=int, default=32, help="Input batch size.") parser.add_argument('-s', "--seq-length", type=int, default=1048, help="Input sequence length.") - parser.add_argument('-n', "--num-heads", type=int, default=64, + parser.add_argument('-n', "--num-heads", type=int, default=16, help="Number of attention heads.") - parser.add_argument('-d', "--head-dim", type=int, default=512, + parser.add_argument('-d', "--head-dim", type=int, default=128, help="Dimension of each attention head (number of KV channels).") parser.add_argument('-l', "--num-layers", type=int, default=1, help="Number of modules chained together with nn.Sequential.") parser.add_argument("--seed", type=int, default=1234, help="PyTorch RNG seed.") parser.add_argument("--defer-init", action="store_true", - help="Use 'meta' device to defer module initialization until FSDP sharding.") + help="Defer module parameter initialization until after FSDP sharding.") + parser.add_argument('-v', "--verbose", action="store_true", default=False, + help="Print out information from all GPUs instead of only the root GPU-0.") parser.add_argument("--dtype", type=torch_dtype, default=torch.bfloat16, help="Data type for input tensor and Transformer Engine module parameters.") return parser.parse_args() @@ -64,59 +106,51 @@ def train(args): torch.cuda.set_device(local_rank) if local_rank == 0: print(f"[GPU-0] WORLD_SIZE = {world_size}\n\n", end='') - - dtype = args.dtype - seq_len = args.seq_length - batch_size = args.batch_size - num_heads = args.num_heads - head_dim = args.head_dim - hidden_size = num_heads * head_dim - ffn_hidden_size = 4 * hidden_size torch.manual_seed(args.seed) # Construct a simple homogeneous model (only one layer type) with NO PARALLELISM - # NOTE: Change the layer type, args and kwargs variables below to test different models. - layer_type = te.Linear - layer_args = (hidden_size, ffn_hidden_size) - layer_kwargs = { - 'bias': True, - 'params_dtype': dtype, - 'device': 'meta' if args.defer_init else 'cuda' - } - + layer_args, layer_kwargs = get_layer_args(args) if args.num_layers > 1: - # Multi-layer model via torch.nn.Sequential - if layer_type is te.Linear: - # When layers are linear, input and output feature sizes need to be identical - ffn_hidden_size = hidden_size - layer_args = (hidden_size, hidden_size) - te_layer_list = [ layer_type(*layer_args, **layer_kwargs) for i in range(args.num_layers) ] + te_layer_list = [] + for i in range(args.num_layers): + if args.layer_type in [te.MultiheadAttention, te.TransformerLayer]: + layer_kwargs['layer_number'] = i+1 + te_layer_list.append(args.layer_type(*layer_args, **layer_kwargs)) te_model = nn.Sequential(*te_layer_list) else: # Single layer model - te_model = layer_type(*layer_args, **layer_kwargs) + te_model = args.layer_type(*layer_args, **layer_kwargs) + if local_rank == 0: + print(f"[GPU-0] TransformerEngine Model:\n{te_model}\n", end='') # Print out allocated device memory before the model parameters are sharded by FSDP pre_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 - print(f"[GPU-{local_rank}] Pre-FSDP memory use = {pre_mem_use}MiB\n", end='') + if local_rank == 0 or args.verbose: + print(f"[GPU-{local_rank}] Pre-FSDP memory use = {pre_mem_use}MiB\n", end='') # Wrap the model with FSDP # NOTE: The TE model itself has no inherent parallelism. FSDP shards model parameters and # controls all communication. all_gpus = dist.new_group(backend='nccl') + fsdp_wrap_policy = always_wrap_policy + if args.layer_type == te.TransformerLayer: + # NOTE: FSDP causes illegal memory access without this special policy for Transformers + fsdp_wrap_policy = partial(transformer_auto_wrap_policy, + transformer_layer_cls={te.TransformerLayer}) te_model = FullyShardedDataParallel(te_model, process_group=all_gpus, use_orig_params=True, mixed_precision=MixedPrecision( - param_dtype=dtype, + param_dtype=args.dtype, reduce_dtype=torch.float32, ), sync_module_states=True, - auto_wrap_policy=always_wrap_policy) + auto_wrap_policy=fsdp_wrap_policy) # Print out allocated device memory after the model parameters are sharded post_mem_use = torch.cuda.memory_allocated(device=f"cuda:{local_rank}") * 1e-6 - print(f"[GPU-{local_rank}] Post-FSDP memory use = {post_mem_use}MiB\n", end='') + if local_rank == 0 or args.verbose: + print(f"[GPU-{local_rank}] Post-FSDP memory use = {post_mem_use}MiB\n", end='') # Fp8 setup for TE fp8_format = Format.HYBRID @@ -132,7 +166,8 @@ def train(args): start.record() for i in range(args.num_iters): # Generate a random input batch - x = torch.rand(seq_len, batch_size, hidden_size).to(dtype=dtype).cuda() + x = torch.rand(args.seq_length, args.batch_size, + args.num_heads*args.head_dim).to(dtype=args.dtype).cuda() # fp8_autocast needs to be given the FSDP process group for amax reductions with te.fp8_autocast(enabled=not args.no_fp8, fp8_recipe=fp8_recipe, fp8_group=all_gpus): y = te_model(x) @@ -149,9 +184,10 @@ def train(args): # Print out "training" time and peak memory use stats train_time = start.elapsed_time(end)/1000. max_memory_alloc = int(torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}") * 1e-6) - print(f"\n[GPU-{local_rank}] Training Time: {train_time}s\n" + - f"[GPU-{local_rank}] Avg. Iter. Time: {train_time /args.num_iters}s\n" + - f"[GPU-{local_rank}] Peak memory use = {max_memory_alloc}MiB\n\n", end='') + if local_rank == 0 or args.verbose: + print(f"[GPU-{local_rank}] Training Time: {train_time}s\n" + + f"[GPU-{local_rank}] Avg. Iter. Time: {train_time /args.num_iters}s\n" + + f"[GPU-{local_rank}] Peak memory use = {max_memory_alloc}MiB\n\n", end='') if __name__ == "__main__": diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index a76e9a422a..114c91f979 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -10,9 +10,7 @@ import torch from .. import cpp_extensions as tex -from ..fp8 import get_fp8_te_dtype, FP8GlobalStateManager -from ..float8_tensor import Float8Tensor -from ..distributed import initialize_affine_weight_gpu +from ..fp8 import get_fp8_te_dtype from ..utils import get_default_init_method def _get_normalization_func(normalization: str, @@ -195,47 +193,17 @@ def _noop_cat( @dataclass class _ParameterInitMeta: + """ + Stores essential metadata needed to support deferred parameter initialization. + """ parent: torch.nn.Module init_fn: Optional[Callable] = get_default_init_method() get_rng_state_tracker: Optional[Callable] = None - partition_dim: Optional[int] = 0, - stride: Optional[int] = 1, fp8_meta_index: Optional[int] = None def __post_init__(self): + """Safeguard reference to the parameter's parent module and initialization function.""" assert self.parent is not None if self.init_fn is None: self.init_fn = get_default_init_method() - def init_as_weight(self, param: torch.Tensor, set_tp_attributes: bool = False) -> torch.Tensor: - if param.device == torch.device('meta'): - return param - - initialize_affine_weight_gpu( - param, - self.init_fn, - self.get_rng_state_tracker, - partition_dim=self.partition_dim, - stride=self.stride, - set_tp_attributes=set_tp_attributes - ) - - if FP8GlobalStateManager.with_fp8_parameters(): - self.parent.init_fp8_metadata() - self.parent.fp8_meta["update_amax_and_scale_fwd"] = True - param = Float8Tensor.to_float8( - param, - fp8_meta=self.parent.fp8_meta, - fp8_meta_index=self.fp8_meta_index - ) - - return param - - def init_as_bias(self, param: torch.Tensor) -> torch.Tensor: - if param.device == torch.device('meta'): - return param - - with torch.no_grad(): - param.zero_() - - return param diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index efeccd3a77..d3d5fc38f5 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -236,6 +236,7 @@ def __init__(self) -> None: int(os.getenv("NVTE_ASYNC_AMAX_REDUCTION", "0")) ) self.param_init_meta = {} + self.primary_weights_in_fp8 = FP8GlobalStateManager.with_fp8_parameters() def set_meta_tensor(self, fwd: bool) -> None: """Init scales and amaxes for fwd | bwd.""" @@ -770,12 +771,23 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: if param.device == torch.device('meta'): param = param.to(device='cuda') - if 'weight' in name: - # Initialize weight values on device - param = self.param_init_meta[name].init_as_weight(param) - elif 'bias' in name: - # Reset biases to zero - param = self.param_init_meta[name].init_as_bias(param) + # Initialize the parameter values on device + init_fn = self.param_init_meta[name].init_fn + get_rng_state_tracker = self.param_init_meta[name].get_rng_state_tracker + if get_rng_state_tracker is None: + init_fn(param) + else: + with get_rng_state_tracker().fork(): + init_fn(param) + + # If primary weights are in fp8, wrap the parameter as Float8Tensor + fp8_meta_index = self.param_init_meta[name].fp8_meta_index + if self.primary_weights_in_fp8 and fp8_meta_index is not None: + param = Float8Tensor.to_float8( + param, + fp8_meta=self.fp8_meta, + fp8_meta_index=fp8_meta_index + ) # Redo parameter wrap in case we broke it above setattr(self, name, torch.nn.Parameter(param)) diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py index 653e23f4f3..67fd1e8118 100644 --- a/transformer_engine/pytorch/module/layernorm.py +++ b/transformer_engine/pytorch/module/layernorm.py @@ -4,6 +4,7 @@ """LayerNorm API""" import os +import warnings from typing import Union, Tuple, Optional import torch @@ -139,7 +140,8 @@ def __init__( ) setattr(self.weight, "sequence_parallel", sequence_parallel) setattr(self.bias, "sequence_parallel", sequence_parallel) - self.reset_layer_norm_parameters() + + self.reset_parameters(defer_init=(device == 'meta')) # These many SMs are subtracted from the total SM count when calling forward # and backward LayerNorm C APIs. These envvars can be used to prevent the LN @@ -150,12 +152,24 @@ def __init__( def reset_layer_norm_parameters(self) -> None: """Init LN params""" + warnings.warn( + ("This method will be deprecated in an upcoming release. " + "Update your code to use LayerNorm.reset_parameters() instead."), + DeprecationWarning, + stacklevel=2 + ) if not self.zero_centered_gamma: init.ones_(self.weight) else: init.zeros_(self.weight) init.zeros_(self.bias) + def reset_parameters(self, defer_init=False) -> None: + if defer_init: + return + init.constant_(self.weight, float(not self.zero_centered_gamma)) + init.zeros_(self.bias) + @no_torch_dynamo() def forward(self, inp: torch.Tensor) -> torch.Tensor: """LayerNorm FWD""" diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index d36d5a9923..b6b7f38718 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -25,6 +25,7 @@ from ..utils import ( divide, get_default_init_method, + init_method_constant, cast_if_needed, assert_dim_for_fp8_exec, clear_tensor_data, @@ -749,43 +750,25 @@ def __init__( self.sequence_parallel = (self.tp_size > 1) and sequence_parallel self.eps = eps - self.layer_norm_weight = torch.nn.Parameter( + layer_norm_weight = torch.nn.Parameter( torch.empty(in_features, device=device, dtype=params_dtype) ) + self.register_parameter('layer_norm_weight', layer_norm_weight, + init_fn=init_method_constant(float(not self.zero_centered_gamma))) setattr(self.layer_norm_weight, "sequence_parallel", self.sequence_parallel) if self.normalization != "RMSNorm": - self.layer_norm_bias = torch.nn.Parameter( + layer_norm_bias = torch.nn.Parameter( torch.empty(in_features, device=device, dtype=params_dtype) ) + self.register_parameter('layer_norm_bias', layer_norm_bias) setattr(self.layer_norm_bias, "sequence_parallel", self.sequence_parallel) else: self.layer_norm_bias = None - self.reset_layer_norm_parameters() - temp_weight = torch.empty( + self.weight_tensor = torch.empty( self.out_features, self.in_features, device=device, dtype=params_dtype) - initialize_affine_weight_gpu( - temp_weight, - init_method, - get_rng_state_tracker, - partition_dim=1 if self.parallel_mode == "row" else 0, - stride=1, - ) - - if self.primary_weights_in_fp8: - self.init_fp8_metadata() - self.fp8_meta["update_amax_and_scale_fwd"] = True - - self.weight_tensor = Float8Tensor.to_float8( - temp_weight, - fp8_meta=self.fp8_meta, - fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, - ) - else: - self.weight_tensor = temp_weight - if self.use_bias: self.bias_tensor = torch.empty( self.out_features, @@ -794,9 +777,6 @@ def __init__( else: self.bias_tensor = torch.Tensor().to(dtype=params_dtype, device=device) - with torch.no_grad(): - self.bias_tensor.zero_() - # Configure parameter splits self.weight_names = [] self.bias_names = [] @@ -861,7 +841,10 @@ def __init__( if is_subview: weight = weight[split_start:split_end] weight = torch.nn.Parameter(weight) - self.register_parameter(self.weight_names[i], weight) + self.register_parameter(self.weight_names[i], weight, + init_fn=init_method, + get_rng_state_tracker=get_rng_state_tracker, + fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT) # Construct bias parameter if needed if self.use_bias: @@ -892,8 +875,13 @@ def __init__( del self.weight_tensor del self.bias_tensor - self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features))) + if self.primary_weights_in_fp8: + self.init_fp8_metadata() + self.fp8_meta["update_amax_and_scale_fwd"] = True + + self.reset_parameters(defer_init=(device == 'meta')) + self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features))) # For RPL, bias has to be added after TP collectives # So it cannot be fused with the GEMM @@ -911,6 +899,12 @@ def __init__( def reset_layer_norm_parameters(self) -> None: """Init LN params""" + warnings.warn( + ("This method will be deprecated in an upcoming release. " + "Update your code to use LayerNormLinear.reset_parameters() instead."), + DeprecationWarning, + stacklevel=2 + ) if not self.zero_centered_gamma: init.ones_(self.layer_norm_weight) else: diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index e5e884cd22..54236f21c1 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -30,6 +30,7 @@ from ..utils import ( divide, get_default_init_method, + init_method_constant, cast_if_needed, assert_dim_for_fp8_exec, clear_tensor_data, @@ -1170,91 +1171,76 @@ def __init__( # LN init self.eps = eps - self.layer_norm_weight = Parameter( + layer_norm_weight = Parameter( torch.empty(hidden_size, device=device, dtype=params_dtype) ) + self.register_parameter('layer_norm_weight', layer_norm_weight, + init_fn=init_method_constant(float(not self.zero_centered_gamma))) setattr(self.layer_norm_weight, "sequence_parallel", self.sequence_parallel) if self.normalization != "RMSNorm": - self.layer_norm_bias = Parameter( + layer_norm_bias = Parameter( torch.empty(hidden_size, device=device, dtype=params_dtype) ) + self.register_parameter('layer_norm_bias', layer_norm_bias) setattr(self.layer_norm_bias, "sequence_parallel", self.sequence_parallel) else: self.layer_norm_bias = None - self.reset_layer_norm_parameters() + # FC1 init if self.activation in ['reglu', 'geglu', 'swiglu']: fc1_output_features = 2 * self.size_per_partition else: fc1_output_features = self.size_per_partition - # FC1 init - fc1_temp_weight = torch.empty( - fc1_output_features, hidden_size, device=device, dtype=params_dtype) - - initialize_affine_weight_gpu( - fc1_temp_weight, - init_method, - get_rng_state_tracker, - set_tp_attributes=False, - ) - if self.primary_weights_in_fp8: - self.init_fp8_metadata(num_gemms=2) - self.fp8_meta["update_amax_and_scale_fwd"] = True - - fc1_temp_weight = Float8Tensor.to_float8( - fc1_temp_weight, - fp8_meta=self.fp8_meta, - fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, + fc1_weight = Parameter( + torch.empty( + fc1_output_features, hidden_size, device=device, dtype=params_dtype ) - - self.fc1_weight = Parameter(fc1_temp_weight) + ) + self.register_parameter('fc1_weight', fc1_weight, + init_fn=init_method, + get_rng_state_tracker=get_rng_state_tracker, + fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT) set_tensor_model_parallel_attributes(self.fc1_weight, True, 0, 1) self.fp8_weight_shapes.append(self.fc1_weight.shape) if self.use_bias: - self.fc1_bias = Parameter( + fc1_bias = Parameter( torch.empty(fc1_output_features, device=device, dtype=params_dtype) ) + self.register_parameter('fc1_bias', fc1_bias) set_tensor_model_parallel_attributes(self.fc1_bias, True, 0, 1) else: self.fc1_bias = torch.Tensor().to(dtype=params_dtype, device=device) - with torch.no_grad(): - self.fc1_bias.zero_() - # FC2 init - fc2_temp_weight = torch.empty( - hidden_size, self.size_per_partition, device=device, dtype=params_dtype) - - initialize_affine_weight_gpu( - fc2_temp_weight, - output_layer_init_method, - get_rng_state_tracker, - set_tp_attributes=False, + fc2_weight = Parameter( + torch.empty(hidden_size, self.size_per_partition, device=device, dtype=params_dtype) ) - - if self.primary_weights_in_fp8: - fc2_temp_weight = Float8Tensor.to_float8( - fc2_temp_weight, - fp8_meta=self.fp8_meta, - fp8_meta_index=tex.FP8FwdTensors.GEMM2_WEIGHT, - ) - - self.fc2_weight = Parameter(fc2_temp_weight) + self.register_parameter('fc2_weight', fc2_weight, + init_fn=output_layer_init_method, + get_rng_state_tracker=get_rng_state_tracker, + fp8_meta_index=tex.FP8FwdTensors.GEMM2_WEIGHT) set_tensor_model_parallel_attributes(self.fc2_weight, True, 1, 1) self.fp8_weight_shapes.append(self.fc2_weight.shape) if self.use_bias: - self.fc2_bias = Parameter( + fc2_bias = Parameter( torch.empty(hidden_size, device=device, dtype=params_dtype) ) + self.register_parameter('fc2_bias', fc2_bias) # RPL if self.set_parallel_mode: setattr(self.fc2_bias, "sequence_parallel", sequence_parallel) else: self.fc2_bias = torch.Tensor().to(dtype=params_dtype, device=device) + if self.primary_weights_in_fp8: + self.init_fp8_metadata(num_gemms=2) + self.fp8_meta["update_amax_and_scale_fwd"] = True + + self.reset_parameters(defer_init=(device == 'meta')) + # For RPL, bias has to be added after TP collectives # So it cannot be fused with the GEMM if self.set_parallel_mode and self.apply_bias: @@ -1262,9 +1248,6 @@ def __init__( else: self.gemm_bias_unfused_add = False - with torch.no_grad(): - self.fc2_bias.zero_() - if self.bias_gelu_nvfusion: set_jit_fusion_options() if seq_length and micro_batch_size: @@ -1281,6 +1264,12 @@ def __init__( def reset_layer_norm_parameters(self) -> None: """Init LN params""" + warnings.warn( + ("This method will be deprecated in an upcoming release. " + "Update your code to use LayerNormMLP.reset_parameters() instead."), + DeprecationWarning, + stacklevel=2 + ) if not self.zero_centered_gamma: init.ones_(self.layer_norm_weight) else: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 90bc3bb8af..2cad516881 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -766,6 +766,10 @@ def __init__( del self.weight_tensor del self.bias_tensor + if self.primary_weights_in_fp8: + self.init_fp8_metadata() + self.fp8_meta["update_amax_and_scale_fwd"] = True + self.reset_parameters(defer_init=(device == 'meta')) self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features))) diff --git a/transformer_engine/pytorch/module/rmsnorm.py b/transformer_engine/pytorch/module/rmsnorm.py index 8da16d1c38..a3f65dffcd 100644 --- a/transformer_engine/pytorch/module/rmsnorm.py +++ b/transformer_engine/pytorch/module/rmsnorm.py @@ -4,6 +4,7 @@ """RMSNorm API""" import os +import warnings from typing import Union, Tuple, Optional import torch @@ -141,7 +142,8 @@ def __init__( ) ) setattr(self.weight, "sequence_parallel", sequence_parallel) - self.reset_rms_norm_parameters() + + self.reset_parameters(defer_init=(device == 'meta')) # These many SMs are subtracted from the total SM count when calling forward # and backward RMSNorm C APIs. These envvars can be used to prevent the LN @@ -152,11 +154,21 @@ def __init__( def reset_rms_norm_parameters(self) -> None: """Init RMSNorm params""" + warnings.warn( + ("This method will be deprecated in an upcoming release. " + "Update your code to use RMSNorm.reset_parameters() instead."), + DeprecationWarning, + stacklevel=2 + ) if not self.zero_centered_gamma: init.ones_(self.weight) else: init.zeros_(self.weight) + def reset_parameters(self, defer_init=False) -> None: + if defer_init: + return + init.constant_(self.weight, float(not self.zero_centered_gamma)) @no_torch_dynamo() def forward(self, inp: torch.Tensor) -> torch.Tensor: diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 6250c07d60..819b3d4827 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -40,6 +40,21 @@ def get_default_init_method() -> Callable: return init_method_normal(0.023) +def init_method_constant(val: float) -> Callable: + """Init method to set all tensor elements to a constant value.""" + if val == 1.0: + def init_(tensor: torch.Tensor) -> Callable: + return torch.nn.init.ones_(tensor) + elif val == 0.0: + def init_(tensor: torch.Tensor) -> Callable: + return torch.nn.init.zeros_(tensor) + else: + def init_(tensor: torch.Tensor) -> Callable: + return torch.nn.init.constant_(tensor, val) + + return init_ + + def init_method_normal(sigma: float) -> Callable: """Init method based on N(0, sigma).""" From df7e75b852cc073687aaf4264256569ebd67ea94 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Tue, 16 Jan 2024 17:38:51 +0000 Subject: [PATCH 8/9] fixed linting errors Signed-off-by: Alp Dener --- transformer_engine/pytorch/module/_common.py | 1 - transformer_engine/pytorch/module/layernorm.py | 1 + transformer_engine/pytorch/module/layernorm_linear.py | 5 ++--- transformer_engine/pytorch/module/layernorm_mlp.py | 7 +++---- transformer_engine/pytorch/module/rmsnorm.py | 1 + 5 files changed, 7 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 114c91f979..c0172f1fba 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -206,4 +206,3 @@ def __post_init__(self): assert self.parent is not None if self.init_fn is None: self.init_fn = get_default_init_method() - diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py index 67fd1e8118..fac941306f 100644 --- a/transformer_engine/pytorch/module/layernorm.py +++ b/transformer_engine/pytorch/module/layernorm.py @@ -165,6 +165,7 @@ def reset_layer_norm_parameters(self) -> None: init.zeros_(self.bias) def reset_parameters(self, defer_init=False) -> None: + """Init LayerNorm parameters""" if defer_init: return init.constant_(self.weight, float(not self.zero_centered_gamma)) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index b6b7f38718..2e6803f992 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -34,7 +34,6 @@ set_tensor_model_parallel_attributes, get_distributed_world_size, allreduce, - initialize_affine_weight_gpu, reduce_scatter_along_first_dim, gather_along_first_dim, ) @@ -755,13 +754,13 @@ def __init__( ) self.register_parameter('layer_norm_weight', layer_norm_weight, init_fn=init_method_constant(float(not self.zero_centered_gamma))) - setattr(self.layer_norm_weight, "sequence_parallel", self.sequence_parallel) + setattr(self.layer_norm_weight, "sequence_parallel", self.sequence_parallel) # pylint: disable=access-member-before-definition if self.normalization != "RMSNorm": layer_norm_bias = torch.nn.Parameter( torch.empty(in_features, device=device, dtype=params_dtype) ) self.register_parameter('layer_norm_bias', layer_norm_bias) - setattr(self.layer_norm_bias, "sequence_parallel", self.sequence_parallel) + setattr(self.layer_norm_bias, "sequence_parallel", self.sequence_parallel) # pylint: disable=access-member-before-definition else: self.layer_norm_bias = None diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 54236f21c1..8f88d725ad 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -39,7 +39,6 @@ set_tensor_model_parallel_attributes, get_distributed_world_size, allreduce, - initialize_affine_weight_gpu, reduce_scatter_along_first_dim, gather_along_first_dim, ) @@ -1182,7 +1181,7 @@ def __init__( torch.empty(hidden_size, device=device, dtype=params_dtype) ) self.register_parameter('layer_norm_bias', layer_norm_bias) - setattr(self.layer_norm_bias, "sequence_parallel", self.sequence_parallel) + setattr(self.layer_norm_bias, "sequence_parallel", self.sequence_parallel) # pylint: disable=access-member-before-definition else: self.layer_norm_bias = None @@ -1209,7 +1208,7 @@ def __init__( torch.empty(fc1_output_features, device=device, dtype=params_dtype) ) self.register_parameter('fc1_bias', fc1_bias) - set_tensor_model_parallel_attributes(self.fc1_bias, True, 0, 1) + set_tensor_model_parallel_attributes(self.fc1_bias, True, 0, 1) # pylint: disable=access-member-before-definition else: self.fc1_bias = torch.Tensor().to(dtype=params_dtype, device=device) @@ -1231,7 +1230,7 @@ def __init__( self.register_parameter('fc2_bias', fc2_bias) # RPL if self.set_parallel_mode: - setattr(self.fc2_bias, "sequence_parallel", sequence_parallel) + setattr(self.fc2_bias, "sequence_parallel", sequence_parallel) # pylint: disable=access-member-before-definition else: self.fc2_bias = torch.Tensor().to(dtype=params_dtype, device=device) diff --git a/transformer_engine/pytorch/module/rmsnorm.py b/transformer_engine/pytorch/module/rmsnorm.py index a3f65dffcd..cad357de04 100644 --- a/transformer_engine/pytorch/module/rmsnorm.py +++ b/transformer_engine/pytorch/module/rmsnorm.py @@ -166,6 +166,7 @@ def reset_rms_norm_parameters(self) -> None: init.zeros_(self.weight) def reset_parameters(self, defer_init=False) -> None: + """Reset RMSNorm parameters""" if defer_init: return init.constant_(self.weight, float(not self.zero_centered_gamma)) From 63581daff3062fa934a794aedaef8384239db82b Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Wed, 17 Jan 2024 17:03:40 +0000 Subject: [PATCH 9/9] removed unnecessary reference to the parent module of parameter, added clarifying comments in parameter reset Signed-off-by: Alp Dener --- transformer_engine/pytorch/module/_common.py | 2 -- transformer_engine/pytorch/module/base.py | 5 ++++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index c0172f1fba..d2ab776288 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -196,13 +196,11 @@ class _ParameterInitMeta: """ Stores essential metadata needed to support deferred parameter initialization. """ - parent: torch.nn.Module init_fn: Optional[Callable] = get_default_init_method() get_rng_state_tracker: Optional[Callable] = None fp8_meta_index: Optional[int] = None def __post_init__(self): """Safeguard reference to the parameter's parent module and initialization function.""" - assert self.parent is not None if self.init_fn is None: self.init_fn = get_default_init_method() diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index d3d5fc38f5..ad1f383617 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -755,7 +755,7 @@ def register_parameter(self, name, param, **kwargs): metedata used in deferred initialization. """ super().register_parameter(name, param) - self.param_init_meta[name] = _ParameterInitMeta(self, **kwargs) + self.param_init_meta[name] = _ParameterInitMeta(**kwargs) def reset_parameters(self, defer_init: Optional[bool] = False) -> None: """ @@ -790,6 +790,9 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: ) # Redo parameter wrap in case we broke it above + # NOTE: Currently this can only be broken when primary weights are in Fp8 but + # re-applying the nn.Parameter() wrap is a no-op when the input is already + # a parameter so we always re-apply it just for extra safety. setattr(self, name, torch.nn.Parameter(param)) @abstractmethod