HuggingFace의 Accelerate 라이브러리 소개

from accelerate import Accelerator 

accelerator = Accelerator()

# dummy code to get dataloaders, model & optimizer
train_dataloader, eval_dataloader, model, optimizer = get_everything()

# prepare for DDP using accelerator
train_dataloader, eval_dataloader, model, optimizer  = accelerator.prepare(
	train_dataloader, eval_dataloader, model, optimizer
 )

class Accelerator:
    def __init__(
        self,
        device_placement: bool = True,
        split_batches: bool = False,
        fp16: bool = None,
        cpu: bool = False,
        rng_types: Optional[List[Union[str, RNGType]]] = None,
        kwargs_handlers: Optional[List[KwargsHandler]] = None,
    ):
	# initialize state
        self.state = AcceleratorState(fp16=fp16, cpu=cpu, _from_accelerator=True)

class AcceleratorState:
    _shared_state = {}

    def __init__(self, fp16: bool = None, cpu: bool = False, _from_accelerator: bool = False):
        self.__dict__ = self._shared_state
        if not getattr(self, "initialized", False):
	    if is_tpu_available() and not cpu:
		# setup all TPU related variables
            elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
		# setup all MULTI-GPU related variables
            else:
		# setup single GPU or CPU related variables depending on whether CUDA is available
            self.initialized = True

if is_tpu_available() and not cpu:
   # setup variables for TPU 
elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
   # setup variables for MULTI_GPU
else:
   # setup variables for SINGLE_GPU or CPU depending on whether CUDA is available 

if is_tpu_available() and not cpu:
   # setup variables for TPU 
elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
    self.distributed_type = DistributedType.MULTI_GPU
    if not torch.distributed.is_initialized():
        torch.distributed.init_process_group(backend="nccl")
    self.num_processes = torch.distributed.get_world_size()
    self.process_index = torch.distributed.get_rank()
    self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
    self.device = torch.device("cuda", self.local_process_index)
    torch.cuda.set_device(self.device)
    self.use_fp16 = parse_flag_from_env("USE_FP16", False) if fp16 is None else fp16
else:
   # setup variables for SINGLE_GPU or CPU depending on whether CUDA is available 

    def prepare_model(self, model):
        if self.device_placement:
            model = model.to(self.device)
        if self.distributed_type == DistributedType.MULTI_GPU:
            kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
            model = torch.nn.parallel.DistributedDataParallel(
                model,
                device_ids=[self.local_process_index],
                output_device=self.local_process_index,
                **kwargs,
            )
        if self.native_amp:
            model.forward = torch.cuda.amp.autocast()(model.forward)
        return model

new_dataset = dataloader.dataset
new_batch_sampler = dataloader.batch_sampler
generator = getattr(dataloader, "generator", None)

new_batch_sampler = BatchSamplerShard(
        dataloader.batch_sampler,
        num_processes=num_processes,
        process_index=process_index,
        split_batches=split_batches,
)


class BatchSamplerShard(BatchSampler):
    def __init__(
        self,
        batch_sampler: BatchSampler,
        num_processes: int = 1,
        process_index: int = 0,
        split_batches: bool = False,
    ):
        if split_batches and batch_sampler.batch_size % num_processes != 0:
            raise ValueError(
                f"To use `BatchSamplerShard` in `split_batches` mode, the batch size ({batch_sampler.batch_size}) "
                f"needs to be a round multiple of the number of processes ({num_processes})."
            )
        self.batch_sampler = batch_sampler
        self.num_processes = num_processes
        self.process_index = process_index
        self.split_batches = split_batches
        self.batch_size = batch_sampler.batch_size
        self.drop_last = batch_sampler.drop_last

    def __len__(self):
        if len(self.batch_sampler) % self.num_processes == 0:
            return len(self.batch_sampler) // self.num_processes
        length = len(self.batch_sampler) // self.num_processes
        return length if self.drop_last else length + 1

    def __iter__(self):
        return self._iter_with_split() if self.split_batches else self._iter_with_no_split()

    def _iter_with_no_split(self):
        initial_data = []
        batch_to_yield = []
        for idx, batch in enumerate(self.batch_sampler):
            # We gather the initial indices in case we need to circle back at the end.
            if not self.drop_last and idx < self.num_processes:
                initial_data += batch
            # We identify the batch to yield but wait until we ar sure every process gets a full batch before actually
            # yielding it.
            if idx % self.num_processes == self.process_index:
                batch_to_yield = batch
            if idx % self.num_processes == self.num_processes - 1 and len(batch) == self.batch_size:
                yield batch_to_yield
                batch_to_yield = []


if idx % self.num_processes == self.process_index:
    batch_to_yield = batch

class DataLoaderShard(DataLoader):
    def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwargs):
        super().__init__(dataset, **kwargs)
        self.device = device
        self.rng_types = rng_types
        self.generator = generator

    def __iter__(self):
        if self.rng_types is not None:
            synchronize_rng_states(self.rng_types, self.generator)
        state = AcceleratorState()
        for batch in super().__iter__():
            if state.distributed_type == DistributedType.TPU:
                xm.mark_step()
            yield batch if self.device is None else send_to_device(batch, self.device)

        for batch in super().__iter__():
	    # Ignoring TPU operations
            yield batch if self.device is None else send_to_device(batch, self.device)

def send_to_device(tensor, device):
    if isinstance(tensor, (list, tuple)):
        return type(tensor)(send_to_device(t, device) for t in tensor)
    elif isinstance(tensor, dict):
        return type(tensor)({k: send_to_device(v, device) for k, v in tensor.items()})
    elif not hasattr(tensor, "to"):
        return tensor
    return tensor.to(device)