chendl/compositional_test

1import logging2import os3from typing import List, Tuple4 5import numpy as np6import psutil7import torch8import torch.distributed as dist9 10from transformers import RagRetriever11 12 13logger = logging.getLogger(__name__)14 15 16class RagPyTorchDistributedRetriever(RagRetriever):17    """18    A distributed retriever built on top of the ``torch.distributed`` communication package. During training all workers19    initialize their own instance of the retriever, however, only the main worker loads the index into memory. The index is stored20    in cpu memory. The index will also work well in a non-distributed setup.21 22    Args:23        config (:class:`~transformers.RagConfig`):24            The configuration of the RAG model this Retriever is used with. Contains parameters indicating which ``Index`` to build.25        question_encoder_tokenizer (:class:`~transformers.PreTrainedTokenizer`):26            The tokenizer that was used to tokenize the question.27            It is used to decode the question and then use the generator_tokenizer.28        generator_tokenizer (:class:`~transformers.PreTrainedTokenizer`):29            The tokenizer used for the generator part of the RagModel.30        index (:class:`~transformers.models.rag.retrieval_rag.Index`, optional, defaults to the one defined by the configuration):31            If specified, use this index instead of the one built using the configuration32    """33 34    def __init__(self, config, question_encoder_tokenizer, generator_tokenizer, index=None):35        super().__init__(36            config,37            question_encoder_tokenizer=question_encoder_tokenizer,38            generator_tokenizer=generator_tokenizer,39            index=index,40            init_retrieval=False,41        )42        self.process_group = None43 44    def init_retrieval(self, distributed_port: int):45        """46        Retriever initialization function, needs to be called from the training process. The function sets some common parameters47        and environment variables. On top of that, (only) the main process in the process group loads the index into memory.48 49        Args:50            distributed_port (:obj:`int`):51                The port on which the main communication of the training run is carried out. We set the port for retrieval-related52                communication as ``distributed_port + 1``.53        """54 55        logger.info("initializing retrieval")56 57        # initializing a separate process group for retrieval as the default58        # nccl backend doesn't support gather/scatter operations while gloo59        # is too slow to replace nccl for the core gpu communication60        if dist.is_initialized():61            logger.info("dist initialized")62            # needs to be set manually63            os.environ["GLOO_SOCKET_IFNAME"] = self._infer_socket_ifname()64            # avoid clash with the NCCL port65            os.environ["MASTER_PORT"] = str(distributed_port + 1)66            self.process_group = dist.new_group(ranks=None, backend="gloo")67 68        # initialize retriever only on the main worker69        if not dist.is_initialized() or self._is_main():70            logger.info("dist not initialized / main")71            self.index.init_index()72 73        # all processes wait untill the retriever is initialized by the main process74        if dist.is_initialized():75            torch.distributed.barrier(group=self.process_group)76 77    def _is_main(self):78        return dist.get_rank(group=self.process_group) == 079 80    def _scattered(self, scatter_list, target_shape, target_type=torch.float32):81        target_tensor = torch.empty(target_shape, dtype=target_type)82        dist.scatter(target_tensor, src=0, scatter_list=scatter_list, group=self.process_group)83        return target_tensor84 85    def _infer_socket_ifname(self):86        addrs = psutil.net_if_addrs()87        # a hacky way to deal with varying network interface names88        ifname = next((addr for addr in addrs if addr.startswith("e")), None)89        return ifname90 91    def retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, List[dict]]:92        """93        Retrieves documents for specified ``question_hidden_states``. The main process, which has the access to the index stored in memory, gathers queries94        from all the processes in the main training process group, performs the retrieval and scatters back the results.95 96        Args:97            question_hidden_states (:obj:`np.ndarray` of shape :obj:`(batch_size, vector_size)`):98                A batch of query vectors to retrieve with.99            n_docs (:obj:`int`):100                The number of docs retrieved per query.101 102        Output:103            retrieved_doc_embeds (:obj:`np.ndarray` of shape :obj:`(batch_size, n_docs, dim)`104                The retrieval embeddings of the retrieved docs per query.105            doc_ids (:obj:`np.ndarray` of shape :obj:`batch_size, n_docs`)106                The ids of the documents in the index107            doc_dicts (:obj:`List[dict]`):108                The retrieved_doc_embeds examples per query.109        """110 111        # single GPU training112        if not dist.is_initialized():113            doc_ids, retrieved_doc_embeds = self._main_retrieve(question_hidden_states, n_docs)114            return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(doc_ids)115 116        # distributed training117        world_size = dist.get_world_size(group=self.process_group)118 119        # gather logic120        gather_list = None121        if self._is_main():122            gather_list = [torch.empty(question_hidden_states.shape, dtype=torch.float32) for _ in range(world_size)]123        dist.gather(torch.tensor(question_hidden_states), dst=0, gather_list=gather_list, group=self.process_group)124 125        # scatter logic126        n_queries = question_hidden_states.shape[0]127        scatter_ids = []128        scatter_vectors = []129        if self._is_main():130            assert len(gather_list) == world_size131            ids, vectors = self._main_retrieve(torch.cat(gather_list).numpy(), n_docs)132            ids, vectors = torch.tensor(ids), torch.tensor(vectors)133            scatter_ids = self._chunk_tensor(ids, n_queries)134            scatter_vectors = self._chunk_tensor(vectors, n_queries)135        doc_ids = self._scattered(scatter_ids, [n_queries, n_docs], target_type=torch.int64)136        retrieved_doc_embeds = self._scattered(scatter_vectors, [n_queries, n_docs, question_hidden_states.shape[1]])137 138        return retrieved_doc_embeds.numpy(), doc_ids.numpy(), self.index.get_doc_dicts(doc_ids)139