Cherry-pick useful PRs from fleety (#11051)

* [LLM] fix normal sl (#10609)

* [LLM] support disable monkey patch (#10617)

* [LLM] use safer grad sync method when enabling sp (#10714)

* [RL] disable aistudio download and fix qwen bug (#10819)

* [RL] fix qwen load when fuse is enabled and modify gate presion to fp32 (#10842)

* 修复qwen3moe的一系列报错(justin) (#10818)

* [LLM] zcc support rng states (#10430) (#10485)

* [DLTP-85730] optimize save cost for zcc

* [LLM] fix moe using on tensor parallelism

* cherry-pick useful PRs from fleety and fix some typos

---------

Co-authored-by: aiyinyuedejustin <[email protected]>

Author: SylarTiaNII

paddlenlp/trainer/trainer.py

@@ -140,6 +140,7 @@
     DefaultFlowCallback,
     PrinterCallback,
     ProgressCallback,
+    SPGradSyncCallback,
     TrainerCallback,
     TrainerControl,
     TrainerState,
@@ -444,9 +445,8 @@ def _save_ckpt_func(state_dict, path, signal_path=None):
             ), "should_save_sharding_stage1_model should be True when using zero cost checkpoint"
             assert (
                 ShardingOption.FULL_SHARD not in self.args.sharding
-            ), "FULL_SHARD is not supported when using zero cost checkpoint"
-            assert not self.args.save_tokenizer, "save_tokenizer is not supported when using zero cost checkpoint"
-            assert not self.args.save_rng_states, "save_rng_states is not supported when using zero cost checkpoint"
+            ), "FULL_SHARD is not supported when using flash save mode"
+            assert not self.args.save_tokenizer, "save_tokenizer is not supported when using flash save mode"
 
             # init attributes for zero cost checkpoint mode
             self.zcc_manager = None
@@ -2021,34 +2021,18 @@ def _load_rng_state(self, checkpoint):
         if checkpoint is None:
             return
 
-        # if use distributed training
-        if self.args.world_size > 1:
-            process_index = self.args.process_index
-            rng_file_list = [None for x in range(self.args.world_size)]
-            if self.args.should_save:
-                rng_file = os.path.join(checkpoint, f"rng_state_{self.args.world_size}.pth")
-                if os.path.isfile(rng_file):
-                    rng_file_list = paddle.load(rng_file, return_numpy=True)
-            paddle.distributed.broadcast_object_list(rng_file_list, src=0)
-            # if rng_file_list still empty, not log rng state.
-            if rng_file_list[0] is None:
-                logger.info(
-                    f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
-                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
-                )
-                return
-            else:
-                checkpoint_rng_state = rng_file_list[process_index]
-        else:
-            rng_file = os.path.join(checkpoint, "rng_state.pth")
-            if not os.path.isfile(rng_file):
-                logger.info(
-                    "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
-                    "fashion, reproducibility is not guaranteed."
-                )
-                return
+        rng_file = os.path.join(checkpoint, f"rng_state_{dist.get_rank()}.pth")
+        if not os.path.isfile(rng_file):
+            logger.info(
+                "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
+                "fashion, reproducibility is not guaranteed."
+            )
+            return
 
-            checkpoint_rng_state = paddle.load(rng_file, return_numpy=True)
+        checkpoint_rng_state = paddle.load(rng_file, return_numpy=True)
+        if checkpoint_rng_state.get("world_size", None) != self.args.world_size:
+            logger.warn("Cannot load rng states when changing world size of training job.")
+            return
 
         random.setstate(checkpoint_rng_state["python"])
         np.random.set_state(checkpoint_rng_state["numpy"])
@@ -2210,11 +2194,6 @@ def _wrap_model(self, model, training=True):
             else:
                 model, self.optimizer = decorated
 
-        if self.args.tensor_parallel_degree > 1 and self.args.sequence_parallel:
-            register_sequence_parallel_allreduce_hooks(
-                model, self.args.gradient_accumulation_steps, self.args.fuse_sequence_parallel_allreduce
-            )
-
         if self.args.world_size == 1:
             if self.args.amp_master_grad:
                 mix_precision_utils.MixPrecisionLayer(model, dtype=self.amp_dtype)
@@ -2403,6 +2382,17 @@ def get_expected_keys(inputs, keys):
                 ):
                     self.optimizer._set_broadcast_overlap(True, model)
 
+        # use callback for sp grad sync in case of unexpected behaviour (except sharding stage 2&3)
+        if self.args.tensor_parallel_degree > 1 and self.args.sequence_parallel:
+            if ShardingOption.SHARD_GRAD_OP in self.args.sharding or ShardingOption.FULL_SHARD in self.args.sharding:
+                register_sequence_parallel_allreduce_hooks(
+                    unwrap_model(model),
+                    self.args.gradient_accumulation_steps,
+                    self.args.fuse_sequence_parallel_allreduce,
+                )
+            else:
+                self.add_callback(SPGradSyncCallback(model._layers))
+
         return model
 
     def _prepare_input(self, data: Union[paddle.Tensor, Any]) -> Union[paddle.Tensor, Any]:
@@ -2739,28 +2729,24 @@ def _save_checkpoint(self, model, metrics=None):
         if self.args.should_save:
             self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
 
-        # Save RNG state in non-distributed training
-        rng_states = {
-            "python": random.getstate(),
-            "numpy": np.random.get_state(),
-            "cuda": paddle.get_rng_state(),
-            "cpu": paddle.framework.core.default_cpu_generator().get_state(),
-        }
-        if self.args.use_hybrid_parallel:
-            rng_states[
-                "hybrid_parallel_rng_state_tracker"
-            ] = fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
+        if self.args.save_rng_states:
+            # Save RNG state in non-distributed training
+            rng_states = {
+                "python": random.getstate(),
+                "numpy": np.random.get_state(),
+                "cuda": paddle.get_rng_state(),
+                "cpu": paddle.framework.core.default_cpu_generator().get_state(),
+                "world_size": self.args.world_size,
+            }
+            if self.args.use_hybrid_parallel:
+                rng_states[
+                    "hybrid_parallel_rng_state_tracker"
+                ] = fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
 
         if self.args.save_rng_states:
-            if self.args.world_size > 1:
-                rng_states_list = []
-                paddle.distributed.all_gather_object(rng_states_list, rng_states)
-                if self.args.should_save:
-                    os.makedirs(output_dir, exist_ok=True)
-                    paddle.save(rng_states_list, os.path.join(output_dir, f"rng_state_{self.args.world_size}.pth"))
-            else:
-                os.makedirs(output_dir, exist_ok=True)
-                paddle.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
+            rng_state_file = os.path.join(output_dir, f"rng_state_{dist.get_rank()}.pth")
+            os.makedirs(output_dir, exist_ok=True)
+            paddle.save(rng_states, rng_state_file)
 
         # only save model state dict, ignore optimizer and scheduler
         if not self.args.ignore_save_lr_and_optim:

paddlenlp/trainer/trainer_callback.py

@@ -20,10 +20,18 @@
 """
 import dataclasses
 import json
+import time
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 import numpy as np
+from paddle.distributed.fleet import fleet
+from paddle.distributed.fleet.utils.hybrid_parallel_util import (
+    fused_allreduce_gradients_with_group,
+)
+from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+    is_sequence_parallel_parameter,
+)
 from tqdm.auto import tqdm
 
 from paddlenlp.utils.log import logger
@@ -609,3 +617,33 @@ def on_evaluate(self, args, state, control, metrics, **kwargs):
         self.check_metric_value(args, state, control, metric_value)
         if self.early_stopping_patience_counter >= self.early_stopping_patience:
             control.should_training_stop = True
+
+
+class SPGradSyncCallback(TrainerCallback):
+    """
+    SPGradSyncCallback
+    只能在非 sharding stage2 的情况下使用。
+    开启sharding stage2 时，在 `on_optimizer_begin` 的时候 grad 已经被清空了
+    """
+
+    def __init__(self, model):
+        assert hasattr(fleet, "_hcg"), "must use MP when calling this Callback"
+        logger.info("using sp callback")
+        params = []
+        self.model = model
+        for n, p in model.named_parameters():
+            if is_sequence_parallel_parameter(p):
+                logger.info(f"register bw hook for:{n}")
+                params.append(p)
+
+        logger.info(f"#-sp-sync param:{len(params)}")
+        self._sp_params = params
+
+    def on_optimizer_begin(self, args, state, control, **kwargs):
+        """on_optimizer_begin"""
+        if self._sp_params:
+            now = time.time()
+            mp_group = fleet.get_hybrid_communicate_group().get_model_parallel_group()
+            fused_allreduce_gradients_with_group(self._sp_params, group=mp_group, scale=1.0)  # sum not mean
+            another_time = time.time()
+            logger.info(f"sync gradients takes {another_time - now} time")

paddlenlp/trainer/utils/zero_cost_checkpoint.py

@@ -18,10 +18,12 @@
 import json
 import multiprocessing
 import os
+import random
 import time
 from collections import OrderedDict
 from enum import Enum
 
+import numpy as np
 import paddle
 import paddle.autograd as imperative_base
 import paddle.distributed as dist
@@ -414,10 +416,26 @@ def on_step_end(self, args, state, control, model, lr_scheduler, optimizer, **kw
             self.maybe_update_zcc_worker(args, model, optimizer, state.global_step)
             checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
             save_infos = self._get_save_infos_based_on_steps(state, args, checkpoint_folder)
-            non_cached_objects = (lr_scheduler.state_dict(), copy.deepcopy(state))
+            non_cached_objects = (lr_scheduler.state_dict(), state, self.get_rng_states(args))
             self.manager.get_idle_worker_for_saving((save_infos, non_cached_objects))
             self.runtime_timer.stop()
 
+    def get_rng_states(self, args):
+        if not args.save_rng_states:
+            return None
+        rng_states = {
+            "python": random.getstate(),
+            "numpy": np.random.get_state(),
+            "cuda": paddle.get_rng_state(),
+            "cpu": paddle.framework.core.default_cpu_generator().get_state(),
+            "world_size": args.world_size,
+        }
+        if args.use_hybrid_parallel:
+            rng_states[
+                "hybrid_parallel_rng_state_tracker"
+            ] = dist.fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
+        return rng_states
+
     def _get_save_infos_based_on_steps(self, state, args, checkpoint_folder):
         flash_device_checkpoint_dir = None
         persistent_checkpoint_dir = None
@@ -701,6 +719,7 @@ def __init__(
         # TODO(@gexiao): remove lr scheduler saves
         self.lr_scheduler = None
         self.trainer_state = None
+        self.rng_state = None
 
         # for dumping
         self.flash_device_save_dir = None
@@ -734,7 +753,7 @@ def process_prepare_task(self, prepares):
             return
         save_infos, non_cached_objects = prepares
         self.flash_device_save_dir, self.persistent_save_dir = save_infos
-        self.lr_scheduler, self.trainer_state = non_cached_objects
+        self.lr_scheduler, self.trainer_state, self.rng_state = non_cached_objects
 
     def process_offload_task(self, dump, global_step):
         """
@@ -897,6 +916,11 @@ def process_dump_task_impl(self, output_dir):
         if self.device_id == 0:
             self.trainer_state.save_to_json(trainer_state_name_path)
 
+        # Step2.5: save RNG State
+        if self.rng_state is not None:
+            rng_state_name_path = os.path.join(output_dir, f"rng_state_{dist.get_rank()}.pth")
+            paddle.save(self.rng_state, rng_state_name_path)
+
         # Step3: dump save signals
         saved_signal_path = os.path.join(output_dir, f"saved_signal_{self.global_rank}")
         with open(saved_signal_path, mode="w+") as f:

paddlenlp/transformers/moe_gate.py

@@ -326,6 +326,7 @@ def top1gating(
             logits += self.gumbel_rsample(logits.shape)
 
         gates = self.gate_score_func(logits=logits)
+
         capacity = self._capacity(gates, self.capacity_factor, self.max_capacity, self.min_capacity)
 
         # Create a mask for 1st's expert per token
@@ -396,6 +397,7 @@ def top2gating(
         logits: paddle.Tensor,
     ) -> Tuple[int, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]:
         # everything is in fp32 in this function
+
         gates = self.gate_score_func(logits=logits)
 
         # Create a mask for 1st's expert per token.