Add Non ZCC EMA callback (#2923)

Author: sneaxiy

paddleformers/trainer/trainer.py

@@ -195,12 +195,13 @@
 
 try:
     from .utils.zero_cost_checkpoint import (
+        NonZCCEMACallback,
         ZeroCostCheckpointCallback,
         ZeroCostCheckpointManager,
         get_fused_param_mappings,
     )
 except (ImportError, ModuleNotFoundError):
-    ZeroCostCheckpointManager, get_fused_param_mappings = None, None
+    ZeroCostCheckpointManager, NonZCCEMACallback, get_fused_param_mappings = None, None, None
 from .utils.helper import (  # nested_truncate,
     broadcast_dataset_rank0_model,
     broadcast_dp_optimizer,
@@ -873,6 +874,9 @@ def create_zcc_manager(self, unwrapped_model, resume_from_checkpoint=None):
 
         logger.info("Create zero cost checkpoint manager done.")
 
+    def add_non_zcc_ema_callback(self, resume_from_checkpoint):
+        self.add_callback(NonZCCEMACallback(resume_from_checkpoint, self.args, self.sharding_io))
+
     def _save_flex_model_state(self, output_dir):
         model_sharded_state_dict = self.model.sharded_state_dict()
         model_state_dict_path = os.path.join(output_dir, MODEL_STATE_DIC)
@@ -1135,6 +1139,8 @@ def train(
 
         if self.args.enable_zero_cost_checkpoint:
             self.create_zcc_manager(model, resume_from_checkpoint)
+        elif self.args.zcc_save_ema_coef is not None:
+            self.add_non_zcc_ema_callback(resume_from_checkpoint)
 
         logger.info(f"{self.runtime_timer.log()}")
         logger.info("***** Running training *****")
@@ -1365,6 +1371,16 @@ def _inner_training_loop(
                         self._skip_steps_since_last_logged += 1
 
                         self.state.epoch = epoch + (step + 1) / steps_in_epoch
+
+                        # For ZCC EMA
+                        if self.args.enable_zero_cost_checkpoint or self.args.zcc_save_ema_coef is not None:
+                            tr_loss_for_zcc = tr_loss.clone()
+                            dist.all_reduce(
+                                tr_loss_for_zcc, dist.ReduceOp.SUM
+                            )  # 3级并行时，每个pp下的loss会广播，全局reduce-mean的时候，分子分母都会乘以pp_world_size，结果会被约掉
+                            tr_loss_for_zcc_scalar = tr_loss_for_zcc.item() / dist.get_world_size()
+                            self.state.loss = tr_loss_for_zcc_scalar
+
                         self.state.consumed_samples = (
                             self.state.global_step
                             * args.per_device_train_batch_size

paddleformers/trainer/training_args.py

@@ -1090,6 +1090,10 @@ class TrainingArguments:
         default=1,
         metadata={"help": "Interval between updating EMA parameters."},
     )
+    zcc_ema_loss_threshold: Optional[float] = field(
+        default=None,
+        metadata={"help": "If set not None, only do EMA when the training loss is smaller than the threshold value"},
+    )
     save_tokenizer: Optional[bool] = field(
         default=True,
         metadata={"help": "Save tokenizer to output_dir."},
@@ -2099,7 +2103,7 @@ def is_context_parallel_supported():
         assert (
             self.save_steps % self.zcc_ema_interval == 0
         ), f"save_steps[{self.save_steps}] must be divisible by zcc_ema_interval[{self.zcc_ema_interval}]"
-        if self.zcc_save_ema_coef is not None:
+        if self.enable_zero_cost_checkpoint and self.zcc_save_ema_coef is not None:
             assert (
                 self.zcc_workers_num == 1
             ), "EMA function in zero cost checkpoint mode does not support zcc_workers_num > 1 for now."

paddleformers/trainer/utils/zero_cost_checkpoint.py

@@ -173,10 +173,11 @@ def ema_reset(self):
         self.ema_buffer_modele_params = None
 
     @imperative_base.no_grad()
-    def ema_accumulate(self):
+    def ema_accumulate(self, global_step, loss, zcc_ema_loss_threshold):
         """
         perform ema update : ` \alpha * EMA + (1-\alpha) + model`
         build `self.ema_buffer` if necessary
+        when loss < threshold, do ema update
         """
         # logger.info(f'[ZCC EMA] wait all done, doing EMA w/ coef: {self.ema_coef}, status:{self.status()}')
         # do update: ema = alpha * ema + (1-alpha) * model
@@ -185,14 +186,19 @@ def ema_accumulate(self):
             cpu_master_weights = self.optimizer_fusion_storage_helper.cpu_buffer._slice(
                 self.master_min_offset, self.master_max_offset
             ).cpu()
-            self.ema_buffer = self.ema_coef * self.ema_buffer + (1 - self.ema_coef) * cpu_master_weights
-            # logger.info(f'[ZCC EMA2] wait all done, doing EMA w/ coef: {self.ema_coef}, status:{self.status()}')
-            for index, ema_buf in self.ema_buffer_model_params.items():
-                _, cpu_buf = self.param_fusion_storage_helper.inited_buffers[index]
-                updated_ema = self.ema_coef * ema_buf + (1 - self.ema_coef) * cpu_buf.cpu()
-                self.ema_buffer_model_params[index] = updated_ema
-
-        logger.info(f"[ZCC EMA] accumulating, buffer type:{self.ema_buffer.place} {self.ema_buffer.dtype}, done")
+            if zcc_ema_loss_threshold is None or loss < zcc_ema_loss_threshold:
+                self.ema_buffer = self.ema_coef * self.ema_buffer + (1 - self.ema_coef) * cpu_master_weights
+                for index, ema_buf in self.ema_buffer_model_params.items():
+                    _, cpu_buf = self.param_fusion_storage_helper.inited_buffers[index]
+                    updated_ema = self.ema_coef * ema_buf + (1 - self.ema_coef) * cpu_buf
+                    self.ema_buffer_model_params[index] = updated_ema
+                logger.info(
+                    f"[ZCC EMA] accmulating, buffer type:{self.ema_buffer.place} {self.ema_buffer.dtype}, done"
+                )
+            else:
+                logger.info(
+                    f"[ZCC EMA] accmulating SKIP for global_step:{global_step}, because loss:{loss} > threshold:{zcc_ema_loss_threshold}"
+                )
 
     @imperative_base.no_grad()
     def ema_state_dict(self):
@@ -790,7 +796,11 @@ def process_offload_task(self, dump, global_step):
             self.global_step.value = global_step
 
             if self.ema_coef is not None:
-                self.zcc_ema_processor.ema_accumulate()
+                self.zcc_ema_processor.ema_accumulate(
+                    self.trainer_state.global_step,
+                    self.trainer_state.loss,
+                    self.training_args_content.zcc_ema_loss_threshold,
+                )
 
         # continue to process dumping task at the last chunk
         if self.offloaded_numels == self.all_numel:
@@ -1006,3 +1016,86 @@ def manage_offload_chunk(self):
         logger.info(
             f"[ZCC Worker{self.worker_id}] All numel: {self.all_numel}, Offload chunks: {self.offload_chunks}, Chunk size: {self.chunk_size_in_numel}]"
         )
+
+
+class EMABuffer:
+    def __init__(self, resume_from_checkpoint, args, sharding_io, offload=True):
+        assert sharding_io is not None, "EMA should be only enabled when save_sharded_model is True"
+        self.master_weights = {}
+        self.model_params = {}
+        self.args = args
+        self.sharding_io = sharding_io
+        self.offload = offload
+        if resume_from_checkpoint is not None:
+            self._load(resume_from_checkpoint)
+
+    def _ema_path(self, base_path):
+        path = _add_variant(PADDLE_OPTIMIZER_NAME, self.args.optimizer_name_suffix)
+        path = path.replace("optimizer", "ema")
+        return os.path.join(base_path, path)
+
+    def _load(self, resume_from_checkpoint):
+        ema_path = self._ema_path(resume_from_checkpoint)
+        if not os.path.exists(ema_path):
+            return
+
+        logger.info(f"Loading EMA checkpoint from {resume_from_checkpoint} ...")
+        with device_guard("cpu"):
+            ema_state_dict = paddle.load(ema_path)
+        logger.info(f"Load EMA checkpoint from {resume_from_checkpoint} done")
+
+        self.master_weights = ema_state_dict.pop("master_weights")
+        self.model_params = ema_state_dict
+
+    def save(self, global_step):
+        base_path = os.path.join(self.args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{global_step}")
+        ema_path = self._ema_path(base_path)
+        ema_state_dict = {"master_weights": self.master_weights}
+        ema_state_dict.update(self.model_params)
+        os.makedirs(base_path, exist_ok=True)
+        logger.info(f"Saving EMA checkpoint to {base_path} ...")
+        paddle.save(ema_state_dict, ema_path)
+        logger.info(f"Save EMA checkpoint to {base_path} done")
+
+    def ema_accumulate(self, global_step, loss, ema_loss_threshold):
+        if ema_loss_threshold is None or loss < ema_loss_threshold:
+            logger.info(f"EMA accumulating for step {global_step} ...")
+            self._ema_impl(
+                state_dict=self.sharding_io.optimizer.state_dict()["master_weights"],
+                ema_state_dict=self.master_weights,
+            )
+            self._ema_impl(
+                state_dict=self.sharding_io.manipulate_state_dict_and_config(
+                    unwrap_model(self.sharding_io.model),
+                    merge_tensor_parallel=False,
+                )[0],
+                ema_state_dict=self.model_params,
+            )
+            logger.info(f"EMA accumulate done for step {global_step}")
+
+    def _ema_impl(self, state_dict, ema_state_dict):
+        ema_coef = self.args.zcc_save_ema_coef
+        for k, v in state_dict.items():
+            if k in ema_state_dict:
+                ema_tensor = ema_state_dict[k]
+                ema_tensor = ema_coef * ema_tensor.cuda() + (1 - ema_coef) * v.cuda()
+                ema_tensor.name = v.name
+                v = ema_tensor
+                del ema_tensor
+
+            if self.offload:
+                v_pin = v.pin_memory()
+                v_pin.name = v.name
+                v = v_pin
+            ema_state_dict[k] = v
+
+
+class NonZCCEMACallback(TrainerCallback):
+    def __init__(self, resume_from_checkpoint, args, sharding_io, offload=True):
+        self.buffer = EMABuffer(resume_from_checkpoint, args, sharding_io, offload)
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if state.global_step % args.zcc_ema_interval == 0:
+            self.buffer.ema_accumulate(state.global_step, state.loss, args.zcc_ema_loss_threshold)
+        if control.should_save:
+            self.buffer.save(state.global_step)