U
0-e @ s d dl Z d dlmZ d dlmZmZmZmZ d dlZd dl m
Z
m Z m Z ddl
mZ ddlmZmZmZmZ ddlmZmZmZmZmZmZmZm Z m Z m Z m Z ee Z!dd ddd dd d d dddd
d d Z"i Z#e#$ D ] \Z%Z&e d e%re"'e& qG d
d de
Z(G dd de Z)G dd dZ*G dd de e*Z+ed drdd dl,m- m.Z/ G dd de/j0Z1G dd de e*Z2d$e eej3 ee4 ee4 e5e5eeee6ef ee5 e5ee e d d d Z7G d d d e
Z8G d d! d!e Z9d%d"d#Z:dS )& N)suppress)CallableListOptionalUnion) BatchSampler
DataLoaderIterableDataset )
get_logger)AcceleratorStateDistributedType
GradientStateis_tpu_available) RNGType broadcastbroadcast_object_list concatenatefind_batch_sizeget_data_structureinitialize_tensorsis_torch_versionsend_to_device
slice_tensorssynchronize_rng_statesF )
batch_sizeshufflesampler
batch_samplerZ num_workersZ
collate_fnZ
pin_memory drop_lasttimeoutZworker_init_fnZmultiprocessing_context generatorZprefetch_factorZpersistent_workers>=c @ sT e Zd ZdZdeeeeedddZed d
Z d d Z
d
d Z dd Z dd Z
dS )BatchSamplerSharda
Wraps a PyTorch `BatchSampler` to generate batches for one of the processes only. Instances of this class will
always yield a number of batches that is a round multiple of `num_processes` and that all have the same size.
Depending on the value of the `drop_last` attribute of the batch sampler passed, it will either stop the iteration
at the first batch that would be too small / not present on all processes or loop with indices from the beginning.
Args:
batch_sampler (`torch.utils.data.sampler.BatchSampler`):
The batch sampler to split in several shards.
num_processes (`int`, *optional*, defaults to 1):
The number of processes running concurrently.
process_index (`int`, *optional*, defaults to 0):
The index of the current process.
split_batches (`bool`, *optional*, defaults to `False`):
Whether the shards should be created by splitting a batch to give a piece of it on each process, or by
yielding different full batches on each process.
On two processes with a sampler of `[[0, 1, 2, 3], [4, 5, 6, 7]]`, this will result in:
- the sampler on process 0 to yield `[0, 1, 2, 3]` and the sampler on process 1 to yield `[4, 5, 6, 7]` if
this argument is set to `False`.
- the sampler on process 0 to yield `[0, 1]` then `[4, 5]` and the sampler on process 1 to yield `[2, 3]`
then `[6, 7]` if this argument is set to `True`.
even_batches (`bool`, *optional*, defaults to `True`):
Whether or not to loop back at the beginning of the sampler when the number of samples is not a round
multiple of (original batch size / number of processes).
`BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
equal to `False`
r
r FT)r
num_processes
process_index
split_batches even_batchesc C s |r*|j | dkr*td|j d| d|| _|| _|| _|| _|| _t|dd | _ t|dd| _| j d kr|| jr|tdd S ) Nr zDTo use `BatchSamplerShard` in `split_batches` mode, the batch size (;) needs to be a round multiple of the number of processes ().r r FzNYou need to use `even_batches=False` when the batch sampler has no batch size.) r
ValueErrorr r% r& r' r( getattrr )selfr r% r& r' r( r. W/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/accelerate/data_loader.py__init__f s zBatchSamplerShard.__init__c C s
t | jS Nlenr r- r. r. r/ total_length} s z BatchSamplerShard.total_lengthc C s | j rt| jS t| j| j dkr4t| j| j S t| j| j }| jrN|S | jr\|d S | jt| j| j k rz|d S |S d S Nr r
)r' r3 r r% r r( r& r- lengthr. r. r/ __len__ s
zBatchSamplerShard.__len__c C s | j r| S | S r1 )r' _iter_with_split_iter_with_no_splitr4 r. r. r/ __iter__ s zBatchSamplerShard.__iter__c c s g }| j j| j }t| j D ]@\}}|dkr0|}t|| jkr ||| j || jd V q | jst|dkrt|| jk r| jst||| j kr||| j || jd V n>t|| jk r||7 }q|| }||| j || jd V d S r6 )r r r% enumerater3 r& r r( )r- initial_dataZ batch_lengthidxbatchr. r. r/ r: s
z"BatchSamplerShard._iter_with_splitc c sl g }g }t | jD ]j\}}| js2|| jk r2||7 }|| j | jkrF|}|| j | jd kr| jd ksrt|| jkr|V g }q| jsht|dkrh| jst|dkr|V nt|| jkr|V t|| j| j k r||7 }qt|| jkrg }|d7 }d}|| j dks t|dkrh|| j t| }|||| 7 }|| j | jkrV|V |}g }|d7 }qd S Nr
r )r= r r r% r& r r3 r( )r- r> Zbatch_to_yieldr? r@ Z cycle_indexZ end_indexr. r. r/ r; sD
z%BatchSamplerShard._iter_with_no_splitN)r
r FT)__name__
__module__ __qualname____doc__r intboolr0 propertyr5 r9 r< r: r; r. r. r. r/ r$ C s$ %
r$ c @ s2 e Zd ZdZd eeeeeedddZdd Zd
S ) IterableDatasetSharda
Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will
always yield a number of samples that is a round multiple of the actual batch size (depending of the value of
`split_batches`, this is either `batch_size` or `batch_size x num_processes`). Depending on the value of the
`drop_last` attribute of the batch sampler passed, it will either stop the iteration at the first batch that would
be too small or loop with indices from the beginning.
Args:
dataset (`torch.utils.data.dataset.IterableDataset`):
The batch sampler to split in several shards.
batch_size (`int`, *optional*, defaults to 1):
The size of the batches per shard (if `split_batches=False`) or the size of the batches (if
`split_batches=True`).
drop_last (`bool`, *optional*, defaults to `False`):
Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the
beginning.
num_processes (`int`, *optional*, defaults to 1):
The number of processes running concurrently.
process_index (`int`, *optional*, defaults to 0):
The index of the current process.
split_batches (`bool`, *optional*, defaults to `False`):
Whether the shards should be created by splitting a batch to give a piece of it on each process, or by
yielding different full batches on each process.
On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7]`, this will result in:
- the shard on process 0 to yield `[0, 1, 2, 3]` and the shard on process 1 to yield `[4, 5, 6, 7]` if this
argument is set to `False`.
- the shard on process 0 to yield `[0, 1, 4, 5]` and the sampler on process 1 to yield `[2, 3, 6, 7]` if
this argument is set to `True`.
r
Fr )datasetr r r% r& r' c C sV |r.|dkr.|| dkr.t d| d| d|| _|| _|| _|| _|| _|| _d S )Nr
r zGTo use `IterableDatasetShard` in `split_batches` mode, the batch size (r) r* )r+ rJ r r r% r& r' )r- rJ r r r% r& r' r. r. r/ r0 s z IterableDatasetShard.__init__c c s | j r | jn
| j| j }| j r*| j| j n| j}t| j| | jd | }d }g }| jD ]B}|| t||krX|D ]}|| V qv|d kr| }g }qX| j st|dkr|d kr| }t||k r||7 }q|D ]}|| V qd S rA )
r' r r% ranger& rJ appendr3 copyr )r- Zreal_batch_sizeZprocess_batch_sizeZ
process_slice first_batch
current_batchelementir. r. r/ r< s(
z IterableDatasetShard.__iter__N)r
Fr
r F) rB rC rD rE r rF rG r0 r< r. r. r. r/ rI s # rI c @ s0 e Zd ZdZdd Zdd Zdd Zdd Zd
S ) DataLoaderStateMixina
Mixin class that adds a state to a `DataLoader` to keep track of the status inside the dataloader such as at the
end of the iteration, the number of items in the dataset in the last batch relative to the batch size, and other
useful information that might be needed.
**Available attributes:**
- **end_of_dataloader** (`bool`) -- Whether at the last iteration or batch
- **remainder** (`int`) -- The number of items that are remaining in the last batch, relative to the total
batch size
c K s d| _ d| _d S NFend_of_dataloader remainder)clskwargsr. r. r/ __init_subclass__; s z&DataLoaderStateMixin.__init_subclass__c C s d| _ d| _d S rS rU r4 r. r. r/ reset? s zDataLoaderStateMixin.resetc C sL | tt& t| jdt| j}|| j | _W 5 Q R X | j | dS )z6Prepares the gradient state for the current dataloadertotal_dataset_lengthN)
r[ r Exceptionr, rJ r3 total_batch_sizerW gradient_stateZ_add_dataloaderr7 r. r. r/ beginC s
zDataLoaderStateMixin.beginc C s | j | dS )z9Cleans up the gradient state after exiting the dataloaderN)r_ Z_remove_dataloaderr4 r. r. r/ endK s zDataLoaderStateMixin.endN)rB rC rD rE rZ r[ r` ra r. r. r. r/ rR - s
rR c sF e Zd ZdZd fdd Z fddZedd Zed
d Z Z S )
DataLoaderSharda
Subclass of a PyTorch `DataLoader` that will deal with device placement and current distributed setup.
Args:
dataset (`torch.utils.data.dataset.Dataset`):
The dataset to use to build this datalaoder.
device (`torch.device`, *optional*):
If passed, the device to put all batches on.
rng_types (list of `str` or [`~utils.RNGType`]):
The list of random number generators to synchronize at the beginning of each iteration. Should be one or
several of:
- `"torch"`: the base torch random number generator
- `"cuda"`: the CUDA random number generator (GPU only)
- `"xla"`: the XLA random number generator (TPU only)
- `"generator"`: an optional `torch.Generator`
synchronized_generator (`torch.Generator`, *optional*):
A random number generator to keep synchronized across processes.
split_batches (`int`, *optional*, defaults to 0):
The number of batches to skip at the beginning.
kwargs:
All other keyword arguments to pass to the regular `DataLoader` initialization.
**Available attributes:**
- **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
number of processes
- **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
Nr c s4 t j|f| || _|| _|| _|| _t | _d S r1 )superr0 device rng_typessynchronized_generator skip_batchesr r_ )r- rJ rd re rf rg rY __class__r. r/ r0 q s zDataLoaderShard.__init__c # s | j d k rt| j | j | t }z t|}W n tk
rP d V Y nX d}z>| jd k rnt || j}t|}|| j
kr|V |d7 }|}W qV tk
r d| _ || j
kr|V Y qY qVX qV| d S )Nr r
T)
re r rf r` rc r< next
StopIterationrd r rg rV ra )r- Zdataloader_iterrO batch_index
next_batchrh r. r/ r< y s.
zDataLoaderShard.__iter__c C s< t | jtr| jn| j}t|ddr*|jS |jt|dd S )Nr' Fr% r
)
isinstancer r r r, r )r- r r. r. r/ r^ s
z DataLoaderShard.total_batch_sizec C s" t | jdr| jjS t| jS d S )Nr5 )hasattrrJ r5 r3 r4 r. r. r/ r\ s z$DataLoaderShard.total_dataset_length)NNNr )
rB rC rD rE r0 r< rH r^ r\
__classcell__r. r. rh r/ rb P s
rb )Z check_devicec sN e Zd ZdZeejd fdd Z fddZe dd Z
e d d
Z Z S ) MpDeviceLoaderWrappera
Wrapper for the xpl.MpDeviceLoader class that knows the total batch size.
XLA preloading threads will all call DataLoaderShard's __iter__(). Remove rng_types from DataLoaderShard to
prevent it from using the XLA device in the preloading threads, and synchronize the RNG once from the main
thread only.
**Available attributes:**
- **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
number of processes
- **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
)
dataloaderrd c s$ t || | jj| _d | j_d S r1 )rc r0 _loaderre
_rng_types)r- rr rd rh r. r/ r0 s
z MpDeviceLoaderWrapper.__init__c s$ | j d k rt| j | jj t S r1 )rt r rs rf rc r< r4 rh r. r/ r< s
z MpDeviceLoaderWrapper.__iter__c C s | j jS r1 )rs r^ r4 r. r. r/ r^ s z&MpDeviceLoaderWrapper.total_batch_sizec C s | j jS r1 )rs r\ r4 r. r. r/ r\ s z*MpDeviceLoaderWrapper.total_dataset_length)
rB rC rD rE rb torchrd r0 r< rH r^ r\ rp r. r. rh r/ rq s
rq c sb e Zd ZdZdeed fdd
Zdd Z fd
d Z fd d
Ze dd Z
e dd Z Z S )DataLoaderDispatchera
Subclass of a PyTorch `DataLoader` that will iterate and preprocess on process 0 only, then dispatch on each
process their part of the batch.
Args:
split_batches (`bool`, *optional*, defaults to `False`):
Whether the resulting `DataLoader` should split the batches of the original data loader across devices or
yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of
`num_processes` batches at each iteration). Another way to see this is that the observed batch size will be
the same as the initial `dataloader` if this option is set to `True`, the batch size of the initial
`dataloader` multiplied by `num_processes` otherwise. Setting this option to `True` requires that the batch
size of the `dataloader` is a round multiple of `batch_size`.
skip_batches (`int`, *optional*, defaults to 0):
The number of batches to skip at the beginning of an iteration.
**Available attributes:**
- **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes.
Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total
number of processes
- **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes.
Fr N)r'
_drop_lastc s d}t ddr*ddlm} t||r*|j}t j|f| || _|rXtj j
j j ||d t
| _t | _|| _|| _|d krtn|| _d S )NFr# z1.11.0r )ShufflerIterDataPipe)r )r Z-torch.utils.data.datapipes.iter.combinatoricsrx rn Z_shuffle_enabledrc r0 r' ru utilsdataZgraph_settingsZapply_shuffle_settingsr r_ r staterw rg r slice_fn) r- rJ r' rg rw r| rY r rx rh r. r/ r0 s
z DataLoaderDispatcher.__init__c C s d\}}| j jdkrzP| jr&t|}n0g }t| j jD ]}|t| q6t|dd}t|dg}W q t k
r d dg}Y qX n
d | j
g}t | |d | _
| j
r| js| j s| j jdkrt
|dkrt|dd}t|dg}nd dg}t | ||fS )N)NNr dimFTr
)r{ r& r' rj rK r% rL r r rk _stop_iterationr rw r3 )r- iteratorZbatchesr@ _
batch_infor. r. r/ _fetch_batches s.
z#DataLoaderDispatcher._fetch_batchesc # s | d }tddr"t }n| jjdkr8t }d}d| _d }| |\}}d}|s|| }}| jjdkrt|d }t || jj
}t |dd}| j s|d kr| j
|td| jj| jj| jjd}|d krtd| dt|} | | jj }
| j}|s,| |\}}| jr,|d d kr,d }| j sd|rd| | jj dkrdt||gdd
}|
d 7 }
t| jj|
| jjd |
} | j
|| | jj| jjd}|rd | _| | _|| jkr|V |d 7 }qX| d S ) Nr# z2.0.1r F)Z from_process)r& r% z"Batch does not contain any data (`zM`). At the end of all iterable data available before expected stop iteration.Tr} r
)r` r rc r< r{ r& r r r r rd r rw r| slicer% r+ r r rV rW rg ra ) r- Z
main_iteratorZstop_iterationrN rm Znext_batch_inforl r@ r Zobserved_batch_sizer Z
data_slicerh r. r/ r< # sh
z DataLoaderDispatcher.__iter__c s< t }| jr|S | jr&|| jj S t|| jj S d S r1 )rc r9 r' rw r{ r% mathceil)r- Z whole_lengthrh r. r/ r9 l s
z DataLoaderDispatcher.__len__c C s | j r| jjS | jj| jj S r1 )r' rJ r r% r4 r. r. r/ r^ u s z%DataLoaderDispatcher.total_batch_sizec C s
t | jS r1 )r3 rJ r4 r. r. r/ r\ { s z)DataLoaderDispatcher.total_dataset_length)Fr FN)
rB rC rD rE rG r0 r r< r9 rH r^ r\ rp r. r. rh r/ rv s % I
rv T) rr rd r% r& r'
put_on_devicere dispatch_batchesr( slice_fn_for_dispatchreturnc
s |dkr |sd}n t jt}|r.|s.tdt }
|dkrB|
j}|dkrP|
j}|r jdkr j| dkrtd j d| d j} t | ts jnd} d}
d}|dks|
j t
j krr|srt | trt jd ddk r jj
}t| j j|||d
} npt jt}
|
r jj}n jj}t|d rN|j
dkrHt |_
|j
}|
rZ jn j}t|||||d } d d
dddg|dk r|dkrd |kr|d fddtD }| dkr j|d< |r|s j| n j|d < |r |d t| f|| j| d| np|
r\t| f|rB|
j t
jkrB|nd| j||d| n2t| f|rz|
j t
jkrz|nd| ||d| |
j t
jkrt |S S )a
Wraps a PyTorch `DataLoader` to generate batches for one of the processes only.
Depending on the value of the `drop_last` attribute of the `dataloader` passed, it will either stop the iteration
at the first batch that would be too small / not present on all processes or loop with indices from the beginning.
Args:
dataloader (`torch.utils.data.dataloader.DataLoader`):
The data loader to split across several devices.
device (`torch.device`):
The target device for the returned `DataLoader`.
num_processes (`int`, *optional*):
The number of processes running concurrently. Will default to the value given by
[`~state.AcceleratorState`].
process_index (`int`, *optional*):
The index of the current process. Will default to the value given by [`~state.AcceleratorState`].
split_batches (`bool`, *optional*, defaults to `False`):
Whether the resulting `DataLoader` should split the batches of the original data loader across devices or
yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of
`num_processes` batches at each iteration).
Another way to see this is that the observed batch size will be the same as the initial `dataloader` if
this option is set to `True`, the batch size of the initial `dataloader` multiplied by `num_processes`
otherwise.
Setting this option to `True` requires that the batch size of the `dataloader` is a round multiple of
`batch_size`.
put_on_device (`bool`, *optional*, defaults to `False`):
Whether or not to put the batches on `device` (only works if the batches are nested list, tuples or
dictionaries of tensors).
rng_types (list of `str` or [`~utils.RNGType`]):
The list of random number generators to synchronize at the beginning of each iteration. Should be one or
several of:
- `"torch"`: the base torch random number generator
- `"cuda"`: the CUDA random number generator (GPU only)
- `"xla"`: the XLA random number generator (TPU only)
- `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your
dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type.
dispatch_batches (`bool`, *optional*):
If set to `True`, the datalaoder prepared is only iterated through on the main process and then the batches
are split and broadcast to each process. Will default to `True` when the underlying dataset is an
`IterableDataset`, `False` otherwise.
even_batches (`bool`, *optional*, defaults to `True`):
If set to `True`, in cases where the total batch size across all processes does not exactly divide the
dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
all workers.
slice_fn_for_dispatch (`Callable`, *optional*`):
If passed, this function will be used to slice tensors across `num_processes`. Will default to
[`~utils.slice_tensors`]. This argument is used only when `dispatch_batches` is set to `True` and will be
ignored otherwise.
Returns:
`torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
`BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
equal to `False`
NFz s z'prepare_data_loader..)r' r rw r| )rd r r re rf )rd r re rf ) rn rJ r r+ r r% r& r r Zdistributed_typer
Z MEGATRON_LMr, r" rI r r r ro ru Generatorr$ remover poprv rb ZTPUrq )rr rd r% r& r' r re r r( r r{ Z new_datasetnew_batch_samplersampler_is_batch_samplerrf r r rY r. r r/ prepare_data_loader s K
r c @ s6 e Zd ZdZd ddZdd Zedd Zd d
Zd S )
SkipBatchSamplerzx
A `torch.utils.data.BatchSampler` that skips the first `n` batches of another `torch.utils.data.BatchSampler`.
r c C s || _ || _d S r1 )r rg )r- r rg r. r. r/ r0 I s zSkipBatchSampler.__init__c c s( t | jD ]\}}|| jkr
|V q
d S r1 )r= r rg )r- indexZsamplesr. r. r/ r< M s
zSkipBatchSampler.__iter__c C s
t | jS r1 r2 r4 r. r. r/ r5 R s z SkipBatchSampler.total_lengthc C s t | j| j S r1 )r3 r rg r4 r. r. r/ r9 V s zSkipBatchSampler.__len__N)r ) rB rC rD rE r0 r< rH r5 r9 r. r. r. r/ r D s
r c s. e Zd ZdZd fdd Z fddZ ZS )SkipDataLoadera
Subclass of a PyTorch `DataLoader` that will skip the first batches.
Args:
dataset (`torch.utils.data.dataset.Dataset`):
The dataset to use to build this datalaoder.
skip_batches (`int`, *optional*, defaults to 0):
The number of batches to skip at the beginning.
kwargs:
All other keyword arguments to pass to the regular `DataLoader` initialization.
r c s t j|f| || _d S r1 )rc r0 rg )r- rJ rg rY rh r. r/ r0 g s zSkipDataLoader.__init__c # s, t t D ]\}}|| jkr|V qd S r1 )r= rc r< rg )r- r r@ rh r. r/ r< k s
zSkipDataLoader.__iter__)r )rB rC rD rE r0 r< rp r. r. rh r/ r Z s r c sN j }d}t|trd}n(t jt}|r0 jn j}t||d}dddddg fd d
tD }|dkr j|d< j |d< t t
r|dkr||d < t
|f j | j d | nt t
r|dkr||d < n |r||d< j |d< n||d< t
|f j j jd
| n4|dkr6t|fd |i| nt|fd|i| S )ze
Creates a `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`.
FN)rg r r r r r c s& i | ] }|kr|t |t| qS r. r r r r. r/ r s z&skip_first_batches..rg )r' r rw )rd re rf )rJ rn r r r r r r r r rv r' rw rb rd re rf r r )rr Z num_batchesrJ r r r rY r. r r/ skip_first_batchesq sf
r ) NNNFFNNTN)r );r
contextlibr typingr r r r ru Ztorch.utils.datar r r loggingr r{ r r
r r ry r r r r r r r r r r r rB loggerr Z%_PYTORCH_DATALOADER_ADDITIONAL_KWARGSitemsvZadditional_kwargsupdater$ rI rR rb Z%torch_xla.distributed.parallel_loader distributedZparallel_loaderZxplZMpDeviceLoaderrq rv rd rF rG strr r r r r. r. r. r/ s~ 4
Q#W % 4 E