Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,44 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+Utils/JDC/bst.t7 filter=lfs diff=lfs merge=lfs -text
+Utils/PLBERT/step_1050000.t7 filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/01008270.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/kaede_san.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/riamu_zeroshot_02.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/sample_ref01.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/sample_ref02.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/shiki_fine05.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/syuukovoice_200918_3_01.wav filter=lfs diff=lfs merge=lfs -text

.venv/.gitignore

@@ -0,0 +1 @@
+*

.venv/CACHEDIR.TAG

@@ -0,0 +1 @@
+Signature: 8a477f597d28d172789f06886806bc55

.venv/CHANGES.rst

@@ -0,0 +1,76 @@
+CHANGES
+=======
+
+0.4.0 (2024-7-26)
+-------------------
+- Add stub files according to PEP 561 for mypy (thanks @ernix)
+
+0.3.4 (2023-2-18)
+-------------------
+- Fix to support Python2.7 ~ 3.4 (thanks @manjuu-eater)
+- Support Python 3.11
+
+0.3.3 (2022-12-31)
+-------------------
+- Support Python 3.10
+- Re-support Python2.7 ~ 3.4 (thanks @manjuu-eater)
+- Fix z2h, h2z all flag off bug (thanks @manjuu-eater)
+
+0.3.1 (2022-12-14)
+-------------------
+- Fix alpha2kana infinite loop bug (thanks @frog42)
+
+0.3 (2021-03-29)
+-------------------
+- Fix bug (alphabet2kana) thanks @Cuddlemuffin007
+- Support Python 3.8 and 3.9
+- Add handy functions: alphabet2kata and kata2alphabet. thanks @kokimame
+- Add function for julius: hiragana2julius
+
+0.2.4 (2018-02-04)
+-------------------
+- Fix bug (kana2alphabet)
+- Support Python 3.7
+- No longer support Python 2.6
+- Add aliases of z2h -> zenkaku2hankaku and h2z -> hankaku2zenkaku
+
+0.2.3 (2018-02-03)
+-------------------
+- Fix bugs (alphabet2kana, kana2alphabet) thanks @letuananh
+
+0.2.2 (2018-01-22)
+-------------------
+- Fix bug (kana2alphabet) thanks @kokimame
+- Support Python 3.6
+
+0.2.1 (2017-09-14)
+-------------------
+- Fix bugs (alphabet2kana, kana2alphabet)
+
+0.2 (2015-04-02)
+------------------
+
+- Change module name jctconv -> jaconv
+- Add alphabet and hiragana interconvert (alphabet2kana, kana2alphabet)
+
+0.1.1 (2015-03-12)
+------------------
+
+- Support Windows
+- Support Python 3.5
+
+
+0.1 (2014-11-24)
+------------------
+
+- Add some Japanese characters to convert table (ゝゞ・「」。、)
+- Decresing memory usage
+- Some function names are deprecated (hankaku2zenkaku, zenkaku2hankaku, H2K, H2hK, K2H)
+
+
+0.0.7 (2014-03-22)
+------------------
+
+z2h and h2z allow mojimoji-like target character type determination.
+Bug fix about Half Kana conversion.
+

.venv/Lib/site-packages/_soundfile.py

