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

Code before:

    def __init__(self, module: Module, precision: str, shard_config: ShardConfig, dp_group: ProcessGroup,
                 use_ddp: bool) -> None:
        self.stage_manager = shard_config.pipeline_stage_manager
        self.dp_group = dp_group
        shardformer = ShardFormer(shard_config)
        module, self.shared_params = shardformer.optimize(module)
        # TODO(ver217): add input type cast
        self.shared_param_process_groups = []
        for shared_param in self.shared_params:
            if len(shared_param) > 0:
                self.shared_param_process_groups.append(
                    self.stage_manager.init_process_group_by_stages(list(shared_param.keys())))
        if precision == 'fp16':
            module = module.half().cuda()
        elif precision == 'bf16':
            module = module.to(dtype=torch.bfloat16).cuda()
        else:
            module = module.cuda()    # train without AMP

        if use_ddp:

            ddp_config = dict(process_group=dp_group,
                              broadcast_buffers=True,
                              bucket_cap_mb=25,
                              find_unused_parameters=False,
                              check_reduction=False,
                              gradient_as_bucket_view=False,
                              static_graph=False)

            # convert model to sync bn
            module = SyncBatchNorm.convert_sync_batchnorm(module, dp_group)

            # wrap the model with PyTorch DDP
            module = DDP(module, **ddp_config)

        super().__init__(module)

    def sync_shared_params(self):
        for shared_param, group in zip(self.shared_params, self.shared_param_process_groups):
            if self.stage_manager.stage in shared_param:
                param = shared_param[self.stage_manager.stage]
                dist.all_reduce(param.grad, group=group)
            dist.barrier()

    def no_sync(self) -> Iterator[None]:
        # no sync grads across data parallel
        return nullcontext()

    def sync_grads(self):
        # sync grad across data parallel

Review comment (suggestion):
`find_unused_parameters=False` and `static_graph=False` may lead to correctness issue when there are unused params. This is commonly observed when layer drop is used.

Code after:

    def __init__(self, module: Module, precision: str, shard_config: ShardConfig, dp_group: ProcessGroup, use_ddp: bool,
                 ddp_config: dict) -> None:
        self.stage_manager = shard_config.pipeline_stage_manager
        self.dp_group = dp_group
        shardformer = ShardFormer(shard_config)
        module, self.shared_params = shardformer.optimize(module)
        # TODO(ver217): add input type cast
        self.shared_param_process_groups = []
        for shared_param in self.shared_params:
            if len(shared_param) > 0:
                self.shared_param_process_groups.append(
                    self.stage_manager.init_process_group_by_stages(list(shared_param.keys())))
        if precision == 'fp16':
            module = module.half().cuda()
        elif precision == 'bf16':
            module = module.to(dtype=torch.bfloat16).cuda()
        else:
            module = module.cuda()    # train without AMP

        if use_ddp:

            # convert model to sync bn
            module = SyncBatchNorm.convert_sync_batchnorm(module, dp_group)

            # wrap the model with PyTorch DDP
            module = DDP(module, process_group=dp_group, **ddp_config)

        super().__init__(module)

    def sync_shared_params(self):
        for shared_param, group in zip(self.shared_params, self.shared_param_process_groups):
            if self.stage_manager.stage in shared_param:
                param = shared_param[self.stage_manager.stage]
                dist.all_reduce(param.grad, group=group)
            dist.barrier()

    def no_sync(self) -> Iterator[None]:
        # no sync grads across data parallel
        return nullcontext()

    def sync_grads(self):
        # sync grad across data parallel
        if self.dp_group.size() == 1:
            return
        for p in self.module.parameters():
            if p.grad is not None:
                dist.all_reduce(p.grad, group=self.dp_group)
                p.grad.div_(self.dp_group.size())

    def unwrap(self):