Helpful Utilities
Below are a variety of utility functions that π€ Accelerate provides, broken down by use-case.
Constants
Constants used throughout π€ Accelerate for reference
The following are constants used when utilizing Accelerator.save_state()
utils.MODEL_NAME
: "pytorch_model"
utils.OPTIMIZER_NAME
: "optimizer"
utils.RNG_STATE_NAME
: "random_states"
utils.SCALER_NAME
: "scaler.pt
utils.SCHEDULER_NAME
: "scheduler
The following are constants used when utilizing Accelerator.save_model()
utils.WEIGHTS_NAME
: "pytorch_model.bin"
utils.SAFE_WEIGHTS_NAME
: "model.safetensors"
utils.WEIGHTS_INDEX_NAME
: "pytorch_model.bin.index.json"
utils.SAFE_WEIGHTS_INDEX_NAME
: "model.safetensors.index.json"
Data Classes
These are basic dataclasses used throughout π€ Accelerate and they can be passed in as parameters.
Standalone
These are standalone dataclasses used for checks, such as the type of distributed system being used
class accelerate.utils.ComputeEnvironment
< source >( valuenames = Nonemodule = Nonequalname = Nonetype = Nonestart = 1 )
Represents a type of the compute environment.
Values:
- LOCAL_MACHINE β private/custom cluster hardware.
- AMAZON_SAGEMAKER β Amazon SageMaker as compute environment.
class accelerate.DistributedType
< source >( valuenames = Nonemodule = Nonequalname = Nonetype = Nonestart = 1 )
Represents a type of distributed environment.
Values:
- NO β Not a distributed environment, just a single process.
- MULTI_CPU β Distributed on multiple CPU nodes.
- MULTI_GPU β Distributed on multiple GPUs.
- MULTI_MLU β Distributed on multiple MLUs.
- MULTI_MUSA β Distributed on multiple MUSAs.
- MULTI_NPU β Distributed on multiple NPUs.
- MULTI_XPU β Distributed on multiple XPUs.
- DEEPSPEED β Using DeepSpeed.
- XLA β Using TorchXLA.
- TPU β This field will be deprecated in v0.27.0. Use XLA instead.
class accelerate.utils.DynamoBackend
< source >( valuenames = Nonemodule = Nonequalname = Nonetype = Nonestart = 1 )
Represents a dynamo backend (see https://pytorch.org/docs/stable/torch.compiler.html).
Values:
- NO β Do not use torch dynamo.
- EAGER β Uses PyTorch to run the extracted GraphModule. This is quite useful in debugging TorchDynamo issues.
- AOT_EAGER β Uses AotAutograd with no compiler, i.e, just using PyTorch eager for the AotAutogradβs extracted forward and backward graphs. This is useful for debugging, and unlikely to give speedups.
- INDUCTOR β Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton kernels. Read more
- AOT_TS_NVFUSER β nvFuser with AotAutograd/TorchScript. Read more
- NVPRIMS_NVFUSER β nvFuser with PrimTorch. Read more
- CUDAGRAPHS β cudagraphs with AotAutograd. Read more
- OFI β Uses Torchscript optimize_for_inference. Inference only. Read more
- FX2TRT β Uses Nvidia TensorRT for inference optimizations. Inference only. Read more
- ONNXRT β Uses ONNXRT for inference on CPU/GPU. Inference only. Read more
- TENSORRT β Uses ONNXRT to run TensorRT for inference optimizations. Read more
- AOT_TORCHXLA_TRACE_ONCE β Uses Pytorch/XLA with TorchDynamo optimization, for training. Read more
- TORCHXLA_TRACE_ONCE β Uses Pytorch/XLA with TorchDynamo optimization, for inference. Read more
- IPEX β Uses IPEX for inference on CPU. Inference only. Read more.
- TVM β Uses Apach TVM for inference optimizations. Read more
class accelerate.utils.LoggerType
< source >( valuenames = Nonemodule = Nonequalname = Nonetype = Nonestart = 1 )
Represents a type of supported experiment tracker
Values:
- ALL β all available trackers in the environment that are supported
- TENSORBOARD β TensorBoard as an experiment tracker
- WANDB β wandb as an experiment tracker
- COMETML β comet_ml as an experiment tracker
- DVCLIVE β dvclive as an experiment tracker
class accelerate.utils.PrecisionType
< source >( valuenames = Nonemodule = Nonequalname = Nonetype = Nonestart = 1 )
Represents a type of precision used on floating point values
Values:
- NO β using full precision (FP32)
- FP16 β using half precision
- BF16 β using brain floating point precision
class accelerate.utils.RNGType
< source >( valuenames = Nonemodule = Nonequalname = Nonetype = Nonestart = 1 )
An enumeration.
class accelerate.utils.SageMakerDistributedType
< source >( valuenames = Nonemodule = Nonequalname = Nonetype = Nonestart = 1 )
Represents a type of distributed environment.
Values:
- NO β Not a distributed environment, just a single process.
- DATA_PARALLEL β using sagemaker distributed data parallelism.
- MODEL_PARALLEL β using sagemaker distributed model parallelism.
Kwargs
These are configurable arguments for specific interactions throughout the PyTorch ecosystem that Accelerate handles under the hood.
Use this object in your Accelerator to customize how torch.autocast
behaves. Please refer to the
documentation of this context manager for more
information on each argument.
class accelerate.DistributedDataParallelKwargs
< source >( dim: int = 0broadcast_buffers: bool = Truebucket_cap_mb: int = 25find_unused_parameters: bool = Falsecheck_reduction: bool = Falsegradient_as_bucket_view: bool = Falsestatic_graph: bool = Falsecomm_hook: DDPCommunicationHookType = <DDPCommunicationHookType.NO: 'no'>comm_wrapper: Literal = <DDPCommunicationHookType.NO: 'no'>comm_state_option: dict = <factory> )
Use this object in your Accelerator to customize how your model is wrapped in a
torch.nn.parallel.DistributedDataParallel
. Please refer to the documentation of this
wrapper for more
information on each argument.
gradient_as_bucket_view
is only available in PyTorch 1.7.0 and later versions.
static_graph
is only available in PyTorch 1.11.0 and later versions.
class accelerate.utils.FP8RecipeKwargs
< source >( backend: Literal = 'MSAMP'opt_level: Literal = 'O2'margin: int = 0interval: int = 1fp8_format: Literal = 'E4M3'amax_history_len: int = 1amax_compute_algo: Literal = 'most_recent'override_linear_precision: Tuple = (False, False, False) )
Parameters
- backend (
str
, optional, defaults to βmsampβ) β Which FP8 engine to use. Must be one of"msamp"
(MS-AMP) or"te"
(TransformerEngine). - margin (
int
, optional, default to 0) β The margin to use for the gradient scaling. - interval (
int
, optional, default to 1) β The interval to use for how often the scaling factor is recomputed. - fp8_format (
str
, optional, default to βE4M3β) β The format to use for the FP8 recipe. Must be one ofE4M3
orHYBRID
. - amax_history_len (
int
, optional, default to 1024) β The length of the history to use for the scaling factor computation - amax_compute_algo (
str
, optional, default to βmost_recentβ) β The algorithm to use for the scaling factor computation. Must be one ofmax
ormost_recent
. - override_linear_precision (
tuple
of threebool
, optional, default to(False, False, False)
) β Whether or not to executefprop
,dgrad
, andwgrad
GEMMS in higher precision. - optimization_level (
str
), one ofO1
,O2
. (default isO2
) β What level of 8-bit collective communication should be used with MS-AMP. In general:- O1: Weight gradients and
all_reduce
communications are done in fp8, reducing GPU memory usage and communication bandwidth - O2: First-order optimizer states are in 8-bit, and second order states are in FP16. Only available when using Adam or AdamW. This maintains accuracy and can potentially save the highest memory.
- 03: Specifically for DeepSpeed, implements capabilities so weights and master weights of models
are stored in FP8. If
fp8
is selected and deepspeed is enabled, will be used by default. (Not available currently).
- O1: Weight gradients and
Use this object in your Accelerator to customize the initialization of the recipe for FP8 mixed precision
training with transformer-engine
or ms-amp
.
For more information on transformer-engine
args, please refer to the API
documentation.
For more information on the ms-amp
args, please refer to the Optimization Level
documentation.
class accelerate.GradScalerKwargs
< source >( init_scale: float = 65536.0growth_factor: float = 2.0backoff_factor: float = 0.5growth_interval: int = 2000enabled: bool = True )
Use this object in your Accelerator to customize the behavior of mixed precision, specifically how the
torch.cuda.amp.GradScaler
used is created. Please refer to the documentation of this
scaler for more information on each argument.
GradScaler
is only available in PyTorch 1.5.0 and later versions.
class accelerate.InitProcessGroupKwargs
< source >( backend: Optional = 'nccl'init_method: Optional = Nonetimeout: Optional = None )
Use this object in your Accelerator to customize the initialization of the distributed processes. Please refer to the documentation of this method for more information on each argument.
Note: If timeout
is set to None
, the default will be based upon how backend
is set.
Internal mixin that implements a to_kwargs()
method for a dataclass.
Returns a dictionary containing the attributes with values different from the default of this class.
Plugins
These are plugins that can be passed to the Accelerator object. While they are defined elsewhere in the documentation, for convenience all of them are available to see here:
class accelerate.DeepSpeedPlugin
< source >( hf_ds_config: Any = Nonegradient_accumulation_steps: int = Nonegradient_clipping: float = Nonezero_stage: int = Noneis_train_batch_min: bool = Trueoffload_optimizer_device: str = Noneoffload_param_device: str = Noneoffload_optimizer_nvme_path: str = Noneoffload_param_nvme_path: str = Nonezero3_init_flag: bool = Nonezero3_save_16bit_model: bool = Nonetransformer_moe_cls_names: str = None )
This plugin is used to integrate DeepSpeed.
deepspeed_config_process
< source >( prefix = ''mismatches = Noneconfig = Nonemust_match = True**kwargs )
Process the DeepSpeed config with the values from the kwargs.
class accelerate.FullyShardedDataParallelPlugin
< source >( sharding_strategy: typing.Any = Nonebackward_prefetch: typing.Any = Nonemixed_precision_policy: typing.Any = Noneauto_wrap_policy: Optional = Nonecpu_offload: typing.Any = Noneignored_modules: Optional = Nonestate_dict_type: typing.Any = Nonestate_dict_config: typing.Any = Noneoptim_state_dict_config: typing.Any = Nonelimit_all_gathers: bool = Trueuse_orig_params: bool = Trueparam_init_fn: Optional = Nonesync_module_states: bool = Trueforward_prefetch: bool = Falseactivation_checkpointing: bool = False )
This plugin is used to enable fully sharded data parallelism.
class accelerate.utils.GradientAccumulationPlugin
< source >( num_steps: int = Noneadjust_scheduler: bool = Truesync_with_dataloader: bool = Truesync_each_batch: bool = False )
Parameters
- num_steps (
int
) β The number of steps to accumulate gradients for. - adjust_scheduler (
bool
, optional, defaults toTrue
) β Whether to adjust the scheduler steps to account for the number of steps being accumulated. Should beTrue
if the used scheduler was not adjusted for gradient accumulation. - sync_with_dataloader (
bool
, optional, defaults toTrue
) β Whether to synchronize setting the gradients when at the end of the dataloader. - sync_each_batch (
bool
, optional) β Whether to synchronize setting the gradients at each data batch. Seting toTrue
may reduce memory requirements when using gradient accumulation with distributed training, at expense of speed.
A plugin to configure gradient accumulation behavior. You can only pass one of gradient_accumulation_plugin
or
gradient_accumulation_steps
to Accelerator. Passing both raises an error.
class accelerate.utils.MegatronLMPlugin
< source >( tp_degree: int = Nonepp_degree: int = Nonenum_micro_batches: int = Nonegradient_clipping: float = Nonesequence_parallelism: bool = Nonerecompute_activations: bool = Noneuse_distributed_optimizer: bool = Nonepipeline_model_parallel_split_rank: int = Nonenum_layers_per_virtual_pipeline_stage: int = Noneis_train_batch_min: str = Truetrain_iters: int = Nonetrain_samples: int = Noneweight_decay_incr_style: str = 'constant'start_weight_decay: float = Noneend_weight_decay: float = Nonelr_decay_style: str = 'linear'lr_decay_iters: int = Nonelr_decay_samples: int = Nonelr_warmup_iters: int = Nonelr_warmup_samples: int = Nonelr_warmup_fraction: float = Nonemin_lr: float = 0consumed_samples: List = Noneno_wd_decay_cond: Optional = Nonescale_lr_cond: Optional = Nonelr_mult: float = 1.0megatron_dataset_flag: bool = Falseseq_length: int = Noneencoder_seq_length: int = Nonedecoder_seq_length: int = Nonetensorboard_dir: str = Noneset_all_logging_options: bool = Falseeval_iters: int = 100eval_interval: int = 1000return_logits: bool = Falsecustom_train_step_class: Optional = Nonecustom_train_step_kwargs: Optional = Nonecustom_model_provider_function: Optional = Nonecustom_prepare_model_function: Optional = Nonecustom_megatron_datasets_provider_function: Optional = Nonecustom_get_batch_function: Optional = Nonecustom_loss_function: Optional = Noneother_megatron_args: Optional = None )
Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective activation recomputation and optimized fused kernels.
class accelerate.utils.TorchDynamoPlugin
< source >( backend: DynamoBackend = Nonemode: str = Nonefullgraph: bool = Nonedynamic: bool = Noneoptions: Any = Nonedisable: bool = False )
This plugin is used to compile a model with PyTorch 2.0
Configurations
These are classes which can be configured and passed through to the appropriate integration
class accelerate.utils.BnbQuantizationConfig
< source >( load_in_8bit: bool = Falsellm_int8_threshold: float = 6.0load_in_4bit: bool = Falsebnb_4bit_quant_type: str = 'fp4'bnb_4bit_use_double_quant: bool = Falsebnb_4bit_compute_dtype: bool = 'fp16'torch_dtype: dtype = Noneskip_modules: List = Nonekeep_in_fp32_modules: List = None )
A plugin to enable BitsAndBytes 4bit and 8bit quantization
class accelerate.DataLoaderConfiguration
< source >( split_batches: bool = Falsedispatch_batches: bool = Noneeven_batches: bool = Trueuse_seedable_sampler: bool = Falsenon_blocking: bool = False )
Configuration for dataloader-related items when calling accelerator.prepare
.
class accelerate.utils.ProjectConfiguration
< source >( project_dir: str = Nonelogging_dir: str = Noneautomatic_checkpoint_naming: bool = Falsetotal_limit: int = Noneiteration: int = 0save_on_each_node: bool = False )
Configuration for the Accelerator object based on inner-project needs.
Sets self.project_dir
and self.logging_dir
to the appropriate values.
Environmental Variables
These are environmental variables that can be enabled for different use cases
ACCELERATE_DEBUG_MODE
(str
): Whether to run accelerate in debug mode. More info available here.
Data Manipulation and Operations
These include data operations that mimic the same torch
ops but can be used on distributed processes.
accelerate.utils.broadcast
< source >( tensorfrom_process: int = 0 )
Recursively broadcast tensor in a nested list/tuple/dictionary of tensors to all devices.
accelerate.utils.broadcast_object_list
< source >( object_listfrom_process: int = 0 )
Broadcast a list of picklable objects form one process to the others.
accelerate.utils.concatenate
< source >( datadim = 0 )
Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.
accelerate.utils.convert_to_fp32
< source >( tensor )
Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32.
accelerate.utils.gather
< source >( tensor )
Recursively gather tensor in a nested list/tuple/dictionary of tensors from all devices.
accelerate.utils.gather_object
< source >( object: Any )
Recursively gather object in a nested list/tuple/dictionary of objects from all devices.
accelerate.utils.listify
< source >( data )
Recursively finds tensors in a nested list/tuple/dictionary and converts them to a list of numbers.
accelerate.utils.pad_across_processes
< source >( tensordim = 0pad_index = 0pad_first = False )
Parameters
- tensor (nested list/tuple/dictionary of
torch.Tensor
) β The data to gather. - dim (
int
, optional, defaults to 0) β The dimension on which to pad. - pad_index (
int
, optional, defaults to 0) β The value with which to pad. - pad_first (
bool
, optional, defaults toFalse
) β Whether to pad at the beginning or the end.
Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so they can safely be gathered.
accelerate.utils.recursively_apply
< source >( funcdata*argstest_type = <function is_torch_tensor at 0x7fb387603490>error_on_other_type = False**kwargs )
Parameters
- func (
callable
) β The function to recursively apply. - data (nested list/tuple/dictionary of
main_type
) β The data on which to applyfunc
*args β Positional arguments that will be passed tofunc
when applied on the unpacked data. - main_type (
type
, optional, defaults totorch.Tensor
) β The base type of the objects to which applyfunc
. - error_on_other_type (
bool
, optional, defaults toFalse
) β Whether to return an error or not if after unpackingdata
, we get on an object that is not of typemain_type
. IfFalse
, the function will leave objects of types different thanmain_type
unchanged. - **kwargs (additional keyword arguments, optional) β
Keyword arguments that will be passed to
func
when applied on the unpacked data.
Recursively apply a function on a data structure that is a nested list/tuple/dictionary of a given base type.
Recursively reduce the tensors in a nested list/tuple/dictionary of lists of tensors across all processes by the mean of a given operation.
accelerate.utils.send_to_device
< source >( tensordevicenon_blocking = Falseskip_keys = None )
Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
accelerate.utils.slice_tensors
< source >( datatensor_sliceprocess_index = Nonenum_processes = None )
Recursively takes a slice in a nested list/tuple/dictionary of tensors.
Environment Checks
These functionalities check the state of the current working environment including information about the operating system itself, what it can support, and if particular dependencies are installed.
Checks if bf16 is supported, optionally ignoring the TPU
Checks if torch_npu
is installed and potentially if a NPU is in the environment
accelerate.utils.is_torch_version
< source >( operation: strversion: str )
Compares the current PyTorch version to a given reference with an operation.
Check if torch_xla
is available. To train a native pytorch job in an environment with torch xla installed, set
the USE_TORCH_XLA to false.
Checks if XPU acceleration is available either via intel_extension_for_pytorch
or via stock PyTorch (>=2.4) and
potentially if a XPU is in the environment
Environment Manipulation
A context manager that will add each keyword argument passed to os.environ
and remove them when exiting.
Will convert the values in kwargs
to strings and upper-case all the keys.
A context manager that will temporarily clear environment variables.
When this context exits, the previous environment variables will be back.
accelerate.commands.config.default.write_basic_config
< source >( mixed_precision = 'no'save_location: str = '/github/home/.cache/huggingface/accelerate/default_config.yaml'use_xpu: bool = False )
Parameters
- mixed_precision (
str
, optional, defaults to βnoβ) β Mixed Precision to use. Should be one of βnoβ, βfp16β, or βbf16β - save_location (
str
, optional, defaults todefault_json_config_file
) β Optional custom save location. Should be passed to--config_file
when usingaccelerate launch
. Default location is inside the huggingface cache folder (~/.cache/huggingface
) but can be overriden by setting theHF_HOME
environmental variable, followed byaccelerate/default_config.yaml
. - use_xpu (
bool
, optional, defaults toFalse
) β Whether to use XPU if available.
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.
When setting up π€ Accelerate for the first time, rather than running accelerate config
[~utils.write_basic_config] can be used as an alternative for quick configuration.
Assigns the current process to a specific NUMA node. Ideally most efficient when having at least 2 cpus per node.
This result is cached between calls. If you want to override it, please use
accelerate.utils.environment.override_numa_afifnity
.
accelerate.utils.environment.override_numa_affinity
< source >( local_process_index: intverbose: Optional = None )
Overrides whatever NUMA affinity is set for the current process. This is very taxing and requires recalculating the
affinity to set, ideally you should use utils.environment.set_numa_affinity
instead.
Memory
accelerate.find_executable_batch_size
< source >( function: callable = Nonestarting_batch_size: int = 128 )
A basic decorator that will try to execute function
. If it fails from exceptions related to out-of-memory or
CUDNN, the batch size is cut in half and passed to function
function
must take in a batch_size
parameter as its first argument.
Modeling
These utilities relate to interacting with PyTorch models
Computes the total size of the model and its largest layer
accelerate.utils.compute_module_sizes
< source >( model: Moduledtype: Union = Nonespecial_dtypes: Optional = Nonebuffers_only: bool = False )
Compute the size of each submodule of a given model.
accelerate.utils.extract_model_from_parallel
< source >( modelkeep_fp32_wrapper: bool = Truerecursive: bool = False ) β torch.nn.Module
Parameters
- model (
torch.nn.Module
) β The model to extract. - keep_fp32_wrapper (
bool
, optional) β Whether to remove mixed precision hooks from the model. - recursive (
bool
, optional, defaults toFalse
) β Whether to recursively extract all cases ofmodule.module
frommodel
as well as unwrap child sublayers recursively, not just the top-level distributed containers.
Returns
torch.nn.Module
The extracted model.
Extract a model from its distributed containers.
accelerate.utils.get_balanced_memory
< source >( model: Modulemax_memory: Optional = Noneno_split_module_classes: Optional = Nonedtype: Union = Nonespecial_dtypes: Optional = Nonelow_zero: bool = False )
Parameters
- model (
torch.nn.Module
) β The model to analyze. - max_memory (
Dict
, optional) β A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset. Example:max_memory={0: "1GB"}
. - 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). - dtype (
str
ortorch.dtype
, optional) β If provided, the weights will be converted to that type when loaded. - special_dtypes (
Dict[str, Union[str, torch.device]]
, optional) β If provided, special dtypes to consider for some specific weights (will override dtype used as default for all weights). - low_zero (
bool
, optional) β Minimizes the number of weights on GPU 0, which is convenient when itβs used for other operations (like the Transformers generate function).
Compute a max_memory
dictionary for infer_auto_device_map() that will balance the use of each available GPU.
All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
meta device (as it would if initialized within the init_empty_weights
context manager).
accelerate.utils.get_max_layer_size
< source >( modules: Listmodule_sizes: Dictno_split_module_classes: List ) β Tuple[int, List[str]]
Parameters
- modules (
List[Tuple[str, torch.nn.Module]]
) β The list of named modules where we want to determine the maximum layer size. - module_sizes (
Dict[str, int]
) β A dictionary mapping each layer name to its size (as generated bycompute_module_sizes
). - no_split_module_classes (
List[str]
) β A list of class names for layers we donβt want to be split.
Returns
Tuple[int, List[str]]
The maximum size of a layer with the list of layer names realizing that maximum size.
Utility function that will scan a list of named modules and return the maximum size used by one full layer. The definition of a layer being:
- a module with no direct children (just parameters and buffers)
- a module whose class name is in the list
no_split_module_classes
accelerate.infer_auto_device_map
< source >( model: Modulemax_memory: Optional = Noneno_split_module_classes: Optional = Nonedtype: Union = Nonespecial_dtypes: Optional = Noneverbose: bool = Falseclean_result: bool = Trueoffload_buffers: bool = False )
Parameters
- model (
torch.nn.Module
) β The model to analyze. - max_memory (
Dict
, optional) β A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset. Example:max_memory={0: "1GB"}
. - 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). - dtype (
str
ortorch.dtype
, optional) β If provided, the weights will be converted to that type when loaded. - special_dtypes (
Dict[str, Union[str, torch.device]]
, optional) β If provided, special dtypes to consider for some specific weights (will override dtype used as default for all weights). - verbose (
bool
, optional, defaults toFalse
) β Whether or not to provide debugging statements as the function builds the device_map. - clean_result (
bool
, optional, defaults toTrue
) β Clean the resulting device_map by grouping all submodules that go on the same device together. - offload_buffers (
bool
, optional, defaults toFalse
) β 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.
Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk, such that:
- we donβt exceed the memory available of any of the GPU.
- if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that has the largest size.
- if offload to the CPU is needed,we donβt exceed the RAM available on the CPU.
- if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk that has the largest size.
All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the
meta device (as it would if initialized within the init_empty_weights
context manager).
accelerate.load_checkpoint_in_model
< source >( model: Modulecheckpoint: Uniondevice_map: Optional = Noneoffload_folder: Union = Nonedtype: Union = Noneoffload_state_dict: bool = Falseoffload_buffers: bool = Falsekeep_in_fp32_modules: List = Noneoffload_8bit_bnb: bool = Falsestrict: bool = False )
Parameters
- model (
torch.nn.Module
) β The model in which we want to load a checkpoint. - checkpoint (
str
oros.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. - a path to a folder containing a unique pytorch_model.bin or a model.safetensors file.
- 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. - offload_folder (
str
oros.PathLike
, optional) β If thedevice_map
contains any value"disk"
, the folder where we will offload weights. - dtype (
str
ortorch.dtype
, optional) β If provided, the weights will be converted to that type when loaded. - offload_state_dict (
bool
, optional, defaults toFalse
) β IfTrue
, 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. - offload_buffers (
bool
, optional, defaults toFalse
) β Whether or not to include the buffers in the weights offloaded to disk. - keep_in_fp32_modules(
List[str]
, optional) β A list of the modules that we keep intorch.float32
dtype. - offload_8bit_bnb (
bool
, optional) β Whether or not to enable offload of 8-bit modules on cpu/disk. - strict (
bool
, optional, defaults toFalse
) β Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the modelβs state_dict.
Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are loaded.
Once loaded across devices, you still need to call dispatch_model() on your model to make it able to run. To group the checkpoint loading and dispatch in one single call, use load_checkpoint_and_dispatch().
accelerate.utils.load_offloaded_weights
< source >( modelindexoffload_folder )
Loads the weights from the offload folder into the model.
accelerate.utils.load_state_dict
< source >( checkpoint_filedevice_map = None )
Parameters
- checkpoint_file (
str
) β The path to the checkpoint to load. - 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.
Load a checkpoint from a given file. If the checkpoint is in the safetensors format and a device map is passed, the weights can be fast-loaded directly on the GPU.
accelerate.utils.offload_state_dict
< source >( save_dir: Unionstate_dict: Dict )
Offload a state dict in a given folder.
accelerate.utils.retie_parameters
< source >( modeltied_params )
Reties tied parameters in a given model if the link was broken (for instance when adding hooks).
accelerate.utils.set_module_tensor_to_device
< source >( module: Moduletensor_name: strdevice: Unionvalue: Optional = Nonedtype: Union = Nonefp16_statistics: Optional = Nonetied_params_map: Optional = None )
Parameters
- module (
torch.nn.Module
) β The module in which the tensor we want to move lives. - tensor_name (
str
) β The full name of the parameter/buffer. - device (
int
,str
ortorch.device
) β The device on which to set the tensor. - value (
torch.Tensor
, optional) β The value of the tensor (useful when going from the meta device to any other device). - dtype (
torch.dtype
, optional) β If passed along the value of the parameter will be cast to thisdtype
. Otherwise,value
will be cast to the dtype of the existing parameter in the model. - fp16_statistics (
torch.HalfTensor
, optional) β The list of fp16 statistics to set on the module, used for 8 bit model serialization. - tied_params_map (Dict[int, Dict[torch.device, torch.Tensor]], optional, defaults to
None
) β A map of current 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 on the device for all others, instead of duplicating memory.
A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
param.to(device)
creates a new tensor not linked to the parameter, which is why we need this function).
accelerate.utils.shard_checkpoint
< source >( state_dict: Dictmax_shard_size: Union = '10GB'weights_name: str = 'pytorch_model.bin' )
Parameters
- state_dict (
Dict[str, torch.Tensor]
) β The state dictionary of a model to save. - max_shard_size (
int
orstr
, optional, defaults to"10GB"
) β The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit (like"5MB"
). - weights_name (
str
, optional, defaults to"pytorch_model.bin"
) β The name of the model save file.
Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a given size.
The sub-checkpoints are determined by iterating through the state_dict
in the order of its keys, so there is no
optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For example, if the
limit is 10GB and we have weights of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB],
[6+2+2GB] and not [6+2+2GB], [6+2GB], [6GB].
If one of the modelβs weight is bigger that max_sahrd_size
, it will end up in its own sub-checkpoint which will
have a size greater than max_shard_size
.
Parallel
These include general utilities that should be used when working in parallel.
accelerate.utils.extract_model_from_parallel
< source >( modelkeep_fp32_wrapper: bool = Truerecursive: bool = False ) β torch.nn.Module
Parameters
- model (
torch.nn.Module
) β The model to extract. - keep_fp32_wrapper (
bool
, optional) β Whether to remove mixed precision hooks from the model. - recursive (
bool
, optional, defaults toFalse
) β Whether to recursively extract all cases ofmodule.module
frommodel
as well as unwrap child sublayers recursively, not just the top-level distributed containers.
Returns
torch.nn.Module
The extracted model.
Extract a model from its distributed containers.
accelerate.utils.save
< source >( objfsave_on_each_node: bool = Falsesafe_serialization: bool = False )
Save the data to disk. Use in place of torch.save()
.
Introduces a blocking point in the script, making sure all processes have reached this point before continuing.
Make sure all processes will reach this instruction otherwise one of your processes will hang forever.
Random
These utilities relate to setting and synchronizing of all the random states.
accelerate.utils.set_seed
< source >( seed: intdevice_specific: bool = Falsedeterministic: bool = False )
Parameters
Helper function for reproducible behavior to set the seed in random
, numpy
, torch
.
accelerate.utils.synchronize_rng_state
< source >( rng_type: Optional = Nonegenerator: Optional = None )
PyTorch XLA
These include utilities that are useful while using PyTorch with XLA.
accelerate.utils.install_xla
< source >( upgrade: bool = False )
Helper function to install appropriate xla wheels based on the torch
version in Google Colaboratory.
Loading model weights
These include utilities that are useful to load checkpoints.
accelerate.load_checkpoint_in_model
< source >( model: Modulecheckpoint: Uniondevice_map: Optional = Noneoffload_folder: Union = Nonedtype: Union = Noneoffload_state_dict: bool = Falseoffload_buffers: bool = Falsekeep_in_fp32_modules: List = Noneoffload_8bit_bnb: bool = Falsestrict: bool = False )
Parameters
- model (
torch.nn.Module
) β The model in which we want to load a checkpoint. - checkpoint (
str
oros.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. - a path to a folder containing a unique pytorch_model.bin or a model.safetensors file.
- 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. - offload_folder (
str
oros.PathLike
, optional) β If thedevice_map
contains any value"disk"
, the folder where we will offload weights. - dtype (
str
ortorch.dtype
, optional) β If provided, the weights will be converted to that type when loaded. - offload_state_dict (
bool
, optional, defaults toFalse
) β IfTrue
, 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. - offload_buffers (
bool
, optional, defaults toFalse
) β Whether or not to include the buffers in the weights offloaded to disk. - keep_in_fp32_modules(
List[str]
, optional) β A list of the modules that we keep intorch.float32
dtype. - offload_8bit_bnb (
bool
, optional) β Whether or not to enable offload of 8-bit modules on cpu/disk. - strict (
bool
, optional, defaults toFalse
) β Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the modelβs state_dict.
Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are loaded.
Once loaded across devices, you still need to call dispatch_model() on your model to make it able to run. To group the checkpoint loading and dispatch in one single call, use load_checkpoint_and_dispatch().
Quantization
These include utilities that are useful to quantize model.
accelerate.utils.load_and_quantize_model
< source >( model: Modulebnb_quantization_config: BnbQuantizationConfigweights_location: Union = Nonedevice_map: Optional = Noneno_split_module_classes: Optional = Nonemax_memory: Optional = Noneoffload_folder: Union = Noneoffload_state_dict: bool = False ) β torch.nn.Module
Parameters
- model (
torch.nn.Module
) β Input model. The model can be already loaded or on the meta device - bnb_quantization_config (
BnbQuantizationConfig
) β The bitsandbytes quantization parameters - weights_location (
str
oros.PathLike
) β The folder weights_location 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. - a path to a folder containing a unique pytorch_model.bin file.
- 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. - 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). - max_memory (
Dict
, optional) β A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset. - offload_folder (
str
oros.PathLike
, optional) β If thedevice_map
contains any value"disk"
, the folder where we will offload weights. - offload_state_dict (
bool
, optional, defaults toFalse
) β IfTrue
, 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.
Returns
torch.nn.Module
The quantized model
This function will quantize the input model with the associated config passed in bnb_quantization_config
. If the
model is in the meta device, we will load and dispatch the weights according to the device_map
passed. If the
model is already loaded, we will quantize the model and put the model on the GPU,