diff --git a/apex/_autocast_utils.py b/apex/_autocast_utils.py index e86c6c6a5..3a92a83f3 100644 --- a/apex/_autocast_utils.py +++ b/apex/_autocast_utils.py @@ -3,6 +3,9 @@ import torch +__all__ = ["_cast_if_autocast_enabled"] + + def _get_autocast_dtypes() -> Sequence[torch.dtype]: if torch.cuda.is_bf16_supported(): return [torch.half, torch.bfloat16] diff --git a/apex/contrib/optimizers/distributed_fused_adam.py b/apex/contrib/optimizers/distributed_fused_adam.py index 550068022..117d7294b 100644 --- a/apex/contrib/optimizers/distributed_fused_adam.py +++ b/apex/contrib/optimizers/distributed_fused_adam.py @@ -170,6 +170,8 @@ def __init__(self): self.status = DistributedFusedAdam.GradientStatus.READY # Request object for asynchronous communication self.sync_request = None + # Params that have generated grads + self.grads_generated = set() def sync_wait(self): """Wait for asynchronous communication to finish""" @@ -298,7 +300,6 @@ def __init__(self, # Objects for gradient synchronization self._grads_buckets = collections.defaultdict(self.GradientBucket) - self._grads_generated = set() self._pipeline_streams = [torch.cuda.Stream() for _ in range(self.pipeline_size)] # Divide gradients by factor before optimizer step. Used for @@ -322,16 +323,39 @@ def __init__(self, # Attach hooks for gradient synchronization self._register_post_backward_hooks() + # Allocate contiguous gradient buffer if needed + if self.contiguous_grad_buffer: + self._init_grad_buffer() + + def _make_post_backward_hook(self, param, param_group_id, param_id): + """Create callback function to call after param generates grad + + Lazily initialize parameter and try launching grad sync. + + """ + def hook(*unused): + with self._lock: + need_to_initialize = 'fragments' not in self.state[param] + if need_to_initialize: + self._init_param_state(param, param_group_id, param_id) + if self.greedy_grad_copy: + self._grad_copy(param) + if self.overlap_grad_sync: + self._try_start_bucket_grad_sync( + params=[param], + ignore_last_bucket=need_to_initialize, + ) + return hook + def _register_post_backward_hooks(self): """Attach hooks for gradient synchronization - Optimizer state for parameters are initialized lazily as they - are encountered in the backward pass. + Also synchronizes param values between processes and counts + number of parameters being optimized. """ self._num_grads = 0 - grad_buffer_size = 0 - self._lock = threading.Lock() + self._lock = threading.Lock() # Not sure if needed self._grad_accs = [] for param_group_id, group in enumerate(self.param_groups): for param_id, param in enumerate(group['params']): @@ -343,40 +367,34 @@ def _register_post_backward_hooks(self): if param.requires_grad: self._num_grads += 1 - # Callback after gradient is generated - def wrapper(p, p_group_id, p_id): - p_tmp = p.expand_as(p) - grad_acc = p_tmp.grad_fn.next_functions[0][0] - def reduction_hook(*unused): - with self._lock: - if 'fragments' not in self.state[p]: - self._init_param_state(p, p_group_id, p_id) - if self.greedy_grad_copy: - self._grad_copy(p) - if self.overlap_grad_sync: - self._try_start_bucket_grad_sync( - params=[p], - ignore_last_bucket=True, - ) - grad_acc.register_hook(reduction_hook) - self._grad_accs.append(grad_acc) - wrapper(param, param_group_id, param_id) - - # Gradient size, with padding for alignment + # Register callback for after grad is generated + param_tmp = param.expand_as(param) + grad_acc = param_tmp.grad_fn.next_functions[0][0] + hook = self._make_post_backward_hook( + param, + param_group_id, + param_id, + ) + grad_acc.register_hook(hook) + self._grad_accs.append(grad_acc) + + def _init_grad_buffer(self): + """Allocate contiguous buffer for grad buckets""" + grad_buffer_size = 0 + for group in self.param_groups: + for param in group['params']: + if param.requires_grad: grad_size = _round_to_multiple(param.numel(), self.alignment) grad_buffer_size += grad_size - - # Allocate contiguous gradient buffer if needed - if self.contiguous_grad_buffer: - grad_buffer_size = _round_to_multiple( - grad_buffer_size, - self.bucket_size, - ) - self._grad_buffer = torch.zeros( - [grad_buffer_size], - dtype=self.dtype, - device=self.device, - ) + grad_buffer_size = _round_to_multiple( + grad_buffer_size, + self.bucket_size, + ) + self._grad_buffer = torch.zeros( + [grad_buffer_size], + dtype=self.dtype, + device=self.device, + ) def init_params(self, params=None): """Initialize optimizer state for parameters @@ -516,8 +534,7 @@ def zero_grad(self, set_to_none=True): param.grad.zero_() # Reset other state - self._grads_generated = set() - self._inv_grad_scale = torch.full([1], 1.0, dtype=self.dtype, device=self.device) + self._grad_scale = torch.full([], 1.0, dtype=self.dtype, device=self.device) self._grad_norm = None def _grad_copy(self, param): @@ -606,13 +623,10 @@ def _force_bucket_grad_sync(self): device=self.device, ) - # Reset set of generated gradients - self._grads_generated = set() - def _try_start_bucket_grad_sync( self, params=[], - ignore_last_bucket=True, + ignore_last_bucket=False, ): """Launches gradient synchronization if enough buckets are ready @@ -632,38 +646,28 @@ def _try_start_bucket_grad_sync( # Register params that have generated grads for param in params: - self._grads_generated.add(param) for fragment in self.state[param]['fragments']: bucket_id = fragment.bucket_id + bucket = self._grads_buckets[bucket_id] bucket_fragments = self.state['buckets'][bucket_id].fragments - is_filled = True - for other_fragment in reversed(bucket_fragments): - param_group_id = other_fragment.param_group_id - param_id = other_fragment.param_id - other_param = self.param_groups[param_group_id]['params'][param_id] - if other_param not in self._grads_generated: - is_filled = False - break - if is_filled: - bucket = self._grads_buckets[bucket_id] + bucket.grads_generated.add(param) + if len(bucket.grads_generated) == len(bucket_fragments): bucket.status = self.GradientStatus.FULLY_FILLED # Launch reductions if enough buckets are ready - if len(self._grads_generated) == self._num_grads: - self._force_bucket_grad_sync() - else: - filled_buckets = [] - for bucket_id, bucket in sorted(self._grads_buckets.items()): - if ignore_last_bucket and bucket_id == len(self.state['buckets'])-1: - continue - if bucket.status == self.GradientStatus.FULLY_FILLED: - filled_buckets.append(bucket) - pipeline_size = _round_to_multiple( - len(filled_buckets), - self.pipeline_size, - ) - if pipeline_size > 0: - self._start_bucket_grad_sync(filled_buckets[:pipeline_size]) + filled_buckets = [] + for bucket_id, bucket in sorted(self._grads_buckets.items()): + if ignore_last_bucket and bucket_id == len(self.state['buckets'])-1: + continue + if bucket.status == self.GradientStatus.FULLY_FILLED: + filled_buckets.append(bucket) + pipeline_size = _round_to_multiple( + len(filled_buckets), + self.pipeline_size, + round_up=False, + ) + if pipeline_size > 0: + self._start_bucket_grad_sync(filled_buckets[:pipeline_size]) def _start_bucket_grad_sync(self, buckets): """Synchronize gradient buckets @@ -692,13 +696,22 @@ def _start_bucket_grad_sync(self, buckets): stream.wait_stream(main_stream) for i, bucket in enumerate(buckets): bucket.status = self.GradientStatus.SYNCING - stream = self._pipeline_streams[i % self.pipeline_size] - with torch.cuda.stream(stream): - - # Reduce-scatter over distributed process group - bucket.sync_wait() - if self.distributed_size == 1: - bucket.sync_grads_shard = bucket.grads_bucket + bucket.grads_generated.clear() + bucket.sync_wait() + if self.distributed_size == 1: + bucket.sync_grads_shard = bucket.grads_bucket + else: + bucket.sync_grads_shard = torch.zeros( + [self.shard_size], + dtype=self.grad_sync_dtype, + device=self.device, + ) + grads_bucket_shards = [ + bucket.grads_bucket[i*self.shard_size:(i+1)*self.shard_size] + for i in range(self.distributed_size) + ] + if self._reduce_scatter_no_copy: + no_copy_kwarg = { 'no_copy': True } else: with torch.cuda.stream(main_stream): bucket.sync_grads_shard = torch.zeros( @@ -805,7 +818,7 @@ def grad_sync(self): ) self._force_bucket_grad_sync() - def _local_grad_norm(self, parameters=[], norm_type=2.0): + def _local_grad_norm(self, parameters=None, norm_type=2.0): """Local contribution to parameter gradient norm Returns square of 2-norm. Other norms are not yet supported. @@ -821,7 +834,7 @@ def _local_grad_norm(self, parameters=[], norm_type=2.0): # Make sure that gradients have been reduced self.grad_sync() - if not parameters or len(parameters) == self._num_grads: + if parameters is None or len(parameters) == self._num_grads: # Compute norm of all local gradients dummy_overflow_buf = torch.zeros([1], dtype=torch.int32, device='cuda') grad_norm_sq = multi_tensor_applier( @@ -852,7 +865,7 @@ def _local_grad_norm(self, parameters=[], norm_type=2.0): return grad_norm_sq.detach().view([]) - def grad_norm(self, parameters=[], norm_type=2.0, force=False): + def grad_norm(self, parameters=None, norm_type=2.0, force=False): """Gradient norm of parameters in optimizer The norm is computed over all gradients together, as if they @@ -886,7 +899,7 @@ def grad_norm(self, parameters=[], norm_type=2.0, force=False): self._grad_norm = grad_norm_sq.sqrt() return self._grad_norm.detach() - def clip_grad_norm(self, max_norm, parameters=[], norm_type=2.0): + def clip_grad_norm(self, max_norm, parameters=None, norm_type=2.0): """Clips gradient norm of parameters in optimizer The norm is computed over all gradients together, as if they @@ -1180,28 +1193,25 @@ def state_dict(self, gather_on_root=True): # Split data into chunks and gather on root rank # Note: Assuming we are using the NCCL backend, communication # must happen on the GPU. We split the data into fixed-size - # chunks so that the GPU memory usage is limited to - # (chunk_size * distributed_size) bytes. + # chunks to limit GPU memory usage. # TODO: Avoid chunking with direct communication between CPUs main_stream = torch.cuda.current_stream() for stream in self._pipeline_streams: stream.wait_stream(main_stream) for stream_id, offset in enumerate(range(0, max_state_size, chunk_size)): stream_id %= self.pipeline_size - - # Buffers for chunk - if self.distributed_rank == 0: - gathered_chunks = [ - gathered_chunks_buffers[stream_id][i*chunk_size:(i+1)*chunk_size] - for i in range(self.distributed_size) - ] - else: - chunk = chunk_buffers[stream_id] - - # Perform communication on parallel stream stream = self._pipeline_streams[stream_id] with torch.cuda.stream(stream): + # Buffers for chunk + if self.distributed_rank == 0: + gathered_chunks = [ + gathered_chunks_buffers[stream_id][i*chunk_size:(i+1)*chunk_size] + for i in range(self.distributed_size) + ] + else: + chunk = chunk_buffers[stream_id] + # Copy to GPU if self.distributed_rank != 0 and offset < local_state_size: local_chunk_size = min(chunk_size, local_state_size-offset) @@ -1216,24 +1226,30 @@ def state_dict(self, gather_on_root=True): ) # Gather on root - if self.distributed_rank == 0: - if self._gather_no_copy: - no_copy_kwarg = { 'no_copy': True } + # Note: Call in main stream to avoid memory pool + # overheads from internal memory allocations in + # gather. + main_stream.wait_stream(stream) + with torch.cuda.stream(main_stream): + if self.distributed_rank == 0: + if self._gather_no_copy: + no_copy_kwarg = { 'no_copy': True } + else: + no_copy_kwarg = {} + torch.distributed.gather( + gathered_chunks[0], + gathered_chunks, + dst=self._process_group_ranks[0], + group=self.process_group, + **no_copy_kwarg, + ) else: - no_copy_kwarg = {} - torch.distributed.gather( - gathered_chunks[0], - gathered_chunks, - dst=self._process_group_ranks[0], - group=self.process_group, - **no_copy_kwarg, - ) - else: - torch.distributed.gather( - chunk, - dst=self._process_group_ranks[0], - group=self.process_group, - ) + torch.distributed.gather( + chunk, + dst=self._process_group_ranks[0], + group=self.process_group, + ) + stream.wait_stream(main_stream) # Copy back to CPU if self.distributed_rank == 0: diff --git a/apex/contrib/test/optimizers/test_dist_adam.py b/apex/contrib/test/optimizers/test_dist_adam.py index bd23ce2ae..6040e8eb4 100644 --- a/apex/contrib/test/optimizers/test_dist_adam.py +++ b/apex/contrib/test/optimizers/test_dist_adam.py @@ -6,6 +6,7 @@ from torch.testing._internal import common_utils from apex.contrib.optimizers.distributed_fused_adam import DistributedFusedAdam from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase +from apex.testing.common_utils import skipIfRocm class SimpleModel(torch.nn.Module): @@ -250,6 +251,7 @@ def test_clip_grad_norm(self): dist_model.parameters()): torch.testing.assert_close(dist_param, ref_param) + @skipIfRocm def test_grad_scaler(self): torch.manual_seed(self.seed + self.rank) diff --git a/apex/contrib/test/run_rocm_extensions.py b/apex/contrib/test/run_rocm_extensions.py index c7801988b..ae5e7d01f 100644 --- a/apex/contrib/test/run_rocm_extensions.py +++ b/apex/contrib/test/run_rocm_extensions.py @@ -1,26 +1,76 @@ +import argparse +import os import unittest import sys -test_dirs = ["groupbn", "fused_dense", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", "."] # "." for test_label_smoothing.py -ROCM_BLACKLIST = [ - "layer_norm" +#test_dirs = ["fused_dense", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", "optimizers", ".", "groupbn"] # "." for test_label_smoothing.py +#ROCM_BLACKLIST = [ +# "layer_norm" +#] + +TEST_ROOT = os.path.dirname(os.path.abspath(__file__)) +TEST_DIRS = [ + "fused_dense", + "layer_norm", # not fully supported on ROCm + "conv_bias_relu",# not fully supported on ROCm + "fmha", # not fully supported on ROCm + #"cudnn_gbn", # not fully supported on ROCm + #"bottleneck", # not fully supported on ROCm + "multihead_attn", + "transducer", + "focal_loss", + "index_mul_2d", + "optimizers", + "xentropy", + "clip_grad", + "groupbn", +] + +DEFAULT_TEST_DIRS = [ + "fused_dense", + "multihead_attn", + "transducer", + "focal_loss", + "index_mul_2d", + "optimizers", + "xentropy", + "clip_grad", + "groupbn", ] -runner = unittest.TextTestRunner(verbosity=2) +def parse_args(): + parser = argparse.ArgumentParser( + description="Extension test runner", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--include", + nargs="+", + choices=TEST_DIRS, + default=DEFAULT_TEST_DIRS, + help="select a set of tests to run (defaults to ALL tests).", + ) + args, _ = parser.parse_known_args() + return args -errcode = 0 +def main(args): + runner = unittest.TextTestRunner(verbosity=2) + errcode = 0 + for test_dir in args.include: + test_dir = os.path.join(TEST_ROOT, test_dir) + print(test_dir) + suite = unittest.TestLoader().discover(test_dir) -for test_dir in test_dirs: - if test_dir in ROCM_BLACKLIST: - continue - suite = unittest.TestLoader().discover(test_dir) + print("\nExecuting tests from " + test_dir) - print("\nExecuting tests from " + test_dir) + result = runner.run(suite) - result = runner.run(suite) + if not result.wasSuccessful(): + errcode = 1 - if not result.wasSuccessful(): - errcode = 1 + sys.exit(errcode) -sys.exit(errcode) +if __name__ == '__main__': + args = parse_args() + main(args) diff --git a/apex/contrib/test/run_rocm_extensions.sh b/apex/contrib/test/run_rocm_extensions.sh new file mode 100644 index 000000000..302b42189 --- /dev/null +++ b/apex/contrib/test/run_rocm_extensions.sh @@ -0,0 +1,2 @@ +#!/bin/bash +APEX_TEST_WITH_ROCM=1 APEX_SKIP_FLAKY_TEST=1 python3 run_rocm_extensions.py diff --git a/apex/contrib/test/test_label_smoothing.py b/apex/contrib/test/xentropy/test_label_smoothing.py similarity index 100% rename from apex/contrib/test/test_label_smoothing.py rename to apex/contrib/test/xentropy/test_label_smoothing.py diff --git a/apex/mlp/mlp.py b/apex/mlp/mlp.py index bae38f3f8..be0901eb8 100644 --- a/apex/mlp/mlp.py +++ b/apex/mlp/mlp.py @@ -1,9 +1,12 @@ from copy import copy import math + import torch from torch import nn + +from apex._autocast_utils import _cast_if_autocast_enabled import mlp_cuda -from .. import amp + class MlpFunction(torch.autograd.Function): @staticmethod @@ -21,7 +24,11 @@ def backward(ctx, grad_o): del ctx.outputs return (None, None, *grads) -mlp_function = amp.half_function(MlpFunction.apply) + +def mlp_function(bias, activation, *args): + autocast_args = _cast_if_autocast_enabled(bias, activation, *args) + return MlpFunction.apply(*autocast_args) + class MLP(torch.nn.Module): """Launch MLP in C++ @@ -32,16 +39,16 @@ class MLP(torch.nn.Module): relu (bool): Default True """ def __init__(self, mlp_sizes, bias=True, activation='relu'): - super(MLP, self).__init__() + super().__init__() self.num_layers = len(mlp_sizes) - 1 self.mlp_sizes = copy(mlp_sizes) self.bias = 1 if bias else 0 - if activation is 'none': + if activation == 'none': self.activation = 0 - elif activation is 'relu': + elif activation == 'relu': self.activation = 1 - elif activation is 'sigmoid': + elif activation == 'sigmoid': self.activation = 2 else: raise TypeError("activation must be relu or none.") diff --git a/apex/transformer/_ucc_util.py b/apex/transformer/_ucc_util.py new file mode 100644 index 000000000..a286e2ab1 --- /dev/null +++ b/apex/transformer/_ucc_util.py @@ -0,0 +1,9 @@ +from torch import distributed as dist + +HAS_UCC = hasattr(dist, "is_ucc_available") and dist.is_ucc_available() +if not HAS_UCC: + try: + import torch_ucc + HAS_UCC = True + except ImportError: + HAS_UCC = False diff --git a/apex/transformer/functional/fused_softmax.py b/apex/transformer/functional/fused_softmax.py index 8ceaffef9..b0bfebb93 100644 --- a/apex/transformer/functional/fused_softmax.py +++ b/apex/transformer/functional/fused_softmax.py @@ -92,10 +92,73 @@ def backward(ctx, output_grads): def scaled_masked_softmax(inputs, mask, scale): + # input is 4D tensor (b, np, sq, sk) + if mask is not None: + args = _cast_if_autocast_enabled(inputs, mask, scale) + with torch.cuda.amp.autocast(enabled=False): + return ScaledMaskedSoftmax.apply(*args) + else: + args = _cast_if_autocast_enabled(inputs, scale) + with torch.cuda.amp.autocast(enabled=False): + return ScaledSoftmax.apply(*args) + + +class GenericScaledMaskedSoftmax(torch.autograd.Function): + @staticmethod + def forward(ctx, inputs, mask, scale): + import generic_scaled_masked_softmax_cuda + + scale_t = torch.tensor([scale]) + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, mask, scale_t[0]) + ctx.save_for_backward(softmax_results, scale_t) + return softmax_results + + @staticmethod + def backward(ctx, output_grads): + import generic_scaled_masked_softmax_cuda_new + + softmax_results, scale_t = ctx.saved_tensors + + input_grads = generic_scaled_masked_softmax_cuda.backward(output_grads, softmax_results, scale_t[0]) + return input_grads, None, None + + +def generic_scaled_masked_softmax(inputs, mask, scale): # input is 4D tensor (b, np, sq, sk) args = _cast_if_autocast_enabled(inputs, mask, scale) with torch.cuda.amp.autocast(enabled=False): - return ScaledMaskedSoftmax.apply(*args) + return GenericScaledMaskedSoftmax.apply(*args) + + +class ScaledSoftmax(torch.autograd.Function): + """ + Fused operation which performs following two operations in sequence + 1. Scale the tensor. + 2. Perform softmax. + """ + + @staticmethod + def forward(ctx, inputs, scale): + import scaled_softmax_cuda + + scale_t = torch.tensor([scale]) + + softmax_results = scaled_softmax_cuda.forward( + inputs, scale_t[0] + ) + ctx.save_for_backward(softmax_results, scale_t) + return softmax_results + + @staticmethod + def backward(ctx, output_grads): + import scaled_softmax_cuda + + softmax_results, scale_t = ctx.saved_tensors + + input_grads = scaled_softmax_cuda.backward( + output_grads, softmax_results, scale_t[0] + ) + return input_grads, None, None class FusedScaleMaskSoftmax(torch.nn.Module): @@ -164,14 +227,14 @@ def is_kernel_available(self, mask, b, np, sq, sk): and self.input_in_float16 # input must be fp16 and ( self.attn_mask_type == AttnMaskType.causal - or (self.attn_mask_type == AttnMaskType.padding and mask is not None) + or self.attn_mask_type == AttnMaskType.padding ) - and 16 < sk <= 2048 # sk must be 16 ~ 2048 + and 16 < sk <= 4096 # sk must be 16 ~ 4096 and sq % 4 == 0 # sq must be divisor of 4 and sk % 4 == 0 # sk must be divisor of 4 and attn_batches % 4 == 0 # np * b must be divisor of 4 ): - if 0 <= sk <= 2048: + if 0 <= sk <= 4096: batch_per_block = self.get_batch_per_block(sq, sk, b, np) if self.attn_mask_type == AttnMaskType.causal: @@ -209,3 +272,30 @@ def get_batch_per_block(sq, sk, b, np): import scaled_masked_softmax_cuda return scaled_masked_softmax_cuda.get_batch_per_block(sq, sk, b, np) + +class GenericFusedScaleMaskSoftmax(FusedScaleMaskSoftmax): + """ + Generic version of FusedSacleMaskSoftmax. + It removes the seq-len limitations and has slight performance degragation compared with FusedScaleMaskSoftmax + + fused operation: scaling + mask + softmax + + Arguments: + input_in_fp16: flag to indicate if input in fp16 data format. + input_in_bf16: flag to indicate if input in bf16 data format. + scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion + mask_func: mask function to be applied. + softmax_in_fp32: if true, softmax in performed at fp32 precision. + scale: scaling factor used in input tensor scaling. + """ + + def __init__( + self, input_in_fp16, input_in_bf16, scaled_masked_softmax_fusion, mask_func, softmax_in_fp32, scale, + ): + super().__init__(input_in_fp16, input_in_bf16, AttnMaskType.padding, scaled_masked_softmax_fusion, mask_func, softmax_in_fp32, scale) + self.scaled_masked_softmax_fusion = generic_scaled_masked_softmax + + def is_kernel_available(self, mask, b, np, sq, sk): + if self.scaled_masked_softmax_fusion and 0 < sk: # user want to fuse # sk must be 1 ~ + return True + return False \ No newline at end of file diff --git a/apex/transformer/parallel_state.py b/apex/transformer/parallel_state.py index a8d16bfd3..f5821c2c2 100644 --- a/apex/transformer/parallel_state.py +++ b/apex/transformer/parallel_state.py @@ -21,6 +21,7 @@ import torch from apex.transformer.log_util import get_transformer_logger +from apex.transformer._ucc_util import HAS_UCC _logger = get_transformer_logger(__name__) @@ -126,7 +127,8 @@ def initialize_model_parallel( assert default_backend is None or default_backend in ("nccl", "ucc") assert p2p_backend is None or p2p_backend in ("nccl", "ucc") if "ucc" in (default_backend, p2p_backend): - check_torch_ucc_availability() + if not HAS_UCC: + raise ImportError("UCC backend requires pytorch source build with UCC installed and enabled") warnings.warn("`ucc` backend support is experimental", ExperimentalWarning) if default_backend == "ucc": warnings.warn("The UCC's functionality as `default_backend` is not well verified", ExperimentalWarning) @@ -671,12 +673,3 @@ def destroy_model_parallel(): # Used to warn when the UCC is specified. class ExperimentalWarning(Warning): pass - - -def check_torch_ucc_availability() -> None: - try: - import torch_ucc # NOQA - except ImportError: - raise ImportError( - "UCC backend requires [torch_ucc](https://github.com/facebookresearch/torch_ucc) but not found" - ) diff --git a/apex/transformer/pipeline_parallel/p2p_communication.py b/apex/transformer/pipeline_parallel/p2p_communication.py index 6c4b0d93d..aa8470c5f 100644 --- a/apex/transformer/pipeline_parallel/p2p_communication.py +++ b/apex/transformer/pipeline_parallel/p2p_communication.py @@ -128,6 +128,7 @@ def _communicate( fp32_residual_connection: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> Tuple[Union[torch.Tensor, FutureTensor, None], Union[torch.Tensor, FutureTensor, None]]: """Base function for communication of tensors between stages. @@ -159,6 +160,8 @@ def _communicate( sequence_parallel_enabled: Set to :obj:`True` if sequence parallel is enabled. This argument is here for consistency with Megatron-LM. This argument has an effect on the communication optimization, not on tensor_shape update. + sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. + To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. Returns: tuple containing @@ -267,8 +270,9 @@ def tensor_recv_next_wait(): torch.cuda.synchronize() tensor_recv_next_waitfunc = tensor_recv_next_wait else: - # To protect against race condition when using batch_isend_irecv(). - torch.cuda.synchronize() + if sync_batch_comm: + # To protect against race condition when using batch_isend_irecv(). + torch.cuda.synchronize() # If using scatter-gather optimization, gather smaller chunks. if scatter_gather_optimization_doable: @@ -325,6 +329,7 @@ def recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Receive tensor from previous rank in pipeline (forward receive).""" @@ -342,6 +347,7 @@ def recv_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-recv").stop() @@ -354,6 +360,7 @@ def recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Receive tensor from next rank in pipeline (backward receive).""" @@ -370,6 +377,7 @@ def recv_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-recv").stop() @@ -384,6 +392,7 @@ def send_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> None: """Send tensor to next rank in pipeline (forward send).""" @@ -401,6 +410,7 @@ def send_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-send").stop() @@ -413,6 +423,7 @@ def send_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> None: """Send tensor to previous rank in pipeline (backward send).""" @@ -429,6 +440,7 @@ def send_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-send").stop() @@ -441,6 +453,7 @@ def send_forward_recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Batched send and recv with next rank in pipeline.""" @@ -457,6 +470,7 @@ def send_forward_recv_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-send-backward-recv").stop() @@ -470,6 +484,7 @@ def send_backward_recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Batched send and recv with previous rank in pipeline.""" @@ -486,6 +501,7 @@ def send_backward_recv_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-send-forward-recv").stop() @@ -500,6 +516,7 @@ def send_forward_recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor]: """Batched recv from previous rank and send to next rank in pipeline.""" @@ -514,6 +531,7 @@ def send_forward_recv_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-send-forward-recv").stop() @@ -528,6 +546,7 @@ def send_backward_recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor]: """Batched recv from next rank and send to previous rank in pipeline.""" @@ -542,6 +561,7 @@ def send_backward_recv_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-send-backward-recv").stop() @@ -558,6 +578,7 @@ def send_forward_backward_recv_forward_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Tuple[Union[torch.Tensor, FutureTensor], Union[torch.Tensor, FutureTensor]]: """Batched send and recv with previous and next ranks in pipeline.""" @@ -572,6 +593,7 @@ def send_forward_backward_recv_forward_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-backward-send-forward-backward-recv").stop() diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py index 5500085a4..6bd3ff139 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py @@ -1,4 +1,4 @@ -from contextlib import contextmanager +import contextlib from typing import List, Union, Optional import torch @@ -20,14 +20,6 @@ _logger = get_transformer_logger(__name__) -@contextmanager -def placeholder_handler(): - try: - yield - finally: - pass - - def forward_backward_no_pipelining( forward_step_func: FwdStepFunc, batch: Batch, @@ -59,7 +51,7 @@ def forward_backward_no_pipelining( disable_autocast: Turn off `enabled` flag of `torch.cuda.amp.autocast` if :obj:`True`. Should be used when your forward and loss computation is in the autocast context to avoid unnecesarily nest autocast context. - custom_sync_context_handler: + custom_sync_context_handler: Context manager to disable asynchronous gradient reductions. **kwargs: Added to handle `tensor_shape` which has no effect on this function. Returns: @@ -77,7 +69,7 @@ def forward_backward_no_pipelining( elif isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel): context_handler = model.no_sync else: - context_handler = placeholder_handler + context_handler = contextlib.nullcontext losses_reduced = [] input_tensor, output_tensor_grad = None, None diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py index 17ad8334d..748b23b4a 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py @@ -1,4 +1,5 @@ -from typing import List, Union, Optional, Sequence +import contextlib +from typing import Any, List, Optional, Sequence, Union import warnings import torch @@ -36,6 +37,8 @@ def _forward_backward_pipelining_with_interleaving( deallocate_pipeline_outputs: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, + custom_sync_context_handler: Optional[Any] = None, + sync_batch_comm: bool = True, **kwargs, ) -> List[Union[torch.Tensor, Sequence[torch.Tensor]]]: """Run interleaved 1F1B schedule with communication between pipeline stages as needed. @@ -70,9 +73,18 @@ def _forward_backward_pipelining_with_interleaving( sequence_parallel_enabled: Set to :obj:`True` for this function to handle sequence length. When :obj:`True`, the sequence length on each tensor model parallel rank is updated to :math:`original\_sequence\_length / tensor\_model\_parallel\_world\_size`. + custom_sync_context_handler: Does nothing if ``None`` (default + value). Otherwise, this is treated as a function to + construct a context manager to disable asynchronous + gradient reductions. Asynchronous gradient reductions are + only enabled in the final backward pass of each model + chunk. + sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. + To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. Returns: a list of loss `torch.Tensor`s if the last stage, empty list otherwise. + """ if not isinstance(model, list): raise RuntimeError("`model` must be a list of `nn.Module`'s'") @@ -83,6 +95,26 @@ def _forward_backward_pipelining_with_interleaving( "This option is not recommended." ) + # Construct helper functions for async grad reductions + if custom_sync_context_handler is not None: + sync_context_handler = custom_sync_context_handler + else: + sync_context_handler = contextlib.nullcontext + sync_context = None + def disable_grad_sync(): + """Disable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is None: + sync_context = sync_context_handler() + sync_context.__enter__() + def enable_grad_sync(): + """Enable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is not None: + sync_context.__exit__(None, None, None) + sync_context = None + disable_grad_sync() + # mypy will blame the following if statement if sequence_parallel_enabled: seq_length, batch_size, hidden = tensor_shape @@ -142,17 +174,41 @@ def _forward_backward_pipelining_with_interleaving( # Helper function definitions. ################################################################################################################### def get_model_chunk_id(microbatch_id: int, forward: bool) -> int: - """Helper function to get the model chunk ID given the iteration number.""" - pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size() - microbatch_id_in_group = microbatch_id % ( - pipeline_parallel_size * num_model_chunks - ) + """Helper function to get the model chunk ID given the iteration number. + + Each model chunk processes pipeline_parallel_size microbatches + at a time. We assume that the number of microbatches is a + multiple of pipeline_parallel_size*num_model_chunks. + """ + microbatch_group_size = pipeline_parallel_size * num_model_chunks + microbatch_id_in_group = microbatch_id % microbatch_group_size model_chunk_id = microbatch_id_in_group // pipeline_parallel_size if not forward: model_chunk_id = num_model_chunks - model_chunk_id - 1 return model_chunk_id - def forward_step_helper(microbatch_id: int, curr_iters: List[int]) -> torch.Tensor: + def is_last_microbatch_for_model_chunk(microbatch_id: int) -> bool: + """Helper function to check if an iteration is the last for a model + chunk. + + Each model chunk processes pipeline_parallel_size microbatches + at a time. We assume that the number of microbatches is a + multiple of pipeline_parallel_size*num_model_chunks. + """ + microbatch_group_size = pipeline_parallel_size * num_model_chunks + num_microbatch_groups = num_microbatches // microbatch_group_size + microbatch_group_id = microbatch_id // microbatch_group_size + microbatch_id_in_group = microbatch_id % microbatch_group_size + if microbatch_group_id == num_microbatch_groups - 1: + return microbatch_id_in_group % pipeline_parallel_size == pipeline_parallel_size - 1 + else: + return False + + def forward_step_helper( + microbatch_id: int, + curr_iters: List[int], + checkpoint_activations_micro_batch: Optional[bool] = None, + ) -> torch.Tensor: """Helper method to run forward step with model split into chunks (run set_virtual_pipeline_model_parallel_rank() before calling forward_step()). @@ -194,6 +250,8 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: model_type = get_model_type(model[model_chunk_id]) parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id) + if is_last_microbatch_for_model_chunk(microbatch_id): + enable_grad_sync() if parallel_state.is_pipeline_last_stage(): if len(output_tensor_grads[model_chunk_id]) == 0: output_tensor_grads[model_chunk_id].append(None) @@ -208,6 +266,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: grad_scaler=grad_scaler, deallocate_pipeline_outputs=deallocate_pipeline_outputs, ) + disable_grad_sync() return input_tensor_grad @@ -221,6 +280,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) _logger.info("Warmup phase") @@ -269,6 +329,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad) else: @@ -280,6 +341,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) input_tensors[next_forward_model_chunk_id].append(input_tensor) free_output_tensor(output_tensor, deallocate_pipeline_outputs) @@ -365,6 +427,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) free_output_tensor(output_tensor, deallocate_pipeline_outputs) @@ -387,6 +450,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) for k in range(num_microbatches_remaining, num_microbatches): @@ -409,7 +473,11 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) + # Make sure to exit context handler for async grad reductions + enable_grad_sync() + return losses_reduced diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py index 5dc2933ba..500a1d4e1 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py @@ -1,4 +1,5 @@ -from typing import Union, List, Optional, Sequence +import contextlib +from typing import Any, List, Optional, Sequence, Union import warnings import torch @@ -89,6 +90,7 @@ def recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: input_tensors = [] for tensor_shape in tensor_shapes: @@ -101,6 +103,7 @@ def recv_forward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) return input_tensors @@ -112,6 +115,7 @@ def recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: output_tensor_grads = [] for tensor_shape in tensor_shapes: @@ -124,6 +128,7 @@ def recv_backward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) return output_tensor_grads @@ -136,6 +141,7 @@ def send_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> None: if not isinstance(output_tensors, list): output_tensors = [output_tensors] @@ -148,6 +154,7 @@ def send_forward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) @@ -158,6 +165,7 @@ def send_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> None: if not isinstance(input_tensor_grads, list): input_tensor_grads = [input_tensor_grads] @@ -170,6 +178,7 @@ def send_backward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) @@ -180,6 +189,7 @@ def send_forward_recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: if not isinstance(output_tensors, list): output_tensors = [output_tensors] @@ -194,6 +204,7 @@ def send_forward_recv_backward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) output_tensor_grads.append(output_tensor_grad) return output_tensor_grads @@ -206,6 +217,7 @@ def send_backward_recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: if not isinstance(input_tensor_grads, list): input_tensor_grads = [input_tensor_grads] @@ -220,6 +232,7 @@ def send_backward_recv_forward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) input_tensors.append(input_tensor) return input_tensors @@ -239,6 +252,9 @@ def forward_backward_pipelining_without_interleaving( deallocate_pipeline_outputs: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, + custom_sync_context_handler: Optional[Any] = None, + custom_grad_sync_func: Optional[Any] = None, + sync_batch_comm: bool = True, **kwargs, ) -> List[Union[torch.Tensor, Sequence[torch.Tensor]]]: """Run non-interleaved 1F1B schedule, with communication between pipeline stages. @@ -269,9 +285,21 @@ def forward_backward_pipelining_without_interleaving( sequence_parallel_enabled: Set to :obj:`True` for this function to handle sequence length. When :obj:`True`, the sequence length on each tensor model parallel rank is updated to :math:`original\_sequence\_length / tensor\_model\_parallel\_world\_size`. + custom_sync_context_handler: Does nothing if ``None`` (default + value). Otherwise, a function to construct a context + manager that disable asynchronous gradient reductions. + Asynchronous gradient reductions are only enabled in the + first pipeline stage, during the last backward pass. + custom_grad_sync_func: Does nothing if ``None`` (default + value). Otherwise, a function to perform gradient + reductions. This is called in all pipeline stages except + the first, during the bubble overhead. + sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. + To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. Returns: a list of loss `torch.Tensor`s if the last stage, empty list otherwise. + """ # timers = get_timers() @@ -287,6 +315,26 @@ def forward_backward_pipelining_without_interleaving( raise RuntimeError(msg) model: torch.nn.Module = model[0] + # Disable async grad reductions + if custom_sync_context_handler is not None: + sync_context_handler = custom_sync_context_handler + else: + sync_context_handler = contextlib.nullcontext + sync_context = None + def disable_grad_sync(): + """Disable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is None: + sync_context = sync_context_handler() + sync_context.__enter__() + def enable_grad_sync(): + """Enable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is not None: + sync_context.__exit__(None, None, None) + sync_context = None + disable_grad_sync() + # Compute number of warmup microbatches. num_microbatches: int = get_num_microbatches() num_warmup_microbatches: int = ( @@ -334,6 +382,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) cur_microbatch: Optional[torch.Tensor] = get_kth_microbatch(batch, i) output_tensor = forward_step( @@ -352,6 +401,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) if not forward_only: @@ -364,7 +414,12 @@ def forward_backward_pipelining_without_interleaving( # receive this tensor here. if num_microbatches_remaining > 0: _logger.debug("recv_forward before steady state start") - input_tensor: List[Union[None, torch.Tensor, FutureTensor]] = recv_forward(tensor_shapes=recv_tensor_shapes, dtype=dtype, async_comm=async_comm) + input_tensor: List[Union[None, torch.Tensor, FutureTensor]] = recv_forward( + tensor_shapes=recv_tensor_shapes, + dtype=dtype, + async_comm=async_comm, + sync_batch_comm=sync_batch_comm, + ) ################################################################################################################### # Run 1F1B in steady state. @@ -392,6 +447,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) if not last_iteration: @@ -401,6 +457,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) else: @@ -411,6 +468,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # Add input_tensor and output_tensor to end of list. @@ -440,6 +498,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) else: _logger.debug("send bwd and receive fwd") @@ -449,6 +508,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ################################################################################################################### # Run cooldown backward passes. @@ -457,6 +517,12 @@ def forward_backward_pipelining_without_interleaving( if not forward_only: for i in range(num_warmup_microbatches): _logger.debug(f"cooldown iter: {i} / {num_warmup_microbatches}") + + if i == num_warmup_microbatches-1 and rank == 0: + # Async grad reduction in first pipeline stage, during + # last backward pass + enable_grad_sync() + input_tensor = input_tensors.pop(0) output_tensor = output_tensors.pop(0) @@ -466,6 +532,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) input_tensor_grad = backward_step( @@ -484,6 +551,13 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) + # Grad reduction in all pipeline stages except the first, during + # the bubble overhead + enable_grad_sync() + if rank != 0 and custom_grad_sync_func is not None: + custom_grad_sync_func() + return losses_reduced diff --git a/apex/transformer/tensor_parallel/cross_entropy.py b/apex/transformer/tensor_parallel/cross_entropy.py index 3918645c7..b14d7c19c 100644 --- a/apex/transformer/tensor_parallel/cross_entropy.py +++ b/apex/transformer/tensor_parallel/cross_entropy.py @@ -22,7 +22,7 @@ class _VocabParallelCrossEntropy(torch.autograd.Function): @staticmethod - def forward(ctx, vocab_parallel_logits, target): + def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0): # Maximum value along vocab dimension across all GPUs. logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] @@ -72,6 +72,27 @@ def forward(ctx, vocab_parallel_logits, target): # Store softmax, target-mask and masked-target for backward pass. exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + + vocab_size = exp_logits.size(-1) + if label_smoothing > 0: + """ + We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth. + = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt}) + = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i + = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K + From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py + """ + assert 1.0 > label_smoothing > 0.0 + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + + # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs. + log_probs = torch.log(exp_logits) + mean_log_probs = log_probs.mean(dim=-1) + loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs + + ctx.label_smoothing, ctx.vocab_size = label_smoothing, vocab_size ctx.save_for_backward(exp_logits, target_mask, masked_target_1d) return loss @@ -81,6 +102,7 @@ def backward(ctx, grad_output): # Retreive tensors from the forward path. softmax, target_mask, masked_target_1d = ctx.saved_tensors + label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size # All the inputs have softmax as thier gradient. grad_input = softmax @@ -90,14 +112,23 @@ def backward(ctx, grad_output): # Add the gradient from matching classes. arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) - grad_2d[arange_1d, masked_target_1d] -= 1.0 - target_mask.view(-1).float() + + softmax_update = 1.0 - target_mask.view(-1).float() + + if label_smoothing > 0: + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update + average_grad = 1 / vocab_size + grad_2d[arange_1d, :] -= smoothing * average_grad + else: + grad_2d[arange_1d, masked_target_1d] -= softmax_update # Finally elementwise multiplication with the output gradients. grad_input.mul_(grad_output.unsqueeze(dim=-1)) - return grad_input, None + return grad_input, None, None -def vocab_parallel_cross_entropy(vocab_parallel_logits, target): +def vocab_parallel_cross_entropy(vocab_parallel_logits, target, label_smoothing=0.0): """Helper function for the cross entropy.""" - return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target) + return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, label_smoothing) diff --git a/apex/transformer/testing/arguments.py b/apex/transformer/testing/arguments.py index f39c288ce..5c74054ca 100644 --- a/apex/transformer/testing/arguments.py +++ b/apex/transformer/testing/arguments.py @@ -20,7 +20,7 @@ import torch -def parse_args(extra_args_provider=None, defaults={}, +def parse_args(extra_args_provider=None, defaults={}, override_args={}, ignore_unknown_args=False): """Parse all arguments.""" parser = argparse.ArgumentParser(description='Megatron-LM Arguments', @@ -59,16 +59,22 @@ def parse_args(extra_args_provider=None, defaults={}, # Distributed args. args.rank = int(os.getenv('RANK', '0')) args.world_size = int(os.getenv("WORLD_SIZE", '1')) + + for key in override_args: + setattr(args, key, override_args[key]) + # Tensor model parallel size. args.tensor_model_parallel_size = min( args.tensor_model_parallel_size, args.world_size) assert args.world_size % args.tensor_model_parallel_size == 0, 'world size'\ ' ({}) is not divisible by tensor model parallel size ({})'.format( args.world_size, args.tensor_model_parallel_size) + # Pipeline model parallel size. args.pipeline_model_parallel_size = min( args.pipeline_model_parallel_size, (args.world_size // args.tensor_model_parallel_size)) + args.transformer_pipeline_model_parallel_size = ( args.pipeline_model_parallel_size - 1 if args.standalone_embedding_stage else diff --git a/apex/transformer/testing/commons.py b/apex/transformer/testing/commons.py index 226e449db..83f894cb5 100644 --- a/apex/transformer/testing/commons.py +++ b/apex/transformer/testing/commons.py @@ -35,7 +35,7 @@ Batch, ) from apex.transformer.testing import global_vars - +from apex.transformer._ucc_util import HAS_UCC TEST_SUCCESS_MESSAGE = ">> passed the test :-)" @@ -257,7 +257,8 @@ def initialize_distributed(backend="nccl"): if backend not in ("nccl", "ucc"): raise RuntimeError(f"Currently only nccl & ucc are supported but {backend}") if backend == "ucc": - import torch_ucc # NOQA + if not HAS_UCC: + raise ImportError("UCC backend requires pytorch source build with UCC installed and enabled") args = global_vars.get_args() local_rank = args.local_rank diff --git a/apex/transformer/testing/distributed_test_base.py b/apex/transformer/testing/distributed_test_base.py index 7a8168759..bcfbfc032 100644 --- a/apex/transformer/testing/distributed_test_base.py +++ b/apex/transformer/testing/distributed_test_base.py @@ -9,12 +9,7 @@ from torch.testing._internal import common_utils from torch.testing._internal import common_distributed -HAS_TORCH_UCC = None -try: - import torch_ucc - HAS_TORCH_UCC = True -except ImportError: - HAS_TORCH_UCC = False +from apex.transformer._ucc_util import HAS_UCC # NOTE(mkozuki): Version guard for ucc. ref: https://github.com/openucx/ucc/issues/496 _TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION = Version("470.42.01") @@ -38,6 +33,7 @@ def setUp(self) -> None: self._spawn_processes() def tearDown(self) -> None: + torch.cuda.empty_cache() super().tearDown() @property @@ -88,16 +84,16 @@ class NcclDistributedTestBase(DistributedTestBase): DISTRIBUTED_BACKEND = "nccl" - @unittest.skipUnless( - HAS_TORCH_UCC, - "Requires [`torch_ucc`](https://github.com/facebookresearch/torch_ucc)", + HAS_UCC, + "Requires either torch ucc or pytorch build from source with native ucc installed and enabled", ) @unittest.skipUnless( HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER, f"`torch_ucc` requires NVIDIA driver >= {_TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION} but {_driver_version} found. " "See https://github.com/openucx/ucc/issues/496", ) + class UccDistributedTestBase(DistributedTestBase): DISTRIBUTED_BACKEND = "ucc" diff --git a/apex/transformer/testing/global_vars.py b/apex/transformer/testing/global_vars.py index 6b8537456..fa68f2c14 100644 --- a/apex/transformer/testing/global_vars.py +++ b/apex/transformer/testing/global_vars.py @@ -84,11 +84,12 @@ def get_timers(): return _GLOBAL_TIMERS -def set_global_variables(extra_args_provider=None, args_defaults={}, +def set_global_variables(extra_args_provider=None, args_defaults={}, override_args={}, ignore_unknown_args=False): """Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers.""" args = _parse_args(extra_args_provider=extra_args_provider, defaults=args_defaults, + override_args=override_args, ignore_unknown_args=ignore_unknown_args) # _build_num_microbatches_calculator(args) # if args.vocab_file: @@ -98,13 +99,14 @@ def set_global_variables(extra_args_provider=None, args_defaults={}, _set_timers() -def _parse_args(extra_args_provider=None, defaults={}, +def _parse_args(extra_args_provider=None, defaults={}, override_args={}, ignore_unknown_args=False): """Parse entire arguments.""" global _GLOBAL_ARGS _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args') _GLOBAL_ARGS = parse_args(extra_args_provider=extra_args_provider, defaults=defaults, + override_args=override_args, ignore_unknown_args=ignore_unknown_args) return _GLOBAL_ARGS diff --git a/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu b/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu index 60d1e8d1f..878a08075 100644 --- a/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu +++ b/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu @@ -30,6 +30,34 @@ void gemmex_wrapper_fp16( const float* beta, at::BFloat16* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f16_r, + ldc, + C, + rocblas_datatype_f16_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard /*algo*/, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -50,6 +78,7 @@ void gemmex_wrapper_fp16( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } // FP16 inputs and FP16 accumulation @@ -68,6 +97,34 @@ void gemmex_wrapper_fp16( const float* beta, at::Half* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f16_r, + ldc, + C, + rocblas_datatype_f16_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard /*algo*/, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -88,6 +145,7 @@ void gemmex_wrapper_fp16( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } template diff --git a/csrc/megatron/fused_weight_gradient_dense_cuda.cu b/csrc/megatron/fused_weight_gradient_dense_cuda.cu index dfaa1345d..5d0b08876 100644 --- a/csrc/megatron/fused_weight_gradient_dense_cuda.cu +++ b/csrc/megatron/fused_weight_gradient_dense_cuda.cu @@ -30,6 +30,34 @@ void gemmex_wrapper( const float* beta, float* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f32_r, + ldc, + C, + rocblas_datatype_f32_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -50,6 +78,7 @@ void gemmex_wrapper( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } // FP16 Tensor core wrapper around cublas GEMMEx @@ -68,6 +97,34 @@ void gemmex_wrapper( const float* beta, float* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f32_r, + ldc, + C, + rocblas_datatype_f32_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -88,6 +145,7 @@ void gemmex_wrapper( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } // FP32 wrapper around cublas GEMMEx @@ -106,6 +164,34 @@ void gemmex_wrapper( const float *beta, float *C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f32_r, + lda, + B, + rocblas_datatype_f32_r, + ldb, + beta, + C, + rocblas_datatype_f32_r, + ldc, + C, + rocblas_datatype_f32_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -126,6 +212,7 @@ void gemmex_wrapper( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } template diff --git a/csrc/megatron/generic_scaled_masked_softmax.cpp b/csrc/megatron/generic_scaled_masked_softmax.cpp new file mode 100644 index 000000000..4bd9f0e7f --- /dev/null +++ b/csrc/megatron/generic_scaled_masked_softmax.cpp @@ -0,0 +1,82 @@ +/* coding=utf-8 + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace multihead_attn +{ + namespace fused_softmax + { + namespace generic_scaled_masked_softmax + { + + torch::Tensor fwd_cuda( + torch::Tensor const &input, + torch::Tensor const &mask, + float scale_factor); + + torch::Tensor bwd_cuda( + torch::Tensor const &output_grads, + torch::Tensor const &softmax_results, + float scale_factor); + + torch::Tensor fwd( + torch::Tensor const &input, + torch::Tensor const &mask, + float scale_factor) + { + AT_ASSERTM(input.dim() == 4, "expected 4D tensor"); + AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) || + (input.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + AT_ASSERTM(mask.dim() == 4, "expected 4D tensor"); + + return fwd_cuda(input, mask, scale_factor); + } + + torch::Tensor bwd( + torch::Tensor const &output_grads, + torch::Tensor const &softmax_results, + float scale_factor) + { + + AT_ASSERTM(output_grads.dim() == 4, "expected 3D tensor"); + AT_ASSERTM(softmax_results.dim() == 4, "expected 3D tensor"); + + AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) || + (output_grads.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) || + (softmax_results.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + + return bwd_cuda(output_grads, softmax_results, scale_factor); + } + + } // end namespace generic_scaled_masked_softmax + } // end namespace fused_softmax +} // end namespace multihead_attn + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", + &multihead_attn::fused_softmax::generic_scaled_masked_softmax::fwd, + "Self Multihead Attention scaled, time masked softmax -- Forward."); + + m.def("backward", + &multihead_attn::fused_softmax::generic_scaled_masked_softmax::bwd, + "Self Multihead Attention scaled, time masked softmax -- Backward."); +} diff --git a/csrc/megatron/generic_scaled_masked_softmax.h b/csrc/megatron/generic_scaled_masked_softmax.h new file mode 100644 index 000000000..555ff462f --- /dev/null +++ b/csrc/megatron/generic_scaled_masked_softmax.h @@ -0,0 +1,384 @@ +/* coding=utf-8 + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +template +struct Add { + __device__ __forceinline__ T operator()(T a, T b) const { + return a + b; + } +}; + +template +struct Max { + __device__ __forceinline__ T operator()(T a, T b) const { + return a < b ? b : a; + } +}; + +template +__device__ __forceinline__ T WARP_SHFL_DOWN_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) +{ +#if CUDA_VERSION >= 9000 + return __shfl_down_sync(mask, value, laneMask, width); +#else + return __shfl_down(value, laneMask, width); +#endif +} + +template class ReduceOp> +__device__ __forceinline__ acc_t warp_reduce_new(acc_t val) { + ReduceOp r; + #pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) + { + val = r(val, WARP_SHFL_DOWN_NATIVE(val, offset, WARP_SIZE)); + } + return val; +} + + +template +__global__ void scaled_masked_softmax_warp_backward_new( + output_t *gradInput, //[batches, attn_heads, q_len, k_len] + input_t *grad, + const input_t *output, //[batches, attn_heads, q_len, k_len] + acc_t scale, + int element_count) +{ + int threads_per_block = blockDim.x; + //the first element_count*2 elements are used for cache, the last 128 is used for reduction + extern __shared__ acc_t shared_data[]; + input_t *local_data = (input_t *)shared_data; + input_t *output_data = &local_data[element_count]; + // maximum shared cached 128, enough for 4096 elements reduction into 4096/32= 128 elements + acc_t *shared = (acc_t *)(&(local_data[element_count*2])); + + int num_reductions = (element_count - 1) / threads_per_block + 1; + + int offset = blockIdx.x * element_count; + + int local_idx = threadIdx.x; + int lane = threadIdx.x % C10_WARP_SIZE; + int wid = threadIdx.x / C10_WARP_SIZE; + int warps_per_thread_block = threads_per_block / C10_WARP_SIZE; + + // load the data to local data + acc_t val = 0.0; + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = output[offset + i*threads_per_block + local_idx]; + output_data[i*threads_per_block + local_idx] = val; + local_data[i*threads_per_block + local_idx] = val * grad[offset + i*threads_per_block + local_idx]; + } + __syncthreads(); + } + + // find the sum + for (int i = local_idx; i < (element_count - 1) / C10_WARP_SIZE + 1; i += threads_per_block){ + shared[i] = 0.0; + } + __syncthreads(); + + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = local_data[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0 && wid + warps_per_thread_block * i < (element_count - 1) / C10_WARP_SIZE + 1) { + shared[wid + warps_per_thread_block*i] = val; + } + __syncthreads(); + } + + // final shared reduction + + int shared_mem_len = (element_count - 1) / C10_WARP_SIZE + 1; + int num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + while ( shared_mem_len > 1 ){ + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < shared_mem_len){ + val = shared[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0) { + shared[wid + warps_per_thread_block * i] = val; + } + __syncthreads(); + } + shared_mem_len = num_warps; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + } + val = shared[0]; + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + gradInput[offset + i] = (output_t)(scale*(local_data[i] - output_data[i]*val)); + } +} + +} // end of anonymous namespace + +template +void dispatch_scaled_masked_softmax_backward_new( + output_t *grad_input, + input_t *grad, + const input_t *output, + const acc_t scale, + int query_seq_len, + int key_seq_len, + int batches, + int attn_heads) +{ + if (key_seq_len == 0) + { + return; + } + else + { + int batch_count = batches * attn_heads * query_seq_len; + // use 128 threads per block to maximimize gpu utilization + constexpr int threads_per_block = 128; + int num_warps = (key_seq_len - 1) / C10_WARP_SIZE + 1; + dim3 blocks(batch_count, 1, 1); + dim3 threads(threads_per_block, 1, 1); + + scaled_masked_softmax_warp_backward_new + <<>>(grad_input, grad, output, scale, key_seq_len); + } +} + +/* + * Extended softmax (from native aten pytorch) with following additional features + * 1) input scaling + * 2) Explicit masking + */ +template +__global__ void scaled_masked_softmax_warp_forward_new( + output_t *dst, + const input_t *src, + const uint8_t *mask, + const acc_t scale, + int query_len, // query_len + int attn_heads, + int element_count, // key_len + int pad_batches) // mask batch size +{ + // min threawds_per_block has to be bigger than 128 + int threads_per_block = blockDim.x; + // the first element_count is used for cache, the last 128 is used for reduction + extern __shared__ acc_t local_data[]; + // maximum shared cached 128, enough for 4096 elements reduction into 4096/32= 128 elements + acc_t *shared = &(local_data[element_count]); + // number of 1024 threads reductions + int num_reductions = (element_count - 1) / threads_per_block + 1; + + int offset = blockIdx.x * element_count; + int mask_offset; + int query_id = blockIdx.x % query_len; + if (pad_batches == 1){ + // broadcaste the mask tensor + mask_offset = query_id * element_count; + } + else{ + int mask_batch_id = blockIdx.x / attn_heads / query_len; + mask_offset = (mask_batch_id * query_len + query_id) * element_count; + } + + int local_idx = threadIdx.x; + int lane = threadIdx.x % C10_WARP_SIZE; + int wid = threadIdx.x / C10_WARP_SIZE; + int warps_per_thread_block = threads_per_block / C10_WARP_SIZE; + + // load the data to local data + for (int i = local_idx; i < element_count; i += threads_per_block) + { + // TODO, use the copy vector method + if (mask[mask_offset + i] == 1) + { + local_data[i] = -10000.0; + } + else + { + local_data[i] = src[offset + i] * scale; + } + } + + // first find the max value + for (int i = local_idx; i < (element_count - 1) / C10_WARP_SIZE + 1; i += threads_per_block){ + shared[i] = -10000.0; + } + __syncthreads(); + acc_t val = -10000.0; + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = local_data[i*threads_per_block + local_idx]; + } + else{ + val = -10000.0; + } + __syncthreads(); + val = warp_reduce_new(val); + + if (lane==0 && wid + warps_per_thread_block * i < (element_count - 1) / C10_WARP_SIZE + 1) { + shared[wid + warps_per_thread_block*i] = val; + } + __syncthreads(); + } + + // final shared reduction + int shared_mem_len = (element_count - 1) / C10_WARP_SIZE + 1; + int num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + while ( shared_mem_len > 1 ){ + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < shared_mem_len){ + val = shared[i*threads_per_block + local_idx]; + } + else{ + val = -10000.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0) { + shared[wid + warps_per_thread_block * i] = val; + } + __syncthreads(); + } + shared_mem_len = num_warps; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + } + + acc_t reduced_val = shared[0]; + if (reduced_val < -10000.0 + 0.1){ + // if everything is masked, pay attention to nothing + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + dst[offset + i] = 0.0; + } + return; + } + + // update the values + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + local_data[i] = std::exp(local_data[i] - reduced_val); + } + + // find the sum + for (int i = local_idx; i < (element_count - 1) / C10_WARP_SIZE + 1; i += threads_per_block){ + shared[i] = 0.0; + } + __syncthreads(); + + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = local_data[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + + val = warp_reduce_new(val); + if (lane==0 && wid + warps_per_thread_block * i < (element_count - 1) / C10_WARP_SIZE + 1) { + shared[wid + warps_per_thread_block*i] = val; + } + __syncthreads(); + } + + shared_mem_len = (element_count - 1) / C10_WARP_SIZE + 1; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + while ( shared_mem_len > 1 ){ + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < shared_mem_len){ + val = shared[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0) { + shared[wid + warps_per_thread_block * i] = val; + } + __syncthreads(); + } + shared_mem_len = num_warps; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + } + + reduced_val = shared[0]; + + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + dst[offset + i] = local_data[i] / reduced_val; + } +} + + +template +void dispatch_scaled_masked_softmax_forward_new( + output_t *dst, + const input_t *src, + const uint8_t *mask, + const input_t scale, + int query_seq_len, + int key_seq_len, + int batches, + int attn_heads, + int pad_batches) +{ + if (key_seq_len == 0) { + return; + } else { + int batch_count = batches * attn_heads * query_seq_len; + + // use 128 threads per block to maximimize gpu utilization + constexpr int threads_per_block = 128; + + // calculate the needed shared memory + int num_warps = (key_seq_len - 1) / C10_WARP_SIZE + 1; + + dim3 blocks(batch_count, 1, 1); + dim3 threads(threads_per_block, 1, 1); + scaled_masked_softmax_warp_forward_new + <<>>(dst, src, mask, scale, query_seq_len, attn_heads, key_seq_len, pad_batches); + } +} diff --git a/csrc/megatron/generic_scaled_masked_softmax_cuda.cu b/csrc/megatron/generic_scaled_masked_softmax_cuda.cu new file mode 100644 index 000000000..5af9aa9cc --- /dev/null +++ b/csrc/megatron/generic_scaled_masked_softmax_cuda.cu @@ -0,0 +1,116 @@ +/* coding=utf-8 + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#ifndef __HIP_PLATFORM_HCC__ +#include +#endif +#include +#include +#include "generic_scaled_masked_softmax.h" +#include "type_shim.h" + +namespace multihead_attn { +namespace fused_softmax { +namespace generic_scaled_masked_softmax { + +torch::Tensor fwd_cuda( + torch::Tensor const& input, + torch::Tensor const& mask, + float scale_factor) +{ + // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = input.size(0); + const int pad_batches = mask.size(0); + const int attn_heads = input.size(1); + const int query_seq_len = input.size(2); + const int key_seq_len = input.size(3); + TORCH_INTERNAL_ASSERT(pad_batches == 1 || pad_batches == batches); + TORCH_INTERNAL_ASSERT(mask.size(1) == 1); + TORCH_INTERNAL_ASSERT(mask.size(2) == query_seq_len); + TORCH_INTERNAL_ASSERT(mask.size(3) == key_seq_len); + + // Output + auto act_options = input.options().requires_grad(false); + torch::Tensor softmax_results = + torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options); + + // Softmax Intermediate Result Ptr + void* input_ptr = static_cast(input.data_ptr()); + void* mask_ptr = static_cast(mask.data_ptr()); + void* softmax_results_ptr = static_cast(softmax_results.data_ptr()); + + DISPATCH_HALF_AND_BFLOAT( + input.scalar_type(), + "dispatch_scaled_masked_softmax_forward", + dispatch_scaled_masked_softmax_forward_new( + reinterpret_cast(softmax_results_ptr), + reinterpret_cast(input_ptr), + reinterpret_cast(mask_ptr), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads, + pad_batches); + ); + return softmax_results; +} + +torch::Tensor bwd_cuda( + torch::Tensor const& output_grads_, + torch::Tensor const& softmax_results_, + float scale_factor) { + + auto output_grads = output_grads_.contiguous(); + auto softmax_results = softmax_results_.contiguous(); + + //output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = output_grads.size(0); + const int attn_heads = output_grads.size(1); + const int query_seq_len = output_grads.size(2); + const int key_seq_len = output_grads.size(3); + + auto act_options = output_grads.options(); + torch::Tensor input_grad = + torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options); + + void* output_grads_ptr = static_cast(output_grads.data_ptr()); + + //Softmax Grad + DISPATCH_HALF_AND_BFLOAT( + output_grads_.scalar_type(), + "dispatch_scaled_masked_softmax_backward", + dispatch_scaled_masked_softmax_backward_new( + reinterpret_cast(static_cast(input_grad.data_ptr())), + reinterpret_cast(output_grads_ptr), + reinterpret_cast(softmax_results.data_ptr()), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads); + ); + + //backward pass is completely in-place + return input_grad; +} +} +} +} diff --git a/csrc/megatron/scaled_masked_softmax.h b/csrc/megatron/scaled_masked_softmax.h index 78a29cf3b..ad35b3baa 100644 --- a/csrc/megatron/scaled_masked_softmax.h +++ b/csrc/megatron/scaled_masked_softmax.h @@ -90,6 +90,117 @@ __device__ __forceinline__ void warp_reduce(acc_t* sum) { } } + +/* + * Extended softmax (from native aten pytorch) with following additional features + * 1) input scaling + */ +template +__global__ void scaled_softmax_warp_forward( + output_t *dst, + const input_t *src, + const acc_t scale, + int micro_batch_size, + int element_count) +{ + // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and + // warp_size of method warp_softmax_forward_kernel. + constexpr int next_power_of_two = 1 << log2_elements; + constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE; + constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE; + constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1; + constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4; + + // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, ) + // gridDim/blockIdx = (seq_len, attn_heads, batches) + int first_batch = (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z))+ threadIdx.y) * WARP_BATCH; + + // micro_batch_size might not be a multiple of WARP_BATCH. Check how + // many batches have to computed within this WARP. + int local_batches = micro_batch_size - first_batch; + if (local_batches > WARP_BATCH) + local_batches = WARP_BATCH; + + // there might be multiple batches per warp. compute the index within the batch + int local_idx = threadIdx.x; + + src += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx; + dst += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx; + + // load data from global memory + acc_t elements[WARP_BATCH][WARP_ITERATIONS]; + input_t temp_data[ELEMENTS_PER_LDG_STG]; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + int batch_element_count = (i >= local_batches) ? 0 : element_count; + + #pragma unroll + for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) { + int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE; + + if (element_index < batch_element_count) { + int itr_idx = i*element_count+it*WARP_SIZE; + copy_vector(temp_data, src + itr_idx); + + #pragma unroll + for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { + elements[i][it + element] = (acc_t)temp_data[element] * scale; + } + } else { + #pragma unroll + for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { + elements[i][it + element] = -std::numeric_limits::infinity(); + } + } + } + } + + // compute max_value + acc_t max_value[WARP_BATCH]; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + max_value[i] = elements[i][0]; + #pragma unroll + for (int it = 1; it < WARP_ITERATIONS; ++it) { + max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it]; + } + } + warp_reduce(max_value); + + acc_t sum[WARP_BATCH] { 0.0f }; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + #pragma unroll + for (int it = 0; it < WARP_ITERATIONS; ++it) { + elements[i][it] = std::exp((elements[i][it] - max_value[i])); + sum[i] += elements[i][it]; + } + } + warp_reduce(sum); + + // store result + output_t out[ELEMENTS_PER_LDG_STG]; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + if (i >= local_batches) + break; + #pragma unroll + for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) { + int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE; + if (element_index < element_count) { + #pragma unroll + for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { + out[element] = elements[i][it + element] / sum[i]; + } + copy_vector(dst + i * element_count + it * WARP_SIZE, out); + } else { + break; + } + } + } +} + + /* * Extended softmax (from native aten pytorch) with following additional features * 1) input scaling @@ -326,6 +437,98 @@ int get_batch_per_block(int query_seq_len, int key_seq_len, int batches, int att return batches_per_block; } +template +void dispatch_scaled_softmax_forward( + output_t *dst, + const input_t *src, + const input_t scale, + int query_seq_len, + int key_seq_len, + int batches, + int attn_heads) +{ + TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 4096 ); + if (key_seq_len == 0) { + return; + } else { + int log2_elements = log2_ceil(key_seq_len); + const int next_power_of_two = 1 << log2_elements; + int batch_count = batches * attn_heads * query_seq_len; + + // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_forward. + int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE; + + // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_forward. + int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1; + + // use 128 threads per block to maximimize gpu utilization + constexpr int threads_per_block = 128; + + int warps_per_block = (threads_per_block / warp_size); + int batches_per_block = warps_per_block * batches_per_warp; + TORCH_INTERNAL_ASSERT(query_seq_len%batches_per_block == 0); + dim3 blocks(query_seq_len/batches_per_block, attn_heads, batches); + dim3 threads(warp_size, warps_per_block, 1); + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 1: // 2 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 2: // 4 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 3: // 8 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 4: // 16 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 5: // 32 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 6: // 64 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 7: // 128 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 8: // 256 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 9: // 512 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 10: // 1024 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 11: // 2048 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 12: // 4096 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + default: + break; + } + } +} + template void dispatch_scaled_masked_softmax_forward( output_t *dst, @@ -338,7 +541,7 @@ void dispatch_scaled_masked_softmax_forward( int attn_heads, int pad_batches) { - TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 2048 ); + TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 4096 ); if (key_seq_len == 0) { return; } else { @@ -410,6 +613,10 @@ void dispatch_scaled_masked_softmax_forward( scaled_masked_softmax_warp_forward <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches); break; + case 12: // 4096 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches); + break; default: break; } @@ -427,7 +634,7 @@ void dispatch_scaled_masked_softmax_backward( int batches, int attn_heads) { - TORCH_INTERNAL_ASSERT( key_seq_len >= 0 && key_seq_len <= 2048 ); + TORCH_INTERNAL_ASSERT( key_seq_len >= 0 && key_seq_len <= 4096 ); if (key_seq_len == 0) { return; } else { @@ -498,6 +705,11 @@ void dispatch_scaled_masked_softmax_backward( scaled_masked_softmax_warp_backward <<>>(grad_input, grad, output, scale, batch_count, key_seq_len); break; + case 12: // 4096 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, key_seq_len); + break; + default: break; } diff --git a/csrc/megatron/scaled_masked_softmax_cuda.cu b/csrc/megatron/scaled_masked_softmax_cuda.cu index 60966706b..de911a446 100644 --- a/csrc/megatron/scaled_masked_softmax_cuda.cu +++ b/csrc/megatron/scaled_masked_softmax_cuda.cu @@ -18,7 +18,9 @@ #include #include #include -//#include +#ifndef __HIP_PLATFORM_HCC__ +#include +#endif #include #include #include "scaled_masked_softmax.h" @@ -44,7 +46,7 @@ torch::Tensor fwd_cuda( const int attn_heads = input.size(1); const int query_seq_len = input.size(2); const int key_seq_len = input.size(3); - TORCH_INTERNAL_ASSERT(key_seq_len <= 2048); + TORCH_INTERNAL_ASSERT(key_seq_len <= 4096); TORCH_INTERNAL_ASSERT(query_seq_len > 1); TORCH_INTERNAL_ASSERT(pad_batches == 1 || pad_batches == batches); TORCH_INTERNAL_ASSERT(mask.size(1) == 1); diff --git a/csrc/megatron/scaled_softmax.cpp b/csrc/megatron/scaled_softmax.cpp new file mode 100644 index 000000000..a25cfa328 --- /dev/null +++ b/csrc/megatron/scaled_softmax.cpp @@ -0,0 +1,74 @@ +/* coding=utf-8 + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace multihead_attn { +namespace fused_softmax { +namespace scaled_softmax { + +torch::Tensor fwd_cuda( + torch::Tensor const& input, + float scale_factor); + +torch::Tensor bwd_cuda( + torch::Tensor const& output_grads, + torch::Tensor const& softmax_results, + float scale_factor); + +torch::Tensor fwd( + torch::Tensor const& input, + float scale_factor) { + AT_ASSERTM(input.dim() == 4, "expected 4D tensor"); + AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) || + (input.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + + return fwd_cuda(input, scale_factor); +} + +torch::Tensor bwd( + torch::Tensor const& output_grads, + torch::Tensor const& softmax_results, + float scale_factor) { + + AT_ASSERTM(output_grads.dim() == 4, "expected 3D tensor"); + AT_ASSERTM(softmax_results.dim() == 4, "expected 3D tensor"); + + AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) || + (output_grads.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) || + (softmax_results.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + + return bwd_cuda(output_grads, softmax_results, scale_factor); +} + +} // end namespace scaled_softmax +} // end namespace fused_softmax +} // end namespace multihead_attn + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", + &multihead_attn::fused_softmax::scaled_softmax::fwd, + "Self Multihead Attention scaled, softmax -- Forward."); + m.def("backward", + &multihead_attn::fused_softmax::scaled_softmax::bwd, + "Self Multihead Attention scaled, softmax -- Backward."); +} + diff --git a/csrc/megatron/scaled_softmax_cuda.cu b/csrc/megatron/scaled_softmax_cuda.cu new file mode 100644 index 000000000..7dcbca68e --- /dev/null +++ b/csrc/megatron/scaled_softmax_cuda.cu @@ -0,0 +1,104 @@ +/* coding=utf-8 + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +//#include +#include +#include +#include "scaled_masked_softmax.h" +#include "type_shim.h" + +namespace multihead_attn { +namespace fused_softmax { +namespace scaled_softmax { + +torch::Tensor fwd_cuda( + torch::Tensor const& input, + float scale_factor) +{ + // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = input.size(0); + const int attn_heads = input.size(1); + const int query_seq_len = input.size(2); + const int key_seq_len = input.size(3); + TORCH_INTERNAL_ASSERT(key_seq_len <= 4096); + TORCH_INTERNAL_ASSERT(query_seq_len > 1); + + // Output + auto act_options = input.options().requires_grad(false); + torch::Tensor softmax_results = + torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options); + + // Softmax Intermediate Result Ptr + void* input_ptr = static_cast(input.data_ptr()); + void* softmax_results_ptr = static_cast(softmax_results.data_ptr()); + + DISPATCH_HALF_AND_BFLOAT( + input.scalar_type(), + "dispatch_scaled_softmax_forward", + dispatch_scaled_softmax_forward( + reinterpret_cast(softmax_results_ptr), + reinterpret_cast(input_ptr), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads); + ); + return softmax_results; +} + +torch::Tensor bwd_cuda( + torch::Tensor const& output_grads_, + torch::Tensor const& softmax_results_, + float scale_factor) { + + auto output_grads = output_grads_.contiguous(); + auto softmax_results = softmax_results_.contiguous(); + + //output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = output_grads.size(0); + const int attn_heads = output_grads.size(1); + const int query_seq_len = output_grads.size(2); + const int key_seq_len = output_grads.size(3); + + void* output_grads_ptr = static_cast(output_grads.data_ptr()); + + //Softmax Grad + DISPATCH_HALF_AND_BFLOAT( + output_grads_.scalar_type(), + "dispatch_scaled_masked_softmax_backward", + dispatch_scaled_masked_softmax_backward( + reinterpret_cast(output_grads_ptr), + reinterpret_cast(output_grads_ptr), + reinterpret_cast(softmax_results.data_ptr()), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads); + ); + + //backward pass is completely in-place + return output_grads; +} +} +} +} + diff --git a/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu b/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu index df022cbbf..c4cbbf17c 100644 --- a/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu +++ b/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu @@ -18,7 +18,9 @@ #include #include #include -//#include +#ifndef __HIP_PLATFORM_HCC__ +#include +#endif #include #include #include "scaled_upper_triang_masked_softmax.h" diff --git a/setup.py b/setup.py index d86a601a4..9a4e7b0b2 100644 --- a/setup.py +++ b/setup.py @@ -292,6 +292,13 @@ def check_if_rocm_pytorch(): include_dirs=[os.path.join(this_dir, 'csrc')], extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) + ext_modules.append( + CUDAExtension(name='generic_scaled_masked_softmax_cuda', + sources=['csrc/megatron/generic_scaled_masked_softmax.cpp', + 'csrc/megatron/generic_scaled_masked_softmax_cuda.cu'], + include_dirs=[os.path.join(this_dir, "csrc")], + extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, + 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) ext_modules.append( CUDAExtension(name='scaled_masked_softmax_cuda', sources=['csrc/megatron/scaled_masked_softmax.cpp', @@ -300,6 +307,38 @@ def check_if_rocm_pytorch(): os.path.join(this_dir, 'csrc/megatron')], extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) + ext_modules.append( + CUDAExtension(name='scaled_softmax_cuda', + sources=['csrc/megatron/scaled_softmax.cpp', + 'csrc/megatron/scaled_softmax_cuda.cu'], + include_dirs=[os.path.join(this_dir, 'csrc'), + os.path.join(this_dir, 'csrc/megatron')], + extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, + 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) + if not IS_ROCM_PYTORCH: + if bare_metal_version >= Version("11.0"): + cc_flag = [] + cc_flag.append("-gencode") + cc_flag.append("arch=compute_70,code=sm_70") + cc_flag.append("-gencode") + cc_flag.append("arch=compute_80,code=sm_80") + if bare_metal_version >= Version("11.1"): + cc_flag.append("-gencode") + cc_flag.append("arch=compute_86,code=sm_86") + if bare_metal_version >= Version("11.8"): + cc_flag.append("-gencode") + cc_flag.append("arch=compute_90,code=sm_90") + + ext_modules.append( + CUDAExtension(name='fused_weight_gradient_mlp_cuda', + sources=['csrc/megatron/fused_weight_gradient_dense.cpp', + 'csrc/megatron/fused_weight_gradient_dense_cuda.cu', + 'csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu'], + include_dirs=[os.path.join(this_dir, 'csrc'), + os.path.join(this_dir, 'csrc/megatron')], + extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, + 'nvcc':append_nvcc_threads(nvcc_args_transformer + ['--use_fast_math'] + cc_flag) + if not IS_ROCM_PYTORCH else hipcc_args_transformer})) if "--bnp" in sys.argv or "--cuda_ext" in sys.argv: @@ -555,7 +594,7 @@ def check_if_rocm_pytorch(): if "--transducer" in sys.argv or "--cuda_ext" in sys.argv: if "--transducer" in sys.argv: sys.argv.remove("--transducer") - + if not IS_ROCM_PYTORCH: raise_if_cuda_home_none("--transducer") diff --git a/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py b/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py index 18219522c..a3826253a 100644 --- a/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py +++ b/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py @@ -3,8 +3,12 @@ import torch -import apex from apex.testing.common_utils import skipFlakyTest +from apex.normalization import FusedLayerNorm +from apex.normalization import FusedRMSNorm +from apex.normalization import MixedFusedLayerNorm +from apex.normalization import MixedFusedRMSNorm + class TestFusedLayerNorm(unittest.TestCase): dtype = torch.float @@ -18,15 +22,15 @@ class TestFusedLayerNorm(unittest.TestCase): def setUp(self): # bias and weight are set to 0 and 1 respectively, so no need to copy parameters from cpu module to the gpu one if not self.mixed_fused: - self.module_cpu_ = apex.normalization.FusedLayerNorm( + self.module_cpu_ = FusedLayerNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).cpu() - self.module_cuda_ = apex.normalization.FusedLayerNorm( + self.module_cuda_ = FusedLayerNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).to(device="cuda", dtype=self.dtype) else: assert self.elementwise_affine - self.module_cpu_ = apex.normalization.MixedFusedLayerNorm( + self.module_cpu_ = MixedFusedLayerNorm( normalized_shape=self.normalized_shape).cpu() - self.module_cuda_ = apex.normalization.MixedFusedLayerNorm( + self.module_cuda_ = MixedFusedLayerNorm( normalized_shape=self.normalized_shape).to(device="cuda", dtype=self.dtype) @@ -65,8 +69,7 @@ def _check_same_output(self, batch_size, contiguous): def _test_same_output(self, batch_size): for contiguous in (True, False): - with self.subTest(contiguous=contiguous): - self._check_same_output(batch_size, contiguous) + self._check_same_output(batch_size, contiguous) def test_layer_norm(self): self._test_same_output(16) @@ -87,15 +90,15 @@ class TestFusedRMSNorm(unittest.TestCase): def setUp(self): # bias and weight are set to 0 and 1 respectively, so no need to copy parameters from cpu module to the gpu one if not self.mixed_fused: - self.module_cpu_ = apex.normalization.FusedRMSNorm( + self.module_cpu_ = FusedRMSNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).cpu() - self.module_cuda_ = apex.normalization.FusedRMSNorm( + self.module_cuda_ = FusedRMSNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).to(device="cuda", dtype=self.dtype) else: assert self.elementwise_affine - self.module_cpu_ = apex.normalization.MixedFusedRMSNorm( + self.module_cpu_ = MixedFusedRMSNorm( normalized_shape=self.normalized_shape).cpu() - self.module_cuda_ = apex.normalization.MixedFusedRMSNorm( + self.module_cuda_ = MixedFusedRMSNorm( normalized_shape=self.normalized_shape).to(device="cuda", dtype=self.dtype) def _check_same_output(self, batch_size, contiguous): @@ -136,8 +139,7 @@ def _check_same_output(self, batch_size, contiguous): def _test_same_output(self, batch_size): for contiguous in (True, False): - with self.subTest(contiguous=contiguous): - self._check_same_output(batch_size, contiguous) + self._check_same_output(batch_size, contiguous) def test_layer_norm(self): self._test_same_output(16) @@ -206,17 +208,17 @@ def _prep_layers(normalized_shape, elementwise_affine, dtype): native = torch.nn.LayerNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ).to(device="cuda", dtype=dtype) - fused = apex.normalization.FusedLayerNorm( + fused = FusedLayerNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ).cuda() return native, fused def _prep_rms_layers(normalized_shape, elementwise_affine, dtype): - native = apex.normalization.FusedRMSNorm( + native = FusedRMSNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ) - fused = apex.normalization.FusedRMSNorm( + fused = FusedRMSNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ).cuda() return native, fused @@ -261,8 +263,7 @@ def _run_test(self, dtype, elementwise_affine): def test_autocast(self): for (dtype, elementwise_affine) in itertools.product(autocast_dtypes, (True, False)): - with self.subTest(f"{dtype}-{elementwise_affine}"): - self._run_test(dtype, elementwise_affine) + self._run_test(dtype, elementwise_affine) @unittest.skip("Skipped on ROCm5.2 due to the failure of reproducing the issue locally. (Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!) Please refer to https://github.com/ROCmSoftwarePlatform/apex/pull/78") class TestAutocastFusedRMSNorm(unittest.TestCase): @@ -294,5 +295,8 @@ def _run_test(self, dtype, elementwise_affine): def test_autocast(self): for (dtype, elementwise_affine) in itertools.product(autocast_dtypes, (True, False)): - with self.subTest(f"{dtype}-{elementwise_affine}"): - self._run_test(dtype, elementwise_affine) + self._run_test(dtype, elementwise_affine) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/L0/run_mlp/test_mlp.py b/tests/L0/run_mlp/test_mlp.py index 615dec95c..dcc97aefb 100644 --- a/tests/L0/run_mlp/test_mlp.py +++ b/tests/L0/run_mlp/test_mlp.py @@ -1,20 +1,24 @@ """Tests for c++ MLP""" -import unittest +from itertools import product from time import time -import numpy as np import torch from torch import nn +from torch.testing._internal import common_utils +from torch.testing._internal.common_device_type import instantiate_device_type_tests +from torch.testing._internal.common_device_type import onlyCUDA from apex.mlp import MLP from apex.testing.common_utils import skipFlakyTest + batch_size = 1024 mlp_sizes = [480, 1024, 1024, 512, 256, 1] num_iters = 10 -class TestMLP(unittest.TestCase): +# note(crcrpar): On Ampere, this test should be run without TF32 enabled. +class TestMLP(common_utils.TestCase): def test_creation(self): MLP(mlp_sizes) @@ -25,114 +29,85 @@ def test_numeric(self): mlp_layers = [] for i in range(mlp.num_layers): linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1]) - mlp.weights[i].data.copy_(linear.weight) - mlp.biases[i].data.copy_(linear.bias) + with torch.no_grad(): + mlp.weights[i].copy_(linear.weight) + mlp.biases[i].copy_(linear.bias) mlp_layers.append(linear) - mlp_layers.append(nn.ReLU(inplace=True)) + mlp_layers.append(nn.ReLU()) ref_mlp = nn.Sequential(*mlp_layers).cuda() - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_() + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda") + .uniform_(-1.0, 1.0) + .requires_grad_() + ) ref_input = test_input.clone().detach().requires_grad_() mlp_out = mlp(test_input) ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + self.assertEqual(mlp_out, ref_out) # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - test_input.grad.detach().cpu().numpy(), - ref_input.grad.detach().cpu().numpy(), - atol=0, rtol=1e-5) - np.testing.assert_allclose( - mlp.biases[0].grad.detach().cpu().numpy(), - ref_mlp[0].bias.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + mlp_out.mean().mul(10.0).backward() + ref_out.mean().mul(10.0).backward() + self.assertEqual(test_input.grad, ref_input.grad) + self.assertEqual(mlp.biases[0].grad, ref_mlp[0].bias.grad) - @skipFlakyTest - def test_no_bias(self): - for use_activation in ['none', 'relu', 'sigmoid']: - mlp = MLP(mlp_sizes, bias=False, activation=use_activation).cuda() - - mlp_layers = [] - for i in range(mlp.num_layers): - linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=False) - mlp.weights[i].data.copy_(linear.weight) - mlp_layers.append(linear) - if use_activation == 'relu': - mlp_layers.append(nn.ReLU(inplace=True)) - if use_activation == 'sigmoid': - mlp_layers.append(nn.Sigmoid()) - - ref_mlp = nn.Sequential(*mlp_layers).cuda() - - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_() - ref_input = test_input.clone().detach().requires_grad_() - mlp_out = mlp(test_input) - ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + def _test_mlp_impl(self, use_activation: str, bias: bool, enable_autocast: bool): + mlp = MLP(mlp_sizes, bias=bias, activation=use_activation).cuda() - # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - test_input.grad.detach().cpu().numpy(), - ref_input.grad.detach().cpu().numpy(), - atol=0, rtol=100) - np.testing.assert_allclose( - mlp.weights[0].grad.detach().cpu().numpy(), - ref_mlp[0].weight.grad.detach().cpu().numpy(), - atol=1e-7, rtol=100) + mlp_layers = [] + for i in range(mlp.num_layers): + linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=bias) + with torch.no_grad(): + mlp.weights[i].copy_(linear.weight) + if bias: + mlp.biases[i].copy_(linear.bias) + mlp_layers.append(linear) + if use_activation == "relu": + mlp_layers.append(nn.ReLU()) + if use_activation == "sigmoid": + mlp_layers.append(nn.Sigmoid()) - @skipFlakyTest - def test_with_bias(self): - for use_activation in ['none', 'relu', 'sigmoid']: - mlp = MLP(mlp_sizes, bias=True, activation=use_activation).cuda() - - mlp_layers = [] - for i in range(mlp.num_layers): - linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=True) - mlp.weights[i].data.copy_(linear.weight) - mlp.biases[i].data.copy_(linear.bias) - mlp_layers.append(linear) - if use_activation == 'relu': - mlp_layers.append(nn.ReLU(inplace=True)) - if use_activation == 'sigmoid': - mlp_layers.append(nn.Sigmoid()) - - ref_mlp = nn.Sequential(*mlp_layers).cuda() - - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_() - ref_input = test_input.clone().detach().requires_grad_() - mlp_out = mlp(test_input) - ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + ref_mlp = nn.Sequential(*mlp_layers).cuda() + + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda") + .uniform_(-1.0, 1.0) + .requires_grad_() + ) + ref_input = test_input.clone().detach().requires_grad_() + with torch.cuda.amp.autocast_mode.autocast(enabled=enable_autocast): + mlp_out = mlp(test_input) + mlp_loss = mlp_out.mean().mul(10.0) # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - test_input.grad.detach().cpu().numpy(), - ref_input.grad.detach().cpu().numpy(), - atol=0, rtol=1) - np.testing.assert_allclose( - mlp.weights[0].grad.detach().cpu().numpy(), - ref_mlp[0].weight.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1) - np.testing.assert_allclose( - mlp.biases[0].grad.detach().cpu().numpy(), - ref_mlp[0].bias.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + ref_out = ref_mlp(ref_input) + ref_loss = ref_out.mean().mul(10.0) + + mlp_loss.backward() + ref_loss.backward() + if enable_autocast: + self.assertEqual(mlp_out.dtype, torch.float16) + self.assertEqual(ref_out.dtype, torch.float16) + else: + self.assertEqual(mlp_out, ref_out) + self.assertEqual(test_input.grad, ref_input.grad) + self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad) + + @common_utils.parametrize( + "use_activation,bias", + list(product(("none", "relu", "sigmoid"), (True, False))), + ) + def test_mlp(self, use_activation: str, bias: bool): + self._test_mlp_impl(use_activation, bias, enable_autocast=False) + + @common_utils.parametrize( + "use_activation,bias", + list(product(("none", "relu", "sigmoid"), (True, False))), + ) + def test_mlp_autocast_fp16(self, use_activation: str, bias: bool): + self._test_mlp_impl(use_activation, bias, enable_autocast=True) @skipFlakyTest def test_no_grad(self): @@ -141,30 +116,25 @@ def test_no_grad(self): mlp_layers = [] for i in range(mlp.num_layers): linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1]) - mlp.weights[i].data.copy_(linear.weight) - mlp.biases[i].data.copy_(linear.bias) + with torch.no_grad(): + mlp.weights[i].copy_(linear.weight) + mlp.biases[i].copy_(linear.bias) mlp_layers.append(linear) mlp_layers.append(nn.ReLU(inplace=True)) ref_mlp = nn.Sequential(*mlp_layers).cuda() - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.) + test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1.0, 1.0) ref_input = test_input.clone().detach() mlp_out = mlp(test_input) ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + self.assertEqual(mlp_out, ref_out) # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - mlp.weights[0].grad.detach().cpu().numpy(), - ref_mlp[0].weight.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + ref_out.mean().mul(10.0).backward() + self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad) + @skipFlakyTest def test_performance_half(self): mlp = MLP(mlp_sizes).cuda().half() @@ -178,10 +148,16 @@ def test_performance_half(self): ref_mlp = nn.Sequential(*mlp_layers).cuda().half() - test_input = torch.empty( - batch_size, mlp_sizes[0], device="cuda", dtype=torch.half).fill_(10.).requires_grad_() - ref_input = torch.empty( - batch_size, mlp_sizes[0], device="cuda", dtype=torch.half).fill_(10.).requires_grad_() + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda", dtype=torch.half) + .fill_(10.0) + .requires_grad_() + ) + ref_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda", dtype=torch.half) + .fill_(10.0) + .requires_grad_() + ) # Warm up GPU for _ in range(100): @@ -204,7 +180,8 @@ def test_performance_half(self): ref_loss.backward() torch.cuda.synchronize() stop_time = time() - print(F"\nPytorch MLP time {(stop_time - start_time) * 1000. / num_iters:.4f} ms") + ref_time = (stop_time - start_time) * 1000.0 / num_iters + print(f"\nPytorch MLP time {ref_time:.4f} ms") torch.cuda.synchronize() start_time = time() @@ -215,8 +192,19 @@ def test_performance_half(self): test_loss.backward() torch.cuda.synchronize() stop_time = time() - print(F"C++ MLP time {(stop_time - start_time) * 1000. / num_iters:.4f} ms") - #torch.cuda.profiler.stop() + actual_time = (stop_time - start_time) * 1000.0 / num_iters + print(f"C++ MLP time {actual_time:.4f} ms") + torch.cuda.profiler.stop() + self.assertLessEqual( + actual_time, + ref_time, + msg=f"Custom extension took {actual_time:.4f} while PyTorch took {ref_time:.4f}", + ) + + +instantiate_device_type_tests(TestMLP, globals(), only_for=("cuda",)) + + +if __name__ == "__main__": + common_utils.run_tests() -if __name__ == '__main__': - unittest.main() diff --git a/tests/L0/run_test.py b/tests/L0/run_test.py index e87a1e8e9..4b8f5c0ad 100644 --- a/tests/L0/run_test.py +++ b/tests/L0/run_test.py @@ -2,7 +2,9 @@ How to run this script? -1. Run all the tests: `python /path/to/apex/tests/L0/run_test.py` +1. Run all the tests: `python /path/to/apex/tests/L0/run_test.py` If you want an xml report, + pass `--xml-report`, i.e. `python /path/to/apex/tests/L0/run_test.py --xml-report` and + the file is created in `/path/to/apex/tests/L0`. 2. Run one of the tests (e.g. fused layer norm): `python /path/to/apex/tests/L0/run_test.py --include run_fused_layer_norm` 3. Run two or more of the tests (e.g. optimizers and fused layer norm): @@ -46,16 +48,35 @@ def parse_args(): default=DEFAULT_TEST_DIRS, help="select a set of tests to run (defaults to ALL tests).", ) + parser.add_argument( + "--xml-report", + action="store_true", + help="pass this argument to get a junit xml report. (requires `xmlrunner`)", + ) args, _ = parser.parse_known_args() return args -def main(args): - runner = unittest.TextTestRunner(verbosity=2) +def main(args: argparse.Namespace) -> None: + test_runner_kwargs = {"verbosity": 2} + Runner = unittest.TextTestRunner + if args.xml_report: + import xmlrunner + from datetime import date # NOQA + Runner = xmlrunner.XMLTestRunner + errcode = 0 for test_dir in args.include: + if args.xml_report: + this_dir = os.path.abspath(os.path.dirname(__file__)) + xml_output = os.path.join( + this_dir, + f"""TEST_{test_dir}_{date.today().strftime("%y%m%d")}""", + ) + test_runner_kwargs["output"] = xml_output + + runner = Runner(**test_runner_kwargs) test_dir = os.path.join(TEST_ROOT, test_dir) - print(test_dir) suite = unittest.TestLoader().discover(test_dir) print("\nExecuting tests from " + test_dir) diff --git a/tests/L0/run_transformer/run_bert_minimal_test.py b/tests/L0/run_transformer/run_bert_minimal_test.py deleted file mode 100644 index 639c31e77..000000000 --- a/tests/L0/run_transformer/run_bert_minimal_test.py +++ /dev/null @@ -1,260 +0,0 @@ -import random -import torch -try: - import torch_ucc -except ImportError: - HAS_TORCH_UCC = False -else: - HAS_TORCH_UCC = True - print("Use UCC as backend of Pipeline Parallel ProcessGroups") - -from apex.transformer.enums import ModelType -from apex.transformer import tensor_parallel -from apex.transformer import parallel_state -from apex.transformer.log_util import set_logging_level -from apex.transformer.tensor_parallel import vocab_parallel_cross_entropy -from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import unwrap_model -from apex.transformer.pipeline_parallel.utils import ( - average_losses_across_data_parallel_group, -) -from apex.transformer.pipeline_parallel.schedules import get_forward_backward_func -from apex.transformer.pipeline_parallel.schedules.common import build_model -from apex.transformer.pipeline_parallel.schedules.common import ( - _get_params_for_weight_decay_optimization, -) - -from apex.transformer.testing.standalone_bert import bert_model_provider -from apex.transformer.testing import global_vars -from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE -from apex.transformer.testing.commons import initialize_distributed -from apex.transformer.testing.commons import print_separator - -import warnings - - -class DebugWarning(Warning): - pass - - -set_logging_level("WARNING") -mode = None -MANUAL_SEED = 42 -inds = None -masks = None -data_idx = 0 -MASK_PROB = 0.1 -EASY_MODE = False -EASY_MODE_SIZ = 32 -ONCE = False - - -def download_fancy_data(): - # import requests - # response = requests.get('https://internet.com/book.txt') - # text = ' '.join(response.text.split()) - text = """ - An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. - """ - text = text * 1024 - encoded = text.encode("ascii", "replace") - ints = [int(encoded[i]) for i in range(len(encoded))] - return torch.tensor(ints) - - -# build a batch given sequence_len and batch size -def generate_fancy_data_labels(sequence_len, batch_size): - global data_idx - global inds - global masks - global MANUAL_SEED - temps = [] - for i in range(batch_size): - if inds is None or data_idx >= len(inds): - # hack as use of RNG will fall out of sync due to pipelines being different - torch.manual_seed(MANUAL_SEED) - inds = torch.randperm(effective_length, device="cuda") - masks = ( - torch.rand( - len(inds) // batch_size + 1, batch_size, sequence_len, device="cuda" - ) - >= MASK_PROB - ).long() - MANUAL_SEED += 1 - print("new epoch", len(inds)) - data_idx = 0 - print("my start", inds[0:5]) - print("masks_checksum:", torch.sum(masks)) - if EASY_MODE: - data_idx_ = data_idx % EASY_MODE_SIZ - else: - data_idx_ = data_idx - offset = inds[data_idx_] # * SEQUENCE_LEN - data_idx += 1 - - curr = fancy_data[offset : offset + sequence_len].clone().detach() - temps.append(curr) - temp = torch.stack(temps, dim=0).cuda() - mask = masks[data_idx // batch_size] - mask_not = torch.logical_not(mask).long() - data = mask * temp + mask_not * 124 - label = temp - if parallel_state.get_tensor_model_parallel_rank() == 0: - data_dict = {"text": data, "label": label, "mask_not": mask_not} - else: - data_dict = None - keys = ["text", "label", "mask_not"] - dtype = torch.int64 - broadcasted_data = tensor_parallel.broadcast_data(keys, data_dict, torch.long) - return ( - broadcasted_data["text"].long(), - broadcasted_data["label"].long(), - broadcasted_data["mask_not"], - ) - - -easy_data = None - - -def fwd_step_func(batch, model): - data, label, loss_mask = batch - y = model(data, torch.ones_like(data), lm_labels=label) - - def loss_func(output_tensor): - global ONCE - output_tensor, _ = output_tensor - lm_loss_ = output_tensor.float() - lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum() - averaged_loss = average_losses_across_data_parallel_group([lm_loss]) - if data_idx >= 1536: - assert averaged_loss < 4.8 - if not ONCE: - print("LOSS OK") - ONCE = True - return lm_loss, {"avg": averaged_loss} - - return y, loss_func - - -def train( - model, optim, virtual_pipeline_model_parallel_size, pipeline_model_parallel_size, async_comm -): - sequence_len = global_vars.get_args().seq_length - micro_batch_size = global_vars.get_args().micro_batch_size - hidden_size = global_vars.get_args().hidden_size - forward_backward_func = get_forward_backward_func( - virtual_pipeline_model_parallel_size, pipeline_model_parallel_size - ) - tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size) - for _ in range(16): - batch = generate_fancy_data_labels(sequence_len, batch_size) - optim.zero_grad() - forward_backward_func( - fwd_step_func, - batch, - model, - forward_only=False, - tensor_shape=tensor_shape, - async_comm=async_comm, - sequence_parallel_enabled=global_vars.get_args().sequence_parallel, - ) - # All-reduce layernorm parameters across model parallel nodes - # when sequence parallelism is used - if parallel_state.get_tensor_model_parallel_world_size() > 1 and global_vars.get_args().sequence_parallel: - for model_module in model: - unwrapped_model = unwrap_model(model_module) - for param in unwrapped_model.parameters(): - if getattr(param, 'sequence_parallel_enabled', False): - grad = param.grad - torch.distributed.all_reduce(grad, group=parallel_state.get_tensor_model_parallel_group()) - - optim.step() - - -if __name__ == "__main__": - global fancy_data - global effective_length - - global_vars.set_global_variables() - - fancy_data = download_fancy_data() - effective_length = fancy_data.size(0) // global_vars.get_args().seq_length - effective_length = fancy_data.size(0) - global_vars.get_args().seq_length - - initialize_distributed("nccl") - world_size = torch.distributed.get_world_size() - failure = None - init = True - try: - virtual_pipeline_model_parallel_sizes = (None, 2,) - if HAS_TORCH_UCC: - # Deliberately skipping test with interleaved schedule for BERT model. - # It deadlocks on hybrid UCC/NCCL backend. - virtual_pipeline_model_parallel_sizes = (None,) - for virtual_pipeline_model_parallel_size in virtual_pipeline_model_parallel_sizes: - args = global_vars.get_args() - async_comm = not args.sequence_parallel and virtual_pipeline_model_parallel_size is None - data_idx = 0 - ONCE = False - if init: - init = False - args = global_vars.get_args() - args.padded_vocab_size = 128 # needed in standalone gpt - args.model_type = ModelType.encoder_or_decoder - batch_size = args.global_batch_size - micro_batch_size = args.micro_batch_size - setup_microbatch_calculator( - args.rank, - args.rampup_batch_size, - args.global_batch_size, - args.micro_batch_size, - args.data_parallel_size, - ) - else: - parallel_state.destroy_model_parallel() - parallel_state.initialize_model_parallel( - args.tensor_model_parallel_size, - args.pipeline_model_parallel_size, - virtual_pipeline_model_parallel_size, - default_backend="nccl", - p2p_backend="ucc" if HAS_TORCH_UCC else "nccl", - ) - pipeline_model_parallel_size = ( - parallel_state.get_pipeline_model_parallel_world_size() - ) - - tensor_parallel.random.model_parallel_cuda_manual_seed(0) - model = build_model( - bert_model_provider, - wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, - virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, - cpu_offload=args.cpu_offload, - ) - assert isinstance(model, list) - assert len(model) == ( - 1 - if virtual_pipeline_model_parallel_size is None - else virtual_pipeline_model_parallel_size - ) - _param_groups = _get_params_for_weight_decay_optimization(model) - optim = torch.optim.Adam(_param_groups) - print(effective_length) - print(fancy_data.size(0)) - train( - model, - optim, - virtual_pipeline_model_parallel_size, - args.pipeline_model_parallel_size, - async_comm, - ) - except Exception as e: - failure = str(e) - finally: - parallel_state.destroy_model_parallel() - if failure is not None: - warnings.warn( - f"Minimal BERT Pipeline Parallel Failed with: {failure}", DebugWarning - ) - print(f"Minimal BERT Pipeline Parallel Failed with: {failure}") - torch.distributed.barrier() - print(TEST_SUCCESS_MESSAGE) diff --git a/tests/L0/run_transformer/run_gpt_minimal_test.py b/tests/L0/run_transformer/run_gpt_minimal_test.py deleted file mode 100644 index b3674c898..000000000 --- a/tests/L0/run_transformer/run_gpt_minimal_test.py +++ /dev/null @@ -1,223 +0,0 @@ -from functools import partial -from typing import List -import time - -import torch -try: - import torch_ucc -except ImportError: - HAS_TORCH_UCC = False -else: - HAS_TORCH_UCC = True - print("Use UCC as backend of Pipeline Parallel ProcessGroups") - -from apex.transformer import parallel_state -from apex.transformer.enums import ModelType -from apex.transformer.tensor_parallel import model_parallel_cuda_manual_seed -from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import unwrap_model -from apex.transformer.pipeline_parallel.utils import ( - average_losses_across_data_parallel_group, -) -from apex.transformer.pipeline_parallel.utils import get_ltor_masks_and_position_ids -from apex.transformer.pipeline_parallel.schedules.common import build_model -from apex.transformer.pipeline_parallel.schedules.common import ( - _get_params_for_weight_decay_optimization, -) -from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import ( - forward_backward_pipelining_without_interleaving, -) -from apex.transformer.testing.standalone_gpt import gpt_model_provider -from apex.transformer.testing import global_vars -from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE -from apex.transformer.testing.commons import initialize_distributed - -MANUAL_SEED = 42 -inds = None -data_idx = 0 -N_VOCAB = 128 - - -def download_fancy_data(): - # import requests - # response = requests.get('https://internet.com/book.txt') - # text = ' '.join(response.text.split()) - text = """ - An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. - """ - text = text * 1024 - encoded = text.encode("ascii", "replace") - ints = [int(encoded[i]) for i in range(len(encoded))] - return torch.tensor(ints) - - -# build a batch given sequence_len and batch size -def generate_fancy_data_labels(sequence_len, batch_size): - global data_idx - global inds - global MANUAL_SEED - temps = list() - for i in range(batch_size): - if inds is None or data_idx >= len(inds): - # hack as use of RNG will fall out of sync due to pipelines being different - model_parallel_cuda_manual_seed(MANUAL_SEED) - inds = torch.randperm(effective_length, device="cuda") - MANUAL_SEED += 1 - data_idx = 0 - data_idx_ = data_idx - offset = inds[data_idx_] - data_idx += 1 - curr = fancy_data[offset : offset + sequence_len + 1].clone().detach() - temps.append(curr) - temp = torch.stack(temps, dim=0).cuda() - return temp - - -easy_data = None - - -def get_batch(int_tensors: List[torch.Tensor]): - data = int_tensors[0] - # Unpack. - tokens_ = data.long() - labels = tokens_[:, 1:].contiguous() - tokens = tokens_[:, :-1].contiguous() - # Get the masks and position ids. - attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( - tokens, - N_VOCAB, # tokenizer.eod, - False, # args.reset_position_ids, - False, # args.reset_attention_mask, - False, # args.eod_mask_loss, - ) - return tokens, labels, loss_mask, attention_mask, position_ids - - -# Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L75 -def loss_func(loss_mask, output_tensor): - losses = output_tensor.float() - loss_mask = loss_mask.view(-1).float() - loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() - - # Reduce loss for logging. - averaged_loss = average_losses_across_data_parallel_group([loss]) - - return loss, {"lm loss": averaged_loss[0]} - - -# Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L86 -def fwd_step_func(batch, model): - """Forward step.""" - tokens, labels, loss_mask, attention_mask, position_ids = get_batch(batch) - output_tensor = model(tokens, position_ids, attention_mask, labels=labels) - return output_tensor, partial(loss_func, loss_mask) - - -def train(model, optim, pipeline_model_parallel_size, async_comm): - sequence_len = global_vars.get_args().seq_length - micro_batch_size = global_vars.get_args().micro_batch_size - hidden_size = global_vars.get_args().hidden_size - fwd_bwd_func = forward_backward_pipelining_without_interleaving - - tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size) - runtime = 0 - # training loop - for i in range(3): - since = time.time() - if torch.distributed.get_rank() == 0: - print("begin iter", i) - batch = [ - generate_fancy_data_labels(args.seq_length, args.global_batch_size) - for _ in range(pipeline_model_parallel_size) - ] - if torch.distributed.get_rank() == 0: - print("finished making batch...") - optim.zero_grad() - fwd_bwd_func( - fwd_step_func, - batch, - model, - forward_only=False, - tensor_shape=tensor_shape, - async_comm=async_comm, - sequence_parallel_enabled=args.sequence_parallel, - ) - if torch.distributed.get_rank() == 0: - print("finished forward step") - # All-reduce layernorm parameters across model parallel nodes - # when sequence parallelism is used - if parallel_state.get_tensor_model_parallel_world_size() > 1 and global_vars.get_args().sequence_parallel: - for model_module in model: - unwrapped_model = unwrap_model(model_module) - for param in unwrapped_model.parameters(): - if getattr(param, 'sequence_parallel_enabled', False): - grad = param.grad - torch.distributed.all_reduce(grad, group=parallel_state.get_tensor_model_parallel_group()) - optim.step() - if torch.distributed.get_rank() == 0: - print("finished iter", i) - runtime += time.time() - since - return runtime / 3.0 - - -if __name__ == "__main__": - init = True - global_vars.set_global_variables() - for async_comm in (False,) if global_vars.get_args().sequence_parallel else (False, True): - global fancy_data - global effective_length - - if init: - init = False - - fancy_data = download_fancy_data() - args = global_vars.get_args() - args.model_type = ModelType.encoder_or_decoder - effective_length = fancy_data.size(0) // args.seq_length - effective_length = fancy_data.size(0) - args.seq_length - - initialize_distributed("nccl") - world_size = torch.distributed.get_world_size() - - failure = None - args.padded_vocab_size = 128 - batch_size = args.global_batch_size - micro_batch_size = args.micro_batch_size - setup_microbatch_calculator( - args.rank, - args.rampup_batch_size, - args.global_batch_size, - args.micro_batch_size, - args.data_parallel_size, # args.data_parallel_size, - ) - world_size = torch.distributed.get_world_size() - - print(args.tensor_model_parallel_size, "MODEL PARALLEL SIZE") - - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=args.tensor_model_parallel_size, - pipeline_model_parallel_size_=args.pipeline_model_parallel_size, - default_backend="nccl", - p2p_backend="ucc" if HAS_TORCH_UCC else "nccl", - ) - - pipeline_model_parallel_size = ( - parallel_state.get_pipeline_model_parallel_world_size() - ) - model_parallel_cuda_manual_seed(0) - model = build_model( - gpt_model_provider, - wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, - virtual_pipeline_model_parallel_size=None, - cpu_offload=args.cpu_offload, - ) - assert isinstance(model, list), model - _param_groups = _get_params_for_weight_decay_optimization(model) - optim = torch.optim.Adam(_param_groups) - runtime = train(model, optim, args.pipeline_model_parallel_size, async_comm) - - parallel_state.destroy_model_parallel() - torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print(TEST_SUCCESS_MESSAGE) - print("Average Iteration Time:", runtime) diff --git a/tests/L0/run_transformer/test_batch_sampler.py b/tests/L0/run_transformer/test_batch_sampler.py index 52175d53a..99b0f1aeb 100644 --- a/tests/L0/run_transformer/test_batch_sampler.py +++ b/tests/L0/run_transformer/test_batch_sampler.py @@ -1,10 +1,6 @@ -from itertools import product - import torch from torch.testing._internal import common_utils from torch.utils.data import Dataset -from torch.utils.data import RandomSampler -from torch.utils.data import BatchSampler from torch.utils.data import DataLoader from apex.transformer.pipeline_parallel.utils import _split_batch_into_microbatch as split_batch_into_microbatch @@ -81,27 +77,30 @@ def __iter__(self): # Samples 8 tensors in total. # First sample 4 tensors twice, then sample 2 tensors fourth. class TestBatchSamplerBehavior(common_utils.TestCase): + def tearDown(self) -> None: + torch.cuda.empty_cache() + super().tearDown() + def test_batch_sampler_behavior(self): dataset = MyIterableDataset(0, 100) for num_workers in (1, 2, 4): - with self.subTest(f"{num_workers}"): - torch.manual_seed(42) - loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 4, 0, 1), num_workers=num_workers) - samples = [] - for i, batch in enumerate(loader): - samples.append(batch) - if i == 2 - 1: - break - - torch.manual_seed(42) - loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 2, 0, 1), num_workers=num_workers) - samples2 = [] - for i, batch in enumerate(loader): - samples2.append(batch) - if i == 4 - 1: - break - self.assertEqual(torch.cat(samples), torch.cat(samples2)) + torch.manual_seed(42) + loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 4, 0, 1), num_workers=num_workers) + samples = [] + for i, batch in enumerate(loader): + samples.append(batch) + if i == 2 - 1: + break + + torch.manual_seed(42) + loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 2, 0, 1), num_workers=num_workers) + samples2 = [] + for i, batch in enumerate(loader): + samples2.append(batch) + if i == 4 - 1: + break + self.assertEqual(torch.cat(samples), torch.cat(samples2), msg=f"num_workers={num_workers}") def test_split_batch(self): diff --git a/tests/L0/run_transformer/test_bert_minimal.py b/tests/L0/run_transformer/test_bert_minimal.py new file mode 100644 index 000000000..887f7e259 --- /dev/null +++ b/tests/L0/run_transformer/test_bert_minimal.py @@ -0,0 +1,250 @@ +import torch +import unittest +from apex.transformer.testing import global_vars +from apex.transformer.testing.standalone_bert import bert_model_provider +from apex.transformer.pipeline_parallel.schedules.common import ( + _get_params_for_weight_decay_optimization, build_model +) +from apex.transformer.pipeline_parallel.schedules import get_forward_backward_func +from apex.transformer.pipeline_parallel.utils import ( + average_losses_across_data_parallel_group, unwrap_model, setup_microbatch_calculator +) +from apex.transformer.log_util import set_logging_level +from apex.transformer import tensor_parallel, parallel_state +from apex.transformer.enums import ModelType +from apex.transformer._ucc_util import HAS_UCC +from apex.transformer.testing.distributed_test_base import UccDistributedTestBase, NcclDistributedTestBase +import logging + +from torch.testing._internal import common_utils + +logging.getLogger("torch").setLevel(logging.WARNING) + + +logging.getLogger("apex").setLevel(logging.WARNING) + + +set_logging_level("WARNING") + + +class BertTestBase: + + def _download_fancy_data(self): + text = """ + An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. + """ + text = text * 1024 + encoded = text.encode("ascii", "replace") + ints = [int(encoded[i]) for i in range(len(encoded))] + return torch.tensor(ints) + + # build a batch given sequence_len and batch size + def _generate_fancy_data_labels(self, sequence_len, batch_size): + temps = [] + for i in range(batch_size): + if self.inds is None or self.data_idx >= len(self.inds): + # hack as use of RNG will fall out of sync due to pipelines being different + torch.manual_seed(self.MANUAL_SEED) + self.inds = torch.randperm( + self.effective_length, device="cuda") + self.masks = ( + torch.rand( + len(self.inds) // batch_size + 1, batch_size, sequence_len, device="cuda" + ) + >= self.MASK_PROB + ).long() + self.MANUAL_SEED += 1 + self.data_idx = 0 + if self.rank == 0: + print("new epoch", len(self.inds)) + print("my start", self.inds[0:5]) + print("masks_checksum:", torch.sum(self.masks)) + if self.EASY_MODE: + data_idx_ = self.data_idx % self.EASY_MODE_SIZ + else: + data_idx_ = self.data_idx + offset = self.inds[data_idx_] # * SEQUENCE_LEN + self.data_idx += 1 + + curr = self.fancy_data[offset: offset + + sequence_len].clone().detach() + temps.append(curr) + temp = torch.stack(temps, dim=0).cuda() + mask = self.masks[self.data_idx // batch_size] + mask_not = torch.logical_not(mask).long() + data = mask * temp + mask_not * 124 + label = temp + if parallel_state.get_tensor_model_parallel_rank() == 0: + data_dict = {"text": data, "label": label, "mask_not": mask_not} + else: + data_dict = None + keys = ["text", "label", "mask_not"] + broadcasted_data = tensor_parallel.broadcast_data( + keys, data_dict, torch.long) + return ( + broadcasted_data["text"].long(), + broadcasted_data["label"].long(), + broadcasted_data["mask_not"], + ) + + def _fwd_step_func(self, batch, model): + data, label, loss_mask = batch + y = model(data, torch.ones_like(data), lm_labels=label) + + def loss_func(output_tensor): + output_tensor, _ = output_tensor + lm_loss_ = output_tensor.float() + lm_loss = torch.sum(lm_loss_.view(-1) * + loss_mask.reshape(-1)) / loss_mask.sum() + averaged_loss = average_losses_across_data_parallel_group([ + lm_loss]) + if self.data_idx >= 1536: + # NOTE (patwang): Loss cutoff might be excessively high but roughly one in five + # unlucky random seeds do cause loss to spike to just under 8.0 + self.assertLess(averaged_loss, 8.0) + return lm_loss, {"avg": averaged_loss} + + return y, loss_func + + def _train( + self, model, optim, virtual_pipeline_model_parallel_size, pipeline_model_parallel_size, async_comm + ): + args = global_vars.get_args() + sequence_len = args.seq_length + micro_batch_size = args.micro_batch_size + hidden_size = args.hidden_size + global_batch_size = args.global_batch_size + forward_backward_func = get_forward_backward_func( + virtual_pipeline_model_parallel_size, pipeline_model_parallel_size + ) + tensor_shape = (sequence_len, micro_batch_size, hidden_size) + for _ in range(16): + batch = self._generate_fancy_data_labels( + sequence_len, global_batch_size) + optim.zero_grad() + forward_backward_func( + self._fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=tensor_shape, + async_comm=async_comm, + sequence_parallel_enabled=args.sequence_parallel, + ) + # All-reduce layernorm parameters across model parallel nodes + # when sequence parallelism is used + if parallel_state.get_tensor_model_parallel_world_size() > 1 and args.sequence_parallel: + for model_module in model: + unwrapped_model = unwrap_model(model_module) + for param in unwrapped_model.parameters(): + if getattr(param, 'sequence_parallel_enabled', False): + grad = param.grad + torch.distributed.all_reduce( + grad, group=parallel_state.get_tensor_model_parallel_group()) + + optim.step() + + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_bert_without_interleaving(self): + self._test_bert(virtual_pipeline_model_parallel_size=None) + + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_bert_with_interleaving(self): + if self.DISTRIBUTED_BACKEND == 'ucc': + self.skipTest('skip interleaving with ucc') + self._test_bert(virtual_pipeline_model_parallel_size=2) + + def _test_bert(self, virtual_pipeline_model_parallel_size): + + self.MANUAL_SEED = 42 + self.inds = None + self.masks = None + self.data_idx = 0 + self.MASK_PROB = 0.1 + self.EASY_MODE = False + self.EASY_MODE_SIZ = 32 + + num_devices = torch.cuda.device_count() + tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices > 4 else 1 + pipeline_model_parallel_size = num_devices // tensor_model_parallel_size + + override_args = { + "micro_batch_size": 2, + "num_layers": 16, + "hidden_size": 256, + "num_attention_heads": 8, + "max_position_embeddings": 512, + "seq_length": 512, + "global_batch_size": 128, + "pipeline_model_parallel_size": pipeline_model_parallel_size, + "tensor_model_parallel_size": tensor_model_parallel_size, + "bert_binary_head": False, + "world_size": self.world_size, + "rank": self.rank, + } + + global_vars.set_global_variables(override_args=override_args, ignore_unknown_args=True) + args = global_vars.get_args() + + self.fancy_data = self._download_fancy_data() + self.effective_length = self.fancy_data.size(0) // args.seq_length + self.effective_length = self.fancy_data.size(0) - args.seq_length + + if self.rank == 0: + print( + f'testing backend: {self.DISTRIBUTED_BACKEND} with virtual_pipeline_model_parallel_size: {virtual_pipeline_model_parallel_size}') + async_comm = not args.sequence_parallel and virtual_pipeline_model_parallel_size is None + self.data_idx = 0 + args.padded_vocab_size = 128 # needed in standalone gpt + args.model_type = ModelType.encoder_or_decoder + setup_microbatch_calculator( + args.rank, + args.rampup_batch_size, + args.global_batch_size, + args.micro_batch_size, + args.data_parallel_size, + ) + parallel_state.initialize_model_parallel( + args.tensor_model_parallel_size, + args.pipeline_model_parallel_size, + virtual_pipeline_model_parallel_size, + default_backend="nccl", + p2p_backend=self.DISTRIBUTED_BACKEND, + ) + + tensor_parallel.random.model_parallel_cuda_manual_seed(0) + model = build_model( + bert_model_provider, + wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, + virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, + cpu_offload=args.cpu_offload, + ) + assert isinstance(model, list) + assert len(model) == ( + 1 + if virtual_pipeline_model_parallel_size is None + else virtual_pipeline_model_parallel_size + ) + _param_groups = _get_params_for_weight_decay_optimization(model) + optim = torch.optim.Adam(_param_groups) + self._train( + model, + optim, + virtual_pipeline_model_parallel_size, + args.pipeline_model_parallel_size, + async_comm, + ) + torch.distributed.barrier() + + +class NcclBertTest(BertTestBase, NcclDistributedTestBase): + pass + + +@unittest.skipUnless(HAS_UCC, "requires pytorch to be built with native ucc") +class UccBertTest(BertTestBase, UccDistributedTestBase): + pass + + +if __name__ == "__main__": + common_utils.run_tests() diff --git a/tests/L0/run_transformer/test_cross_entropy.py b/tests/L0/run_transformer/test_cross_entropy.py index 1f5162876..d1e1f9b1e 100644 --- a/tests/L0/run_transformer/test_cross_entropy.py +++ b/tests/L0/run_transformer/test_cross_entropy.py @@ -18,7 +18,7 @@ def torch_cross_entropy( - batch_size: int, seq_length: int, vocab_size: int, logits_scale: float, seed: int, + batch_size: int, seq_length: int, vocab_size: int, logits_scale: float, seed: int, label_smoothing: float = 0.0 ) -> Tuple[torch.Tensor, torch.Tensor]: set_random_seed(seed) identity = IdentityLayer( @@ -28,7 +28,7 @@ def torch_cross_entropy( target = torch.cuda.LongTensor(size=(batch_size, seq_length)).random_(0, vocab_size) loss = ( F.cross_entropy( - logits.view(-1, logits.size()[-1]), target.view(-1), reduction="none" + logits.view(-1, logits.size()[-1]), target.view(-1), reduction="none", label_smoothing=label_smoothing ) .view_as(target) .mean() @@ -38,7 +38,7 @@ def torch_cross_entropy( def tensor_sharded_cross_entropy( - batch_size, seq_length, vocab_size, logits_scale, seed + batch_size, seq_length, vocab_size, logits_scale, seed, label_smoothing=0.0 ): set_random_seed(seed) identity = IdentityLayer( @@ -48,7 +48,7 @@ def tensor_sharded_cross_entropy( logits_parallel = tensor_parallel.scatter_to_tensor_model_parallel_region(logits) target = torch.cuda.LongTensor(size=(batch_size, seq_length)).random_(0, vocab_size) logits_parallel_ = logits_parallel.clone().detach() - loss = cross_entropy.vocab_parallel_cross_entropy(logits_parallel, target).mean() + loss = cross_entropy.vocab_parallel_cross_entropy(logits_parallel, target, label_smoothing=label_smoothing).mean() loss.backward() # check for mutation assert torch.equal(logits_parallel_, logits_parallel) @@ -63,27 +63,30 @@ def test_cross_entropy(self): for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - vocab_size = vocab_size_per_partition * tensor_model_parallel_world_size - loss_torch, grad_torch = torch_cross_entropy( - batch_size, sequence_length, vocab_size, logits_scale, seed - ) - ( - loss_tensor_parallel, - grad_tensor_parallel, - ) = tensor_sharded_cross_entropy( - batch_size, sequence_length, vocab_size, logits_scale, seed - ) - - self.assertEqual(loss_torch, loss_tensor_parallel) - self.assertEqual(grad_torch, grad_tensor_parallel) - - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + vocab_size = vocab_size_per_partition * tensor_model_parallel_world_size + loss_torch, grad_torch = torch_cross_entropy( + batch_size, sequence_length, vocab_size, logits_scale, seed + ) + ( + loss_tensor_parallel, + grad_tensor_parallel, + ) = tensor_sharded_cross_entropy( + batch_size, sequence_length, vocab_size, logits_scale, seed + ) + + self.assertEqual( + loss_torch, loss_tensor_parallel, + msg=f"tensor_model_parallel_size: {tensor_model_parallel_world_size}", + ) + self.assertEqual( + grad_torch, grad_tensor_parallel, + msg=f"tensor_model_parallel_size: {tensor_model_parallel_world_size}", + ) + + parallel_state.destroy_model_parallel() class NcclVocabParallelCrossEntropyTest(VocabParallelCrossEntropyTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/run_dynamic_batchsize_test.py b/tests/L0/run_transformer/test_dynamic_batchsize.py similarity index 54% rename from tests/L0/run_transformer/run_dynamic_batchsize_test.py rename to tests/L0/run_transformer/test_dynamic_batchsize.py index b2c020a41..327aa4d83 100644 --- a/tests/L0/run_transformer/run_dynamic_batchsize_test.py +++ b/tests/L0/run_transformer/test_dynamic_batchsize.py @@ -1,45 +1,38 @@ from typing import Tuple, List import torch +import unittest from apex.transformer import parallel_state from apex.transformer.pipeline_parallel.utils import get_num_microbatches from apex.transformer.pipeline_parallel.schedules.common import ( - _get_params_for_weight_decay_optimization, + _get_params_for_weight_decay_optimization, build_model ) -from apex.transformer.pipeline_parallel.schedules.common import build_model from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_with_interleaving import ( _forward_backward_pipelining_with_interleaving, ) -from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import _reconfigure_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import update_num_microbatches +from apex.transformer.pipeline_parallel.utils import ( + setup_microbatch_calculator, _reconfigure_microbatch_calculator, update_num_microbatches +) from apex.transformer.testing import global_vars -from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE -from apex.transformer.testing.commons import initialize_distributed -from apex.transformer.testing.commons import print_separator -from apex.transformer.testing.commons import fwd_step_func -from apex.transformer.log_util import get_transformer_logger, set_logging_level -from apex.transformer.testing.commons import model_provider_func -from apex.transformer._data import MegatronPretrainingRandomSampler -from apex.transformer._data import MegatronPretrainingSampler +from apex.transformer.testing.commons import ( + print_separator, fwd_step_func, model_provider_func +) +from apex.transformer.log_util import get_transformer_logger +from apex.transformer._data import MegatronPretrainingRandomSampler, MegatronPretrainingSampler +from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase +from torch.testing._internal import common_utils # note(mkozuki): To see warmup, steady, cooldown iterations, uncomment the line below # set_logging_level("INFO") _logger = get_transformer_logger("pipeline_parallel_test") # note(mkozuki): To see if local batch size increases, uncomment the line below # _logger.setLevel("INFO") -global_vars.set_global_variables( - args_defaults={"global_batch_size": 512, "rampup_batch_size": [64, 64, 1000],}, - ignore_unknown_args=True, -) -RAMPUP_BATCH_SIZE = [] NUM_ITERATIONS = 20 NUM_SAMPLES = 16384 // 2 -batch_size, micro_batch_size = None, None HIDDEN_SIZE = 16 @@ -88,9 +81,10 @@ def run_interleaved_with_dynamic_batch_size( ) assert isinstance(model, list) assert len(model) == virtual_pipeline_model_parallel_size - optimizer = torch.optim.Adam(_get_params_for_weight_decay_optimization(model)) + optimizer = torch.optim.Adam( + _get_params_for_weight_decay_optimization(model)) - initial_local_minibatch_size = get_num_microbatches() * micro_batch_size + initial_local_minibatch_size = get_num_microbatches() * args.micro_batch_size dataset = Dataset(NUM_SAMPLES) data_loader = torch.utils.data.DataLoader( dataset, @@ -110,11 +104,11 @@ def get_num_samples(batch): assert isinstance(batch, (list, tuple)) return [get_num_samples(b) for b in batch] - tensor_shape = [micro_batch_size, HIDDEN_SIZE, HIDDEN_SIZE] + tensor_shape = [args.micro_batch_size, HIDDEN_SIZE, HIDDEN_SIZE] consumed_samples = 0 for i in range(NUM_ITERATIONS): update_num_microbatches(consumed_samples, consistency_check=False) - local_batch_size = get_num_microbatches() * micro_batch_size + local_batch_size = get_num_microbatches() * args.micro_batch_size data_iter._index_sampler.local_minibatch_size = local_batch_size local_mini_batch = next(data_iter) @@ -134,7 +128,7 @@ def get_num_samples(batch): consumed_samples += ( parallel_state.get_data_parallel_world_size() * get_num_microbatches() - * micro_batch_size + * args.micro_batch_size ) if not forward_only: @@ -146,57 +140,84 @@ def get_num_samples(batch): optimizer.zero_grad(set_to_none=True) torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print(TEST_SUCCESS_MESSAGE) + + +class DynamicBatchsizeTestBase: + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_dynamic_batchsize(self): + + n_tests = 0 + failures = [] + + override_args = { + "micro_batch_size": 2, + "num_layers": 16, + "hidden_size": 256, + "num_attention_heads": 8, + "max_position_embeddings": 512, + "seq_length": 512, + "global_batch_size": 128, + "use_cpu_initialization": True, + "world_size": self.world_size, + "rank": self.rank, + } + + global_vars.set_global_variables( + args_defaults={"global_batch_size": 512, + "rampup_batch_size": [64, 64, 1000], }, + ignore_unknown_args=True, + override_args=override_args, + ) + + args = global_vars.get_args() + + setup_microbatch_calculator( + args.rank, + args.rampup_batch_size, + args.global_batch_size, + args.micro_batch_size, + 1, # args.data_parallel_size, + ) + for BatchSamplerCls in ( + MegatronPretrainingSampler, + MegatronPretrainingRandomSampler, + ): + for forward_only in (False, True): + n_tests += 1 + pipeline_model_parallel_size = self.world_size + try: + run_interleaved_with_dynamic_batch_size( + pipeline_model_parallel_size, forward_only, BatchSamplerCls, + ) + except Exception as e: + msg = ( + f"\tforward_only: {forward_only}\n" + f"pipeline rank: {parallel_state.get_pipeline_model_parallel_rank()}, " + f"virtual pipeline rank: {parallel_state.get_virtual_pipeline_model_parallel_rank()}\n" + f"{str(e)}" + ) + raise RuntimeError(msg) + finally: + parallel_state.destroy_model_parallel() + print_separator("TEST RESULT") + if failures: + torch.distributed.barrier() + if torch.distributed.get_rank() == 0: + print("\n".join(failures)) + msg = f"{len(failures)} / {n_tests} cases failed" + raise RuntimeError(msg) + else: + torch.distributed.barrier() + if torch.distributed.get_rank() == 0: + print("### PASS!") + + +class NcclDynamicBatchsizeTest(DynamicBatchsizeTestBase, NcclDistributedTestBase): + pass + +# TODO: (Fuzzkatt) UCC still doesn't work with fwd_bwd_pipelining_with_interleaving if __name__ == "__main__": torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - n_tests = 0 - failures = [] - - initialize_distributed() - world_size = torch.distributed.get_world_size() - args = global_vars.get_args() - batch_size = args.global_batch_size - micro_batch_size = args.micro_batch_size - setup_microbatch_calculator( - args.rank, - args.rampup_batch_size, - args.global_batch_size, - args.micro_batch_size, - 1, # args.data_parallel_size, - ) - for BatchSamplerCls in ( - MegatronPretrainingSampler, - MegatronPretrainingRandomSampler, - ): - for forward_only in (False, True): - n_tests += 1 - pipeline_model_parallel_size = world_size - try: - run_interleaved_with_dynamic_batch_size( - pipeline_model_parallel_size, forward_only, BatchSamplerCls, - ) - except Exception as e: - msg = ( - f"\tforward_only: {forward_only}\n" - f"pipeline rank: {parallel_state.get_pipeline_model_parallel_rank()}, " - f"virtual pipeline rank: {parallel_state.get_virtual_pipeline_model_parallel_rank()}\n" - f"{str(e)}" - ) - raise RuntimeError(msg) - finally: - parallel_state.destroy_model_parallel() - print_separator("TEST RESULT") - if failures: - torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print("\n".join(failures)) - msg = f"{len(failures)} / {n_tests} cases failed" - raise RuntimeError(msg) - else: - torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print("### PASS!") + common_utils.run_tests() diff --git a/tests/L0/run_transformer/test_fused_softmax.py b/tests/L0/run_transformer/test_fused_softmax.py index 278df69f7..c49967073 100644 --- a/tests/L0/run_transformer/test_fused_softmax.py +++ b/tests/L0/run_transformer/test_fused_softmax.py @@ -14,6 +14,14 @@ def attention_mask_func(attention_scores, attention_mask): return attention_scores.masked_fill(attention_mask, -10000.0) +def forward_torch_softmax(input, mask, scale): + input = input * scale + mask_output = attention_mask_func(input, mask) if mask is not None else input + probs = torch.nn.Softmax(dim=-1)(mask_output) + all_k_masked = mask.all(axis=-1) + zero_attention_mask = (1.0 - all_k_masked.float())[:, :, :, None] + probs = probs * zero_attention_mask + return probs autocast_dtypes = ( (torch.half, torch.bfloat16) if torch.cuda.is_bf16_supported() else (torch.half,) @@ -49,6 +57,10 @@ def _setup_fused_softmax( ) return fused_fn, torch_fn + def tearDown(self) -> None: + torch.cuda.empty_cache() + super().tearDown() + def test_fused_scale_mask_softmax(self): """ attention_scores.shape = [4, 12, 24, 24] @@ -57,75 +69,153 @@ def test_fused_scale_mask_softmax(self): for (dtype, scale, softmax_in_fp32, shape) in itertools.product( (torch.half, torch.bfloat16), (None, 2.0), (False, True), ((4, 12, 24, 24), (32, 12, 4, 214)) ): - with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - if not (scale is None or softmax_in_fp32): - with self.assertRaises(RuntimeError): - self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.padding, - ) - return - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.padding, - ) + msg = f"{dtype}-{scale}-{softmax_in_fp32}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + if not (scale is None or softmax_in_fp32): + with self.assertRaises(RuntimeError, msg=msg): + self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, + ) + return + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, + ) - attention_scores_0 = ( - torch.randn(shape) - .to(device="cuda", dtype=dtype) - .requires_grad_(True) - ) - with torch.no_grad(): - attention_scores_1 = attention_scores_0.clone().requires_grad_(True) - mask_shape = (shape[0],) + (1,) + shape[2:] - mask = torch.randint(0, 2, mask_shape, device="cuda").bool() - expected = fused_fn(attention_scores_0, mask) - actual = torch_fn(attention_scores_1, mask) - self.assertEqual(actual, expected) - - g0 = torch.rand_like(actual) - with torch.no_grad(): - g1 = g0.clone() - expected.backward(g0) - actual.backward(g1) + attention_scores_0 = ( + torch.randn(shape) + .to(device="cuda", dtype=dtype) + .requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = attention_scores_0.clone().requires_grad_(True) + mask_shape = (shape[0],) + (1,) + shape[2:] + mask = torch.randint(0, 2, mask_shape, device="cuda").bool() + expected = fused_fn(attention_scores_0, mask) + actual = torch_fn(attention_scores_1, mask) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) def test_autocast_fused_scale_mask_softmax(self): for dtype in autocast_dtypes: - with self.subTest(f"{dtype}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding - ) + msg = f"dtype: {dtype}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding + ) - attention_scores_0 = ( - torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + attention_scores_0 = ( + torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = ( + attention_scores_0.clone().to(dtype).requires_grad_(True) ) - with torch.no_grad(): - attention_scores_1 = ( - attention_scores_0.clone().to(dtype).requires_grad_(True) + mask = torch.randint(0, 2, (4, 1, 24, 24)).bool().cuda() + + expected = torch_fn(attention_scores_1, mask) + with torch.cuda.amp.autocast(dtype=dtype): + actual = fused_fn(attention_scores_0, mask) + self.assertEqual(actual.dtype, dtype, msg=msg) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) + + def test_fused_scale_softmax(self): + """ + attention_scores.shape = [4, 12, 24, 24] + mask = None + """ + for (dtype, scale, softmax_in_fp32, shape) in itertools.product( + (torch.half, torch.bfloat16), (None, 2.0), (False, True), ((4, 12, 24, 24), (32, 12, 4, 214)) + ): + msg = f"{dtype}-{scale}-{softmax_in_fp32}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + if not (scale is None or softmax_in_fp32): + with self.assertRaises(RuntimeError, msg=msg): + self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, ) - mask = torch.randint(0, 2, (4, 1, 24, 24)).bool().cuda() + return + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, + ) + + attention_scores_0 = ( + torch.randn(shape) + .to(device="cuda", dtype=dtype) + .requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = attention_scores_0.clone().requires_grad_(True) + mask = None - expected = torch_fn(attention_scores_1, mask) - with torch.cuda.amp.autocast(dtype=dtype): - actual = fused_fn(attention_scores_0, mask) - self.assertEqual(actual.dtype, dtype) - self.assertEqual(actual, expected) + expected = fused_fn(attention_scores_0, mask) + actual = torch_fn(attention_scores_1, mask) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) + + def test_autocast_fused_scale_softmax(self): + for dtype in autocast_dtypes: + msg = f"dtype: {dtype}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding + ) + + attention_scores_0 = ( + torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = ( + attention_scores_0.clone().to(dtype).requires_grad_(True) + ) + mask = None - g0 = torch.rand_like(actual) - with torch.no_grad(): - g1 = g0.clone() - expected.backward(g0) - actual.backward(g1) + expected = torch_fn(attention_scores_1, mask) + with torch.cuda.amp.autocast(dtype=dtype): + actual = fused_fn(attention_scores_0, mask) + self.assertEqual(actual.dtype, dtype, msg=msg) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) def test_fused_upper_triangle_mask_softmax(self): """ @@ -138,83 +228,144 @@ def test_fused_upper_triangle_mask_softmax(self): for (dtype, scale, softmax_in_fp32) in itertools.product( (torch.half, torch.bfloat16), (None, 2.0), (False, True), ): - with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - if not (scale is None or softmax_in_fp32): - with self.assertRaises(RuntimeError): - self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.causal, - ) - return - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.causal, - ) + msg = f"{dtype}-{scale}-{softmax_in_fp32}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + if not (scale is None or softmax_in_fp32): + with self.assertRaises(RuntimeError, msg=msg): + self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.causal, + ) + return + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.causal, + ) - attn_weights_0 = ( - torch.randn((4, 12, 24, 24)) - .to(device="cuda", dtype=dtype) - .requires_grad_(True) - ) - with torch.no_grad(): - attn_weights_1 = attn_weights_0.clone().requires_grad_(True) - total_mask = ( - ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) - .unsqueeze(0) - .unsqueeze(0) - ) - total_mask = total_mask.repeat((4, 1, 1, 1)) - expected = fused_fn(attn_weights_0, total_mask) - actual = torch_fn(attn_weights_1, total_mask) - self.assertEqual(actual, expected) + attn_weights_0 = ( + torch.randn((4, 12, 24, 24)) + .to(device="cuda", dtype=dtype) + .requires_grad_(True) + ) + with torch.no_grad(): + attn_weights_1 = attn_weights_0.clone().requires_grad_(True) + total_mask = ( + ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) + .unsqueeze(0) + .unsqueeze(0) + ) + total_mask = total_mask.repeat((4, 1, 1, 1)) + expected = fused_fn(attn_weights_0, total_mask) + actual = torch_fn(attn_weights_1, total_mask) + self.assertEqual(actual, expected, msg=msg) - g0 = torch.randn_like(actual) - with torch.no_grad(): - g1 = g0.clone() - actual.backward(g0) - expected.backward(g1) + g0 = torch.randn_like(actual) + with torch.no_grad(): + g1 = g0.clone() + actual.backward(g0) + expected.backward(g1) def test_autocast_fused_upper_triangle_mask_softmax(self): for dtype in autocast_dtypes: - with self.subTest(f"{dtype}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.causal - ) + msg = f"dtype: {dtype}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.causal + ) - attn_weights_0 = ( - torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) - ) - with torch.no_grad(): - attn_weights_1 = ( - attn_weights_0.clone().to(dtype).requires_grad_(True) - ) - total_mask = ( - ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) - .unsqueeze(0) - .unsqueeze(0) + attn_weights_0 = ( + torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + ) + with torch.no_grad(): + attn_weights_1 = ( + attn_weights_0.clone().to(dtype).requires_grad_(True) ) + total_mask = ( + ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) + .unsqueeze(0) + .unsqueeze(0) + ) + + with torch.cuda.amp.autocast(dtype=dtype): + actual = fused_fn(attn_weights_0, total_mask) + self.assertEqual(actual.dtype, dtype, msg=msg) + expected = torch_fn(attn_weights_1, total_mask) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.randn_like(actual) + with torch.no_grad(): + g1 = g0.clone() + actual.backward(g0) + expected.backward(g1) + + +class TestGenericFusedSoftmaxKernel(common_utils.TestCase): + + def setUp(self): + super().setUp() + self.batch = 2 + self.attn = 16 + self.scale_t = 1.0 + self.dtype = torch.float16 + self.device = torch.cuda.current_device() + self.thresh = {"atol": 1e-3, "rtol": 1e-3} + + qlen = [1, 2] + klen = [1, 2, 3, 4, 5, 8, 10, 11, 13, 128, 256, 1200, 1234] + available_cuda_mem = torch.cuda.memory.mem_get_info(self.device)[0] / (1024 ** 3) + if available_cuda_mem > 40: + qlen.extend([1234, 2322, 2348]) + klen.extend([2048, 3123, 4096, 4128, 7234, 8192]) + + self.q_k_lens = itertools.product(qlen, klen) + + def tearDown(self) -> None: + torch.cuda.empty_cache() + super().tearDown() + + def test_forward(self, allmasked: bool=False): + import generic_scaled_masked_softmax_cuda + for qlen, klen in self.q_k_lens: + inputs = torch.normal(0, 2, (self.batch, self.attn, qlen, klen), dtype=self.dtype, device=self.device) + masks = ( + torch.randint(0, 2, (self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + if not allmasked else torch.ones((self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + ) + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, self.scale_t) + softmax_results_torch = forward_torch_softmax(inputs, masks, self.scale_t) + self.assertEqual( + softmax_results_torch.to(self.dtype), softmax_results, **self.thresh, msg=f"(q, k) = ({qlen, klen})") + + def test_backward(self, allmasked: bool=False): + import generic_scaled_masked_softmax_cuda + for qlen, klen in self.q_k_lens: + inputs = torch.normal(0, 2, (self.batch, self.attn, qlen, klen), dtype=self.dtype, device=self.device) + backward = torch.rand_like(inputs, dtype=torch.float16, device=self.device) + masks = ( + torch.randint(0, 2, (self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + if not allmasked else torch.ones((self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + ) + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, self.scale_t) + back_grad = generic_scaled_masked_softmax_cuda.backward(backward, softmax_results, self.scale_t) + inputs.requires_grad = True + softmax_results_torch = forward_torch_softmax(inputs, masks, self.scale_t) + softmax_results_torch.backward(backward) + self.assertEqual(back_grad, inputs.grad, **self.thresh, msg=f"(q, k) = ({qlen, klen})") + + def test_allmasked(self): + self.test_forward(True) + + def test_allmask_backward(self): + self.test_backward(True) - with torch.cuda.amp.autocast(dtype=dtype): - actual = fused_fn(attn_weights_0, total_mask) - self.assertEqual(actual.dtype, dtype) - expected = torch_fn(attn_weights_1, total_mask) - self.assertEqual(actual, expected) - - g0 = torch.randn_like(actual) - with torch.no_grad(): - g1 = g0.clone() - actual.backward(g0) - expected.backward(g1) if __name__ == "__main__": common_utils.run_tests() diff --git a/tests/L0/run_transformer/test_gpt_minimal.py b/tests/L0/run_transformer/test_gpt_minimal.py new file mode 100644 index 000000000..4f9b5e5eb --- /dev/null +++ b/tests/L0/run_transformer/test_gpt_minimal.py @@ -0,0 +1,232 @@ +from functools import partial +from typing import List +import time + +import torch + +import unittest + +from apex.transformer._ucc_util import HAS_UCC +from apex.transformer import parallel_state +from apex.transformer.enums import ModelType +from apex.transformer.tensor_parallel import model_parallel_cuda_manual_seed +from apex.transformer.pipeline_parallel.utils import ( + average_losses_across_data_parallel_group, unwrap_model, setup_microbatch_calculator, + get_ltor_masks_and_position_ids +) +from apex.transformer.pipeline_parallel.schedules.common import ( + _get_params_for_weight_decay_optimization, build_model +) +from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import ( + forward_backward_pipelining_without_interleaving, +) +from apex.transformer.testing.standalone_gpt import gpt_model_provider +from apex.transformer.testing import global_vars + +from apex.transformer.testing.distributed_test_base import UccDistributedTestBase, NcclDistributedTestBase + +from torch.testing._internal import common_utils +from torch.testing._internal.common_device_type import instantiate_device_type_tests + + +class GptTestBase: + + def _download_fancy_data(self): + text = """ + An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. + """ + text = text * 1024 + encoded = text.encode("ascii", "replace") + ints = [int(encoded[i]) for i in range(len(encoded))] + return torch.tensor(ints) + + # build a batch given sequence_len and batch size + def _generate_fancy_data_labels(self, sequence_len, batch_size): + temps = list() + for i in range(batch_size): + if self.inds is None or self.data_idx >= len(self.inds): + # hack as use of RNG will fall out of sync due to pipelines being different + model_parallel_cuda_manual_seed(self.MANUAL_SEED) + self.inds = torch.randperm(effective_length, device="cuda") + self.MANUAL_SEED += 1 + self.data_idx = 0 + data_idx_ = self.data_idx + offset = self.inds[data_idx_] + self.data_idx += 1 + curr = fancy_data[offset: offset + + sequence_len + 1].clone().detach() + temps.append(curr) + temp = torch.stack(temps, dim=0).cuda() + return temp + + def _get_batch(self, int_tensors: List[torch.Tensor]): + data = int_tensors[0] + # Unpack. + tokens_ = data.long() + labels = tokens_[:, 1:].contiguous() + tokens = tokens_[:, :-1].contiguous() + # Get the masks and position ids. + attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( + tokens, + self.N_VOCAB, # tokenizer.eod, + False, # args.reset_position_ids, + False, # args.reset_attention_mask, + False, # args.eod_mask_loss, + ) + return tokens, labels, loss_mask, attention_mask, position_ids + + # Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L75 + def _loss_func(self, loss_mask, output_tensor): + losses = output_tensor.float() + loss_mask = loss_mask.view(-1).float() + loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() + + # Reduce loss for logging. + averaged_loss = average_losses_across_data_parallel_group([loss]) + + return loss, {"lm loss": averaged_loss[0]} + + # Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L86 + def _fwd_step_func(self, batch, model): + """Forward step.""" + tokens, labels, loss_mask, attention_mask, position_ids = self._get_batch( + batch) + output_tensor = model(tokens, position_ids, + attention_mask, labels=labels) + return output_tensor, partial(self._loss_func, loss_mask) + + def _train(self, model, optim, pipeline_model_parallel_size, async_comm): + args = global_vars.get_args() + fwd_bwd_func = forward_backward_pipelining_without_interleaving + + tensor_shape = (args.seq_length, args.micro_batch_size, + args.hidden_size) + runtime = 0 + # training loop + for i in range(3): + since = time.time() + if torch.distributed.get_rank() == 0: + print("begin iter", i) + batch = [ + self._generate_fancy_data_labels( + args.seq_length, args.global_batch_size) + for _ in range(pipeline_model_parallel_size) + ] + if torch.distributed.get_rank() == 0: + print("finished making batch...") + optim.zero_grad() + fwd_bwd_func( + self._fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=tensor_shape, + async_comm=async_comm, + sequence_parallel_enabled=args.sequence_parallel, + ) + if torch.distributed.get_rank() == 0: + print("finished forward step") + # All-reduce layernorm parameters across model parallel nodes + # when sequence parallelism is used + if parallel_state.get_tensor_model_parallel_world_size() > 1 and global_vars.get_args().sequence_parallel: + for model_module in model: + unwrapped_model = unwrap_model(model_module) + for param in unwrapped_model.parameters(): + if getattr(param, 'sequence_parallel_enabled', False): + grad = param.grad + torch.distributed.all_reduce( + grad, group=parallel_state.get_tensor_model_parallel_group()) + optim.step() + if torch.distributed.get_rank() == 0: + print("finished iter", i) + runtime += time.time() - since + return runtime / 3.0 + + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_gpt(self): + self.MANUAL_SEED = 42 + self.inds = None + self.data_idx = 0 + self.N_VOCAB = 128 + init = True + + num_devices = torch.cuda.device_count() + tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices >= 4 else 1 + pipeline_model_parallel_size = num_devices // tensor_model_parallel_size + + override_args = { + "micro_batch_size": 2, + "num_layers": 16, + "hidden_size": 256, + "num_attention_heads": 8, + "max_position_embeddings": 512, + "seq_length": 512, + "global_batch_size": 128, + "pipeline_model_parallel_size": pipeline_model_parallel_size, + "tensor_model_parallel_size": tensor_model_parallel_size, + "world_size": self.world_size, + "rank": self.rank, + } + + global_vars.set_global_variables(override_args=override_args, ignore_unknown_args=True) + args = global_vars.get_args() + + for async_comm in (False,) if args.sequence_parallel else (False, True): + global fancy_data + global effective_length + + if init: + init = False + + fancy_data = self._download_fancy_data() + args = global_vars.get_args() + args.model_type = ModelType.encoder_or_decoder + effective_length = fancy_data.size(0) // args.seq_length + effective_length = fancy_data.size(0) - args.seq_length + + args.padded_vocab_size = 128 + setup_microbatch_calculator( + args.rank, + args.rampup_batch_size, + args.global_batch_size, + args.micro_batch_size, + args.data_parallel_size, + ) + + print(args.tensor_model_parallel_size, "MODEL PARALLEL SIZE") + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=args.tensor_model_parallel_size, + pipeline_model_parallel_size_=args.pipeline_model_parallel_size, + default_backend="nccl", + p2p_backend=self.DISTRIBUTED_BACKEND, + ) + + model_parallel_cuda_manual_seed(0) + model = build_model( + gpt_model_provider, + wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, + virtual_pipeline_model_parallel_size=None, + cpu_offload=args.cpu_offload, + ) + assert isinstance(model, list), model + _param_groups = _get_params_for_weight_decay_optimization(model) + optim = torch.optim.Adam(_param_groups) + runtime = self._train( + model, optim, args.pipeline_model_parallel_size, async_comm) + + parallel_state.destroy_model_parallel() + torch.distributed.barrier() + + +class NcclGptTest(GptTestBase, NcclDistributedTestBase): + pass + + +@unittest.skipUnless(HAS_UCC, "requires pytorch to be built with native ucc") +class UccGptTest(GptTestBase, UccDistributedTestBase): + pass + + +if __name__ == "__main__": + common_utils.run_tests() diff --git a/tests/L0/run_transformer/test_layers.py b/tests/L0/run_transformer/test_layers.py index b3b2eb2fc..e2df2acfc 100644 --- a/tests/L0/run_transformer/test_layers.py +++ b/tests/L0/run_transformer/test_layers.py @@ -47,43 +47,43 @@ def test_all_gather_parity(self) -> None: for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest(tensor_model_parallel_world_size=tensor_model_parallel_world_size): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() - cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") - with torch.no_grad(): - tensor = tensor_model_parallel_rank * torch.ones( - self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) - numel = tensor.numel() - numel_gathered = tensor_model_parallel_world_size * numel - gathered = torch.empty( - torch.Size((numel_gathered,)), - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - chunks = [ - gathered[i * numel : (i + 1) * numel] - for i in range(tensor_model_parallel_world_size) - ] - all_gather(chunks, tensor, group=parallel_state.get_tensor_model_parallel_group()) - - gathered_for_base = torch.empty( - torch.Size((numel_gathered,)), - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - _all_gather_base( - gathered_for_base, - tensor, - group=parallel_state.get_tensor_model_parallel_group(), - ) - - self.assertEqual(gathered, gathered_for_base) - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() + cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") + with torch.no_grad(): + tensor = tensor_model_parallel_rank * torch.ones( + self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) + numel = tensor.numel() + numel_gathered = tensor_model_parallel_world_size * numel + gathered = torch.empty( + torch.Size((numel_gathered,)), + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + chunks = [ + gathered[i * numel : (i + 1) * numel] + for i in range(tensor_model_parallel_world_size) + ] + all_gather(chunks, tensor, group=parallel_state.get_tensor_model_parallel_group()) + + gathered_for_base = torch.empty( + torch.Size((numel_gathered,)), + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + _all_gather_base( + gathered_for_base, + tensor, + group=parallel_state.get_tensor_model_parallel_group(), + ) + + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + self.assertEqual(gathered, gathered_for_base, msg=msg) + parallel_state.destroy_model_parallel() @torch.no_grad() @unittest.skipIf(torch.cuda.device_count() < 2, "Requires >=2 GPUs") @@ -95,113 +95,111 @@ def test_reduce_scatter_parity(self) -> None: for tensor_model_parallel_world_size in range(2, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest(tensor_model_parallel_world_size=tensor_model_parallel_world_size): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() - cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") - with torch.no_grad(): - input = torch.cat([ - i * torch.ones(self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) - for i in range(tensor_model_parallel_world_size) - ]) - input_list = [t.clone() for t in input.chunk(tensor_model_parallel_world_size)] - output = torch.empty( - self.tensor_shape, - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - reduce_scatter( - output, input_list, - group=parallel_state.get_tensor_model_parallel_group(), - ) - - output_for_base = torch.empty( - self.tensor_shape, - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - _reduce_scatter_base( - output_for_base, - input, - group=parallel_state.get_tensor_model_parallel_group(), - ) - - self.assertEqual(output, output_for_base) - self.assertEqual(input, torch.cat(input_list)) - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() + cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") + with torch.no_grad(): + input = torch.cat([ + i * torch.ones(self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) + for i in range(tensor_model_parallel_world_size) + ]) + input_list = [t.clone() for t in input.chunk(tensor_model_parallel_world_size)] + output = torch.empty( + self.tensor_shape, + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + reduce_scatter( + output, input_list, + group=parallel_state.get_tensor_model_parallel_group(), + ) + + output_for_base = torch.empty( + self.tensor_shape, + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + _reduce_scatter_base( + output_for_base, + input, + group=parallel_state.get_tensor_model_parallel_group(), + ) + + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + self.assertEqual(output, output_for_base, msg=msg) + self.assertEqual(input, torch.cat(input_list), msg=msg) + parallel_state.destroy_model_parallel() def test_parallel_embedding(self) -> None: for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - set_random_seed(self.SEED + 1) - input_tensor = torch.randint( - 0, - self.VOCAB_SIZE, - ( - self.BATCH_SIZE, - self.SEQUENCE_LENGTH, - ), - device="cuda", - ) - loss_weight = torch.randn( - ( - self.BATCH_SIZE, - self.SEQUENCE_LENGTH, - self.HIDDEN_SIZE, - ), - device="cuda", - ) - - set_random_seed(self.SEED) - embedding_torch = nn.Embedding( - self.VOCAB_SIZE, + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + set_random_seed(self.SEED + 1) + input_tensor = torch.randint( + 0, + self.VOCAB_SIZE, + ( + self.BATCH_SIZE, + self.SEQUENCE_LENGTH, + ), + device="cuda", + ) + loss_weight = torch.randn( + ( + self.BATCH_SIZE, + self.SEQUENCE_LENGTH, self.HIDDEN_SIZE, - ).cuda() - output_torch = embedding_torch(input_tensor) - loss_torch = torch.mul(output_torch, loss_weight).sum() - loss_torch.backward() - - # N.B.(mkozuki): With affine weight initialization on GPU, - # it's super difficult to keep the consistency with nn.Embedding. - # Thus, turning on `use_cpu_initialization`. - set_random_seed(self.SEED) - embedding_vocab_parallel = layers.VocabParallelEmbedding( - self.VOCAB_SIZE, - self.HIDDEN_SIZE, - init_method=nn.init.normal_, - use_cpu_initialization=True, - ).cuda() - output_vocab_parallel = embedding_vocab_parallel(input_tensor) - loss_vocab_parallel = torch.mul( - output_vocab_parallel, loss_weight - ).sum() - loss_vocab_parallel.backward() - - self.assertEqual(output_torch, output_vocab_parallel) - self.assertEqual(loss_torch, loss_vocab_parallel) - - splitted_weight_torch = torch.split( - embedding_torch.weight.grad, - self.VOCAB_SIZE - // tensor_model_parallel_world_size, - 0, - )[parallel_state.get_tensor_model_parallel_rank()] - self.assertEqual( - splitted_weight_torch, embedding_vocab_parallel.weight.grad - ) - - parallel_state.destroy_model_parallel() + ), + device="cuda", + ) + + set_random_seed(self.SEED) + embedding_torch = nn.Embedding( + self.VOCAB_SIZE, + self.HIDDEN_SIZE, + ).cuda() + output_torch = embedding_torch(input_tensor) + loss_torch = torch.mul(output_torch, loss_weight).sum() + loss_torch.backward() + + # N.B.(mkozuki): With affine weight initialization on GPU, + # it's super difficult to keep the consistency with nn.Embedding. + # Thus, turning on `use_cpu_initialization`. + set_random_seed(self.SEED) + embedding_vocab_parallel = layers.VocabParallelEmbedding( + self.VOCAB_SIZE, + self.HIDDEN_SIZE, + init_method=nn.init.normal_, + use_cpu_initialization=True, + ).cuda() + output_vocab_parallel = embedding_vocab_parallel(input_tensor) + loss_vocab_parallel = torch.mul( + output_vocab_parallel, loss_weight + ).sum() + loss_vocab_parallel.backward() + + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + self.assertEqual(output_torch, output_vocab_parallel, msg=msg) + self.assertEqual(loss_torch, loss_vocab_parallel, msg=msg) + + splitted_weight_torch = torch.split( + embedding_torch.weight.grad, + self.VOCAB_SIZE + // tensor_model_parallel_world_size, + 0, + )[parallel_state.get_tensor_model_parallel_rank()] + self.assertEqual( + splitted_weight_torch, embedding_vocab_parallel.weight.grad, msg=msg, + ) + + parallel_state.destroy_model_parallel() def _affine_weight_init_test_impl( self, init_device: str, is_column_parallel: bool @@ -210,56 +208,55 @@ def _affine_weight_init_test_impl( for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size + ) + input_size: int = self.INPUT_SIZE_COEFF * tensor_model_parallel_world_size + output_size: int = self.OUTPUT_SIZE_COEFF * tensor_model_parallel_world_size + + weight_shape = ( + (self.OUTPUT_SIZE_COEFF, input_size) + if is_column_parallel + else (output_size, self.INPUT_SIZE_COEFF) + ) + weight = torch.empty(weight_shape) + set_random_seed(self.SEED) + + sharding_dim_size = ( + self.OUTPUT_SIZE_COEFF + if is_column_parallel + else self.INPUT_SIZE_COEFF + ) + + if init_device == "cpu": + layers._initialize_affine_weight_cpu( + weight, + output_size, + input_size, + sharding_dim_size, + dim, + nn.init.normal_, + params_dtype=torch.float32, ) - input_size: int = self.INPUT_SIZE_COEFF * tensor_model_parallel_world_size - output_size: int = self.OUTPUT_SIZE_COEFF * tensor_model_parallel_world_size - - weight_shape = ( - (self.OUTPUT_SIZE_COEFF, input_size) - if is_column_parallel - else (output_size, self.INPUT_SIZE_COEFF) + else: + layers._initialize_affine_weight_gpu( + weight, torch.nn.init.normal_, dim ) - weight = torch.empty(weight_shape) - set_random_seed(self.SEED) + # Target + set_random_seed(self.SEED) + if init_device == "cpu": + main_weight = torch.empty(output_size, input_size) + nn.init.normal_(main_weight) + curr_weight = torch.split(main_weight, sharding_dim_size, dim=dim)[ + parallel_state.get_tensor_model_parallel_rank() + ] + else: + curr_weight = torch.empty(*weight_shape) + nn.init.normal_(curr_weight) - sharding_dim_size = ( - self.OUTPUT_SIZE_COEFF - if is_column_parallel - else self.INPUT_SIZE_COEFF - ) - - if init_device == "cpu": - layers._initialize_affine_weight_cpu( - weight, - output_size, - input_size, - sharding_dim_size, - dim, - nn.init.normal_, - params_dtype=torch.float32, - ) - else: - layers._initialize_affine_weight_gpu( - weight, torch.nn.init.normal_, dim - ) - # Target - set_random_seed(self.SEED) - if init_device == "cpu": - main_weight = torch.empty(output_size, input_size) - nn.init.normal_(main_weight) - curr_weight = torch.split(main_weight, sharding_dim_size, dim=dim)[ - parallel_state.get_tensor_model_parallel_rank() - ] - else: - curr_weight = torch.empty(*weight_shape) - nn.init.normal_(curr_weight) - self.assertEqual(curr_weight, weight) - parallel_state.destroy_model_parallel() + self.assertEqual( + curr_weight, weight, msg=f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}") + parallel_state.destroy_model_parallel() def test_affine_weight_init_column_parallel_cpu(self) -> None: self._affine_weight_init_test_impl(init_device="cpu", is_column_parallel=True) @@ -282,6 +279,7 @@ def test_row_parallel_linear_gradient_accumulation_fusion(self) -> None: def test_row_parallel_linear_gradient_accumulation_fusion_in_fp16(self) -> None: self._row_parallel_linear_test_impl(True, True, False) + # fails on native ucc and torch ucc: ucc does not support reduce scatter @unittest.skipIf(torch.cuda.device_count() < 2, "Sequence Parallel requires >=2 GPUs") def test_row_parallel_linear_sequence_parallel(self) -> None: self._row_parallel_linear_test_impl(False, False, True) @@ -304,103 +302,105 @@ def _row_parallel_linear_test_impl( ): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size, - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - set_random_seed(self.SEED) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + set_random_seed(self.SEED) + + linear = layers.RowParallelLinear( + self.HIDDEN_SIZE, + self.HIDDEN_SIZE, + keep_master_weight_for_test=True, + params_dtype=torch.float32, + use_cpu_initialization=True, + gradient_accumulation_fusion=gradient_accumulation_fusion, + accumulation_in_fp16=accumulation_in_fp16, + sequence_parallel_enabled=sequence_parallel_enabled, + # n.b.(mkozuki): RowParallelLinear is constructed with `input_is_parallel=True` + # by default, e.g. https://github.com/NVIDIA/NeMo/blob/782b4e1652aaa43c8be390d9\ + # db0dc89544afa080/nemo/collections/nlp/modules/common/megatron/transformer.py#L204 + input_is_parallel=True, + ).cuda() + if accumulation_in_fp16: + linear = linear.half() + # Simulate the situation where fusion of weight grad calculation and gradient accumulation is enabled. + if gradient_accumulation_fusion: + with torch.no_grad(): + linear.weight.main_grad = torch.zeros_like(linear.weight) - linear = layers.RowParallelLinear( - self.HIDDEN_SIZE, - self.HIDDEN_SIZE, - keep_master_weight_for_test=True, - params_dtype=torch.float32, - use_cpu_initialization=True, - gradient_accumulation_fusion=gradient_accumulation_fusion, - accumulation_in_fp16=accumulation_in_fp16, - sequence_parallel_enabled=sequence_parallel_enabled, - # n.b.(mkozuki): RowParallelLinear is constructed with `input_is_parallel=True` - # by default, e.g. https://github.com/NVIDIA/NeMo/blob/782b4e1652aaa43c8be390d9\ - # db0dc89544afa080/nemo/collections/nlp/modules/common/megatron/transformer.py#L204 - input_is_parallel=True, - ).cuda() - if accumulation_in_fp16: - linear = linear.half() - # Simulate the situation where fusion of weight grad calculation and gradient accumulation is enabled. - if gradient_accumulation_fusion: - with torch.no_grad(): - linear.weight.main_grad = torch.zeros_like(linear.weight) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" - with torch.no_grad(): - orig_input_tensor = torch.randn(tensor_shape, requires_grad=True, device="cuda") - orig_loss_weight = torch.randn(tensor_shape, device="cuda") - input_tensor = orig_input_tensor.chunk( + with torch.no_grad(): + orig_input_tensor = torch.randn(tensor_shape, requires_grad=True, device="cuda") + orig_loss_weight = torch.randn(tensor_shape, device="cuda") + input_tensor = orig_input_tensor.chunk( + chunks=tensor_model_parallel_world_size, + dim=2, + )[parallel_state.get_tensor_model_parallel_rank()].contiguous() + if sequence_parallel_enabled: + loss_weight = orig_loss_weight.chunk( chunks=tensor_model_parallel_world_size, - dim=2, - )[parallel_state.get_tensor_model_parallel_rank()].contiguous() - if sequence_parallel_enabled: - loss_weight = orig_loss_weight.chunk( - chunks=tensor_model_parallel_world_size, - dim=0, - )[parallel_state.get_tensor_model_parallel_rank()] - else: - loss_weight = orig_loss_weight - if accumulation_in_fp16: - orig_input_tensor = orig_input_tensor.half() - input_tensor = input_tensor.half() - loss_weight = loss_weight.half() - input_tensor.requires_grad_() - output, _ = linear(input_tensor) - loss = torch.mul(output, loss_weight).sum() - loss.backward() - self.assertIsNotNone(input_tensor.grad) + dim=0, + )[parallel_state.get_tensor_model_parallel_rank()] + else: + loss_weight = orig_loss_weight + if accumulation_in_fp16: + orig_input_tensor = orig_input_tensor.half() + input_tensor = input_tensor.half() + loss_weight = loss_weight.half() + input_tensor.requires_grad_() + output, _ = linear(input_tensor) + loss = torch.mul(output, loss_weight).sum() + loss.backward() + self.assertIsNotNone(input_tensor.grad, msg=msg) + + ref_linear = nn.Linear( + in_features=self.HIDDEN_SIZE, + out_features=self.HIDDEN_SIZE, + bias=False, + device="cuda", + ) + with torch.no_grad(): + dldy = orig_loss_weight.clone() + x = orig_input_tensor.clone() + ref_linear.weight.copy_(linear.master_weight) + if accumulation_in_fp16: + ref_linear = ref_linear.half() + x.requires_grad_() + expected_output = ref_linear(x) + expected_loss = torch.mul(expected_output, dldy).sum() + expected_loss.backward() - ref_linear = nn.Linear( - in_features=self.HIDDEN_SIZE, - out_features=self.HIDDEN_SIZE, - bias=False, - device="cuda", - ) - with torch.no_grad(): - dldy = orig_loss_weight.clone() - x = orig_input_tensor.clone() - ref_linear.weight.copy_(linear.master_weight) - if accumulation_in_fp16: - ref_linear = ref_linear.half() - x.requires_grad_() - expected_output = ref_linear(x) - expected_loss = torch.mul(expected_output, dldy).sum() - expected_loss.backward() - - if not accumulation_in_fp16: - if sequence_parallel_enabled: - self.assertEqual( - x=output, - y=expected_output.chunk( - chunks=tensor_model_parallel_world_size, - dim=0, - )[parallel_state.get_tensor_model_parallel_rank()], - ) - else: - self.assertEqual( - x=output, - y=expected_output, - ) - - grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" - # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. - if tensor_model_parallel_world_size == 1: + if not accumulation_in_fp16: + if sequence_parallel_enabled: self.assertEqual( - x=getattr(linear.weight, grad_attr_name), - y=ref_linear.weight.grad.chunk( + x=output, + y=expected_output.chunk( chunks=tensor_model_parallel_world_size, dim=0, )[parallel_state.get_tensor_model_parallel_rank()], + msg=msg, + ) + else: + self.assertEqual( + x=output, + y=expected_output, + msg=msg, ) - parallel_state.destroy_model_parallel() + grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" + # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. + if tensor_model_parallel_world_size == 1: + self.assertEqual( + x=getattr(linear.weight, grad_attr_name), + y=ref_linear.weight.grad.chunk( + chunks=tensor_model_parallel_world_size, + dim=0, + )[parallel_state.get_tensor_model_parallel_rank()], + msg=msg, + ) + + parallel_state.destroy_model_parallel() def test_column_parallel_linear(self): self._column_parallel_linear_test_impl(False, False, False, False) @@ -438,112 +438,115 @@ def _column_parallel_linear_test_impl( if async_tensor_model_parallel_allreduce and sequence_parallel_enabled: if tensor_model_parallel_world_size == 1: continue - with self.subTest(tensor_model_parallel_world_size=tensor_model_parallel_world_size): - if self.world_size % tensor_model_parallel_world_size: - continue - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - - input_tensor_shape = self.tensor_shape - expected_output_shape = self.tensor_shape - # When sequence parallel, `gather_output` is disabled, i.e., - # output of matmul isn't gathered in dimension of feature/hidden (last dim). - if sequence_parallel_enabled: - expected_output_shape[-1] //= tensor_model_parallel_world_size - - # tensor's shape is [sequence length, batch size, hidden size] - set_random_seed(self.SEED) - linear = layers.ColumnParallelLinear( - self.HIDDEN_SIZE, - self.HIDDEN_SIZE, + if self.world_size % tensor_model_parallel_world_size: + continue + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + + input_tensor_shape = self.tensor_shape + expected_output_shape = self.tensor_shape + # When sequence parallel, `gather_output` is disabled, i.e., + # output of matmul isn't gathered in dimension of feature/hidden (last dim). + if sequence_parallel_enabled: + expected_output_shape[-1] //= tensor_model_parallel_world_size + + # tensor's shape is [sequence length, batch size, hidden size] + set_random_seed(self.SEED) + linear = layers.ColumnParallelLinear( + self.HIDDEN_SIZE, + self.HIDDEN_SIZE, + bias=False, + keep_master_weight_for_test=True, + params_dtype=torch.float32, + use_cpu_initialization=True, + gather_output=not sequence_parallel_enabled, + no_async_tensor_model_parallel_allreduce=not async_tensor_model_parallel_allreduce, + gradient_accumulation_fusion=gradient_accumulation_fusion, + accumulation_in_fp16=accumulation_in_fp16, + sequence_parallel_enabled=sequence_parallel_enabled, + ).cuda() + if accumulation_in_fp16: + linear = linear.half() + + # Simulate the situation where fusion of weight grad calculation and gradient accumulation happens. + if gradient_accumulation_fusion: + with torch.no_grad(): + linear.weight.main_grad = torch.zeros_like(linear.weight) + + orig_input_tensor = torch.randn(input_tensor_shape, device="cuda", requires_grad=True) + if accumulation_in_fp16: + orig_input_tensor = orig_input_tensor.half() + if sequence_parallel_enabled: + input_tensor = list( + orig_input_tensor.chunk(tensor_model_parallel_world_size, dim=0) + )[parallel_state.get_tensor_model_parallel_rank()] + else: + input_tensor = orig_input_tensor + output, _ = linear(input_tensor) + # The order of dimension is expected to be (sequence, batch, hidden) + self.assertEqual(output.shape, expected_output_shape, msg=msg) + + orig_loss_weight = torch.randn(input_tensor_shape, device="cuda") + if accumulation_in_fp16: + orig_loss_weight = orig_loss_weight.half() + if sequence_parallel_enabled: + loss_weight = orig_loss_weight.chunk( + tensor_model_parallel_world_size, dim=2, + )[parallel_state.get_tensor_model_parallel_rank()] + else: + loss_weight = orig_loss_weight + loss = torch.mul(output, loss_weight).sum() + loss.backward() + + with torch.no_grad(): + dldy = orig_loss_weight.clone() + x = orig_input_tensor.clone() + ref_linear = nn.Linear( + in_features=self.HIDDEN_SIZE, + out_features=self.HIDDEN_SIZE, bias=False, - keep_master_weight_for_test=True, - params_dtype=torch.float32, - use_cpu_initialization=True, - gather_output=not sequence_parallel_enabled, - no_async_tensor_model_parallel_allreduce=not async_tensor_model_parallel_allreduce, - gradient_accumulation_fusion=gradient_accumulation_fusion, - accumulation_in_fp16=accumulation_in_fp16, - sequence_parallel_enabled=sequence_parallel_enabled, - ).cuda() - if accumulation_in_fp16: - linear = linear.half() - - # Simulate the situation where fusion of weight grad calculation and gradient accumulation happens. - if gradient_accumulation_fusion: - with torch.no_grad(): - linear.weight.main_grad = torch.zeros_like(linear.weight) - - orig_input_tensor = torch.randn(input_tensor_shape, device="cuda", requires_grad=True) - if accumulation_in_fp16: - orig_input_tensor = orig_input_tensor.half() - if sequence_parallel_enabled: - input_tensor = list( - orig_input_tensor.chunk(tensor_model_parallel_world_size, dim=0) - )[parallel_state.get_tensor_model_parallel_rank()] - else: - input_tensor = orig_input_tensor - output, _ = linear(input_tensor) - # The order of dimension is expected to be (sequence, batch, hidden) - self.assertEqual(output.shape, expected_output_shape) - - orig_loss_weight = torch.randn(input_tensor_shape, device="cuda") + device="cuda", + ) if accumulation_in_fp16: - orig_loss_weight = orig_loss_weight.half() - if sequence_parallel_enabled: - loss_weight = orig_loss_weight.chunk( - tensor_model_parallel_world_size, dim=2, - )[parallel_state.get_tensor_model_parallel_rank()] - else: - loss_weight = orig_loss_weight - loss = torch.mul(output, loss_weight).sum() - loss.backward() - - with torch.no_grad(): - dldy = orig_loss_weight.clone() - x = orig_input_tensor.clone() - ref_linear = nn.Linear( - in_features=self.HIDDEN_SIZE, - out_features=self.HIDDEN_SIZE, - bias=False, - device="cuda", - ) - if accumulation_in_fp16: - ref_linear = ref_linear.half() - # NOTE(mkozuki): `master_weight` is available because `keep_master_weight_for_test` is set. - ref_linear.weight.copy_(linear.master_weight) - x.requires_grad_() - expected_output = ref_linear(x) - if sequence_parallel_enabled: - chunk = expected_output.chunk( - tensor_model_parallel_world_size, - dim=2, - )[parallel_state.get_tensor_model_parallel_rank()] - self.assertEqual( - x=output, - y=chunk, - ) - else: - self.assertEqual( - x=output, - y=expected_output, - ) + ref_linear = ref_linear.half() + # NOTE(mkozuki): `master_weight` is available because `keep_master_weight_for_test` is set. + ref_linear.weight.copy_(linear.master_weight) + x.requires_grad_() + expected_output = ref_linear(x) + if sequence_parallel_enabled: + chunk = expected_output.chunk( + tensor_model_parallel_world_size, + dim=2, + )[parallel_state.get_tensor_model_parallel_rank()] + self.assertEqual( + x=output, + y=chunk, + msg=msg, + ) + else: + self.assertEqual( + x=output, + y=expected_output, + msg=msg, + ) - expected_loss = torch.mul(expected_output, dldy).sum() - expected_loss.backward() - grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" - # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. - if tensor_model_parallel_world_size == 1: - self.assertEqual( - x=getattr(linear.weight, grad_attr_name), - y=ref_linear.weight.grad.chunk( - chunks=tensor_model_parallel_world_size, - dim=0, - )[parallel_state.get_tensor_model_parallel_rank()], - ) + expected_loss = torch.mul(expected_output, dldy).sum() + expected_loss.backward() + grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" + # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. + if tensor_model_parallel_world_size == 1: + self.assertEqual( + x=getattr(linear.weight, grad_attr_name), + y=ref_linear.weight.grad.chunk( + chunks=tensor_model_parallel_world_size, + dim=0, + )[parallel_state.get_tensor_model_parallel_rank()], + msg=msg, + ) - parallel_state.destroy_model_parallel() + parallel_state.destroy_model_parallel() class NcclTensorParallelLayerTest(TensorParallelLayerTestBase, NcclDistributedTestBase): diff --git a/tests/L0/run_transformer/test_mapping.py b/tests/L0/run_transformer/test_mapping.py index 9ebda6dc6..18fa07636 100644 --- a/tests/L0/run_transformer/test_mapping.py +++ b/tests/L0/run_transformer/test_mapping.py @@ -18,67 +18,65 @@ def test_reduce(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) - t = torch.full((10, 10, 10, 10), 50, device=f"cuda:{self.rank}") - expected = torch.full( - (10, 10, 10, 10), - 50 * tensor_model_paralell_world_size, - device=f"cuda:{self.rank}", - ) - self.assertTrue(torch.equal(mappings._reduce(t), expected)) - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) + t = torch.full((10, 10, 10, 10), 50, device=f"cuda:{self.rank}") + expected = torch.full( + (10, 10, 10, 10), + 50 * tensor_model_paralell_world_size, + device=f"cuda:{self.rank}", + ) + self.assertTrue( + torch.equal(mappings._reduce(t), expected), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}", + ) + parallel_state.destroy_model_parallel() def test_split(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) - tensors = [ - torch.randn(10, 1) - for rank in range(tensor_model_paralell_world_size) - ] - x = torch.cat(tensors, 1) - out = mappings._split_along_last_dim(x) - self.assertTrue( - torch.equal( - out, tensors[parallel_state.get_tensor_model_parallel_rank()] - ) - ) - parallel_state.destroy_model_parallel() + tensors = [ + torch.randn(10, 1) + for _ in range(tensor_model_paralell_world_size) + ] + x = torch.cat(tensors, 1) + out = mappings._split_along_last_dim(x) + self.assertTrue( + torch.equal( + out, tensors[parallel_state.get_tensor_model_parallel_rank()] + ), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}" + ) + parallel_state.destroy_model_parallel() def test_gather(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) - device = f"cuda:{self.rank}" - gathered = mappings._gather_along_last_dim( - torch.tensor( - [parallel_state.get_tensor_model_parallel_rank()], device=device - ) - ) - expected = torch.tensor( - [rank for rank in range(tensor_model_paralell_world_size)], - device=device, + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) + device = f"cuda:{self.rank}" + gathered = mappings._gather_along_last_dim( + torch.tensor( + [parallel_state.get_tensor_model_parallel_rank()], device=device ) - self.assertTrue(torch.equal(gathered, expected)) - parallel_state.destroy_model_parallel() + ) + expected = torch.tensor( + [rank for rank in range(tensor_model_paralell_world_size)], + device=device, + ) + self.assertTrue( + torch.equal(gathered, expected), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}", + ) + parallel_state.destroy_model_parallel() class NcclMappingTest(MappingTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/test_microbatches.py b/tests/L0/run_transformer/test_microbatches.py index 0d4b50972..13e64881a 100644 --- a/tests/L0/run_transformer/test_microbatches.py +++ b/tests/L0/run_transformer/test_microbatches.py @@ -39,36 +39,36 @@ def _test(self, rampup_batch_size: Optional[List[int]]) -> None: continue if self.world_size % data_parallel_size != 0: continue - with self.subTest(data_parallel_size=data_parallel_size): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=self.world_size // data_parallel_size, - pipeline_model_parallel_size_=1, - ) - self.assertEqual(data_parallel_size, parallel_state.get_data_parallel_world_size()) - - _reconfigure_microbatch_calculator( - self.rank, - rampup_batch_size, - self.GLOBAL_BATCH_SIZE, - self.MICRO_BATCH_SIZE, - data_parallel_size, - ) - - self.assertEqual(get_micro_batch_size(), expected_micro_batch_size) - self.assertEqual(get_num_microbatches(), expected_global_batch_size / expected_micro_batch_size / data_parallel_size) - current_global_batch_size = get_current_global_batch_size() - self.assertEqual(current_global_batch_size, expected_global_batch_size) - - # Make sure `global_batch_size` equals to the final global batch size after - # certain number of updates. - if rampup_batch_size: - update_num_microbatches(current_global_batch_size) - for i in range(100): - current_global_batch_size = get_current_global_batch_size() - update_num_microbatches(current_global_batch_size) + msg = f"data_parallel_size: {data_parallel_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=self.world_size // data_parallel_size, + pipeline_model_parallel_size_=1, + ) + self.assertEqual(data_parallel_size, parallel_state.get_data_parallel_world_size(), msg=msg) + + _reconfigure_microbatch_calculator( + self.rank, + rampup_batch_size, + self.GLOBAL_BATCH_SIZE, + self.MICRO_BATCH_SIZE, + data_parallel_size, + ) + + self.assertEqual(get_micro_batch_size(), expected_micro_batch_size, msg=msg) + self.assertEqual(get_num_microbatches(), expected_global_batch_size / expected_micro_batch_size / data_parallel_size, msg=msg) + current_global_batch_size = get_current_global_batch_size() + self.assertEqual(current_global_batch_size, expected_global_batch_size, msg=msg) + + # Make sure `global_batch_size` equals to the final global batch size after + # certain number of updates. + if rampup_batch_size: + update_num_microbatches(current_global_batch_size) + for i in range(100): current_global_batch_size = get_current_global_batch_size() - self.assertEqual(get_current_global_batch_size(), self.GLOBAL_BATCH_SIZE) - parallel_state.destroy_model_parallel() + update_num_microbatches(current_global_batch_size) + current_global_batch_size = get_current_global_batch_size() + self.assertEqual(get_current_global_batch_size(), self.GLOBAL_BATCH_SIZE, msg=msg) + parallel_state.destroy_model_parallel() def test_constant_microbatch_calculator(self): self._test(rampup_batch_size=None) diff --git a/tests/L0/run_transformer/test_parallel_state.py b/tests/L0/run_transformer/test_parallel_state.py index 0314895ea..a86c2fb36 100644 --- a/tests/L0/run_transformer/test_parallel_state.py +++ b/tests/L0/run_transformer/test_parallel_state.py @@ -28,42 +28,43 @@ def test_initialize_model_parallel(self) -> None: self.assertFalse(parallel_state.model_parallel_is_initialized()) for tensor_model_parallel_world_size in range(1, self.world_size + 1): - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - if self.world_size % tensor_model_parallel_world_size: - continue - - pipeline_model_parallel_world_size = ( - self.world_size // tensor_model_parallel_world_size - ) - - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - pipeline_model_parallel_size_=pipeline_model_parallel_world_size, - ) - self.assertEqual( - tensor_model_parallel_world_size, - parallel_state.get_tensor_model_parallel_world_size(), - ) - expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( - self.rank, tensor_model_parallel_world_size - ) - self.assertEqual( - expected_tensor_model_parallel_rank, - parallel_state.get_tensor_model_parallel_rank(), - ) - - expected_tensor_model_parallel_src_rank = ( - self.rank // tensor_model_parallel_world_size - ) * tensor_model_parallel_world_size - self.assertEqual( - expected_tensor_model_parallel_src_rank, - parallel_state.get_tensor_model_parallel_src_rank(), - ) - - parallel_state.destroy_model_parallel() - self.assertFalse(parallel_state.model_parallel_is_initialized()) + msg = f"tensor_model_parallel_world_siz: {tensor_model_parallel_world_size}" + if self.world_size % tensor_model_parallel_world_size: + continue + + pipeline_model_parallel_world_size = ( + self.world_size // tensor_model_parallel_world_size + ) + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + ) + self.assertEqual( + tensor_model_parallel_world_size, + parallel_state.get_tensor_model_parallel_world_size(), + msg=msg, + ) + expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( + self.rank, tensor_model_parallel_world_size + ) + self.assertEqual( + expected_tensor_model_parallel_rank, + parallel_state.get_tensor_model_parallel_rank(), + msg=msg, + ) + + expected_tensor_model_parallel_src_rank = ( + self.rank // tensor_model_parallel_world_size + ) * tensor_model_parallel_world_size + self.assertEqual( + expected_tensor_model_parallel_src_rank, + parallel_state.get_tensor_model_parallel_src_rank(), + msg=msg, + ) + + parallel_state.destroy_model_parallel() + self.assertFalse(parallel_state.model_parallel_is_initialized(), msg=msg) def test_initialize_model_parallel_with_virtual_and_split(self) -> None: if self.world_size < 4: @@ -138,43 +139,44 @@ def test_initialize_model_parallel_decoder_only(self) -> None: self.assertFalse(parallel_state.model_parallel_is_initialized()) for tensor_model_parallel_world_size in range(1, self.world_size + 1): - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - if self.world_size % tensor_model_parallel_world_size: - continue - - pipeline_model_parallel_world_size = ( - self.world_size // tensor_model_parallel_world_size - ) - - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - pipeline_model_parallel_size_=pipeline_model_parallel_world_size, - pipeline_model_parallel_split_rank_=0, - ) - self.assertEqual( - tensor_model_parallel_world_size, - parallel_state.get_tensor_model_parallel_world_size(), - ) - expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( - self.rank, tensor_model_parallel_world_size - ) - self.assertEqual( - expected_tensor_model_parallel_rank, - parallel_state.get_tensor_model_parallel_rank(), - ) - - expected_tensor_model_parallel_src_rank = ( - self.rank // tensor_model_parallel_world_size - ) * tensor_model_parallel_world_size - self.assertEqual( - expected_tensor_model_parallel_src_rank, - parallel_state.get_tensor_model_parallel_src_rank(), - ) - - parallel_state.destroy_model_parallel() - self.assertFalse(parallel_state.model_parallel_is_initialized()) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + if self.world_size % tensor_model_parallel_world_size: + continue + + pipeline_model_parallel_world_size = ( + self.world_size // tensor_model_parallel_world_size + ) + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + pipeline_model_parallel_split_rank_=0, + ) + self.assertEqual( + tensor_model_parallel_world_size, + parallel_state.get_tensor_model_parallel_world_size(), + msg=msg, + ) + expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( + self.rank, tensor_model_parallel_world_size + ) + self.assertEqual( + expected_tensor_model_parallel_rank, + parallel_state.get_tensor_model_parallel_rank(), + msg=msg, + ) + + expected_tensor_model_parallel_src_rank = ( + self.rank // tensor_model_parallel_world_size + ) * tensor_model_parallel_world_size + self.assertEqual( + expected_tensor_model_parallel_src_rank, + parallel_state.get_tensor_model_parallel_src_rank(), + msg=msg, + ) + + parallel_state.destroy_model_parallel() + self.assertFalse(parallel_state.model_parallel_is_initialized(), msg=msg) class NcclParallelStateTest(ParallelStateTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index a409c40f2..a1193f1c6 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -1,12 +1,14 @@ +import contextlib import logging import itertools +import os +from packaging.version import parse, Version import re from typing import Optional, Tuple, List import unittest import torch from torch.testing._internal import common_utils -from torch.testing._internal import common_cuda from apex._autocast_utils import _get_autocast_dtypes from apex.transformer import parallel_state @@ -28,16 +30,25 @@ ) from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase from apex.transformer.testing.distributed_test_base import UccDistributedTestBase -from apex.transformer.testing.distributed_test_base import HAS_TORCH_UCC from apex.transformer.testing.distributed_test_base import HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER from apex.transformer.testing import commons as testing_utils - +from apex.transformer._ucc_util import HAS_UCC logging.getLogger("torch").setLevel(logging.WARNING) logging.getLogger("apex").setLevel(logging.WARNING) weight_coeff = 1024 +# Guard for https://github.com/pytorch/pytorch/pull/82450 +CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = False +ngc_container_2209, pytorch_113 = Version("22.09"), Version("1.13") +if parse(os.getenv("NVIDIA_PYTORCH_VERSION", "22.08")) >= ngc_container_2209: + CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = True +elif parse(torch.__version__) >= pytorch_113: + CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = True +else: + CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = False + def get_init_weights_func(offset: int = 0): @torch.no_grad() @@ -113,6 +124,7 @@ def _forward_backward_test_impl( *, default_backend: Optional[str] = None, p2p_backend: Optional[str] = None, + sync_batch_comm: bool = True, ) -> None: if fwd_bwd_func == _forward_backward_pipelining_with_interleaving: self.assertIsNotNone(virtual_pipeline_model_parallel_size) @@ -166,7 +178,6 @@ def _forward_backward_test_impl( hidden_size=self.HIDDEN_SIZE, ) - offset = pipeline_model_parallel_world_size if virtual_pipeline_model_parallel_size is not None else 0 for idx, model_module in enumerate(model): model_module = model_module.to(dtype) @@ -192,6 +203,7 @@ def _forward_backward_test_impl( async_comm=async_comm, grad_scaler=grad_scaler, deallocate_pipeline_output=deallocate_pipeline_outputs, + sync_batch_comm=sync_batch_comm, ) if dtype == get_dtype_for_comparison(): @@ -235,48 +247,58 @@ def test_learning_no_pipelining(self): def test_inference_no_pipelining(self): self._forward_backward_test_impl(True, forward_backward_no_pipelining, 1, None) - def test_learning_pipelining_without_interleaving(self): + def test_learning_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, forward_backward_pipelining_without_interleaving, None, None + False, forward_backward_pipelining_without_interleaving, None, None, sync_batch_comm=sync_batch_comm, ) - def test_inference_pipelining_without_interleaving(self): + def test_inference_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, forward_backward_pipelining_without_interleaving, None, None + True, forward_backward_pipelining_without_interleaving, None, None, sync_batch_comm=sync_batch_comm, ) - def test_learning_async_pipelining_without_interleaving(self): + def test_learning_async_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, forward_backward_pipelining_without_interleaving, None, None, async_comm=True + False, forward_backward_pipelining_without_interleaving, None, None, async_comm=True, + sync_batch_comm=sync_batch_comm, ) - def test_inference_async_pipelining_without_interleaving(self): + def test_inference_async_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, forward_backward_pipelining_without_interleaving, None, None, async_comm=True + True, forward_backward_pipelining_without_interleaving, None, None, async_comm=True, + sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_learning_pipelining_with_interleaving(self): + def test_learning_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2 + False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, + sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_inference_pipelining_with_interleaving(self): + def test_inference_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2 + True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, + sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_learning_async_pipelining_with_interleaving(self): + def test_learning_async_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True + False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True, + sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_inference_async_pipelining_with_interleaving(self): + def test_inference_async_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True + True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True, + sync_batch_comm=sync_batch_comm, ) @@ -294,12 +316,12 @@ def _run_hybrid_distributed_backend(self, forward_only: bool) -> None: @unittest.skipUnless(HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER, "Needs driver >= 470.42.01") def _test_hybrid_backends(self, forward_only: bool) -> None: - if HAS_TORCH_UCC: + if HAS_UCC: self._run_hybrid_distributed_backend(forward_only) else: with self.assertRaisesRegex( ImportError, - re.escape("UCC backend requires [torch_ucc](https://github.com/facebookresearch/torch_ucc) but not found"), + re.escape("UCC backend requires pytorch source build with UCC installed and enabled"), ): self._run_hybrid_distributed_backend(forward_only) @@ -309,6 +331,38 @@ def test_learning_pipelining_without_interleaving_ucc_for_p2p(self): def test_inference_pipelining_without_interleaving_ucc_for_p2p(self): self._test_hybrid_backends(True) + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_pipelining_without_interleaving_skyp_sync_after_batch_isend_irecv(self): + self.test_learning_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_pipelining_without_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_async_pipelining_without_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_learning_async_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_async_pipelining_without_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_async_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_learning_pipelining_with_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_pipelining_with_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_async_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_learning_async_pipelining_with_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_async_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_async_pipelining_with_interleaving(sync_batch_comm=False) + # n.b.(mkozuki): pipeline parallel w/o interleaving with UCX_TLS=tcp,sm fails. class UccPipelineParallelForwardBackwardTest(UccDistributedTestBase, PipelineParallelForwardBackwardTestBase): @@ -327,25 +381,22 @@ def world_size(self) -> int: @unittest.skipIf(torch.cuda.device_count() < 4, "Requires >= 4 GPUs") class NcclPipelineParallelWithToyParallelMLP(NcclDistributedTestBase): - GLOBAL_BATCH_SIZE = 16 - MICRO_BATCH_SIZE = 2 - HIDDEN_SIZE = 64 + GLOBAL_BATCH_SIZE: int = 16 + MICRO_BATCH_SIZE: int = 2 + HIDDEN_SIZE: int = 64 # TODO(mkozuki): Change `DECODER_SEQUENCE_LENGTH` to a value different from `ENCODER_SEQUENCE_LENGTH`. # To test forward_backward_pipelining_without_interleaving with `model_type=ModelType.encoder_and_decoder`, # `decoder_seq_length` is necessary and ideally should be different from `encoder_sequence_length` # but my laziness let me use the same value. # Note that you may have to either update `MyModel` def or define another `MyModel`. # to support different `DECODER_SEQUENCE_LENGTH`. - ENCODER_SEQUENCE_LENGTH = 32 - DECODER_SEQUENCE_LENGTH = 32 + ENCODER_SEQUENCE_LENGTH: int = 32 + DECODER_SEQUENCE_LENGTH: int = 32 @property def world_size(self) -> int: return min(torch.cuda.device_count(), 8) - # TODO(mkozuki): Add cases of async_comm=True - # TODO(mkozuki): Add loss check. - # TODO(mkozuki): Call `build_model` with `model_type`. # TODO(mkozuki): Set `tensor_model_parallel>1` for encoder_and_decoder as well if there's enough GPUs # in order to let `sequence_parallel_enabled` have an effect on tensor shape logic. def _forward_backward_test_impl( @@ -380,6 +431,7 @@ def _forward_backward_test_impl( micro_batch_size=self.MICRO_BATCH_SIZE, data_parallel_size=parallel_state.get_data_parallel_world_size(), ) + # TODO(mkozuki): Call `build_model` with `model_type`. model = build_model( testing_utils.mlp_provider_func, wrap_with_ddp=False, @@ -443,5 +495,222 @@ def test_pipelining_without_interleaving_sequence_parallel_encoder_or_decoder_ha self._forward_backward_test_impl(forward_only=False, sequence_parallel_enabled=True, model_type=ModelType.encoder_or_decoder, dtype=torch.half) +class NcclPipelineParallelWithCustomSyncContextHandler(NcclDistributedTestBase): + + GLOBAL_BATCH_SIZE = 32 + MICRO_BATCH_SIZE = 1 + HIDDEN_SIZE = 1 + + @property + def world_size(self) -> int: + return min(torch.cuda.device_count(), 8) + + @unittest.skipIf(torch.cuda.device_count() < 2 or torch.cuda.device_count() % 2 != 0, "Requires >= 2 GPUs") + def test_pipelining_without_interleaving_with_custom_sync_context_handler(self) -> None: + + # Parallel configuration + world_size = torch.cuda.device_count() + tensor_model_parallel_world_size = 1 + data_parallel_size = 2 if world_size > 2 else 1 + pipeline_model_parallel_world_size = world_size // data_parallel_size + + # Initialize pipeline parallelism + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + ) + pp_utils._reconfigure_microbatch_calculator( + rank=parallel_state.get_tensor_model_parallel_rank(), + rampup_batch_size=None, + global_batch_size=self.GLOBAL_BATCH_SIZE, + micro_batch_size=self.MICRO_BATCH_SIZE, + data_parallel_size=parallel_state.get_data_parallel_world_size(), + ) + pp_utils.update_num_microbatches(0) + + # Construct synthetic data + dtype = get_dtype_for_comparison() + hidden_size = self.HIDDEN_SIZE + microbatch_size = self.MICRO_BATCH_SIZE + global_batch_shape = ( + self.GLOBAL_BATCH_SIZE + // parallel_state.get_data_parallel_world_size(), + hidden_size, + hidden_size, + ) + batch = None + if parallel_state.is_pipeline_first_stage(): + batch = (torch.ones(global_batch_shape, dtype=dtype).cuda(), ) + + # Construct model + model = build_model( + testing_utils.model_provider_func, + wrap_with_ddp=True, + hidden_size=hidden_size, + )[0] + model = model.to(dtype) + model.module.apply(get_init_weights_func(0)) + + # Construct context that destroys all grads on exit + has_entered_grad_sync_context = False + has_exited_grad_sync_context = False + has_called_grad_sync_func = False + @contextlib.contextmanager + def custom_grad_sync_context(): + try: + nonlocal has_entered_grad_sync_context + has_entered_grad_sync_context = True + yield + finally: + nonlocal has_exited_grad_sync_context + has_exited_grad_sync_context = True + for param in model.parameters(): + param.grad = None + def custom_grad_sync_func(): + nonlocal has_called_grad_sync_func + has_called_grad_sync_func = True + + # Training step with pipeline parallelism + loss = forward_backward_pipelining_without_interleaving( + testing_utils.fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=(microbatch_size, hidden_size, hidden_size), + dtype=dtype, + async_comm=False, + grad_scaler=None, + deallocate_pipeline_outputs=False, + sequence_parallel_enabled=False, + custom_sync_context_handler=custom_grad_sync_context, + custom_grad_sync_func=custom_grad_sync_func, + ) + torch.cuda.synchronize() + + # Check if model has initialized gradients + has_any_grads = any(param.grad is not None for param in model.parameters()) + has_all_grads = all(param.grad is not None for param in model.parameters()) + + # Check context behavior + self.assertTrue(has_entered_grad_sync_context, 'Has not entered custom sync context') + self.assertTrue(has_exited_grad_sync_context, 'Has not exited custom sync context') + self.assertEqual( + has_any_grads, + has_all_grads, + 'Expected gradients to all be uninitialized or all be initialized', + ) + self.assertEqual( + has_all_grads, + parallel_state.is_pipeline_first_stage(), + 'Expected gradients to be initialized only in first pipeline stage', + ) + + # Clean up + parallel_state.destroy_model_parallel() + + @unittest.skipIf(torch.cuda.device_count() < 4 or torch.cuda.device_count() % 2 != 0, "Requires >= 4 GPUs") + def test_pipelining_with_interleaving_with_custom_sync_context_handler(self) -> None: + + # Parallel configuration + world_size = torch.cuda.device_count() + tensor_model_parallel_world_size = 1 + data_parallel_size = 2 if world_size > 4 else 1 + pipeline_model_parallel_world_size = world_size // data_parallel_size + virtual_pipeline_model_parallel_size = 2 + + # Initialize pipeline parallelism + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + virtual_pipeline_model_parallel_size_=virtual_pipeline_model_parallel_size, + ) + pp_utils._reconfigure_microbatch_calculator( + rank=parallel_state.get_tensor_model_parallel_rank(), + rampup_batch_size=None, + global_batch_size=self.GLOBAL_BATCH_SIZE, + micro_batch_size=self.MICRO_BATCH_SIZE, + data_parallel_size=parallel_state.get_data_parallel_world_size(), + ) + pp_utils.update_num_microbatches(0) + + # Construct synthetic data + dtype = get_dtype_for_comparison() + hidden_size = self.HIDDEN_SIZE + microbatch_size = self.MICRO_BATCH_SIZE + global_batch_shape = ( + self.GLOBAL_BATCH_SIZE + // parallel_state.get_data_parallel_world_size(), + hidden_size, + hidden_size, + ) + batch = None + if parallel_state.is_pipeline_first_stage(): + batch = (torch.ones(global_batch_shape, dtype=dtype).cuda(), ) + + # Construct model + model = build_model( + testing_utils.model_provider_func, + wrap_with_ddp=True, + virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, + hidden_size=hidden_size, + ) + for module in model: + module.to(dtype) + module.module.apply(get_init_weights_func(0)) + + # Construct context that keeps track whenever entered/exited + grad_sync_context_enter_count = 0 + grad_sync_context_exit_count = 0 + @contextlib.contextmanager + def custom_grad_sync_context(): + try: + nonlocal grad_sync_context_enter_count + grad_sync_context_enter_count += 1 + yield + finally: + nonlocal grad_sync_context_exit_count + grad_sync_context_exit_count += 1 + for module in model: + for param in module.parameters(): + param.grad = None + + # Training step with pipeline parallelism + loss = _forward_backward_pipelining_with_interleaving( + testing_utils.fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=(microbatch_size, hidden_size, hidden_size), + dtype=dtype, + async_comm=False, + grad_scaler=None, + deallocate_pipeline_outputs=False, + sequence_parallel_enabled=False, + custom_sync_context_handler=custom_grad_sync_context, + ) + torch.cuda.synchronize() + + # Check context behavior + self.assertTrue( + grad_sync_context_enter_count > 0, + 'Has not entered custom sync context', + ) + self.assertEqual( + grad_sync_context_enter_count, + grad_sync_context_exit_count, + 'Has not entered and exited custom sync context ' + 'the same number of times', + ) + self.assertEqual( + grad_sync_context_exit_count, + virtual_pipeline_model_parallel_size + 1, + 'Expected to exit custom sync context once per model chunk ' + 'and once at the function end', + ) + + # Clean up + parallel_state.destroy_model_parallel() + + if __name__ == "__main__": common_utils.run_tests() diff --git a/tests/L0/run_transformer/test_random.py b/tests/L0/run_transformer/test_random.py index 6060f9ed9..bd737ee5e 100644 --- a/tests/L0/run_transformer/test_random.py +++ b/tests/L0/run_transformer/test_random.py @@ -18,98 +18,95 @@ def test_set_cuda_rng_state(self): for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size - ) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size + ) - size, seed = 123, 1234 - torch.cuda.manual_seed(seed) - tensor = torch.cuda.FloatTensor(size) + size, seed = 123, 1234 + torch.cuda.manual_seed(seed) + tensor = torch.cuda.FloatTensor(size) - rng_state = torch.cuda.get_rng_state() - rng_state_clone = rng_state.clone() + rng_state = torch.cuda.get_rng_state() + rng_state_clone = rng_state.clone() - for _ in range(5): - torch.randn(size, out=tensor) - result_1 = tensor.clone() + for _ in range(5): + torch.randn(size, out=tensor) + result_1 = tensor.clone() - self.assertEqual(rng_state.sub(rng_state_clone).max(), 0) - self.assertGreater( - torch.cuda.get_rng_state().sub(rng_state_clone).max(), 0 - ) + self.assertEqual(rng_state.sub(rng_state_clone).max(), 0, msg=msg) + self.assertGreater( + torch.cuda.get_rng_state().sub(rng_state_clone).max(), 0, + msg=msg, + ) - new_rng_state = torch.cuda.get_rng_state() - self.assertGreater(new_rng_state.sub(rng_state).max(), 0) + new_rng_state = torch.cuda.get_rng_state() + self.assertGreater(new_rng_state.sub(rng_state).max(), 0, msg=msg) - tensor_parallel.random._set_cuda_rng_state(rng_state) - for _ in range(5): - torch.randn(size, out=tensor) - tensor_parallel.random._set_cuda_rng_state(rng_state) - for _ in range(5): - torch.randn(size, out=tensor) - result_2 = tensor.clone() + tensor_parallel.random._set_cuda_rng_state(rng_state) + for _ in range(5): + torch.randn(size, out=tensor) + tensor_parallel.random._set_cuda_rng_state(rng_state) + for _ in range(5): + torch.randn(size, out=tensor) + result_2 = tensor.clone() - self.assertEqual(result_2, result_1) + self.assertEqual(result_2, result_1, msg=msg) - self.assertEqual(rng_state.sub(rng_state_clone).max(), 0) + self.assertEqual(rng_state.sub(rng_state_clone).max(), 0, msg=msg) - parallel_state.destroy_model_parallel() + parallel_state.destroy_model_parallel() def test_cuda_rng_tracker(self): for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size - ) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size + ) - seed_1, seed_2, size = 1234, 4321, [12, 21] - tensor = torch.cuda.FloatTensor(size) + seed_1, seed_2, size = 1234, 4321, [12, 21] + tensor = torch.cuda.FloatTensor(size) - torch.cuda.manual_seed(seed_1) - torch.randn(size, out=tensor) - target_11 = tensor.clone() - torch.randn(size, out=tensor) - target_12 = tensor.clone() + torch.cuda.manual_seed(seed_1) + torch.randn(size, out=tensor) + target_11 = tensor.clone() + torch.randn(size, out=tensor) + target_12 = tensor.clone() - torch.cuda.manual_seed(seed_2) - torch.randn(size, out=tensor) - targt_21 = tensor.clone() - torch.randn(size, out=tensor) - target_22 = tensor.clone() + torch.cuda.manual_seed(seed_2) + torch.randn(size, out=tensor) + targt_21 = tensor.clone() + torch.randn(size, out=tensor) + target_22 = tensor.clone() - torch.cuda.manual_seed(seed_1) - tensor_parallel.random.get_cuda_rng_tracker().add("test", seed_2) + torch.cuda.manual_seed(seed_1) + tensor_parallel.random.get_cuda_rng_tracker().add("test", seed_2) + torch.randn(size, out=tensor) + result_11 = tensor.clone() + + with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): torch.randn(size, out=tensor) - result_11 = tensor.clone() + result_21 = tensor.clone() - with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): - torch.randn(size, out=tensor) - result_21 = tensor.clone() + torch.randn(size, out=tensor) + result_12 = tensor.clone() + with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): torch.randn(size, out=tensor) - result_12 = tensor.clone() - - with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): - torch.randn(size, out=tensor) - result_22 = tensor.clone() + result_22 = tensor.clone() - self.assertEqual(target_11, result_11) - self.assertEqual(target_12, result_12) - self.assertEqual(targt_21, result_21) - self.assertEqual(target_22, result_22) - self.assertNotEqual(result_11, result_21) - self.assertNotEqual(result_21, result_22) + self.assertEqual(target_11, result_11, msg=msg) + self.assertEqual(target_12, result_12, msg=msg) + self.assertEqual(targt_21, result_21, msg=msg) + self.assertEqual(target_22, result_22, msg=msg) + self.assertNotEqual(result_11, result_21, msg=msg) + self.assertNotEqual(result_21, result_22, msg=msg) - tensor_parallel.random.get_cuda_rng_tracker().reset() - parallel_state.destroy_model_parallel() + tensor_parallel.random.get_cuda_rng_tracker().reset() + parallel_state.destroy_model_parallel() class NcclTransformerRandomTest(TransformerRandomTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/test_transformer_module.py b/tests/L0/run_transformer/test_transformer_module.py deleted file mode 100644 index 77ce67edb..000000000 --- a/tests/L0/run_transformer/test_transformer_module.py +++ /dev/null @@ -1,107 +0,0 @@ -from typing import Tuple -import os -import subprocess -import sys -import unittest - - -SEVERALGPU_TEST = [ - "bert_minimal_test", - "gpt_minimal_test", - "dynamic_batchsize_test", -] - - -def get_multigpu_launch_option(min_gpu): - should_skip = False - import torch - - num_devices = torch.cuda.device_count() - if num_devices < min_gpu: - should_skip = True - distributed_run_options = f"-m torch.distributed.run --nproc_per_node={num_devices}" - return should_skip, distributed_run_options - - -def get_launch_option(test_filename) -> Tuple[bool, str]: - should_skip = False - for severalgpu_test in SEVERALGPU_TEST: - if severalgpu_test in test_filename: - return get_multigpu_launch_option(3) - return should_skip, "" - - -def run_transformer_tests(): - python_executable_path = sys.executable - directory = os.path.dirname(__file__) - files = [ - os.path.join(directory, f) - for f in os.listdir(directory) - if f.startswith("run_") and os.path.isfile(os.path.join(directory, f)) - ] - print("#######################################################") - print(f"# Python executable path: {python_executable_path}") - print(f"# {len(files)} tests: {files}") - print("#######################################################") - errors = [] - for i, test_file in enumerate(files, 1): - is_denied = False - should_skip, launch_option = get_launch_option(test_file) - if should_skip: - print( - f"### {i} / {len(files)}: {test_file} skipped. Requires multiple GPUs." - ) - continue - test_run_cmd = ( - f"{python_executable_path} {launch_option} {test_file} " - "--micro-batch-size 2 --num-layers 16 --hidden-size 256 --num-attention-heads 8 --max-position-embeddings " - "512 --seq-length 512 --global-batch-size 128" - ) - if "bert" in test_file or "gpt" in test_file: - import torch - - num_devices = torch.cuda.device_count() - if "bert" in test_file: - # "bert" uses the interleaving. - tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices > 4 else 1 - if "gpt" in test_file: - # "gpt" uses the non-interleaving. - tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices >= 4 else 1 - pipeline_model_parallel_size = num_devices // tensor_model_parallel_size - test_run_cmd += f" --pipeline-model-parallel-size {pipeline_model_parallel_size} --tensor-model-parallel-size {tensor_model_parallel_size}" - - if "bert" in test_file: - test_run_cmd += f" --bert-no-binary-head" - else: - test_run_cmd += f" --use-cpu-initialization" - print(f"### {i} / {len(files)}: cmd: {test_run_cmd}") - try: - output = ( - subprocess.check_output(test_run_cmd, shell=True) - .decode(sys.stdout.encoding) - .strip() - ) - except Exception as e: - errors.append((test_file, str(e))) - else: - if ">> passed the test :-)" not in output: - errors.append((test_file, output)) - else: - if not errors: - print("### PASSED") - else: - print("### FAILED") - short_msg = f"{len(errors)} out of {len(files)} tests failed" - print(short_msg) - for (filename, log) in errors: - print(f"File: {filename}\nLog: {log}") - raise RuntimeError(short_msg) - - -class TestTransformer(unittest.TestCase): - def test_transformer(self): - run_transformer_tests() - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/L0/run_transformer/test_transformer_utils.py b/tests/L0/run_transformer/test_transformer_utils.py index d5d16080d..4164eca87 100644 --- a/tests/L0/run_transformer/test_transformer_utils.py +++ b/tests/L0/run_transformer/test_transformer_utils.py @@ -17,23 +17,23 @@ def test_split_tensor_along_last_dim(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) - - device = "cpu" - input_tensor = torch.randn((100, 100, 100), device=device) - splits = utils.split_tensor_along_last_dim(input_tensor, 10) - last_dim_shapes = torch.tensor( - [int(split.size()[-1]) for split in splits] - ) - - self.assertTrue(torch.equal(last_dim_shapes, torch.full((10,), 10),)) - - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) + + device = "cpu" + input_tensor = torch.randn((100, 100, 100), device=device) + splits = utils.split_tensor_along_last_dim(input_tensor, 10) + last_dim_shapes = torch.tensor( + [int(split.size()[-1]) for split in splits] + ) + + self.assertTrue( + torch.equal(last_dim_shapes, torch.full((10,), 10),), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}", + ) + + parallel_state.destroy_model_parallel() if __name__ == "__main__":