Repository: microsoft/DeepSpeed
File: deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py (Python)

Code before:

        manager = Manager()
        self.shared_overlap_grad_map = manager.dict()
        self.shared_stale_param_map = manager.dict()

        for group_no, group in enumerate(self.optimizer.param_groups):
            for param_idx, param in enumerate(group['params']):
                param.data.share_memory_()
                if not hasattr(param, 'stale_param'):
                    param.stale_param = torch.zeros_like(param.data, dtype=param.dtype, device=param.device)
                    param.stale_param.data.share_memory_()
                    key = (group_no, param_idx)
                    self.shared_stale_param_map[key] = param.stale_param
                if param.overlap_grad is not None:
                    param.overlap_grad[0].data.share_memory_()
                    param.overlap_grad[1].data.share_memory_()
                    key = (group_no, param_idx)
                    self.shared_overlap_grad_map[key] = param.overlap_grad

        param_groups_data = self.optimizer.param_groups
        curr_rank = dist.get_rank()
        total_rank = dist.get_world_size()

        current_process = psutil.Process()
        current_affinity = current_process.cpu_affinity()
        ds_num_cores = 2
        zf_affinity = current_affinity[ds_num_cores:]
        ds_affinity = current_affinity[:ds_num_cores]
        current_process.cpu_affinity(ds_affinity)
        self.process = ctx.Process(
            target=zenflow_optimizer_process,
            args=(self.child_conn, curr_rank, total_rank, param_groups_data, self.shared_overlap_grad_map,
                  self.shared_stale_param_map, zf_affinity),
        )
        self.process.daemon = True
        self.process.start()

        msg = self.parent_conn.recv()
        assert msg["type"] == "ready", "Optimizer process did not initialize correctly."

        self.process_optimizer_established = True

    def wait_last_update_and_copy(self):

        if not hasattr(self, 'parent_conn'):
            return

        if self.micro_step + 1 > self.full_warm_up_rounds and self.first_update_round_after_warmup:
            self.first_update_round_after_warmup = False
            return


Review comment (performance):
Hi, @delock, thanks for this great patch, I am working on validating this binding policy. Glad to see this binding policy achieve better results. However, why should we choose `ds_num_cores = 2`? I see some slow down from main process from my side.

Code after:
        ctx = get_context("spawn")
        self.parent_conn, self.child_conn = Pipe()

        manager = Manager()
        self.shared_overlap_grad_map = manager.dict()
        self.shared_stale_param_map = manager.dict()

        for group_no, group in enumerate(self.optimizer.param_groups):
            for param_idx, param in enumerate(group['params']):
                param.data.share_memory_()
                if not hasattr(param, 'stale_param'):
                    param.stale_param = torch.zeros_like(param.data, dtype=param.dtype, device=param.device)
                    param.stale_param.data.share_memory_()
                    key = (group_no, param_idx)
                    self.shared_stale_param_map[key] = param.stale_param
                if param.overlap_grad is not None:
                    param.overlap_grad[0].data.share_memory_()
                    param.overlap_grad[1].data.share_memory_()
                    key = (group_no, param_idx)
                    self.shared_overlap_grad_map[key] = param.overlap_grad

        param_groups_data = self.optimizer.param_groups
        curr_rank = dist.get_rank()
        total_rank = dist.get_world_size()

        current_process = psutil.Process()
        current_affinity = current_process.cpu_affinity()
        all_affinities = [
            torch.zeros(len(current_affinity),
                        dtype=type(current_affinity[0]),
                        device=get_accelerator().current_device_name()) for _ in range(total_rank)
        ]
        dist.all_gather(
            all_affinities,
            torch.tensor(current_affinity,
                         dtype=type(current_affinity[0]),
                         device=get_accelerator().current_device_name()))
        # When affinity across all ranks are the same, the workers are not binded.  Do a soft bind here
        if self.all_tensors_equal(all_affinities):
            num_phy_cores = psutil.cpu_count(logical=False)
            available_phy_cores = [i for i in current_affinity if i < num_phy_cores]
            num_available_phy_cores = len(available_phy_cores)
            my_rank = curr_rank
            my_size = total_rank
            cores_per_rank = num_available_phy_cores // my_size
            current_affinity = available_phy_cores[my_rank * cores_per_rank:(my_rank + 1) * cores_per_rank]
        pt_num_cores = math.ceil(self.pt_reserved_cores_perc * len(current_affinity))
        if pt_num_cores > 0 and pt_num_cores < len(current_affinity):
            zf_affinity = current_affinity[pt_num_cores:]
            pt_affinity = current_affinity[:pt_num_cores]
        else: