chatlearn/models/base_module.py

# Copyright 2024 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""base module"""

from collections import defaultdict
from itertools import cycle
from pathlib import Path
import math
import time
import os

import torch

import ray
import ray.util.collective as col
from ray.util.collective.collective_group.base_collective_group import BaseGroup
from ray.util.collective.collective_group.nccl_collective_group import NCCLGroup
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors

from chatlearn.data.sampler import MultiDatasetSampler
from chatlearn.data.data import RLHFDataLoader
from chatlearn.checkpoint.checkpoint_manager import CheckpointManager
from chatlearn.utils import future
from chatlearn.utils.constant import LOG_START
from chatlearn.utils.dist_utils import bucket_tensors, coalesced_comm_dense
from chatlearn.utils.dist_utils import bucket_tensors_two_stage_generator, coalesced_comm_dense_two_stage
from chatlearn.utils.global_vars import get_args
from chatlearn.utils.global_vars import set_global_variables
from chatlearn.utils.logger import logger
from chatlearn.utils.logger import log_rank_0, debug_rank_0, setup_logger
from chatlearn.utils.timer import Timers
from chatlearn.utils.utils import get_host_addr, map_reduce_metrics
from chatlearn.launcher import dlc_utils


class BaseModule:
    """BaseModule is the base class for Base models.

    Args
    ----
    name : str
        model name
    """

    def __init__(self, name, args=None, replica_id=0):
        logger.info(f"{LOG_START} basemodule {name} init start")
        self.name = name
        if args is None:
            global_args = get_args()
        else:
            global_args = args
            set_global_variables(args)
        self.global_args = global_args
        args = global_args.models[name]
        self.total_gpu = args.num_gpu
        self.total_cpu = args.num_cpu
        self.gpu_per_process = args.gpu_per_process
        self.trainable = args.trainable
        self._runtime_args = self.global_args.runtime_args
        self._module_args = args
        self.replica_id = replica_id
        self.config_dir = args.config_dir
        self._is_colocate = False

        if self.total_gpu > 0:
            self._num_gpu_per_replica = (
                args.tensor_model_parallel_size
                * args.pipeline_model_parallel_size
                * args.expert_model_parallel_size
                * args.zero_size
                * args.fsdp_size
            )
            assert self._num_gpu_per_replica <= self.total_gpu, \
                f"_num_gpu_per_replica {self._num_gpu_per_replica} larger than total_gpu {self.total_gpu} " + \
                f"tp_size: {args.tensor_model_parallel_size} pp_size: {args.pipeline_model_parallel_size} " + \
                f"ep_size: {args.expert_model_parallel_size} zero_size: {args.zero_size}"
            assert self.total_gpu % self._num_gpu_per_replica == 0
            if not self.trainable:
                self._num_replica = args.num_gpu // self._num_gpu_per_replica
            else:
                # For trainable models, perform the DP inside DistActor
                self._num_replica = 1
                self._num_gpu_per_replica = self.total_gpu
        else:
            self._num_gpu_per_replica = 0
            self._num_replica = args.num_replica

        assert self._num_replica >= 1
        self._param_ranks = None
        self._named_parameters = None
        self._param_to_name = None
        self._parameters = None
        self._coalesced_parameters = None
        self.error_signal = None
        self._rank = None
        self._world_size = None
        self._group_names = []
        self._dataloader = None
        self._eval_dataloader = None
        self._kl_coef = None
        self._padding_config = {}
        self._storage = None
        self._timers = None
        self._data_iter = None
        self._eval_data_iter = None
        self.call_funcs = []
        self.trainable_funcs = []
        self._data_ckpt_manager = None
        self._peak_memory = 0
        self._parameters_to_sync = defaultdict(list)
        self._parameters_to_send = defaultdict(list)
        self._parameters_to_recv = defaultdict(list)
        self._parameters_shape = []
        # current compute iteration
        self._iteration = 0
        self._train_iteration = 0
        self._episode_id = 0
        self.enable_lora = self._module_args.lora.enable_lora
        self._finalized = False
        self._resume_training = False
        self._address = dlc_utils.get_addr() if dlc_utils.in_dlc_env() else get_host_addr()
        self._is_master_node = os.environ.get("RANK", '0') == '0'
        self._logger = setup_logger(model_name=self.name, ip_addr=self._address)
        # parameter sync from src_model
        self._src_parameter_model = None
        self.profiler = None
        self._buffer_num = {}
        self._tp_division = {}
        self._tp_num_mapping = 1
        self._sync_buffer = defaultdict(list)
        self._sync_dst_rank_to_src_ranks = {}
        self._expert_sync_buffer = {}
        self._synchronizer = None
        self._metric_prefix = ""
        self._metric_list = []
        self._stage_resume_done = False
        logger.info(f"{LOG_START} basemodule {name} init done")

    def get_sync_buffer(self):
        return self._sync_buffer

    def set_tp_num_mapping(self, _tp_num_mapping):
        self._tp_num_mapping = _tp_num_mapping

    @property
    def tp_num_mapping(self):
        return self._tp_num_mapping

    def set_buffer_num(self, buffer_num):
        self._buffer_num.update(buffer_num)

    def get_buffer_num(self, param_names):
        return [self._buffer_num[name] for name in param_names]

    def set_tp_division(self, tp_division):
        self._tp_division.update(tp_division)

    def get_tp_division(self, param_names):
        return [self._tp_division[name] for name in param_names]

    @property
    def is_colocate(self):
        return self._is_colocate

    def set_colocate(self, flag):
        self._is_colocate = flag

    def finalize(self):
        """
        finalize the class, any change from user after finalize will not work.

        :meta private:
        """
        self._finalized = True

    def _assert_not_finalized(self):
        """
        :meta private:
        """
        assert not self._finalized, f"{self} is finalized, any change to the class should happen before finalize."

    def get_runtime_args(self):
        return self.runtime_args

    @property
    def runtime_args(self):
        """
        Return the arguments related to alignment training,
        the settings that are specified under the "runtime" section of the YAML configuration file.
        """
        return self._runtime_args

    @property
    def model_args(self):
        """
        Return model arguments, such as those related to Megatron,
        should be specified in a separate configuration yaml file for the model being used.
        """
        return self._module_args.args_dict

    @property
    def module_args(self):
        """
        Return module arguments. module_args include `num_gpu`, `gpu_per_process`, `model_config_file`, etc.
        """
        return self._module_args

    @property
    def parameter_sync_frequency(self):
        return self.module_args.sync_frequency

    def set_env(self, args):
        """
        set system env, private

        :meta private:
        """

    def set_error_signal(self, error_signal):
        """
        signal for handling errors

        :meta private:
        """
        self.error_signal = error_signal

    def error(self, error_msg=None):
        """
        :meta private:
        """
        future.wait(self.error_signal.set.remote(error_msg))

    def init(self):
        """
        Init env.
        """

    def setup(self):
        """
        Create model / optimizer / opt_param_scheduler / etc.
        """

    @property
    def data_ckpt_manager(self):
        """
        :meta private:
        """
        if self.runtime_args.data_checkpoint_path is not None:
            assert self._data_ckpt_manager is not None
        return self._data_ckpt_manager

    def model_setup(self):
        """
        :meta private:
        """
        self.global_args.active_module_args = self._module_args
        if self.runtime_args.data_checkpoint_path is not None:
            self._data_ckpt_manager = CheckpointManager(self, self.runtime_args.data_checkpoint_path,
                                                       self.runtime_args.max_data_ckpt_nums,
                                                       self.runtime_args.load_data_checkpoint_iteration)
            if self.runtime_args.enable_resume_training:
                meta = self._data_ckpt_manager.resume()
                if meta:
                    self._resume_training = self.runtime_args.consumed_samples > 0
                    start_episode = meta["episode"] + 1
                    self._episode_id = start_episode
                    self._iteration = start_episode * math.ceil(self.runtime_args.sample_per_episode / \
                        self._num_replica / self.module_args.generation_batch_size)

                    log_rank_0(
                        f"{self.name} resume training {self._resume_training}: "
                        f"set start iteration to {self._iteration} and episode id to {self._episode_id}",
                        self._logger)
        self.setup()

    def forward_step(self, data, iteration):
        """
        Perform forward step for one batch.

        Args
        ----
        data : dict
            data for forward_step
        iteration : int
            local forward iteration
        
        Returns
        -------
        Dict
            A dict of results, where key is the string type, and the value is the tensor or a list,
            where the first dim of tensor or the len of list equals to batch size
        """

    def train_step(self, data, iteration):
        """
        Perform train_step for one batch, including a list of micro-batches.

        Args
        ----
        data : [Dict]
            A list of micro-batch for train_step, type of each micro-batch is dict
        iteration : int
            local train iteration
        """

    def eval_step(self, data):
        """
        Perform eval_step for one batch

        Args
        ----
            data: Dict
                Data for eval_step.

        Returns
        -------
            Dict
                A dict of results, where key is the string type, and the value is the tensor or a list,
                where the first dim of tensor or the len of list equals to batch size
        """

    def save_checkpoint(self, iteration):
        """
        Save checkpoint given iteration.

        Args
        ----
            iteration: int
                Current training iteration
        """

    def save_data_checkpoint(self, replica_id, iteration, episode_id):
        """
        Save checkpoint for dataloader.

        :meta private:
        """
        if self.data_ckpt_manager is not None:
            consumed_samples = self.runtime_args.consumed_samples
            self.data_ckpt_manager.save_checkpoint(replica_id, iteration, episode_id, consumed_samples)

    def put(self, key, data):
        """
        Put the data to shared storage.

        Args
        ----
            key: Str
                Use key to put.
            data
                data to save
        """
        self._storage.put.remote(key, data)

    def get(self, key):
        """
        Get data from shared storage using key

        Args
        ----
            key: Str
                use key to get
        """
        ref = self._storage.get.remote(key)
        return future.get(ref)

    def validate(self):
        """
        :meta private:
        """

    def before_episode(self):
        """
        Operations before one episode.
        """

    def after_episode(self):
        """
        Operations after one episode.
        """
        self._episode_id += 1

    def build_dataset(self, train_prompts, is_eval=False):
        """
        Build prompt dataset

        Args
        ----
            train_prompts: [Str]
                A list of prompt string.
        Returns
        -------
            torch.utils.data.Dataset
                Dataset with user-defined collate_fn
        """

    def build_all_dataset(self, train_prompts_list, is_eval=False):
        """
        Build all prompt datasets

        Args
        ----
            train_prompts_list: List[List[Str]]
                A list of prompt string lists.
        Returns
        -------
            List[torch.utils.data.Dataset]
                A list of Dataset with user-defined collate_fn
        """
        all_datasets = []
        for train_prompts in train_prompts_list:
            all_datasets.append(
                self.build_dataset(train_prompts, is_eval)
            )
        return all_datasets

    def _build_dataloader(self, data, sample_per_episode, is_eval=False):
        """
        build and set the dataloader for the model

        Args:
            data: a list of string
            is_eval: set to `True` to build a dataloader for evaluation (default: `False`)

        :meta private:
        """
        all_datasets = self.build_all_dataset(data, is_eval) # pylint: disable=assignment-from-no-return
        consumed_samples = 0
        data_ratio = self.runtime_args.data_ratio
        shuffle = self.runtime_args.data_shuffle
        data_rerank = self.runtime_args.data_rerank
        if not is_eval:
            if self.data_ckpt_manager is not None:
                consumed_samples = self.runtime_args.consumed_samples
        collate_fn = all_datasets[0].collate_fn if hasattr(all_datasets[0], 'collate_fn') else None
        drop_last = self.model_args['drop_last'] if 'drop_last' in self.model_args else False
        dataloader = self.build_dataloader(all_datasets,
                                           sample_per_episode=sample_per_episode,
                                           collate_fn=collate_fn,
                                           is_eval=is_eval,
                                           consumed_samples=consumed_samples,
                                           data_ratio=data_ratio,
                                           shuffle=shuffle,
                                           drop_last=drop_last,
                                           data_rerank=data_rerank)

        if is_eval:
            self._eval_dataloader = dataloader
            self._eval_data_iter = iter(self._eval_dataloader)
        else:
            self._data_iter = iter(dataloader)
            self._data_iter = cycle(self._data_iter)
            self._dataloader = dataloader

    def build_dataloader(self,
                         all_datasets,
                         sample_per_episode,
                         collate_fn=None,
                         is_eval=False,
                         consumed_samples=0,
                         data_ratio=None,
                         shuffle=True,
                         drop_last=False,
                         data_rerank=True):
        """
        build the dataloader for the model
        Args:
            all_datasets: a list of torch.utils.data.Dataset objects
            batch_size: how many samples per batch to load
            collate_fn: set when loading from an map-style dataset (defulat: `None`)
            is_eval: set to `True` to build a dataloader for evaluation (default: `False`)
            consumed_samples: consumed samples (default: `0`)
            data_ratio: ratio of samples for each dataset (default: `None`)
            drop_last: whether to drop last samples (default: `False`)

        :meta private:
        """
        log_rank_0(
            f"Creating DataLoader... consumed_samples: {consumed_samples}, "
            f"data_ratio: {data_ratio}",
            self._logger
        )
        if "num_inference_per_prompt" in self.model_args:
            num_inference_per_prompt = self.model_args["num_inference_per_prompt"]
        else:
            num_inference_per_prompt = 1
        vllm_prompt_key = self.model_args["vllm_prompt_key"] \
            if "vllm_prompt_key" in self.model_args else "prompt"
        self._logger.info(f"====Data Rerank: {data_rerank}")
        if is_eval:
            batch_sampler = MultiDatasetSampler(
                dataset_sizes=[len(dataset) for dataset in all_datasets],
                sample_per_episode=sample_per_episode,
                shuffle=False,
                is_eval=True,
                data_parallel_rank=self.replica_id,
                data_parallel_size=self._num_replica
            )
        else:
            batch_sampler = MultiDatasetSampler(
                dataset_sizes=[len(dataset) for dataset in all_datasets],
                sample_per_episode=sample_per_episode,
                data_ratio=data_ratio,
                consumed_samples=consumed_samples,
                num_inference_per_prompt=num_inference_per_prompt,
                shuffle=shuffle,
                is_eval=False,
                data_parallel_rank=self.replica_id,
                data_parallel_size=self._num_replica,
                drop_last="drop" if drop_last else "cycle",
                data_rerank=data_rerank
            )
        return RLHFDataLoader(
            all_datasets,
            batch_sampler,
            collate_fn=collate_fn,
            add_uid=True,
            data_parallel_rank=self.replica_id,
            data_parallel_size=self._num_replica,
            num_inference_per_prompt=num_inference_per_prompt,
            vllm_prompt_key=vllm_prompt_key
        )

    def reset_eval_data_iter(self):
        """
        :meta private:
        """
        if self._eval_dataloader is not None:
            self._eval_data_iter = iter(self._eval_dataloader)

    def next_batch(self, is_eval=False):
        """
        :meta private:
        """
        if is_eval:
            return next(self._eval_data_iter)
        else:
            return next(self._data_iter)

    @property
    def num_replica(self):
        """
        :meta private:
        """
        return self._num_replica

    @property
    def num_gpu_per_replica(self):
        """
        :meta private:
        """
        return self._num_gpu_per_replica

    def setup_collective_group(self, rank, world_size, backend, group_name):
        """
        :meta private:
        """
        self._group_names.append(group_name)
        self._world_size = world_size
        col.init_collective_group(
            world_size, rank, backend=backend, group_name=group_name)

    def broadcast_dummy_tensor_send(self, src_rank, group_name):
        x = torch.zeros(1, device="cuda")
        col.broadcast(x, src_rank=src_rank, group_name=group_name)
        del x

    def broadcast_dummy_tensor_recv(self, src_rank, group_name):
        x = torch.zeros(1, device="cuda")
        col.broadcast(x, src_rank=src_rank, group_name=group_name)
        del x

    def _destroy_collective_group(self, group_name):
        """
        :meta private:
        """
        from ray.util.collective.collective import _group_mgr # pylint: disable=import-outside-toplevel
        rank = col.get_rank(group_name)
        saved_group: BaseGroup = _group_mgr.get_group_by_name(group_name)
        saved_comm_keys = []
        if isinstance(saved_group, (NCCLGroup, )):
            saved_comm_keys = list(saved_group._dev_comm_map.keys())

        try:
            col.destroy_collective_group(group_name)
        except Exception as e:
            self._logger.warning(f"_destroy_collective_group {group_name} {e}")

        if isinstance(saved_group, (NCCLGroup, )):
            for comm_key in saved_comm_keys:
                group_key = saved_group._generate_group_key(comm_key)
                from ray.util.collective.const import get_store_name # pylint: disable=import-outside-toplevel
                store_name = get_store_name(group_key)
                try:
                    store = ray.get_actor(store_name)
                    if rank == 0:
                        raise RuntimeError(f'{store_name} in group {group_name} should be killed on rank {rank}.')
                    self._logger.debug(f'Kill {store_name} in group {group_name} on rank {rank}')
                    ray.kill(store)
                except ValueError:
                    ...

    def destroy_collective_group(self):
        for group_name in self._group_names:
            self._destroy_collective_group(group_name)
        self._group_names = []

    def get_local_param_ranks(self):
        """
        :meta private:
        """

    def fuse_lora_layer(self):
        """
        :meta private:
        """
        from chatlearn.models.megatron.lora import fuse_lora_layer # pylint: disable=import-outside-toplevel
        fuse_lora_layer(self.model)

    def unfuse_lora_layer(self):
        """
        :meta private:
        """
        from chatlearn.models.megatron.lora import unfuse_lora_layer # pylint: disable=import-outside-toplevel
        unfuse_lora_layer(self.model)

    @property
    def rank(self):
        """
        :meta private:
        """
        return self._rank

    def get_rank(self):
        """
        :meta private:
        """
        return self.rank

    def is_last_rank(self):
        """
        Is last rank.
        """
        return True

    @property
    def parameters(self):
        """
        :meta private:
        """
        if self._parameters is None:
            if not isinstance(self.model, list):
                model = [self.model]
            else:
                model = self.model
            self._parameters = []
            for partition in model:
                for item in partition.parameters():
                    self._parameters.append(item)
        return self._parameters

    @property
    def named_parameters(self):
        """
        :meta private:
        """
        if self._named_parameters is None:
            if not isinstance(self.model, list):
                model = [self.model]
            else:
                model = self.model
            self._named_parameters = {}
            for partition in model:
                for item in partition.named_parameters():
                    self._named_parameters[item[0]] = item[1]
        return self._named_parameters

    @property
    def param_to_name(self):
        """
        :meta private:
        """
        if self._param_to_name is None:
            if not isinstance(self.model, list):
                model = [self.model]
            else:
                model = self.model
            self._param_to_name = {}
            for partition in model:
                for item in partition.named_parameters():
                    self._param_to_name[item[1]] = item[0]
        return self._param_to_name

    def _set_sync_parameters(self, trainable_param_names, pipe_stage=0, parameters_to_sync=None):
        if parameters_to_sync is None:
            parameters_to_sync = defaultdict(list)
        assert pipe_stage not in parameters_to_sync or len(parameters_to_sync[pipe_stage])==0
        params_to_sync_list = [(name, self.named_parameters[name]) for name in trainable_param_names]
        if self._synchronizer is not None:
            params_to_sync_list = self._synchronizer.transform_parameters(params_to_sync_list)
        parameters_to_sync[pipe_stage] = params_to_sync_list
        return parameters_to_sync

    def set_sync_parameters(self, trainable_param_names, pipe_stage=0, parameters_to_sync=None):
        """
        :meta private:
        """
        if parameters_to_sync is None:
            parameters_to_sync = self._parameters_to_sync
        if pipe_stage not in parameters_to_sync or len(parameters_to_sync[pipe_stage]) == 0:
            self._set_sync_parameters(trainable_param_names, pipe_stage, parameters_to_sync)

    def reset_sync_parameters(self, trainable_param_names, pipe_stage=0):
        self._parameters_to_sync[pipe_stage] = []
        self._set_sync_parameters(trainable_param_names, pipe_stage, self._parameters_to_sync)

    def set_send_parameters(self, trainable_param_names, pipe_stage=0):
        """
        :meta private:
        """
        return self.set_sync_parameters(trainable_param_names, pipe_stage, self._parameters_to_send)

    def set_recv_parameters(self, to_rank, trainable_param_names, pipe_stage=0):
        """
        :meta private:
        """
        parameters_to_recv = defaultdict(list)
        self._parameters_to_recv[to_rank] = parameters_to_recv
        return self.set_sync_parameters(trainable_param_names, pipe_stage, parameters_to_recv)

    def clear_sync_parameters(self):
        self._parameters_to_sync = defaultdict(list)

    def clear_send_recv_parameters(self):
        self._parameters_to_send = defaultdict(list)
        self._parameters_to_recv = defaultdict(list)

    def clear_sync_send_recv_parameters(self):
        self.clear_sync_parameters()
        self.clear_send_recv_parameters()

    def get_parameter_names(self, requires_grad=True):
        """
        :meta private:
        """
        param_to_name = self.param_to_name
        if requires_grad:
            return [param_to_name[param] for param in self.parameters if param.requires_grad]
        else:
            return [param_to_name[param] for param in self.parameters]

    def get_parameter_shape(self, pipe_stage=0, parameters_to_sync=None):
        """
        :meta private:
        """
        if parameters_to_sync is None:
            parameters_to_sync = self._parameters_to_sync
        parameters_shape = []
        for name, param in parameters_to_sync[pipe_stage]:
            if self._expert_sync_buffer and name in self._expert_sync_buffer and \
                    self._synchronizer and self._synchronizer.is_parameter_changed:
                parameters_shape.append((name, self._expert_sync_buffer[name].shape))
            else:
                parameters_shape.append((name, param.shape))
        return parameters_shape

    def get_parameter(self, name):
        """
        :meta private:
        """
        if name not in self.named_parameters:
            raise Exception(f"parameter {name} not exits")
        return self.named_parameters[name]

    def get_parameter_to_sync(self, name, pipe_stage, to_cpu=False, regroup=False):
        assert pipe_stage in self._parameters_to_sync and len(self._parameters_to_sync[pipe_stage]) > 0
        for name0, param in self._parameters_to_sync[pipe_stage]:
            if name0 == name:
                if name in self._expert_sync_buffer and self._synchronizer and \
                        self._synchronizer.is_parameter_changed:
                    param = self._expert_sync_buffer[name]
                    regroup_routed_experts = True
                else:
                    regroup_routed_experts = False
                if regroup and self._synchronizer:
                    param = self._synchronizer.regroup_params_to_sync(
                        name,
                        param.data,
                        self._tp_division[name],
                        regroup_routed_experts
                    )
                if to_cpu:
                    param = param.cpu()
                else:
                    param = param.cuda()
                return param

    def get_parameter_to_sync_names(self, pipe_stage):
        return [items[0] for items in self._parameters_to_sync[pipe_stage]]

    def exist_parameter(self, name):
        """
        :meta private:
        """
        return name in self.named_parameters

    def parameter_shape(self, name):
        """
        :meta private:
        """
        return self.get_parameter(name).shape

    def send_recv_parameter(self, rank, group_name, func, pipe_stage=0):
        """
        :meta private:
        """
        tensors = [param.data for _, param in self._parameters_to_sync[pipe_stage]]
        dense_buckets, sparse_bucket = bucket_tensors(tensors, bucket_size_mb=self.runtime_args.coalesced_buffer_mb)
        debug_rank_0(f"{self.name} Got dense_buckets {len(dense_buckets)}, spase_bucket {len(sparse_bucket)}", self._logger)
        for bucket in dense_buckets:
            tensor_changed = func is col.recv
            coalesced_comm_dense(bucket, func, extra_args=(rank, group_name), tensor_changed=tensor_changed)
        for param in sparse_bucket:
            func(param, rank, group_name)

    def alltoall_routed_expert_parameter(self, pipe_stage=0):
        assert self._synchronizer is not None
        for name, param in self._parameters_to_sync[pipe_stage]:
            param, state = self._synchronizer.alltoall_routed_experts(
                name,
                param,
                self.tensor_and_expert_parallel_group()
            )
            if state:
                self._expert_sync_buffer.pop(name, "Not Found.")
                self._expert_sync_buffer[name] = param

    def allgather_routed_expert_parameter(self, group_name, pipe_stage=0):
        assert self._synchronizer is not None
        for name, param in self._parameters_to_sync[pipe_stage]:
            param, state = self._synchronizer.allgather_routed_experts(
                name,
                param,
                group_name,
                tp_rank=self.tensor_parallel_rank()
            )
            if state:
                self._expert_sync_buffer.pop(name, "Not Found.")
                self._expert_sync_buffer[name] = param

    def broadcast_parameter(self, rank, src_rank, group_name, pipe_stage=0):
        """
        :meta private:
        """
        tensors = []
        for name, param in self._parameters_to_sync[pipe_stage]:
            if self._expert_sync_buffer and name in self._expert_sync_buffer and \
                    (self._synchronizer and self._synchronizer.is_parameter_changed):
                tensors.append(self._expert_sync_buffer[name])
            else:
                tensors.append(param.data)

        assert len(tensors) > 0
        dense_buckets, sparse_bucket = bucket_tensors(tensors, bucket_size_mb=self.runtime_args.coalesced_buffer_mb)
        debug_rank_0(f"{self.name} Got dense_buckets {len(dense_buckets)}, spase_bucket {len(sparse_bucket)}", self._logger)
        tensor_changed = rank != src_rank

        for bucket in dense_buckets:
            coalesced_comm_dense(bucket, col.broadcast, extra_args=(src_rank, group_name), tensor_changed=tensor_changed)

        for param in sparse_bucket:
            col.broadcast(param, src_rank, group_name)

    def broadcast_parameter_two_stage(self, to_rank, buffer_rank, rank, src_rank, group_name, pipe_stage=0, stage2=False):
        """
        Arguments:
            to_rank: receive rank in mapping from trainer to inference model.
            buffer_rank: index which tensors of sync buffer to be sended in stage2.
            rank: destination rank in communication group which enumerate receive ranks.
            src_rank: source rank in communication group. always 0.
            group_name: communication group name.
            pipe_stage: pipeline stage. default 0.
            stage2: bool. whether stage2 or not. default False.
        Example: trainer_tp = 4, inference_tp = 8. pipeline_size = 1
            stage1: [(from_rank, to_rank), ...] = [(0, 8), (1, 10), (2, 12), (3, 14)]
            stage2: [(from_rank, to_rank), ...] = [(8, 9), (10, 11), (12, 13), (14, 15)]

            For stage1 pair (0, 8):
                1. call broadcast func: (0 -> 0). src_rank: 0, rank: 0.
                2. call broadcast func: (0 -> 8). src_rank: 0, rank: 1.

                After (0, 8), to_rank 8 received tensor slices of 8 and 9.

            For stage2 pair (8, 9):
                1. call broadcast func: (8 -> 8). src_rank: 0, rank: 0.
                2. call broadcast func: (8 -> 9). src_rank: 0, rank: 1.
                In (8 -> 8), we need to send tp_slice of 'to_rank' 9, so set buffer_rank 9 to fetch tensors in sync buffer.
        """
        tensor_changed = rank != src_rank
        start = time.time()
        arguments = f"{to_rank}_{buffer_rank}_{rank}_{src_rank}_{group_name}_{pipe_stage}_{stage2}"

        if stage2:
            if tensor_changed:
                parameters_to_sync = self._parameters_to_recv[to_rank]
            else:
                parameters_to_sync = self._parameters_to_send
        else:
            if rank not in self._sync_dst_rank_to_src_ranks:
                self._sync_dst_rank_to_src_ranks.update({rank:[src_rank]})
                del self._sync_buffer
                self._sync_buffer = defaultdict(list)
            else:
                self._sync_dst_rank_to_src_ranks[rank].append(src_rank)
            parameters_to_sync = self._parameters_to_sync

        def tensor_generator():
            if stage2 and not tensor_changed and self._sync_buffer:# pylint: disable=too-many-nested-blocks
                idx = 0
                for name, param in parameters_to_sync[pipe_stage]:
                    value = self._sync_buffer[buffer_rank % self.tp_num_mapping][idx].cuda() # restore from cpu
                    self._logger.debug(
                        f"Adding {name}({value.shape}) to sync for if branch from "
                        f"src_rank: {src_rank} to rank: {rank} in pipe_stage {pipe_stage}"
                    )
                    buffer_num = 1
                    idx += 1
                    yield value, buffer_num
                del self._sync_buffer[buffer_rank % self.tp_num_mapping]
            else:
                idx = 0
                for name, param in parameters_to_sync[pipe_stage]:
                    idx += 1
                    param_data = param.data
                    if rank and self._buffer_num and not stage2:
                        assert name in self._buffer_num, f"{name} in self._buffer_num for rank {rank}"
                        buffer_num = self._buffer_num[name]
                    elif stage2:
                        buffer_num = 1
                    else:
                        if self._expert_sync_buffer and name in self._expert_sync_buffer:
                            param_data = self._expert_sync_buffer[name]
                            regroup_routed_experts = True # For routed experts in Qwen2vLLM
                        else:
                            regroup_routed_experts = False
                        # regroup src_tensor by tp_rank
                        param_data = self._synchronizer.regroup_params_to_sync(
                            name,
                            param_data,
                            self._tp_division[name],
                            regroup_routed_experts
                        )
                        # move self._expert_sync_buffer[name] to cpu mem to save gpu mem
                        if regroup_routed_experts and name in self._expert_sync_buffer:
                            cpu_expert = self._expert_sync_buffer[name].cpu()
                            del self._expert_sync_buffer[name]
                            self._expert_sync_buffer[name] = cpu_expert
                        buffer_num = 1
                    self._logger.debug(
                        f"Adding {name}({param_data.shape}) to sync for else branch from "
                        f"src_rank: {src_rank} to rank: {rank} in pipe_stage {pipe_stage}"
                    )
                    yield param_data, buffer_num

        bucket_generator = bucket_tensors_two_stage_generator(
            tensor_generator, bucket_size_mb=self.runtime_args.coalesced_buffer_mb,
            stage2=stage2, tensor_changed=tensor_changed and not stage2
        )
        dense_bucket_num = 0
        sparse_bucket_num = 0
        for bucket_or_tensor, is_dense in bucket_generator:
            if is_dense:
                index = 0 if stage2 else (to_rank % self.tp_num_mapping)
                all_buffers = coalesced_comm_dense_two_stage(
                    bucket_or_tensor, col.broadcast, rank,
                    extra_args=(src_rank, group_name), tensor_changed=tensor_changed,
                    stage2=stage2, index=index)
                if tensor_changed and not stage2:
                    for key, value in all_buffers.items():
                        cpu_value = []
                        for tensor in value:
                            cpu_value.append(tensor.cpu().pin_memory()) # save gpu memory
                        del value
                        self._sync_buffer[key] += cpu_value
                    del all_buffers
                dense_bucket_num += 1
            else:
                col.broadcast(bucket_or_tensor, src_rank, group_name)
                sparse_bucket_num += 1

        if stage2:
            self._sync_dst_rank_to_src_ranks = {}

        self._logger.debug(f"broadcast_parameter_two_stage {arguments} done using {time.time()-start} seconds")
        debug_rank_0(f"{self.name} Got dense_buckets {dense_bucket_num}, sparse_bucket {sparse_bucket_num}", self._logger)

    def send_parameter(self, dst_rank, group_name, pipe_stage=0):
        """
        :meta private:
        """
        self.send_recv_parameter(dst_rank, group_name, col.send, pipe_stage)

    def recv_parameter(self, src_rank, group_name, pipe_stage=0):
        """
        :meta private:
        """
        self.send_recv_parameter(src_rank, group_name, col.recv, pipe_stage)

    def ray_put_parameter(self, group_name, pipe_stage=0):
        """
        :meta private:
        """
        name2ref = {}
        tensors = [param.data for _, param in self._parameters_to_sync[pipe_stage]]
        dense_buckets, sparse_bucket = bucket_tensors(tensors, bucket_size_mb=self.runtime_args.coalesced_buffer_mb)
        debug_rank_0(f"{self.name} Put dense_buckets {len(dense_buckets)}, spase_bucket {len(sparse_bucket)}", self._logger)
        for bucket_id, bucket in enumerate(dense_buckets):
            flat_tensors = _flatten_dense_tensors(bucket)
            flat_tensors_ref = ray.put(flat_tensors)
            name2ref[group_name + ":dense_bucket_" + str(bucket_id)] = flat_tensors_ref
        for param_id, param in enumerate(sparse_bucket):
            param_ref = ray.put(param)
            name2ref[group_name + ":sparse_bucket_" + str(param_id)] = param_ref
        return name2ref

    def ray_get_parameter(self, group_name, name2ref, pipe_stage=0):
        """
        :meta private:
        """
        tensors = [param.data for _, param in self._parameters_to_sync[pipe_stage]]
        dense_buckets, sparse_bucket = bucket_tensors(tensors, bucket_size_mb=self.runtime_args.coalesced_buffer_mb)
        debug_rank_0(f"{self.name} Get dense_buckets {len(dense_buckets)}, spase_bucket {len(sparse_bucket)}", self._logger)
        for bucket_id, bucket in enumerate(dense_buckets):
            put_ref = name2ref[group_name + ":dense_bucket_" + str(bucket_id)]
            flat_tensors = ray.get(put_ref)
            for tensor, synced in zip(
                bucket, _unflatten_dense_tensors(flat_tensors, bucket)):
                tensor.copy_(synced)
        for param_id, param in enumerate(sparse_bucket):
            put_ref = name2ref[group_name + ":sparse_bucket_" + str(param_id)]
            param.copy_(ray.get(put_ref))

    def pipeline_model_parallel_size(self):
        """
        :meta private:
        """
        return self.module_args.pipeline_model_parallel_size

    def tensor_model_parallel_size(self):
        """
        :meta private:
        """
        return self.module_args.tensor_model_parallel_size

    def expert_model_parallel_size(self):
        """
        :meta private:
        """
        return self.module_args.expert_model_parallel_size

    def num_layers(self):
        """
        :meta private:
        """

    def set_storage(self, storage):
        """
        :meta private:
        """
        self._storage = storage

    def timers(self, name):
        """
        :meta private:
        """
        if self._timers is None:
            self._timers = Timers()
        return self._timers(name)

    def timer_summary(self, e2e_cost=None):
        """
        :meta private:
        """
        if self._timers:
            return self._timers.log(return_dict=True, e2e_cost=e2e_cost)

    def get_and_clear_metrics(self):
        """
        get logging metrics
        """
        if self._metric_list is None or len(self._metric_list) == 0:
            return self._metric_prefix, {}

        reduced_metrics = map_reduce_metrics(self._metric_list)
        self._metric_list = []
        return self._metric_prefix, reduced_metrics

    def add_padding_config(self, key, padding_value=0.0, padding_type="right"):
        """
        Add spectial padding config for certain value.

        Args
        ----
        key: str
            The key for data to be padded.
        padding_value: float
            Padding value, default is 0.
        padding_type: str
            Default right, can be right/left.
        """
        self._padding_config[key] = {"padding_value": padding_value, "padding_type": padding_type}

    def padding_config(self):
        """
        :meta private:
        """
        return self._padding_config

    def peak_memory(self):
        """
        :meta private:
        """
        return 0.0

    @property
    def resume_training(self):
        """
        resume training from last checkpoint.
        """
        return self._resume_training

    def get_address(self):
        """
        Get node address

        :meta private:
        """
        return self._address

    def is_master_node(self):
        """
        Whether this node is master node.
        :meta private:
        """
        return self._is_master_node

    def set_src_parameter_model(self, src_model):
        """
        src_model that sync parameter to current model
        :meta private:
        """
        self._src_parameter_model = src_model

    @property
    def src_parameter_model(self):
        """
        src_model that sync parameter to current model
        """
        return self._src_parameter_model

    def offload_optimizer_states(self):
        """
        offload optimizer states
        """

    def onload_optimizer_states(self):
        """
        onload optimizer states
        """

    def offload_main_weights(self):
        """
        offload main weights
        """

    def onload_main_weights(self):
        """
        onload main weights
        """

    def offload_weights(self):
        """
        offload weights
        """

    def onload_weights(self):
        """
        onload weights
        """

    def free_grad_buffers(self):
        """
        free grad buffers and related tensors
        """

    def build_grad_buffers(self):
        """
        build grad buffers and related tensors
        """

    def onload(self):
        pass

    def offload(self):
        pass

    @property
    def world_size(self):
        pass

    @property
    def data_parallel_size(self):
        """
        data parallel size

        :meta private:
        """

    @property
    def data_parallel_rank(self):
        """
        data parallel rank

        :meta private:
        """

    def empty_cache(self):
        """
        :meta private:
        """

    def get_data_parallel_rank(self):
        return self.data_parallel_rank

    def get_data_parallel_size(self):
        return self.data_parallel_size

    def get_pipeline_stage_layer_num(self):
        pass

    def get_pipeline_stage_layer_offset(self):
        return 0

    def set_synchronizer(self, synchronizer):
        self._synchronizer = synchronizer

    def expert_parallel_rank(self):
        """
        :meta private:
        """
        return 0

    def enable_stage_resume(self, is_eval):
        """
        check whether to resume stage outputs.
        """
        if is_eval:
            return False
        if self.model_args.get("enable_stage_resume", False):
            assert self.runtime_args.data_checkpoint_path, \
                "data_checkpoint_path must be set for stage resume."
            return True
        return False

    def get_stage_outputs_path(self, iteration):
        """
        get path for stage outputs.
        """
        save_dir = self.runtime_args.data_checkpoint_path
        save_path = f"{save_dir}/{iteration}/{self.name}_replica_{self.replica_id}.pt"
        save_path_meta = f"{save_dir}/{iteration}/{self.name}_replica_{self.replica_id}_meta.txt"
        return save_path, save_path_meta

    def load_stage_outputs(self, is_eval, iteration):
        """
        load stage outputs for resume.
        """
        outputs = None
        # only load once for each launching.
        if self.enable_stage_resume(is_eval) and not self._stage_resume_done:
            self._stage_resume_done = True
            save_path, save_path_meta=self.get_stage_outputs_path(iteration)
            if os.path.exists(save_path) and os.path.exists(save_path_meta):
                try:
                    with open(save_path_meta, "r", encoding='utf-8') as f:
                        replica_id = int(f.readline())
                    if replica_id == self.replica_id:
                        outputs = torch.load(save_path)
                        logger.info(f"resume stage outputs for model:{self.name}, path:{save_path}")
                except ValueError:
                    logger.warning(f"ignore incomplete stage outputs, path:{save_path}")
        return outputs

    def save_stage_outputs(self, is_eval, outputs, iteration):
        """
        save stage outputs for resume.
        """
        if self.enable_stage_resume(is_eval):
            save_path, save_path_meta=self.get_stage_outputs_path(iteration)
            logger.info(f"Start to save stage outputs:{save_path}")
            Path(save_path).parent.mkdir(parents=True, exist_ok=True)
            torch.save(outputs, save_path)
            # save meta
            with open(save_path_meta, "w", encoding='utf-8') as f:
                f.write(f"{self.replica_id}")
            logger.info(f"Finished to save stage outputs:{save_path}")