@@ -0,0 +1,11 @@
+# auto-generated file
+import _cffi_backend
+
+ffi = _cffi_backend.FFI('_soundfile',
+    _version = 0x2601,
+    _types = b'\x00\x00\x17\x0D\x00\x00\x6D\x03\x00\x00\x07\x01\x00\x00\x6C\x03\x00\x00\x7A\x03\x00\x00\x00\x0F\x00\x00\x17\x0D\x00\x00\x6F\x03\x00\x00\x07\x01\x00\x00\x03\x11\x00\x00\x00\x0F\x00\x00\x17\x0D\x00\x00\x07\x01\x00\x00\x07\x01\x00\x00\x03\x11\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x17\x0D\x00\x00\x7B\x03\x00\x00\x07\x01\x00\x00\x03\x11\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x6E\x03\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x17\x11\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x6C\x03\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x17\x11\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x17\x11\x00\x00\x6F\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x17\x11\x00\x00\x07\x01\x00\x00\x07\x11\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x17\x11\x00\x00\x07\x01\x00\x00\x04\x11\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x17\x11\x00\x00\x70\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x17\x11\x00\x00\x74\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x17\x11\x00\x00\x02\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x17\x11\x00\x00\x17\x01\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x17\x11\x00\x00\x79\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x17\x11\x00\x00\x04\x11\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x17\x01\x00\x00\x07\x01\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x04\x11\x00\x00\x17\x01\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x3B\x0D\x00\x00\x7A\x03\x00\x00\x17\x01\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x7A\x0D\x00\x00\x17\x11\x00\x00\x00\x0F\x00\x00\x00\x09\x00\x00\x01\x09\x00\x00\x02\x09\x00\x00\x03\x09\x00\x00\x02\x01\x00\x00\x0E\x01\x00\x00\x00\x0B\x00\x00\x01\x0B\x00\x00\x02\x0B\x00\x00\x0D\x01\x00\x00\x56\x03\x00\x00\x5B\x03\x00\x00\x5E\x03\x00\x00\x63\x03\x00\x00\x05\x01\x00\x00\x00\x01\x00\x00\x10\x01',
+    _globals = (b'\xFF\xFF\xFF\x0BSFC_FILE_TRUNCATE',4224,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_INFO',4136,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_MAJOR',4145,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_MAJOR_COUNT',4144,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_SUBTYPE',4147,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_SUBTYPE_COUNT',4146,b'\xFF\xFF\xFF\x0BSFC_GET_LIB_VERSION',4096,b'\xFF\xFF\xFF\x0BSFC_GET_LOG_INFO',4097,b'\xFF\xFF\xFF\x0BSFC_SET_CLIPPING',4288,b'\xFF\xFF\xFF\x0BSFC_SET_SCALE_FLOAT_INT_READ',4116,b'\xFF\xFF\xFF\x0BSFC_SET_SCALE_INT_FLOAT_WRITE',4117,b'\xFF\xFF\xFF\x0BSFM_RDWR',48,b'\xFF\xFF\xFF\x0BSFM_READ',16,b'\xFF\xFF\xFF\x0BSFM_WRITE',32,b'\xFF\xFF\xFF\x0BSF_FALSE',0,b'\xFF\xFF\xFF\x0BSF_FORMAT_ENDMASK',805306368,b'\xFF\xFF\xFF\x0BSF_FORMAT_SUBMASK',65535,b'\xFF\xFF\xFF\x0BSF_FORMAT_TYPEMASK',268369920,b'\xFF\xFF\xFF\x0BSF_TRUE',1,b'\x00\x00\x25\x23sf_close',0,b'\x00\x00\x32\x23sf_command',0,b'\x00\x00\x25\x23sf_error',0,b'\x00\x00\x1D\x23sf_error_number',0,b'\x00\x00\x28\x23sf_error_str',0,b'\x00\x00\x22\x23sf_format_check',0,b'\x00\x00\x19\x23sf_get_string',0,b'\x00\x00\x06\x23sf_open',0,b'\x00\x00\x0B\x23sf_open_fd',0,b'\x00\x00\x00\x23sf_open_virtual',0,b'\x00\x00\x25\x23sf_perror',0,b'\x00\x00\x38\x23sf_read_double',0,b'\x00\x00\x3D\x23sf_read_float',0,b'\x00\x00\x42\x23sf_read_int',0,b'\x00\x00\x51\x23sf_read_raw',0,b'\x00\x00\x4C\x23sf_read_short',0,b'\x00\x00\x51\x23sf_readf_double',0,b'\x00\x00\x51\x23sf_readf_float',0,b'\x00\x00\x51\x23sf_readf_int',0,b'\x00\x00\x51\x23sf_readf_short',0,b'\x00\x00\x47\x23sf_seek',0,b'\x00\x00\x2D\x23sf_set_string',0,b'\x00\x00\x16\x23sf_strerror',0,b'\x00\x00\x20\x23sf_version_string',0,b'\x00\x00\x11\x23sf_wchar_open',0,b'\x00\x00\x38\x23sf_write_double',0,b'\x00\x00\x3D\x23sf_write_float',0,b'\x00\x00\x42\x23sf_write_int',0,b'\x00\x00\x51\x23sf_write_raw',0,b'\x00\x00\x4C\x23sf_write_short',0,b'\x00\x00\x68\x23sf_write_sync',0,b'\x00\x00\x51\x23sf_writef_double',0,b'\x00\x00\x51\x23sf_writef_float',0,b'\x00\x00\x51\x23sf_writef_int',0,b'\x00\x00\x51\x23sf_writef_short',0),
+    _struct_unions = ((b'\x00\x00\x00\x6B\x00\x00\x00\x02SF_FORMAT_INFO',b'\x00\x00\x02\x11format',b'\x00\x00\x07\x11name',b'\x00\x00\x07\x11extension'),(b'\x00\x00\x00\x6C\x00\x00\x00\x02SF_INFO',b'\x00\x00\x3B\x11frames',b'\x00\x00\x02\x11samplerate',b'\x00\x00\x02\x11channels',b'\x00\x00\x02\x11format',b'\x00\x00\x02\x11sections',b'\x00\x00\x02\x11seekable'),(b'\x00\x00\x00\x6D\x00\x00\x00\x02SF_VIRTUAL_IO',b'\x00\x00\x76\x11get_filelen',b'\x00\x00\x75\x11seek',b'\x00\x00\x77\x11read',b'\x00\x00\x78\x11write',b'\x00\x00\x76\x11tell'),(b'\x00\x00\x00\x6E\x00\x00\x00\x10SNDFILE_tag',)),
+    _enums = (b'\x00\x00\x00\x71\x00\x00\x00\x16$1\x00SF_FORMAT_SUBMASK,SF_FORMAT_TYPEMASK,SF_FORMAT_ENDMASK',b'\x00\x00\x00\x72\x00\x00\x00\x16$2\x00SFC_GET_LIB_VERSION,SFC_GET_LOG_INFO,SFC_GET_FORMAT_INFO,SFC_GET_FORMAT_MAJOR_COUNT,SFC_GET_FORMAT_MAJOR,SFC_GET_FORMAT_SUBTYPE_COUNT,SFC_GET_FORMAT_SUBTYPE,SFC_FILE_TRUNCATE,SFC_SET_CLIPPING,SFC_SET_SCALE_FLOAT_INT_READ,SFC_SET_SCALE_INT_FLOAT_WRITE',b'\x00\x00\x00\x73\x00\x00\x00\x16$3\x00SF_FALSE,SF_TRUE,SFM_READ,SFM_WRITE,SFM_RDWR'),
+    _typenames = (b'\x00\x00\x00\x6BSF_FORMAT_INFO',b'\x00\x00\x00\x6CSF_INFO',b'\x00\x00\x00\x6DSF_VIRTUAL_IO',b'\x00\x00\x00\x6ESNDFILE',b'\x00\x00\x00\x3Bsf_count_t',b'\x00\x00\x00\x76sf_vio_get_filelen',b'\x00\x00\x00\x77sf_vio_read',b'\x00\x00\x00\x75sf_vio_seek',b'\x00\x00\x00\x76sf_vio_tell',b'\x00\x00\x00\x78sf_vio_write'),
+)

.venv/Lib/site-packages/_virtualenv.py

@@ -0,0 +1,101 @@
+"""Patches that are applied at runtime to the virtual environment."""
+
+import os
+import sys
+
+VIRTUALENV_PATCH_FILE = os.path.join(__file__)
+
+
+def patch_dist(dist):
+    """
+    Distutils allows user to configure some arguments via a configuration file:
+    https://docs.python.org/3.11/install/index.html#distutils-configuration-files.
+
+    Some of this arguments though don't make sense in context of the virtual environment files, let's fix them up.
+    """  # noqa: D205
+    # we cannot allow some install config as that would get packages installed outside of the virtual environment
+    old_parse_config_files = dist.Distribution.parse_config_files
+
+    def parse_config_files(self, *args, **kwargs):
+        result = old_parse_config_files(self, *args, **kwargs)
+        install = self.get_option_dict("install")
+
+        if "prefix" in install:  # the prefix governs where to install the libraries
+            install["prefix"] = VIRTUALENV_PATCH_FILE, os.path.abspath(sys.prefix)
+        for base in ("purelib", "platlib", "headers", "scripts", "data"):
+            key = f"install_{base}"
+            if key in install:  # do not allow global configs to hijack venv paths
+                install.pop(key, None)
+        return result
+
+    dist.Distribution.parse_config_files = parse_config_files
+
+
+# Import hook that patches some modules to ignore configuration values that break package installation in case
+# of virtual environments.
+_DISTUTILS_PATCH = "distutils.dist", "setuptools.dist"
+# https://docs.python.org/3/library/importlib.html#setting-up-an-importer
+
+
+class _Finder:
+    """A meta path finder that allows patching the imported distutils modules."""
+
+    fullname = None
+
+    # lock[0] is threading.Lock(), but initialized lazily to avoid importing threading very early at startup,
+    # because there are gevent-based applications that need to be first to import threading by themselves.
+    # See https://github.com/pypa/virtualenv/issues/1895 for details.
+    lock = []  # noqa: RUF012
+
+    def find_spec(self, fullname, path, target=None):  # noqa: ARG002
+        if fullname in _DISTUTILS_PATCH and self.fullname is None:
+            # initialize lock[0] lazily
+            if len(self.lock) == 0:
+                import threading
+
+                lock = threading.Lock()
+                # there is possibility that two threads T1 and T2 are simultaneously running into find_spec,
+                # observing .lock as empty, and further going into hereby initialization. However due to the GIL,
+                # list.append() operation is atomic and this way only one of the threads will "win" to put the lock
+                # - that every thread will use - into .lock[0].
+                # https://docs.python.org/3/faq/library.html#what-kinds-of-global-value-mutation-are-thread-safe
+                self.lock.append(lock)
+
+            from functools import partial
+            from importlib.util import find_spec
+
+            with self.lock[0]:
+                self.fullname = fullname
+                try:
+                    spec = find_spec(fullname, path)
+                    if spec is not None:
+                        # https://www.python.org/dev/peps/pep-0451/#how-loading-will-work
+                        is_new_api = hasattr(spec.loader, "exec_module")
+                        func_name = "exec_module" if is_new_api else "load_module"
+                        old = getattr(spec.loader, func_name)
+                        func = self.exec_module if is_new_api else self.load_module
+                        if old is not func:
+                            try:  # noqa: SIM105
+                                setattr(spec.loader, func_name, partial(func, old))
+                            except AttributeError:
+                                pass  # C-Extension loaders are r/o such as zipimporter with <3.7
+                        return spec
+                finally:
+                    self.fullname = None
+        return None
+
+    @staticmethod
+    def exec_module(old, module):
+        old(module)
+        if module.__name__ in _DISTUTILS_PATCH:
+            patch_dist(module)
+
+    @staticmethod
+    def load_module(old, name):
+        module = old(name)
+        if module.__name__ in _DISTUTILS_PATCH:
+            patch_dist(module)
+        return module
+
+
+sys.meta_path.insert(0, _Finder())

.venv/Lib/site-packages/accelerate/__init__.py

@@ -0,0 +1,50 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+__version__ = "1.2.0"
+
+from .accelerator import Accelerator
+from .big_modeling import (
+    cpu_offload,
+    cpu_offload_with_hook,
+    disk_offload,
+    dispatch_model,
+    init_empty_weights,
+    init_on_device,
+    load_checkpoint_and_dispatch,
+)
+from .data_loader import skip_first_batches
+from .inference import prepare_pippy
+from .launchers import debug_launcher, notebook_launcher
+from .state import PartialState
+from .utils import (
+    AutocastKwargs,
+    DataLoaderConfiguration,
+    DDPCommunicationHookType,
+    DeepSpeedPlugin,
+    DistributedDataParallelKwargs,
+    DistributedType,
+    FullyShardedDataParallelPlugin,
+    GradScalerKwargs,
+    InitProcessGroupKwargs,
+    ProfileKwargs,
+    find_executable_batch_size,
+    infer_auto_device_map,
+    is_rich_available,
+    load_checkpoint_in_model,
+    synchronize_rng_states,
+)
+
+
+if is_rich_available():
+    from .utils import rich

.venv/Lib/site-packages/accelerate/big_modeling.py

@@ -0,0 +1,637 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import os
+from contextlib import contextmanager
+from functools import wraps
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .hooks import (
+    AlignDevicesHook,
+    CpuOffload,
+    UserCpuOffloadHook,
+    add_hook_to_module,
+    attach_align_device_hook,
+    attach_align_device_hook_on_blocks,
+)
+from .utils import (
+    OffloadedWeightsLoader,
+    check_cuda_p2p_ib_support,
+    check_device_map,
+    extract_submodules_state_dict,
+    find_tied_parameters,
+    get_balanced_memory,
+    infer_auto_device_map,
+    is_bnb_available,
+    is_mlu_available,
+    is_musa_available,
+    is_npu_available,
+    is_torch_version,
+    is_xpu_available,
+    load_checkpoint_in_model,
+    offload_state_dict,
+    parse_flag_from_env,
+    retie_parameters,
+)
+from .utils.other import recursive_getattr
+
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def init_empty_weights(include_buffers: bool = None):
+    """
+    A context manager under which models are initialized with all parameters on the meta device, therefore creating an
+    empty model. Useful when just initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```python
+    import torch.nn as nn
+    from accelerate import init_empty_weights
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+
+    <Tip warning={true}>
+
+    Any model created under this context manager has no weights. As such you can't do something like
+    `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].
+    Make sure to overwrite the default device_map param for [`load_checkpoint_and_dispatch`], otherwise dispatch is not
+    called.
+
+    </Tip>
+    """
+    if include_buffers is None:
+        include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False)
+    with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f:
+        yield f
+
+
+@contextmanager
+def init_on_device(device: torch.device, include_buffers: bool = None):
+    """
+    A context manager under which models are initialized with all parameters on the specified device.
+
+    Args:
+        device (`torch.device`):
+            Device to initialize all parameters on.
+        include_buffers (`bool`, *optional*):
+            Whether or not to also put all buffers on the meta device while initializing.
+
+    Example:
+
+    ```python
+    import torch.nn as nn
+    from accelerate import init_on_device
+
+    with init_on_device(device=torch.device("cuda")):
+        tst = nn.Linear(100, 100)  # on `cuda` device
+    ```
+    """
+    if include_buffers is None:
+        include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False)
+
+    # TODO(shingjan): remove the torch version check once older versions are deprecated
+    if is_torch_version(">=", "2.0") and include_buffers:
+        with device:
+            yield
+        return
+
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            kwargs["requires_grad"] = param.requires_grad
+            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+
+    def register_empty_buffer(module, name, buffer, persistent=True):
+        old_register_buffer(module, name, buffer, persistent=persistent)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(device)
+
+    # Patch tensor creation
+    if include_buffers:
+        tensor_constructors_to_patch = {
+            torch_function_name: getattr(torch, torch_function_name)
+            for torch_function_name in ["empty", "zeros", "ones", "full"]
+        }
+    else:
+        tensor_constructors_to_patch = {}
+
+    def patch_tensor_constructor(fn):
+        def wrapper(*args, **kwargs):
+            kwargs["device"] = device
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        for torch_function_name in tensor_constructors_to_patch.keys():
+            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+        for torch_function_name, old_torch_function in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)
+
+
+def cpu_offload(
+    model: nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    preload_module_classes: Optional[List[str]] = None,
+):
+    """
+    Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one
+    copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that
+    state dict and put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the model that will be kept on CPU.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+    """
+    if execution_device is None:
+        execution_device = next(iter(model.parameters())).device
+    if state_dict is None:
+        state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()}
+
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=state_dict,
+        preload_module_classes=preload_module_classes,
+    )
+
+    return model
+
+
+def cpu_offload_with_hook(
+    model: torch.nn.Module,
+    execution_device: Optional[Union[int, str, torch.device]] = None,
+    prev_module_hook: Optional[UserCpuOffloadHook] = None,
+):
+    """
+    Offloads a model on the CPU and puts it back to an execution device when executed. The difference with
+    [`cpu_offload`] is that the model stays on the execution device after the forward and is only offloaded again when
+    the `offload` method of the returned `hook` is called. Useful for pipelines running a model in a loop.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device(`str`, `int` or `torch.device`, *optional*):
+            The device on which the model should be executed. Will default to the MPS device if it's available, then
+            GPU 0 if there is a GPU, and finally to the CPU.
+        prev_module_hook (`UserCpuOffloadHook`, *optional*):
+            The hook sent back by this function for a previous model in the pipeline you are running. If passed, its
+            offload method will be called just before the forward of the model to which this hook is attached.
+
+    Example:
+
+    ```py
+    model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device)
+    model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1)
+    model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_device, prev_module_hook=hook_2)
+
+    hid_1 = model_1(input)
+    for i in range(50):
+        # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
+        hid_2 = model_2(hid_1)
+    # model2 is offloaded to the CPU just before this forward.
+    hid_3 = model_3(hid_3)
+
+    # For model3, you need to manually call the hook offload method.
+    hook_3.offload()
+    ```
+    """
+    hook = CpuOffload(execution_device=execution_device, prev_module_hook=prev_module_hook)
+    add_hook_to_module(model, hook, append=True)
+    user_hook = UserCpuOffloadHook(model, hook)
+    return model, user_hook
+
+
+def disk_offload(
+    model: nn.Module,
+    offload_dir: Union[str, os.PathLike],
+    execution_device: Optional[torch.device] = None,
+    offload_buffers: bool = False,
+    preload_module_classes: Optional[List[str]] = None,
+):
+    """
+    Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as
+    memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and
+    put on the execution device passed as they are needed, then offloaded again.
+
+    Args:
+        model (`torch.nn.Module`): The model to offload.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        execution_device (`torch.device`, *optional*):
+            The device on which the forward pass of the model will be executed (should be a GPU). Will default to the
+            model's first parameter device.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+    """
+    if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
+        offload_state_dict(offload_dir, model.state_dict())
+    if execution_device is None:
+        execution_device = next(iter(model.parameters())).device
+    weights_map = OffloadedWeightsLoader(save_folder=offload_dir)
+
+    add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True)
+    attach_align_device_hook(
+        model,
+        execution_device=execution_device,
+        offload=True,
+        offload_buffers=offload_buffers,
+        weights_map=weights_map,
+        preload_module_classes=preload_module_classes,
+    )
+
+    return model
+
+
+def dispatch_model(
+    model: nn.Module,
+    device_map: Dict[str, Union[str, int, torch.device]],
+    main_device: Optional[torch.device] = None,
+    state_dict: Optional[Dict[str, torch.Tensor]] = None,
+    offload_dir: Optional[Union[str, os.PathLike]] = None,
+    offload_index: Optional[Dict[str, str]] = None,
+    offload_buffers: bool = False,
+    skip_keys: Optional[Union[str, List[str]]] = None,
+    preload_module_classes: Optional[List[str]] = None,
+    force_hooks: bool = False,
+):
+    """
+    Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
+    the CPU or even the disk.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to dispatch.
+        device_map (`Dict[str, Union[str, int, torch.device]]`):
+            A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that
+            `"disk"` is accepted even if it's not a proper value for `torch.device`.
+        main_device (`str`, `int` or `torch.device`, *optional*):
+            The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or
+            `"disk"`.
+        state_dict (`Dict[str, torch.Tensor]`, *optional*):
+            The state dict of the part of the model that will be kept on CPU.
+        offload_dir (`str` or `os.PathLike`):
+            The folder in which to offload the model weights (or where the model weights are already offloaded).
+        offload_index (`Dict`, *optional*):
+            A dictionary from weight name to their information (`dtype`/ `shape` or safetensors filename). Will default
+            to the index saved in `save_folder`.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to offload the buffers with the model parameters.
+        skip_keys (`str` or `List[str]`, *optional*):
+            A list of keys to ignore when moving inputs or outputs between devices.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        force_hooks (`bool`, *optional*, defaults to `False`):
+            Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
+            single device.
+    """
+    # Error early if the device map is incomplete.
+    check_device_map(model, device_map)
+
+    # We need to force hook for quantized model that can't be moved with to()
+    if getattr(model, "quantization_method", "bitsandbytes") == "bitsandbytes":
+        # since bnb 0.43.2, we can move 4-bit model
+        if getattr(model, "is_loaded_in_8bit", False) or (
+            getattr(model, "is_loaded_in_4bit", False) and not is_bnb_available(min_version="0.43.2")
+        ):
+            force_hooks = True
+
+    # We attach hooks if the device_map has at least 2 different devices or if
+    # force_hooks is set to `True`. Otherwise, the model in already loaded
+    # in the unique device and the user can decide where to dispatch the model.
+    # If the model is quantized, we always force-dispatch the model
+    if (len(set(device_map.values())) > 1) or force_hooks:
+        if main_device is None:
+            if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
+                main_device = "cpu"
+            else:
+                main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+
+        if main_device != "cpu":
+            cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+            if state_dict is None and len(cpu_modules) > 0:
+                state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+
+        disk_modules = [name for name, device in device_map.items() if device == "disk"]
+        if offload_dir is None and offload_index is None and len(disk_modules) > 0:
+            raise ValueError(
+                "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
+                f"need to be offloaded: {', '.join(disk_modules)}."
+            )
+        if (
+            len(disk_modules) > 0
+            and offload_index is None
+            and (not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")))
+        ):
+            disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
+            offload_state_dict(offload_dir, disk_state_dict)
+
+        execution_device = {
+            name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
+        }
+        execution_device[""] = main_device
+        offloaded_devices = ["disk"] if main_device == "cpu" or main_device == "mps" else ["cpu", "disk"]
+        offload = {name: device in offloaded_devices for name, device in device_map.items()}
+        save_folder = offload_dir if len(disk_modules) > 0 else None
+        if state_dict is not None or save_folder is not None or offload_index is not None:
+            device = main_device if offload_index is not None else None
+            weights_map = OffloadedWeightsLoader(
+                state_dict=state_dict, save_folder=save_folder, index=offload_index, device=device
+            )
+        else:
+            weights_map = None
+
+        # When dispatching the model's parameters to the devices specified in device_map, we want to avoid allocating memory several times for the
+        # tied parameters. The dictionary tied_params_map keeps track of the already allocated data for a given tied parameter (represented by its
+        # original pointer) on each devices.
+        tied_params = find_tied_parameters(model)
+
+        tied_params_map = {}
+        for group in tied_params:
+            for param_name in group:
+                # data_ptr() is enough here, as `find_tied_parameters` finds tied params simply by comparing `param1 is param2`, so we don't need
+                # to care about views of tensors through storage_offset.
+                data_ptr = recursive_getattr(model, param_name).data_ptr()
+                tied_params_map[data_ptr] = {}
+
+                # Note: To handle the disk offloading case, we can not simply use weights_map[param_name].data_ptr() as the reference pointer,
+                # as we have no guarantee that safetensors' `file.get_tensor()` will always give the same pointer.
+
+        attach_align_device_hook_on_blocks(
+            model,
+            execution_device=execution_device,
+            offload=offload,
+            offload_buffers=offload_buffers,
+            weights_map=weights_map,
+            skip_keys=skip_keys,
+            preload_module_classes=preload_module_classes,
+            tied_params_map=tied_params_map,
+        )
+
+        # warn if there is any params on the meta device
+        offloaded_devices_str = " and ".join(
+            [device for device in set(device_map.values()) if device in ("cpu", "disk")]
+        )
+        if len(offloaded_devices_str) > 0:
+            logger.warning(
+                f"Some parameters are on the meta device because they were offloaded to the {offloaded_devices_str}."
+            )
+
+        # Attaching the hook may break tied weights, so we retie them
+        retie_parameters(model, tied_params)
+
+        # add warning to cuda and to method
+        def add_warning(fn, model):
+            @wraps(fn)
+            def wrapper(*args, **kwargs):
+                warning_msg = "You shouldn't move a model that is dispatched using accelerate hooks."
+                if str(fn.__name__) == "to":
+                    to_device = torch._C._nn._parse_to(*args, **kwargs)[0]
+                    if to_device is not None:
+                        logger.warning(warning_msg)
+                else:
+                    logger.warning(warning_msg)
+                for param in model.parameters():
+                    if param.device == torch.device("meta"):
+                        raise RuntimeError("You can't move a model that has some modules offloaded to cpu or disk.")
+                return fn(*args, **kwargs)
+
+            return wrapper
+
+        # Make sure to update _accelerate_added_attributes in hooks.py if you add any hook
+        model.to = add_warning(model.to, model)
+        if is_npu_available():
+            model.npu = add_warning(model.npu, model)
+        elif is_mlu_available():
+            model.mlu = add_warning(model.mlu, model)
+        elif is_musa_available():
+            model.musa = add_warning(model.musa, model)
+        elif is_xpu_available():
+            model.xpu = add_warning(model.xpu, model)
+        else:
+            model.cuda = add_warning(model.cuda, model)
+
+        # Check if we are using multi-gpus with RTX 4000 series
+        use_multi_gpu = len([device for device in set(device_map.values()) if device not in ("cpu", "disk")]) > 1
+        if use_multi_gpu and not check_cuda_p2p_ib_support():
+            logger.warning(
+                "We've detected an older driver with an RTX 4000 series GPU. These drivers have issues with P2P. "
+                "This can affect the multi-gpu inference when using accelerate device_map."
+                "Please make sure to update your driver to the latest version which resolves this."
+            )
+    else:
+        device = list(device_map.values())[0]
+        # `torch.Tensor.to(<int num>)` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)).
+        if is_npu_available() and isinstance(device, int):
+            device = f"npu:{device}"
+        elif is_mlu_available() and isinstance(device, int):
+            device = f"mlu:{device}"
+        elif is_musa_available() and isinstance(device, int):
+            device = f"musa:{device}"
+        elif is_xpu_available() and isinstance(device, int):
+            device = f"xpu:{device}"
+        if device != "disk":
+            model.to(device)
+        else:
+            raise ValueError(
+                "You are trying to offload the whole model to the disk. Please use the `disk_offload` function instead."
+            )
+    # Convert OrderedDict back to dict for easier usage
+    model.hf_device_map = dict(device_map)
+    return model
+
+
+def load_checkpoint_and_dispatch(
+    model: nn.Module,
+    checkpoint: Union[str, os.PathLike],
+    device_map: Optional[Union[str, Dict[str, Union[int, str, torch.device]]]] = None,
+    max_memory: Optional[Dict[Union[int, str], Union[int, str]]] = None,
+    no_split_module_classes: Optional[List[str]] = None,
+    offload_folder: Optional[Union[str, os.PathLike]] = None,
+    offload_buffers: bool = False,
+    dtype: Optional[Union[str, torch.dtype]] = None,
+    offload_state_dict: Optional[bool] = None,
+    skip_keys: Optional[Union[str, List[str]]] = None,
+    preload_module_classes: Optional[List[str]] = None,
+    force_hooks: bool = False,
+    strict: bool = False,
+):
+    """
+    Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
+    loaded and adds the various hooks that will make this model run properly (even if split across devices).
+
+    Args:
+        model (`torch.nn.Module`): The model in which we want to load a checkpoint.
+        checkpoint (`str` or `os.PathLike`):
+            The folder checkpoint to load. It can be:
+            - a path to a file containing a whole model state dict
+            - a path to a `.json` file containing the index to a sharded checkpoint
+            - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint.
+        device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
+            A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
+            name, once a given module name is inside, every submodule of it will be sent to the same device.
+
+            To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For more
+            information about each option see [here](../concept_guides/big_model_inference#designing-a-device-map).
+            Defaults to None, which means [`dispatch_model`] will not be called.
+        max_memory (`Dict`, *optional*):
+            A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU
+            and the available CPU RAM if unset.
+        no_split_module_classes (`List[str]`, *optional*):
+            A list of layer class names that should never be split across device (for instance any layer that has a
+            residual connection).
+        offload_folder (`str` or `os.PathLike`, *optional*):
+            If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as
+            well as the parameters.
+        dtype (`str` or `torch.dtype`, *optional*):
+            If provided, the weights will be converted to that type when loaded.
+        offload_state_dict (`bool`, *optional*):
+            If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
+            the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map
+            picked contains `"disk"` values.
+        skip_keys (`str` or `List[str]`, *optional*):
+            A list of keys to ignore when moving inputs or outputs between devices.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        force_hooks (`bool`, *optional*, defaults to `False`):
+            Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
+            single device.
+        strict (`bool`, *optional*, defaults to `False`):
+            Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's
+            state_dict.
+
+    Example:
+
+    ```python
+    >>> from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+    >>> from huggingface_hub import hf_hub_download
+    >>> from transformers import AutoConfig, AutoModelForCausalLM
+
+    >>> # Download the Weights
+    >>> checkpoint = "EleutherAI/gpt-j-6B"
+    >>> weights_location = hf_hub_download(checkpoint, "pytorch_model.bin")
+
+    >>> # Create a model and initialize it with empty weights
+    >>> config = AutoConfig.from_pretrained(checkpoint)
+    >>> with init_empty_weights():
+    ...     model = AutoModelForCausalLM.from_config(config)
+
+    >>> # Load the checkpoint and dispatch it to the right devices
+    >>> model = load_checkpoint_and_dispatch(
+    ...     model, weights_location, device_map="auto", no_split_module_classes=["GPTJBlock"]
+    ... )
+    ```
+    """
+    if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
+        raise ValueError(
+            "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or "
+            "'sequential'."
+        )
+    if isinstance(device_map, str):
+        if device_map != "sequential":
+            max_memory = get_balanced_memory(
+                model,
+                max_memory=max_memory,
+                no_split_module_classes=no_split_module_classes,
+                dtype=dtype,
+                low_zero=(device_map == "balanced_low_0"),
+            )
+        device_map = infer_auto_device_map(
+            model,
+            max_memory=max_memory,
+            no_split_module_classes=no_split_module_classes,
+            dtype=dtype,
+            offload_buffers=offload_buffers,
+        )
+    if offload_state_dict is None and device_map is not None and "disk" in device_map.values():
+        offload_state_dict = True
+    load_checkpoint_in_model(
+        model,
+        checkpoint,
+        device_map=device_map,
+        offload_folder=offload_folder,
+        dtype=dtype,
+        offload_state_dict=offload_state_dict,
+        offload_buffers=offload_buffers,
+        strict=strict,
+    )
+    if device_map is None:
+        return model
+    return dispatch_model(
+        model,
+        device_map=device_map,
+        offload_dir=offload_folder,
+        offload_buffers=offload_buffers,
+        skip_keys=skip_keys,
+        preload_module_classes=preload_module_classes,
+        force_hooks=force_hooks,
+    )

.venv/Lib/site-packages/accelerate/checkpointing.py

@@ -0,0 +1,306 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+from pathlib import Path
+from typing import List
+
+import numpy as np
+import torch
+from safetensors.torch import load_model
+from torch.cuda.amp import GradScaler
+
+from .utils import (
+    MODEL_NAME,
+    OPTIMIZER_NAME,
+    RNG_STATE_NAME,
+    SAFE_MODEL_NAME,
+    SAFE_WEIGHTS_NAME,
+    SAMPLER_NAME,
+    SCALER_NAME,
+    SCHEDULER_NAME,
+    WEIGHTS_NAME,
+    get_pretty_name,
+    is_mlu_available,
+    is_torch_xla_available,
+    is_xpu_available,
+    load,
+    save,
+)
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+from .logging import get_logger
+from .state import PartialState
+
+
+logger = get_logger(__name__)
+
+
+def save_accelerator_state(
+    output_dir: str,
+    model_states: List[dict],
+    optimizers: list,
+    schedulers: list,
+    dataloaders: list,
+    process_index: int,
+    step: int,
+    scaler: GradScaler = None,
+    save_on_each_node: bool = False,
+    safe_serialization: bool = True,
+):
+    """
+    Saves the current states of the models, optimizers, scaler, and RNG generators to a given directory.
+
+    <Tip>
+
+    If `safe_serialization` is `True`, models will be saved with `safetensors` while the rest are saved using native
+    `pickle`.
+
+    </Tip>
+
+    Args:
+        output_dir (`str` or `os.PathLike`):
+            The name of the folder to save all relevant weights and states.
+        model_states (`List[torch.nn.Module]`):
+            A list of model states
+        optimizers (`List[torch.optim.Optimizer]`):
+            A list of optimizer instances
+        schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`):
+            A list of learning rate schedulers
+        dataloaders (`List[torch.utils.data.DataLoader]`):
+            A list of dataloader instances to save their sampler states
+        process_index (`int`):
+            The current process index in the Accelerator state
+        step (`int`):
+            The current step in the internal step tracker
+        scaler (`torch.amp.GradScaler`, *optional*):
+            An optional gradient scaler instance to save;
+        save_on_each_node (`bool`, *optional*):
+            Whether to save on every node, or only the main node.
+        safe_serialization (`bool`, *optional*, defaults to `True`):
+            Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+    """
+    output_dir = Path(output_dir)
+    # Model states
+    for i, state in enumerate(model_states):
+        weights_name = WEIGHTS_NAME if not safe_serialization else SAFE_WEIGHTS_NAME
+        if i > 0:
+            weights_name = weights_name.replace(".", f"_{i}.")
+        output_model_file = output_dir.joinpath(weights_name)
+        save(state, output_model_file, save_on_each_node=save_on_each_node, safe_serialization=safe_serialization)
+        logger.info(f"Model weights saved in {output_model_file}")
+    # Optimizer states
+    for i, opt in enumerate(optimizers):
+        state = opt.state_dict()
+        optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
+        output_optimizer_file = output_dir.joinpath(optimizer_name)
+        save(state, output_optimizer_file, save_on_each_node=save_on_each_node, safe_serialization=False)
+        logger.info(f"Optimizer state saved in {output_optimizer_file}")
+    # Scheduler states
+    for i, scheduler in enumerate(schedulers):
+        state = scheduler.state_dict()
+        scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
+        output_scheduler_file = output_dir.joinpath(scheduler_name)
+        save(state, output_scheduler_file, save_on_each_node=save_on_each_node, safe_serialization=False)
+        logger.info(f"Scheduler state saved in {output_scheduler_file}")
+    # DataLoader states
+    for i, dataloader in enumerate(dataloaders):
+        sampler_name = f"{SAMPLER_NAME}.bin" if i == 0 else f"{SAMPLER_NAME}_{i}.bin"
+        output_sampler_file = output_dir.joinpath(sampler_name)
+        # Only save if we have our custom sampler
+        from .data_loader import IterableDatasetShard, SeedableRandomSampler
+
+        if isinstance(dataloader.dataset, IterableDatasetShard):
+            sampler = dataloader.get_sampler()
+            if isinstance(sampler, SeedableRandomSampler):
+                save(sampler, output_sampler_file, save_on_each_node=save_on_each_node, safe_serialization=False)
+        if getattr(dataloader, "use_stateful_dataloader", False):
+            dataloader_state_dict_name = "dl_state_dict.bin" if i == 0 else f"dl_state_dict_{i}.bin"
+            output_dataloader_state_dict_file = output_dir.joinpath(dataloader_state_dict_name)
+            state_dict = dataloader.state_dict()
+            torch.save(state_dict, output_dataloader_state_dict_file)
+        logger.info(f"Sampler state for dataloader {i} saved in {output_sampler_file}")
+
+    # GradScaler state
+    if scaler is not None:
+        state = scaler.state_dict()
+        output_scaler_file = output_dir.joinpath(SCALER_NAME)
+        torch.save(state, output_scaler_file)
+        logger.info(f"Gradient scaler state saved in {output_scaler_file}")
+    # Random number generator states
+    states = {}
+    states_name = f"{RNG_STATE_NAME}_{process_index}.pkl"
+    states["step"] = step
+    states["random_state"] = random.getstate()
+    states["numpy_random_seed"] = np.random.get_state()
+    states["torch_manual_seed"] = torch.get_rng_state()
+    if is_xpu_available():
+        states["torch_xpu_manual_seed"] = torch.xpu.get_rng_state_all()
+    if is_mlu_available():
+        states["torch_mlu_manual_seed"] = torch.mlu.get_rng_state_all()
+    else:
+        states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
+    if is_torch_xla_available():
+        states["xm_seed"] = xm.get_rng_state()
+    output_states_file = output_dir.joinpath(states_name)
+    torch.save(states, output_states_file)
+    logger.info(f"Random states saved in {output_states_file}")
+    return output_dir
+
+
+def load_accelerator_state(
+    input_dir,
+    models,
+    optimizers,
+    schedulers,
+    dataloaders,
+    process_index,
+    scaler=None,
+    map_location=None,
+    **load_model_func_kwargs,
+):
+    """
+    Loads states of the models, optimizers, scaler, and RNG generators from a given directory.
+
+    Args:
+        input_dir (`str` or `os.PathLike`):
+            The name of the folder to load all relevant weights and states.
+        models (`List[torch.nn.Module]`):
+            A list of model instances
+        optimizers (`List[torch.optim.Optimizer]`):
+            A list of optimizer instances
+        schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`):
+            A list of learning rate schedulers
+        process_index (`int`):
+            The current process index in the Accelerator state
+        scaler (`torch.amp.GradScaler`, *optional*):
+            An optional *GradScaler* instance to load
+        map_location (`str`, *optional*):
+            What device to load the optimizer state onto. Should be one of either "cpu" or "on_device".
+        load_model_func_kwargs (`dict`, *optional*):
+            Additional arguments that can be passed to the model's `load_state_dict` method.
+
+    Returns:
+        `dict`: Contains the `Accelerator` attributes to override while loading the state.
+    """
+    # stores the `Accelerator` attributes to override
+    override_attributes = dict()
+    if map_location not in [None, "cpu", "on_device"]:
+        raise TypeError(
+            "Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"
+        )
+    if map_location is None:
+        map_location = "cpu"
+    elif map_location == "on_device":
+        map_location = PartialState().device
+
+    input_dir = Path(input_dir)
+    # Model states
+    for i, model in enumerate(models):
+        ending = f"_{i}" if i > 0 else ""
+        input_model_file = input_dir.joinpath(f"{SAFE_MODEL_NAME}{ending}.safetensors")
+        if input_model_file.exists():
+            load_model(model, input_model_file, device=str(map_location), **load_model_func_kwargs)
+        else:
+            # Load with torch
+            input_model_file = input_dir.joinpath(f"{MODEL_NAME}{ending}.bin")
+            state_dict = load(input_model_file, map_location=map_location)
+            model.load_state_dict(state_dict, **load_model_func_kwargs)
+    logger.info("All model weights loaded successfully")
+
+    # Optimizer states
+    for i, opt in enumerate(optimizers):
+        optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
+        input_optimizer_file = input_dir.joinpath(optimizer_name)
+        optimizer_state = load(input_optimizer_file, map_location=map_location)
+        optimizers[i].load_state_dict(optimizer_state)
+    logger.info("All optimizer states loaded successfully")
+
+    # Scheduler states
+    for i, scheduler in enumerate(schedulers):
+        scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
+        input_scheduler_file = input_dir.joinpath(scheduler_name)
+        scheduler_state = load(input_scheduler_file)
+        scheduler.load_state_dict(scheduler_state)
+    logger.info("All scheduler states loaded successfully")
+
+    for i, dataloader in enumerate(dataloaders):
+        sampler_name = f"{SAMPLER_NAME}.bin" if i == 0 else f"{SAMPLER_NAME}_{i}.bin"
+        input_sampler_file = input_dir.joinpath(sampler_name)
+        # Only load if we have our custom sampler
+        from .data_loader import IterableDatasetShard, SeedableRandomSampler
+
+        if isinstance(dataloader.dataset, IterableDatasetShard):
+            sampler = dataloader.get_sampler()
+            if isinstance(sampler, SeedableRandomSampler):
+                sampler = dataloader.set_sampler(load(input_sampler_file))
+        if getattr(dataloader, "use_stateful_dataloader", False):
+            dataloader_state_dict_name = "dl_state_dict.bin" if i == 0 else f"dl_state_dict_{i}.bin"
+            input_dataloader_state_dict_file = input_dir.joinpath(dataloader_state_dict_name)
+            if input_dataloader_state_dict_file.exists():
+                state_dict = load(input_dataloader_state_dict_file)
+                dataloader.load_state_dict(state_dict)
+    logger.info("All dataloader sampler states loaded successfully")
+
+    # GradScaler state
+    if scaler is not None:
+        input_scaler_file = input_dir.joinpath(SCALER_NAME)
+        scaler_state = load(input_scaler_file)
+        scaler.load_state_dict(scaler_state)
+        logger.info("GradScaler state loaded successfully")
+
+    # Random states
+    try:
+        states = load(input_dir.joinpath(f"{RNG_STATE_NAME}_{process_index}.pkl"))
+        if "step" in states:
+            override_attributes["step"] = states["step"]
+        random.setstate(states["random_state"])
+        np.random.set_state(states["numpy_random_seed"])
+        torch.set_rng_state(states["torch_manual_seed"])
+        if is_xpu_available():
+            torch.xpu.set_rng_state_all(states["torch_xpu_manual_seed"])
+        if is_mlu_available():
+            torch.mlu.set_rng_state_all(states["torch_mlu_manual_seed"])
+        else:
+            torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"])
+        if is_torch_xla_available():
+            xm.set_rng_state(states["xm_seed"])
+        logger.info("All random states loaded successfully")
+    except Exception:
+        logger.info("Could not load random states")
+
+    return override_attributes
+
+
+def save_custom_state(obj, path, index: int = 0, save_on_each_node: bool = False):
+    """
+    Saves the state of `obj` to `{path}/custom_checkpoint_{index}.pkl`
+    """
+    # Should this be the right way to get a qual_name type value from `obj`?
+    save_location = Path(path) / f"custom_checkpoint_{index}.pkl"
+    logger.info(f"Saving the state of {get_pretty_name(obj)} to {save_location}")
+    save(obj.state_dict(), save_location, save_on_each_node=save_on_each_node)
+
+
+def load_custom_state(obj, path, index: int = 0):
+    """
+    Loads the state of `obj` at `{path}/custom_checkpoint_{index}.pkl`. Will always set `weights_only=False` when
+    loading the state.
+    """
+    load_location = f"{path}/custom_checkpoint_{index}.pkl"
+    logger.info(f"Loading the state of {get_pretty_name(obj)} from {load_location}")
+    obj.load_state_dict(load(load_location, map_location="cpu", weights_only=False))

.venv/Lib/site-packages/accelerate/commands/accelerate_cli.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from accelerate.commands.config import get_config_parser
+from accelerate.commands.env import env_command_parser
+from accelerate.commands.estimate import estimate_command_parser
+from accelerate.commands.launch import launch_command_parser
+from accelerate.commands.merge import merge_command_parser
+from accelerate.commands.test import test_command_parser
+from accelerate.commands.tpu import tpu_command_parser
+from accelerate.commands.utils import CustomArgumentParser
+
+
+def main():
+    parser = CustomArgumentParser("Accelerate CLI tool", usage="accelerate <command> [<args>]", allow_abbrev=False)
+    subparsers = parser.add_subparsers(help="accelerate command helpers")
+
+    # Register commands
+    get_config_parser(subparsers=subparsers)
+    estimate_command_parser(subparsers=subparsers)
+    env_command_parser(subparsers=subparsers)
+    launch_command_parser(subparsers=subparsers)
+    merge_command_parser(subparsers=subparsers)
+    tpu_command_parser(subparsers=subparsers)
+    test_command_parser(subparsers=subparsers)
+
+    # Let's go
+    args = parser.parse_args()
+
+    if not hasattr(args, "func"):
+        parser.print_help()
+        exit(1)
+
+    # Run
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

.venv/Lib/site-packages/accelerate/commands/config/__init__.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+from .config import config_command_parser
+from .config_args import default_config_file, load_config_from_file  # noqa: F401
+from .default import default_command_parser
+from .update import update_command_parser
+
+
+def get_config_parser(subparsers=None):
+    parent_parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
+    # The main config parser
+    config_parser = config_command_parser(subparsers)
+    # The subparser to add commands to
+    subcommands = config_parser.add_subparsers(title="subcommands", dest="subcommand")
+
+    # Then add other parsers with the parent parser
+    default_command_parser(subcommands, parents=[parent_parser])
+    update_command_parser(subcommands, parents=[parent_parser])
+
+    return config_parser
+
+
+def main():
+    config_parser = get_config_parser()
+    args = config_parser.parse_args()
+
+    if not hasattr(args, "func"):
+        config_parser.print_help()
+        exit(1)
+
+    # Run
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

.venv/Lib/site-packages/accelerate/commands/config/config_args.py

@@ -0,0 +1,252 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional, Union
+
+import yaml
+
+from ...utils import ComputeEnvironment, DistributedType, SageMakerDistributedType
+from ...utils.constants import SAGEMAKER_PYTHON_VERSION, SAGEMAKER_PYTORCH_VERSION, SAGEMAKER_TRANSFORMERS_VERSION
+
+
+hf_cache_home = os.path.expanduser(
+    os.environ.get("HF_HOME", os.path.join(os.environ.get("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.yaml")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    if config_file is not None:
+        if not os.path.isfile(config_file):
+            raise FileNotFoundError(
+                f"The passed configuration file `{config_file}` does not exist. "
+                "Please pass an existing file to `accelerate launch`, or use the default one "
+                "created through `accelerate config` and run `accelerate launch` "
+                "without the `--config_file` argument."
+            )
+    else:
+        config_file = default_config_file
+    with open(config_file, encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if (
+                json.load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE)
+                == ComputeEnvironment.LOCAL_MACHINE
+            ):
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if (
+                yaml.safe_load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE)
+                == ComputeEnvironment.LOCAL_MACHINE
+            ):
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
+    compute_environment: ComputeEnvironment
+    distributed_type: Union[DistributedType, SageMakerDistributedType]
+    mixed_precision: str
+    use_cpu: bool
+    debug: bool
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+
+        def _convert_enums(value):
+            if isinstance(value, Enum):
+                return value.value
+            if isinstance(value, dict):
+                if not bool(value):
+                    return None
+                for key1, value1 in value.items():
+                    value[key1] = _convert_enums(value1)
+            return value
+
+        for key, value in result.items():
+            result[key] = _convert_enums(value)
+        result = {k: v for k, v in result.items() if v is not None}
+        return result
+
+    @staticmethod
+    def process_config(config_dict):
+        """
+        Processes `config_dict` and sets default values for any missing keys
+        """
+        if "compute_environment" not in config_dict:
+            config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE
+        if "distributed_type" not in config_dict:
+            raise ValueError("A `distributed_type` must be specified in the config file.")
+        if "num_processes" not in config_dict and config_dict["distributed_type"] == DistributedType.NO:
+            config_dict["num_processes"] = 1
+        if "mixed_precision" not in config_dict:
+            config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else None
+        if "fp16" in config_dict:  # Convert the config to the new format.
+            del config_dict["fp16"]
+        if "dynamo_backend" in config_dict:  # Convert the config to the new format.
+            dynamo_backend = config_dict.pop("dynamo_backend")
+            config_dict["dynamo_config"] = {} if dynamo_backend == "NO" else {"dynamo_backend": dynamo_backend}
+        if "use_cpu" not in config_dict:
+            config_dict["use_cpu"] = False
+        if "debug" not in config_dict:
+            config_dict["debug"] = False
+        if "enable_cpu_affinity" not in config_dict:
+            config_dict["enable_cpu_affinity"] = False
+        return config_dict
+
+    @classmethod
+    def from_json_file(cls, json_file=None):
+        json_file = default_json_config_file if json_file is None else json_file
+        with open(json_file, encoding="utf-8") as f:
+            config_dict = json.load(f)
+        config_dict = cls.process_config(config_dict)
+        extra_keys = sorted(set(config_dict.keys()) - set(cls.__dataclass_fields__.keys()))
+        if len(extra_keys) > 0:
+            raise ValueError(
+                f"The config file at {json_file} had unknown keys ({extra_keys}), please try upgrading your `accelerate`"
+                " version or fix (and potentially remove) these keys from your config file."
+            )
+
+        return cls(**config_dict)
+
+    def to_json_file(self, json_file):
+        with open(json_file, "w", encoding="utf-8") as f:
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, encoding="utf-8") as f:
+            config_dict = yaml.safe_load(f)
+        config_dict = cls.process_config(config_dict)
+        extra_keys = sorted(set(config_dict.keys()) - set(cls.__dataclass_fields__.keys()))
+        if len(extra_keys) > 0:
+            raise ValueError(
+                f"The config file at {yaml_file} had unknown keys ({extra_keys}), please try upgrading your `accelerate`"
+                " version or fix (and potentially remove) these keys from your config file."
+            )
+        return cls(**config_dict)
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if isinstance(self.compute_environment, str):
+            self.compute_environment = ComputeEnvironment(self.compute_environment)
+        if isinstance(self.distributed_type, str):
+            if self.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
+                self.distributed_type = SageMakerDistributedType(self.distributed_type)
+            else:
+                self.distributed_type = DistributedType(self.distributed_type)
+        if getattr(self, "dynamo_config", None) is None:
+            self.dynamo_config = {}
+
+
+@dataclass
+class ClusterConfig(BaseConfig):
+    num_processes: int = -1  # For instance if we use SLURM and the user manually passes it in
+    machine_rank: int = 0
+    num_machines: int = 1
+    gpu_ids: Optional[str] = None
+    main_process_ip: Optional[str] = None
+    main_process_port: Optional[int] = None
+    rdzv_backend: Optional[str] = "static"
+    same_network: Optional[bool] = False
+    main_training_function: str = "main"
+    enable_cpu_affinity: bool = False
+
+    # args for FP8 training
+    fp8_config: dict = None
+    # args for deepspeed_plugin
+    deepspeed_config: dict = None
+    # args for fsdp
+    fsdp_config: dict = None
+    # args for megatron_lm
+    megatron_lm_config: dict = None
+    # args for ipex
+    ipex_config: dict = None
+    # args for mpirun
+    mpirun_config: dict = None
+    # args for TPU
+    downcast_bf16: bool = False
+
+    # args for TPU pods
+    tpu_name: str = None
+    tpu_zone: str = None
+    tpu_use_cluster: bool = False
+    tpu_use_sudo: bool = False
+    command_file: str = None
+    commands: List[str] = None
+    tpu_vm: List[str] = None
+    tpu_env: List[str] = None
+
+    # args for dynamo
+    dynamo_config: dict = None
+
+    def __post_init__(self):
+        if self.deepspeed_config is None:
+            self.deepspeed_config = {}
+        if self.fsdp_config is None:
+            self.fsdp_config = {}
+        if self.megatron_lm_config is None:
+            self.megatron_lm_config = {}
+        if self.ipex_config is None:
+            self.ipex_config = {}
+        if self.mpirun_config is None:
+            self.mpirun_config = {}
+        if self.fp8_config is None:
+            self.fp8_config = {}
+        return super().__post_init__()
+
+
+@dataclass
+class SageMakerConfig(BaseConfig):
+    ec2_instance_type: str
+    iam_role_name: str
+    image_uri: Optional[str] = None
+    profile: Optional[str] = None
+    region: str = "us-east-1"
+    num_machines: int = 1
+    gpu_ids: str = "all"
+    base_job_name: str = f"accelerate-sagemaker-{num_machines}"
+    pytorch_version: str = SAGEMAKER_PYTORCH_VERSION
+    transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION
+    py_version: str = SAGEMAKER_PYTHON_VERSION
+    sagemaker_inputs_file: str = None
+    sagemaker_metrics_file: str = None
+    additional_args: dict = None
+    dynamo_config: dict = None
+    enable_cpu_affinity: bool = False

.venv/Lib/site-packages/accelerate/commands/config/default.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+import torch
+
+from ...utils import is_mlu_available, is_musa_available, is_npu_available, is_xpu_available
+from .config_args import ClusterConfig, default_json_config_file
+from .config_utils import SubcommandHelpFormatter
+
+
+description = "Create a default config file for Accelerate with only a few flags set."
+
+
+def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file, use_xpu: bool = False):
+    """
+    Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also
+    set CPU if it is a CPU-only machine.
+
+    Args:
+        mixed_precision (`str`, *optional*, defaults to "no"):
+            Mixed Precision to use. Should be one of "no", "fp16", or "bf16"
+        save_location (`str`, *optional*, defaults to `default_json_config_file`):
+            Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default
+            location is inside the huggingface cache folder (`~/.cache/huggingface`) but can be overriden by setting
+            the `HF_HOME` environmental variable, followed by `accelerate/default_config.yaml`.
+        use_xpu (`bool`, *optional*, defaults to `False`):
+            Whether to use XPU if available.
+    """
+    path = Path(save_location)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.exists():
+        print(
+            f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."
+        )
+        return False
+    mixed_precision = mixed_precision.lower()
+    if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
+        raise ValueError(
+            f"`mixed_precision` should be one of 'no', 'fp16', 'bf16', or 'fp8'. Received {mixed_precision}"
+        )
+    config = {
+        "compute_environment": "LOCAL_MACHINE",
+        "mixed_precision": mixed_precision,
+    }
+    if is_mlu_available():
+        num_mlus = torch.mlu.device_count()
+        config["num_processes"] = num_mlus
+        config["use_cpu"] = False
+        if num_mlus > 1:
+            config["distributed_type"] = "MULTI_MLU"
+        else:
+            config["distributed_type"] = "NO"
+    elif is_musa_available():
+        num_musas = torch.musa.device_count()
+        config["num_processes"] = num_musas
+        config["use_cpu"] = False
+        if num_musas > 1:
+            config["distributed_type"] = "MULTI_MUSA"
+        else:
+            config["distributed_type"] = "NO"
+    elif torch.cuda.is_available():
+        num_gpus = torch.cuda.device_count()
+        config["num_processes"] = num_gpus
+        config["use_cpu"] = False
+        if num_gpus > 1:
+            config["distributed_type"] = "MULTI_GPU"
+        else:
+            config["distributed_type"] = "NO"
+    elif is_xpu_available() and use_xpu:
+        num_xpus = torch.xpu.device_count()
+        config["num_processes"] = num_xpus
+        config["use_cpu"] = False
+        if num_xpus > 1:
+            config["distributed_type"] = "MULTI_XPU"
+        else:
+            config["distributed_type"] = "NO"
+    elif is_npu_available():
+        num_npus = torch.npu.device_count()
+        config["num_processes"] = num_npus
+        config["use_cpu"] = False
+        if num_npus > 1:
+            config["distributed_type"] = "MULTI_NPU"
+        else:
+            config["distributed_type"] = "NO"
+    else:
+        num_xpus = 0
+        config["use_cpu"] = True
+        config["num_processes"] = 1
+        config["distributed_type"] = "NO"
+    config["debug"] = False
+    config["enable_cpu_affinity"] = False
+    config = ClusterConfig(**config)
+    config.to_json_file(path)
+    return path
+
+
+def default_command_parser(parser, parents):
+    parser = parser.add_parser("default", parents=parents, help=description, formatter_class=SubcommandHelpFormatter)
+    parser.add_argument(
+        "--config_file",
+        default=default_json_config_file,
+        help=(
+            "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+        dest="save_location",
+    )
+
+    parser.add_argument(
+        "--mixed_precision",
+        choices=["no", "fp16", "bf16"],
+        type=str,
+        help="Whether or not to use mixed precision training. "
+        "Choose between FP16 and BF16 (bfloat16) training. "
+        "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.",
+        default="no",
+    )
+    parser.set_defaults(func=default_config_command)
+    return parser
+
+
+def default_config_command(args):
+    config_file = write_basic_config(args.mixed_precision, args.save_location)
+    if config_file:
+        print(f"accelerate configuration saved at {config_file}")

.venv/Lib/site-packages/accelerate/commands/config/sagemaker.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+
+from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES
+from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
+from ...utils.imports import is_boto3_available
+from .config_args import SageMakerConfig
+from .config_utils import (
+    DYNAMO_BACKENDS,
+    _ask_field,
+    _ask_options,
+    _convert_dynamo_backend,
+    _convert_mixed_precision,
+    _convert_sagemaker_distributed_mode,
+    _convert_yes_no_to_bool,
+)
+
+
+if is_boto3_available():
+    import boto3  # noqa: F401
+
+
+def _create_iam_role_for_sagemaker(role_name):
+    iam_client = boto3.client("iam")
+
+    sagemaker_trust_policy = {
+        "Version": "2012-10-17",
+        "Statement": [
+            {"Effect": "Allow", "Principal": {"Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"}
+        ],
+    }
+    try:
+        # create the role, associated with the chosen trust policy
+        iam_client.create_role(
+            RoleName=role_name, AssumeRolePolicyDocument=json.dumps(sagemaker_trust_policy, indent=2)
+        )
+        policy_document = {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Effect": "Allow",
+                    "Action": [
+                        "sagemaker:*",
+                        "ecr:GetDownloadUrlForLayer",
+                        "ecr:BatchGetImage",
+                        "ecr:BatchCheckLayerAvailability",
+                        "ecr:GetAuthorizationToken",
+                        "cloudwatch:PutMetricData",
+                        "cloudwatch:GetMetricData",
+                        "cloudwatch:GetMetricStatistics",
+                        "cloudwatch:ListMetrics",
+                        "logs:CreateLogGroup",
+                        "logs:CreateLogStream",
+                        "logs:DescribeLogStreams",
+                        "logs:PutLogEvents",
+                        "logs:GetLogEvents",
+                        "s3:CreateBucket",
+                        "s3:ListBucket",
+                        "s3:GetBucketLocation",
+                        "s3:GetObject",
+                        "s3:PutObject",
+                    ],
+                    "Resource": "*",
+                }
+            ],
+        }
+        # attach policy to role
+        iam_client.put_role_policy(
+            RoleName=role_name,
+            PolicyName=f"{role_name}_policy_permission",
+            PolicyDocument=json.dumps(policy_document, indent=2),
+        )
+    except iam_client.exceptions.EntityAlreadyExistsException:
+        print(f"role {role_name} already exists. Using existing one")
+
+
+def _get_iam_role_arn(role_name):
+    iam_client = boto3.client("iam")
+    return iam_client.get_role(RoleName=role_name)["Role"]["Arn"]
+
+
+def get_sagemaker_input():
+    credentials_configuration = _ask_options(
+        "How do you want to authorize?",
+        ["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "],
+        int,
+    )
+    aws_profile = None
+    if credentials_configuration == 0:
+        aws_profile = _ask_field("Enter your AWS Profile name: [default] ", default="default")
+        os.environ["AWS_PROFILE"] = aws_profile
+    else:
+        print(
+            "Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with,"
+            "`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`"
+        )
+        aws_access_key_id = _ask_field("AWS Access Key ID: ")
+        os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
+
+        aws_secret_access_key = _ask_field("AWS Secret Access Key: ")
+        os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
+
+    aws_region = _ask_field("Enter your AWS Region: [us-east-1]", default="us-east-1")
+    os.environ["AWS_DEFAULT_REGION"] = aws_region
+
+    role_management = _ask_options(
+        "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?",
+        ["Provide IAM Role name", "Create new IAM role using credentials"],
+        int,
+    )
+    if role_management == 0:
+        iam_role_name = _ask_field("Enter your IAM role name: ")
+    else:
+        iam_role_name = "accelerate_sagemaker_execution_role"
+        print(f'Accelerate will create an iam role "{iam_role_name}" using the provided credentials')
+        _create_iam_role_for_sagemaker(iam_role_name)
+
+    is_custom_docker_image = _ask_field(
+        "Do you want to use custom Docker image? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    docker_image = None
+    if is_custom_docker_image:
+        docker_image = _ask_field("Enter your Docker image: ", lambda x: str(x).lower())
+
+    is_sagemaker_inputs_enabled = _ask_field(
+        "Do you want to provide SageMaker input channels with data locations? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    sagemaker_inputs_file = None
+    if is_sagemaker_inputs_enabled:
+        sagemaker_inputs_file = _ask_field(
+            "Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): ",
+            lambda x: str(x).lower(),
+        )
+
+    is_sagemaker_metrics_enabled = _ask_field(
+        "Do you want to enable SageMaker metrics? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    sagemaker_metrics_file = None
+    if is_sagemaker_metrics_enabled:
+        sagemaker_metrics_file = _ask_field(
+            "Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): ",
+            lambda x: str(x).lower(),
+        )
+
+    distributed_type = _ask_options(
+        "What is the distributed mode?",
+        ["No distributed training", "Data parallelism"],
+        _convert_sagemaker_distributed_mode,
+    )
+    dynamo_config = {}
+    use_dynamo = _ask_field(
+        "Do you wish to optimize your script with torch dynamo?[yes/NO]:",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+    if use_dynamo:
+        prefix = "dynamo_"
+        dynamo_config[prefix + "backend"] = _ask_options(
+            "Which dynamo backend would you like to use?",
+            [x.lower() for x in DYNAMO_BACKENDS],
+            _convert_dynamo_backend,
+            default=2,
+        )
+        use_custom_options = _ask_field(
+            "Do you want to customize the defaults sent to torch.compile? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+
+        if use_custom_options:
+            dynamo_config[prefix + "mode"] = _ask_options(
+                "Which mode do you want to use?",
+                TORCH_DYNAMO_MODES,
+                lambda x: TORCH_DYNAMO_MODES[int(x)],
+                default="default",
+            )
+            dynamo_config[prefix + "use_fullgraph"] = _ask_field(
+                "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            dynamo_config[prefix + "use_dynamic"] = _ask_field(
+                "Do you want to enable dynamic shape tracing? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+    ec2_instance_query = "Which EC2 instance type you want to use for your training?"
+    if distributed_type != SageMakerDistributedType.NO:
+        ec2_instance_type = _ask_options(
+            ec2_instance_query, SAGEMAKER_PARALLEL_EC2_INSTANCES, lambda x: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(x)]
+        )
+    else:
+        ec2_instance_query += "? [ml.p3.2xlarge]:"
+        ec2_instance_type = _ask_field(ec2_instance_query, lambda x: str(x).lower(), default="ml.p3.2xlarge")
+
+    debug = False
+    if distributed_type != SageMakerDistributedType.NO:
+        debug = _ask_field(
+            "Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+
+    num_machines = 1
+    if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
+        num_machines = _ask_field(
+            "How many machines do you want use? [1]: ",
+            int,
+            default=1,
+        )
+
+    mixed_precision = _ask_options(
+        "Do you wish to use FP16 or BF16 (mixed precision)?",
+        ["no", "fp16", "bf16", "fp8"],
+        _convert_mixed_precision,
+    )
+
+    if use_dynamo and mixed_precision == "no":
+        print(
+            "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts."
+        )
+
+    return SageMakerConfig(
+        image_uri=docker_image,
+        compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER,
+        distributed_type=distributed_type,
+        use_cpu=False,
+        dynamo_config=dynamo_config,
+        ec2_instance_type=ec2_instance_type,
+        profile=aws_profile,
+        region=aws_region,
+        iam_role_name=iam_role_name,
+        mixed_precision=mixed_precision,
+        num_machines=num_machines,
+        sagemaker_inputs_file=sagemaker_inputs_file,
+        sagemaker_metrics_file=sagemaker_metrics_file,
+        debug=debug,
+    )

.venv/Lib/site-packages/accelerate/commands/config/update.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+from .config_args import default_config_file, load_config_from_file
+from .config_utils import SubcommandHelpFormatter
+
+
+description = "Update an existing config file with the latest defaults while maintaining the old configuration."
+
+
+def update_config(args):
+    """
+    Update an existing config file with the latest defaults while maintaining the old configuration.
+    """
+    config_file = args.config_file
+    if config_file is None and Path(default_config_file).exists():
+        config_file = default_config_file
+    elif not Path(config_file).exists():
+        raise ValueError(f"The passed config file located at {config_file} doesn't exist.")
+    config = load_config_from_file(config_file)
+
+    if config_file.endswith(".json"):
+        config.to_json_file(config_file)
+    else:
+        config.to_yaml_file(config_file)
+    return config_file
+
+
+def update_command_parser(parser, parents):
+    parser = parser.add_parser("update", parents=parents, help=description, formatter_class=SubcommandHelpFormatter)
+    parser.add_argument(
+        "--config_file",
+        default=None,
+        help=(
+            "The path to the config file to update. Will default to a file named default_config.yaml in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+    )
+
+    parser.set_defaults(func=update_config_command)
+    return parser
+
+
+def update_config_command(args):
+    config_file = update_config(args)
+    print(f"Sucessfully updated the configuration file at {config_file}.")

.venv/Lib/site-packages/accelerate/commands/env.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import platform
+import subprocess
+
+import numpy as np
+import psutil
+import torch
+
+from accelerate import __version__ as version
+from accelerate.commands.config import default_config_file, load_config_from_file
+
+from ..utils import is_mlu_available, is_musa_available, is_npu_available, is_xpu_available
+
+
+def env_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("env")
+    else:
+        parser = argparse.ArgumentParser("Accelerate env command")
+
+    parser.add_argument(
+        "--config_file", default=None, help="The config file to use for the default values in the launching script."
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=env_command)
+    return parser
+
+
+def env_command(args):
+    pt_version = torch.__version__
+    pt_cuda_available = torch.cuda.is_available()
+    pt_xpu_available = is_xpu_available()
+    pt_mlu_available = is_mlu_available()
+    pt_musa_available = is_musa_available()
+    pt_npu_available = is_npu_available()
+
+    accelerate_config = "Not found"
+    # Get the default from the config file.
+    if args.config_file is not None or os.path.isfile(default_config_file):
+        accelerate_config = load_config_from_file(args.config_file).to_dict()
+
+    # if we can run which, get it
+    command = None
+    bash_location = "Not found"
+    if os.name == "nt":
+        command = ["where", "accelerate"]
+    elif os.name == "posix":
+        command = ["which", "accelerate"]
+    if command is not None:
+        bash_location = subprocess.check_output(command, text=True, stderr=subprocess.STDOUT).strip()
+    info = {
+        "`Accelerate` version": version,
+        "Platform": platform.platform(),
+        "`accelerate` bash location": bash_location,
+        "Python version": platform.python_version(),
+        "Numpy version": np.__version__,
+        "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+        "PyTorch XPU available": str(pt_xpu_available),
+        "PyTorch NPU available": str(pt_npu_available),
+        "PyTorch MLU available": str(pt_mlu_available),
+        "PyTorch MUSA available": str(pt_musa_available),
+        "System RAM": f"{psutil.virtual_memory().total / 1024 ** 3:.2f} GB",
+    }
+    if pt_cuda_available:
+        info["GPU type"] = torch.cuda.get_device_name()
+    if pt_mlu_available:
+        info["MLU type"] = torch.mlu.get_device_name()
+    if pt_npu_available:
+        info["CANN version"] = torch.version.cann
+
+    print("\nCopy-and-paste the text below in your GitHub issue\n")
+    print("\n".join([f"- {prop}: {val}" for prop, val in info.items()]))
+
+    print("- `Accelerate` default config:" if args.config_file is None else "- `Accelerate` config passed:")
+    accelerate_config_str = (
+        "\n".join([f"\t- {prop}: {val}" for prop, val in accelerate_config.items()])
+        if isinstance(accelerate_config, dict)
+        else f"\t{accelerate_config}"
+    )
+    print(accelerate_config_str)
+
+    info["`Accelerate` configs"] = accelerate_config
+
+    return info
+
+
+def main() -> int:
+    parser = env_command_parser()
+    args = parser.parse_args()
+    env_command(args)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.venv/Lib/site-packages/accelerate/data_loader.py

@@ -0,0 +1,1323 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from contextlib import suppress
+from typing import Callable, List, Optional, Union
+
+import torch
+from torch.utils.data import BatchSampler, DataLoader, IterableDataset, RandomSampler
+
+from .logging import get_logger
+from .state import DistributedType, GradientState, PartialState, is_torch_xla_available
+from .utils import (
+    RNGType,
+    broadcast,
+    broadcast_object_list,
+    concatenate,
+    find_batch_size,
+    get_data_structure,
+    initialize_tensors,
+    is_torch_version,
+    is_torchdata_stateful_dataloader_available,
+    send_to_device,
+    slice_tensors,
+    synchronize_rng_states,
+)
+
+
+logger = get_logger(__name__)
+
+# kwargs of the DataLoader in min version 1.4.0.
+_PYTORCH_DATALOADER_KWARGS = {
+    "batch_size": 1,
+    "shuffle": False,
+    "sampler": None,
+    "batch_sampler": None,
+    "num_workers": 0,
+    "collate_fn": None,
+    "pin_memory": False,
+    "drop_last": False,
+    "timeout": 0,
+    "worker_init_fn": None,
+    "multiprocessing_context": None,
+    "generator": None,
+    "prefetch_factor": 2,
+    "persistent_workers": False,
+}
+
+# kwargs added after by version
+_PYTORCH_DATALOADER_ADDITIONAL_KWARGS = {}
+
+for v, additional_kwargs in _PYTORCH_DATALOADER_ADDITIONAL_KWARGS.items():
+    if is_torch_version(">=", v):
+        _PYTORCH_DATALOADER_KWARGS.update(additional_kwargs)
+
+
+class SeedableRandomSampler(RandomSampler):
+    """
+    Same as a random sampler, except that in `__iter__` a seed can be used.
+
+    Needed specifically in distributed cases, when the random generator for each GPU needs to start from the same seed
+    and be fully reproducable on multiple iterations.
+
+    If a custom `generator` is passed, it will rely on its initial seed as well as the current iteration it is on
+    (stored in `self.epoch`).
+    """
+
+    def __init__(self, *args, **kwargs):
+        data_seed = kwargs.pop("data_seed", None)
+        super().__init__(*args, **kwargs)
+
+        self.initial_seed = data_seed if data_seed is not None else torch.random.initial_seed()
+        self.epoch = 0
+
+    def __iter__(self):
+        if self.generator is None:
+            self.generator = torch.Generator()
+            self.generator.manual_seed(self.initial_seed)
+
+        # Allow `self.epoch` to modify the seed of the generator
+        seed = self.epoch + self.initial_seed
+        # print("Setting seed at epoch", self.epoch, seed)
+        self.generator.manual_seed(seed)
+        yield from super().__iter__()
+        self.set_epoch(self.epoch + 1)
+
+    def set_epoch(self, epoch: int):
+        "Sets the current iteration of the sampler."
+        self.epoch = epoch
+
+
+class BatchSamplerShard(BatchSampler):
+    """
+    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).
+
+    <Tip warning={true}>
+
+    `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
+    equal to `False`
+
+    </Tip>"""
+
+    def __init__(
+        self,
+        batch_sampler: BatchSampler,
+        num_processes: int = 1,
+        process_index: int = 0,
+        split_batches: bool = False,
+        even_batches: bool = True,
+    ):
+        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.even_batches = even_batches
+        self.batch_size = getattr(batch_sampler, "batch_size", None)
+        self.drop_last = getattr(batch_sampler, "drop_last", False)
+        if self.batch_size is None and self.even_batches:
+            raise ValueError(
+                "You need to use `even_batches=False` when the batch sampler has no batch size. If you "
+                "are not calling this method directly, set `accelerator.even_batches=False` instead."
+            )
+
+    @property
+    def total_length(self):
+        return len(self.batch_sampler)
+
+    def __len__(self):
+        if self.split_batches:
+            # Split batches does not change the length of the batch sampler
+            return len(self.batch_sampler)
+        if len(self.batch_sampler) % self.num_processes == 0:
+            # If the length is a round multiple of the number of processes, it's easy.
+            return len(self.batch_sampler) // self.num_processes
+        length = len(self.batch_sampler) // self.num_processes
+        if self.drop_last:
+            # Same if we drop the remainder.
+            return length
+        elif self.even_batches:
+            # When we even batches we always get +1
+            return length + 1
+        else:
+            # Otherwise it depends on the process index.
+            return length + 1 if self.process_index < len(self.batch_sampler) % self.num_processes else length
+
+    def __iter__(self):
+        return self._iter_with_split() if self.split_batches else self._iter_with_no_split()
+
+    def _iter_with_split(self):
+        initial_data = []
+        batch_length = self.batch_sampler.batch_size // self.num_processes
+        for idx, batch in enumerate(self.batch_sampler):
+            if idx == 0:
+                initial_data = batch
+            if len(batch) == self.batch_size:
+                # If the batch is full, we yield the part of it this process is responsible of.
+                yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)]
+
+        # If drop_last is True of the last batch was full, iteration is over, otherwise...
+        if not self.drop_last and len(initial_data) > 0 and len(batch) < self.batch_size:
+            if not self.even_batches:
+                if len(batch) > batch_length * self.process_index:
+                    yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)]
+            else:
+                # For degenerate cases where the dataset has less than num_process * batch_size samples
+                while len(initial_data) < self.batch_size:
+                    initial_data += initial_data
+                batch = batch + initial_data
+                yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)]
+
+    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 (
+                self.batch_size is None or len(batch) == self.batch_size
+            ):
+                yield batch_to_yield
+                batch_to_yield = []
+
+        # If drop_last is True, iteration is over, otherwise...
+        if not self.drop_last and len(initial_data) > 0:
+            if not self.even_batches:
+                if len(batch_to_yield) > 0:
+                    yield batch_to_yield
+            else:
+                # ... we yield the complete batch we had saved before if it has the proper length
+                if len(batch_to_yield) == self.batch_size:
+                    yield batch_to_yield
+
+                # For degenerate cases where the dataset has less than num_process * batch_size samples
+                while len(initial_data) < self.num_processes * self.batch_size:
+                    initial_data += initial_data
+
+                # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next
+                if len(batch) == self.batch_size:
+                    batch = []
+                    idx += 1
+
+                # Make sure we yield a multiple of self.num_processes batches
+                cycle_index = 0
+                while idx % self.num_processes != 0 or len(batch) > 0:
+                    end_index = cycle_index + self.batch_size - len(batch)
+                    batch += initial_data[cycle_index:end_index]
+                    if idx % self.num_processes == self.process_index:
+                        yield batch
+                    cycle_index = end_index
+                    batch = []
+                    idx += 1
+
+
+class IterableDatasetShard(IterableDataset):
+    """
+    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`.
+    """
+
+    def __init__(
+        self,
+        dataset: IterableDataset,
+        batch_size: int = 1,
+        drop_last: bool = False,
+        num_processes: int = 1,
+        process_index: int = 0,
+        split_batches: bool = False,
+    ):
+        if split_batches and batch_size > 1 and batch_size % num_processes != 0:
+            raise ValueError(
+                f"To use `IterableDatasetShard` in `split_batches` mode, the batch size ({batch_size}) "
+                f"needs to be a round multiple of the number of processes ({num_processes})."
+            )
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+        self.num_processes = num_processes
+        self.process_index = process_index
+        self.split_batches = split_batches
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+        if hasattr(self.dataset, "set_epoch"):
+            self.dataset.set_epoch(epoch)
+
+    def __len__(self):
+        # We will just raise the downstream error if the underlying dataset is not sized
+        if self.drop_last:
+            return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size
+        else:
+            return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size
+
+    def __iter__(self):
+        if (
+            not hasattr(self.dataset, "set_epoch")
+            and hasattr(self.dataset, "generator")
+            and isinstance(self.dataset.generator, torch.Generator)
+        ):
+            self.dataset.generator.manual_seed(self.epoch)
+        real_batch_size = self.batch_size if self.split_batches else (self.batch_size * self.num_processes)
+        process_batch_size = (self.batch_size // self.num_processes) if self.split_batches else self.batch_size
+        process_slice = range(self.process_index * process_batch_size, (self.process_index + 1) * process_batch_size)
+
+        first_batch = None
+        current_batch = []
+        for element in self.dataset:
+            current_batch.append(element)
+            # Wait to have a full batch before yielding elements.
+            if len(current_batch) == real_batch_size:
+                for i in process_slice:
+                    yield current_batch[i]
+                if first_batch is None:
+                    first_batch = current_batch.copy()
+                current_batch = []
+
+        # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning.
+        if not self.drop_last and len(current_batch) > 0:
+            if first_batch is None:
+                first_batch = current_batch.copy()
+            while len(current_batch) < real_batch_size:
+                current_batch += first_batch
+            for i in process_slice:
+                yield current_batch[i]
+
+
+class DataLoaderStateMixin:
+    """
+    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
+
+    <Tip warning={true}>
+
+        Inheriters of this class should ensure that the class creates a `GradientState()` instance, stored in
+        `self.gradient_state`.
+
+    </Tip>
+
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        cls.end_of_dataloader = False
+        cls.remainder = -1
+
+    def reset(self):
+        self.end_of_dataloader = False
+        self.remainder = -1
+
+    def begin(self):
+        "Prepares the gradient state for the current dataloader"
+        self.reset()
+        with suppress(Exception):
+            if not self._drop_last:
+                length = getattr(self.dataset, "total_dataset_length", len(self.dataset))
+                self.remainder = length % self.total_batch_size
+        self.gradient_state._add_dataloader(self)
+
+    def end(self):
+        "Cleans up the gradient state after exiting the dataloader"
+        self.gradient_state._remove_dataloader(self)
+
+
+class DataLoaderAdapter:
+    """
+    A class which wraps around a PyTorch `DataLoader` (or variants of it) to be used with the `Accelerator`. For
+    compatability reasons, this class inherits from the class it wraps around, so it can be used as a drop-in.
+    """
+
+    def __init__(self, dataset, use_stateful_dataloader=False, batch_sampler=None, **kwargs):
+        self.use_stateful_dataloader = use_stateful_dataloader
+        if is_torchdata_stateful_dataloader_available():
+            from torchdata.stateful_dataloader import StatefulDataLoader
+
+        if use_stateful_dataloader and not is_torchdata_stateful_dataloader_available():
+            raise ImportError(
+                "StatefulDataLoader is not available. Please install torchdata version 0.8.0 or higher to use it."
+            )
+        if use_stateful_dataloader:
+            self.base_dataloader = StatefulDataLoader(dataset, batch_sampler=batch_sampler, **kwargs)
+        else:
+            self.base_dataloader = DataLoader(dataset, batch_sampler=batch_sampler, **kwargs)
+
+        if hasattr(self.base_dataloader, "state_dict"):
+            self.dl_state_dict = self.base_dataloader.state_dict()
+
+    def __getattr__(self, name):
+        # Avoid infinite recursion if we try to access a nonexistent base_dataloader attribute.
+        if name == "base_dataloader":
+            raise AttributeError()
+        # Delegate attribute access to the internal dataloader
+        return getattr(self.base_dataloader, name)
+
+    def state_dict(self):
+        return self.dl_state_dict
+
+    def load_state_dict(self, state_dict):
+        self.base_dataloader.load_state_dict(state_dict)
+
+    @property
+    def __class__(self):
+        """
+        In order to maintain backwards compatability with other code, we need to ensure `isinstance(obj, DataLoader)`
+        returs true. This is because some downstream code assumes that the `DataLoader` is the base class of the
+        object.
+        """
+        return self.base_dataloader.__class__
+
+    def __len__(self):
+        return len(self.base_dataloader)
+
+    def adjust_state_dict_for_prefetch(self):
+        """
+        Adjusts the state dict for prefetching. Natively, this will adjust all of the iters yielded keys in
+        `self.dl_state_dict` by a factor of `num_processes - 1`, however if a custom correction is needed, this can be
+        overridden.
+
+        This should modify `self.dl_state_dict` directly
+        """
+        # The state dict will be off by a factor of `n-1` batch too many during DDP,
+        # so we need to adjust it here
+        if PartialState().distributed_type != DistributedType.NO:
+            factor = PartialState().num_processes - 1
+            if self.dl_state_dict["_sampler_iter_yielded"] > 0:
+                self.dl_state_dict["_sampler_iter_yielded"] -= factor
+            if self.dl_state_dict["_num_yielded"] > 0:
+                self.dl_state_dict["_num_yielded"] -= factor
+            if self.dl_state_dict["_index_sampler_state"] is not None:
+                if (
+                    "samples_yielded" in self.dl_state_dict["_index_sampler_state"]
+                    and self.dl_state_dict["_index_sampler_state"]["samples_yielded"] > 0
+                ):
+                    self.dl_state_dict["_index_sampler_state"]["samples_yielded"] -= self.batch_size * factor
+
+    def _update_state_dict(self):
+        # The state_dict of the underlying base_dataloader may be ahead of what is currently being yielded.
+        # E.g. the implementation of DataLoaderShard involves having an underlying iterator 1 element ahead of
+        # what it wants to yield.
+        #
+        # _update_state_dict is called to snapshot the state_dict that would properly recover the DataLoaderAdapter.
+        if hasattr(self.base_dataloader, "state_dict"):
+            self.dl_state_dict = self.base_dataloader.state_dict()
+            # Potentially modify the state_dict to adjust for prefetching
+            self.adjust_state_dict_for_prefetch()
+            # Then tag if we are at the end of the dataloader
+            self.dl_state_dict["_iterator_finished"] = self.end_of_dataloader
+
+
+class DataLoaderShard(DataLoaderAdapter, DataLoaderStateMixin):
+    """
+    Subclass of `DataLoaderAdapter` that will deal with device placement and current distributed setup.
+
+    Args:
+        dataset (`torch.utils.data.dataset.Dataset`):
+            The dataset to use to build this dataloader.
+        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.
+        skip_batches (`int`, *optional*, defaults to 0):
+            The number of batches to skip at the beginning.
+        use_stateful_dataloader (`bool`, *optional*, defaults to `False`):
+            Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`.
+        **kwargs (additional keyword arguments, *optional*):
+            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.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        device=None,
+        rng_types=None,
+        synchronized_generator=None,
+        skip_batches=0,
+        use_stateful_dataloader=False,
+        _drop_last: bool = False,
+        _non_blocking: bool = False,
+        **kwargs,
+    ):
+        super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs)
+        self.device = device
+        self.rng_types = rng_types
+        self.synchronized_generator = synchronized_generator
+        self.skip_batches = skip_batches
+        self.gradient_state = GradientState()
+        self._drop_last = _drop_last
+        self._non_blocking = _non_blocking
+        self.iteration = 0
+
+    def __iter__(self):
+        if self.rng_types is not None:
+            synchronize_rng_states(self.rng_types, self.synchronized_generator)
+        self.begin()
+
+        self.set_epoch(self.iteration)
+        dataloader_iter = self.base_dataloader.__iter__()
+        # We iterate one batch ahead to check when we are at the end
+        try:
+            current_batch = next(dataloader_iter)
+        except StopIteration:
+            yield
+
+        batch_index = 0
+        while True:
+            try:
+                # But we still move it to the device so it is done before `StopIteration` is reached
+                if self.device is not None:
+                    current_batch = send_to_device(current_batch, self.device, non_blocking=self._non_blocking)
+                self._update_state_dict()
+                next_batch = next(dataloader_iter)
+                if batch_index >= self.skip_batches:
+                    yield current_batch
+                batch_index += 1
+                current_batch = next_batch
+            except StopIteration:
+                self.end_of_dataloader = True
+                self._update_state_dict()
+                if batch_index >= self.skip_batches:
+                    yield current_batch
+                break
+
+        self.iteration += 1
+        self.end()
+
+    def __reduce__(self):
+        """
+        Define the `__reduce__` method to ensure a `DataLoaderShard` can be pickled and unpickled. This needs to be
+        explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its
+        `__class__` member.
+        """
+        args = super().__reduce__()
+        return (DataLoaderShard, *args[1:])
+
+    def set_epoch(self, epoch: int):
+        # In case it is manually passed in, the user can set it to what they like
+        if self.iteration != epoch:
+            self.iteration = epoch
+        if hasattr(self.batch_sampler, "set_epoch"):
+            self.batch_sampler.set_epoch(epoch)
+        if hasattr(self.batch_sampler, "sampler") and hasattr(self.batch_sampler.sampler, "set_epoch"):
+            self.batch_sampler.sampler.set_epoch(epoch)
+        # We support if a custom `Dataset` implementation has `set_epoch`
+        # or in general HF datasets `Datasets`
+        elif hasattr(self.dataset, "set_epoch"):
+            self.dataset.set_epoch(epoch)
+
+    @property
+    def total_batch_size(self):
+        batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
+        return (
+            batch_sampler.batch_size
+            if getattr(batch_sampler, "split_batches", False)
+            else (batch_sampler.batch_size * getattr(batch_sampler, "num_processes", 1))
+        )
+
+    @property
+    def total_dataset_length(self):
+        if hasattr(self.dataset, "total_length"):
+            return self.dataset.total_length
+        else:
+            return len(self.dataset)
+
+    def get_sampler(self):
+        return get_sampler(self)
+
+    def set_sampler(self, sampler):
+        sampler_is_batch_sampler = isinstance(self.sampler, BatchSampler)
+        if sampler_is_batch_sampler:
+            self.sampler.sampler = sampler
+        else:
+            self.batch_sampler.sampler = sampler
+            if hasattr(self.batch_sampler, "batch_sampler"):
+                self.batch_sampler.batch_sampler.sampler = sampler
+
+
+if is_torch_xla_available():
+    import torch_xla.distributed.parallel_loader as xpl
+
+    class MpDeviceLoaderWrapper(xpl.MpDeviceLoader):
+        """
+        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.
+        """
+
+        def __init__(self, dataloader: DataLoaderShard, device: torch.device):
+            super().__init__(dataloader, device)
+            self._rng_types = self._loader.rng_types
+            self._loader.rng_types = None
+            self.device = device
+
+        def __iter__(self):
+            if self._rng_types is not None:
+                synchronize_rng_states(self._rng_types, self._loader.synchronized_generator)
+
+            return super().__iter__()
+
+        def set_epoch(self, epoch: int):
+            if hasattr(self.dataloader, "set_epoch"):
+                self.dataloader.set_epoch(epoch)
+
+        @property
+        def total_batch_size(self):
+            return self._loader.total_batch_size
+
+        @property
+        def total_dataset_length(self):
+            return self._loader.total_dataset_length
+
+        @property
+        def batch_sampler(self):
+            return self._loader.batch_sampler
+
+        @property
+        def dataloader(self):
+            return self._loader
+
+
+class DataLoaderDispatcher(DataLoaderAdapter, DataLoaderStateMixin):
+    """
+    Subclass of `DataLoaderAdapter` 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.
+        use_stateful_dataloader (`bool`, *optional*, defaults to `False`):
+            Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`.
+
+    **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.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        split_batches: bool = False,
+        skip_batches=0,
+        use_stateful_dataloader=False,
+        _drop_last: bool = False,
+        _non_blocking: bool = False,
+        slice_fn=None,
+        **kwargs,
+    ):
+        shuffle = False
+        if is_torch_version(">=", "1.11.0"):
+            from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe
+
+            # We need to save the shuffling state of the DataPipe
+            if isinstance(dataset, ShufflerIterDataPipe):
+                shuffle = dataset._shuffle_enabled
+        super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs)
+        self.split_batches = split_batches
+        if shuffle:
+            torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle)
+
+        self.gradient_state = GradientState()
+        self.state = PartialState()
+        self._drop_last = _drop_last
+        self._non_blocking = _non_blocking
+        self.skip_batches = skip_batches
+
+        self.slice_fn = slice_tensors if slice_fn is None else slice_fn
+        self.iteration = 0
+
+    def _fetch_batches(self, iterator):
+        batches, batch = None, None
+        # On process 0, we gather the batch to dispatch.
+        if self.state.process_index == 0:
+            try:
+                if self.split_batches:
+                    # One batch of the main iterator is dispatched and split.
+                    self._update_state_dict()
+                    batch = next(iterator)
+                else:
+                    # num_processes batches of the main iterator are concatenated then dispatched and split.
+                    # We add the batches one by one so we have the remainder available when drop_last=False.
+                    batches = []
+                    for _ in range(self.state.num_processes):
+                        self._update_state_dict()
+                        batches.append(next(iterator))
+                    try:
+                        batch = concatenate(batches, dim=0)
+                    except RuntimeError as e:
+                        raise RuntimeError(
+                            "You can't use batches of different size with `dispatch_batches=True` or when using an `IterableDataset`."
+                            "either pass `dispatch_batches=False` and have each process fetch its own batch "
+                            " or pass `split_batches=True`. By doing so, the main process will fetch a full batch and "
+                            "slice it into `num_processes` batches for each process."
+                        ) from e
+                # In both cases, we need to get the structure of the batch that we will broadcast on other
+                # processes to initialize the tensors with the right shape.
+                # data_structure, stop_iteration
+                batch_info = [get_data_structure(batch), False]
+            except StopIteration:
+                batch_info = [None, True]
+        else:
+            batch_info = [None, self._stop_iteration]
+        # This is inplace, so after this instruction, every process has the same `batch_info` as process 0.
+        broadcast_object_list(batch_info)
+        self._stop_iteration = batch_info[1]
+        if self._stop_iteration:
+            # If drop_last is False and split_batches is False, we may have a remainder to take care of.
+            if not self.split_batches and not self._drop_last:
+                if self.state.process_index == 0 and len(batches) > 0:
+                    batch = concatenate(batches, dim=0)
+                    batch_info = [get_data_structure(batch), False]
+                else:
+                    batch_info = [None, True]
+                broadcast_object_list(batch_info)
+        return batch, batch_info
+
+    def __iter__(self):
+        self.begin()
+        self.set_epoch(self.iteration)
+        main_iterator = None
+        if is_torch_version(">=", "2.0.1"):
+            # NOTE PyTorch DataLoader adds forward compatibilities for DataPipes, which broadcasts
+            # shared seed to all dist processes. Thus, we need to create iterator for all dist processes.
+            # But, we only iterate through the DataLoader on process 0.
+            main_iterator = self.base_dataloader.__iter__()
+        elif self.state.process_index == 0:
+            main_iterator = self.base_dataloader.__iter__()
+        stop_iteration = False
+        self._stop_iteration = False
+        first_batch = None
+        next_batch, next_batch_info = self._fetch_batches(main_iterator)
+        batch_index = 0
+        while not stop_iteration:
+            batch, batch_info = next_batch, next_batch_info
+
+            if self.state.process_index != 0:
+                # Initialize tensors on other processes than process 0.
+                batch = initialize_tensors(batch_info[0])
+            batch = send_to_device(batch, self.state.device, non_blocking=self._non_blocking)
+            # Broadcast the batch before splitting it.
+            batch = broadcast(batch, from_process=0)
+
+            if not self._drop_last and first_batch is None:
+                # We keep at least num processes elements of the first batch to be able to complete the last batch
+                first_batch = self.slice_fn(
+                    batch,
+                    slice(0, self.state.num_processes),
+                    process_index=self.state.process_index,
+                    num_processes=self.state.num_processes,
+                )
+
+            if batch is None:
+                raise ValueError(
+                    f"Batch does not contain any data (`{batch}`). At the end of all iterable data available before expected stop iteration."
+                )
+
+            observed_batch_size = find_batch_size(batch)
+            batch_size = observed_batch_size // self.state.num_processes
+
+            stop_iteration = self._stop_iteration
+            if not stop_iteration:
+                # We may still be at the end of the dataloader without knowing it yet: if there is nothing left in
+                # the dataloader since the number of batches is a round multiple of the number of processes.
+                next_batch, next_batch_info = self._fetch_batches(main_iterator)
+                # next_batch_info[0] is None when there are no more batches, otherwise we still need to process them.
+                if self._stop_iteration and next_batch_info[0] is None:
+                    stop_iteration = True
+
+            if not self._drop_last and stop_iteration and observed_batch_size % self.state.num_processes != 0:
+                # If the last batch is not complete, let's add the first batch to it.
+                batch = concatenate([batch, first_batch], dim=0)
+                # Batch size computation above is wrong, it's off by 1 so we fix it.
+                batch_size += 1
+
+            data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size)
+            batch = self.slice_fn(
+                batch,
+                data_slice,
+                process_index=self.state.process_index,
+                num_processes=self.state.num_processes,
+            )
+
+            if stop_iteration:
+                self.end_of_dataloader = True
+                self._update_state_dict()
+                self.remainder = observed_batch_size
+            if batch_index >= self.skip_batches:
+                yield batch
+            batch_index += 1
+        self.iteration += 1
+        self.end()
+
+    def set_epoch(self, epoch: int):
+        # In case it is manually passed in, the user can set it to what they like
+        if self.iteration != epoch:
+            self.iteration = epoch
+        if hasattr(self.batch_sampler, "sampler") and hasattr(self.batch_sampler.sampler, "set_epoch"):
+            self.batch_sampler.sampler.set_epoch(epoch)
+        elif hasattr(self.dataset, "set_epoch"):
+            self.dataset.set_epoch(epoch)
+
+    def __len__(self):
+        whole_length = len(self.base_dataloader)
+        if self.split_batches:
+            return whole_length
+        elif self._drop_last:
+            return whole_length // self.state.num_processes
+        else:
+            return math.ceil(whole_length / self.state.num_processes)
+
+    def __reduce__(self):
+        """
+        Define the `__reduce__` method to ensure a `DataLoaderDispatcher` can be pickled and unpickled. This needs to
+        be explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its
+        `__class__` member.
+        """
+        args = super().__reduce__()
+        return (DataLoaderDispatcher, *args[1:])
+
+    @property
+    def total_batch_size(self):
+        return (
+            self.dataset.batch_size if self.split_batches else (self.dataset.batch_size * self.dataset.num_processes)
+        )
+
+    @property
+    def total_dataset_length(self):
+        return len(self.dataset)
+
+    def get_sampler(self):
+        return get_sampler(self)
+
+    def set_sampler(self, sampler):
+        sampler_is_batch_sampler = isinstance(self.sampler, BatchSampler)
+        if sampler_is_batch_sampler:
+            self.sampler.sampler = sampler
+        else:
+            self.batch_sampler.sampler = sampler
+            if hasattr(self.batch_sampler, "batch_sampler"):
+                self.batch_sampler.batch_sampler.sampler = sampler
+
+
+def get_sampler(dataloader):
+    """
+    Get the sampler associated to the dataloader
+
+    Args:
+        dataloader (`torch.utils.data.dataloader.DataLoader`):
+            The data loader to split across several devices.
+    Returns:
+        `torch.utils.data.Sampler`: The sampler associated to the dataloader
+    """
+    sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler)
+    if sampler_is_batch_sampler:
+        sampler = getattr(dataloader.sampler, "sampler", None)
+    else:
+        sampler = getattr(dataloader.batch_sampler, "sampler", None)
+    return sampler
+
+
+def prepare_data_loader(
+    dataloader: DataLoader,
+    device: Optional[torch.device] = None,
+    num_processes: Optional[int] = None,
+    process_index: Optional[int] = None,
+    split_batches: bool = False,
+    put_on_device: bool = False,
+    rng_types: Optional[List[Union[str, RNGType]]] = None,
+    dispatch_batches: Optional[bool] = None,
+    even_batches: bool = True,
+    slice_fn_for_dispatch: Optional[Callable] = None,
+    use_seedable_sampler: bool = False,
+    data_seed: Optional[int] = None,
+    non_blocking: bool = False,
+    use_stateful_dataloader: bool = False,
+) -> DataLoader:
+    """
+    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.PartialState`].
+        process_index (`int`, *optional*):
+            The index of the current process. Will default to the value given by [`~state.PartialState`].
+        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 dataloader 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.
+        use_seedable_sampler (`bool`, *optional*, defaults to `False`):
+            Whether to use the [`~data_loader.SeedableRandomSampler`] instead of a `RandomSampler` for better
+            reproducability. Comes at a cost of potentially different performances due to different shuffling
+            algorithms but ensures results will be the *exact* same. Should be paired with `set_seed()` at every
+            `self.set_epoch`
+        data_seed (`int`, *optional*, defaults to `None`):
+            The seed to use for the underlying generator when using `use_seedable_sampler`. If `None`, the generator
+            will use the current default seed from torch.
+        non_blocking (`bool`, *optional*, defaults to `False`):
+            If set to `True`, dataloader will utilize non-blocking host-to-device transfers. If the dataloader has
+            `pin_memory` set to `True`, this will help to increase overlap between data transfer and computations.
+        use_stateful_dataloader (`bool`, *optional*, defaults to `False`):
+            "If set to true, the dataloader prepared by the Accelerator will be backed by "
+            "[torchdata.StatefulDataLoader](https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader).
+            This requires `torchdata` version 0.8.0 or higher that supports StatefulDataLoader to be installed."
+
+
+    Returns:
+        `torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches
+
+    <Tip warning={true}>
+
+    `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches`
+    equal to `False`
+
+    </Tip>
+    """
+    if dispatch_batches is None:
+        if not put_on_device:
+            dispatch_batches = False
+        else:
+            dispatch_batches = isinstance(dataloader.dataset, IterableDataset)
+
+    if dispatch_batches and not put_on_device:
+        raise ValueError("Using `dispatch_batches=True` requires `put_on_device=True`.")
+    # Grab defaults from PartialState
+    state = PartialState()
+    if num_processes is None:
+        num_processes = state.num_processes
+    if process_index is None:
+        process_index = state.process_index
+
+    # Sanity check
+    if split_batches:
+        if dataloader.batch_size is not None:
+            batch_size_for_check = dataloader.batch_size
+        else:
+            # For custom batch_sampler
+            if hasattr(dataloader.batch_sampler, "batch_size"):
+                batch_size_for_check = dataloader.batch_sampler.batch_size
+            else:
+                raise ValueError(
+                    "In order to use `split_batches==True` you must have a `batch_size` attribute either in the passed "
+                    "`dataloader` or `dataloader.batch_sampler` objects, and it has to return a natural number. "
+                    "Your `dataloader.batch_size` is None and `dataloader.batch_sampler` "
+                    f"(`{type(dataloader.batch_sampler)}`) does not have the `batch_size` attribute set."
+                )
+
+        if batch_size_for_check > 1 and batch_size_for_check % num_processes != 0:
+            raise ValueError(
+                f"To use a `DataLoader` in `split_batches` mode, the batch size ({dataloader.batch_size}) "
+                f"needs to be a round multiple of the number of processes ({num_processes})."
+            )
+
+    new_dataset = dataloader.dataset
+    # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
+    new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
+    sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler)
+    synchronized_generator = None
+
+    sampler = get_sampler(dataloader)
+    if isinstance(sampler, RandomSampler) and use_seedable_sampler:
+        # When iterating through the dataloader during distributed processes
+        # we want to ensure that on each process we are iterating through the same
+        # samples in the same order if a seed is set. This requires a tweak
+        # to the `torch.utils.data.RandomSampler` class (if used).
+        sampler = SeedableRandomSampler(
+            data_source=sampler.data_source,
+            replacement=sampler.replacement,
+            num_samples=sampler._num_samples,
+            generator=getattr(sampler, "generator", torch.Generator()),
+            data_seed=data_seed,
+        )
+
+    if isinstance(dataloader.sampler, RandomSampler) and state.distributed_type == DistributedType.XLA:
+        # isinstance(dataloader.sampler, RandomSampler) indicates the original dataloader has `shuffle` enabled.
+        generator = torch.Generator().manual_seed(42)
+        dataloader.generator = generator
+        dataloader.sampler.generator = generator
+    # No change if no multiprocess
+    if (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM) and not dispatch_batches:
+        if isinstance(new_dataset, IterableDataset):
+            if getattr(dataloader.dataset, "generator", None) is not None:
+                synchronized_generator = dataloader.dataset.generator
+            new_dataset = IterableDatasetShard(
+                new_dataset,
+                batch_size=dataloader.batch_size,
+                drop_last=dataloader.drop_last,
+                num_processes=num_processes,
+                process_index=process_index,
+                split_batches=split_batches,
+            )
+        else:
+            if not use_seedable_sampler and hasattr(sampler, "generator"):
+                if sampler.generator is None:
+                    sampler.generator = torch.Generator()
+                synchronized_generator = sampler.generator
+            batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler
+            new_batch_sampler = BatchSamplerShard(
+                batch_sampler,
+                num_processes=num_processes,
+                process_index=process_index,
+                split_batches=split_batches,
+                even_batches=even_batches,
+            )
+
+    # We ignore all of those since they are all dealt with by our new_batch_sampler
+    ignore_kwargs = [
+        "batch_size",
+        "shuffle",
+        "sampler",
+        "batch_sampler",
+        "drop_last",
+    ]
+
+    if rng_types is not None and synchronized_generator is None and "generator" in rng_types:
+        rng_types.remove("generator")
+
+    kwargs = {
+        k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k])
+        for k in _PYTORCH_DATALOADER_KWARGS
+        if k not in ignore_kwargs
+    }
+
+    # Need to provide batch_size as batch_sampler is None for Iterable dataset
+    if new_batch_sampler is None:
+        kwargs["drop_last"] = dataloader.drop_last
+        kwargs["batch_size"] = (
+            dataloader.batch_size // num_processes if split_batches and not dispatch_batches else dataloader.batch_size
+        )
+    if dispatch_batches:
+        kwargs.pop("generator")
+        dataloader = DataLoaderDispatcher(
+            new_dataset,
+            split_batches=split_batches,
+            batch_sampler=new_batch_sampler,
+            _drop_last=dataloader.drop_last,
+            _non_blocking=non_blocking,
+            slice_fn=slice_fn_for_dispatch,
+            use_stateful_dataloader=use_stateful_dataloader,
+            **kwargs,
+        )
+    elif sampler_is_batch_sampler:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device and state.distributed_type != DistributedType.XLA else None,
+            sampler=new_batch_sampler,
+            batch_size=dataloader.batch_size,
+            rng_types=rng_types,
+            _drop_last=dataloader.drop_last,
+            _non_blocking=non_blocking,
+            synchronized_generator=synchronized_generator,
+            use_stateful_dataloader=use_stateful_dataloader,
+            **kwargs,
+        )
+    else:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device and state.distributed_type != DistributedType.XLA else None,
+            batch_sampler=new_batch_sampler,
+            rng_types=rng_types,
+            synchronized_generator=synchronized_generator,
+            _drop_last=dataloader.drop_last,
+            _non_blocking=non_blocking,
+            use_stateful_dataloader=use_stateful_dataloader,
+            **kwargs,
+        )
+
+    if isinstance(sampler, SeedableRandomSampler) and use_seedable_sampler:
+        dataloader.set_sampler(sampler)
+    if state.distributed_type == DistributedType.XLA:
+        return MpDeviceLoaderWrapper(dataloader, device)
+    return dataloader
+
+
+class SkipBatchSampler(BatchSampler):
+    """
+    A `torch.utils.data.BatchSampler` that skips the first `n` batches of another `torch.utils.data.BatchSampler`.
+    Should not be used if the original dataloader is a `StatefulDataLoader`.
+    """
+
+    def __init__(self, batch_sampler, skip_batches=0):
+        self.batch_sampler = batch_sampler
+        self.skip_batches = skip_batches
+
+    def __iter__(self):
+        for index, samples in enumerate(self.batch_sampler):
+            if index >= self.skip_batches:
+                yield samples
+
+    @property
+    def total_length(self):
+        return len(self.batch_sampler)
+
+    def __len__(self):
+        return len(self.batch_sampler) - self.skip_batches
+
+
+class SkipDataLoader(DataLoaderAdapter, DataLoaderStateMixin):
+    """
+    Subclass of a PyTorch `DataLoader` that will skip the first batches. Generally it's preferable to use
+    `skip_first_batches`/`torchdata.StatefulDataLoader` instead of this class.
+
+    Args:
+        dataset (`torch.utils.data.dataset.Dataset`):
+            The dataset to use to build this dataloader.
+        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.
+    """
+
+    def __init__(self, dataset, skip_batches=0, use_stateful_dataloader=False, **kwargs):
+        super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs)
+        self.skip_batches = skip_batches
+        self.gradient_state = GradientState()
+
+    def __iter__(self):
+        self.begin()
+        for index, batch in enumerate(self.base_dataloader.__iter__()):
+            if index >= self.skip_batches:
+                self._update_state_dict()
+                yield batch
+        self.end()
+
+    def __len__(self):
+        return len(self.base_dataloader) - self.skip_batches
+
+    def __reduce__(self):
+        """
+        Define the `__reduce__` method to ensure a `SkipDataLoader` can be pickled and unpickled. This needs to be
+        explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its
+        `__class__` member.
+        """
+        args = super().__reduce__()
+        return (SkipDataLoader, *args[1:])
+
+
+def skip_first_batches(dataloader, num_batches=0):
+    """
+    Creates a `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`. Should not be used if
+    the original dataloader is a `StatefulDataLoader`.
+    """
+    state = PartialState()
+    if state.distributed_type == DistributedType.XLA:
+        device = dataloader.device
+        dataloader = dataloader.dataloader
+
+    dataset = dataloader.dataset
+    sampler_is_batch_sampler = False
+    if isinstance(dataset, IterableDataset):
+        new_batch_sampler = None
+    else:
+        sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler)
+        batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler
+        new_batch_sampler = SkipBatchSampler(batch_sampler, skip_batches=num_batches)
+
+    # We ignore all of those since they are all dealt with by our new_batch_sampler
+    ignore_kwargs = [
+        "batch_size",
+        "shuffle",
+        "sampler",
+        "batch_sampler",
+        "drop_last",
+    ]
+
+    kwargs = {
+        k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k])
+        for k in _PYTORCH_DATALOADER_KWARGS
+        if k not in ignore_kwargs
+    }
+
+    # Need to provide batch_size as batch_sampler is None for Iterable dataset
+    if new_batch_sampler is None:
+        kwargs["drop_last"] = dataloader.drop_last
+        kwargs["batch_size"] = dataloader.batch_size
+
+    if isinstance(dataloader, DataLoaderDispatcher):
+        if new_batch_sampler is None:
+            # Need to manually skip batches in the dataloader
+            kwargs["skip_batches"] = num_batches
+        dataloader = DataLoaderDispatcher(
+            dataset,
+            split_batches=dataloader.split_batches,
+            batch_sampler=new_batch_sampler,
+            _drop_last=dataloader._drop_last,
+            **kwargs,
+        )
+    elif isinstance(dataloader, DataLoaderShard):
+        if new_batch_sampler is None:
+            # Need to manually skip batches in the dataloader
+            kwargs["skip_batches"] = num_batches
+        elif sampler_is_batch_sampler:
+            kwargs["sampler"] = new_batch_sampler
+            kwargs["batch_size"] = dataloader.batch_size
+        else:
+            kwargs["batch_sampler"] = new_batch_sampler
+        dataloader = DataLoaderShard(
+            dataset,
+            device=dataloader.device,
+            rng_types=dataloader.rng_types,
+            synchronized_generator=dataloader.synchronized_generator,
+            **kwargs,
+        )
+    else:
+        if new_batch_sampler is None:
+            # Need to manually skip batches in the dataloader
+            dataloader = SkipDataLoader(dataset, skip_batches=num_batches, **kwargs)
+        else:
+            dataloader = DataLoader(dataset, batch_sampler=new_batch_sampler, **kwargs)
+
+    if state.distributed_type == DistributedType.XLA:
+        dataloader = MpDeviceLoaderWrapper(dataloader, device)
+
+    return dataloader

