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

Code before:
                    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=torch.int32, device=get_accelerator().current_device_name())
            for _ in range(total_rank)
        ]
        dist.all_gather(
            all_affinities,
            torch.tensor(current_affinity, dtype=torch.int32, 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):
            print(f"Recompute worker 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]
        ds_num_cores = self.pt_reserved_cores
        if ds_num_cores > 0 and ds_num_cores < len(current_affinity):
            zf_affinity = current_affinity[ds_num_cores:]
            pt_affinity = current_affinity[:ds_num_cores]
        else:
            zf_affinity = current_affinity
            pt_affinity = current_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()
        print(f"Setting pytorch affinity to {pt_affinity}, OMP_NUM_THREADS={len(pt_affinity)}")
        current_process.cpu_affinity(pt_affinity)
        os.environ['OMP_NUM_THREADS'] = str(len(pt_affinity))

        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'):

Review comment (suggestion):
For clarity, is `DT cores == PT cores`?
```suggestion
        pt_num_cores = self.pt_reserved_cores
```

Code after:

        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:
            zf_affinity = current_affinity
            pt_affinity = current_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()
        current_process.cpu_affinity(pt_affinity)
        os.environ['OMP_NUM_THREADS'] = str(len(pt_affinity))

        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'):