[Cherry-pick] Cherry-pick from fleety (PaddlePaddle#11047)

* add timer log in trainer (PaddlePaddle#10880)

* add layer norm backward (PaddlePaddle#10886)

* add memory usage message in tensorboard (PaddlePaddle#10887)

Author: ForFishes

paddlenlp/trainer/trainer.py

@@ -39,6 +39,7 @@
 import paddle.amp.auto_cast as autocast
 import paddle.distributed as dist
 import paddle.nn as nn
+import psutil
 from packaging import version
 from paddle import framework
 from paddle.distributed.fleet.meta_parallel import PipelineLayer
@@ -3204,6 +3205,14 @@ def log(self, logs: Dict[str, float], **kwargs) -> None:
 
         if self.state.epoch is not None:
             logs["progress_or_epoch"] = round(self.state.epoch, 4)
+
+        if self.timers:
+            logs.update(self.timers.info(self.timers.timers.keys()))
+
+        mem_info = psutil.virtual_memory()
+        logs["cpu_used_memory"] = round(mem_info.used / (1024**3), 2)
+        logs["cpu_available_memory"] = round(mem_info.available / (1024**3), 2)
+
         self.state.log_history = []
         self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs, **kwargs)
 

slm/model_zoo/gpt-3/external_ops/fused_ln/layer_norm_cuda.cu

@@ -237,5 +237,16 @@ PD_BUILD_GRAD_OP(fused_rms_norm)
 #endif
     ;
 
+PD_BUILD_OP(fused_rms_norm_grad_func)
+    .Inputs({"x", "scale", "invvar", "dy"})
+    .Outputs({"dx", "d_scale"})
+    .Attrs({"epsilon: float"})
+    .SetKernelFn(PD_KERNEL(RMSLnBwd))
+    .SetInferShapeFn(PD_INFER_SHAPE(RMSLnBwdInferShape))
+#ifdef CUSTOM_OP_WITH_SPMD
+    .SetInferSpmdFn(PD_INFER_SPMD_RULE(phi::distributed::RmsNormGradInferSpmd))
+#endif
+    ;
+
 
 // https://github.com/NVIDIA/apex/blob/85e9eddece9d4ac72b48c2407f8162f2173e1bf4/csrc/layer_norm_cuda_kernel.cu#L679