.venv/Lib/site-packages/accelerate/hooks.py

@@ -0,0 +1,726 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+from typing import Dict, List, Mapping, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .state import PartialState
+from .utils import (
+    PrefixedDataset,
+    find_device,
+    named_module_tensors,
+    send_to_device,
+    set_module_tensor_to_device,
+)
+from .utils.memory import clear_device_cache
+from .utils.modeling import get_non_persistent_buffers
+from .utils.other import recursive_getattr
+
+
+_accelerate_added_attributes = ["to", "cuda", "npu", "xpu", "mlu", "musa"]
+
+
+class ModelHook:
+    """
+    A hook that contains callbacks to be executed just before and after the forward method of a model. The difference
+    with PyTorch existing hooks is that they get passed along the kwargs.
+
+    Class attribute:
+    - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under
+      the `torch.no_grad()` context manager.
+    """
+
+    no_grad = False
+
+    def init_hook(self, module):
+        """
+        To be executed when the hook is attached to the module.
+
+        Args:
+            module (`torch.nn.Module`): The module attached to this hook.
+        """
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        """
+        To be executed just before the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass will be executed just after this event.
+            args (`Tuple[Any]`): The positional arguments passed to the module.
+            kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module.
+
+        Returns:
+            `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`.
+        """
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        """
+        To be executed just after the forward method of the model.
+
+        Args:
+            module (`torch.nn.Module`): The module whose forward pass been executed just before this event.
+            output (`Any`): The output of the module.
+
+        Returns:
+            `Any`: The processed `output`.
+        """
+        return output
+
+    def detach_hook(self, module):
+        """
+        To be executed when the hook is detached from a module.
+
+        Args:
+            module (`torch.nn.Module`): The module detached from this hook.
+        """
+        return module
+
+
+class SequentialHook(ModelHook):
+    """
+    A hook that can contain several hooks and iterates through them at each event.
+    """
+
+    def __init__(self, *hooks):
+        self.hooks = hooks
+
+    def init_hook(self, module):
+        for hook in self.hooks:
+            module = hook.init_hook(module)
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        for hook in self.hooks:
+            args, kwargs = hook.pre_forward(module, *args, **kwargs)
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        for hook in self.hooks:
+            output = hook.post_forward(module, output)
+        return output
+
+    def detach_hook(self, module):
+        for hook in self.hooks:
+            module = hook.detach_hook(module)
+        return module
+
+
+def add_hook_to_module(module: nn.Module, hook: ModelHook, append: bool = False):
+    """
+    Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove
+    this behavior and restore the original `forward` method, use `remove_hook_from_module`.
+
+    <Tip warning={true}>
+
+    If the module already contains a hook, this will replace it with the new hook passed by default. To chain two hooks
+    together, pass `append=True`, so it chains the current and new hook into an instance of the `SequentialHook` class.
+
+    </Tip>
+
+    Args:
+        module (`torch.nn.Module`):
+            The module to attach a hook to.
+        hook (`ModelHook`):
+            The hook to attach.
+        append (`bool`, *optional*, defaults to `False`):
+            Whether the hook should be chained with an existing one (if module already contains a hook) or not.
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can
+        be discarded).
+    """
+
+    if append and (getattr(module, "_hf_hook", None) is not None):
+        old_hook = module._hf_hook
+        remove_hook_from_module(module)
+        hook = SequentialHook(old_hook, hook)
+
+    if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"):
+        # If we already put some hook on this module, we replace it with the new one.
+        old_forward = module._old_forward
+    else:
+        old_forward = module.forward
+        module._old_forward = old_forward
+
+    module = hook.init_hook(module)
+    module._hf_hook = hook
+
+    def new_forward(module, *args, **kwargs):
+        args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
+        if module._hf_hook.no_grad:
+            with torch.no_grad():
+                output = module._old_forward(*args, **kwargs)
+        else:
+            output = module._old_forward(*args, **kwargs)
+        return module._hf_hook.post_forward(module, output)
+
+    # Overriding a GraphModuleImpl forward freezes the forward call and later modifications on the graph will fail.
+    # Reference: https://pytorch.slack.com/archives/C3PDTEV8E/p1705929610405409
+    if "GraphModuleImpl" in str(type(module)):
+        module.__class__.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward)
+    else:
+        module.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward)
+
+    return module
+
+
+def remove_hook_from_module(module: nn.Module, recurse=False):
+    """
+    Removes any hook attached to a module via `add_hook_to_module`.
+
+    Args:
+        module (`torch.nn.Module`): The module to attach a hook to.
+        recurse (`bool`, **optional**): Whether to remove the hooks recursively
+
+    Returns:
+        `torch.nn.Module`: The same module, with the hook detached (the module is modified in place, so the result can
+        be discarded).
+    """
+
+    if hasattr(module, "_hf_hook"):
+        module._hf_hook.detach_hook(module)
+        delattr(module, "_hf_hook")
+
+    if hasattr(module, "_old_forward"):
+        # Overriding a GraphModuleImpl forward freezes the forward call and later modifications on the graph will fail.
+        # Reference: https://pytorch.slack.com/archives/C3PDTEV8E/p1705929610405409
+        if "GraphModuleImpl" in str(type(module)):
+            module.__class__.forward = module._old_forward
+        else:
+            module.forward = module._old_forward
+        delattr(module, "_old_forward")
+
+    # Remove accelerate added warning hooks from dispatch_model
+    for attr in _accelerate_added_attributes:
+        module.__dict__.pop(attr, None)
+
+    if recurse:
+        for child in module.children():
+            remove_hook_from_module(child, recurse)
+
+    return module
+
+
+class AlignDevicesHook(ModelHook):
+    """
+    A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the
+    associated module, potentially offloading the weights after the forward pass.
+
+    Args:
+        execution_device (`torch.device`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass.
+        offload (`bool`, *optional*, defaults to `False`):
+            Whether or not the weights should be offloaded after the forward pass.
+        io_same_device (`bool`, *optional*, defaults to `False`):
+            Whether or not the output should be placed on the same device as the input was.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        place_submodules (`bool`, *optional*, defaults to `False`):
+            Whether to place the submodules on `execution_device` during the `init_hook` event.
+    """
+
+    def __init__(
+        self,
+        execution_device: Optional[Union[int, str, torch.device]] = None,
+        offload: bool = False,
+        io_same_device: bool = False,
+        weights_map: Optional[Mapping] = None,
+        offload_buffers: bool = False,
+        place_submodules: bool = False,
+        skip_keys: Optional[Union[str, List[str]]] = None,
+        tied_params_map: Optional[Dict[int, Dict[torch.device, torch.Tensor]]] = None,
+    ):
+        self.execution_device = execution_device
+        self.offload = offload
+        self.io_same_device = io_same_device
+        self.weights_map = weights_map
+        self.offload_buffers = offload_buffers
+        self.place_submodules = place_submodules
+        self.skip_keys = skip_keys
+
+        # Will contain the input device when `io_same_device=True`.
+        self.input_device = None
+        self.param_original_devices = {}
+        self.buffer_original_devices = {}
+        self.tied_params_names = set()
+
+        # The hook pre_forward/post_forward need to have knowledge of this dictionary, as with offloading we want to avoid duplicating memory
+        # for tied weights already loaded on the target execution device.
+        self.tied_params_map = tied_params_map
+
+    def __repr__(self):
+        return (
+            f"AlignDevicesHook(execution_device={self.execution_device}, offload={self.offload}, "
+            f"io_same_device={self.io_same_device}, offload_buffers={self.offload_buffers}, "
+            f"place_submodules={self.place_submodules}, skip_keys={repr(self.skip_keys)})"
+        )
+
+    def init_hook(self, module):
+        # In case the AlignDevicesHook is on meta device, ignore tied weights as data_ptr() is then always zero.
+        if self.execution_device == "meta" or self.execution_device == torch.device("meta"):
+            self.tied_params_map = None
+
+        if not self.offload and self.execution_device is not None:
+            for name, _ in named_module_tensors(module, recurse=self.place_submodules):
+                set_module_tensor_to_device(module, name, self.execution_device, tied_params_map=self.tied_params_map)
+        elif self.offload:
+            self.original_devices = {
+                name: param.device for name, param in named_module_tensors(module, recurse=self.place_submodules)
+            }
+            if self.weights_map is None:
+                self.weights_map = {
+                    name: param.to("cpu")
+                    for name, param in named_module_tensors(
+                        module, include_buffers=self.offload_buffers, recurse=self.place_submodules
+                    )
+                }
+            for name, _ in named_module_tensors(
+                module, include_buffers=self.offload_buffers, recurse=self.place_submodules, remove_non_persistent=True
+            ):
+                # When using disk offloading, we can not rely on `weights_map[name].data_ptr()` as the reference pointer,
+                # as we have no guarantee that safetensors' `file.get_tensor()` will always give the same pointer.
+                # As we have no reliable way to track the shared data pointer of tied weights in this case, we use tied_params_names: List[str]
+                # to add on the fly pointers to `tied_params_map` in the pre_forward call.
+                if (
+                    self.tied_params_map is not None
+                    and recursive_getattr(module, name).data_ptr() in self.tied_params_map
+                ):
+                    self.tied_params_names.add(name)
+
+                set_module_tensor_to_device(module, name, "meta")
+
+            if not self.offload_buffers and self.execution_device is not None:
+                for name, _ in module.named_buffers(recurse=self.place_submodules):
+                    set_module_tensor_to_device(
+                        module, name, self.execution_device, tied_params_map=self.tied_params_map
+                    )
+            elif self.offload_buffers and self.execution_device is not None:
+                for name in get_non_persistent_buffers(module, recurse=self.place_submodules):
+                    set_module_tensor_to_device(
+                        module, name, self.execution_device, tied_params_map=self.tied_params_map
+                    )
+
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        if self.io_same_device:
+            self.input_device = find_device([args, kwargs])
+        if self.offload:
+            self.tied_pointers_to_remove = set()
+
+            for name, _ in named_module_tensors(
+                module,
+                include_buffers=self.offload_buffers,
+                recurse=self.place_submodules,
+                remove_non_persistent=True,
+            ):
+                fp16_statistics = None
+                value = self.weights_map[name]
+                if "weight" in name and name.replace("weight", "SCB") in self.weights_map.keys():
+                    if value.dtype == torch.int8:
+                        fp16_statistics = self.weights_map[name.replace("weight", "SCB")]
+
+                # In case we are using offloading with tied weights, we need to keep track of the offloaded weights
+                # that are loaded on device at this point, as we will need to remove them as well from the dictionary
+                # self.tied_params_map in order to allow to free memory.
+                if name in self.tied_params_names and value.data_ptr() not in self.tied_params_map:
+                    self.tied_params_map[value.data_ptr()] = {}
+
+                if (
+                    value is not None
+                    and self.tied_params_map is not None
+                    and value.data_ptr() in self.tied_params_map
+                    and self.execution_device not in self.tied_params_map[value.data_ptr()]
+                ):
+                    self.tied_pointers_to_remove.add((value.data_ptr(), self.execution_device))
+
+                set_module_tensor_to_device(
+                    module,
+                    name,
+                    self.execution_device,
+                    value=value,
+                    fp16_statistics=fp16_statistics,
+                    tied_params_map=self.tied_params_map,
+                )
+
+        return send_to_device(args, self.execution_device), send_to_device(
+            kwargs, self.execution_device, skip_keys=self.skip_keys
+        )
+
+    def post_forward(self, module, output):
+        if self.offload:
+            for name, _ in named_module_tensors(
+                module,
+                include_buffers=self.offload_buffers,
+                recurse=self.place_submodules,
+                remove_non_persistent=True,
+            ):
+                set_module_tensor_to_device(module, name, "meta")
+                if type(module).__name__ == "Linear8bitLt":
+                    module.state.SCB = None
+                    module.state.CxB = None
+
+            # We may have loaded tied weights into self.tied_params_map (avoiding to load them several times in e.g. submodules): remove them from
+            # this dictionary to allow the garbage collector to do its job.
+            for value_pointer, device in self.tied_pointers_to_remove:
+                del self.tied_params_map[value_pointer][device]
+            self.tied_pointers_to_remove = set()
+
+        if self.io_same_device and self.input_device is not None:
+            output = send_to_device(output, self.input_device, skip_keys=self.skip_keys)
+
+        return output
+
+    def detach_hook(self, module):
+        if self.offload:
+            for name, device in self.original_devices.items():
+                if device != torch.device("meta"):
+                    set_module_tensor_to_device(module, name, device, value=self.weights_map.get(name, None))
+        return module
+
+
+def attach_execution_device_hook(
+    module: torch.nn.Module,
+    execution_device: Union[int, str, torch.device],
+    skip_keys: Optional[Union[str, List[str]]] = None,
+    preload_module_classes: Optional[List[str]] = None,
+    tied_params_map: Optional[Dict[int, Dict[torch.device, torch.Tensor]]] = None,
+):
+    """
+    Recursively attaches `AlignDevicesHook` to all submodules of a given model to make sure they have the right
+    execution device
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`int`, `str` or `torch.device`):
+            The device on which inputs and model weights should be placed before the forward pass.
+        skip_keys (`str` or `List[str]`, *optional*):
+            A list of keys to ignore when moving inputs or outputs between devices.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`):
+            A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution
+            device, this parameter is useful to reuse the first available pointer of a shared weight for all others,
+            instead of duplicating memory.
+    """
+    if not hasattr(module, "_hf_hook") and len(module.state_dict()) > 0:
+        add_hook_to_module(
+            module,
+            AlignDevicesHook(execution_device, skip_keys=skip_keys, tied_params_map=tied_params_map),
+        )
+
+    # Break the recursion if we get to a preload module.
+    if preload_module_classes is not None and module.__class__.__name__ in preload_module_classes:
+        return
+
+    for child in module.children():
+        attach_execution_device_hook(
+            child,
+            execution_device,
+            skip_keys=skip_keys,
+            preload_module_classes=preload_module_classes,
+            tied_params_map=tied_params_map,
+        )
+
+
+def attach_align_device_hook(
+    module: torch.nn.Module,
+    execution_device: Optional[torch.device] = None,
+    offload: bool = False,
+    weights_map: Optional[Mapping] = None,
+    offload_buffers: bool = False,
+    module_name: str = "",
+    skip_keys: Optional[Union[str, List[str]]] = None,
+    preload_module_classes: Optional[List[str]] = None,
+    tied_params_map: Optional[Dict[int, Dict[torch.device, torch.Tensor]]] = None,
+):
+    """
+    Recursively attaches `AlignDevicesHook` to all submodules of a given model that have direct parameters and/or
+    buffers.
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`torch.device`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass.
+        offload (`bool`, *optional*, defaults to `False`):
+            Whether or not the weights should be offloaded after the forward pass.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        module_name (`str`, *optional*, defaults to `""`):
+            The name of the module.
+        skip_keys (`str` or `List[str]`, *optional*):
+            A list of keys to ignore when moving inputs or outputs between devices.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`):
+            A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution
+            device, this parameter is useful to reuse the first available pointer of a shared weight for all others,
+            instead of duplicating memory.
+    """
+    # Attach the hook on this module if it has any direct tensor.
+    directs = named_module_tensors(module)
+    full_offload = (
+        offload and preload_module_classes is not None and module.__class__.__name__ in preload_module_classes
+    )
+
+    if len(list(directs)) > 0 or full_offload:
+        if weights_map is not None:
+            prefix = f"{module_name}." if len(module_name) > 0 else ""
+            prefixed_weights_map = PrefixedDataset(weights_map, prefix)
+        else:
+            prefixed_weights_map = None
+        hook = AlignDevicesHook(
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=prefixed_weights_map,
+            offload_buffers=offload_buffers,
+            place_submodules=full_offload,
+            skip_keys=skip_keys,
+            tied_params_map=tied_params_map,
+        )
+        add_hook_to_module(module, hook, append=True)
+
+    # We stop the recursion in case we hit the full offload.
+    if full_offload:
+        return
+
+    # Recurse on all children of the module.
+    for child_name, child in module.named_children():
+        child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name
+        attach_align_device_hook(
+            child,
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=weights_map,
+            offload_buffers=offload_buffers,
+            module_name=child_name,
+            preload_module_classes=preload_module_classes,
+            skip_keys=skip_keys,
+            tied_params_map=tied_params_map,
+        )
+
+
+def remove_hook_from_submodules(module: nn.Module):
+    """
+    Recursively removes all hooks attached on the submodules of a given model.
+
+    Args:
+        module (`torch.nn.Module`): The module on which to remove all hooks.
+    """
+    remove_hook_from_module(module)
+    for child in module.children():
+        remove_hook_from_submodules(child)
+
+
+def attach_align_device_hook_on_blocks(
+    module: nn.Module,
+    execution_device: Optional[Union[torch.device, Dict[str, torch.device]]] = None,
+    offload: Union[bool, Dict[str, bool]] = False,
+    weights_map: Mapping = None,
+    offload_buffers: bool = False,
+    module_name: str = "",
+    skip_keys: Optional[Union[str, List[str]]] = None,
+    preload_module_classes: Optional[List[str]] = None,
+    tied_params_map: Optional[Dict[int, Dict[torch.device, torch.Tensor]]] = None,
+):
+    """
+    Attaches `AlignDevicesHook` to all blocks of a given model as needed.
+
+    Args:
+        module (`torch.nn.Module`):
+            The module where we want to attach the hooks.
+        execution_device (`torch.device` or `Dict[str, torch.device]`, *optional*):
+            The device on which inputs and model weights should be placed before the forward pass. It can be one device
+            for the whole module, or a dictionary mapping module name to device.
+        offload (`bool`, *optional*, defaults to `False`):
+            Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole
+            module, or a dictionary mapping module name to boolean.
+        weights_map (`Mapping[str, torch.Tensor]`, *optional*):
+            When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
+        offload_buffers (`bool`, *optional*, defaults to `False`):
+            Whether or not to include the associated module's buffers when offloading.
+        module_name (`str`, *optional*, defaults to `""`):
+            The name of the module.
+        skip_keys (`str` or `List[str]`, *optional*):
+            A list of keys to ignore when moving inputs or outputs between devices.
+        preload_module_classes (`List[str]`, *optional*):
+            A list of classes whose instances should load all their weights (even in the submodules) at the beginning
+            of the forward. This should only be used for classes that have submodules which are registered but not
+            called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
+            `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`):
+            A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution
+            device, this parameter is useful to reuse the first available pointer of a shared weight for all others,
+            instead of duplicating memory.
+    """
+    # If one device and one offload, we've got one hook.
+    if not isinstance(execution_device, Mapping) and not isinstance(offload, dict):
+        if not offload:
+            hook = AlignDevicesHook(
+                execution_device=execution_device,
+                io_same_device=True,
+                skip_keys=skip_keys,
+                place_submodules=True,
+                tied_params_map=tied_params_map,
+            )
+            add_hook_to_module(module, hook)
+        else:
+            attach_align_device_hook(
+                module,
+                execution_device=execution_device,
+                offload=True,
+                weights_map=weights_map,
+                offload_buffers=offload_buffers,
+                module_name=module_name,
+                skip_keys=skip_keys,
+                tied_params_map=tied_params_map,
+            )
+        return
+
+    if not isinstance(execution_device, Mapping):
+        execution_device = {key: execution_device for key in offload.keys()}
+    if not isinstance(offload, Mapping):
+        offload = {key: offload for key in execution_device.keys()}
+
+    if module_name in execution_device and module_name in offload and not offload[module_name]:
+        hook = AlignDevicesHook(
+            execution_device=execution_device[module_name],
+            offload_buffers=offload_buffers,
+            io_same_device=(module_name == ""),
+            place_submodules=True,
+            skip_keys=skip_keys,
+            tied_params_map=tied_params_map,
+        )
+        add_hook_to_module(module, hook)
+        attach_execution_device_hook(
+            module, execution_device[module_name], skip_keys=skip_keys, tied_params_map=tied_params_map
+        )
+    elif module_name in execution_device and module_name in offload:
+        attach_align_device_hook(
+            module,
+            execution_device=execution_device[module_name],
+            offload=True,
+            weights_map=weights_map,
+            offload_buffers=offload_buffers,
+            module_name=module_name,
+            skip_keys=skip_keys,
+            preload_module_classes=preload_module_classes,
+            tied_params_map=tied_params_map,
+        )
+        if not hasattr(module, "_hf_hook"):
+            hook = AlignDevicesHook(
+                execution_device=execution_device[module_name],
+                io_same_device=(module_name == ""),
+                skip_keys=skip_keys,
+                tied_params_map=tied_params_map,
+            )
+            add_hook_to_module(module, hook)
+        attach_execution_device_hook(
+            module,
+            execution_device[module_name],
+            preload_module_classes=preload_module_classes,
+            skip_keys=skip_keys,
+            tied_params_map=tied_params_map,
+        )
+    elif module_name == "":
+        hook = AlignDevicesHook(
+            execution_device=execution_device.get(""),
+            io_same_device=True,
+            skip_keys=skip_keys,
+            tied_params_map=tied_params_map,
+        )
+        add_hook_to_module(module, hook)
+
+    for child_name, child in module.named_children():
+        child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name
+        attach_align_device_hook_on_blocks(
+            child,
+            execution_device=execution_device,
+            offload=offload,
+            weights_map=weights_map,
+            offload_buffers=offload_buffers,
+            module_name=child_name,
+            preload_module_classes=preload_module_classes,
+            skip_keys=skip_keys,
+            tied_params_map=tied_params_map,
+        )
+
+
+class CpuOffload(ModelHook):
+    """
+    Offloads a model on the CPU until its forward pass is called. The model will not be offloaded back to the CPU after
+    the forward, the user needs to call the `init_hook` method again for this.
+
+    Args:
+        execution_device(`str`, `int` or `torch.device`, *optional*):
+            The device on which the model should be executed. Will default to the MPS device if it's available, then
+            GPU 0 if there is a GPU, and finally to the CPU.
+        prev_module_hook (`UserCpuOffloadHook`, *optional*):
+            The hook sent back by [`cpu_offload_with_hook`] for a previous model in the pipeline you are running. If
+            passed, its offload method will be called just before the forward of the model to which this hook is
+            attached.
+    """
+
+    def __init__(
+        self,
+        execution_device: Optional[Union[str, int, torch.device]] = None,
+        prev_module_hook: Optional["UserCpuOffloadHook"] = None,
+    ):
+        self.prev_module_hook = prev_module_hook
+
+        self.execution_device = execution_device if execution_device is not None else PartialState().default_device
+
+    def init_hook(self, module):
+        return module.to("cpu")
+
+    def pre_forward(self, module, *args, **kwargs):
+        if self.prev_module_hook is not None:
+            self.prev_module_hook.offload()
+            clear_device_cache()
+        module.to(self.execution_device)
+        return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)
+
+
+class UserCpuOffloadHook:
+    """
+    A simple hook grouping a model and a `ModelHook`, which provides easy APIs for to call the init method of the hook
+    or remove it entirely.
+    """
+
+    def __init__(self, model, hook):
+        self.model = model
+        self.hook = hook
+
+    def offload(self):
+        self.hook.init_hook(self.model)
+
+    def remove(self):
+        remove_hook_from_module(self.model)

