[transformer] Port Sequence Parallelism (takeover of NVIDIA#1396) (NV…

…IDIA#1400) * it looks possible to remove this file * add communication collectives * update Column|RowParallelLinear * update checkpoint function * update function name * parity between public and private collectives * row parallel linear * column parallel linear * sequence parallel: p2p comm fix typo * sequence parallel: pipeline parallel * fix typo * add layernorm with sequence_parallel_enabled attr * class variable -> member variable * fix col parallel test with sequence parallel * Initial test of `forward_backward_pipelining_without_interleaving` with `model_type=ModelType.encoder_and_decoder` * add cases pretending to test sequence_parallel * Apply 2 suggestion(s) to 1 file(s) * update sequence_parallel_enabled docstring * update docstring: order of tensor dimensions, sequence_parallel_enabled behavior * Divide sequence_length if sequence parallel tensor shape should be updated if sequence parallel is enabled. * cherry-pick NVIDIA/Megatron-LM@8474e6e * type annotation * Fix matmul call in RowParallelLinear Fix `sequence_parallel_enabled` to `False` as you can see in https://github.com/NVIDIA/Megatron-LM/blob/d898a8991d1a08d29074f87819d1bf41517e35f5/megatron/mpu/layers.py#L511-L514 * update rowparallellinear test * fix `loss_weight` is not defined in test_layers * @eqy's comment * mixed fused layer norm * fix typo * misc * test_layers cleanup * Skip Bert/GPT script Since these two models haven't gotten updated for sequence parallle, e.g. the update of the order of dimension from (batch, sequence, feature) to (sequence, batch, feature) and global variables of arguments * debug part 1/N: comment out `x.retain_grad` * debug part 2/N: [ColumnParallelLinear] comment out overriding of sequence_parallel_enabled * debug 3/N: add pipeline test with parallel mlp * Fix handling `self.input_tensor` and argument * tp2pp4 ModelType.encoder_or_decoder is failing, which can be at my fault because the backward is blaming the output and the grad_ouptut shape don't match * revert debug 1/N * defer tensor model parallel size > 1 * split tensor in sequence dim * cosmetic * cosmetic: remove archaic comment * enable TP>1 for encoder_and_decoder as well * set requires_grad=True always... * Set `scatter_gather_tensors_in_pipeline` to :obj:`False` for the sake of nemo megatron's GPT works with sequence parallel enabled. * brush up comment of `requires_grad()` There's a possibility that PyTorch DistributedDataParallel hangs when some tensor (or parameter) doesn't require grad according to @ptrblck. This forced `requires_grad` in my understanding is different from that. * misc changes of scatter_gather_tensors_in_pipeline comment * guard for torch_ucc * cosmetic changes related to tests * update command line arguments * update TransformerLanguageModel * rename * move gpt to gpt.py * update bert * add all_gather for params in sequence parallel region * misc. some diffs were lost during rebasing... * updates for non sequence parallel execution * gpt with sequence parallel * Apply 2 suggestion(s) to 2 file(s) * update tensor&pipeline parallel size * why `sequence_parallel_enabled` is not supplied!? Did I messed up when rebasing? * cosmetic fix * correct key is sequence_parallel_enabled
Vahiid65 · Jun 23, 2022 · 3ff1a10 · 3ff1a10
1 parent 57f890a
commit 3ff1a10
Show file tree

Hide file tree

Showing 22 changed files with 3,383 additions and 1,862 deletions.
diff --git a/apex/transformer/layers/__init__.py b/apex/transformer/layers/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+from apex.transformer.layers.layer_norm import FastLayerNorm
+from apex.transformer.layers.layer_norm import FusedLayerNorm
+from apex.transformer.layers.layer_norm import MixedFusedLayerNorm
+
+
+__all__ = [
+    "FastLayerNorm",
+    "FusedLayerNorm",
+    "MixedFusedLayerNorm",
+]
diff --git a/apex/transformer/layers/layer_norm.py b/apex/transformer/layers/layer_norm.py
@@ -0,0 +1,99 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+# NOTE(mkozuki): This file defines two LayerNorm that are compatible with Megatron-LM.
+# while avoiding introducing the breaking change of `"sequence_parallel_enabled"` attribute into apex.normalization.FusedLayerNorm
+# and apex.contrib.layer_norm.FastLayerNorm.
+import warnings
+
+import torch
+
+from apex.normalization import FusedLayerNorm as OrigFusedLayerNorm
+from apex.normalization import MixedFusedLayerNorm as OrigMixedFusedLayerNorm
+try:
+    from apex.contrib.layer_norm import FastLayerNorm as OrigFastLayerNorm
+except ImportError:
+    HAS_FAST_LAYER_NORM = False
+else:
+    HAS_FAST_LAYER_NORM = True
+
+
+__all__ = [
+    "FusedLayerNorm",
+    "FastLayerNorm",
+    "MixedFusedLayerNorm",
+]
+
+
+def _set_sequence_parallel_enabled(
+    param: torch.Tensor,
+    sequence_parallel_enabled: bool,
+) -> None:
+    setattr(param, "sequence_parallel_enabled", sequence_parallel_enabled)
+
+
+class FusedLayerNorm(OrigFusedLayerNorm):
+    def __init__(
+        self,
+        normalized_shape,
+        eps: float = 1e-5,
+        elementwise_affine: bool = True,
+        *,
+        sequence_parallel_enabled: bool = False,
+    ):
+        super().__init__(
+            normalized_shape=normalized_shape,
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+        )
+        self.sequence_parallel_enabled = sequence_parallel_enabled
+        if self.elementwise_affine:
+            _set_sequence_parallel_enabled(self.weight, self.sequence_parallel_enabled)
+            _set_sequence_parallel_enabled(self.bias, self.sequence_parallel_enabled)
+
+
+# note: MixedFusedLayerNorm is no different from FusedLayerNorm if it's used in `torch.cuda.amp`.
+class MixedFusedLayerNorm(OrigMixedFusedLayerNorm):
+    def __init__(
+        self,
+        normalized_shape,
+        eps: float = 1e-5,
+        **kwargs,
+    ) -> None:
+        self.sequence_parallel_enabled = kwargs.get("sequence_parallel_enabled", False)
+        super().__init__(normalized_shape=normalized_shape, eps=eps, **kwargs)
+        if self.sequence_parallel_enabled:
+            _set_sequence_parallel_enabled(self.weight, self.sequence_parallel_enabled)
+            _set_sequence_parallel_enabled(self.bias, self.sequence_parallel_enabled)
+
+
+if HAS_FAST_LAYER_NORM:
+    class FastLayerNorm(OrigFastLayerNorm):
+        def __init__(
+            self,
+            hidden_size,
+            eps: float = 1e-5,
+            *,
+            sequence_parallel_enabled: bool = False,
+        ):
+            super().__init__(
+                hidden_size=hidden_size,
+                eps=eps
+            )
+            self.sequence_parallel_enabled = sequence_parallel_enabled
+            _set_sequence_parallel_enabled(self.weight, self.sequence_parallel_enabled)
+            _set_sequence_parallel_enabled(self.bias, self.sequence_parallel_enabled)
+else:
+    class FastLayerNorm(FusedLayerNorm):
+        def __init__(
+            self,
+            hidden_size,
+            eps: float = 1e-5,
+            *,
+            sequence_parallel_enabled: bool = False,
+        ):
+            warnings.warn("`apex.contrib.layer_norm.FastLayerNorm` isn't available thus falling back to `apex.normalization.FusedLayerNorm`")
+            super().__init__(
+                normalized_shape=hidden_size,
+                eps=eps,
+                elementwise_affine=True,
+                sequence_parallel_enabled=sequence_parallel_enabled,
+            )