Sync layer norm

.github/workflows/main.yml

@@ -186,7 +186,9 @@ jobs:
           pip install pytest-timeout
 
       - name: Run tests
-        run: pytest --timeout=600 tests
+        # run: pytest --timeout=600 tests
+        # run just the test we want for now
+        run: pytest --timeout=600 tests/test_training.py::MegDSTestTraining::test_layer_norm_consistent_0_bf16
 
   stop-runner:
     name: Stop self-hosted EC2 runner

megatron/arguments.py

@@ -738,6 +738,7 @@ def _add_distributed_args(parser):
     group.add_argument('--use-cpu-initialization', action='store_true',
                        default=None, help='If set, affine parallel weights '
                        'initialization uses CPU' )
+    group.add_argument('--force-sync-layer-norm-parameters', action="store_true")
     return parser
 
 

megatron/checkpointing.py

@@ -169,6 +169,69 @@ def save_checkpoint(iteration, model, optimizer, lr_scheduler):
         # Trim off the filename and mp_rank_* directory.
         for _ in range(3):
             checkpoint_name = os.path.dirname(checkpoint_name)
+
+            # Debug
+            layer_norms_params_end_with = [
+                "word_embeddings.norm.weight", "word_embeddings.norm.bias",
+                "input_layernorm.weight", "input_layernorm.bias",
+                "post_attention_layernorm.weight", "post_attention_layernorm.bias",
+                "self_attention.dense.bias", "mlp.dense_4h_to_h.bias",
+            ]
+            for n,p in model[0].named_parameters():
+                # Here is how you can access fp32 version of the bf16 param and fp32 optim states
+                #
+                # Note that there is an all_reduce called on all dp ranks when `get_full_hp_param` is called -
+                # so it's not free
+                #
+                # a. fp32 param
+                for end in layer_norms_params_end_with:
+                    if n.endswith(end):
+                        fp32_param = p.get_full_hp_param()
+
+                        fp32_params_acculumator = [
+                            torch.zeros_like(fp32_param)
+                            for _ in range(mpu.get_tensor_model_parallel_world_size())
+                        ]
+                        torch.distributed.gather(
+                            fp32_param,
+                            fp32_params_acculumator,
+                            dst=0,
+                            group=mpu.get_tensor_model_parallel_group()
+                        )
+                        if mpu.get_tensor_model_parallel_rank() == 0:
+                            square = torch.tensor([
+                                [
+                                    torch.max(torch.abs(c1 - c2))
+                                    for c2 in fp32_params_acculumator
+                                ] for c1 in fp32_params_acculumator
+                            ])
+                            print(f"Parameter name = {n}")
+                            print(square)
+
+                        # b. fp32 optim states
+                        for key in ['exp_avg', 'exp_avg_sq']:
+                            full_optim_state = p.get_full_hp_param(optim_state_key=key)
+
+                            full_optim_state_acculumator = [
+                                torch.zeros_like(fp32_param)
+                                for _ in range(mpu.get_tensor_model_parallel_world_size())
+                            ]
+                            torch.distributed.gather(
+                                full_optim_state,
+                                full_optim_state_acculumator,
+                                dst=0,
+                                group=mpu.get_tensor_model_parallel_group()
+                            )
+                            if mpu.get_tensor_model_parallel_rank() == 0:
+                                square = torch.tensor([
+                                    [
+                                        torch.max(torch.abs(c1 - c2))
+                                        for c2 in full_optim_state_acculumator
+                                    ] for c1 in full_optim_state_acculumator
+                                ])
+                                print(f"Optimizer state: parameter name = {n}, key = {key}")
+                                print(square)
+
         model[0].save_checkpoint(checkpoint_name, client_state=state_dict)
 
     # Wait so everyone is done (necessary)

megatron/data/data_samplers.py

@@ -52,6 +52,7 @@ def build_pretraining_data_loader(dataset, consumed_samples):
     return torch.utils.data.DataLoader(dataset,
                                        batch_sampler=batch_sampler,
                                        num_workers=args.num_workers,
+                                       generator=torch.Generator().manual_seed(args.seed),
                                        pin_memory=True)
 
 class MegatronPretrainingSampler:

megatron/model/fused_layer_norm.py

@@ -19,10 +19,11 @@
 
 import numbers
 import torch
-from megatron import mpu
+from megatron import mpu, get_args
 from torch.nn.parameter import Parameter
 from torch.nn import init
 import importlib
+from megatron import mpu
 
 global fused_mix_prec_layer_norm_cuda
 fused_mix_prec_layer_norm_cuda = None