.venv/Lib/site-packages/accelerate/inference.py

@@ -0,0 +1,184 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from types import MethodType
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from .state import PartialState
+from .utils import (
+    calculate_maximum_sizes,
+    convert_bytes,
+    copy_tensor_to_devices,
+    ignorant_find_batch_size,
+    infer_auto_device_map,
+    is_pippy_available,
+    pad_input_tensors,
+    send_to_device,
+)
+
+
+def generate_device_map(model, num_processes: int = 1, no_split_module_classes=None, max_memory: dict = None):
+    """
+    Calculates the device map for `model` with an offset for PiPPy
+    """
+    if num_processes == 1:
+        return infer_auto_device_map(model, no_split_module_classes=no_split_module_classes, clean_result=False)
+    if max_memory is None:
+        model_size, shared = calculate_maximum_sizes(model)
+
+        # Split into `n` chunks for each GPU
+        memory = (model_size + shared[0]) / num_processes
+        memory = convert_bytes(memory)
+        value, ending = memory.split(" ")
+
+        # Add a chunk to deal with potential extra shared memory instances
+        memory = math.ceil(float(value)) * 1.1
+        memory = f"{memory} {ending}"
+        max_memory = {i: memory for i in range(num_processes)}
+    device_map = infer_auto_device_map(
+        model,
+        max_memory=max_memory,
+        no_split_module_classes=no_split_module_classes,
+        clean_result=False,
+    )
+    return device_map
+
+
+def find_pippy_batch_size(args, kwargs):
+    found_batch_size = None
+    if args is not None:
+        for arg in args:
+            found_batch_size = ignorant_find_batch_size(arg)
+            if found_batch_size is not None:
+                break
+    if kwargs is not None and found_batch_size is None:
+        for kwarg in kwargs.values():
+            found_batch_size = ignorant_find_batch_size(kwarg)
+            if found_batch_size is not None:
+                break
+    return found_batch_size
+
+
+def build_pipeline(model, split_points, args, kwargs, num_chunks):
+    """
+    Attaches the split points to the model based on `self.device_map` and generates a `PipelineStage`. Requires passing
+    in needed `args` and `kwargs` as the model needs on the CPU.
+
+    Users can pass in custom `num_chunks` as an optional hyper-parameter. By default will use
+    `AcceleratorState.num_processes`
+    """
+    # Note: We import here to reduce import time from general modules, and isolate outside dependencies
+    from torch.distributed.pipelining import ScheduleGPipe, SplitPoint, pipeline
+
+    # We need to annotate the split points in the model for PiPPy
+    state = PartialState()
+    split_spec = {split_point: SplitPoint.BEGINNING for split_point in split_points}
+    pipe = pipeline(
+        model,
+        mb_args=args,
+        mb_kwargs=kwargs,
+        split_spec=split_spec,
+    )
+    stage = pipe.build_stage(state.local_process_index, device=state.device)
+    schedule = ScheduleGPipe(stage, num_chunks)
+
+    return schedule
+
+
+def pippy_forward(forward, num_chunks, gather_output, *args, **kwargs):
+    state = PartialState()
+    output = None
+
+    if state.num_processes == 1:
+        output = forward(*args, **kwargs)
+    elif state.is_local_main_process:
+        found_batch_size = find_pippy_batch_size(args, kwargs)
+        if found_batch_size is None:
+            raise ValueError("Could not find batch size from args or kwargs")
+        else:
+            if found_batch_size != num_chunks:
+                args = pad_input_tensors(args, found_batch_size, num_chunks)
+                kwargs = pad_input_tensors(kwargs, found_batch_size, num_chunks)
+        forward(*args, **kwargs)
+    elif state.is_last_process:
+        output = forward()
+    else:
+        forward()
+    if gather_output:
+        # Each node will get a copy of the full output which is only on the last GPU
+        output = copy_tensor_to_devices(output)
+    return output
+
+
+def prepare_pippy(
+    model,
+    split_points: Optional[Union[str, List[str]]] = "auto",
+    no_split_module_classes: Optional[List[str]] = None,
+    example_args: Optional[Tuple[Any]] = (),
+    example_kwargs: Optional[Dict[str, Any]] = None,
+    num_chunks: Optional[int] = None,
+    gather_output: Optional[bool] = False,
+):
+    """
+    Wraps `model` for pipeline parallel inference.
+
+    Args:
+        model (`torch.nn.Module`):
+            A model we want to split for pipeline-parallel inference
+        split_points (`str` or `List[str]`, defaults to 'auto'):
+            How to generate the split points and chunk the model across each GPU. 'auto' will find the best balanced
+            split given any model. Should be a list of layer names in the model to split by otherwise.
+        no_split_module_classes (`List[str]`):
+            A list of class names for layers we don't want to be split.
+        example_args (tuple of model inputs):
+            The expected inputs for the model that uses order-based inputs for a *single process*. Recommended to use
+            this method if possible.
+        example_kwargs (dict of model inputs)
+            The expected inputs for the model that uses dictionary-based inputs for a *single process*. This is a
+            *highly* limiting structure that requires the same keys be present at *all* inference calls. Not
+            recommended unless the prior condition is true for all cases.
+        num_chunks (`int`, defaults to the number of available GPUs):
+            The number of different stages the Pipeline will have. By default it will assign one chunk per GPU, but
+            this can be tuned and played with. In general one should have num_chunks >= num_gpus.
+        gather_output (`bool`, defaults to `False`):
+            If `True`, the output from the last GPU (which holds the true outputs) is sent across to all GPUs.
+    """
+    if not is_pippy_available():
+        raise ImportError("Using `torch.distributed.pipelining` requires PyTorch 2.4.0 or later.")
+    state = PartialState()
+    example_args = send_to_device(example_args, "cpu")
+    example_kwargs = send_to_device(example_kwargs, "cpu")
+    if num_chunks is None:
+        num_chunks = state.num_processes
+    if split_points == "auto":
+        device_map = generate_device_map(model, num_chunks, no_split_module_classes=no_split_module_classes)
+        split_points = []
+        for i in range(1, num_chunks):
+            split_points.append(next(k for k, v in device_map.items() if v == i))
+    model.hf_split_points = split_points
+    stage = build_pipeline(model, split_points, example_args, example_kwargs, num_chunks)
+    model._original_forward = model.forward
+    model._original_call = model.__call__
+    model.pippy_stage = stage
+    model.hf_split_points = split_points
+
+    def forward(*args, **kwargs):
+        return pippy_forward(stage.step, num_chunks, gather_output, *args, **kwargs)
+
+    # To act like a decorator so that it can be popped when doing `extract_model_from_parallel`
+    # Note: creates an infinite recursion loop with `generate`
+    model_forward = MethodType(forward, model)
+    forward.__wrapped__ = model_forward
+    model.forward = forward
+    return model

.venv/Lib/site-packages/accelerate/launchers.py

@@ -0,0 +1,302 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import tempfile
+
+import torch
+
+from .state import AcceleratorState, PartialState
+from .utils import (
+    PrecisionType,
+    PrepareForLaunch,
+    are_libraries_initialized,
+    check_cuda_p2p_ib_support,
+    get_gpu_info,
+    is_mps_available,
+    is_torch_version,
+    patch_environment,
+)
+from .utils.constants import ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION
+
+
+def test_launch():
+    "Verify a `PartialState` can be initialized."
+    _ = PartialState()
+
+
+def notebook_launcher(
+    function,
+    args=(),
+    num_processes=None,
+    mixed_precision="no",
+    use_port="29500",
+    master_addr="127.0.0.1",
+    node_rank=0,
+    num_nodes=1,
+    rdzv_backend="static",
+    rdzv_endpoint="",
+    rdzv_conf=None,
+    rdzv_id="none",
+    max_restarts=0,
+    monitor_interval=0.1,
+    log_line_prefix_template=None,
+):
+    """
+    Launches a training function, using several processes or multiple nodes if it's possible in the current environment
+    (TPU with multiple cores for instance).
+
+    <Tip warning={true}>
+
+    To use this function absolutely zero calls to a CUDA device must be made in the notebook session before calling. If
+    any have been made, you will need to restart the notebook and make sure no cells use any CUDA capability.
+
+    Setting `ACCELERATE_DEBUG_MODE="1"` in your environment will run a test before truly launching to ensure that none
+    of those calls have been made.
+
+    </Tip>
+
+    Args:
+        function (`Callable`):
+            The training function to execute. If it accepts arguments, the first argument should be the index of the
+            process run.
+        args (`Tuple`):
+            Tuple of arguments to pass to the function (it will receive `*args`).
+        num_processes (`int`, *optional*):
+            The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to
+            the number of GPUs available otherwise.
+        mixed_precision (`str`, *optional*, defaults to `"no"`):
+            If `fp16` or `bf16`, will use mixed precision training on multi-GPU.
+        use_port (`str`, *optional*, defaults to `"29500"`):
+            The port to use to communicate between processes when launching a multi-GPU training.
+        master_addr (`str`, *optional*, defaults to `"127.0.0.1"`):
+            The address to use for communication between processes.
+        node_rank (`int`, *optional*, defaults to 0):
+            The rank of the current node.
+        num_nodes (`int`, *optional*, defaults to 1):
+            The number of nodes to use for training.
+        rdzv_backend (`str`, *optional*, defaults to `"static"`):
+            The rendezvous method to use, such as 'static' (the default) or 'c10d'
+        rdzv_endpoint (`str`, *optional*, defaults to `""`):
+            The endpoint of the rdzv sync. storage.
+        rdzv_conf (`Dict`, *optional*, defaults to `None`):
+            Additional rendezvous configuration.
+        rdzv_id (`str`, *optional*, defaults to `"none"`):
+            The unique run id of the job.
+        max_restarts (`int`, *optional*, defaults to 0):
+            The maximum amount of restarts that elastic agent will conduct on workers before failure.
+        monitor_interval (`float`, *optional*, defaults to 0.1):
+            The interval in seconds that is used by the elastic_agent as a period of monitoring workers.
+        log_line_prefix_template (`str`, *optional*, defaults to `None`):
+            The prefix template for elastic launch logging. Available from PyTorch 2.2.0.
+
+    Example:
+
+    ```python
+    # Assume this is defined in a Jupyter Notebook on an instance with two GPUs
+    from accelerate import notebook_launcher
+
+
+    def train(*args):
+        # Your training function here
+        ...
+
+
+    notebook_launcher(train, args=(arg1, arg2), num_processes=2, mixed_precision="fp16")
+    ```
+    """
+    # Are we in a google colab or a Kaggle Kernel?
+    in_colab = False
+    in_kaggle = False
+    if any(key.startswith("KAGGLE") for key in os.environ.keys()):
+        in_kaggle = True
+    elif "IPython" in sys.modules:
+        in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython())
+
+    try:
+        mixed_precision = PrecisionType(mixed_precision.lower())
+    except ValueError:
+        raise ValueError(
+            f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."
+        )
+
+    if (in_colab or in_kaggle) and (os.environ.get("TPU_NAME", None) is not None):
+        # TPU launch
+        import torch_xla.distributed.xla_multiprocessing as xmp
+
+        if len(AcceleratorState._shared_state) > 0:
+            raise ValueError(
+                "To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside "
+                "your training function. Restart your notebook and make sure no cells initializes an "
+                "`Accelerator`."
+            )
+        if num_processes is None:
+            num_processes = 8
+
+        launcher = PrepareForLaunch(function, distributed_type="XLA")
+        print(f"Launching a training on {num_processes} TPU cores.")
+        xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
+    elif in_colab and get_gpu_info()[1] < 2:
+        # No need for a distributed launch otherwise as it's either CPU or one GPU.
+        if torch.cuda.is_available():
+            print("Launching training on one GPU.")
+        else:
+            print("Launching training on one CPU.")
+        function(*args)
+    else:
+        if num_processes is None:
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call."
+            )
+        if node_rank >= num_nodes:
+            raise ValueError("The node_rank must be less than the number of nodes.")
+        if num_processes > 1:
+            # Multi-GPU launch
+            from torch.distributed.launcher.api import LaunchConfig, elastic_launch
+            from torch.multiprocessing import start_processes
+            from torch.multiprocessing.spawn import ProcessRaisedException
+
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+            # Check for specific libraries known to initialize CUDA that users constantly use
+            problematic_imports = are_libraries_initialized("bitsandbytes")
+            if len(problematic_imports) > 0:
+                err = (
+                    "Could not start distributed process. Libraries known to initialize CUDA upon import have been "
+                    "imported already. Please keep these imports inside your training function to try and help with this:"
+                )
+                for lib_name in problematic_imports:
+                    err += f"\n\t* `{lib_name}`"
+                raise RuntimeError(err)
+
+            patched_env = dict(
+                nproc=num_processes,
+                node_rank=node_rank,
+                world_size=num_nodes * num_processes,
+                master_addr=master_addr,
+                master_port=use_port,
+                mixed_precision=mixed_precision,
+            )
+
+            # Check for CUDA P2P and IB issues
+            if not check_cuda_p2p_ib_support():
+                patched_env["nccl_p2p_disable"] = "1"
+                patched_env["nccl_ib_disable"] = "1"
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            with patch_environment(**patched_env):
+                # First dummy launch
+                if os.environ.get("ACCELERATE_DEBUG_MODE", "false").lower() == "true":
+                    launcher = PrepareForLaunch(test_launch, distributed_type="MULTI_GPU")
+                    try:
+                        start_processes(launcher, args=(), nprocs=num_processes, start_method="fork")
+                    except ProcessRaisedException as e:
+                        err = "An issue was found when verifying a stable environment for the notebook launcher."
+                        if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
+                            raise RuntimeError(
+                                f"{err}"
+                                "This likely stems from an outside import causing issues once the `notebook_launcher()` is called. "
+                                "Please review your imports and test them when running the `notebook_launcher()` to identify "
+                                "which one is problematic and causing CUDA to be initialized."
+                            ) from e
+                        else:
+                            raise RuntimeError(f"{err} The following error was raised: {e}") from e
+                # Now the actual launch
+                launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
+                print(f"Launching training on {num_processes} GPUs.")
+                try:
+                    if rdzv_conf is None:
+                        rdzv_conf = {}
+                    if rdzv_backend == "static":
+                        rdzv_conf["rank"] = node_rank
+                        if not rdzv_endpoint:
+                            rdzv_endpoint = f"{master_addr}:{use_port}"
+                    launch_config_kwargs = dict(
+                        min_nodes=num_nodes,
+                        max_nodes=num_nodes,
+                        nproc_per_node=num_processes,
+                        run_id=rdzv_id,
+                        rdzv_endpoint=rdzv_endpoint,
+                        rdzv_backend=rdzv_backend,
+                        rdzv_configs=rdzv_conf,
+                        max_restarts=max_restarts,
+                        monitor_interval=monitor_interval,
+                        start_method="fork",
+                    )
+                    if is_torch_version(">=", ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION):
+                        launch_config_kwargs["log_line_prefix_template"] = log_line_prefix_template
+                    elastic_launch(config=LaunchConfig(**launch_config_kwargs), entrypoint=function)(*args)
+                except ProcessRaisedException as e:
+                    if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
+                        raise RuntimeError(
+                            "CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. "
+                            "This likely stems from an outside import causing issues once the `notebook_launcher()` is called. "
+                            "Please review your imports and test them when running the `notebook_launcher()` to identify "
+                            "which one is problematic and causing CUDA to be initialized."
+                        ) from e
+                    else:
+                        raise RuntimeError(f"An issue was found when launching the training: {e}") from e
+
+        else:
+            # No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
+            if is_mps_available():
+                os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+                print("Launching training on MPS.")
+            elif torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
+
+
+def debug_launcher(function, args=(), num_processes=2):
+    """
+    Launches a training function using several processes on CPU for debugging purposes.
+
+    <Tip warning={true}>
+
+    This function is provided for internal testing and debugging, but it's not intended for real trainings. It will
+    only use the CPU.
+
+    </Tip>
+
+    Args:
+        function (`Callable`):
+            The training function to execute.
+        args (`Tuple`):
+            Tuple of arguments to pass to the function (it will receive `*args`).
+        num_processes (`int`, *optional*, defaults to 2):
+            The number of processes to use for training.
+    """
+    from torch.multiprocessing import start_processes
+
+    with tempfile.NamedTemporaryFile() as tmp_file:
+        # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+        # process here (the other ones will be set be the launcher).
+        with patch_environment(
+            world_size=num_processes,
+            master_addr="127.0.0.1",
+            master_port="29500",
+            accelerate_mixed_precision="no",
+            accelerate_debug_rdv_file=tmp_file.name,
+            accelerate_use_cpu="yes",
+        ):
+            launcher = PrepareForLaunch(function, debug=True)
+            start_processes(launcher, args=args, nprocs=num_processes, start_method="fork")

.venv/Lib/site-packages/accelerate/local_sgd.py

@@ -0,0 +1,104 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+
+from accelerate import Accelerator, DistributedType
+
+
+class LocalSGD:
+    """
+    A helper class to support local SGD on top of Accelerator. It simply runs a given number of updates independently
+    on each device, and averages model weights every K synchronization step.
+
+    It should be used only in the multi-GPU (or multi-CPU) setup without extensions such as DeepSpeed. In particular,
+    this is a simple implementation that cannot support scenarios such as model parallelism.
+
+
+    Although we are not aware of the true origins of this simple approach, the idea of local SGD is quite old and goes
+    back to at least:
+
+    Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). [Parallel SGD: When does averaging help?. arXiv preprint
+    arXiv:1606.07365.](https://arxiv.org/abs/1606.07365)
+
+    We credit the term Local SGD to the following paper (but there might be earlier references we are not aware of).
+
+    Stich, Sebastian Urban. ["Local SGD Converges Fast and Communicates Little." ICLR 2019-International Conference on
+    Learning Representations. No. CONF. 2019.](https://arxiv.org/abs/1805.09767)
+
+    """
+
+    def __enter__(self):
+        if self.enabled:
+            self.model_sync_obj = self.model.no_sync()
+            self.model_sync_obj.__enter__()
+
+        return self
+
+    def __exit__(self, type, value, tb):
+        if self.enabled:
+            # Average all models on exit
+            self._sync_and_avg_model_params()
+            self.model_sync_obj.__exit__(type, value, tb)
+
+    def __init__(self, accelerator: Accelerator, model: torch.nn.Module, local_sgd_steps: int, enabled: bool = True):
+        """
+        Constructor.
+
+        Args:
+            model (`torch.nn.Module):
+                The model whose parameters we need to average.
+            accelerator (`Accelerator`):
+                Accelerator object.
+            local_sgd_steps (`int`):
+                A number of local SGD steps (before model parameters are synchronized).
+            enabled (`bool):
+                Local SGD is disabled if this parameter set to `False`.
+        """
+        if accelerator.distributed_type not in [
+            DistributedType.NO,
+            DistributedType.MULTI_CPU,
+            DistributedType.MULTI_GPU,
+            DistributedType.MULTI_XPU,
+            DistributedType.MULTI_MLU,
+            DistributedType.MULTI_MUSA,
+            DistributedType.MULTI_NPU,
+        ]:
+            raise NotImplementedError("LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)")
+        self.enabled = enabled and accelerator.distributed_type != DistributedType.NO
+        self.num_steps = 0
+        if self.enabled:
+            self.accelerator = accelerator
+            self.model = model
+            self.local_sgd_steps = local_sgd_steps
+
+    def step(self):
+        """
+        This function makes a "step" and synchronizes model parameters if necessary.
+        """
+        self.num_steps += 1
+        if not self.enabled:
+            return
+
+        if self.num_steps % self.local_sgd_steps == 0:
+            self._sync_and_avg_model_params()
+
+    def _sync_and_avg_model_params(self):
+        """
+        Synchronize + Average model parameters across all GPUs
+        """
+
+        self.accelerator.wait_for_everyone()
+        with self.accelerator.autocast():
+            for param in self.model.parameters():
+                param.data = self.accelerator.reduce(param.data, reduction="mean")

.venv/Lib/site-packages/accelerate/logging.py

@@ -0,0 +1,125 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+import logging
+import os
+
+from .state import PartialState
+
+
+class MultiProcessAdapter(logging.LoggerAdapter):
+    """
+    An adapter to assist with logging in multiprocess.
+
+    `log` takes in an additional `main_process_only` kwarg, which dictates whether it should be called on all processes
+    or only the main executed one. Default is `main_process_only=True`.
+
+    Does not require an `Accelerator` object to be created first.
+    """
+
+    @staticmethod
+    def _should_log(main_process_only):
+        "Check if log should be performed"
+        state = PartialState()
+        return not main_process_only or (main_process_only and state.is_main_process)
+
+    def log(self, level, msg, *args, **kwargs):
+        """
+        Delegates logger call after checking if we should log.
+
+        Accepts a new kwarg of `main_process_only`, which will dictate whether it will be logged across all processes
+        or only the main executed one. Default is `True` if not passed
+
+        Also accepts "in_order", which if `True` makes the processes log one by one, in order. This is much easier to
+        read, but comes at the cost of sometimes needing to wait for the other processes. Default is `False` to not
+        break with the previous behavior.
+
+        `in_order` is ignored if `main_process_only` is passed.
+        """
+        if PartialState._shared_state == {}:
+            raise RuntimeError(
+                "You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility."
+            )
+        main_process_only = kwargs.pop("main_process_only", True)
+        in_order = kwargs.pop("in_order", False)
+        # set `stacklevel` to exclude ourself in `Logger.findCaller()` while respecting user's choice
+        kwargs.setdefault("stacklevel", 2)
+
+        if self.isEnabledFor(level):
+            if self._should_log(main_process_only):
+                msg, kwargs = self.process(msg, kwargs)
+                self.logger.log(level, msg, *args, **kwargs)
+
+            elif in_order:
+                state = PartialState()
+                for i in range(state.num_processes):
+                    if i == state.process_index:
+                        msg, kwargs = self.process(msg, kwargs)
+                        self.logger.log(level, msg, *args, **kwargs)
+                    state.wait_for_everyone()
+
+    @functools.lru_cache(None)
+    def warning_once(self, *args, **kwargs):
+        """
+        This method is identical to `logger.warning()`, but will emit the warning with the same message only once
+
+        Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the
+        cache. The assumption here is that all warning messages are unique across the code. If they aren't then need to
+        switch to another type of cache that includes the caller frame information in the hashing function.
+        """
+        self.warning(*args, **kwargs)
+
+
+def get_logger(name: str, log_level: str = None):
+    """
+    Returns a `logging.Logger` for `name` that can handle multiprocessing.
+
+    If a log should be called on all processes, pass `main_process_only=False` If a log should be called on all
+    processes and in order, also pass `in_order=True`
+
+    Args:
+        name (`str`):
+            The name for the logger, such as `__file__`
+        log_level (`str`, *optional*):
+            The log level to use. If not passed, will default to the `LOG_LEVEL` environment variable, or `INFO` if not
+
+    Example:
+
+    ```python
+    >>> from accelerate.logging import get_logger
+    >>> from accelerate import Accelerator
+
+    >>> logger = get_logger(__name__)
+
+    >>> accelerator = Accelerator()
+    >>> logger.info("My log", main_process_only=False)
+    >>> logger.debug("My log", main_process_only=True)
+
+    >>> logger = get_logger(__name__, log_level="DEBUG")
+    >>> logger.info("My log")
+    >>> logger.debug("My second log")
+
+    >>> array = ["a", "b", "c", "d"]
+    >>> letter_at_rank = array[accelerator.process_index]
+    >>> logger.info(letter_at_rank, in_order=True)
+    ```
+    """
+    if log_level is None:
+        log_level = os.environ.get("ACCELERATE_LOG_LEVEL", None)
+    logger = logging.getLogger(name)
+    if log_level is not None:
+        logger.setLevel(log_level.upper())
+        logger.root.setLevel(log_level.upper())
+    return MultiProcessAdapter(logger, {})

