Repository: microsoft/DeepSpeed
File: deepspeed/runtime/zero/partition_parameters.py (Python)

Code before:
        if self.fwd_module is not None:
            # is a no-op if already registered
            for p in self.params:
                register_external_parameter(self.fwd_module, p)

    def __enter__(self):
        if not self.enabled:
            return
        self.params[0].all_gather(param_list=self.params)
        if self.src_rank is None:
            self._param_versions = [(p, p.data.data_ptr(), p._version) for p in self.params]

    def __exit__(self, *exc):
        if not self.enabled:
            return
        if self.src_rank is None:
            if self._param_versions:
                modified_params = [
                    p for p, data_ptr, version in self._param_versions
                    if p.data.data_ptr() != data_ptr or p._version != version
                ]
                modified_local = bool(modified_params)
                modified_global = modified_local
                if dist.is_initialized():
                    modified_flag = torch.tensor(
                        int(modified_local),
                        device=get_accelerator().current_device_name(),
                    )
                    dist.all_reduce(modified_flag, op=dist.ReduceOp.MAX, group=self.params[0].ds_process_group)
                    modified_global = bool(modified_flag.item())
                if modified_global:
                    self.params[0].partition(param_list=self.params, has_been_updated=False)
                    raise RuntimeError(
                        "Detected in-place modification of ZeRO-3 parameters inside GatheredParameters with "
                        "modifier_rank=None. Use modifier_rank=<rank> to broadcast updates across ranks.")
            self.params[0].partition(param_list=self.params, has_been_updated=False)
            return

        # Broadcast parameters from modifier_rank to all other ranks.
        # NCCL backend requires tensors to be on GPU. If parameters have been moved to a different
        # device (e.g., CPU) inside the context, broadcasting will fail. Users should use
        # modifier_rank=None if they don't need to broadcast updates across ranks.
        expected_device = torch.device(get_accelerator().current_device_name())
        for p in self.params:
            if p.data.device != expected_device:
                raise RuntimeError(
                    f"Parameter {p.ds_id} is on {p.data.device} but broadcast requires it to be on {expected_device}. "
                    f"When using GatheredParameters with modifier_rank set, parameters must remain on "
                    f"the accelerator device. If you don't need to broadcast updates, use modifier_rank=None.")

        handles = [dist.broadcast(p.data, self.src_rank, group=p.ds_process_group, async_op=True) for p in self.params]

Review comment (bug):
I have two concerns about detecting unexpected in-place modifications
1. The perf overhead, especially from all-reduce call. 
2. The use of `_version` which is a pytorch internal API not recommended for developers.

How about if we enable this feature only in debug-mode of the GatheredParameters API? For example, reuse the `sanity_checks` config: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/runtime/zero/config.py#L52

Code after:
            if self.params[0].ds_process_group == dist.get_world_group():
                self.src_rank = modifier_rank
            else:
                # A group was specified; convert DP rank to global rank
                self.src_rank = dist.get_global_rank(self.params[0].ds_process_group, modifier_rank)
        self.fwd_module = fwd_module
        if self.fwd_module is not None:
            # is a no-op if already registered
            for p in self.params:
                register_external_parameter(self.fwd_module, p)

    def __enter__(self):
        if not self.enabled:
            return
        self.params[0].all_gather(param_list=self.params)
        if self.src_rank is None and self.enable_sanity_checks:
            self._param_versions = [(p, p.data.data_ptr(), p._version) for p in self.params]

    def __exit__(self, *exc):
        if not self.enabled:
            return
        if self.src_rank is None:
            if self._param_versions:
                modified_params = [
                    p for p, data_ptr, version in self._param_versions
                    if p.data.data_ptr() != data_ptr or p._version != version
                ]
                modified_local = bool(modified_params)
                modified_global = modified_local
                if dist.is_initialized():
                    modified_flag = torch.tensor(
                        int(modified_local),
                        device=get_accelerator().current_device_name(),
                    )
                    dist.all_reduce(modified_flag, op=dist.ReduceOp.MAX, group=self.params[0].ds_process_group)
                    modified_global = bool(modified_flag.item())
                if modified_global:
                    self.params[0].partition(param_list=self.params, has_been_updated=False)
                    raise RuntimeError(
                        "Detected in-place modification of ZeRO-3 parameters inside GatheredParameters with "
                        "modifier_rank=None. Use modifier_rank=<rank> to broadcast updates across ranks.")
            self.params[0].partition(param_list=self.params, has_been_updated=False)
            return

        # Broadcast parameters from modifier_rank to all other ranks.
        # NCCL backend requires tensors to be on GPU. If parameters have been moved to a different
        # device (e.g., CPU) inside the context, broadcasting will fail. Users should use
        # modifier_rank=None if they don't need to broadcast updates across ranks.
        expected_device = torch.device(get_accelerator().current_device_name())
        for p in self.params:
            if p.data.device != expected_device: