def _compute_sensitivity_impl_multiple_workers(
        self,
        rank: int,
        num_workers: int,
        dataset: _List,
        loss_fn: _Callable,
        sensitivity_path: _Optional[str] = None,
        fsdp_auto_wrap_policy: _Optional[_FSDPAutoWrapPolicy] = None,
    ):
        """
        Computes sensitivity for the model weights using multiple processes.
        This mode is useful for large models for which computing gradients on a single
        process is infeasible because the model does not fit on a single GPU. The model is
        sharded on multiple GPUs using :py:class:`FullyShardedDataParallel`, which enables
        distributed computation of gradients.

        If ``sensitivity_path`` is passed as ``None``, the sensitivity matrices are
        stored temporarily and deleted after model compression. Otherwise, they are
        saved at the location specified by ``sensitivity_path``.

        Args:
            rank (:obj:`int`):  Rank of the worker process on which this function is executed
            num_workers (:obj:`int`): Number of workers used for computing sensitivity
            dataset (:py:class:`Iterable`): An iterable where each element
                is an input to the model to be compressed. Used for computing gradients of model weights.
            loss_fn (:obj:`Callable`): A callable which takes the model and data as input and performs
                a forward pass on the model and computes the training loss
            sensitivity_path (:obj:`str` or ``None``): An optional path for saving the sensitivity
                of weights. Defaults to ``None``.
            fsdp_auto_wrap_policy (:py:class:`_FSDPAutoWrapPolicy` or ``None``): Policy to apply
                :py:class:`FullyShardedDataParallel` to submodules of ``model``. Defaults to ``None``.
        """
        _ddp_setup(rank, num_workers)
        auto_wrap_policy = (
            fsdp_auto_wrap_policy.get_policy() if fsdp_auto_wrap_policy is not None else None
        )
        model = _FSDP(
            module=self._model,
            auto_wrap_policy=auto_wrap_policy,
            sharding_strategy=_ShardingStrategy.FULL_SHARD,
            use_orig_params=False,
            device_id=_torch.cuda.current_device(),
            sync_module_states=True,
        )

        # We want to compute squares of gradients of the un-sharded parameters
        # to use later for k-means. However, parameters are sharded and gradients
        # are also computed in the sharded state. And there is no efficient way
        # to un-shard them, hence we use an optimizer to add the sharded gradients
        # to the parameters, which can later be un-sharded when we save the state dict.
        optim = _torch.optim.SGD(
            [param for param in model.parameters() if param.requires_grad], lr=1.0
        )
        optim.zero_grad()

        with self._register_grad_square_hooks(model):
            for didx, data in enumerate(dataset):
                if _is_leader():
                    _logger.info(f"Computing sensitivity using sample {didx}")
                loss = loss_fn(model, data)
                loss.backward()

            # we set the parameters to zero so that when we call optim.step,
            # the parameter values are equal to the square of the gradient
            with _torch.no_grad():
                for param in model.parameters():
                    param.data.zero_()

            optim.step()

            cfg = _FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
            with _FSDP.state_dict_type(model, _StateDictType.FULL_STATE_DICT, cfg):
                sensitivity_dict = model.state_dict()

            if _is_leader():
                _torch.save(sensitivity_dict, self._get_sensitivity_path(sensitivity_path))

        Computes sensitivity for the model weights using multiple processes.
        This mode is useful for large models for which computing gradients on a single
        process is infeasible because the model does not fit on a single GPU. The model is
        sharded on multiple GPUs using :py:class:`FullyShardedDataParallel`, which enables
        distributed computation of gradients.

        If ``sensitivity_path`` is passed as ``None``, the sensitivity matrices are
        stored temporarily and deleted after model compression. Otherwise, they are
        saved at the location specified by ``sensitivity_path``.

        Args:
            rank (:obj:`int`):  Rank of the worker process on which this function is executed
            num_workers (:obj:`int`): Number of workers used for computing sensitivity
            dataset (:py:class:`Iterable`): An iterable where each element
                is an input to the model to be compressed. Used for computing gradients of model weights.
            loss_fn (:obj:`Callable`): A callable which takes the model and data as input and performs
                a forward pass on the model and computes the training loss
            sensitivity_path (:obj:`str` or ``None``): An optional path for saving the sensitivity
                of weights. Defaults to ``None``.
            fsdp_auto_wrap_policy (:py:class:`_FSDPAutoWrapPolicy` or ``None``): Policy to apply
                :py:class:`FullyShardedDataParallel` to submodules of ``model``. Defaults to ``None``.
        
coremltools/optimize/torch/palettization/sensitive_k_means.py