.venv/Lib/site-packages/accelerate/memory_utils.py

@@ -0,0 +1,22 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+
+
+warnings.warn(
+    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: "
+    "`from accelerate import find_executable_batch_size` to avoid this warning.",
+    FutureWarning,
+)

.venv/Lib/site-packages/accelerate/optimizer.py

@@ -0,0 +1,212 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import torch
+
+from .state import AcceleratorState, GradientState
+from .utils import DistributedType, honor_type, is_lomo_available, is_torch_xla_available
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+
+def move_to_device(state, device):
+    if isinstance(state, (list, tuple)):
+        return honor_type(state, (move_to_device(t, device) for t in state))
+    elif isinstance(state, dict):
+        return type(state)({k: move_to_device(v, device) for k, v in state.items()})
+    elif isinstance(state, torch.Tensor):
+        return state.to(device)
+    return state
+
+
+class AcceleratedOptimizer(torch.optim.Optimizer):
+    """
+    Internal wrapper around a torch optimizer.
+
+    Conditionally will perform `step` and `zero_grad` if gradients should be synchronized when performing gradient
+    accumulation.
+
+    Args:
+        optimizer (`torch.optim.optimizer.Optimizer`):
+            The optimizer to wrap.
+        device_placement (`bool`, *optional*, defaults to `True`):
+            Whether or not the optimizer should handle device placement. If so, it will place the state dictionary of
+            `optimizer` on the right device.
+        scaler (`torch.cuda.amp.grad_scaler.GradScaler`, *optional*):
+            The scaler to use in the step function if training with mixed precision.
+    """
+
+    def __init__(self, optimizer, device_placement=True, scaler=None):
+        self.optimizer = optimizer
+        self.scaler = scaler
+        self.accelerator_state = AcceleratorState()
+        self.gradient_state = GradientState()
+        self.device_placement = device_placement
+        self._is_overflow = False
+
+        if self.scaler is not None:
+            self._accelerate_step_called = False
+            self._optimizer_original_step_method = self.optimizer.step
+            self._optimizer_patched_step_method = patch_optimizer_step(self, self.optimizer.step)
+
+        # Handle device placement
+        if device_placement:
+            state_dict = self.optimizer.state_dict()
+            if self.accelerator_state.distributed_type == DistributedType.XLA:
+                xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device)
+            else:
+                state_dict = move_to_device(state_dict, self.accelerator_state.device)
+            self.optimizer.load_state_dict(state_dict)
+
+    @property
+    def state(self):
+        return self.optimizer.state
+
+    @state.setter
+    def state(self, state):
+        self.optimizer.state = state
+
+    @property
+    def param_groups(self):
+        return self.optimizer.param_groups
+
+    @param_groups.setter
+    def param_groups(self, param_groups):
+        self.optimizer.param_groups = param_groups
+
+    @property
+    def defaults(self):
+        return self.optimizer.defaults
+
+    @defaults.setter
+    def defaults(self, defaults):
+        self.optimizer.defaults = defaults
+
+    def add_param_group(self, param_group):
+        self.optimizer.add_param_group(param_group)
+
+    def load_state_dict(self, state_dict):
+        if self.accelerator_state.distributed_type == DistributedType.XLA and self.device_placement:
+            xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device)
+        self.optimizer.load_state_dict(state_dict)
+
+    def state_dict(self):
+        return self.optimizer.state_dict()
+
+    def zero_grad(self, set_to_none=None):
+        if self.gradient_state.sync_gradients:
+            accept_arg = "set_to_none" in inspect.signature(self.optimizer.zero_grad).parameters
+            if accept_arg:
+                if set_to_none is None:
+                    set_to_none = True
+                self.optimizer.zero_grad(set_to_none=set_to_none)
+            else:
+                if set_to_none is not None:
+                    raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.")
+                self.optimizer.zero_grad()
+
+    def train(self):
+        """
+        Sets the optimizer to "train" mode. Useful for optimizers like `schedule_free`
+        """
+        if hasattr(self.optimizer, "train") and callable(self.optimizer.train):
+            self.optimizer.train()
+        elif (
+            hasattr(self.optimizer, "optimizer")
+            and hasattr(self.optimizer.optimizer, "train")
+            and callable(self.optimizer.optimizer.train)
+        ):
+            # the deepspeed optimizer further wraps the optimizer
+            self.optimizer.optimizer.train()
+
+    def eval(self):
+        """
+        Sets the optimizer to "eval" mode. Useful for optimizers like `schedule_free`
+        """
+        if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval):
+            self.optimizer.eval()
+
+    def step(self, closure=None):
+        if is_lomo_available():
+            from lomo_optim import AdaLomo, Lomo
+
+        if (
+            not self.gradient_state.is_xla_gradients_synced
+            and self.accelerator_state.distributed_type == DistributedType.XLA
+        ):
+            gradients = xm._fetch_gradients(self.optimizer)
+            xm.all_reduce("sum", gradients, scale=1.0 / xm.xrt_world_size())
+            self.gradient_state.is_xla_gradients_synced = True
+
+        if is_lomo_available():
+            #  `step` should be a no-op for LOMO optimizers.
+            if isinstance(self.optimizer, (Lomo, AdaLomo)):
+                return
+
+        if self.gradient_state.sync_gradients:
+            if self.scaler is not None:
+                self.optimizer.step = self._optimizer_patched_step_method
+
+                self.scaler.step(self.optimizer, closure)
+                self.scaler.update()
+
+                if not self._accelerate_step_called:
+                    # If the optimizer step was skipped, gradient overflow was detected.
+                    self._is_overflow = True
+                else:
+                    self._is_overflow = False
+                # Reset the step method to the original one
+                self.optimizer.step = self._optimizer_original_step_method
+                # Reset the indicator
+                self._accelerate_step_called = False
+            else:
+                self.optimizer.step(closure)
+        if self.accelerator_state.distributed_type == DistributedType.XLA:
+            self.gradient_state.is_xla_gradients_synced = False
+
+    def _switch_parameters(self, parameters_map):
+        for param_group in self.optimizer.param_groups:
+            param_group["params"] = [parameters_map.get(p, p) for p in param_group["params"]]
+
+    @property
+    def step_was_skipped(self):
+        """Whether or not the optimizer step was skipped."""
+        return self._is_overflow
+
+    def __getstate__(self):
+        _ignored_keys = [
+            "_accelerate_step_called",
+            "_optimizer_original_step_method",
+            "_optimizer_patched_step_method",
+        ]
+        return {k: v for k, v in self.__dict__.items() if k not in _ignored_keys}
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        if self.scaler is not None:
+            self._accelerate_step_called = False
+            self._optimizer_original_step_method = self.optimizer.step
+            self._optimizer_patched_step_method = patch_optimizer_step(self, self.optimizer.step)
+
+
+def patch_optimizer_step(accelerated_optimizer: AcceleratedOptimizer, method):
+    def patched_step(*args, **kwargs):
+        accelerated_optimizer._accelerate_step_called = True
+        return method(*args, **kwargs)
+
+    return patched_step

.venv/Lib/site-packages/accelerate/scheduler.py

@@ -0,0 +1,98 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation
+
+import warnings
+
+from .state import AcceleratorState, GradientState
+
+
+warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")
+
+
+class AcceleratedScheduler:
+    """
+    A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
+    to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed
+    precision training)
+
+    When performing gradient accumulation scheduler lengths should not be changed accordingly, Accelerate will always
+    step the scheduler to account for it.
+
+    Args:
+        scheduler (`torch.optim.lr_scheduler._LRScheduler`):
+            The scheduler to wrap.
+        optimizers (one or a list of `torch.optim.Optimizer`):
+            The optimizers used.
+        step_with_optimizer (`bool`, *optional*, defaults to `True`):
+            Whether or not the scheduler should be stepped at each optimizer step.
+        split_batches (`bool`, *optional*, defaults to `False`):
+            Whether or not the dataloaders split one batch across the different processes (so batch size is the same
+            regardless of the number of processes) or create batches on each process (so batch size is the original
+            batch size multiplied by the number of processes).
+    """
+
+    def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, split_batches: bool = False):
+        self.scheduler = scheduler
+        self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers]
+        self.split_batches = split_batches
+        self.step_with_optimizer = step_with_optimizer
+        self.gradient_state = GradientState()
+
+    def step(self, *args, **kwargs):
+        if not self.step_with_optimizer:
+            # No link between scheduler and optimizer -> just step
+            self.scheduler.step(*args, **kwargs)
+            return
+
+        # Otherwise, first make sure the optimizer was stepped.
+        if not self.gradient_state.sync_gradients:
+            if self.gradient_state.adjust_scheduler:
+                self.scheduler._step_count += 1
+            return
+
+        for opt in self.optimizers:
+            if opt.step_was_skipped:
+                return
+        if self.split_batches:
+            # Split batches -> the training dataloader batch size is not changed so one step per training step
+            self.scheduler.step(*args, **kwargs)
+        else:
+            # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
+            # num_processes steps per training step
+            num_processes = AcceleratorState().num_processes
+            for _ in range(num_processes):
+                # Special case when using OneCycle and `drop_last` was not used
+                if hasattr(self.scheduler, "total_steps"):
+                    if self.scheduler._step_count <= self.scheduler.total_steps:
+                        self.scheduler.step(*args, **kwargs)
+                else:
+                    self.scheduler.step(*args, **kwargs)
+
+    # Passthroughs
+    def get_last_lr(self):
+        return self.scheduler.get_last_lr()
+
+    def state_dict(self):
+        return self.scheduler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.scheduler.load_state_dict(state_dict)
+
+    def get_lr(self):
+        return self.scheduler.get_lr()
+
+    def print_lr(self, *args, **kwargs):
+        return self.scheduler.print_lr(*args, **kwargs)

.venv/Lib/site-packages/accelerate/state.py

