812339ca3d9b68b5f87f9c391aa918f068d5b004,torch/distributed/optim/zero_redundancy_optimizer.py,ZeroRedundancyOptimizer,__init__,#ZeroRedundancyOptimizer#Any#Any#Any#Any#,124

Before Change

        self.optim = optim(self.partition_parameters()[self.rank], **default)

        // - Sync local and global param_groups keys
        for global_group, local_group in zip(
            self.param_groups, self.optim.param_groups
        ):
            for k, v in local_group.items():
                if k != "params":
                    global_group[k] = v

        //  Optional consolidated optimizer state
        self._all_states: List[Dict[str, Any]] = []

        // Current default device is set by the parameters allocated to this rank
        self._device = list(self.per_device_params.keys())[0]

After Change

        self._partition_parameters_cache: List[List[Dict]] = []
        self._index_to_param_cache: Dict[int, torch.Tensor] = {}
        self._all_params = params
        self._reference_is_trainable_mask = list(map(_is_trainable, self._all_params))

        // Build the wrapped optimizer, responsible for a shard of the params
        self.group = group if group is not None else dist.group.WORLD

Instances