Method __init__

accessory/model/multi_gpu_wrapper.py:172–209

(self,
                 *from_pretrained_args,
                 gpus:Optional[int]=None, gpu_ids:Optional[List[int]]=None,
                 **from_pretrained_kwargs)

    >>> wrapper.generate(["The best programming language is"], max_gen_len=10, temperature=0)
    """
    def __init__(self,
                 *from_pretrained_args,
                 gpus:Optional[int]=None, gpu_ids:Optional[List[int]]=None,
                 **from_pretrained_kwargs):
        if gpus is None and gpu_ids is None:
            raise ValueError('You must specify either gpus or gpu_ids')

        if gpu_ids is None:
            gpu_ids = list(range(gpus))

        self.outer_world = _save_world()
        _reset_world()

        # launch model worker subprocesses and create inner world
        port = accessory.util.misc.find_free_port(10000+int(os.getpid())%100*100, 10100+int(os.getpid())%100*100)
        print(f"Launching {len(gpu_ids)} processes for hosting model with model parallel size {len(gpu_ids)}")
        print("Note that only the output from the FIRST model process will be printed")
        self.model_workers = []
        for rank in range(1, len(gpu_ids) + 1):
            p = subprocess.Popen([sys.executable, __file__, str(port), str(rank), str(len(gpu_ids)+1)])
            self.model_workers.append(p)

        atexit.register(self.on_exit)

        store = dist.TCPStore("127.0.0.1", port, len(gpu_ids) + 1, True)
        dist.init_process_group(
            backend="gloo", rank=0, world_size=len(gpu_ids) + 1,
            # init_method=f"tcp://127.0.0.1:{port}",
            store=store
        )
        self.inner_world = _save_world()


        dist.broadcast_object_list([[gpu_ids, from_pretrained_args, from_pretrained_kwargs]], src=0)
        dist.new_group(ranks=list(range(1, len(gpu_ids) + 1)), backend="nccl")
        dist.barrier()

        _load_world(*self.outer_world)

    def compute_logits(self, *args, **kwargs):
        """