&`iq 6ddlZddlZddlZddlZddlZddlmZmZmZm Z m Z m Z ddl Z ddlZddlmZddlmZmZddlmZddlmZddlmZmZmZmZmZddlmZm Z dd l!m"Z"m#Z#dd l$m%Z%dd l&m'Z'dd l(m)Z)m*Z*dd l+m,Z,ddl-m.Z.m/Z/ej0e1Z2e/ddej3fdZ4e/dde ej3fdZ5e/d d2dej6j7de e8ej3fde e9de ee9efdej6j7f dZ:e/d d3dej;j<jde8de8de8dej;j<jf dZ=d Z>e.d4d"e8ddfd#Z?e.d$ej@jdej@jfd%ZAe.d&ejBddfd'ZCe/dd5d(eDddfd)ZEe.Gd*d+ZFGd,d-e'ZGGd.d/eZHGd0d1eZIdS)6N)AnyCallableDictListOptionalUnion)Version) GradScalerautocast)DistributedDataParallel) Optimizer) DataLoaderDistributedSamplerIterableDataset RandomSamplerSequentialSampler)TagKeyrecord_extra_usage_tag)#get_torch_device_manager_by_context'get_torch_device_manager_by_device_type)session) Accelerator)get_acceleratorset_accelerator)_log_deprecation_warning) Deprecated PublicAPIstable) stabilityreturncvddlm}ttjd|dS)aGets the correct torch device configured for this process. Returns the torch device for the current worker. If more than 1 GPU is requested per worker, returns the device with the minimal device index. .. note:: If you requested multiple GPUs per worker, and want to get the full list of torch devices, please use :meth:`~ray.train.torch.get_devices`. Assumes that `CUDA_VISIBLE_DEVICES` is set and is a superset of the `ray.get_gpu_ids()`. Examples: Example: Launched 2 workers on the current node, each with 1 GPU .. testcode:: :skipif: True os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" ray.get_gpu_ids() == [2] torch.cuda.is_available() == True get_device() == torch.device("cuda:0") Example: Launched 4 workers on the current node, each with 1 GPU .. testcode:: :skipif: True os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" ray.get_gpu_ids() == [2] torch.cuda.is_available() == True get_device() == torch.device("cuda:2") Example: Launched 2 workers on the current node, each with 2 GPUs .. testcode:: :skipif: True os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" ray.get_gpu_ids() == [2,3] torch.cuda.is_available() == True get_device() == torch.device("cuda:2") You can move a model to device by: .. testcode:: :skipif: True model.to(ray.train.torch.get_device()) Instead of manually checking the device type: .. testcode:: :skipif: True model.to("cuda" if torch.cuda.is_available() else "cpu") r torch_utils1)ray.air._internalr#rrTRAIN_TORCH_GET_DEVICE get_devicesr"s t/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/ray/train/torch/train_loop_utils.py get_devicer)$sA~.-----68#>>>  " " $ $Q ''betacjddlm}ttjd|S)aGets the correct torch device list configured for this process. Assumes that `CUDA_VISIBLE_DEVICES` is set and is a superset of the `ray.get_gpu_ids()`. Examples: Example: Launched 2 workers on the current node, each with 1 GPU .. testcode:: :skipif: True os.environ["CUDA_VISIBLE_DEVICES"] == "2,3" ray.get_gpu_ids() == [2] torch.cuda.is_available() == True get_devices() == [torch.device("cuda:0")] Example: Launched 4 workers on the current node, each with 1 GPU .. testcode:: :skipif: True os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" ray.get_gpu_ids() == [2] torch.cuda.is_available() == True get_devices() == [torch.device("cuda:2")] Example: Launched 2 workers on the current node, each with 2 GPUs .. testcode:: :skipif: True os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" ray.get_gpu_ids() == [2,3] torch.cuda.is_available() == True get_devices() == [torch.device("cuda:2"), torch.device("cuda:3")] rr"r$)r%r#rrTRAIN_TORCH_GET_DEVICESr'r"s r(r'r'is<R.-----693???  " " $ $$r*Tddpmodelmove_to_deviceparallel_strategyparallel_strategy_kwargsc |dkr9ttjtdkrtdt t jdtt ||||S)aPrepares the model for distributed execution. This allows you to use the same exact code regardless of number of workers or the device type being used (CPU, GPU). Args: model (torch.nn.Module): A torch model to prepare. move_to_device: Either a boolean indiciating whether to move the model to the correct device or an actual device to move the model to. If set to False, the model needs to manually be moved to the correct device. parallel_strategy ("ddp", "fsdp", or None): Whether to wrap models in ``DistributedDataParallel``, ``FullyShardedDataParallel``, or neither. parallel_strategy_kwargs (Dict[str, Any]): Args to pass into ``DistributedDataParallel`` or ``FullyShardedDataParallel`` initialization if ``parallel_strategy`` is set to "ddp" or "fsdp", respectively. fsdpz1.11.0zsFullyShardedDataParallel requires torch>=1.11.0. Run `pip install 'torch>=1.11.0'` to use FullyShardedDataParallel.r$)r0r1r2) r torch __version__ ImportErrorrrTRAIN_TORCH_PREPARE_MODELr_TorchAccelerator prepare_model)r/r0r1r2s r(r:r:s4F""wu/@'A'AGHDUDU'U'U Q   6;SAAA , - - ; ; %+!9 <  r* data_loaderadd_dist_sampler auto_transfercttjdtt||||S)a Prepares :class:`~torch.utils.data.DataLoader` for distributed execution. This allows you to use the same exact code regardless of number of workers or the device type being used (CPU, GPU). .. note:: This method adds a `DistributedSampler` to the `DataLoader` if the number of training workers is greater than 1. If shuffling is enabled on the original `DataLoader`, then `shuffle=True` will also be passed into the `DistributedSampler` constructor. `shuffle=False` on the original `DataLoader` also means that shuffling is disabled on the sampler. With more than 1 worker, calling the `DistributedSampler.set_epoch` method at the beginning of each epoch before creating the DataLoader iterator is necessary to make shuffling work properly across multiple epochs. Otherwise, the same ordering will be always used. See: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler # noqa: E501 Example: .. testcode: :skipif: True import torch import ray.train.torch train_dataloader = torch.utils.data.DataLoader( ..., batch_size=..., shuffle=True ) train_dataloader = ray.train.torch.prepare_data_loader(train_loader) for epoch in range(10): if ray.train.get_context().get_world_size() > 1: # Required for the distributed sampler to shuffle properly across epochs train_dataloader.sampler.set_epoch(epoch) for X, y in train_loader: # No need to move data to GPU, this is done by `prepare_data_loader`! # X, y = X.to("cuda"), y.to("cuda") ... Args: data_loader (torch.utils.data.DataLoader): The DataLoader to prepare. add_dist_sampler: Whether to add a DistributedSampler to the provided DataLoader. move_to_device: If set, automatically move the data returned by the data loader to the correct device. auto_transfer: If set and device is GPU, another CUDA stream is created to automatically copy data from host (CPU) memory to device (GPU) memory (the default CUDA stream still runs the training procedure). If device is CPU, it will be disabled regardless of the setting. This configuration will be ignored if ``move_to_device`` is False. r$)r<r0r=)rrTRAIN_TORCH_PREPARE_DATALOADERrr9prepare_data_loader)r;r<r0r=s r(r@r@sKB6@#FFF , - - A A)%# B  r*c0ddlm}t|dS)Nr_TORCH_AMP_DEPRECATION_MESSAGE)#ray.train.v2.torch.train_loop_utilsrCrrBs r(_log_amp_deprecation_warningrE s)RRRRRR;<<<<`_. N)rr9enable_reproducibility)rQs r(rSrSAs'&%&&==dCCCCCr*c0eZdZdZdZddeefdZdS)TorchWorkerProfilerzUtility class for running PyTorch Profiler on a Train worker. Args: trace_dir (Optional[str]): The directory to store traces on the worker node. If ``None``, this will use a default temporary dir. pytorch_profiler_worker_tracesN trace_dirc td)NzGThe `ray.train.torch.TorchWorkerProfiler` API is deprecated in Ray 2.0.)DeprecationWarning)selfrWs r(__init__zTorchWorkerProfiler.__init__bs U   r*N)__name__ __module__ __qualname____doc__WORKER_TRACE_DIR_NAMErstrr[r*r(rUrUWsJ=  (3-      r*rUc 0eZdZdZddefdZ ddejjd ed e e d e e e e fd ejjf d Z ddejjjded eded ejjjf dZded efdZdejd dfdZdded dfdZdS)r9zA utility that implements methods to accelerate PyTorch training. Arguments: amp: If true, perform training with automatic mixed precision. Otherwise, use full precision. FrFcv||_|rtnd|_d|_t |_dSr\)amp_is_enabledr scaler_seedrdevice_manager)rZrFs r(r[z_TorchAccelerator.__init__ps9!&)3jlllt  ACCr*Tr.Nr/r0r1r2r ct|pi}tj}t|tjr|}n+t }t|t r|d}|jr|j ||rV|dkrt d|nt d|| |}d}|jrc|j|_t#|j|_t%|dr |j|_t+j|||_tj}|r|dkr|dkr4t0} |jr|jdkr|g|d|}n5tjst7d dd lm} | } |dkr$t d | jd n#t d | jd | |fi|}|S) aPrepares the model for distributed execution. This allows you to use the same exact code regardless of number of workers or the device type being used (CPU, GPU). Args: model (torch.nn.Module): A torch model to prepare. move_to_device: Whether to move the model to the correct device. If set to False, the model needs to manually be moved to the correct device. parallel_strategy ("ddp", "fsdp", or None): Whether to wrap models in ``DistributedDataParallel``, ``FullyShardedDataParallel`` ( Experimental), or neither. parallel_strategy_kwargs (Dict[str, Any]): Args to pass into ``DistributedDataParallel`` or ``FullyShardedDataParallel`` initialization if ``parallel_strategy`` is set to "ddp" or "fsdp", respectively. rzMoving model to device: ct|dr#|}|d|d<|d=n|j}|d=|d|d<|d=|S)N_original_get_state __getstate___unwrapped_forwardforward)hasattrrl__dict__copyrZstates r(model_get_statez8_TorchAccelerator.prepare_model..model_get_statest233 *0022(-.C(Dn%/00 **,,.)$%9:E) *+Lr*rmr.cpu) device_ids output_devicezhFSDP is only available with GPU-enabled training. Set `use_gpu=True` in your Trainer to train with GPUs.)FullyShardedDataParallelzWrapping provided model in .)rget_local_rank isinstancer5devicer)listri is_available set_deviceloggerinfodebugtorfrornr rprmrltypes MethodTypeget_world_sizer typecudarItorch.distributed.fsdprzr]) rZr/r0r1r2rankr~ru world_size DataParallelrzs r(r:z_TorchAccelerator.prepare_modelvs2$<#Ar %'' nel 3 3 ##FF\\F&$'' #   + + - - 3   * *6 2 2 2  %qyy ?v??@@@@ @@@AAAHHV$$E   $   J(-}E $&HJJu}55EM un-- ?,1,>)"'!1/5!I!IE +--  Da E))6 &3355&+:N:N'-h)/0030, z..00&  LKKKKK7 qyy R,:ORRRSSSS S<;PSSSTTT LCC*BCCE r*r;r<r=cRtj}tj|dkrWt|jt s=t |drt|jts|rfd}||}|rt}t|||}|S)aPrepares DataLoader for distributed execution. This allows you to use the same exact code regardless of number of workers or the device type being used (CPU, GPU). Args: data_loader (torch.utils.data.DataLoader): The DataLoader to prepare. add_dist_sampler: Whether to add a DistributedSampler to the provided DataLoader. move_to_device: If set, automatically move the data returned by the data loader to the correct device. auto_transfer: (Experimental) If set and device is GPU, another CUDA stream is created to automatically copy data from host (CPU) memory to device (GPU) memory (the default CUDA stream still runs the training procedure). If device is CPU, it will be disabled regardless of the setting. This configuration will be ignored if ``move_to_device`` is False. rvdatasetcjt|jt }dttt gdffd}|j}|j}j8||}tj }| jt|jttf}|s3dkr-td|jjjd|j|jd|j|j|j|j|j||t1|j|d }t3d i|S) Nworker_init_fnc"dtffd }|S)N worker_idctjdz}tj|tj|r |dSdS)Nl)r5 initial_seednprandomrQ)r worker_seedrs r(wrapperzk_TorchAccelerator.prepare_data_loader..with_sampler..seeded_worker_init_fn..wrapper"sb&+&8&:&:U&B  {333 K000)6*N95555566r*)int)rrs` r(seeded_worker_init_fnzZ_TorchAccelerator.prepare_data_loader..with_sampler..seeded_worker_init_fns/63666666#Nr*rzThe z will be overwritten with a DistributedSampler. You can disable this by setting `with_sampler` to False in `prepare_data_loader`.F)shuffle) r batch_sizer num_workers collate_fn pin_memory drop_lasttimeoutr generatorsamplerrc)r}rrrrrrrrhr5 Generator manual_seedrrwarning __class__r]rrrrrrrrr) loaderrrrrusing_default_samplerdata_loader_argsrZ world_ranks r( with_samplerz;_TorchAccelerator.prepare_data_loader..with_samplerse)9JKKK #$,XseTk-B$C # # # #CIBW7=7G :)%:%:>%J%JN % 1 1I))$*555(2N%6 $F))%-qNNLv~7@LLL &~"("3$#)#5"("3"("3!'!1%~&4!*1&.'RRR $ $ "55$4555r*) rrget_world_rankr}rrrprrr)_WrappedDataLoader) rZr;r<r0r=rrr~rs ` @r(r@z%_TorchAccelerator.prepare_data_loaders6+-- +--  NN{24FGG  Y// {2ODD  ! 8 68 68 68 68 68 6t',{33K  Q\\F,[&-PPKr*rKc.t||jS)zWraps optimizer to support automatic mixed precision. Args: optimizer (torch.optim.Optimizer): The DataLoader to prepare. Returns: A wrapped optimizer. )rg)_WrappedOptimizerrg)rZrKs r(rMz#_TorchAccelerator.prepare_optimizerSs!4;????r*rNc|jr.|j|dS|dS)zComputes the gradient of the specified tensor w.r.t. graph leaves. Args: tensor (torch.Tensor): Tensor of which the derivative will be computed. N)rfrgscalerP)rZrNs r(rPz_TorchAccelerator.backward^sM    K  f % % . . 0 0 0 0 0 OO     r*rrQc||_tj|tj|t j|tjddtjj_ dtj d<dS)z,Limits sources of nondeterministic behavior.TFz:4096:8CUBLAS_WORKSPACE_CONFIGN) rhr5rrrQruse_deterministic_algorithmsbackendscudnn benchmarkosenviron)rZrQs r(rSz(_TorchAccelerator.enable_reproducibilityisq  $ D t *4000).& 1: ,---r*FTr.N)TTFr)r]r^r_r`boolr[r5nnModulerrbrrr:utilsdatarr@r rMTensorrPrrSrcr*r(r9r9hsDDDDDDD $+0=A llxll$C= l #+4S>": l  llllb"&## mm[%0mm m  m   $ mmmm^ @9 @ @ @ @ @ u|     ::3:t::::::r*r9cPeZdZdedejdefdZdZdZ dZ dZ d Z d Z d S) rbase_dataloaderr~r=c|jt|di||_d|_||_t |j|_|jdkr!|j r||_ nd|_ |jdkr!|j r|j |nd|_ d|_ dS)NrqrwF)rqupdategetattr _dataloaderdataloader_iterr~rrrisupports_stream_auto_transfer create_stream_memcpy_stream next_batch)rZrr~r=s r(r[z_WrappedDataLoader.__init__{s W_j"EEFFF*# EfkRR ;%  D$7$G$G$I$I "/D  "'D {e##(;#   - -f 5 5 5  r*c\|dSfd}jj5t|tjjr!fd|D}nt|trtfd|D}nwt|trfd|D}nSt|tj r ||}n-t dt|d|}|cdddS#1swxYwYdS)Nc |jj}n6#t$r)td|djdYnwxYw|S)N) non_blockingzItem z cannot be moved to device r{)rr~rAttributeErrorrr)irZs r(try_move_devicez;_WrappedDataLoader._move_to_device..try_move_devicesw WDD43FDGG! W W W UQUUt{UUUVVVVV WHs!%0AAcBi|]\}}||Src_move_to_device).0kvrZs r( z6_WrappedDataLoader._move_to_device..s-!V!V!VA!T%9%9!%<%z5_WrappedDataLoader._move_to_device..s1&M&M1t';';A'>'>&M&M&M&M&M&Mr*c:g|]}|Srcrrs r( z6_WrappedDataLoader._move_to_device..s'!H!H!Ha$"6"6q"9"9!H!H!Hr*z Data type z' doesn't support being moved to device.)riget_stream_contextrr} collectionsabcMappingitemstuplerr5rrrr)rZitemritem_on_devices` r(rz"_WrappedDataLoader._move_to_devices <4      3 3D4G H H " "$  788 &!V!V!V!V!V!V!VD%(( &!&&M&M&M&M&M&M&M!M!MD$'' &!H!H!H!H4!H!H!HD%,// &!0!6!6 TdTTT"&! " " " " " " " " " " " " " " " " " "sC*D!!D%(D%c|jdS|j}||j|D](} ||#t $rY%wxYwdSr\)rriget_current_stream wait_stream record_streamr)rZr curr_streamrs r(_wait_for_batchz"_WrappedDataLoader._wait_for_batchs   & F )<<>>  3444   A  ,,,,!       sA A%$A%c*t|jSr\)lenrrZs r(__len__z_WrappedDataLoader.__len__s4#$$$r*cdt|jd}|||_dSr\)nextrrrrZrs r(_prefetch_next_batchz'_WrappedDataLoader._prefetch_next_batchs-$.55 ..z::r*c`t|j|_||Sr\)iterrrrrs r(__iter__z_WrappedDataLoader.__iter__s,#D$455 !!### r*cx|j}|t||||Sr\)r StopIterationrrrs r(__next__z_WrappedDataLoader.__next__sA_    Z((( !!###r*N)r]r^r_rr5r~rr[rrrrrrrcr*r(rrzs)38<PT."""8.%%%;;; r*rceZdZddedeefdZedZej dZedZ e j dZ ed Z e j d Z d Z d Z d ZdZddZdS)rNrKrgc"||_||_dSr\)rKrg)rZrKrgs r(r[z_WrappedOptimizer.__init__s" r*c|jjSr\rKrtrs r(rtz_WrappedOptimizer.states ~##r*c||j_dSr\rrss r(rtz_WrappedOptimizer.states$r*c|jjSr\rK param_groupsrs r(rz_WrappedOptimizer.param_groupss ~**r*c||j_dSr\r)rZrs r(rz_WrappedOptimizer.param_groupss&2###r*c|jjSr\rKdefaultsrs r(r z_WrappedOptimizer.defaultss ~&&r*c||j_dSr\r )rZr s r(r z_WrappedOptimizer.defaultss"*r*c:|j|dSr\)rKadd_param_group)rZ param_groups r(rz!_WrappedOptimizer.add_param_groups &&{33333r*c:|j|dSr\)rKload_state_dict)rZ state_dicts r(rz!_WrappedOptimizer.load_state_dicts &&z22222r*c4|jSr\)rKrrs r(rz_WrappedOptimizer.state_dicts~((***r*c8|jdSr\)rK zero_gradrs r(rz_WrappedOptimizer.zero_grads   """""r*c|j;|j|j||jdS|j|dSr\)rgsteprKr)rZclosures r(rz_WrappedOptimizer.stepsZ ; " K  T^W 5 5 5 K   N   ( ( ( ( (r*r\)r]r^r_r rr r[propertyrtsetterrr rrrrrrcr*r(rrs:)Xj5I$$X$ \%%\%++X+333''X'_++_+444333+++###))))))r*rr)TTTrr)Jrloggingrrrtypingrrrrrrnumpyrr5packaging.versionr torch.cuda.ampr r torch.nn.parallelr torch.optimr torch.utils.datarrrrrray._common.usage.usage_librr ray.air._internal.device_managerrrray.train._internalrray.train._internal.acceleratorrray.train._internal.sessionrrray.train.utilsrray.util.annotationsrr getLoggerr]rr~r)r'rrrrbr:rrr@rErJoptimrMrrPrrSrUr9rrrcr*r(r,s$ ================ %%%%%%////////555555!!!!!!GFFFFFFF('''''777777HHHHHHHH44444466666666  8 $ $ XA(ELA(A(A(A(H V+%T%,'+%+%+%+%\ X15',9= %% 8?%$ ,-% }%'tCH~6 %  X_ %%%%P X" FF!,FFF F  [ FFFFR===   D T     *  K!6 K5;;P K K K  K 8U\8d888 8 XDDDTDDDD*            O:O:O:O:O: O:O:O:d]]]]]]]]@.).).).).) .).).).).)r*