@@ -63,6 +64,7 @@ class MixedFusedLayerNorm(torch.nn.Module):
 
   def __init__(self, normalized_shape, eps=1e-5):
         super(MixedFusedLayerNorm, self).__init__()
+        args = get_args()
 
         global fused_mix_prec_layer_norm_cuda
         fused_mix_prec_layer_norm_cuda = importlib.import_module(
@@ -75,6 +77,7 @@ def __init__(self, normalized_shape, eps=1e-5):
         self.weight = Parameter(torch.Tensor(*normalized_shape))
         self.bias = Parameter(torch.Tensor(*normalized_shape))
         self.reset_parameters()
+        self.force_sync_layer_norm_parameters = args.force_sync_layer_norm_parameters
 
 
   def reset_parameters(self):
@@ -84,19 +87,17 @@ def reset_parameters(self):
 
 
   def forward(self, input):
-    weights = [torch.empty_like(self.weight) for tp in range(mpu.get_tensor_model_parallel_world_size())]
-    torch.distributed.all_gather(weights, self.weight, group=mpu.get_tensor_model_parallel_group())
-    biases = [torch.empty_like(self.bias) for tp in range(mpu.get_tensor_model_parallel_world_size())]
-    torch.distributed.all_gather(biases, self.bias, group=mpu.get_tensor_model_parallel_group())
-    if any(torch.any(weight != self.weight) for weight in weights):
-        if mpu.get_tensor_model_parallel_rank() == 0:
-            print("Weight sync failed")
-            print(weights)
-    if any(torch.any(bias != self.bias) for bias in biases):
-        if mpu.get_tensor_model_parallel_rank() == 0:
-            print("Bias sync failed")
-            print(biases)
+    if self.force_sync_layer_norm_parameters:
+        tp_world_size = mpu.get_tensor_model_parallel_world_size()
+        # TODO: hack in order to synchronize all layer norms despite them being unsynched
+        weight = torch.clone(self.weight)
+        bias = torch.clone(self.bias)
+        weight = mpu.reduce_from_tensor_model_parallel_region(weight) / tp_world_size
+        bias = mpu.reduce_from_tensor_model_parallel_region(bias) / tp_world_size
+    else:
+        weight = self.weight
+        bias = self.bias
 
     return FusedLayerNormAffineFunction.apply(
-      input, self.weight, self.bias, self.normalized_shape,self.eps)
+      input, weight, bias, self.normalized_shape,self.eps)
 

megatron/mpu/layers.py

@@ -241,7 +241,7 @@ def forward(self, input_):
                                       self.sparse)
         # Mask the output embedding.
         if self.tensor_model_parallel_size > 1:
-            output_parallel[input_mask, :] = 0.0
+            output_parallel = output_parallel.masked_fill(input_mask[..., None], 0.0)
         # Reduce across all the model parallel GPUs.
         output = reduce_from_tensor_model_parallel_region(output_parallel)
 

megatron/testing_utils.py

@@ -232,9 +232,9 @@ def get_gpu_count():
         return 0
 
 def torch_assert_equal(actual, expected, **kwargs):
-    # assert_equal was added around pt-1.9, it does better checks - e.g will check dimensions match
-    if hasattr(torch.testing, "assert_equal"):
-        return torch.testing.assert_equal(actual, expected, **kwargs)
+    # assert_close was added around pt-1.9, it does better checks - e.g will check dimensions match
+    if hasattr(torch.testing, "assert_close"):
+        return torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, **kwargs)
     else:
         return torch.allclose(actual, expected, rtol=0.0, atol=0.0)
 
@@ -886,4 +886,4 @@ def flatten_arguments(args):
 
     Example: {"arg1": "value1", "arg2": "value2"} -> ["IGNORED", "arg1", "value1", "arg2", "value2"]
     """
-    return ["IGNORED"] + [item for key_value in args.items() for item in key_value if item != ""]
+    return ["IGNORED"] + [item for key_value in args.items() for item in key_value if item != ""]

requirements.txt

@@ -6,9 +6,10 @@ pybind11
 regex
 six
 tensorboard
-torch>=1.7
+torch>=1.11
 transformers
-DeepSpeed @ git+https://github.com/microsoft/DeepSpeed.git
+# for now using this branch for bf16 work
+DeepSpeed @ git+https://github.com/microsoft/DeepSpeed.git@olruwase/bf16-updates
 # versions from HF transformers
 black==21.4b0
 isort>=5.5.4