Repository: hpcaitech/ColossalAI
File: colossalai/booster/plugin/moe_hybrid_parallel_plugin.py (Python)

Code before:
        clip_grad_norm: float = 0.0,  # grad clipping
        verbose: bool = False,
        reduce_bucket_size: int = 1024 * 1024,  # communication
        communication_dtype: Optional[torch.dtype] = None,
        overlap_communication: bool = True,
        partition_grad: bool = False,  # stage 2 flag
        cpu_offload: bool = False,  # cpu offload
        dp_process_group: Optional[ProcessGroup] = None,  # the dp pg for comm
        tp_process_group: Optional[ProcessGroup] = None,  # if using tp
        pp_process_group: Optional[ProcessGroup] = None,
        forced_dtype: Optional[torch.dtype] = None,
        moe_extra_dp_process_group: Optional[ProcessGroup] = None,
    ):
        self.param_info = param_info
        self.stage_manager = model.stage_manager
        self.shared_params = model.shared_params
        self.dp_pg = dp_process_group
        self.tp_pg = tp_process_group
        self.pp_pg = pp_process_group
        if use_pipeline:
            init_pipeline_optimizer(optimizer, model)

        zero_params = list(filter(lambda x: not is_moe_tensor(x), model.parameters()))
        moe_params = list(filter(lambda x: is_moe_tensor(x), model.parameters()))

        optimizer.param_groups.clear()
        optimizer.add_param_group({"params": zero_params})
        optimizer.add_param_group({"params": moe_params})
        strategies = [
            LowLevelOptStrategy(
                param_group=optimizer.param_groups[0],
                process_group=dp_process_group,
                reduce_bucket_size=reduce_bucket_size,
                communication_dtype=communication_dtype,
                overlap_communication=overlap_communication,
                partition_grad=partition_grad,
                cpu_offload=cpu_offload,
            ),
            MoeZeroStrategy(
                param_group=optimizer.param_groups[1],
                process_group=moe_extra_dp_process_group,
                reduce_bucket_size=reduce_bucket_size,
                communication_dtype=communication_dtype,
                overlap_communication=overlap_communication,
                partition_grad=partition_grad,
                cpu_offload=cpu_offload,
            ),
        ]
        super().__init__(
            optimizer=optimizer,
            group_strategies=strategies,

Review comment (bug):
perhaps for now we can throw runtime error here in case that user have more than one torch.Optimizer.param_group?

Code after:
        clip_grad_norm: float = 0.0,  # grad clipping
        verbose: bool = False,
        reduce_bucket_size: int = 1024 * 1024,  # communication
        communication_dtype: Optional[torch.dtype] = None,
        overlap_communication: bool = True,
        partition_grad: bool = False,  # stage 2 flag
        cpu_offload: bool = False,  # cpu offload
        dp_process_group: Optional[ProcessGroup] = None,  # the dp pg for comm
        tp_process_group: Optional[ProcessGroup] = None,  # if using tp
        pp_process_group: Optional[ProcessGroup] = None,
        forced_dtype: Optional[torch.dtype] = None,
        moe_extra_dp_process_group: Optional[ProcessGroup] = None,
    ):
        self.param_info = param_info
        self.stage_manager = model.stage_manager
        self.shared_params = model.shared_params
        self.dp_pg = dp_process_group
        self.tp_pg = tp_process_group
        self.pp_pg = pp_process_group
        if use_pipeline:
            init_pipeline_optimizer(optimizer, model)

        assert (
            len(optimizer.param_groups) == 1
        ), "Currently only one parameter group is supported, and we will support multiple groups later."
        zero_params = list(filter(lambda x: not is_moe_tensor(x), model.parameters()))
        moe_params = list(filter(lambda x: is_moe_tensor(x), model.parameters()))

        optimizer.param_groups.clear()
        optimizer.add_param_group({"params": zero_params})
        optimizer.add_param_group({"params": moe_params})
        strategies = [
            LowLevelOptStrategy(
                param_group=optimizer.param_groups[0],
                process_group=dp_process_group,
                reduce_bucket_size=reduce_bucket_size,
                communication_dtype=communication_dtype,
                overlap_communication=overlap_communication,
                partition_grad=partition_grad,
                cpu_offload=cpu_offload,
            ),
            MoeZeroStrategy(
                param_group=optimizer.param_groups[1],
                process_group=moe_extra_dp_process_group,
                reduce_bucket_size=reduce_bucket_size,
                communication_dtype=communication_dtype,
                overlap_communication=overlap_communication,
                partition_grad=partition_grad,
                cpu_offload=cpu_offload,
            ),
        ]