merge: Resolve merge conflicts pulling in from Microsoft upstream

megatron/arguments.py

@@ -1627,5 +1627,11 @@ def _add_profiler_args(parser):
      type=str,
      default='2,3',
      help="Which steps to profile. Format: <start step>,<end step>")
+
+    group.add_argument("--profile-ranks",
+     type=int,
+     nargs='+',
+     default=None,
+     help="Which ranks to profile. Format: 0 1 2 3")
 
     return parser

megatron/core/tensor_parallel/layers.py

@@ -286,6 +286,7 @@ def forward(ctx, input, weight, bias, gradient_accumulation_fusion,
     @staticmethod
     @custom_bwd
     def backward(ctx, grad_output):
+        args = get_args()
         input, weight = ctx.saved_tensors
         use_bias = ctx.use_bias
 
@@ -367,9 +368,13 @@ def backward(ctx, grad_output):
         #     grad_weight = None
         # else:
         #     grad_weight = grad_output.t().matmul(total_input)
-        from megatron.core.tensor_parallel.weight_grad_store import WeightGradStore
-        WeightGradStore.put(total_input, grad_output, weight, gradientUpdateFunction)
-        grad_weight = None
+        if args.enable_zbh1_pipeline:
+            from megatron.core.tensor_parallel.weight_grad_store import WeightGradStore
+            WeightGradStore.put(total_input, grad_output, weight, gradientUpdateFunction)
+            grad_weight = None
+        else:
+            grad_weight = grad_output.t().matmul(total_input)
+
         grad_bias = grad_output.sum(dim=0) if use_bias else None
 
         if ctx.sequence_parallel:

megatron/model/fused_rmsnorm.py

@@ -1,16 +1,24 @@
+from megatron import get_args
+
 import torch
 from torch.nn.parameter import Parameter
+from torch.nn import init
 import intel_extension_for_pytorch as ipex  # noqa
 
 
 # Taken from facebookresearch/llama
 class RMSNorm(torch.nn.Module):
-    # def __init__(self, dim: int, eps: float = 1e-6):
-    def __init__(self, dim: int, eps: float = 1e-6, sequence_parallel: bool = False):
+    def __init__(
+        self, dim: int, eps: float = 1e-6, sequence_parallel: bool = False
+    ):
         super().__init__()
         self.eps = eps
-        self.weight = Parameter(torch.ones(dim))
+        self.weight = Parameter(torch.ones(dim, dtype=get_args().params_dtype))
+        self.sequence_parallel = sequence_parallel
+        setattr(self.weight, "sequence_parallel", self.sequence_parallel)
 
     def forward(self, x):
-        output = torch.xpu.IpexRmsNorm(x, self.weight.shape, self.weight, self.eps)
+        output = torch.xpu.IpexRmsNorm(
+            x, self.weight.shape, self.weight, self.eps
+        )
         return output

megatron/model/rmsnorm.py

@@ -10,7 +10,9 @@
 
 # Taken from facebookresearch/llama
 class RMSNorm(torch.nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6, sequence_parallel: bool = False):
+    def __init__(
+        self, dim: int, eps: float = 1e-6, sequence_parallel: bool = False
+    ):
         super().__init__()
         self.eps = eps
         init_device = None
@@ -20,7 +22,8 @@ def __init__(self, dim: int, eps: float = 1e-6, sequence_parallel: bool = False)
             torch.empty(dim, device=init_device, dtype=get_args().params_dtype)
         )
         init.ones_(self.weight)
-        setattr(self.weight, "sequence_parallel", sequence_parallel)
+        self.sequence_parallel = sequence_parallel
+        setattr(self.weight, "sequence_parallel", self.sequence_parallel)
 
     def _norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

megatron/profiler.py

@@ -36,7 +36,9 @@ def is_end_step():
     def is_capture_step():
         return cur_step >= start_step and cur_step <= end_step
 
-    if args.profile.startswith('pt'):
+    if args.profile.startswith('pt') and (
+        args.profile_ranks is None or torch.distributed.get_rank() in args.profile_ranks
+    ):
         schedule = torch.profiler.schedule(wait=0, warmup=0, active=active_steps, repeat=1)
         activities = [torch.profiler.ProfilerActivity.CPU]
         activities.extend([torch.profiler.ProfilerActivity.HPU] if device.startswith("hpu") else [])