def METHOD_NAME(module, param, use_bucket):
    ddp_state_for_reversed_params = module._ddp_state_for_reversed_params
    def allreduce_with_bucket(grad):
        buckets = module._buckets
        bucket_tensors = module._bucket_tensors
        ddp_state_for_reversed_params[param][0] = True
        for index, bucket in enumerate(buckets):
            deleted = all(ddp_state_for_reversed_params[x][1] for x in bucket)
            if deleted:
                continue
            assert not any(ddp_state_for_reversed_params[x][1] for x in bucket)
            all_params_in_bucket_ready = all(
                ddp_state_for_reversed_params[x][0] for x in bucket
            )
            if all_params_in_bucket_ready:
                for x in bucket:
                    ddp_state_for_reversed_params[x][1] = True
                # NOTE(jianhao)(higher-order-grad):
                # local allreduce doesn't have gradient function, higher-order grad may be unsupported
                flow._C.local_all_reduce(bucket_tensors[index], inplace=True)
            else:
                break
    def allreduce_without_bucket(grad):
        ddp_state_for_reversed_params[param][0] = True
        for cur_param, (ready, deleted) in ddp_state_for_reversed_params.items():
            if deleted:
                continue
            if ready:
                ddp_state_for_reversed_params[cur_param][1] = True
                # NOTE(jianhao)(higher-order-grad): local allreduce doesn't have gradient function, higher-order grad may be unsupported
                if cur_param is param:
                    flow._C.local_all_reduce(grad, True)
                else:
                    flow._C.local_all_reduce(cur_param.grad, True)
            else:
                break
    return allreduce_with_bucket if use_bucket else allreduce_without_bucket