@@ -0,0 +1,1257 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+import warnings
+from contextlib import contextmanager
+from functools import partial
+from typing import Any, Callable, Optional
+
+import torch
+
+from .utils import (
+    DistributedType,
+    DynamoBackend,
+    GradientAccumulationPlugin,
+    check_cuda_p2p_ib_support,
+    check_fp8_capability,
+    deepspeed_required,
+    get_ccl_version,
+    get_cpu_distributed_information,
+    get_int_from_env,
+    is_ccl_available,
+    is_datasets_available,
+    is_deepspeed_available,
+    is_fp8_available,
+    is_ipex_available,
+    is_mlu_available,
+    is_mps_available,
+    is_musa_available,
+    is_npu_available,
+    is_torch_xla_available,
+    is_xpu_available,
+    parse_choice_from_env,
+    parse_flag_from_env,
+    set_numa_affinity,
+)
+from .utils.dataclasses import SageMakerDistributedType
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+if is_mlu_available(check_device=False):
+    import torch_mlu  # noqa: F401
+
+if is_musa_available(check_device=False):
+    import torch_musa  # noqa: F401
+
+if is_npu_available(check_device=False):
+    import torch_npu  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+def is_initialized() -> bool:
+    """
+    Checks if the `AcceleratorState` has been initialized from `Accelerator`. Same as `AcceleratorState.initialized`,
+    but works as a module method.
+    """
+    return AcceleratorState._shared_state != {}
+
+
+# Lambda function that does nothing
+def do_nothing(*args, **kwargs):
+    return None
+
+
+class ThreadLocalSharedDict(threading.local):
+    """
+    Descriptor that holds a dict shared between instances of a class in the same thread.
+
+    Note: Descriptors have slightly different semantics than just a dict field on its own.
+    `PartialState(...)._shared_state` and `PartialState._shared_state` (instance vs class) give the same value: the
+    underlying _storage dict. Likewise, `PartialState(...)._shared_state = {...}` overrides the _storage dict inside
+    the descriptor as you would expect. However, `PartialState._shared_state = {}` actually replaces the descriptor
+    object with a dict instead Thus, you should modify the _storage dict in-place (e.g. `_shared_state.clear()`).
+
+    See Python documentation for an explanation of descriptors: https://docs.python.org/3/howto/descriptor.html
+
+    This is required for using PyTorch/XLA with PJRT in multithreaded mode (required for TPU v2 and v3).
+
+    See https://github.com/pytorch/xla/blob/r2.0/docs/pjrt.md#multithreading-on-tpu-v2v3
+    """
+
+    def __init__(self, thread_local: bool = False):
+        self._storage = {}
+
+    def __get__(self, obj, objtype=None):
+        return self._storage
+
+    def __set__(self, obj, value):
+        self._storage = value
+
+
+# Prefer global shared dictionary, except when using TPU.
+SharedDict = dict if not is_torch_xla_available() else ThreadLocalSharedDict
+
+
+# Inspired by Alex Martelli's 'Borg'.
+class PartialState:
+    """
+    Singleton class that has information about the current training environment and functions to help with process
+    control. Designed to be used when only process control and device execution states are needed. Does *not* need to
+    be initialized from `Accelerator`.
+
+    Args:
+        cpu (`bool`, *optional*):
+            Whether or not to force the script to execute on CPU. Will ignore any accelerators available if set to
+            `True` and force the execution on the CPU.
+        kwargs (additional keyword arguments, *optional*):
+            Additional keyword arguments to pass to the relevent `init_process_group` function. Valid `kwargs` can be
+            found in [`utils.InitProcessGroupKwargs`]. See the example section for detailed usage.
+
+    **Available attributes:**
+
+        - **device** (`torch.device`) -- The device to use.
+        - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently
+          in use.
+        - **local_process_index** (`int`) -- The index of the current process on the current server.
+        - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type
+          of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8').
+        - **num_processes** (`int`) -- The number of processes currently launched in parallel.
+        - **process_index** (`int`) -- The index of the current process.
+        - **is_last_process** (`bool`) -- Whether or not the current process is the last one.
+        - **is_main_process** (`bool`) -- Whether or not the current process is the main one.
+        - **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node.
+        - **debug** (`bool`) -- Whether or not the current script is being run in debug mode.
+
+    Example:
+    ```python
+    from accelerate.utils import InitProcessGroupKwargs
+
+    # To include `InitProcessGroupKwargs`, init then call `.to_kwargs()`
+    kwargs = InitProcessGroupKwargs(...).to_kwargs()
+    state = PartialState(**kwargs)
+    ```
+    """
+
+    _shared_state = SharedDict()
+    _known_attrs = [
+        "_cpu",
+        "_mixed_precision",
+        "_shared_state",
+        "backend",
+        "debug",
+        "device",
+        "distributed_type",
+        "fork_launched",
+        "local_process_index",
+        "num_processes",
+        "process_index",
+    ]
+
+    def __init__(self, cpu: bool = False, **kwargs):
+        self.__dict__ = self._shared_state
+        if not self.initialized:
+            self._cpu = cpu
+            self.backend = None
+            env_device = os.environ.get("ACCELERATE_TORCH_DEVICE", None)
+            self.device = torch.device(env_device) if env_device is not None else None
+            self.debug = parse_flag_from_env("ACCELERATE_DEBUG_MODE")
+            use_sagemaker_dp = kwargs.pop("_use_sagemaker_dp", None)
+            dist_information = None
+            if use_sagemaker_dp is None:
+                use_sagemaker_dp = (
+                    os.environ.get("ACCELERATE_USE_SAGEMAKER", "false") == "true"
+                    and os.environ.get("ACCELERATE_SAGEMAKER_DISTRIBUTED_TYPE") != SageMakerDistributedType.NO
+                )
+
+            # Sets up self.backend + imports
+            original_backend = kwargs.pop("backend", None)
+            backend, distributed_type = self._prepare_backend(cpu, use_sagemaker_dp, original_backend)
+            if original_backend is not None and backend != original_backend:
+                raise ValueError(f"Your assigned backend {original_backend} is not avaliable, please use {backend}")
+            self.backend = backend
+            self.distributed_type = distributed_type
+            use_deepspeed = False
+            if not cpu and self.backend != "xla":
+                if int(os.environ.get("LOCAL_RANK", -1)) != -1:
+                    # Deal with spawning deepspeed
+                    if os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true":
+                        if not is_deepspeed_available():
+                            raise ImportError(
+                                "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
+                            )
+                        from deepspeed import comm as dist
+
+                        if not dist.is_initialized():
+                            dist.init_distributed(dist_backend=self.backend, auto_mpi_discovery=False, **kwargs)
+                        # We need to flag to `use_deepspeed` to be True to override `distributed_type` later
+                        use_deepspeed = True
+                    # Deal with all other backends but XPU and CPU, that gets handled special later
+                    elif (
+                        self.distributed_type not in (DistributedType.MULTI_XPU, DistributedType.MULTI_CPU)
+                        and not torch.distributed.is_initialized()
+                    ):
+                        torch.distributed.init_process_group(backend=self.backend, **kwargs)
+            # XPU and CPU require special env configs to be set
+            if self.distributed_type in (DistributedType.MULTI_XPU, DistributedType.MULTI_CPU):
+                dist_information = get_cpu_distributed_information()
+                os.environ["RANK"] = str(dist_information.rank)
+                os.environ["WORLD_SIZE"] = str(dist_information.world_size)
+                os.environ["LOCAL_RANK"] = str(dist_information.local_rank)
+                os.environ["LOCAL_WORLD_SIZE"] = str(dist_information.local_world_size)
+                if not os.environ.get("MASTER_PORT", None):
+                    os.environ["MASTER_PORT"] = "29500"
+                if (
+                    not os.environ.get("MASTER_ADDR", None)
+                    and dist_information.local_world_size != dist_information.world_size
+                    and self.backend != "mpi"
+                ):
+                    raise ValueError(
+                        "Tried to launch on distributed with multinode, but `MASTER_ADDR` env was not set, "
+                        "please try exporting rank 0's hostname as `MASTER_ADDR`"
+                    )
+                kwargs["rank"] = dist_information.rank
+                kwargs["world_size"] = dist_information.world_size
+
+                if (
+                    self.distributed_type == DistributedType.MULTI_CPU
+                    and get_int_from_env(["OMP_NUM_THREADS"], 0) == 0
+                ):
+                    import psutil
+
+                    num_cpu_threads_per_process = int(
+                        psutil.cpu_count(logical=False) / dist_information.local_world_size
+                    )
+                    if num_cpu_threads_per_process == 0:
+                        num_cpu_threads_per_process = 1
+                    torch.set_num_threads(num_cpu_threads_per_process)
+                    warnings.warn(
+                        f"OMP_NUM_THREADS/MKL_NUM_THREADS unset, we set it at {num_cpu_threads_per_process} to improve oob"
+                        " performance."
+                    )
+
+                if not torch.distributed.is_initialized():
+                    torch.distributed.init_process_group(backend=self.backend, **kwargs)
+
+            # No backend == no distributed training
+            if self.backend is None:
+                self.distributed_type = DistributedType.NO
+                self.num_processes = 1
+                self.process_index = 0
+                self.local_process_index = 0
+            elif self.backend == "xla":
+                # XLA needs device setting first for `set_replication`
+                self.set_device()
+                xm.set_replication(self.device, xm.get_xla_supported_devices())
+                self.num_processes = xm.xrt_world_size()
+                self.process_index = xm.get_ordinal()
+                if is_torch_xla_available(check_is_tpu=True):
+                    self.local_process_index = xm.get_local_ordinal()
+                else:
+                    self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))
+            else:
+                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)) if dist_information is None else dist_information.local_rank
+                )
+            self.set_device()
+            # Now we can change to deepseed
+            if use_deepspeed:
+                self.distributed_type = DistributedType.DEEPSPEED
+
+            # Set CPU affinity if enabled
+            if parse_flag_from_env("ACCELERATE_CPU_AFFINITY", False):
+                set_numa_affinity(self.local_process_index)
+
+            # Check for old RTX 4000's that can't use P2P or IB and are on old drivers
+            if self.device.type == "cuda" and not check_cuda_p2p_ib_support():
+                if "NCCL_P2P_DISABLE" not in os.environ or "NCCL_IB_DISABLE" not in os.environ:
+                    raise NotImplementedError(
+                        "Using RTX 4000 series doesn't support faster communication broadband via P2P or IB. "
+                        'Please set `NCCL_P2P_DISABLE="1"` and `NCCL_IB_DISABLE="1" or use `accelerate launch` which '
+                        "will do this automatically."
+                    )
+        # Important: This should be the *only* code outside of `self.initialized!`
+        self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
+
+    def __repr__(self) -> str:
+        return (
+            f"Distributed environment: {self.distributed_type}{('  Backend: ' + self.backend) if self.backend else ''}\n"
+            f"Num processes: {self.num_processes}\n"
+            f"Process index: {self.process_index}\n"
+            f"Local process index: {self.local_process_index}\n"
+            f"Device: {self.device}\n"
+        )
+
+    @staticmethod
+    def _reset_state():
+        "Resets `_shared_state`, is used internally and should not be called"
+        PartialState._shared_state.clear()
+
+    @property
+    def initialized(self) -> bool:
+        "Returns whether the `PartialState` has been initialized"
+        return self._shared_state != {}
+
+    @property
+    def use_distributed(self):
+        """
+        Whether the Accelerator is configured for distributed training
+        """
+        return self.distributed_type != DistributedType.NO and self.num_processes > 1
+
+    @property
+    def is_last_process(self) -> bool:
+        "Returns whether the current process is the last one"
+        return self.process_index == self.num_processes - 1
+
+    @property
+    def is_main_process(self) -> bool:
+        "Returns whether the current process is the main process"
+        return (
+            self.process_index == 0 if self.distributed_type != DistributedType.MEGATRON_LM else self.is_last_process
+        )
+
+    @property
+    def is_local_main_process(self) -> bool:
+        "Returns whether the current process is the main process on the local node"
+        return (
+            self.local_process_index == 0
+            if self.distributed_type != DistributedType.MEGATRON_LM
+            else self.is_last_process
+        )
+
+    def wait_for_everyone(self):
+        """
+        Will stop the execution of the current process until every other process has reached that point (so this does
+        nothing when the script is only run in one process). Useful to do before saving a model.
+
+        Example:
+
+        ```python
+        >>> # Assuming two GPU processes
+        >>> import time
+        >>> from accelerate.state import PartialState
+
+        >>> state = PartialState()
+        >>> if state.is_main_process:
+        ...     time.sleep(2)
+        >>> else:
+        ...     print("I'm waiting for the main process to finish its sleep...")
+        >>> state.wait_for_everyone()
+        >>> # Should print on every process at the same time
+        >>> print("Everyone is here")
+        ```
+        """
+        if self.distributed_type in (
+            DistributedType.MULTI_GPU,
+            DistributedType.MULTI_MLU,
+            DistributedType.MULTI_MUSA,
+            DistributedType.MULTI_NPU,
+            DistributedType.MULTI_XPU,
+            DistributedType.MULTI_CPU,
+            DistributedType.DEEPSPEED,
+            DistributedType.FSDP,
+        ):
+            torch.distributed.barrier()
+        elif self.distributed_type == DistributedType.XLA:
+            xm.rendezvous("accelerate.utils.wait_for_everyone")
+
+    def _goes_first(self, is_main: bool):
+        if not is_main:
+            self.wait_for_everyone()
+
+        yield
+
+        if is_main:
+            self.wait_for_everyone()
+
+    @contextmanager
+    def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False):
+        """
+        Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing
+        distributed inference, such as with different prompts.
+
+        Note that when using a `dict`, all keys need to have the same number of elements.
+
+        Args:
+            inputs (`list`, `tuple`, `torch.Tensor`, `dict` of `list`/`tuple`/`torch.Tensor`, or `datasets.Dataset`):
+                The input to split between processes.
+            apply_padding (`bool`, `optional`, defaults to `False`):
+                Whether to apply padding by repeating the last element of the input so that all processes have the same
+                number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing
+                in less inputs than there are processes. If so, just remember to drop the padded elements afterwards.
+
+
+        Example:
+
+        ```python
+        # Assume there are two processes
+        from accelerate import PartialState
+
+        state = PartialState()
+        with state.split_between_processes(["A", "B", "C"]) as inputs:
+            print(inputs)
+        # Process 0
+        ["A", "B"]
+        # Process 1
+        ["C"]
+
+        with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs:
+            print(inputs)
+        # Process 0
+        ["A", "B"]
+        # Process 1
+        ["C", "C"]
+        ```
+        """
+        if self.num_processes == 1:
+            yield inputs
+            return
+        length = len(inputs)
+        # Nested dictionary of any types
+        if isinstance(inputs, dict):
+            length = len(inputs[list(inputs.keys())[0]])
+            if not all(len(v) == length for v in inputs.values()):
+                raise ValueError("All values in the dictionary must have the same length")
+        num_samples_per_process, num_extras = divmod(length, self.num_processes)
+        start_index = self.process_index * num_samples_per_process + min(self.process_index, num_extras)
+        end_index = start_index + num_samples_per_process + (1 if self.process_index < num_extras else 0)
+
+        def _split_values(inputs, start_index, end_index):
+            if isinstance(inputs, (list, tuple, torch.Tensor)):
+                if start_index >= len(inputs):
+                    result = inputs[-1:]
+                else:
+                    result = inputs[start_index:end_index]
+                if apply_padding:
+                    if isinstance(result, torch.Tensor):
+                        from accelerate.utils import pad_across_processes, send_to_device
+
+                        # The tensor needs to be on the device before we can pad it
+                        tensorized_result = send_to_device(result, self.device)
+                        result = pad_across_processes(tensorized_result, pad_index=inputs[-1])
+                    else:
+                        result += [result[-1]] * (num_samples_per_process + 1 - len(result))
+                return result
+            elif isinstance(inputs, dict):
+                for key in inputs.keys():
+                    inputs[key] = _split_values(inputs[key], start_index, end_index)
+                return inputs
+            else:
+                if is_datasets_available():
+                    from datasets import Dataset
+
+                    if isinstance(inputs, Dataset):
+                        if start_index >= len(inputs):
+                            start_index = len(inputs) - 1
+                        if end_index > len(inputs):
+                            end_index = len(inputs)
+                        result_idcs = list(range(start_index, end_index))
+                        if apply_padding:
+                            result_idcs += [end_index - 1] * (num_samples_per_process + 1 - len(result_idcs))
+                        return inputs.select(result_idcs)
+                return inputs
+
+        yield _split_values(inputs, start_index, end_index)
+
+    @contextmanager
+    def main_process_first(self):
+        """
+        Lets the main process go first inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+
+        >>> accelerator = Accelerator()
+        >>> with accelerator.main_process_first():
+        ...     # This will be printed first by process 0 then in a seemingly
+        ...     # random order by the other processes.
+        ...     print(f"This will be printed by process {accelerator.process_index}")
+        ```
+        """
+        yield from self._goes_first(self.is_main_process)
+
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Lets the local main process go inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+
+        Example:
+
+        ```python
+        >>> from accelerate.state import PartialState
+
+        >>> state = PartialState()
+        >>> with state.local_main_process_first():
+        ...     # This will be printed first by local process 0 then in a seemingly
+        ...     # random order by the other processes.
+        ...     print(f"This will be printed by process {state.local_process_index}")
+        ```
+        """
+        yield from self._goes_first(self.is_local_main_process)
+
+    def on_main_process(self, function: Callable[..., Any] = None):
+        """
+        Decorator that only runs the decorated function on the main process.
+
+        Args:
+            function (`Callable`): The function to decorate.
+
+        Example:
+
+        ```python
+        >>> from accelerate.state import PartialState
+
+        >>> state = PartialState()
+
+
+        >>> @state.on_main_process
+        ... def print_something():
+        ...     print("This will be printed by process 0 only.")
+
+
+        >>> print_something()
+        "This will be printed by process 0 only"
+        ```
+        """
+        if not self.initialized:
+            raise ValueError("The `PartialState` or `Accelerator` must be initialized before calling this function.")
+        if self.is_main_process or not self.use_distributed:
+            return function
+        return do_nothing
+
+    def on_local_main_process(self, function: Callable[..., Any] = None):
+        """
+        Decorator that only runs the decorated function on the local main process.
+
+        Args:
+            function (`Callable`): The function to decorate.
+
+        Example:
+        ```python
+        # Assume we have 2 servers with 4 processes each.
+        from accelerate.state import PartialState
+
+        state = PartialState()
+
+
+        @state.on_local_main_process
+        def print_something():
+            print("This will be printed by process 0 only on each server.")
+
+
+        print_something()
+        # On server 1:
+        "This will be printed by process 0 only"
+        # On server 2:
+        "This will be printed by process 0 only"
+        ```
+        """
+        if self.is_local_main_process or not self.use_distributed:
+            return function
+        return do_nothing
+
+    def on_last_process(self, function: Callable[..., Any]):
+        """
+        Decorator that only runs the decorated function on the last process.
+
+        Args:
+            function (`Callable`): The function to decorate.
+
+        Example:
+        ```python
+        # Assume we have 4 processes.
+        from accelerate.state import PartialState
+
+        state = PartialState()
+
+
+        @state.on_last_process
+        def print_something():
+            print(f"Printed on process {state.process_index}")
+
+
+        print_something()
+        "Printed on process 3"
+        ```
+        """
+        if self.is_last_process or not self.use_distributed:
+            return function
+        return do_nothing
+
+    def on_process(self, function: Callable[..., Any] = None, process_index: int = None):
+        """
+        Decorator that only runs the decorated function on the process with the given index.
+
+        Args:
+            function (`Callable`, `optional`):
+                The function to decorate.
+            process_index (`int`, `optional`):
+                The index of the process on which to run the function.
+
+        Example:
+        ```python
+        # Assume we have 4 processes.
+        from accelerate.state import PartialState
+
+        state = PartialState()
+
+
+        @state.on_process(process_index=2)
+        def print_something():
+            print(f"Printed on process {state.process_index}")
+
+
+        print_something()
+        "Printed on process 2"
+        ```
+        """
+        if function is None:
+            return partial(self.on_process, process_index=process_index)
+        if (self.process_index == process_index) or (not self.use_distributed):
+            return function
+        return do_nothing
+
+    def on_local_process(self, function: Callable[..., Any] = None, local_process_index: int = None):
+        """
+        Decorator that only runs the decorated function on the process with the given index on the current node.
+
+        Args:
+            function (`Callable`, *optional*):
+                The function to decorate.
+            local_process_index (`int`, *optional*):
+                The index of the local process on which to run the function.
+
+        Example:
+        ```python
+        # Assume we have 2 servers with 4 processes each.
+        from accelerate import Accelerator
+
+        accelerator = Accelerator()
+
+
+        @accelerator.on_local_process(local_process_index=2)
+        def print_something():
+            print(f"Printed on process {accelerator.local_process_index}")
+
+
+        print_something()
+        # On server 1:
+        "Printed on process 2"
+        # On server 2:
+        "Printed on process 2"
+        ```
+        """
+        if function is None:
+            return partial(self.on_local_process, local_process_index=local_process_index)
+        if (self.local_process_index == local_process_index) or (not self.use_distributed):
+            return function
+        return do_nothing
+
+    def print(self, *args, **kwargs):
+        if self.is_local_main_process:
+            print(*args, **kwargs)
+
+    @property
+    def default_device(self) -> torch.device:
+        """
+        Returns the default device which is:
+        - MPS if `torch.backends.mps.is_available()` and `torch.backends.mps.is_built()` both return True.
+        - CUDA if `torch.cuda.is_available()`
+        - MLU if `is_mlu_available()`
+        - MUSA if `is_musa_available()`
+        - NPU if `is_npu_available()`
+        - CPU otherwise
+        """
+        if is_mps_available():
+            os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+            return torch.device("mps")
+        elif is_mlu_available():
+            return torch.device("mlu")
+        elif is_musa_available():
+            return torch.device("musa")
+        # NPU should be checked before CUDA when using `transfer_to_npu`
+        # See issue #3020: https://github.com/huggingface/accelerate/issues/3020
+        elif is_npu_available():
+            return torch.device("npu")
+        elif torch.cuda.is_available():
+            return torch.device("cuda")
+        elif is_xpu_available():
+            return torch.device("xpu")
+        else:
+            return torch.device("cpu")
+
+    def _prepare_backend(
+        self, cpu: bool = False, sagemaker_dp=False, backend: str = None
+    ) -> tuple[str, DistributedType]:
+        "Prepares any imports needed before initializing the distributed backend and sets `self.backend` properly"
+        distributed_type = None
+        if sagemaker_dp:
+            import smdistributed.dataparallel.torch.torch_smddp  # noqa
+
+            backend = "smddp"
+            distributed_type = DistributedType.MULTI_GPU
+        elif is_torch_xla_available():
+            backend = "xla"
+            distributed_type = DistributedType.XLA
+        elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
+            if is_mlu_available():
+                backend = "cncl"
+                distributed_type = DistributedType.MULTI_MLU
+            elif is_musa_available():
+                backend = "mccl"
+                distributed_type = DistributedType.MULTI_MUSA
+            # NPU should be checked before CUDA when using `transfer_to_npu`
+            # See issue #3020: https://github.com/huggingface/accelerate/issues/3020
+            elif is_npu_available():
+                backend = "hccl"
+                distributed_type = DistributedType.MULTI_NPU
+            elif torch.cuda.is_available():
+                if backend is None:
+                    backend = "nccl"
+                distributed_type = DistributedType.MULTI_GPU
+
+        if distributed_type is None and (
+            int(os.environ.get("LOCAL_RANK", -1)) != -1
+            or get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1
+        ):
+            if not cpu and is_xpu_available():
+                distributed_type = DistributedType.MULTI_XPU
+            else:
+                distributed_type = DistributedType.MULTI_CPU
+
+            if (
+                backend in (None, "ccl")
+                and is_ccl_available()
+                and (get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0 or distributed_type == DistributedType.MULTI_XPU)
+            ):
+                if get_ccl_version() >= "1.12":
+                    import oneccl_bindings_for_pytorch  # noqa: F401
+                else:
+                    import torch_ccl  # noqa: F401
+
+                backend = "ccl"
+            elif backend in (None, "mpi") and torch.distributed.is_mpi_available():
+                backend = "mpi"
+            else:
+                backend = "gloo"
+        if distributed_type is None:
+            distributed_type = DistributedType.NO
+
+        return backend, distributed_type
+
+    def set_device(self):
+        """
+        Sets the device in `self.device` to the current distributed environment.
+        """
+        if self.device is not None:
+            return
+        if self.distributed_type == DistributedType.NO:
+            self.device = torch.device("cpu") if self._cpu else self.default_device
+            return
+        device = str(self.distributed_type).split(".")[-1].replace("MULTI_", "").lower()
+        if device not in ("cpu", "gpu", "mlu", "musa", "npu", "xpu", "xla"):
+            raise ValueError(
+                f"Can't set device for {self.distributed_type} ({device}), verify we should be calling `_set_device()` for it!"
+            )
+        if device == "xla":
+            self.device = xm.xla_device()
+        else:
+            if device == "gpu":
+                device = "cuda"
+            device_module = getattr(torch, device)
+            device_index = self.local_process_index % device_module.device_count()
+            self.device = torch.device(device, device_index)
+            device_module.set_device(self.device)
+
+    def destroy_process_group(self, group=None):
+        """
+        Destroys the process group. If one is not specified, the default process group is destroyed.
+        """
+        if self.fork_launched and group is None:
+            return
+        # needed when using torch.distributed.init_process_group
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group(group)
+
+    def __getattr__(self, name: str):
+        # By this point we know that no attributes of `self` contain `name`,
+        # so we just modify the error message
+        if name in self._known_attrs:
+            raise AttributeError(
+                f"`PartialState` object has no attribute `{name}`. "
+                "This happens if `PartialState._reset_state()` was called and "
+                "an `Accelerator` or `PartialState` was not reinitialized."
+            )
+        # Raise a typical AttributeError
+        raise AttributeError(f"'PartialState' object has no attribute '{name}'")
+
+
+class AcceleratorState:
+    """
+    Singleton class that has information about the current training environment.
+
+    **Available attributes:**
+
+        - **device** (`torch.device`) -- The device to use.
+        - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently
+          in use.
+        - **initialized** (`bool`) -- Whether or not the `AcceleratorState` has been initialized from `Accelerator`.
+        - **local_process_index** (`int`) -- The index of the current process on the current server.
+        - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type
+          of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8').
+        - **num_processes** (`int`) -- The number of processes currently launched in parallel.
+        - **process_index** (`int`) -- The index of the current process.
+        - **is_last_process** (`bool`) -- Whether or not the current process is the last one.
+        - **is_main_process** (`bool`) -- Whether or not the current process is the main one.
+        - **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node.
+        - **debug** (`bool`) -- Whether or not the current script is being run in debug mode.
+    """
+
+    _shared_state = SharedDict()
+    _known_attrs = PartialState._known_attrs + [
+        "deepspeed_plugin",
+        "use_ipex",
+        "fsdp_plugin",
+        "megatron_lm_plugin",
+        "dynamo_plugin",
+    ]
+
+    def __init__(
+        self,
+        mixed_precision: str = None,
+        cpu: bool = False,
+        dynamo_plugin=None,
+        deepspeed_plugin=None,
+        fsdp_plugin=None,
+        megatron_lm_plugin=None,
+        _from_accelerator: bool = False,
+        **kwargs,
+    ):
+        self.__dict__ = self._shared_state
+        if parse_flag_from_env("ACCELERATE_USE_CPU"):
+            cpu = True
+        if PartialState._shared_state == {}:
+            PartialState(cpu, **kwargs)
+        self.__dict__.update(PartialState._shared_state)
+        self._check_initialized(mixed_precision, cpu)
+        if not self.initialized:
+            self.deepspeed_plugins = None
+            self.use_ipex = None
+            mixed_precision = (
+                parse_choice_from_env("ACCELERATE_MIXED_PRECISION", "no")
+                if mixed_precision is None
+                else mixed_precision.lower()
+            )
+            if mixed_precision == "fp8":
+                if not is_fp8_available():
+                    raise ValueError(
+                        "Using `fp8` precision requires `transformer_engine` or `MS-AMP` to be installed."
+                    )
+                elif not check_fp8_capability():
+                    logger.warning(
+                        f"The current device has compute capability of {torch.cuda.get_device_capability()} which is "
+                        "insufficient for FP8 mixed precision training (requires a GPU Hopper/Ada Lovelace "
+                        "or higher, compute capability of 8.9 or higher). Will use FP16 instead."
+                    )
+                    mixed_precision = "fp16"
+
+            self.dynamo_plugin = dynamo_plugin
+            if not _from_accelerator:
+                raise ValueError(
+                    "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
+                    "before using any functionality from the `accelerate` library."
+                )
+            # deepspeed handles mixed_precision using deepspeed_config
+            self._mixed_precision = "no" if self.distributed_type == DistributedType.DEEPSPEED else mixed_precision
+            if self.distributed_type == DistributedType.XLA and is_torch_xla_available(check_is_tpu=True):
+                if mixed_precision == "bf16":
+                    if os.environ.get("ACCELERATE_DOWNCAST_BF16"):
+                        os.environ["XLA_USE_BF16"] = str(0)
+                        os.environ["XLA_DOWNCAST_BF16"] = str(1)
+                        self.downcast_bfloat = True
+                    else:
+                        os.environ["XLA_USE_BF16"] = str(1)
+                        os.environ["XLA_DOWNCAST_BF16"] = str(0)
+                        self.downcast_bfloat = False
+            elif os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true" and not cpu:
+                self.deepspeed_plugins = deepspeed_plugin
+                self.distributed_type = DistributedType.DEEPSPEED
+            elif self.distributed_type in [
+                DistributedType.MULTI_GPU,
+                DistributedType.MULTI_MLU,
+                DistributedType.MULTI_MUSA,
+                DistributedType.MULTI_NPU,
+                DistributedType.MULTI_XPU,
+            ]:
+                if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" or fsdp_plugin is not None:
+                    self.distributed_type = DistributedType.FSDP
+                    if self._mixed_precision != "no":
+                        fsdp_plugin.set_mixed_precision(self._mixed_precision)
+                    self.fsdp_plugin = fsdp_plugin
+                if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false") == "true" and self.distributed_type not in [
+                    DistributedType.MULTI_XPU,
+                ]:
+                    self.distributed_type = DistributedType.MEGATRON_LM
+                    megatron_lm_plugin.set_mixed_precision(self._mixed_precision)
+                    self.megatron_lm_plugin = megatron_lm_plugin
+            elif self.distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_XPU, DistributedType.NO]:
+                if is_ipex_available():
+                    # check if user disables it explicitly
+                    self.use_ipex = parse_flag_from_env("ACCELERATE_USE_IPEX", default=True)
+                else:
+                    self.use_ipex = False
+            if (
+                self.dynamo_plugin.backend != DynamoBackend.NO
+                and self._mixed_precision == "no"
+                and self.device.type == "cuda"
+            ):
+                torch.backends.cuda.matmul.allow_tf32 = True
+            if (
+                self.dynamo_plugin.backend != DynamoBackend.NO
+                and self._mixed_precision == "no"
+                and self.device.type == "musa"
+            ):
+                torch.backends.musa.matmul.allow_tf32 = True
+            PartialState._shared_state["distributed_type"] = self.distributed_type
+
+    @property
+    def initialized(self) -> bool:
+        return self._shared_state != PartialState._shared_state
+
+    def __repr__(self):
+        repr = PartialState().__repr__() + f"\nMixed precision type: {self.mixed_precision}\n"
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n"
+        return repr
+
+    def _check_initialized(self, mixed_precision=None, cpu=None):
+        "Checks if a modification is trying to be made and the `AcceleratorState` has already been initialized"
+        if self.initialized:
+            err = "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `{flag}` to `Accelerator()`."
+            if cpu and self.device.type != "cpu":
+                raise ValueError(err.format(flag="cpu=True"))
+            if (
+                mixed_precision is not None
+                and mixed_precision != self._mixed_precision
+                and self.distributed_type != DistributedType.DEEPSPEED
+            ):
+                raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'"))
+
+    @property
+    def mixed_precision(self):
+        if self.distributed_type == DistributedType.DEEPSPEED:
+            config = self.deepspeed_plugin.deepspeed_config
+            if config.get("fp16", {}).get("enabled", False):
+                mixed_precision = "fp16"
+            elif config.get("bf16", {}).get("enabled", False):
+                mixed_precision = "bf16"
+            else:
+                mixed_precision = "no"
+        else:
+            mixed_precision = self._mixed_precision
+        return mixed_precision
+
+    @staticmethod
+    def _reset_state(reset_partial_state: bool = False):
+        "Resets `_shared_state`, is used internally and should not be called"
+        AcceleratorState._shared_state.clear()
+        if reset_partial_state:
+            PartialState._reset_state()
+
+    def destroy_process_group(self, group=None):
+        """
+        Destroys the process group. If one is not specified, the default process group is destroyed.
+
+        If `self.fork_lauched` is `True` and `group` is `None`, nothing happens.
+        """
+        PartialState().destroy_process_group(group)
+
+    @property
+    def fork_launched(self):
+        return PartialState().fork_launched
+
+    @property
+    def use_distributed(self):
+        """
+        Whether the Accelerator is configured for distributed training
+        """
+        return PartialState().use_distributed
+
+    @property
+    def is_last_process(self) -> bool:
+        "Returns whether the current process is the last one"
+        return PartialState().is_last_process
+
+    @property
+    def is_main_process(self) -> bool:
+        "Returns whether the current process is the main process"
+        return PartialState().is_main_process
+
+    @property
+    def is_local_main_process(self) -> bool:
+        "Returns whether the current process is the main process on the local node"
+        return PartialState().is_local_main_process
+
+    def wait_for_everyone(self):
+        PartialState().wait_for_everyone()
+
+    @contextmanager
+    def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False):
+        """
+        Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing
+        distributed inference, such as with different prompts.
+
+        Note that when using a `dict`, all keys need to have the same number of elements.
+
+        Args:
+            inputs (`list`, `tuple`, `torch.Tensor`, or `dict` of `list`/`tuple`/`torch.Tensor`):
+                The input to split between processes.
+            apply_padding (`bool`, `optional`, defaults to `False`):
+                Whether to apply padding by repeating the last element of the input so that all processes have the same
+                number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing
+                in less inputs than there are processes. If so, just remember to drop the padded elements afterwards.
+
+
+        Example:
+
+        ```python
+        # Assume there are two processes
+        from accelerate.state import AcceleratorState
+
+        state = AcceleratorState()
+        with state.split_between_processes(["A", "B", "C"]) as inputs:
+            print(inputs)
+        # Process 0
+        ["A", "B"]
+        # Process 1
+        ["C"]
+
+        with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs:
+            print(inputs)
+        # Process 0
+        ["A", "B"]
+        # Process 1
+        ["C", "C"]
+        ```
+        """
+        with PartialState().split_between_processes(inputs, apply_padding=apply_padding) as inputs:
+            yield inputs
+
+    @contextmanager
+    def main_process_first(self):
+        """
+        Lets the main process go first inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        with PartialState().main_process_first():
+            yield
+
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Lets the local main process go inside a with block.
+
+        The other processes will enter the with block after the main process exits.
+        """
+        with PartialState().local_main_process_first():
+            yield
+
+    @property
+    def deepspeed_plugin(self):
+        """
+        Returns the currently active DeepSpeedPlugin.
+
+        If not using deepspeed, returns `None`.
+        """
+        # To maintain original behavior, return None if not using deepspeed.
+        if self.distributed_type != DistributedType.DEEPSPEED:
+            return None
+        from accelerate.utils.deepspeed import get_active_deepspeed_plugin
+
+        return get_active_deepspeed_plugin(self)
+
+    @deepspeed_required
+    def get_deepspeed_plugin(self, name: str):
+        """
+        Returns the DeepSpeedPlugin with the given plugin_key.
+        """
+        return self.deepspeed_plugins[name]
+
+    @deepspeed_required
+    def select_deepspeed_plugin(self, name: str = None):
+        """
+        Activates the DeepSpeedPlugin with the given `name`, and will disable all other plugins.
+        """
+        for key, plugin in self.deepspeed_plugins.items():
+            if key != name:
+                plugin._unselect()
+        self.deepspeed_plugins[name].select(_from_accelerator_state=True)
+
+    def print(self, *args, **kwargs):
+        PartialState().print(*args, **kwargs)
+
+    def __getattr__(self, name: str):
+        # By this point we know that no attributes of `self` contain `name`,
+        # so we just modify the error message
+        if name in self._known_attrs:
+            raise AttributeError(
+                f"`AcceleratorState` object has no attribute `{name}`. "
+                "This happens if `AcceleratorState._reset_state()` was called and "
+                "an `Accelerator` or `PartialState` was not reinitialized."
+            )
+        # Raise a typical AttributeError
+        raise AttributeError(f"'AcceleratorState' object has no attribute '{name}'")
+
+
+class GradientState:
+    """
+    Singleton class that has information related to gradient synchronization for gradient accumulation
+
+    **Available attributes:**
+
+        - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader
+        - **remainder** (`int`) -- The number of extra samples that were added from padding the dataloader
+        - **sync_gradients** (`bool`) -- Whether the gradients should be synced across all devices
+        - **active_dataloader** (`Optional[DataLoader]`) -- The dataloader that is currently being iterated over
+        - **dataloader_references** (`List[Optional[DataLoader]]`) -- A list of references to the dataloaders that are
+            being iterated over
+        - **num_steps** (`int`) -- The number of steps to accumulate over
+        - **adjust_scheduler** (`bool`) -- Whether the scheduler should be adjusted to account for the gradient
+            accumulation
+        - **sync_with_dataloader** (`bool`) -- Whether the gradients should be synced at the end of the dataloader
+            iteration and the number of total steps reset
+        - **is_xla_gradients_synced** (`bool`) -- Whether the XLA gradients have been synchronized. It is initialized
+          as false. Once gradients have been reduced before the optimizer step, this flag is set to true. Subsequently,
+            after each step, the flag is reset to false. FSDP will always synchronize the gradients, hence
+            is_xla_gradients_synced is always true.
+    """
+
+    _shared_state = SharedDict()
+
+    def __init__(self, gradient_accumulation_plugin: Optional[GradientAccumulationPlugin] = None):
+        self.__dict__ = self._shared_state
+        if not self.initialized:
+            self.sync_gradients = True
+            self.active_dataloader = None
+            self.dataloader_references = [None]
+            self.plugin_kwargs = (
+                gradient_accumulation_plugin.to_kwargs() if gradient_accumulation_plugin is not None else {}
+            )
+            self._is_xla_gradients_synced = False
+
+        # Plugin args are different and can be updated
+        if gradient_accumulation_plugin is not None and self.plugin_kwargs != gradient_accumulation_plugin.to_kwargs():
+            self.plugin_kwargs = gradient_accumulation_plugin.to_kwargs()
+
+    @property
+    def num_steps(self) -> int:
+        "Returns the number of steps to accumulate over"
+        return self.plugin_kwargs.get("num_steps", 1)
+
+    @property
+    def adjust_scheduler(self) -> bool:
+        "Returns whether the scheduler should be adjusted"
+        return self.plugin_kwargs.get("adjust_scheduler", False)
+
+    @property
+    def sync_with_dataloader(self) -> bool:
+        "Returns whether the gradients should be synced at the end of the dataloader iteration and the number of total steps reset"
+        return self.plugin_kwargs.get("sync_with_dataloader", True)
+
+    @property
+    def initialized(self) -> bool:
+        "Returns whether the `GradientState` has been initialized"
+        return GradientState._shared_state != {}
+
+    @property
+    def end_of_dataloader(self) -> bool:
+        "Returns whether we have reached the end of the current dataloader"
+        if not self.in_dataloader:
+            return False
+        return self.active_dataloader.end_of_dataloader
+
+    @property
+    def remainder(self) -> int:
+        "Returns the number of extra samples that were added from padding the dataloader"
+        if not self.in_dataloader:
+            return -1
+        return self.active_dataloader.remainder
+
+    def __repr__(self):
+        return (
+            f"Sync Gradients: {self.sync_gradients}\n"
+            f"At end of current dataloader: {self.end_of_dataloader}\n"
+            f"Extra samples added: {self.remainder}\n"
+            f"Gradient accumulation plugin: {self.plugin_kwargs}\n"
+        )
+
+    @property
+    def is_xla_gradients_synced(self):
+        "Returns the value of is_xla_gradients_synced. FSDP will always synchronize the gradients, hence is_xla_gradients_synced is always true."
+        if parse_flag_from_env("ACCELERATE_USE_FSDP", default=False):
+            return True
+        return self._is_xla_gradients_synced
+
+    @is_xla_gradients_synced.setter
+    def is_xla_gradients_synced(self, is_synced):
+        "Set the _is_xla_gradients_synced attribute."
+        self._is_xla_gradients_synced = is_synced
+
+    def _set_sync_gradients(self, sync_gradients):
+        "Private function that sets whether gradients should be synchronized. Users should not have to call this."
+        self.sync_gradients = sync_gradients
+        # Allow grad-sync to automatically work on TPUs
+        if (
+            self.sync_gradients
+            and is_torch_xla_available(check_is_tpu=True)
+            and PartialState().distributed_type == DistributedType.XLA
+        ):
+            xm.mark_step()
+
+    def _add_dataloader(self, dataloader):
+        "Private function that adds a dataloader to `self.dataloader_references` and sets `in_dataloader` to `True`. Users should not have to call this."
+        self.active_dataloader = dataloader
+        self.dataloader_references.append(self.active_dataloader)
+
+    def _remove_dataloader(self, dataloader):
+        "Private function that removes a dataloader from `self.dataloader_references` and sets `in_dataloader` to `False` if there are no more dataloaders. Users should not have to call this."
+        self.dataloader_references.remove(dataloader)
+        self.active_dataloader = self.dataloader_references[-1]
+
+    @property
+    def in_dataloader(self) -> bool:
+        "Returns whether the current process is in a dataloader"
+        return self.active_dataloader is not None
+
+    @staticmethod
+    def _reset_state():
+        "Resets `_shared_state`, is used internally and should not be called"
+        GradientState._shared_state.clear()

.venv/Lib/site-packages/accelerate/tracking.py

@@ -0,0 +1,1023 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import json
+import os
+import time
+from functools import wraps
+from typing import Any, Dict, List, Optional, Union
+
+import yaml
+
+from .logging import get_logger
+from .state import PartialState
+from .utils import (
+    LoggerType,
+    is_aim_available,
+    is_clearml_available,
+    is_comet_ml_available,
+    is_dvclive_available,
+    is_mlflow_available,
+    is_tensorboard_available,
+    is_wandb_available,
+    listify,
+)
+
+
+_available_trackers = []
+
+if is_tensorboard_available():
+    _available_trackers.append(LoggerType.TENSORBOARD)
+
+if is_wandb_available():
+    _available_trackers.append(LoggerType.WANDB)
+
+if is_comet_ml_available():
+    _available_trackers.append(LoggerType.COMETML)
+
+if is_aim_available():
+    _available_trackers.append(LoggerType.AIM)
+
+if is_mlflow_available():
+    _available_trackers.append(LoggerType.MLFLOW)
+
+if is_clearml_available():
+    _available_trackers.append(LoggerType.CLEARML)
+
+if is_dvclive_available():
+    _available_trackers.append(LoggerType.DVCLIVE)
+
+logger = get_logger(__name__)
+
+
+def on_main_process(function):
+    """
+    Decorator to selectively run the decorated function on the main process only based on the `main_process_only`
+    attribute in a class.
+
+    Checks at function execution rather than initialization time, not triggering the initialization of the
+    `PartialState`.
+    """
+
+    @wraps(function)
+    def execute_on_main_process(self, *args, **kwargs):
+        if getattr(self, "main_process_only", False):
+            return PartialState().on_main_process(function)(self, *args, **kwargs)
+        else:
+            return function(self, *args, **kwargs)
+
+    return execute_on_main_process
+
+
+def get_available_trackers():
+    "Returns a list of all supported available trackers in the system"
+    return _available_trackers
+
+
+class GeneralTracker:
+    """
+    A base Tracker class to be used for all logging integration implementations.
+
+    Each function should take in `**kwargs` that will automatically be passed in from a base dictionary provided to
+    [`Accelerator`].
+
+    Should implement `name`, `requires_logging_directory`, and `tracker` properties such that:
+
+    `name` (`str`): String representation of the tracker class name, such as "TensorBoard" `requires_logging_directory`
+    (`bool`): Whether the logger requires a directory to store their logs. `tracker` (`object`): Should return internal
+    tracking mechanism used by a tracker class (such as the `run` for wandb)
+
+    Implementations can also include a `main_process_only` (`bool`) attribute to toggle if relevent logging, init, and
+    other functions should occur on the main process or across all processes (by default will use `True`)
+    """
+
+    main_process_only = True
+
+    def __init__(self, _blank=False):
+        if not _blank:
+            err = ""
+            if not hasattr(self, "name"):
+                err += "`name`"
+            if not hasattr(self, "requires_logging_directory"):
+                if len(err) > 0:
+                    err += ", "
+                err += "`requires_logging_directory`"
+
+            # as tracker is a @property that relies on post-init
+            if "tracker" not in dir(self):
+                if len(err) > 0:
+                    err += ", "
+                err += "`tracker`"
+            if len(err) > 0:
+                raise NotImplementedError(
+                    f"The implementation for this tracker class is missing the following "
+                    f"required attributes. Please define them in the class definition: "
+                    f"{err}"
+                )
+
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Implementations should use the experiment configuration
+        functionality of a tracking API.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        pass
+
+    def log(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `values` to the current run. Base `log` implementations of a tracking API should go in here, along with
+        special behavior for the `step parameter.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        pass
+
+    def finish(self):
+        """
+        Should run any finalizing functions within the tracking API. If the API should not have one, just don't
+        overwrite that method.
+        """
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run
+        logging_dir (`str`, `os.PathLike`):
+            Location for TensorBoard logs to be stored.
+        **kwargs (additional keyword arguments, *optional*):
+            Additional key word arguments passed along to the `tensorboard.SummaryWriter.__init__` method.
+    """
+
+    name = "tensorboard"
+    requires_logging_directory = True
+
+    @on_main_process
+    def __init__(self, run_name: str, logging_dir: Union[str, os.PathLike], **kwargs):
+        try:
+            from torch.utils import tensorboard
+        except ModuleNotFoundError:
+            import tensorboardX as tensorboard
+        super().__init__()
+        self.run_name = run_name
+        self.logging_dir = os.path.join(logging_dir, run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir, **kwargs)
+        logger.debug(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}")
+        logger.debug(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the
+        hyperparameters in a yaml file for future use.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        project_run_name = time.time()
+        dir_name = os.path.join(self.logging_dir, str(project_run_name))
+        os.makedirs(dir_name, exist_ok=True)
+        with open(os.path.join(dir_name, "hparams.yml"), "w") as outfile:
+            try:
+                yaml.dump(values, outfile)
+            except yaml.representer.RepresenterError:
+                logger.error("Serialization to store hyperparameters failed")
+                raise
+        logger.debug("Stored initial configuration hyperparameters to TensorBoard and hparams yaml file")
+
+    @on_main_process
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of
+                `str` to `float`/`int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to either `SummaryWriter.add_scaler`,
+                `SummaryWriter.add_text`, or `SummaryWriter.add_scalers` method based on the contents of `values`.
+        """
+        values = listify(values)
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v, global_step=step, **kwargs)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v, global_step=step, **kwargs)
+            elif isinstance(v, dict):
+                self.writer.add_scalars(k, v, global_step=step, **kwargs)
+        self.writer.flush()
+        logger.debug("Successfully logged to TensorBoard")
+
+    @on_main_process
+    def log_images(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `images` to the current run.
+
+        Args:
+            values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`):
+                Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `SummaryWriter.add_image` method.
+        """
+        for k, v in values.items():
+            self.writer.add_images(k, v, global_step=step, **kwargs)
+        logger.debug("Successfully logged images to TensorBoard")
+
+    @on_main_process
+    def finish(self):
+        """
+        Closes `TensorBoard` writer
+        """
+        self.writer.close()
+        logger.debug("TensorBoard writer closed")
+
+
+class WandBTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `wandb`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        **kwargs (additional keyword arguments, *optional*):
+            Additional key word arguments passed along to the `wandb.init` method.
+    """
+
+    name = "wandb"
+    requires_logging_directory = False
+    main_process_only = False
+
+    @on_main_process
+    def __init__(self, run_name: str, **kwargs):
+        super().__init__()
+        self.run_name = run_name
+
+        import wandb
+
+        self.run = wandb.init(project=self.run_name, **kwargs)
+        logger.debug(f"Initialized WandB project {self.run_name}")
+        logger.debug(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.run
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        import wandb
+
+        wandb.config.update(values, allow_val_change=True)
+        logger.debug("Stored initial configuration hyperparameters to WandB")
+
+    @on_main_process
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of
+                `str` to `float`/`int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `wandb.log` method.
+        """
+        self.run.log(values, step=step, **kwargs)
+        logger.debug("Successfully logged to WandB")
+
+    @on_main_process
+    def log_images(self, values: dict, step: Optional[int] = None, **kwargs):
+        """
+        Logs `images` to the current run.
+
+        Args:
+            values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`):
+                Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `wandb.log` method.
+        """
+        import wandb
+
+        for k, v in values.items():
+            self.log({k: [wandb.Image(image) for image in v]}, step=step, **kwargs)
+        logger.debug("Successfully logged images to WandB")
+
+    @on_main_process
+    def log_table(
+        self,
+        table_name: str,
+        columns: List[str] = None,
+        data: List[List[Any]] = None,
+        dataframe: Any = None,
+        step: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Log a Table containing any object type (text, image, audio, video, molecule, html, etc). Can be defined either
+        with `columns` and `data` or with `dataframe`.
+
+        Args:
+            table_name (`str`):
+                The name to give to the logged table on the wandb workspace
+            columns (list of `str`, *optional*):
+                The name of the columns on the table
+            data (List of List of Any data type, *optional*):
+                The data to be logged in the table
+            dataframe (Any data type, *optional*):
+                The data to be logged in the table
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        import wandb
+
+        values = {table_name: wandb.Table(columns=columns, data=data, dataframe=dataframe)}
+        self.log(values, step=step, **kwargs)
+
+    @on_main_process
+    def finish(self):
+        """
+        Closes `wandb` writer
+        """
+        self.run.finish()
+        logger.debug("WandB run closed")
+
+
+class CometMLTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `comet_ml`. Should be initialized at the start of your script.
+
+    API keys must be stored in a Comet config file.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        **kwargs (additional keyword arguments, *optional*):
+            Additional key word arguments passed along to the `Experiment.__init__` method.
+    """
+
+    name = "comet_ml"
+    requires_logging_directory = False
+
+    @on_main_process
+    def __init__(self, run_name: str, **kwargs):
+        super().__init__()
+        self.run_name = run_name
+
+        from comet_ml import Experiment
+
+        self.writer = Experiment(project_name=run_name, **kwargs)
+        logger.debug(f"Initialized CometML project {self.run_name}")
+        logger.debug(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float` or `int`):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, `int`, or `None`.
+        """
+        self.writer.log_parameters(values)
+        logger.debug("Stored initial configuration hyperparameters to CometML")
+
+    @on_main_process
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of
+                `str` to `float`/`int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to either `Experiment.log_metric`, `Experiment.log_other`,
+                or `Experiment.log_metrics` method based on the contents of `values`.
+        """
+        if step is not None:
+            self.writer.set_step(step)
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.log_metric(k, v, step=step, **kwargs)
+            elif isinstance(v, str):
+                self.writer.log_other(k, v, **kwargs)
+            elif isinstance(v, dict):
+                self.writer.log_metrics(v, step=step, **kwargs)
+        logger.debug("Successfully logged to CometML")
+
+    @on_main_process
+    def finish(self):
+        """
+        Closes `comet-ml` writer
+        """
+        self.writer.end()
+        logger.debug("CometML run closed")
+
+
+class AimTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `aim`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`):
+            The name of the experiment run.
+        **kwargs (additional keyword arguments, *optional*):
+            Additional key word arguments passed along to the `Run.__init__` method.
+    """
+
+    name = "aim"
+    requires_logging_directory = True
+
+    @on_main_process
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        self.run_name = run_name
+
+        from aim import Run
+
+        self.writer = Run(repo=logging_dir, **kwargs)
+        self.writer.name = self.run_name
+        logger.debug(f"Initialized Aim project {self.run_name}")
+        logger.debug(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.writer
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        self.writer["hparams"] = values
+
+    @on_main_process
+    def log(self, values: dict, step: Optional[int], **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `Run.track` method.
+        """
+        # Note: replace this with the dictionary support when merged
+        for key, value in values.items():
+            self.writer.track(value, name=key, step=step, **kwargs)
+
+    @on_main_process
+    def log_images(self, values: dict, step: Optional[int] = None, kwargs: Optional[Dict[str, dict]] = None):
+        """
+        Logs `images` to the current run.
+
+        Args:
+            values (`Dict[str, Union[np.ndarray, PIL.Image, Tuple[np.ndarray, str], Tuple[PIL.Image, str]]]`):
+                Values to be logged as key-value pairs. The values need to have type `np.ndarray` or PIL.Image. If a
+                tuple is provided, the first element should be the image and the second element should be the caption.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs (`Dict[str, dict]`):
+                Additional key word arguments passed along to the `Run.Image` and `Run.track` method specified by the
+                keys `aim_image` and `track`, respectively.
+        """
+        import aim
+
+        aim_image_kw = {}
+        track_kw = {}
+
+        if kwargs is not None:
+            aim_image_kw = kwargs.get("aim_image", {})
+            track_kw = kwargs.get("track", {})
+
+        for key, value in values.items():
+            if isinstance(value, tuple):
+                img, caption = value
+            else:
+                img, caption = value, ""
+            aim_image = aim.Image(img, caption=caption, **aim_image_kw)
+            self.writer.track(aim_image, name=key, step=step, **track_kw)
+
+    @on_main_process
+    def finish(self):
+        """
+        Closes `aim` writer
+        """
+        self.writer.close()
+
+
+class MLflowTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script.
+
+    Args:
+        experiment_name (`str`, *optional*):
+            Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument.
+        logging_dir (`str` or `os.PathLike`, defaults to `"."`):
+            Location for mlflow logs to be stored.
+        run_id (`str`, *optional*):
+            If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s
+            end time is unset and its status is set to running, but the run’s other attributes (source_version,
+            source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument.
+        tags (`Dict[str, str]`, *optional*):
+            An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a
+            run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are
+            set on the new run. Environment variable MLFLOW_TAGS has priority over this argument.
+        nested_run (`bool`, *optional*, defaults to `False`):
+            Controls whether run is nested in parent run. True creates a nested run. Environment variable
+            MLFLOW_NESTED_RUN has priority over this argument.
+        run_name (`str`, *optional*):
+            Name of new run (stored as a mlflow.runName tag). Used only when `run_id` is unspecified.
+        description (`str`, *optional*):
+            An optional string that populates the description box of the run. If a run is being resumed, the
+            description is set on the resumed run. If a new run is being created, the description is set on the new
+            run.
+    """
+
+    name = "mlflow"
+    requires_logging_directory = False
+
+    @on_main_process
+    def __init__(
+        self,
+        experiment_name: str = None,
+        logging_dir: Optional[Union[str, os.PathLike]] = None,
+        run_id: Optional[str] = None,
+        tags: Optional[Union[Dict[str, Any], str]] = None,
+        nested_run: Optional[bool] = False,
+        run_name: Optional[str] = None,
+        description: Optional[str] = None,
+    ):
+        experiment_name = os.environ.get("MLFLOW_EXPERIMENT_NAME", experiment_name)
+        run_id = os.environ.get("MLFLOW_RUN_ID", run_id)
+        tags = os.environ.get("MLFLOW_TAGS", tags)
+        if isinstance(tags, str):
+            tags = json.loads(tags)
+
+        nested_run = os.environ.get("MLFLOW_NESTED_RUN", nested_run)
+
+        import mlflow
+
+        exps = mlflow.search_experiments(filter_string=f"name = '{experiment_name}'")
+        if len(exps) > 0:
+            if len(exps) > 1:
+                logger.warning("Multiple experiments with the same name found. Using first one.")
+            experiment_id = exps[0].experiment_id
+        else:
+            experiment_id = mlflow.create_experiment(
+                name=experiment_name,
+                artifact_location=logging_dir,
+                tags=tags,
+            )
+
+        self.active_run = mlflow.start_run(
+            run_id=run_id,
+            experiment_id=experiment_id,
+            run_name=run_name,
+            nested=nested_run,
+            tags=tags,
+            description=description,
+        )
+
+        logger.debug(f"Initialized mlflow experiment {experiment_name}")
+        logger.debug(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    @property
+    def tracker(self):
+        return self.active_run
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        import mlflow
+
+        for name, value in list(values.items()):
+            # internally, all values are converted to str in MLflow
+            if len(str(value)) > mlflow.utils.validation.MAX_PARAM_VAL_LENGTH:
+                logger.warning_once(
+                    f'Accelerate is attempting to log a value of "{value}" for key "{name}" as a parameter. MLflow\'s'
+                    f" log_param() only accepts values no longer than {mlflow.utils.validation.MAX_PARAM_VAL_LENGTH} characters so we dropped this attribute."
+                )
+                del values[name]
+
+        values_list = list(values.items())
+
+        # MLflow cannot log more than 100 values in one go, so we have to split it
+        for i in range(0, len(values_list), mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH):
+            mlflow.log_params(dict(values_list[i : i + mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH]))
+
+        logger.debug("Stored initial configuration hyperparameters to MLflow")
+
+    @on_main_process
+    def log(self, values: dict, step: Optional[int]):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (`dict`):
+                Values to be logged as key-value pairs.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+        """
+        metrics = {}
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                metrics[k] = v
+            else:
+                logger.warning_once(
+                    f'MLflowTracker is attempting to log a value of "{v}" of type {type(v)} for key "{k}" as a metric. '
+                    "MLflow's log_metric() only accepts float and int types so we dropped this attribute."
+                )
+        import mlflow
+
+        mlflow.log_metrics(metrics, step=step)
+        logger.debug("Successfully logged to mlflow")
+
+    @on_main_process
+    def finish(self):
+        """
+        End the active MLflow run.
+        """
+        import mlflow
+
+        mlflow.end_run()
+
+
+class ClearMLTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `clearml`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`, *optional*):
+            Name of the experiment. Environment variables `CLEARML_PROJECT` and `CLEARML_TASK` have priority over this
+            argument.
+        **kwargs (additional keyword arguments, *optional*):
+            Kwargs passed along to the `Task.__init__` method.
+    """
+
+    name = "clearml"
+    requires_logging_directory = False
+
+    @on_main_process
+    def __init__(self, run_name: str = None, **kwargs):
+        from clearml import Task
+
+        current_task = Task.current_task()
+        self._initialized_externally = False
+        if current_task:
+            self._initialized_externally = True
+            self.task = current_task
+            return
+
+        kwargs.setdefault("project_name", os.environ.get("CLEARML_PROJECT", run_name))
+        kwargs.setdefault("task_name", os.environ.get("CLEARML_TASK", run_name))
+        self.task = Task.init(**kwargs)
+
+    @property
+    def tracker(self):
+        return self.task
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Connect configuration dictionary to the Task object. Should be run at the beginning of your experiment.
+
+        Args:
+            values (`dict`):
+                Values to be stored as initial hyperparameters as key-value pairs.
+        """
+        return self.task.connect_configuration(values)
+
+    @on_main_process
+    def log(self, values: Dict[str, Union[int, float]], step: Optional[int] = None, **kwargs):
+        """
+        Logs `values` dictionary to the current run. The dictionary keys must be strings. The dictionary values must be
+        ints or floats
+
+        Args:
+            values (`Dict[str, Union[int, float]]`):
+                Values to be logged as key-value pairs. If the key starts with 'eval_'/'test_'/'train_', the value will
+                be reported under the 'eval'/'test'/'train' series and the respective prefix will be removed.
+                Otherwise, the value will be reported under the 'train' series, and no prefix will be removed.
+            step (`int`, *optional*):
+                If specified, the values will be reported as scalars, with the iteration number equal to `step`.
+                Otherwise they will be reported as single values.
+            kwargs:
+                Additional key word arguments passed along to the `clearml.Logger.report_single_value` or
+                `clearml.Logger.report_scalar` methods.
+        """
+        clearml_logger = self.task.get_logger()
+        for k, v in values.items():
+            if not isinstance(v, (int, float)):
+                logger.warning_once(
+                    "Accelerator is attempting to log a value of "
+                    f'"{v}" of type {type(v)} for key "{k}" as a scalar. '
+                    "This invocation of ClearML logger's  report_scalar() "
+                    "is incorrect so we dropped this attribute."
+                )
+                continue
+            if step is None:
+                clearml_logger.report_single_value(name=k, value=v, **kwargs)
+                continue
+            title, series = ClearMLTracker._get_title_series(k)
+            clearml_logger.report_scalar(title=title, series=series, value=v, iteration=step, **kwargs)
+
+    @on_main_process
+    def log_images(self, values: dict, step: Optional[int] = None, **kwargs):
+        """
+        Logs `images` to the current run.
+
+        Args:
+            values (`Dict[str, List[Union[np.ndarray, PIL.Image]]`):
+                Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `clearml.Logger.report_image` method.
+        """
+        clearml_logger = self.task.get_logger()
+        for k, v in values.items():
+            title, series = ClearMLTracker._get_title_series(k)
+            clearml_logger.report_image(title=title, series=series, iteration=step, image=v, **kwargs)
+
+    @on_main_process
+    def log_table(
+        self,
+        table_name: str,
+        columns: List[str] = None,
+        data: List[List[Any]] = None,
+        dataframe: Any = None,
+        step: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Log a Table to the task. Can be defined eitherwith `columns` and `data` or with `dataframe`.
+
+        Args:
+            table_name (`str`):
+                The name of the table
+            columns (list of `str`, *optional*):
+                The name of the columns on the table
+            data (List of List of Any data type, *optional*):
+                The data to be logged in the table. If `columns` is not specified, then the first entry in data will be
+                the name of the columns of the table
+            dataframe (Any data type, *optional*):
+                The data to be logged in the table
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `clearml.Logger.report_table` method.
+        """
+        to_report = dataframe
+        if dataframe is None:
+            if data is None:
+                raise ValueError(
+                    "`ClearMLTracker.log_table` requires that `data` to be supplied if `dataframe` is `None`"
+                )
+            to_report = [columns] + data if columns else data
+        title, series = ClearMLTracker._get_title_series(table_name)
+        self.task.get_logger().report_table(title=title, series=series, table_plot=to_report, iteration=step, **kwargs)
+
+    @on_main_process
+    def finish(self):
+        """
+        Close the ClearML task. If the task was initialized externally (e.g. by manually calling `Task.init`), this
+        function is a noop
+        """
+        if self.task and not self._initialized_externally:
+            self.task.close()
+
+    @staticmethod
+    def _get_title_series(name):
+        for prefix in ["eval", "test", "train"]:
+            if name.startswith(prefix + "_"):
+                return name[len(prefix) + 1 :], prefix
+        return name, "train"
+
+
+class DVCLiveTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `dvclive`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`, *optional*):
+            Ignored for dvclive. See `kwargs` instead.
+        kwargs:
+            Additional key word arguments passed along to [`dvclive.Live()`](https://dvc.org/doc/dvclive/live).
+
+    Example:
+
+    ```py
+    from accelerate import Accelerator
+
+    accelerator = Accelerator(log_with="dvclive")
+    accelerator.init_trackers(project_name="my_project", init_kwargs={"dvclive": {"dir": "my_directory"}})
+    ```
+    """
+
+    name = "dvclive"
+    requires_logging_directory = False
+
+    @on_main_process
+    def __init__(self, run_name: Optional[str] = None, live: Optional[Any] = None, **kwargs):
+        from dvclive import Live
+
+        super().__init__()
+        self.live = live if live is not None else Live(**kwargs)
+
+    @property
+    def tracker(self):
+        return self.live
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the
+        hyperparameters in a yaml file for future use.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float`, `int`, or a List or Dict of those types):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, or `int`.
+        """
+        self.live.log_params(values)
+
+    @on_main_process
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to `dvclive.Live.log_metric()`.
+        """
+        from dvclive.plots import Metric
+
+        if step is not None:
+            self.live.step = step
+        for k, v in values.items():
+            if Metric.could_log(v):
+                self.live.log_metric(k, v, **kwargs)
+            else:
+                logger.warning_once(
+                    "Accelerator attempted to log a value of "
+                    f'"{v}" of type {type(v)} for key "{k}" as a scalar. '
+                    "This invocation of DVCLive's Live.log_metric() "
+                    "is incorrect so we dropped this attribute."
+                )
+        self.live.next_step()
+
+    @on_main_process
+    def finish(self):
+        """
+        Closes `dvclive.Live()`.
+        """
+        self.live.end()
+
+
+LOGGER_TYPE_TO_CLASS = {
+    "aim": AimTracker,
+    "comet_ml": CometMLTracker,
+    "mlflow": MLflowTracker,
+    "tensorboard": TensorBoardTracker,
+    "wandb": WandBTracker,
+    "clearml": ClearMLTracker,
+    "dvclive": DVCLiveTracker,
+}
+
+
+def filter_trackers(
+    log_with: List[Union[str, LoggerType, GeneralTracker]],
+    logging_dir: Union[str, os.PathLike] = None,
+):
+    """
+    Takes in a list of potential tracker types and checks that:
+        - The tracker wanted is available in that environment
+        - Filters out repeats of tracker types
+        - If `all` is in `log_with`, will return all trackers in the environment
+        - If a tracker requires a `logging_dir`, ensures that `logging_dir` is not `None`
+
+    Args:
+        log_with (list of `str`, [`~utils.LoggerType`] or [`~tracking.GeneralTracker`], *optional*):
+            A list of loggers to be setup for experiment tracking. Should be one or several of:
+
+            - `"all"`
+            - `"tensorboard"`
+            - `"wandb"`
+            - `"comet_ml"`
+            - `"mlflow"`
+            - `"dvclive"`
+            If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
+            also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
+        logging_dir (`str`, `os.PathLike`, *optional*):
+            A path to a directory for storing logs of locally-compatible loggers.
+    """
+    loggers = []
+    if log_with is not None:
+        if not isinstance(log_with, (list, tuple)):
+            log_with = [log_with]
+        if "all" in log_with or LoggerType.ALL in log_with:
+            loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers()
+        else:
+            for log_type in log_with:
+                if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker):
+                    raise ValueError(f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}")
+                if issubclass(type(log_type), GeneralTracker):
+                    loggers.append(log_type)
+                else:
+                    log_type = LoggerType(log_type)
+                    if log_type not in loggers:
+                        if log_type in get_available_trackers():
+                            tracker_init = LOGGER_TYPE_TO_CLASS[str(log_type)]
+                            if tracker_init.requires_logging_directory:
+                                if logging_dir is None:
+                                    raise ValueError(
+                                        f"Logging with `{log_type}` requires a `logging_dir` to be passed in."
+                                    )
+                            loggers.append(log_type)
+                        else:
+                            logger.debug(f"Tried adding logger {log_type}, but package is unavailable in the system.")
+
+    return loggers

.venv/Lib/site-packages/decorator.py

@@ -0,0 +1,451 @@
+# #########################     LICENSE     ############################ #
+
+# Copyright (c) 2005-2021, Michele Simionato
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met:
+
+#   Redistributions of source code must retain the above copyright
+#   notice, this list of conditions and the following disclaimer.
+#   Redistributions in bytecode form must reproduce the above copyright
+#   notice, this list of conditions and the following disclaimer in
+#   the documentation and/or other materials provided with the
+#   distribution.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
+# TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
+# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+# DAMAGE.
+
+"""
+Decorator module, see
+https://github.com/micheles/decorator/blob/master/docs/documentation.md
+for the documentation.
+"""
+import re
+import sys
+import inspect
+import operator
+import itertools
+from contextlib import _GeneratorContextManager
+from inspect import getfullargspec, iscoroutinefunction, isgeneratorfunction
+
+__version__ = '5.1.1'
+
+DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(')
+POS = inspect.Parameter.POSITIONAL_OR_KEYWORD
+EMPTY = inspect.Parameter.empty
+
+
+# this is not used anymore in the core, but kept for backward compatibility
+class FunctionMaker(object):
+    """
+    An object with the ability to create functions with a given signature.
+    It has attributes name, doc, module, signature, defaults, dict and
+    methods update and make.
+    """
+
+    # Atomic get-and-increment provided by the GIL
+    _compile_count = itertools.count()
+
+    # make pylint happy
+    args = varargs = varkw = defaults = kwonlyargs = kwonlydefaults = ()
+
+    def __init__(self, func=None, name=None, signature=None,
+                 defaults=None, doc=None, module=None, funcdict=None):
+        self.shortsignature = signature
+        if func:
+            # func can be a class or a callable, but not an instance method
+            self.name = func.__name__
+            if self.name == '<lambda>':  # small hack for lambda functions
+                self.name = '_lambda_'
+            self.doc = func.__doc__
+            self.module = func.__module__
+            if inspect.isroutine(func):
+                argspec = getfullargspec(func)
+                self.annotations = getattr(func, '__annotations__', {})
+                for a in ('args', 'varargs', 'varkw', 'defaults', 'kwonlyargs',
+                          'kwonlydefaults'):
+                    setattr(self, a, getattr(argspec, a))
+                for i, arg in enumerate(self.args):
+                    setattr(self, 'arg%d' % i, arg)
+                allargs = list(self.args)
+                allshortargs = list(self.args)
+                if self.varargs:
+                    allargs.append('*' + self.varargs)
+                    allshortargs.append('*' + self.varargs)
+                elif self.kwonlyargs:
+                    allargs.append('*')  # single star syntax
+                for a in self.kwonlyargs:
+                    allargs.append('%s=None' % a)
+                    allshortargs.append('%s=%s' % (a, a))
+                if self.varkw:
+                    allargs.append('**' + self.varkw)
+                    allshortargs.append('**' + self.varkw)
+                self.signature = ', '.join(allargs)
+                self.shortsignature = ', '.join(allshortargs)
+                self.dict = func.__dict__.copy()
+        # func=None happens when decorating a caller
+        if name:
+            self.name = name
+        if signature is not None:
+            self.signature = signature
+        if defaults:
+            self.defaults = defaults
+        if doc:
+            self.doc = doc
+        if module:
+            self.module = module
+        if funcdict:
+            self.dict = funcdict
+        # check existence required attributes
+        assert hasattr(self, 'name')
+        if not hasattr(self, 'signature'):
+            raise TypeError('You are decorating a non function: %s' % func)
+
+    def update(self, func, **kw):
+        """
+        Update the signature of func with the data in self
+        """
+        func.__name__ = self.name
+        func.__doc__ = getattr(self, 'doc', None)
+        func.__dict__ = getattr(self, 'dict', {})
+        func.__defaults__ = self.defaults
+        func.__kwdefaults__ = self.kwonlydefaults or None
+        func.__annotations__ = getattr(self, 'annotations', None)
+        try:
+            frame = sys._getframe(3)
+        except AttributeError:  # for IronPython and similar implementations
+            callermodule = '?'
+        else:
+            callermodule = frame.f_globals.get('__name__', '?')
+        func.__module__ = getattr(self, 'module', callermodule)
+        func.__dict__.update(kw)
+
+    def make(self, src_templ, evaldict=None, addsource=False, **attrs):
+        """
+        Make a new function from a given template and update the signature
+        """
+        src = src_templ % vars(self)  # expand name and signature
+        evaldict = evaldict or {}
+        mo = DEF.search(src)
+        if mo is None:
+            raise SyntaxError('not a valid function template\n%s' % src)
+        name = mo.group(1)  # extract the function name
+        names = set([name] + [arg.strip(' *') for arg in
+                              self.shortsignature.split(',')])
+        for n in names:
+            if n in ('_func_', '_call_'):
+                raise NameError('%s is overridden in\n%s' % (n, src))
+
+        if not src.endswith('\n'):  # add a newline for old Pythons
+            src += '\n'
+
+        # Ensure each generated function has a unique filename for profilers
+        # (such as cProfile) that depend on the tuple of (<filename>,
+        # <definition line>, <function name>) being unique.
+        filename = '<decorator-gen-%d>' % next(self._compile_count)
+        try:
+            code = compile(src, filename, 'single')
+            exec(code, evaldict)
+        except Exception:
+            print('Error in generated code:', file=sys.stderr)
+            print(src, file=sys.stderr)
+            raise
+        func = evaldict[name]
+        if addsource:
+            attrs['__source__'] = src
+        self.update(func, **attrs)
+        return func
+
+    @classmethod
+    def create(cls, obj, body, evaldict, defaults=None,
+               doc=None, module=None, addsource=True, **attrs):
+        """
+        Create a function from the strings name, signature and body.
+        evaldict is the evaluation dictionary. If addsource is true an
+        attribute __source__ is added to the result. The attributes attrs
+        are added, if any.
+        """
+        if isinstance(obj, str):  # "name(signature)"
+            name, rest = obj.strip().split('(', 1)
+            signature = rest[:-1]  # strip a right parens
+            func = None
+        else:  # a function
+            name = None
+            signature = None
+            func = obj
+        self = cls(func, name, signature, defaults, doc, module)
+        ibody = '\n'.join('    ' + line for line in body.splitlines())
+        caller = evaldict.get('_call_')  # when called from `decorate`
+        if caller and iscoroutinefunction(caller):
+            body = ('async def %(name)s(%(signature)s):\n' + ibody).replace(
+                'return', 'return await')
+        else:
+            body = 'def %(name)s(%(signature)s):\n' + ibody
+        return self.make(body, evaldict, addsource, **attrs)
+
+
+def fix(args, kwargs, sig):
+    """
+    Fix args and kwargs to be consistent with the signature
+    """
+    ba = sig.bind(*args, **kwargs)
+    ba.apply_defaults()  # needed for test_dan_schult
+    return ba.args, ba.kwargs
+
+
+def decorate(func, caller, extras=(), kwsyntax=False):
+    """
+    Decorates a function/generator/coroutine using a caller.
+    If kwsyntax is True calling the decorated functions with keyword
+    syntax will pass the named arguments inside the ``kw`` dictionary,
+    even if such argument are positional, similarly to what functools.wraps
+    does. By default kwsyntax is False and the the arguments are untouched.
+    """
+    sig = inspect.signature(func)
+    if iscoroutinefunction(caller):
+        async def fun(*args, **kw):
+            if not kwsyntax:
+                args, kw = fix(args, kw, sig)
+            return await caller(func, *(extras + args), **kw)
+    elif isgeneratorfunction(caller):
+        def fun(*args, **kw):
+            if not kwsyntax:
+                args, kw = fix(args, kw, sig)
+            for res in caller(func, *(extras + args), **kw):
+                yield res
+    else:
+        def fun(*args, **kw):
+            if not kwsyntax:
+                args, kw = fix(args, kw, sig)
+            return caller(func, *(extras + args), **kw)
+    fun.__name__ = func.__name__
+    fun.__doc__ = func.__doc__
+    fun.__wrapped__ = func
+    fun.__signature__ = sig
+    fun.__qualname__ = func.__qualname__
+    # builtin functions like defaultdict.__setitem__ lack many attributes
+    try:
+        fun.__defaults__ = func.__defaults__
+    except AttributeError:
+        pass
+    try:
+        fun.__kwdefaults__ = func.__kwdefaults__
+    except AttributeError:
+        pass
+    try:
+        fun.__annotations__ = func.__annotations__
+    except AttributeError:
+        pass
+    try:
+        fun.__module__ = func.__module__
+    except AttributeError:
+        pass
+    try:
+        fun.__dict__.update(func.__dict__)
+    except AttributeError:
+        pass
+    return fun
+
+
+def decoratorx(caller):
+    """
+    A version of "decorator" implemented via "exec" and not via the
+    Signature object. Use this if you are want to preserve the `.__code__`
+    object properties (https://github.com/micheles/decorator/issues/129).
+    """
+    def dec(func):
+        return FunctionMaker.create(
+            func,
+            "return _call_(_func_, %(shortsignature)s)",
+            dict(_call_=caller, _func_=func),
+            __wrapped__=func, __qualname__=func.__qualname__)
+    return dec
+
+
+def decorator(caller, _func=None, kwsyntax=False):
+    """
+    decorator(caller) converts a caller function into a decorator
+    """
+    if _func is not None:  # return a decorated function
+        # this is obsolete behavior; you should use decorate instead
+        return decorate(_func, caller, (), kwsyntax)
+    # else return a decorator function
+    sig = inspect.signature(caller)
+    dec_params = [p for p in sig.parameters.values() if p.kind is POS]
+
+    def dec(func=None, *args, **kw):
+        na = len(args) + 1
+        extras = args + tuple(kw.get(p.name, p.default)
+                              for p in dec_params[na:]
+                              if p.default is not EMPTY)
+        if func is None:
+            return lambda func: decorate(func, caller, extras, kwsyntax)
+        else:
+            return decorate(func, caller, extras, kwsyntax)
+    dec.__signature__ = sig.replace(parameters=dec_params)
+    dec.__name__ = caller.__name__
+    dec.__doc__ = caller.__doc__
+    dec.__wrapped__ = caller
+    dec.__qualname__ = caller.__qualname__
+    dec.__kwdefaults__ = getattr(caller, '__kwdefaults__', None)
+    dec.__dict__.update(caller.__dict__)
+    return dec
+
+
+# ####################### contextmanager ####################### #
+
+
+class ContextManager(_GeneratorContextManager):
+    def __init__(self, g, *a, **k):
+        _GeneratorContextManager.__init__(self, g, a, k)
+
+    def __call__(self, func):
+        def caller(f, *a, **k):
+            with self.__class__(self.func, *self.args, **self.kwds):
+                return f(*a, **k)
+        return decorate(func, caller)
+
+
+_contextmanager = decorator(ContextManager)
+
+
+def contextmanager(func):
+    # Enable Pylint config: contextmanager-decorators=decorator.contextmanager
+    return _contextmanager(func)
+
+
+# ############################ dispatch_on ############################ #
+
+def append(a, vancestors):
+    """
+    Append ``a`` to the list of the virtual ancestors, unless it is already
+    included.
+    """
+    add = True
+    for j, va in enumerate(vancestors):
+        if issubclass(va, a):
+            add = False
+            break
+        if issubclass(a, va):
+            vancestors[j] = a
+            add = False
+    if add:
+        vancestors.append(a)
+
+
+# inspired from simplegeneric by P.J. Eby and functools.singledispatch
+def dispatch_on(*dispatch_args):
+    """
+    Factory of decorators turning a function into a generic function
+    dispatching on the given arguments.
+    """
+    assert dispatch_args, 'No dispatch args passed'
+    dispatch_str = '(%s,)' % ', '.join(dispatch_args)
+
+    def check(arguments, wrong=operator.ne, msg=''):
+        """Make sure one passes the expected number of arguments"""
+        if wrong(len(arguments), len(dispatch_args)):
+            raise TypeError('Expected %d arguments, got %d%s' %
+                            (len(dispatch_args), len(arguments), msg))
+
+    def gen_func_dec(func):
+        """Decorator turning a function into a generic function"""
+
+        # first check the dispatch arguments
+        argset = set(getfullargspec(func).args)
+        if not set(dispatch_args) <= argset:
+            raise NameError('Unknown dispatch arguments %s' % dispatch_str)
+
+        typemap = {}
+
+        def vancestors(*types):
+            """
+            Get a list of sets of virtual ancestors for the given types
+            """
+            check(types)
+            ras = [[] for _ in range(len(dispatch_args))]
+            for types_ in typemap:
+                for t, type_, ra in zip(types, types_, ras):
+                    if issubclass(t, type_) and type_ not in t.mro():
+                        append(type_, ra)
+            return [set(ra) for ra in ras]
+
+        def ancestors(*types):
+            """
+            Get a list of virtual MROs, one for each type
+            """
+            check(types)
+            lists = []
+            for t, vas in zip(types, vancestors(*types)):
+                n_vas = len(vas)
+                if n_vas > 1:
+                    raise RuntimeError(
+                        'Ambiguous dispatch for %s: %s' % (t, vas))
+                elif n_vas == 1:
+                    va, = vas
+                    mro = type('t', (t, va), {}).mro()[1:]
+                else:
+                    mro = t.mro()
+                lists.append(mro[:-1])  # discard t and object
+            return lists
+
+        def register(*types):
+            """
+            Decorator to register an implementation for the given types
+            """
+            check(types)
+
+            def dec(f):
+                check(getfullargspec(f).args, operator.lt, ' in ' + f.__name__)
+                typemap[types] = f
+                return f
+            return dec
+
+        def dispatch_info(*types):
+            """
+            An utility to introspect the dispatch algorithm
+            """
+            check(types)
+            lst = []
+            for anc in itertools.product(*ancestors(*types)):
+                lst.append(tuple(a.__name__ for a in anc))
+            return lst
+
+        def _dispatch(dispatch_args, *args, **kw):
+            types = tuple(type(arg) for arg in dispatch_args)
+            try:  # fast path
+                f = typemap[types]
+            except KeyError:
+                pass
+            else:
+                return f(*args, **kw)
+            combinations = itertools.product(*ancestors(*types))
+            next(combinations)  # the first one has been already tried
+            for types_ in combinations:
+                f = typemap.get(types_)
+                if f is not None:
+                    return f(*args, **kw)
+
+            # else call the default implementation
+            return func(*args, **kw)
+
+        return FunctionMaker.create(
+            func, 'return _f_(%s, %%(shortsignature)s)' % dispatch_str,
+            dict(_f_=_dispatch), register=register, default=func,
+            typemap=typemap, vancestors=vancestors, ancestors=ancestors,
+            dispatch_info=dispatch_info, __wrapped__=func)
+
+    gen_func_dec.__name__ = 'dispatch_on' + dispatch_str
+    return gen_func_dec

.venv/Lib/site-packages/isympy.py

@@ -0,0 +1,342 @@
+"""
+Python shell for SymPy.
+
+This is just a normal Python shell (IPython shell if you have the
+IPython package installed), that executes the following commands for
+the user:
+
+    >>> from __future__ import division
+    >>> from sympy import *
+    >>> x, y, z, t = symbols('x y z t')
+    >>> k, m, n = symbols('k m n', integer=True)
+    >>> f, g, h = symbols('f g h', cls=Function)
+    >>> init_printing()
+
+So starting 'isympy' is equivalent to starting Python (or IPython) and
+executing the above commands by hand.  It is intended for easy and quick
+experimentation with SymPy.  isympy is a good way to use SymPy as an
+interactive calculator. If you have IPython and Matplotlib installed, then
+interactive plotting is enabled by default.
+
+COMMAND LINE OPTIONS
+--------------------
+
+-c CONSOLE, --console=CONSOLE
+
+     Use the specified shell (Python or IPython) shell as the console
+     backend instead of the default one (IPython if present, Python
+     otherwise), e.g.:
+
+        $isympy -c python
+
+    CONSOLE must be one of 'ipython' or 'python'
+
+-p PRETTY, --pretty PRETTY
+
+    Setup pretty-printing in SymPy. When pretty-printing is enabled,
+    expressions can be printed with Unicode or ASCII. The default is
+    to use pretty-printing (with Unicode if the terminal supports it).
+    When this option is 'no', expressions will not be pretty-printed
+    and ASCII will be used:
+
+        $isympy -p no
+
+    PRETTY must be one of 'unicode', 'ascii', or 'no'
+
+-t TYPES, --types=TYPES
+
+    Setup the ground types for the polys.  By default, gmpy ground types
+    are used if gmpy2 or gmpy is installed, otherwise it falls back to python
+    ground types, which are a little bit slower.  You can manually
+    choose python ground types even if gmpy is installed (e.g., for
+    testing purposes):
+
+        $isympy -t python
+
+    TYPES must be one of 'gmpy', 'gmpy1' or 'python'
+
+    Note that the ground type gmpy1 is primarily intended for testing; it
+    forces the use of gmpy version 1 even if gmpy2 is available.
+
+    This is the same as setting the environment variable
+    SYMPY_GROUND_TYPES to the given ground type (e.g.,
+    SYMPY_GROUND_TYPES='gmpy')
+
+    The ground types can be determined interactively from the variable
+    sympy.polys.domains.GROUND_TYPES.
+
+-o ORDER, --order ORDER
+
+    Setup the ordering of terms for printing.  The default is lex, which
+    orders terms lexicographically (e.g., x**2 + x + 1). You can choose
+    other orderings, such as rev-lex, which will use reverse
+    lexicographic ordering (e.g., 1 + x + x**2):
+
+        $isympy -o rev-lex
+
+    ORDER must be one of 'lex', 'rev-lex', 'grlex', 'rev-grlex',
+    'grevlex', 'rev-grevlex', 'old', or 'none'.
+
+    Note that for very large expressions, ORDER='none' may speed up
+    printing considerably but the terms will have no canonical order.
+
+-q, --quiet
+
+    Print only Python's and SymPy's versions to stdout at startup.
+
+-d, --doctest
+
+    Use the same format that should be used for doctests.  This is
+    equivalent to -c python -p no.
+
+-C, --no-cache
+
+    Disable the caching mechanism.  Disabling the cache may slow certain
+    operations down considerably.  This is useful for testing the cache,
+    or for benchmarking, as the cache can result in deceptive timings.
+
+    This is equivalent to setting the environment variable
+    SYMPY_USE_CACHE to 'no'.
+
+-a, --auto-symbols (requires at least IPython 0.11)
+
+    Automatically create missing symbols.  Normally, typing a name of a
+    Symbol that has not been instantiated first would raise NameError,
+    but with this option enabled, any undefined name will be
+    automatically created as a Symbol.
+
+    Note that this is intended only for interactive, calculator style
+    usage. In a script that uses SymPy, Symbols should be instantiated
+    at the top, so that it's clear what they are.
+
+    This will not override any names that are already defined, which
+    includes the single character letters represented by the mnemonic
+    QCOSINE (see the "Gotchas and Pitfalls" document in the
+    documentation). You can delete existing names by executing "del
+    name".  If a name is defined, typing "'name' in dir()" will return True.
+
+    The Symbols that are created using this have default assumptions.
+    If you want to place assumptions on symbols, you should create them
+    using symbols() or var().
+
+    Finally, this only works in the top level namespace. So, for
+    example, if you define a function in isympy with an undefined
+    Symbol, it will not work.
+
+    See also the -i and -I options.
+
+-i, --int-to-Integer (requires at least IPython 0.11)
+
+    Automatically wrap int literals with Integer.  This makes it so that
+    things like 1/2 will come out as Rational(1, 2), rather than 0.5.  This
+    works by preprocessing the source and wrapping all int literals with
+    Integer.  Note that this will not change the behavior of int literals
+    assigned to variables, and it also won't change the behavior of functions
+    that return int literals.
+
+    If you want an int, you can wrap the literal in int(), e.g. int(3)/int(2)
+    gives 1.5 (with division imported from __future__).
+
+-I, --interactive (requires at least IPython 0.11)
+
+    This is equivalent to --auto-symbols --int-to-Integer.  Future options
+    designed for ease of interactive use may be added to this.
+
+-D, --debug
+
+    Enable debugging output.  This is the same as setting the
+    environment variable SYMPY_DEBUG to 'True'.  The debug status is set
+    in the variable SYMPY_DEBUG within isympy.
+
+-- IPython options
+
+    Additionally you can pass command line options directly to the IPython
+    interpreter (the standard Python shell is not supported).  However you
+    need to add the '--' separator between two types of options, e.g the
+    startup banner option and the colors option. You need to enter the
+    options as required by the version of IPython that you are using, too:
+
+    in IPython 0.11,
+
+        $isympy -q -- --colors=NoColor
+
+    or older versions of IPython,
+
+        $isympy -q -- -colors NoColor
+
+See also isympy --help.
+"""
+
+import os
+import sys
+
+# DO NOT IMPORT SYMPY HERE! Or the setting of the sympy environment variables
+# by the command line will break.
+
+def main() -> None:
+    from argparse import ArgumentParser, RawDescriptionHelpFormatter
+
+    VERSION = None
+    if '--version' in sys.argv:
+        # We cannot import sympy before this is run, because flags like -C and
+        # -t set environment variables that must be set before SymPy is
+        # imported. The only thing we need to import it for is to get the
+        # version, which only matters with the --version flag.
+        import sympy
+        VERSION = sympy.__version__
+
+    usage = 'isympy [options] -- [ipython options]'
+    parser = ArgumentParser(
+        usage=usage,
+        description=__doc__,
+        formatter_class=RawDescriptionHelpFormatter,
+    )
+
+    parser.add_argument('--version', action='version', version=VERSION)
+
+    parser.add_argument(
+        '-c', '--console',
+        dest='console',
+        action='store',
+        default=None,
+        choices=['ipython', 'python'],
+        metavar='CONSOLE',
+        help='select type of interactive session: ipython | python; defaults '
+        'to ipython if IPython is installed, otherwise python')
+
+    parser.add_argument(
+        '-p', '--pretty',
+        dest='pretty',
+        action='store',
+        default=None,
+        metavar='PRETTY',
+        choices=['unicode', 'ascii', 'no'],
+        help='setup pretty printing: unicode | ascii | no; defaults to '
+        'unicode printing if the terminal supports it, otherwise ascii')
+
+    parser.add_argument(
+        '-t', '--types',
+        dest='types',
+        action='store',
+        default=None,
+        metavar='TYPES',
+        choices=['gmpy', 'gmpy1', 'python'],
+        help='setup ground types: gmpy | gmpy1 | python; defaults to gmpy if gmpy2 '
+        'or gmpy is installed, otherwise python')
+
+    parser.add_argument(
+        '-o', '--order',
+        dest='order',
+        action='store',
+        default=None,
+        metavar='ORDER',
+        choices=['lex', 'grlex', 'grevlex', 'rev-lex', 'rev-grlex', 'rev-grevlex', 'old', 'none'],
+        help='setup ordering of terms: [rev-]lex | [rev-]grlex | [rev-]grevlex | old | none; defaults to lex')
+
+    parser.add_argument(
+        '-q', '--quiet',
+        dest='quiet',
+        action='store_true',
+        default=False,
+        help='print only version information at startup')
+
+    parser.add_argument(
+        '-d', '--doctest',
+        dest='doctest',
+        action='store_true',
+        default=False,
+        help='use the doctest format for output (you can just copy and paste it)')
+
+    parser.add_argument(
+        '-C', '--no-cache',
+        dest='cache',
+        action='store_false',
+        default=True,
+        help='disable caching mechanism')
+
+    parser.add_argument(
+        '-a', '--auto-symbols',
+        dest='auto_symbols',
+        action='store_true',
+        default=False,
+        help='automatically construct missing symbols')
+
+    parser.add_argument(
+        '-i', '--int-to-Integer',
+        dest='auto_int_to_Integer',
+        action='store_true',
+        default=False,
+        help="automatically wrap int literals with Integer")
+
+    parser.add_argument(
+        '-I', '--interactive',
+        dest='interactive',
+        action='store_true',
+        default=False,
+        help="equivalent to -a -i")
+
+    parser.add_argument(
+        '-D', '--debug',
+        dest='debug',
+        action='store_true',
+        default=False,
+        help='enable debugging output')
+
+    (options, ipy_args) = parser.parse_known_args()
+    if '--' in ipy_args:
+        ipy_args.remove('--')
+
+    if not options.cache:
+        os.environ['SYMPY_USE_CACHE'] = 'no'
+
+    if options.types:
+        os.environ['SYMPY_GROUND_TYPES'] = options.types
+
+    if options.debug:
+        os.environ['SYMPY_DEBUG'] = str(options.debug)
+
+    if options.doctest:
+        options.pretty = 'no'
+        options.console = 'python'
+
+    session = options.console
+
+    if session is not None:
+        ipython = session == 'ipython'
+    else:
+        try:
+            import IPython
+            ipython = True
+        except ImportError:
+            if not options.quiet:
+                from sympy.interactive.session import no_ipython
+                print(no_ipython)
+            ipython = False
+
+    args = {
+        'pretty_print': True,
+        'use_unicode':  None,
+        'use_latex':    None,
+        'order':        None,
+        'argv':         ipy_args,
+    }
+
+    if options.pretty == 'unicode':
+        args['use_unicode'] = True
+    elif options.pretty == 'ascii':
+        args['use_unicode'] = False
+    elif options.pretty == 'no':
+        args['pretty_print'] = False
+
+    if options.order is not None:
+        args['order'] = options.order
+
+    args['quiet'] = options.quiet
+    args['auto_symbols'] = options.auto_symbols or options.interactive
+    args['auto_int_to_Integer'] = options.auto_int_to_Integer or options.interactive
+
+    from sympy.interactive import init_session
+    init_session(ipython, **args)
+
+if __name__ == "__main__":
+    main()

.venv/Lib/site-packages/plac.py

@@ -0,0 +1,37 @@
+# #########################     LICENSE     ###############################
+#
+#   Copyright (c) 2010-2021, Michele Simionato
+#   All rights reserved.
+#
+#   Redistributions of source code must retain the above copyright
+#   notice, this list of conditions and the following disclaimer.
+#   Redistributions in bytecode form must reproduce the above copyright
+#   notice, this list of conditions and the following disclaimer in
+#   the documentation and/or other materials provided with the
+#   distribution.
+
+#   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+#   "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+#   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+#   A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+#   HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+#   INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+#   BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+#   OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+#   ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
+#   TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
+#   USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+#   DAMAGE.
+"""
+See docs/index.html for the documentation.
+"""
+from plac_core import *
+from plac_ext import (import_main, ReadlineInput, Interpreter,
+                      stdout, runp, Monitor, default_help)
+
+__version__ = '1.4.3'
+
+try:
+    from plac_tk import TkMonitor
+except ImportError:
+    pass

.venv/Lib/site-packages/plac_core.py

@@ -0,0 +1,439 @@
+# this module should be kept Python 2.3 compatible
+import re
+import sys
+import time
+import inspect
+import textwrap
+import functools
+import argparse
+from datetime import datetime, date
+from gettext import gettext as _
+
+version = sys.version_info[:2]
+
+if sys.version >= '3':
+    from inspect import getfullargspec
+else:
+    class getfullargspec(object):
+        "A quick and dirty replacement for getfullargspec for Python 2.X"
+        def __init__(self, f):
+            self.args, self.varargs, self.varkw, self.defaults = \
+                inspect.getargspec(f)
+            self.annotations = getattr(f, '__annotations__', {})
+
+
+def to_date(s):
+    """Returns year-month-day"""
+    return date(*time.strptime(s, "%Y-%m-%d")[0:3])
+
+
+def to_datetime(s):
+    """Returns year-month-day hour-minute-second"""
+    return datetime(*time.strptime(s, "%Y-%m-%d %H-%M-%S")[0:6])
+
+
+def getargspec(callableobj):
+    """Given a callable return an object with attributes .args, .varargs,
+    .varkw, .defaults. It tries to do the "right thing" with functions,
+    methods, classes and generic callables."""
+    if inspect.isfunction(callableobj):
+        argspec = getfullargspec(callableobj)
+    elif inspect.ismethod(callableobj):
+        argspec = getfullargspec(callableobj)
+        del argspec.args[0]  # remove first argument
+    elif inspect.isclass(callableobj):
+        if callableobj.__init__ is object.__init__:  # to avoid an error
+            argspec = getfullargspec(lambda self: None)
+        else:
+            argspec = getfullargspec(callableobj.__init__)
+        del argspec.args[0]  # remove first argument
+    elif hasattr(callableobj, '__call__'):
+        argspec = getfullargspec(callableobj.__call__)
+        del argspec.args[0]  # remove first argument
+    else:
+        raise TypeError(_('Could not determine the signature of ') +
+                        str(callableobj))
+    return argspec
+
+
+def annotations(**ann):
+    """
+    Returns a decorator annotating a function with the given annotations.
+    This is a trick to support function annotations in Python 2.X.
+    """
+    def annotate(f):
+        fas = getfullargspec(f)
+        args = fas.args
+        if fas.varargs:
+            args.append(fas.varargs)
+        if fas.varkw:
+            args.append(fas.varkw)
+        for argname in ann:
+            if argname not in args:
+                raise NameError(
+                    _('Annotating non-existing argument: %s') % argname)
+        f.__annotations__ = ann
+        return f
+    return annotate
+
+
+def _annotate(arg, ann, f):
+    try:
+        f.__annotations__[arg] = ann
+    except AttributeError:  # Python 2.7
+        f.__annotations__ = {arg: ann}
+    return f
+
+
+def pos(arg, help=None, type=None, choices=None, metavar=None):
+    """
+    Decorator for annotating positional arguments
+    """
+    return functools.partial(
+        _annotate, arg, (help, 'positional', None, type, choices, metavar))
+
+
+def opt(arg, help=None, type=None, abbrev=None, choices=None, metavar=None):
+    """
+    Decorator for annotating optional arguments
+    """
+    abbrev = abbrev or arg[0]
+    return functools.partial(
+        _annotate, arg, (help, 'option', abbrev, type, choices, metavar))
+
+
+def flg(arg, help=None, abbrev=None):
+    """
+    Decorator for annotating flags
+    """
+    return functools.partial(
+        _annotate, arg, (help, 'flag', abbrev or arg[0], None, None, None))
+
+
+def is_annotation(obj):
+    """
+    An object is an annotation object if it has the attributes
+    help, kind, abbrev, type, choices, metavar.
+    """
+    return (hasattr(obj, 'help') and hasattr(obj, 'kind')
+            and hasattr(obj, 'abbrev') and hasattr(obj, 'type')
+            and hasattr(obj, 'choices') and hasattr(obj, 'metavar'))
+
+
+class Annotation(object):
+    def __init__(self, help=None, kind="positional", abbrev=None, type=None,
+                 choices=None, metavar=None):
+        assert kind in ('positional', 'option', 'flag'), kind
+        if kind == "positional":
+            assert abbrev is None, abbrev
+        self.help = help
+        self.kind = kind
+        self.abbrev = abbrev
+        self.type = type
+        self.choices = choices
+        self.metavar = metavar
+
+    def from_(cls, obj):
+        "Helper to convert an object into an annotation, if needed"
+        if is_annotation(obj):
+            return obj  # do nothing
+        elif inspect.isclass(obj):
+            obj = str(obj)
+        elif iterable(obj):
+            return cls(*obj)
+        return cls(obj)
+    from_ = classmethod(from_)
+
+
+NONE = object()  # sentinel use to signal the absence of a default
+
+PARSER_CFG = getfullargspec(argparse.ArgumentParser.__init__).args[1:]
+# the default arguments accepted by an ArgumentParser object
+
+
+def pconf(obj):
+    """
+    Extracts the configuration of the underlying ArgumentParser from obj
+    """
+    cfg = dict(description=(textwrap.dedent(obj.__doc__.rstrip())
+                            if obj.__doc__ else None),
+               formatter_class=argparse.RawDescriptionHelpFormatter)
+    for name in dir(obj):
+        if name in PARSER_CFG:  # argument of ArgumentParser
+            cfg[name] = getattr(obj, name)
+    return cfg
+
+
+_parser_registry = {}
+
+
+def parser_from(obj, **confparams):
+    """
+    obj can be a callable or an object with a .commands attribute.
+    Returns an ArgumentParser.
+    """
+    try:  # the underlying parser has been generated already
+        return _parser_registry[obj]
+    except KeyError:  # generate a new parser
+        pass
+    conf = pconf(obj).copy()
+    conf.update(confparams)
+    _parser_registry[obj] = parser = ArgumentParser(**conf)
+    parser.obj = obj
+    parser.case_sensitive = confparams.get(
+        'case_sensitive', getattr(obj, 'case_sensitive', True))
+    if hasattr(obj, 'commands') and not inspect.isclass(obj):
+        # a command container instance
+        parser.addsubcommands(obj.commands, obj, 'subcommands')
+    else:
+        parser.populate_from(obj)
+    return parser
+
+
+def _extract_kwargs(args):
+    """
+    Returns two lists: regular args and name=value args
+    """
+    arglist = []
+    kwargs = {}
+    for arg in args:
+        match = re.match(r'([a-zA-Z_]\w*)=', arg)
+        if match:
+            name = match.group(1)
+            kwargs[name] = arg[len(name)+1:]
+        else:
+            arglist.append(arg)
+    return arglist, kwargs
+
+
+def _match_cmd(abbrev, commands, case_sensitive=True):
+    """
+    Extract the command name from an abbreviation or raise a NameError
+    """
+    if not case_sensitive:
+        abbrev = abbrev.upper()
+        commands = [c.upper() for c in commands]
+    perfect_matches = [name for name in commands if name == abbrev]
+    if len(perfect_matches) == 1:
+        return perfect_matches[0]
+    matches = [name for name in commands if name.startswith(abbrev)]
+    n = len(matches)
+    if n == 1:
+        return matches[0]
+    elif n > 1:
+        raise NameError(
+            _('Ambiguous command %r: matching %s' % (abbrev, matches)))
+
+
+class ArgumentParser(argparse.ArgumentParser):
+    """
+    An ArgumentParser with .func and .argspec attributes, and possibly
+    .commands and .subparsers.
+    """
+    case_sensitive = True
+
+    if version < (3, 10):
+        def __init__(self, *args, **kwargs):
+            super(ArgumentParser, self).__init__(*args, **kwargs)
+            if self._action_groups[1].title == _('optional arguments'):
+                self._action_groups[1].title = _('options')
+
+    def alias(self, arg):
+        "Can be overridden to preprocess command-line arguments"
+        return arg
+
+    def consume(self, args):
+        """
+        Call the underlying function with the args. Works also for
+        command containers, by dispatching to the right subparser.
+        """
+        arglist = [self.alias(a) for a in args]
+        cmd = None
+        if hasattr(self, 'subparsers'):
+            subp, cmd = self._extract_subparser_cmd(arglist)
+            if subp is None and cmd is not None:
+                return cmd, self.missing(cmd)
+            elif subp is not None:  # use the subparser
+                self = subp
+        if hasattr(self, 'argspec') and self.argspec.varargs:
+            # ignore unrecognized arguments
+            ns, extraopts = self.parse_known_args(arglist)
+        else:
+            ns, extraopts = self.parse_args(arglist), []  # may raise an exit
+        if not hasattr(self, 'argspec'):
+            raise SystemExit
+        if hasattr(self, 'argspec') and self.argspec.varkw:
+            v = self.argspec.varargs
+            varkw = self.argspec.varkw
+            if v in ns.__dict__:
+                lst = ns.__dict__.pop(v)
+                lst, kwargs = _extract_kwargs(lst)
+                ns.__dict__[v] = lst
+            elif varkw in ns.__dict__:
+                lst = ns.__dict__.pop(varkw)
+                lst, kwargs = _extract_kwargs(lst)
+                ns.__dict__[varkw] = lst
+            if lst and not v:
+                self.error(_('Unrecognized arguments: %s') % arglist)
+        else:
+            kwargs = {}
+        collision = set(self.argspec.args) & set(kwargs)
+        if collision:
+            self.error(
+                _('colliding keyword arguments: %s') % ' '.join(collision))
+        # Correct options with trailing undescores
+        args = [getattr(ns, a.rstrip('_')) for a in self.argspec.args]
+        varargs = getattr(ns, self.argspec.varargs or '', [])
+        return cmd, self.func(*(args + varargs + extraopts), **kwargs)
+
+    def _extract_subparser_cmd(self, arglist):
+        """
+        Extract the right subparser from the first recognized argument
+        """
+        optprefix = self.prefix_chars[0]
+        name_parser_map = self.subparsers._name_parser_map
+        for i, arg in enumerate(arglist):
+            if not arg.startswith(optprefix):
+                cmd = _match_cmd(arg, name_parser_map, self.case_sensitive)
+                del arglist[i]
+                return name_parser_map.get(cmd), cmd or arg
+        return None, None
+
+    def addsubcommands(self, commands, obj, title=None, cmdprefix=''):
+        """
+        Extract a list of subcommands from obj and add them to the parser
+        """
+        if hasattr(obj, cmdprefix) and obj.cmdprefix in self.prefix_chars:
+            raise ValueError(_('The prefix %r is already taken!' % cmdprefix))
+        if not hasattr(self, 'subparsers'):
+            self.subparsers = self.add_subparsers(title=title)
+        elif title:
+            self.add_argument_group(title=title)  # populate ._action_groups
+        prefixlen = len(getattr(obj, 'cmdprefix', ''))
+        add_help = getattr(obj, 'add_help', True)
+        for cmd in commands:
+            func = getattr(obj, cmd[prefixlen:])  # strip the prefix
+            doc = (textwrap.dedent(func.__doc__.rstrip())
+                   if func.__doc__ else None)
+            self.subparsers.add_parser(
+                cmd, add_help=add_help, help=doc, **pconf(func)
+                ).populate_from(func)
+
+    def _set_func_argspec(self, obj):
+        """
+        Extracts the signature from a callable object and adds an .argspec
+        attribute to the parser. Also adds a .func reference to the object.
+        """
+        self.func = obj
+        self.argspec = getargspec(obj)
+        _parser_registry[obj] = self
+
+    def populate_from(self, func):
+        """
+        Extract the arguments from the attributes of the passed function
+        and return a populated ArgumentParser instance.
+        """
+        self._set_func_argspec(func)
+        f = self.argspec
+        defaults = f.defaults or ()
+        n_args = len(f.args)
+        n_defaults = len(defaults)
+        alldefaults = (NONE,) * (n_args - n_defaults) + defaults
+        prefix = self.prefix = getattr(func, 'prefix_chars', '-')[0]
+        for name, default in zip(f.args, alldefaults):
+            ann = f.annotations.get(name, ())
+            a = Annotation.from_(ann)
+            metavar = a.metavar
+            if default is NONE:
+                dflt = None
+            else:
+                dflt = default
+                if a.help is None:
+                    a.help = '[%s]' % str(dflt)  # dflt can be a tuple
+                if a.type is None:
+                    # try to infer the type from the default argument
+                    if isinstance(default, datetime):
+                        a.type = to_datetime
+                    elif isinstance(default, date):
+                        a.type = to_date
+                    elif default is not None:
+                        a.type = type(default)
+                if not metavar and default == '':
+                    metavar = "''"
+            if a.kind in ('option', 'flag'):
+
+                if name.endswith("_"):
+                    # allows reserved words to be specified with underscores
+                    suffix = name.rstrip('_')
+                else:
+                    # convert undescores to dashes.
+                    suffix = name.replace('_', '-')
+
+                if a.abbrev:
+                    shortlong = (prefix + a.abbrev,
+                                 prefix*2 + suffix)
+                else:
+                    shortlong = (prefix + suffix,)
+            elif default is NONE:  # required argument
+                self.add_argument(name, help=a.help, type=a.type,
+                                  choices=a.choices, metavar=metavar)
+            else:  # default argument
+                self.add_argument(
+                    name, nargs='?', help=a.help, default=dflt,
+                    type=a.type, choices=a.choices, metavar=metavar)
+            if a.kind == 'option':
+                if default is not NONE:
+                    metavar = metavar or str(default)
+                self.add_argument(
+                    help=a.help, default=dflt, type=a.type,
+                    choices=a.choices, metavar=metavar, *shortlong)
+            elif a.kind == 'flag':
+                if default is not NONE and default is not False:
+                    raise TypeError(_('Flag %r wants default False, got %r') %
+                                    (name, default))
+                self.add_argument(action='store_true', help=a.help, *shortlong)
+        if f.varargs:
+            a = Annotation.from_(f.annotations.get(f.varargs, ()))
+            self.add_argument(f.varargs, nargs='*', help=a.help, default=[],
+                              type=a.type, metavar=a.metavar)
+        if f.varkw:
+            a = Annotation.from_(f.annotations.get(f.varkw, ()))
+            self.add_argument(f.varkw, nargs='*', help=a.help, default={},
+                              type=a.type, metavar=a.metavar)
+
+    def missing(self, name):
+        "May raise a SystemExit"
+        miss = getattr(self.obj, '__missing__', lambda name:
+                       self.error('No command %r' % name))
+        return miss(name)
+
+    def print_actions(self):
+        "Useful for debugging"
+        print(self)
+        for a in self._actions:
+            print(a)
+
+
+def iterable(obj):
+    "Any object with an __iter__ method which is not a string or class"
+    return hasattr(obj, '__iter__') and not inspect.isclass(obj) and not isinstance(obj, (str, bytes))
+
+
+def call(obj, arglist=None, eager=True, version=None):
+    """
+    If obj is a function or a bound method, parse the given arglist
+    by using the parser inferred from the annotations of obj
+    and call obj with the parsed arguments.
+    If obj is an object with attribute .commands, dispatch to the
+    associated subparser.
+    """
+    if arglist is None:
+        arglist = sys.argv[1:]
+    parser = parser_from(obj)
+    if version:
+        parser.add_argument(
+            '--version', '-v', action='version', version=version)
+    cmd, result = parser.consume(arglist)
+    if iterable(result) and eager:  # listify the result
+        return list(result)
+    return result

.venv/Lib/site-packages/plac_ext.py

@@ -0,0 +1,1205 @@
+# this module requires Python 2.6+
+from __future__ import with_statement
+from contextlib import contextmanager
+from operator import attrgetter
+from gettext import gettext as _
+import inspect
+import os
+import sys
+import cmd
+import shlex
+import subprocess
+import argparse
+import itertools
+import traceback
+import multiprocessing
+import signal
+import threading
+import plac_core
+
+version = sys.version_info[:2]
+
+if version < (3, 5):
+    from imp import load_source
+else:
+    import importlib.util
+
+    def load_source(dotname, path):
+        spec = importlib.util.spec_from_file_location(dotname, path)
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        return mod
+
+
+if sys.version < '3':
+    def exec_(_code_, _globs_=None, _locs_=None):
+        if _globs_ is None:
+            frame = sys._getframe(1)
+            _globs_ = frame.f_globals
+            if _locs_ is None:
+                _locs_ = frame.f_locals
+            del frame
+        elif _locs_ is None:
+            _locs_ = _globs_
+        exec("""exec _code_ in _globs_, _locs_""")
+
+    exec('''
+def raise_(tp, value=None, tb=None):
+    raise tp, value, tb
+''')
+else:
+    exec_ = eval('exec')
+
+    def raise_(tp, value=None, tb=None):
+        """
+        A function that matches the Python 2.x ``raise`` statement. This
+        allows re-raising exceptions with the cls value and traceback on
+        Python 2 and 3.
+        """
+        if value is not None and isinstance(tp, Exception):
+            raise TypeError("instance exception may not have a separate value")
+        if value is not None:
+            exc = tp(value)
+        else:
+            exc = tp
+        if exc.__traceback__ is not tb:
+            raise exc.with_traceback(tb)
+        raise exc
+
+try:
+    raw_input
+except NameError:  # Python 3
+    raw_input = input
+
+
+def decode(val):
+    """
+    Decode an object assuming the encoding is UTF-8.
+    """
+    try:
+        # assume it is an encoded bytes object
+        return val.decode('utf-8')
+    except AttributeError:
+        # it was an already decoded unicode object
+        return str(val)
+
+# ############################ generic utils ############################### #
+
+
+@contextmanager
+def stdout(fileobj):
+    "usage: with stdout(file('out.txt', 'a')): do_something()"
+    orig_stdout = sys.stdout
+    sys.stdout = fileobj
+    try:
+        yield
+    finally:
+        sys.stdout = orig_stdout
+
+
+def write(x):
+    "Write str(x) on stdout and flush, no newline added"
+    sys.stdout.write(str(x))
+    sys.stdout.flush()
+
+
+def gen_val(value):
+    "Return a generator object with a single element"
+    yield value
+
+
+def gen_exc(etype, exc, tb):
+    "Return a generator object raising an exception"
+    raise_(etype, exc, tb)
+    yield
+
+
+def less(text):
+    "Send a text to less via a pipe"
+    # -c clear the screen before starting less
+    po = subprocess.Popen(['less', '-c'], stdin=subprocess.PIPE)
+    try:
+        po.stdin.write(text)
+    except IOError:
+        pass
+    po.stdin.close()
+    po.wait()
+
+
+use_less = (sys.platform != 'win32')  # unices
+
+
+class TerminatedProcess(Exception):
+    pass
+
+
+def terminatedProcess(signum, frame):
+    raise TerminatedProcess
+
+
+# ########################## readline support ############################ #
+
+def read_line(stdin, prompt=''):
+    "Read a line from stdin, using readline when possible"
+    if isinstance(stdin, ReadlineInput):
+        return stdin.readline(prompt)
+    else:
+        write(prompt)
+        return stdin.readline()
+
+
+def read_long_line(stdin, terminator):
+    """
+    Read multiple lines from stdin until the terminator character is found,
+    then yield a single space-separated long line.
+    """
+    while True:
+        lines = []
+        while True:
+            line = stdin.readline()  # ends with \n
+            if not line:  # EOF
+                return
+            line = line.strip()
+            if not line:
+                continue
+            elif line[-1] == terminator:
+                lines.append(line[:-1])
+                break
+            else:
+                lines.append(line)
+        yield ' '.join(lines)
+
+
+class ReadlineInput(object):
+    """
+    An iterable with a .readline method reading from stdin.
+    """
+    def __init__(self, completions, case_sensitive=True, histfile=None):
+        self.completions = completions
+        self.case_sensitive = case_sensitive
+        self.histfile = histfile
+        if not case_sensitive:
+            self.completions = [c.upper() for c in completions]
+        import readline
+        self.rl = readline
+        readline.parse_and_bind("tab: complete")
+        readline.set_completer(self.complete)
+
+    def __enter__(self):
+        self.old_completer = self.rl.get_completer()
+        try:
+            if self.histfile:
+                self.rl.read_history_file(self.histfile)
+        except IOError:  # the first time
+            pass
+        return self
+
+    def __exit__(self, etype, exc, tb):
+        self.rl.set_completer(self.old_completer)
+        if self.histfile:
+            self.rl.write_history_file(self.histfile)
+
+    def complete(self, kw, state):
+        # state is 0, 1, 2, ... and increases by hitting TAB
+        if not self.case_sensitive:
+            kw = kw.upper()
+        try:
+            return [k for k in self.completions if k.startswith(kw)][state]
+        except IndexError:  # no completions
+            return  # exit
+
+    def readline(self, prompt=''):
+        try:
+            return raw_input(prompt) + '\n'
+        except EOFError:
+            return ''
+
+    def __iter__(self):
+        return iter(self.readline, '')
+
+# ################# help functionality in plac interpreters ################# #
+
+
+class HelpSummary(object):
+    "Build the help summary consistently with the cmd module"
+
+    @classmethod
+    def add(cls, obj, specialcommands):
+        p = plac_core.parser_from(obj)
+        c = cmd.Cmd(stdout=cls())
+        c.stdout.write('\n')
+        c.print_topics('special commands',
+                       sorted(specialcommands), 15, 80)
+        c.print_topics('custom commands',
+                       sorted(obj.commands), 15, 80)
+        c.print_topics('commands run in external processes',
+                       sorted(obj.mpcommands), 15, 80)
+        c.print_topics('threaded commands',
+                       sorted(obj.thcommands), 15, 80)
+        p.helpsummary = str(c.stdout)
+
+    def __init__(self):
+        self._ls = []
+
+    def write(self, s):
+        self._ls.append(s)
+
+    def __str__(self):
+        return ''.join(self._ls)
+
+
+class PlacFormatter(argparse.RawDescriptionHelpFormatter):
+    def _metavar_formatter(self, action, default_metavar):
+        'Remove special commands from the usage message'
+        choices = action.choices or {}
+        action.choices = dict((n, c) for n, c in choices.items()
+                              if not n.startswith('.'))
+        return super(PlacFormatter, self)._metavar_formatter(
+            action, default_metavar)
+
+
+def format_help(self):
+    "Attached to plac_core.ArgumentParser for plac interpreters"
+    try:
+        return self.helpsummary
+    except AttributeError:
+        return super(plac_core.ArgumentParser, self).format_help()
+plac_core.ArgumentParser.format_help = format_help
+
+
+def default_help(obj, cmd=None):
+    "The default help functionality in plac interpreters"
+    parser = plac_core.parser_from(obj)
+    if cmd is None:
+        yield parser.format_help()
+        return
+    subp = parser.subparsers._name_parser_map.get(cmd)
+    if subp is None:
+        yield _('Unknown command %s' % cmd)
+    elif getattr(obj, '_interact_', False):  # in interactive mode
+        formatter = subp._get_formatter()
+        formatter._prog = cmd  # remove the program name from the usage
+        formatter.add_usage(
+            subp.usage, [a for a in subp._actions if a.dest != 'help'],
+            subp._mutually_exclusive_groups)
+        formatter.add_text(subp.description)
+        for action_group in subp._action_groups:
+            formatter.start_section(action_group.title)
+            formatter.add_text(action_group.description)
+            formatter.add_arguments(a for a in action_group._group_actions
+                                    if a.dest != 'help')
+            formatter.end_section()
+        yield formatter.format_help()
+    else:  # regular argparse help
+        yield subp.format_help()
+
+# ######################## import management ############################## #
+
+try:
+    PLACDIRS = os.environ.get('PLACPATH', '.').split(':')
+except:
+    raise ValueError(_('Ill-formed PLACPATH: got %PLACPATHs') % os.environ)
+
+
+def partial_call(factory, arglist):
+    "Call a container factory with the arglist and return a plac object"
+    a = plac_core.parser_from(factory).argspec
+    if a.defaults or a.varargs or a.varkw:
+        raise TypeError('Interpreter.call must be invoked on '
+                        'factories with required arguments only')
+    required_args = ', '.join(a.args)
+    if required_args:
+        required_args += ','  # trailing comma
+    code = '''def makeobj(interact, %s *args):
+    obj = factory(%s)
+    obj._interact_ = interact
+    obj._args_ = args
+    return obj\n''' % (required_args, required_args)
+    dic = dict(factory=factory)
+    exec_(code, dic)
+    makeobj = dic['makeobj']
+    makeobj.add_help = False
+    if inspect.isclass(factory):
+        makeobj.__annotations__ = getattr(
+            factory.__init__, '__annotations__', {})
+    else:
+        makeobj.__annotations__ = getattr(
+            factory, '__annotations__', {})
+    makeobj.__annotations__['interact'] = (
+        'start interactive interpreter', 'flag', 'i')
+    return plac_core.call(makeobj, arglist)
+
+
+def import_main(path, *args):
+    """
+    A utility to import the main function of a plac tool. It also
+    works with command container factories.
+    """
+    if ':' in path:  # importing a factory
+        path, factory_name = path.split(':')
+    else:  # importing the main function
+        factory_name = None
+    if not os.path.isabs(path):  # relative path, look at PLACDIRS
+        for placdir in PLACDIRS:
+            fullpath = os.path.join(placdir, path)
+            if os.path.exists(fullpath):
+                break
+        else:  # no break
+            raise ImportError(_('Cannot find %s' % path))
+    else:
+        fullpath = path
+    name, ext = os.path.splitext(os.path.basename(fullpath))
+    module = load_source(name, fullpath)
+    if factory_name:
+        tool = partial_call(getattr(module, factory_name), args)
+    else:
+        tool = module.main
+    return tool
+
+# ############################ Task classes ############################# #
+
+
+# base class not instantiated directly
+class BaseTask(object):
+    """
+    A task is a wrapper over a generator object with signature
+    Task(no, arglist, genobj), attributes
+    .no
+    .arglist
+    .outlist
+    .str
+    .etype
+    .exc
+    .tb
+    .status
+    and methods .run and .kill.
+    """
+    STATES = ('SUBMITTED', 'RUNNING', 'TOBEKILLED',  'KILLED', 'FINISHED',
+              'ABORTED')
+
+    def __init__(self, no, arglist, genobj):
+        self.no = no
+        self.arglist = arglist
+        self._genobj = self._wrap(genobj)
+        self.str, self.etype, self.exc, self.tb = '', None, None, None
+        self.status = 'SUBMITTED'
+        self.outlist = []
+
+    def notify(self, msg):
+        "Notifies the underlying monitor. To be implemented"
+
+    def _wrap(self, genobj, stringify_tb=False):
+        """
+        Wrap the genobj into a generator managing the exceptions,
+        populating the .outlist, setting the .status and yielding None.
+        stringify_tb must be True if the traceback must be sent to a process.
+        """
+        self.status = 'RUNNING'
+        try:
+            for value in genobj:
+                if self.status == 'TOBEKILLED':  # exit from the loop
+                    raise GeneratorExit
+                if value is not None:  # add output
+                    self.outlist.append(value)
+                    self.notify(decode(value))
+                yield
+        except Interpreter.Exit:  # wanted exit
+            self._regular_exit()
+            raise
+        except (GeneratorExit, TerminatedProcess, KeyboardInterrupt):
+            # soft termination
+            self.status = 'KILLED'
+        except Exception:  # unexpected exception
+            self.etype, self.exc, tb = sys.exc_info()
+            self.tb = ''.join(traceback.format_tb(tb)) if stringify_tb else tb
+            self.status = 'ABORTED'
+        else:
+            self._regular_exit()
+
+    def _regular_exit(self):
+        self.status = 'FINISHED'
+        try:
+            self.str = '\n'.join(map(decode, self.outlist))
+        except IndexError:
+            self.str = 'no result'
+
+    def run(self):
+        "Run the inner generator"
+        for none in self._genobj:
+            pass
+
+    def kill(self):
+        "Set a TOBEKILLED status"
+        self.status = 'TOBEKILLED'
+
+    def wait(self):
+        "Wait for the task to finish: to be overridden"
+
+    @property
+    def traceback(self):
+        "Return the traceback as a (possibly empty) string"
+        if self.tb is None:
+            return ''
+        elif isinstance(self.tb, (str, bytes)):
+            return self.tb
+        else:
+            return ''.join(traceback.format_tb(self.tb))
+
+    @property
+    def result(self):
+        self.wait()
+        if self.exc:
+            if isinstance(self.tb, (str, bytes)):
+                raise self.etype(self.tb)
+            else:
+                raise_(self.etype, self.exc, self.tb or None)
+        if not self.outlist:
+            return None
+        return self.outlist[-1]
+
+    def __repr__(self):
+        "String representation containing class name, number, arglist, status"
+        return '<%s %d [%s] %s>' % (
+            self.__class__.__name__, self.no,
+            ' '.join(self.arglist), self.status)
+
+nulltask = BaseTask(0, [], ('skip' for dummy in (1,)))
+
+# ######################## synchronous tasks ############################## #
+
+
+class SynTask(BaseTask):
+    """
+    Synchronous task running in the interpreter loop and displaying its
+    output as soon as available.
+    """
+    def __str__(self):
+        "Return the output string or the error message"
+        if self.etype:  # there was an error
+            return '%s: %s' % (self.etype.__name__, self.exc)
+        else:
+            return '\n'.join(map(str, self.outlist))
+
+
+class ThreadedTask(BaseTask):
+    """
+    A task running in a separated thread.
+    """
+    def __init__(self, no, arglist, genobj):
+        BaseTask.__init__(self, no, arglist, genobj)
+        self.thread = threading.Thread(target=super(ThreadedTask, self).run)
+
+    def run(self):
+        "Run the task into a thread"
+        self.thread.start()
+
+    def wait(self):
+        "Block until the thread ends"
+        self.thread.join()
+
+
+# ######################## multiprocessing tasks ######################### #
+
+def sharedattr(name, on_error):
+    "Return a property to be attached to an MPTask"
+    def get(self):
+        try:
+            return getattr(self.ns, name)
+        except:  # the process was killed or died hard
+            return on_error
+
+    def set(self, value):
+        try:
+            setattr(self.ns, name, value)
+        except:  # the process was killed or died hard
+            pass
+    return property(get, set)
+
+
+class MPTask(BaseTask):
+    """
+    A task running as an external process. The current implementation
+    only works on Unix-like systems, where multiprocessing use forks.
+    """
+    str = sharedattr('str', '')
+    etype = sharedattr('etype', None)
+    exc = sharedattr('exc', None)
+    tb = sharedattr('tb', None)
+    status = sharedattr('status', 'ABORTED')
+
+    @property
+    def outlist(self):
+        try:
+            return self._outlist
+        except:  # the process died hard
+            return []
+
+    def notify(self, msg):
+        self.man.notify_listener(self.no, msg)
+
+    def __init__(self, no, arglist, genobj, manager):
+        """
+        The monitor has a .send method and a .man multiprocessing.Manager
+        """
+        self.no = no
+        self.arglist = arglist
+        self._genobj = self._wrap(genobj, stringify_tb=True)
+        self.man = manager
+        self._outlist = manager.mp.list()
+        self.ns = manager.mp.Namespace()
+        self.status = 'SUBMITTED'
+        self.etype, self.exc, self.tb = None, None, None
+        self.str = repr(self)
+        self.proc = multiprocessing.Process(target=super(MPTask, self).run)
+
+    def run(self):
+        "Run the task into an external process"
+        self.proc.start()
+
+    def wait(self):
+        "Block until the external process ends or is killed"
+        self.proc.join()
+
+    def kill(self):
+        """Kill the process with a SIGTERM inducing a TerminatedProcess
+        exception in the children"""
+        self.proc.terminate()
+
+# ######################## Task Manager ###################### #
+
+
+class TaskManager(object):
+    """
+    Store the given commands into a task registry. Provides methods to
+    manage the submitted tasks.
+    """
+    cmdprefix = '.'
+    specialcommands = set(['.last_tb'])
+
+    def __init__(self, obj):
+        self.obj = obj
+        self.registry = {}  # {taskno : task}
+        if obj.mpcommands or obj.thcommands:
+            self.specialcommands.update(['.kill', '.list', '.output'])
+        interact = getattr(obj, '_interact_', False)
+        self.parser = plac_core.parser_from(
+            obj, prog='' if interact else None, formatter_class=PlacFormatter)
+        HelpSummary.add(obj, self.specialcommands)
+        self.man = Manager() if obj.mpcommands else None
+        signal.signal(signal.SIGTERM, terminatedProcess)
+
+    def close(self):
+        "Kill all the running tasks"
+        for task in self.registry.values():
+            try:
+                if task.status == 'RUNNING':
+                    task.kill()
+                    task.wait()
+            except:  # task killed, nothing to wait
+                pass
+        if self.man:
+            self.man.stop()
+
+    def _get_latest(self, taskno=-1, status=None):
+        "Get the latest submitted task from the registry"
+        assert taskno < 0, 'You must pass a negative number'
+        if status:
+            tasks = [t for t in self.registry.values()
+                     if t.status == status]
+        else:
+            tasks = [t for t in self.registry.values()]
+        tasks.sort(key=attrgetter('no'))
+        if len(tasks) >= abs(taskno):
+            return tasks[taskno]
+
+    # ########################## special commands ######################## #
+
+    @plac_core.annotations(
+        taskno=('task to kill', 'positional', None, int))
+    def kill(self, taskno=-1):
+        'kill the given task (-1 to kill the latest running task)'
+        if taskno < 0:
+            task = self._get_latest(taskno, status='RUNNING')
+            if task is None:
+                yield 'Nothing to kill'
+                return
+        elif taskno not in self.registry:
+            yield 'Unknown task %d' % taskno
+            return
+        else:
+            task = self.registry[taskno]
+        if task.status in ('ABORTED', 'KILLED', 'FINISHED'):
+            yield 'Already finished %s' % task
+            return
+        task.kill()
+        yield task
+
+    @plac_core.annotations(
+        status=('', 'positional', None, str, BaseTask.STATES))
+    def list(self, status='RUNNING'):
+        'list tasks with a given status'
+        for task in self.registry.values():
+            if task.status == status:
+                yield task
+
+    @plac_core.annotations(
+        taskno=('task number', 'positional', None, int))
+    def output(self, taskno=-1, fname=None):
+        'show the output of a given task (and optionally save it to a file)'
+        if taskno < 0:
+            task = self._get_latest(taskno)
+            if task is None:
+                yield 'Nothing to show'
+                return
+        elif taskno not in self.registry:
+            yield 'Unknown task %d' % taskno
+            return
+        else:
+            task = self.registry[taskno]
+        outstr = '\n'.join(map(str, task.outlist))
+        if fname:
+            open(fname, 'w').write(outstr)
+            yield 'saved output of %d into %s' % (taskno, fname)
+            return
+        yield task
+        if len(task.outlist) > 20 and use_less:
+            less(outstr)  # has no meaning for a plac server
+        else:
+            yield outstr
+
+    @plac_core.annotations(
+        taskno=('task number', 'positional', None, int))
+    def last_tb(self, taskno=-1):
+        "show the traceback of a given task, if any"
+        task = self._get_latest(taskno)
+        if task:
+            yield task.traceback
+        else:
+            yield 'Nothing to show'
+
+# ########################## SyncProcess ############################# #
+
+
+class Process(subprocess.Popen):
+    "Start the interpreter specified by the params in a subprocess"
+
+    def __init__(self, params):
+        signal.signal(signal.SIGPIPE, signal.SIG_DFL)
+        # to avoid broken pipe messages
+        code = '''import plac, sys
+sys.argv[0] = '<%s>'
+plac.Interpreter(plac.import_main(*%s)).interact(prompt='i>\\n')
+''' % (params[0], params)
+        subprocess.Popen.__init__(
+            self, [sys.executable, '-u', '-c', code],
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+        self.man = multiprocessing.Manager()
+
+    def close(self):
+        "Close stdin and stdout"
+        self.stdin.close()
+        self.stdout.close()
+        self.man.shutdown()
+
+    def recv(self):  # char-by-char cannot work
+        "Return the output of the subprocess, line-by-line until the prompt"
+        lines = []
+        while True:
+            lines.append(self.stdout.readline())
+            if lines[-1] == 'i>\n':
+                out = ''.join(lines)
+                return out[:-1] + ' '  # remove last newline
+
+    def send(self, line):
+        """Send a line (adding a newline) to the underlying subprocess
+        and wait for the answer"""
+        self.stdin.write(line + os.linesep)
+        return self.recv()
+
+
+class StartStopObject(object):
+    started = False
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass
+
+
+class Monitor(StartStopObject):
+    """
+    Base monitor class with methods add_listener/del_listener/notify_listener
+    read_queue and and start/stop.
+    """
+    def __init__(self, name, queue=None):
+        self.name = name
+        self.queue = queue or multiprocessing.Queue()
+
+    def add_listener(self, taskno):
+        pass
+
+    def del_listener(self, taskno):
+        pass
+
+    def notify_listener(self, taskno, msg):
+        pass
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass
+
+    def read_queue(self):
+        pass
+
+
+class Manager(StartStopObject):
+    """
+    The plac Manager contains a multiprocessing.Manager and a set
+    of slave monitor processes to which we can send commands. There
+    is a manager for each interpreter with mpcommands.
+    """
+    def __init__(self):
+        self.registry = {}
+        self.started = False
+        self.mp = None
+
+    def add(self, monitor):
+        'Add or replace a monitor in the registry'
+        proc = multiprocessing.Process(None, monitor.start, monitor.name)
+        proc.queue = monitor.queue
+        self.registry[monitor.name] = proc
+
+    def delete(self, name):
+        'Remove a named monitor from the registry'
+        del self.registry[name]
+
+    # can be called more than once
+    def start(self):
+        if self.mp is None:
+            self.mp = multiprocessing.Manager()
+        for monitor in self.registry.values():
+            monitor.start()
+        self.started = True
+
+    def stop(self):
+        for monitor in self.registry.values():
+            monitor.queue.close()
+            monitor.terminate()
+        if self.mp:
+            self.mp.shutdown()
+            self.mp = None
+        self.started = False
+
+    def notify_listener(self, taskno, msg):
+        for monitor in self.registry.values():
+            monitor.queue.put(('notify_listener', taskno, msg))
+
+    def add_listener(self, no):
+        for monitor in self.registry.values():
+            monitor.queue.put(('add_listener', no))
+
+# ######################### plac server ############################# #
+
+#
+# Removed in version 1.4.0 due to incompatibility with Python 3.12
+#
+'''
+import asyncore
+import asynchat
+import socket
+
+class _AsynHandler(asynchat.async_chat):
+    "asynchat handler starting a new interpreter loop for each connection"
+
+    terminator = '\r\n'  # the standard one for telnet
+    prompt = 'i> '
+
+    def __init__(self, socket, interpreter):
+        asynchat.async_chat.__init__(self, socket)
+        self.set_terminator(self.terminator)
+        self.i = interpreter
+        self.i.__enter__()
+        self.data = []
+        self.write(self.prompt)
+
+    def write(self, data, *args):
+        "Push a string back to the client"
+        if args:
+            data %= args
+        if data.endswith('\n') and not data.endswith(self.terminator):
+            data = data[:-1] + self.terminator  # fix newlines
+        self.push(data)
+
+    def collect_incoming_data(self, data):
+        "Collect one character at the time"
+        self.data.append(data)
+
+    def found_terminator(self):
+        "Put in the queue the line received from the client"
+        line = ''.join(self.data)
+        self.log('Received line %r from %s' % (line, self.addr))
+        if line == 'EOF':
+            self.i.__exit__(None, None, None)
+            self.handle_close()
+        else:
+            task = self.i.submit(line)
+            task.run()  # synchronous or not
+            if task.etype:  # manage exception
+                error = '%s: %s\nReceived: %s' % (
+                    task.etype.__name__, task.exc, ' '.join(task.arglist))
+                self.log_info(task.traceback + error)  # on the server
+                self.write(error + self.terminator)  # back to the client
+            else:  # no exception
+                self.write(task.str + self.terminator)
+            self.data = []
+            self.write(self.prompt)
+
+
+class _AsynServer(asyncore.dispatcher):
+    "asyncore-based server spawning AsynHandlers"
+
+    def __init__(self, interpreter, newhandler, port, listen=5):
+        self.interpreter = interpreter
+        self.newhandler = newhandler
+        self.port = port
+        asyncore.dispatcher.__init__(self)
+        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.bind(('', port))
+        self.listen(listen)
+
+    def handle_accept(self):
+        clientsock, clientaddr = self.accept()
+        self.log('Connected from %s' % str(clientaddr))
+        i = self.interpreter.__class__(self.interpreter.obj)  # new interpreter
+        self.newhandler(clientsock, i)  # spawn a new handler
+
+'''
+
+# ########################## the Interpreter ############################ #
+
+class Interpreter(object):
+    """
+    A context manager with a .send method and a few utility methods:
+    execute, test and doctest.
+    """
+    class Exit(Exception):
+        pass
+
+    def __init__(self, obj, commentchar='#', split=shlex.split):
+        self.obj = obj
+        try:
+            self.name = obj.__module__
+        except AttributeError:
+            self.name = 'plac'
+        self.commentchar = commentchar
+        self.split = split
+        self._set_commands(obj)
+        self.tm = TaskManager(obj)
+        self.man = self.tm.man
+        self.parser = self.tm.parser
+        if self.commands:
+            self.parser.addsubcommands(
+                self.tm.specialcommands, self.tm, title='special commands')
+        if obj.mpcommands:
+            self.parser.addsubcommands(
+                obj.mpcommands, obj,
+                title='commands run in external processes')
+        if obj.thcommands:
+            self.parser.addsubcommands(
+                obj.thcommands, obj, title='threaded commands')
+        self.parser.error = lambda msg: sys.exit(msg)  # patch the parser
+        self._interpreter = None
+
+    def _set_commands(self, obj):
+        "Make sure obj has the right command attributes as Python sets"
+        for attrname in ('commands', 'mpcommands', 'thcommands'):
+            setattr(self, attrname, set(getattr(self.__class__, attrname, [])))
+            setattr(obj, attrname, set(getattr(obj, attrname, [])))
+        self.commands = obj.commands
+        self.mpcommands.update(obj.mpcommands)
+        self.thcommands.update(obj.thcommands)
+        if (obj.commands or obj.mpcommands or obj.thcommands) and \
+           not hasattr(obj, 'help'):  # add default help
+            obj.help = default_help.__get__(obj, obj.__class__)
+            self.commands.add('help')
+
+    def __enter__(self):
+        "Start the inner interpreter loop"
+        self._interpreter = self._make_interpreter()
+        self._interpreter.send(None)
+        return self
+
+    def __exit__(self, exctype, exc, tb):
+        "Close the inner interpreter and the task manager"
+        self.close(exctype, exc, tb)
+
+    def submit(self, line):
+        "Send a line to the underlying interpreter and return a task object"
+        if self._interpreter is None:
+            raise RuntimeError(_('%r not initialized: probably you forgot to '
+                                 'use the with statement') % self)
+        if isinstance(line, (str, bytes)):
+            arglist = self.split(line, self.commentchar)
+        else:  # expects a list of strings
+            arglist = line
+        if not arglist:
+            return nulltask
+        m = self.tm.man  # manager
+        if m and not m.started:
+            m.start()
+        task = self._interpreter.send(arglist)  # nonblocking
+        if not plac_core._match_cmd(arglist[0], self.tm.specialcommands):
+            self.tm.registry[task.no] = task
+            if m:
+                m.add_listener(task.no)
+        return task
+
+    def send(self, line):
+        """Send a line to the underlying interpreter and return
+        the finished task"""
+        task = self.submit(line)
+        BaseTask.run(task)  # blocking
+        return task
+
+    def tasks(self):
+        "The full lists of the submitted tasks"
+        return self.tm.registry.values()
+
+    def close(self, exctype=None, exc=None, tb=None):
+        "Can be called to close the interpreter prematurely"
+        self.tm.close()
+        if exctype is not None:
+            self._interpreter.throw(exctype, exc, tb)
+        else:
+            self._interpreter.close()
+
+    def _make_interpreter(self):
+        "The interpreter main loop, from lists of arguments to task objects"
+        enter = getattr(self.obj, '__enter__', lambda: None)
+        exit = getattr(self.obj, '__exit__', lambda et, ex, tb: None)
+        enter()
+        task = None
+        try:
+            for no in itertools.count(1):
+                arglist = yield task
+                try:
+                    cmd, result = self.parser.consume(arglist)
+                except SystemExit as e:  # for invalid commands
+                    if e.args == (0,):  # raised as sys.exit(0)
+                        errlist = []
+                    else:
+                        errlist = [str(e)]
+                    task = SynTask(no, arglist, iter(errlist))
+                    continue
+                except:  # anything else
+                    task = SynTask(no, arglist, gen_exc(*sys.exc_info()))
+                    continue
+                if not plac_core.iterable(result):  # atomic result
+                    task = SynTask(no, arglist, gen_val(result))
+                elif cmd in self.obj.mpcommands:
+                    task = MPTask(no, arglist, result, self.tm.man)
+                elif cmd in self.obj.thcommands:
+                    task = ThreadedTask(no, arglist, result)
+                else:  # blocking task
+                    task = SynTask(no, arglist, result)
+        except GeneratorExit:  # regular exit
+            exit(None, None, None)
+        except:  # exceptional exit
+            exit(*sys.exc_info())
+            raise
+
+    def check(self, given_input, expected_output):
+        "Make sure you get the expected_output from the given_input"
+        output = self.send(given_input).str  # blocking
+        ok = (output == expected_output)
+        if not ok:
+            # the message here is not internationalized on purpose
+            msg = 'input: %s\noutput: %s\nexpected: %s' % (
+                given_input, output, expected_output)
+            raise AssertionError(msg)
+
+    def _parse_doctest(self, lineiter):
+        "Returns the lines of input, the lines of output, and the line number"
+        lines = [line.strip() for line in lineiter]
+        inputs = []
+        positions = []
+        for i, line in enumerate(lines):
+            if line.startswith('i> '):
+                inputs.append(line[3:])
+                positions.append(i)
+        positions.append(len(lines) + 1)  # last position
+        outputs = []
+        for i, start in enumerate(positions[:-1]):
+            end = positions[i + 1]
+            outputs.append('\n'.join(lines[start+1:end]))
+        return zip(inputs, outputs, positions)
+
+    def doctest(self, lineiter, verbose=False):
+        """
+        Parse a text containing doctests in a context and tests of all them.
+        Raise an error even if a single doctest if broken. Use this for
+        sequential tests which are logically grouped.
+        """
+        with self:
+            try:
+                for input, output, no in self._parse_doctest(lineiter):
+                    if verbose:
+                        write('i> %s\n' % input)
+                        write('-> %s\n' % output)
+                    task = self.send(input)  # blocking
+                    if not str(task) == output:
+                        msg = ('line %d: input: %s\noutput: %s\nexpected: %s\n'
+                               % (no + 1, input, task, output))
+                        write(msg)
+                        if task.exc:
+                            raise_(task.etype, task.exc, task.tb)
+            except self.Exit:
+                pass
+
+    def execute(self, lineiter, verbose=False):
+        "Execute a lineiter of commands in a context and print the output"
+        with self:
+            try:
+                for line in lineiter:
+                    if verbose:
+                        write('i> ' + line)
+                    task = self.send(line)  # finished task
+                    if task.etype:  # there was an error
+                        raise_(task.etype, task.exc, task.tb)
+                    write('%s\n' % task.str)
+            except self.Exit:
+                pass
+
+    def multiline(self, stdin=sys.stdin, terminator=';', verbose=False):
+        "The multiline mode is especially suited for usage with emacs"
+        with self:
+            try:
+                for line in read_long_line(stdin, terminator):
+                    task = self.submit(line)
+                    task.run()
+                    write('%s\n' % task.str)
+                    if verbose and task.traceback:
+                        write(task.traceback)
+            except self.Exit:
+                pass
+
+    def interact(self, stdin=sys.stdin, prompt='i> ', verbose=False):
+        "Starts an interactive command loop reading commands from the console"
+        try:
+            import readline
+            readline_present = True
+        except ImportError:
+            readline_present = False
+        if stdin is sys.stdin and readline_present:  # use readline
+            histfile = os.path.expanduser('~/.%s.history' % self.name)
+            completions = list(self.commands) + list(self.mpcommands) + \
+                list(self.thcommands) + list(self.tm.specialcommands)
+            self.stdin = ReadlineInput(completions, histfile=histfile)
+        else:
+            self.stdin = stdin
+        self.prompt = prompt
+        self.verbose = verbose
+        intro = self.obj.__doc__ or ''
+        write(intro + '\n')
+        with self:
+            self.obj._interact_ = True
+            if self.stdin is sys.stdin:  # do not close stdin automatically
+                self._manage_input()
+            else:
+                with self.stdin:  # close stdin automatically
+                    self._manage_input()
+
+    def _manage_input(self):
+        "Convert input lines into task which are then executed"
+        try:
+            for line in iter(lambda: read_line(self.stdin, self.prompt), ''):
+                line = line.strip()
+                if not line:
+                    continue
+                task = self.submit(line)
+                task.run()  # synchronous or not
+                write(str(task) + '\n')
+                if self.verbose and task.etype:
+                    write(task.traceback)
+        except self.Exit:
+            pass
+
+    def start_server(self, port=2199, **kw):
+        """Starts an asyncore server reading commands for clients and opening
+        a new interpreter for each connection."""
+        _AsynServer(self, _AsynHandler, port)  # register the server
+        try:
+            asyncore.loop(**kw)
+        except (KeyboardInterrupt, TerminatedProcess):
+            pass
+        finally:
+            asyncore.close_all()
+
+    def add_monitor(self, mon):
+        self.man.add(mon)
+
+    def del_monitor(self, name):
+        self.man.delete(name)
+
+    @classmethod
+    def call(cls, factory, arglist=sys.argv[1:],
+             commentchar='#', split=shlex.split,
+             stdin=sys.stdin, prompt='i> ', verbose=False):
+        """
+        Call a container factory with the arglist and instantiate an
+        interpreter object. If there are remaining arguments, send them to the
+        interpreter, else start an interactive session.
+        """
+        obj = partial_call(factory, arglist)
+        i = cls(obj, commentchar, split)
+        if i.obj._args_:
+            with i:
+                task = i.send(i.obj._args_)  # synchronous
+                if task.exc:
+                    raise_(task.etype, task.exc, task.tb)
+                out = str(task)
+                if out:
+                    print(out)
+        elif i.obj._interact_:
+            i.interact(stdin, prompt, verbose)
+        else:
+            i.parser.print_usage()
+
+# ################################## runp ################################### #
+
+
+class _TaskLauncher(object):
+    "Helper for runp"
+
+    def __init__(self, genseq, mode):
+        if mode == 'p':
+            self.mpcommands = ['rungen']
+        else:
+            self.thcommands = ['rungen']
+        self.genlist = list(genseq)
+
+    def rungen(self, i):
+        for out in self.genlist[int(i) - 1]:
+            yield out
+
+
+def runp(genseq, mode='p'):
+    """Run a sequence of generators in parallel. Mode can be 'p' (use processes)
+    or 't' (use threads). After all of them are finished, return a list of
+    task objects.
+    """
+    assert mode in 'pt', mode
+    launcher = _TaskLauncher(genseq, mode)
+    res = []
+    with Interpreter(launcher) as inter:
+        for i in range(len(launcher.genlist)):
+            inter.submit('rungen %d' % (i + 1)).run()
+        for task in inter.tasks():
+            try:
+                res.append(task.result)
+            except Exception as e:
+                res.append(e)
+    return res

.venv/Lib/site-packages/plac_tk.py

@@ -0,0 +1,64 @@
+from __future__ import print_function
+import os
+import sys
+if sys.version_info < (3,):
+    import Queue as queue
+else:
+    import queue
+import plac_core
+from Tkinter import Tk
+from ScrolledText import ScrolledText
+from plac_ext import Monitor, TerminatedProcess
+
+
+class TkMonitor(Monitor):
+    """
+    An interface over a dictionary {taskno: scrolledtext widget}, with
+    methods add_listener, del_listener, notify_listener and start/stop.
+    """
+    def __init__(self, name, queue=None):
+        Monitor.__init__(self, name, queue)
+        self.widgets = {}
+
+    @plac_core.annotations(taskno=('task number', 'positional', None, int))
+    def add_listener(self, taskno):
+        "There is a ScrolledText for each task"
+        st = ScrolledText(self.root, height=5)
+        st.insert('end', 'Output of task %d\n' % taskno)
+        st.pack()
+        self.widgets[taskno] = st
+
+    @plac_core.annotations(taskno=('task number', 'positional', None, int))
+    def del_listener(self, taskno):
+        del self.widgets[taskno]
+
+    @plac_core.annotations(taskno=('task number', 'positional', None, int))
+    def notify_listener(self, taskno, msg):
+        w = self.widgets[taskno]
+        w.insert('end', msg + '\n')
+        w.update()
+
+    def start(self):
+        'Start the mainloop'
+        self.root = Tk()
+        self.root.title(self.name)
+        self.root.wm_protocol("WM_DELETE_WINDOW", self.stop)
+        self.root.after(0, self.read_queue)
+        try:
+            self.root.mainloop()
+        except KeyboardInterrupt:
+            print('Process %d killed by CTRL-C' % os.getpid(), file=sys.stderr)
+        except TerminatedProcess:
+            pass
+
+    def stop(self):
+        self.root.quit()
+
+    def read_queue(self):
+        try:
+            cmd_args = self.queue.get_nowait()
+        except queue.Empty:
+            pass
+        else:
+            getattr(self, cmd_args[0])(*cmd_args[1:])
+        self.root.after(100, self.read_queue)

.venv/Lib/site-packages/pylab.py

@@ -0,0 +1,3 @@
+from matplotlib.pylab import *  # noqa: F401, F403
+import matplotlib.pylab
+__doc__ = matplotlib.pylab.__doc__