Exllama kernels for faster inference

del = AutoModelForCausalLM.from_pretrained("{your_username}/opt-125m-gptq", device_map="auto", quantization_config = gptq_config)
Note that only 4-bit models are supported for now. Furthermore, it is recommended to deactivate the exllama kernels if you are finetuning a quantized model with peft.
 Fine-tune a quantized model
With the official support of adapters in the Hugging Face ecosystem, you can fine-tune models that have been quantized with GPTQ.
Please have a look at peft library for more details


Example demo

Check out the Google Colab notebook to learn how to quantize your model with GPTQ and how finetune the quantized model with peft.


GPTQConfig

class transformers.GPTQConfig  < source > ( bits: int tokenizer: typing.Any = None dataset: typing.Union[str, typing.List[str], NoneType] = None group_size: int = 128 damp_percent: float = 0.1 desc_act: bool = False sym: bool = True true_sequential: bool = True use_cuda_fp16: bool = False model_seqlen: typing.Optional[int] = None block_name_to_quantize: typing.Optional[str] = None module_name_preceding_first_block: typing.Optional[typing.List[str]] = None batch_size: int = 1 pad_token_id: typing.Optional[int] = None disable_exllama: bool = False max_input_length: typing.Optional[int] = None **kwargs  )   Parameters   bits (int) —
The number of bits to quantize to, supported numbers are (2, 3, 4, 8).   tokenizer (str or PreTrainedTokenizerBase, optional) —
The tokenizer used to process the dataset. You can pass either:
A custom tokenizer object.
A string, the model id of a predefined tokenizer hosted insi


GPTQConfig

de a model repo on huggingface.co.
Valid model ids can be located at the root-level, like bert-base-uncased, or namespaced under a
user or organization name, like dbmdz/bert-base-german-cased.
A path to a directory containing vocabulary files required by the tokenizer, for instance saved
using the save_pretrained() method, e.g., ./my_model_directory/.
   dataset (Union[List[str]], optional) —
The dataset used for quantization. You can provide your own dataset in a list of string or just use the
original datasets used in GPTQ paper [‘wikitext2’,‘c4’,‘c4-new’,‘ptb’,‘ptb-new’]   group_size (int, optional, defaults to 128) —
The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization.   damp_percent (float, optional, defaults to 0.1) —
The percent of the average Hessian diagonal to use for dampening. Recommended value is 0.1.   desc_act (bool, optional, defaults to Fa


GPTQConfig

lse) —
Whether to quantize columns in order of decreasing activation size. Setting it to False can significantly
speed up inference but the perplexity may become slightly worse. Also known as act-order.   sym (bool, optional, defaults to True) —
Whether to use symetric quantization.   true_sequential (bool, optional, defaults to True) —
Whether to perform sequential quantization even within a single Transformer block. Instead of quantizing
the entire block at once, we perform layer-wise quantization. As a result, each layer undergoes
quantization using inputs that have passed through the previously quantized layers.   use_cuda_fp16 (bool, optional, defaults to False) —
Whether or not to use optimized cuda kernel for fp16 model. Need to have model in fp16.   model_seqlen (int, optional) —
The maximum sequence length that the model can take.   block_name_to_quantize (str, optional) —
The transformers block


GPTQConfig

 name to quantize.   module_name_preceding_first_block (List[str], optional) —
The layers that are preceding the first Transformer block.   batch_size (int, optional, defaults to 1) —
The batch size used when processing the dataset   pad_token_id (int, optional) —
The pad token id. Needed to prepare the dataset when batch_size > 1.   disable_exllama (bool, optional, defaults to False) —
Whether to use exllama backend. Only works with bits = 4.   max_input_length (int, optional) —
The maximum input length. This is needed to initialize a buffer that depends on the maximum expected input
length. It is specific to the exllama backend with act-order.    This is a wrapper class about all possible attributes and features that you can play with a model that has been
loaded using optimum api for gptq quantization relying on auto_gptq backend.  post_init  < source > ( )    Safety checker that arguments are correct


bitsandbytes Integration

🤗 Transformers is closely integrated with most used modules on bitsandbytes. You can load your model in 8-bit precision with few lines of code.
This is supported by most of the GPU hardwares since the 0.37.0 release of bitsandbytes.
Learn more about the quantization method in the LLM.int8() paper, or the blogpost about the collaboration.
Since its 0.39.0 release, you can load any model that supports device_map using 4-bit quantization, leveraging FP4 data type.
If you want to quantize your own pytorch model, check out this documentation from 🤗 Accelerate library.
Here are the things you can do using bitsandbytes integration


General usage

You can quantize a model by using the load_in_8bit or load_in_4bit argument when calling the from_pretrained() method as long as your model supports loading with 🤗 Accelerate and contains torch.nn.Linear layers. This should work for any modality as well.
  Copied from transformers import AutoModelForCausalLM
model_8bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", load_in_8bit=True)
model_4bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", load_in_4bit=True)
By default all other modules (e.g. torch.nn.LayerNorm) will be converted in torch.float16, but if you want to change their dtype you can overwrite the torch_dtype argument:
  Copied >>> import torch
>>> from transformers import AutoModelForCausalLM
>>> model_8bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", load_in_8bit=True, torch_dtype=torch.float32)
>>> model_8bit.model.decoder.layers[-1].final_layer_norm.weight.dtype
torch.float32


FP4 quantization

Requirements
Make sure that you have installed the requirements below before running any of the code snippets below.
Latest bitsandbytes library
pip install bitsandbytes>=0.39.0 Install latest accelerate pip install --upgrade accelerate Install latest transformers pip install --upgrade transformers
 Tips and best practices
Advanced usage: Refer to this Google Colab notebook for advanced usage of 4-bit quantization with all the possible options. Faster inference with batch_size=1 : Since the 0.40.0 release of bitsandbytes, for batch_size=1 you can benefit from fast inference. Check out these release notes and make sure to have a version that is greater than 0.40.0 to benefit from this feature out of the box. Training: According to QLoRA paper, for training 4-bit base models (e.g. using LoRA adapters) one should use bnb_4bit_quant_type='nf4'. Inference: For inference, bnb_4bit_quant_type does not


FP4 quantization

 have a huge impact on the performance. However for consistency with the model’s weights, make sure you use the same bnb_4bit_compute_dtype and torch_dtype arguments.
 Load a large model in 4bit
By using load_in_4bit=True when calling the .from_pretrained method, you can divide your memory use by 4 (roughly).
  Copied # pip install transformers accelerate bitsandbytes
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "bigscience/bloom-1b7"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
Note that once a model has been loaded in 4-bit it is currently not possible to push the quantized weights on the Hub. Note also that you cannot train 4-bit weights as this is not supported yet. However you can use 4-bit models to train extra parameters, this will be covered in the next section.


Load a large model in 8bit

You can load a model by roughly halving the memory requirements by using load_in_8bit=True argument when calling .from_pretrained method
  Copied # pip install transformers accelerate bitsandbytes
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "bigscience/bloom-1b7"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True)
Then, use your model as you would usually use a PreTrainedModel.
You can check the memory footprint of your model with get_memory_footprint method.
  Copied print(model.get_memory_footprint())
With this integration we were able to load large models on smaller devices and run them without any issue.
Note that once a model has been loaded in 8-bit it is currently not possible to push the quantized weights on the Hub except if you use the latest transformers and bitsandbytes. Note also that you cannot tr


Load a large model in 8bit

ain 8-bit weights as this is not supported yet. However you can use 8-bit models to train extra parameters, this will be covered in the next section.
Note also that device_map is optional but setting device_map = 'auto' is prefered for inference as it will dispatch efficiently the model on the available ressources.
 Advanced use cases
Here we will cover some advanced use cases you can perform with FP4 quantization
 Change the compute dtype
The compute dtype is used to change the dtype that will be used during computation. For example, hidden states could be in float32 but computation can be set to bf16 for speedups. By default, the compute dtype is set to float32.
  Copied import torch
from transformers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
 Using NF4 (Normal Float 4) data type
You can also use the NF4 data type, which is a new 4bit datatype ad


Load a large model in 8bit

apted for weights that have been initialized using a normal distribution. For that run:
  Copied from transformers import BitsAndBytesConfig
nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)
model_nf4 = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=nf4_config)
 Use nested quantization for more memory efficient inference
We also advise users to use the nested quantization technique. This saves more memory at no additional performance - from our empirical observations, this enables fine-tuning llama-13b model on an NVIDIA-T4 16GB with a sequence length of 1024, batch size of 1 and gradient accumulation steps of 4.
  Copied from transformers import BitsAndBytesConfig
double_quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)
model_double_quant = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=double_quant_config


Push quantized models on the 🤗 Hub

You can push a quantized model on the Hub by naively using push_to_hub method. This will first push the quantization configuration file, then push the quantized model weights.
Make sure to use bitsandbytes>0.37.2 (at this time of writing, we tested it on bitsandbytes==0.38.0.post1) to be able to use this feature.
  Copied from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", device_map="auto", load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
model.push_to_hub("bloom-560m-8bit")
Pushing 8bit models on the Hub is strongely encouraged for large models. This will allow the community to benefit from the memory footprint reduction and loading for example large models on a Google Colab.


Load a quantized model from the 🤗 Hub

You can load a quantized model from the Hub by using from_pretrained method. Make sure that the pushed weights are quantized, by checking that the attribute quantization_config is present in the model configuration object.
  Copied from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("{your_username}/bloom-560m-8bit", device_map="auto")
Note that in this case, you don’t need to specify the arguments load_in_8bit=True, but you need to make sure that bitsandbytes and accelerate are installed.
Note also that device_map is optional but setting device_map = 'auto' is prefered for inference as it will dispatch efficiently the model on the available ressources.


Advanced use cases

This section is intended to advanced users, that want to explore what it is possible to do beyond loading and running 8-bit models.
 Offload between cpu and gpu
One of the advanced use case of this is being able to load a model and dispatch the weights between CPU and GPU. Note that the weights that will be dispatched on CPU will not be converted in 8-bit, thus kept in float32. This feature is intended for users that want to fit a very large model and dispatch the model between GPU and CPU.
First, load a BitsAndBytesConfig from transformers and set the attribute llm_int8_enable_fp32_cpu_offload to True:
  Copied from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
Let’s say you want to load bigscience/bloom


Advanced use cases

-1b7 model, and you have just enough GPU RAM to fit the entire model except the lm_head. Therefore write a custom device_map as follows:
  Copied device_map = {
    "transformer.word_embeddings": 0,
    "transformer.word_embeddings_layernorm": 0,
    "lm_head": "cpu",
    "transformer.h": 0,
    "transformer.ln_f": 0,
}
And load your model as follows:
  Copied model_8bit = AutoModelForCausalLM.from_pretrained(
    "bigscience/bloom-1b7",
    device_map=device_map,
    quantization_config=quantization_config,
)
And that’s it! Enjoy your model!
 Play with llm_int8_threshold
You can play with the llm_int8_threshold argument to change the threshold of the outliers. An “outlier” is a hidden state value that is greater than a certain threshold.
This corresponds to the outlier threshold for outlier detection as describe


Advanced use cases

d in LLM.int8() paper. Any hidden states value that is above this threshold will be considered an outlier and the operation on those values will be done in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], but there are some exceptional systematic outliers that are very differently distributed for large models. These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6, but a lower threshold might be needed for more unstable models (small models, fine-tuning).
This argument can impact the inference speed of the model. We suggest to play with this parameter to find which one is the best for your use case.
  Copied fro


Advanced use cases

m transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
model_id = "bigscience/bloom-1b7"
quantization_config = BitsAndBytesConfig(
    llm_int8_threshold=10,
)
model_8bit = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map=device_map,
    quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
 Skip the conversion of some modules
Some models has several modules that needs to be not converted in 8-bit to ensure stability. For example Jukebox model has several lm_head modules that should be skipped. Play with llm_int8_skip_modules
  Copied from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
model_id = "bigscience/bloom-1b7"
quantization_config = BitsAndBytesConfig(
    llm_int8_skip_modules=["lm_head"],
)
mod


Advanced use cases

el_8bit = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map=device_map,
    quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
 Fine-tune a model that has been loaded in 8-bit
With the official support of adapters in the Hugging Face ecosystem, you can fine-tune models that have been loaded in 8-bit.
This enables fine-tuning large models such as flan-t5-large or facebook/opt-6.7b in a single google Colab. Please have a look at peft library for more details.
Note that you don’t need to pass device_map when loading the model for training. It will automatically load your model on your GPU. You can also set the device map to a specific device if needed (e.g. cuda:0, 0, torch.device('cuda:0')). Please note that device_map=auto should be used for inference on


BitsAndBytesConfig

class transformers.BitsAndBytesConfig  < source > ( load_in_8bit = False load_in_4bit = False llm_int8_threshold = 6.0 llm_int8_skip_modules = None llm_int8_enable_fp32_cpu_offload = False llm_int8_has_fp16_weight = False bnb_4bit_compute_dtype = None bnb_4bit_quant_type = 'fp4' bnb_4bit_use_double_quant = False **kwargs  )   Parameters   load_in_8bit (bool, optional, defaults to False) —
This flag is used to enable 8-bit quantization with LLM.int8().   load_in_4bit (bool, optional, defaults to False) —
This flag is used to enable 4-bit quantization by replacing the Linear layers with FP4/NF4 layers from
bitsandbytes.   llm_int8_threshold (float, optional, defaults to 6.0) —
This corresponds to the outlier threshold for outlier detection as described in LLM.int8() : 8-bit Matrix Multiplication for Transformers at Scale paper: https://arxiv.org/abs/22


BitsAndBytesConfig

08.07339 Any hidden states value
that is above this threshold will be considered an outlier and the operation on those values will be done
in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], but
there are some exceptional systematic outliers that are very differently distributed for large models.
These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of
magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6,
but a lower threshold might be needed for more unstable models (small models, fine-tuning).   llm_int8_skip_modules (List[str], optional) —
An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as
Jukebox that has several heads in different places and n


BitsAndBytesConfig

ot necessarily at the last position. For example
for CausalLM models, the last lm_head is kept in its original dtype.   llm_int8_enable_fp32_cpu_offload (bool, optional, defaults to False) —
This flag is used for advanced use cases and users that are aware of this feature. If you want to split
your model in different parts and run some parts in int8 on GPU and some parts in fp32 on CPU, you can use
this flag. This is useful for offloading large models such as google/flan-t5-xxl. Note that the int8
operations will not be run on CPU.   llm_int8_has_fp16_weight (bool, optional, defaults to False) —
This flag runs LLM.int8() with 16-bit main weights. This is useful for fine-tuning as the weights do not
have to be converted back and forth for the backward pass.   bnb_4bit_compute_dtype (torch.dtype or str, optional, defaults to torch.float32) —
This sets


BitsAndBytesConfig

the computational type which might be different than the input time. For example, inputs might be
fp32, but computation can be set to bf16 for speedups.   bnb_4bit_quant_type (str,  optional, defaults to "fp4") —
This sets the quantization data type in the bnb.nn.Linear4Bit layers. Options are FP4 and NF4 data types
which are specified by fp4 or nf4.   bnb_4bit_use_double_quant (bool, optional, defaults to False) —
This flag is used for nested quantization where the quantization constants from the first quantization are
quantized again.   kwargs (Dict[str, Any], optional) —
Additional parameters from which to initialize the configuration object.    This is a wrapper class about all possible attributes and features that you can play with a model that has been
loaded using bitsandbytes. This replaces load_in_8bit or load_in_4bittherefore both options a


BitsAndBytesConfig

re mutually exclusive. Currently only supports LLM.int8(), FP4, and NF4 quantization. If more methods are added to bitsandbytes,
then more arguments will be added to this class.  is_quantizable  < source > ( )    Returns True if the model is quantizable, False otherwise.  post_init  < source > ( )    Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.  quantization_method  < source > ( )    This method returns the quantization method used for the model. If the model is not quantizable, it returns
None.  to_diff_dict  < source > ( ) → Dict[str, Any]  Returns
Dict[str, Any]

Dictionary of all the attributes that make up this configuration instance,
  Removes all attributes from config which correspond to the default config attributes for better readability and
serializes to a Python dictionary


Quantization with 🤗 optimum

Please have a look at Optimum documentation to learn more about quantization methods that are supported by optimum and see if these are applicable for your use case.


Processors

Processors can mean two different things in the Transformers library:
the objects that pre-process inputs for multi-modal models such as Wav2Vec2 (speech and text)
or CLIP (text and vision) deprecated objects that were used in older versions of the library to preprocess data for GLUE or SQUAD.


Multi-modal processors

Any multi-modal model will require an object to encode or decode the data that groups several modalities (among text,
vision and audio). This is handled by objects called processors, which group together two or more processing objects
such as tokenizers (for the text modality), image processors (for vision) and feature extractors (for audio).
Those processors inherit from the following base class that implements the saving and loading functionality:
 class transformers.ProcessorMixin  < source > ( *args **kwargs  )    This is a mixin used to provide saving/loading functionality for all processor classes.  from_pretrained  < source > ( pretrained_model_name_or_path: typing.Union[str, os.PathLike] cache_dir: typing.Union[str, os.PathLike, NoneType] = None force_download: bool = False local_files_only: bool = False token: typing.Union[bool, str, NoneType] = None revision: str = 'main' **kwargs  )   Parameters   pretrained_model_name


Multi-modal processors

_or_path (str or os.PathLike) —
This can be either:
a string, the model id of a pretrained feature_extractor hosted inside a model repo on
huggingface.co. Valid model ids can be located at the root-level, like bert-base-uncased, or
namespaced under a user or organization name, like dbmdz/bert-base-german-cased.
a path to a directory containing a feature extractor file saved using the
save_pretrained() method, e.g., ./my_model_directory/.
a path or url to a saved feature extractor JSON file, e.g.,
./my_model_directory/preprocessor_config.json.
**kwargs —
Additional keyword arguments passed along to both
from_pretrained() and
~tokenization_utils_base.PreTrainedTokenizer.from_pretrained.
    Instantiate a processor associated with a pretrained model. This class method is simply calling the feature extractor
from_pretrained(), image processor
ImageProcessingMixin and the tokenizer
~tokenization_utils_base.PreTrainedTokenizer.from_pre


Multi-modal processors

trained methods. Please refer to the docstrings of the
methods above for more information.  push_to_hub  < source > ( repo_id: str use_temp_dir: typing.Optional[bool] = None commit_message: typing.Optional[str] = None private: typing.Optional[bool] = None token: typing.Union[bool, str, NoneType] = None max_shard_size: typing.Union[int, str, NoneType] = '10GB' create_pr: bool = False safe_serialization: bool = False revision: str = None **deprecated_kwargs  )   Parameters   repo_id (str) —
The name of the repository you want to push your processor to. It should contain your organization name
when pushing to a given organization.   use_temp_dir (bool, optional) —
Whether or not to use a temporary directory to store the files saved before they are pushed to the Hub.
Will default to True if there is no directory named like repo_id, False otherwise.   commit_message (str, optional) —
Message to commit while pushing. Will default to "U


Multi-modal processors

pload processor".   private (bool, optional) —
Whether or not the repository created should be private.   token (bool or str, optional) —
The token to use as HTTP bearer authorization for remote files. If True, will use the token generated
when running huggingface-cli login (stored in ~/.huggingface). Will default to True if repo_url
is not specified.   max_shard_size (int or str, optional, defaults to "10GB") —
Only applicable for models. The maximum size for a checkpoint before being sharded. Checkpoints shard
will then be each of size lower than this size. If expressed as a string, needs to be digits followed
by a unit (like "5MB").   create_pr (bool, optional, defaults to False) —
Whether or not to create a PR with the uploaded files or directly commit.   safe_serialization (bool, optional, defaults to False) —
Whether or not to convert the model weights in safetensors format for safer serialization.   revision (str, optional


Multi-modal processors

) —
Branch to push the uploaded files to.    Upload the processor files to the 🤗 Model Hub.  Examples:   Copied from transformers import AutoProcessor
processor = AutoProcessor.from_pretrained("bert-base-cased")
# Push the processor to your namespace with the name "my-finetuned-bert".
processor.push_to_hub("my-finetuned-bert")
# Push the processor to an organization with the name "my-finetuned-bert".
processor.push_to_hub("huggingface/my-finetuned-bert")  register_for_auto_class  < source > ( auto_class = 'AutoProcessor'  )   Parameters   auto_class (str or type, optional, defaults to "AutoProcessor") —
The auto class to register this new feature extractor with.    Register this class with a given auto class. This should only be used for custom feature extractors as the ones
in the library are already mapped with AutoProcessor. This API is experimental and may have some slight breaking changes in the next releases.  save_pretrain


Multi-modal processors

ed  < source > ( save_directory push_to_hub: bool = False **kwargs  )   Parameters   save_directory (str or os.PathLike) —
Directory where the feature extractor JSON file and the tokenizer files will be saved (directory will
be created if it does not exist).   push_to_hub (bool, optional, defaults to False) —
Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
repository you want to push to with repo_id (will default to the name of save_directory in your
namespace).   kwargs (Dict[str, Any], optional) —
Additional key word arguments passed along to the push_to_hub() method.    Saves the attributes of this processor (feature extractor, tokenizer…) in the specified directory so that it
can be reloaded using the from_pretrained() method. This class method is simply calling save_pretrained() and
save_pretrained(). Please refer to the docstrings of the
methods above for more information


Deprecated processors

All processors follow the same architecture which is that of the
DataProcessor. The processor returns a list of
InputExample. These
InputExample can be converted to
InputFeatures in order to be fed to the model.
 class transformers.DataProcessor  < source > ( )    Base class for data converters for sequence classification data sets.  get_dev_examples  < source > ( data_dir  )    Gets a collection of InputExample for the dev set.  get_example_from_tensor_dict  < source > ( tensor_dict  )    Gets an example from a dict with tensorflow tensors.  get_labels  < source > ( )    Gets the list of labels for this data set.  get_test_examples  < source > ( data_dir  )    Gets a collection of InputExample for the test set.  get_train_examples  < source > ( data_dir  )    Gets a collection of InputExample for the train set.  tfds_map  < source > ( example  )


Deprecated processors

Some tensorflow_datasets datasets are not formatted the same way the GLUE datasets are. This method converts
examples to the correct format.
 class transformers.InputExample  < source > ( guid: str text_a: str text_b: typing.Optional[str] = None label: typing.Optional[str] = None  )    A single training/test example for simple sequence classification.  to_json_string  < source > ( )    Serializes this instance to a JSON string.
 class transformers.InputFeatures  < source > ( input_ids: typing.List[int] attention_mask: typing.Optional[typing.List[int]] = None token_type_ids: typing.Optional[typing.List[int]] = None label: typing.Union[int, float, NoneType] = None  )    A single set of features of data. Property names are the same names as the corresponding inputs to a model.  to_json_string  < source > ( )    Serializes this instance to a JSON string.


GLUE

General Language Understanding Evaluation (GLUE) is a benchmark that evaluates the
performance of models across a diverse set of existing NLU tasks. It was released together with the paper GLUE: A
multi-task benchmark and analysis platform for natural language understanding
This library hosts a total of 10 processors for the following tasks: MRPC, MNLI, MNLI (mismatched), CoLA, SST2, STSB,
QQP, QNLI, RTE and WNLI.
Those processors are:
~data.processors.utils.MrpcProcessor ~data.processors.utils.MnliProcessor ~data.processors.utils.MnliMismatchedProcessor ~data.processors.utils.Sst2Processor ~data.processors.utils.StsbPr


GLUE

ocessor ~data.processors.utils.QqpProcessor ~data.processors.utils.QnliProcessor ~data.processors.utils.RteProcessor ~data.processors.utils.WnliProcessor
Additionally, the following method can be used to load values from a data file and convert them to a list of
InputExample.
 transformers.glue_convert_examples_to_features  < source > ( examples: typing.Union[typing.List[transformers.data.processors.utils.InputExample], ForwardRef('tf.data.Dataset')] tokenizer: PreTrainedTokenizer max_length: typing.Optional[int] = None task = None label_list = None output_mode = None  )    Loads a data file into a list of InputFeatures


XNLI

The Cross-Lingual NLI Corpus (XNLI) is a benchmark that evaluates the
quality of cross-lingual text representations. XNLI is crowd-sourced dataset based on MultiNLI: pairs of text are labeled with textual entailment annotations for 15
different languages (including both high-resource language such as English and low-resource languages such as Swahili).
It was released together with the paper XNLI: Evaluating Cross-lingual Sentence Representations
This library hosts the processor to load the XNLI data:
~data.processors.utils.XnliProcessor
Please note that since the gold labels are available on the test set, evaluation is performed on the test set.
An example using these processors is given in the run_xnli.py script.


SQuAD

The Stanford Question Answering Dataset (SQuAD) is a benchmark that
evaluates the performance of models on question answering. Two versions are available, v1.1 and v2.0. The first version
(v1.1) was released together with the paper SQuAD: 100,000+ Questions for Machine Comprehension of Text. The second version (v2.0) was released alongside the paper Know What You Don’t
Know: Unanswerable Questions for SQuAD.
This library hosts a processor for each of the two versions:


Processors

Those processors are:
~data.processors.utils.SquadV1Processor ~data.processors.utils.SquadV2Processor
They both inherit from the abstract class ~data.processors.utils.SquadProcessor
 class transformers.data.processors.squad.SquadProcessor  < source > ( )    Processor for the SQuAD data set. overridden by SquadV1Processor and SquadV2Processor, used by the version 1.1 and
version 2.0 of SQuAD, respectively.  get_dev_examples  < source > ( data_dir filename = None  )    Returns the evaluation example from the data directory.  get_examples_from_dataset  < source > ( dataset evaluate = False  )    Creates a list of SquadExample using a TFDS dataset.  Examples:   Copied >>> import tensorflow_datasets a


Processors

s tfds
>>> dataset = tfds.load("squad")
>>> training_examples = get_examples_from_dataset(dataset, evaluate=False)
>>> evaluation_examples = get_examples_from_dataset(dataset, evaluate=True)  get_train_examples  < source > ( data_dir filename = None  )    Returns the training examples from the data directory.
Additionally, the following method can be used to convert SQuAD examples into
~data.processors.utils.SquadFeatures that can be used as model inputs.
 transformers.squad_convert_examples_to_features  < source > ( examples tokenizer max_seq_length doc_stride max_query_length is_training padding_strategy = 'max_length' return_dataset = False threads = 1 tqdm_enabled = True  )    Converts a list


Processors

 of examples into a list of features that can be directly given as input to a model. It is
model-dependant and takes advantage of many of the tokenizer’s features to create the model’s inputs.  Example:   Copied processor = SquadV2Processor()
examples = processor.get_dev_examples(data_dir)
features = squad_convert_examples_to_features(
    examples=examples,
    tokenizer=tokenizer,
    max_seq_length=args.max_seq_length,
    doc_stride=args.doc_stride,
    max_query_length=args.max_query_length,
    is_training=not evaluate,
)
These processors as well as the aforementioned method can be used with files containing the data as well as with the
tensorflow_datasets package. Examples are given below.


Example usage

Here is an example using the processors as well as the conversion method using data files:
  Copied # Loading a V2 processor
processor = SquadV2Processor()
examples = processor.get_dev_examples(squad_v2_data_dir)
# Loading a V1 processor
processor = SquadV1Processor()
examples = processor.get_dev_examples(squad_v1_data_dir)
features = squad_convert_examples_to_features(
    examples=examples,
    tokenizer=tokenizer,
    max_seq_length=max_seq_length,
    doc_stride=args.doc_stride,
    max_query_length=max_query_length,
    is_training=not evalu


Example usage

ate,
)
Using tensorflow_datasets is as easy as using a data file:
  Copied # tensorflow_datasets only handle Squad V1.
tfds_examples = tfds.load("squad")
examples = SquadV1Processor().get_examples_from_dataset(tfds_examples, evaluate=evaluate)
features = squad_convert_examples_to_features(
    examples=examples,
    tokenizer=tokenizer,
    max_seq_length=max_seq_length,
    doc_stride=args.doc_stride,
    max_query_length=max_query_length,
    is_training=not evaluate,
)
Another example using these processors is given in the run_squad.py script.


Optimization

The .optimization module provides:
an optimizer with weight decay fixed that can be used to fine-tuned models, and several schedules in the form of schedule objects that inherit from _LRSchedule: a gradient accumulation class to accumulate the gradients of multiple batches


AdamW (PyTorch)

class transformers.AdamW  < source > ( params: typing.Iterable[torch.nn.parameter.Parameter] lr: float = 0.001 betas: typing.Tuple[float, float] = (0.9, 0.999) eps: float = 1e-06 weight_decay: float = 0.0 correct_bias: bool = True no_deprecation_warning: bool = False  )   Parameters   params (Iterable[nn.parameter.Parameter]) —
Iterable of parameters to optimize or dictionaries defining parameter groups.   lr (float, optional, defaults to 0.001) —
The learning rate to use.   betas (Tuple[float,float], optional, defaults to (0.9, 0.999)) —
Adam’s betas parameters (b1, b2).   eps (float, optional, defaults to 1e-06) —
Adam’s epsilon for numerical stability.   weight_d


AdamW (PyTorch)

ecay (float, optional, defaults to 0.0) —
Decoupled weight decay to apply.   correct_bias (bool, optional, defaults to True) —
Whether or not to correct bias in Adam (for instance, in Bert TF repository they use False).   no_deprecation_warning (bool, optional, defaults to False) —
A flag used to disable the deprecation warning (set to True to disable the warning).    Implements Adam algorithm with weight decay fix as introduced in Decoupled Weight Decay
Regularization.  step  < source > ( closure: typing.Callable = None  )   Parameters   closure (Callable, optional) — A closure that reevaluates the model and returns the loss.    Performs a single optimization step.


AdaFactor (PyTorch)

class transformers.Adafactor  < source > ( params lr = None eps = (1e-30, 0.001) clip_threshold = 1.0 decay_rate = -0.8 beta1 = None weight_decay = 0.0 scale_parameter = True relative_step = True warmup_init = False  )   Parameters   params (Iterable[nn.parameter.Parameter]) —
Iterable of parameters to optimize or dictionaries defining parameter groups.   lr (float, optional) —
The external learning rate.   eps (Tuple[float, float], optional, defaults to (1e-30, 0.001)) —
Regularization constants for square gradient and parameter scale respectively   clip_threshold (float, optional, defaults to 1.0) —
Threshold of root mean square of final gradient update   decay_rate (float, optional, defaults to -0.8) —
Coefficient used to compute running averages of square   beta1 (float, optional) —
Coefficient used for computing running averages of gradient   weight_decay (float


AdaFactor (PyTorch)

, optional, defaults to 0.0) —
Weight decay (L2 penalty)   scale_parameter (bool, optional, defaults to True) —
If True, learning rate is scaled by root mean square   relative_step (bool, optional, defaults to True) —
If True, time-dependent learning rate is computed instead of external learning rate   warmup_init (bool, optional, defaults to False) —
Time-dependent learning rate computation depends on whether warm-up initialization is being used    AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:
https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py Paper: Adafactor: Adaptive Learning Rates with Sublinear Memory Cost https://arxiv.org/abs/1804.04235 Note that
this optimizer internally adjusts the learning rate depending on the scale_parameter, relative_step and
warmup_init options. To use a man


AdaFactor (PyTorch)

ual (external) learning rate schedule you should set scale_parameter=False and
relative_step=False. This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested. Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3): Training without LR warmup or clip_threshold is not recommended. use scheduled LR warm-up to fixed LR use clip_threshold=1.0 (https://arxiv.org/abs/1804.04235) Disable relative updates Use scale_parameter=False Additional optimizer operations like gradient clipping should not be used alongside Adafactor  Example:   Copied Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)  Others reported the following combination to work well:   Copied Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, l


AdaFactor (PyTorch)

r=None) When using lr=None with Trainer you will most likely need to use AdafactorSchedule  scheduler as following:   Copied from transformers.optimization import Adafactor, AdafactorSchedule
optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
lr_scheduler = AdafactorSchedule(optimizer)
trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))  Usage:   Copied # replace AdamW with Adafactor
optimizer = Adafactor(
    model.parameters(),
    lr=1e-3,
    eps=(1e-30, 1e-3),
    clip_threshold=1.0,
    decay_rate=-0.8,
    beta1=None,
    weight_decay=0.0,
    relative_step=False,
    scale_parameter=False,
    warmup_init=False,
)  step  < source > ( closure = None  )   Parameters   closure (callable, optional) — A closure that reevaluates the model
and returns the loss.    Performs a single optimization ste


AdamWeightDecay (TensorFlow)

class transformers.AdamWeightDecay  < source > ( learning_rate: typing.Union[float, keras.src.optimizers.schedules.learning_rate_schedule.LearningRateSchedule] = 0.001 beta_1: float = 0.9 beta_2: float = 0.999 epsilon: float = 1e-07 amsgrad: bool = False weight_decay_rate: float = 0.0 include_in_weight_decay: typing.Optional[typing.List[str]] = None exclude_from_weight_decay: typing.Optional[typing.List[str]] = None name: str = 'AdamWeightDecay' **kwargs  )   Parameters   learning_rate (Union[float, tf.keras.optimizers.schedules.LearningRateSchedule], optional, defaults to 0.001) —
The learning rate to use or a schedule.   beta_1 (float, optional, defaults to 0.9) —
The beta1 parameter in Adam, which is the exponential decay rate for the 1st momentum estimates.   beta_2 (float, optional, defaults to 0.999) —
The beta2 parameter in Adam, which is the exponential decay rate for the 2nd momentum estimates.   epsilon (float, opt


AdamWeightDecay (TensorFlow)

ional, defaults to 1e-07) —
The epsilon parameter in Adam, which is a small constant for numerical stability.   amsgrad (bool, optional, defaults to False) —
Whether to apply AMSGrad variant of this algorithm or not, see On the Convergence of Adam and
Beyond.   weight_decay_rate (float, optional, defaults to 0.0) —
The weight decay to apply.   include_in_weight_decay (List[str], optional) —
List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is
applied to all parameters by default (unless they are in exclude_from_weight_decay).   exclude_from_weight_decay (List[str], optional) —
List of the parameter names (or re patterns) to exclude from applying weight decay to. If a
include_in_weight_decay is passed, the names in it will supersede this list.   name (str, optional, defaults to "AdamWeightDecay") —
Optional name for the operations created when applying gradients.   kwargs (


AdamWeightDecay (TensorFlow)

Dict[str, Any], optional) —
Keyword arguments. Allowed to be {clipnorm, clipvalue, lr, decay}. clipnorm is clip gradients by
norm; clipvalue is clip gradients by value, decay is included for backward compatibility to allow time
inverse decay of learning rate. lr is included for backward compatibility, recommended to use
learning_rate instead.    Adam enables L2 weight decay and clip_by_global_norm on gradients. Just adding the square of the weights to the
loss function is not the correct way of using L2 regularization/weight decay with Adam, since that will interact
with the m and v parameters in strange ways as shown in Decoupled Weight Decay
Regularization. Instead we want to decay the weights in a manner that doesn’t interact with the m/v parameters. This is equivalent
to adding the square of the weights to the loss with plain (non-momentum) SGD.  from_config  < source > ( config  )    Creates an optimizer from its config


AdamWeightDecay (TensorFlow)

 with WarmUp custom object.
 transformers.create_optimizer  < source > ( init_lr: float num_train_steps: int num_warmup_steps: int min_lr_ratio: float = 0.0 adam_beta1: float = 0.9 adam_beta2: float = 0.999 adam_epsilon: float = 1e-08 adam_clipnorm: typing.Optional[float] = None adam_global_clipnorm: typing.Optional[float] = None weight_decay_rate: float = 0.0 power: float = 1.0 include_in_weight_decay: typing.Optional[typing.List[str]] = None  )   Parameters   init_lr (float) —
The desired learning rate at the end of the warmup phase.   num_train_steps (int) —
The total number of training steps.   num_warmup_steps (int) —
The number of warmup steps.   min_lr_ratio (float, optional, defaults to 0) —
The final learning rate at the end of the linear decay will be init_lr * min_lr_ratio.   adam_beta1 (float, optional, defaults to 0.9) —
The beta1 to use in Adam.   adam_beta2 (float, optional, defaults to 0.999) —
The beta2 to u


AdamWeightDecay (TensorFlow)

se in Adam.   adam_epsilon (float, optional, defaults to 1e-8) —
The epsilon to use in Adam.   adam_clipnorm (float, optional, defaults to None) —
If not None, clip the gradient norm for each weight tensor to this value.   adam_global_clipnorm (float, optional, defaults to None) —
If not None, clip gradient norm to this value. When using this argument, the norm is computed over all
weight tensors, as if they were concatenated into a single vector.   weight_decay_rate (float, optional, defaults to 0) —
The weight decay to use.   power (float, optional, defaults to 1.0) —
The power to use for PolynomialDecay.   include_in_weight_decay (List[str], optional) —
List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is
applied to all parameters except bias and layer norm parameters.    Creates an optimizer with a learning rate schedule using a warmup phase followed by a linear decay.


Learning Rate Schedules (Pytorch)

class transformers.SchedulerType  < source > ( value names = None module = None qualname = None type = None start = 1  )    An enumeration.
 transformers.get_scheduler  < source > ( name: typing.Union[str, transformers.trainer_utils.SchedulerType] optimizer: Optimizer num_warmup_steps: typing.Optional[int] = None num_training_steps: typing.Optional[int] = None  )   Parameters   name (str or SchedulerType) —
The name of the scheduler to use.   optimizer (torch.optim.Optimizer) —
The optimizer that will be used during training.   num_warmup_steps (int, optional) —
The number of warmup steps to do. This is not required by all schedulers (hence the argument being
optional), the function will raise an error if it’s unset and the scheduler type requires it.   num_training_steps (`int“, optional) —
The number of training steps to do. This is not required by all schedulers (hence the argument being
opt


Learning Rate Schedules (Pytorch)

ional), the function will raise an error if it’s unset and the scheduler type requires it.    Unified API to get any scheduler from its name.
 transformers.get_constant_schedule  < source > ( optimizer: Optimizer last_epoch: int = -1  )   Parameters   optimizer (~torch.optim.Optimizer) —
The optimizer for which to schedule the learning rate.   last_epoch (int, optional, defaults to -1) —
The index of the last epoch when resuming training.    Create a schedule with a constant learning rate, using the learning rate set in optimizer.
 transformers.get_constant_schedule_with_warmup  < source > ( optimizer: Optimizer num_warmup_steps: int last_epoch: int = -1  )   Parameters   optimizer (~torch.optim.Optimizer) —
The optimizer for which to schedule the learning rate.   num_warmup_steps (int) —
The number of steps for the warmup phase.   last_epoch (int, optional, defaults to -1) —
The index of the l


Learning Rate Schedules (Pytorch)

ast epoch when resuming training.    Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
increases linearly between 0 and the initial lr set in the optimizer.
 transformers.get_cosine_schedule_with_warmup  < source > ( optimizer: Optimizer num_warmup_steps: int num_training_steps: int num_cycles: float = 0.5 last_epoch: int = -1  )   Parameters   optimizer (~torch.optim.Optimizer) —
The optimizer for which to schedule the learning rate.   num_warmup_steps (int) —
The number of steps for the warmup phase.   num_training_steps (int) —
The total number of training steps.   num_cycles (float, optional, defaults to 0.5) —
The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
following a half-cosine).   last_epoch (int, optional, defaults to -1) —
The index of the last epoch when resuming training.


Learning Rate Schedules (Pytorch)

  Create a schedule with a learning rate that decreases following the values of the cosine function between the
initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
initial lr set in the optimizer.
 transformers.get_cosine_with_hard_restarts_schedule_with_warmup  < source > ( optimizer: Optimizer num_warmup_steps: int num_training_steps: int num_cycles: int = 1 last_epoch: int = -1  )   Parameters   optimizer (~torch.optim.Optimizer) —
The optimizer for which to schedule the learning rate.   num_warmup_steps (int) —
The number of steps for the warmup phase.   num_training_steps (int) —
The total number of training steps.   num_cycles (int, optional, defaults to 1) —
The number of hard restarts to use.   last_epoch (int, optional, defaults to -1) —
The index of the last epoch when resuming training.    Create a schedule with a learning


Learning Rate Schedules (Pytorch)

rate that decreases following the values of the cosine function between the
initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
linearly between 0 and the initial lr set in the optimizer.
 transformers.get_linear_schedule_with_warmup  < source > ( optimizer num_warmup_steps num_training_steps last_epoch = -1  )   Parameters   optimizer (~torch.optim.Optimizer) —
The optimizer for which to schedule the learning rate.   num_warmup_steps (int) —
The number of steps for the warmup phase.   num_training_steps (int) —
The total number of training steps.   last_epoch (int, optional, defaults to -1) —
The index of the last epoch when resuming training.    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
a warmup period during which it increases linearly from 0 to the initial


Learning Rate Schedules (Pytorch)

lr set in the optimizer.
 transformers.get_polynomial_decay_schedule_with_warmup  < source > ( optimizer num_warmup_steps num_training_steps lr_end = 1e-07 power = 1.0 last_epoch = -1  )   Parameters   optimizer (~torch.optim.Optimizer) —
The optimizer for which to schedule the learning rate.   num_warmup_steps (int) —
The number of steps for the warmup phase.   num_training_steps (int) —
The total number of training steps.   lr_end (float, optional, defaults to 1e-7) —
The end LR.   power (float, optional, defaults to 1.0) —
Power factor.   last_epoch (int, optional, defaults to -1) —
The index of the last epoch when resuming training.    Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the
optimizer to end lr defined by lr_end, after a warmup period during which it increases linearly from 0 to the
initial lr set in the optimizer. Note: pow


Learning Rate Schedules (Pytorch)

er defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT
implementation at
https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
 transformers.get_inverse_sqrt_schedule  < source > ( optimizer: Optimizer num_warmup_steps: int timescale: int = None last_epoch: int = -1  )   Parameters   optimizer (~torch.optim.Optimizer) —
The optimizer for which to schedule the learning rate.   num_warmup_steps (int) —
The number of steps for the warmup phase.   timescale (int, optional, defaults to num_warmup_steps) —
Time scale.   last_epoch (int, optional, defaults to -1) —
The index of the last epoch when resuming training.    Create a schedule with an inverse square-root learning rate, from the initial lr set in the optimizer, after a
warmup period which increases lr linearly from 0 to the initial lr set in the opti


Warmup (TensorFlow)

class transformers.WarmUp  < source > ( initial_learning_rate: float decay_schedule_fn: typing.Callable warmup_steps: int power: float = 1.0 name: str = None  )   Parameters   initial_learning_rate (float) —
The initial learning rate for the schedule after the warmup (so this will be the learning rate at the end
of the warmup).   decay_schedule_fn (Callable) —
The schedule function to apply after the warmup for the rest of training.   warmup_steps (int) —
The number of steps for the warmup part of training.   power (float, optional, defaults to 1.0) —
The power to use for the polynomial warmup (defaults is a linear warmup).   name (str, optional) —
Optional name prefix for the returned tensors during the schedule.    Applies a warmup schedule on a given learning rate decay schedule.


GradientAccumulator (TensorFlow)

class transformers.GradientAccumulator  < source > ( )    Gradient accumulation utility. When used with a distribution strategy, the accumulator should be called in a
replica context. Gradients will be accumulated locally on each replica and without synchronization. Users should
then call .gradients, scale the gradients if required, and pass the result to apply_gradients.  reset  < source > ( )    Resets the accumulated gradients on the current replica.


Keras callbacks

When training a Transformers model with Keras, there are some library-specific callbacks available to automate common
tasks:


KerasMetricCallback

class transformers.KerasMetricCallback  < source > ( metric_fn: typing.Callable eval_dataset: typing.Union[tensorflow.python.data.ops.dataset_ops.DatasetV2, numpy.ndarray, tensorflow.python.framework.ops.Tensor, tuple, dict] output_cols: typing.Optional[typing.List[str]] = None label_cols: typing.Optional[typing.List[str]] = None batch_size: typing.Optional[int] = None predict_with_generate: bool = False use_xla_generation: bool = False generate_kwargs: typing.Optional[dict] = None  )   Parameters   metric_fn (Callable) —
Metric function provided by the user. It will be called with two arguments - predictions and labels.
These contain the model’s outputs and matching labels from the dataset. It should return a dict mapping
metric names to numerical values.   eval_dataset (tf.data.Dataset or dict or tuple or np.ndarray or tf.Tensor) —
Validation data to be used to gen


KerasMetricCallback

erate predictions for the metric_fn.   output_cols (`List[str], optional) —
A list of columns to be retained from the model output as the predictions. Defaults to all.   label_cols (’List[str], optional’) —
A list of columns to be retained from the input dataset as the labels. Will be autodetected if this is not
supplied.   batch_size (int, optional) —
Batch size. Only used when the data is not a pre-batched tf.data.Dataset.   predict_with_generate (bool, optional, defaults to False) —
Whether we should use model.generate() to get outputs for the model.   use_xla_generation (bool, optional, defaults to False) —
If we’re generating, whether to compile model generation with XLA. This can massively increase the speed of
generation (up to 100X speedup) but will require a new XLA compilation for each input shape. When using XLA
generation, it’s a good idea to pad your inp


KerasMetricCallback

uts to the same size, or to use the pad_to_multiple_of
argument in your tokenizer or DataCollator, which will reduce the number of unique input shapes and
save a lot of compilation time. This option has no effect is predict_with_generate is False.   generate_kwargs (dict, optional) —
Keyword arguments to pass to model.generate() when generating. Has no effect if predict_with_generate
is False.    Callback to compute metrics at the end of every epoch. Unlike normal Keras metrics, these do not need to be
compilable by TF. It is particularly useful for common NLP metrics like BLEU and ROUGE that require string
operations or generation loops that cannot be compiled. Predictions (or generations) will be computed on the
eval_dataset before being passed to the metric_fn in np.ndarray format. The metric_fn should compute
metrics and return a dict mapping metric names to metr


KerasMetricCallback

ic values. We provide an example of a suitable metric_fn that computes ROUGE scores for a summarization model below. Note that
this example skips some post-processing for readability and simplicity, and should probably not be used as-is!    Copied from datasets import load_metric
rouge_metric = load_metric("rouge")
def rouge_fn(predictions, labels):
    decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    result = rouge_metric.compute(predictions=decoded_predictions, references=decoded_labels)
    return {key: value.mid.fmeasure * 100 for key, value in result.items()}  The above function will return a dict containing values which will be logged like any other Keras metric:   Copied {'rouge1': 37.4199, 'rouge2': 13.9768, 'rougeL': 34.361, 'rougeLsum': 35.07


PushToHubCallback

class transformers.PushToHubCallback  < source > ( output_dir: typing.Union[str, pathlib.Path] save_strategy: typing.Union[str, transformers.trainer_utils.IntervalStrategy] = 'epoch' save_steps: typing.Optional[int] = None tokenizer: typing.Optional[transformers.tokenization_utils_base.PreTrainedTokenizerBase] = None hub_model_id: typing.Optional[str] = None hub_token: typing.Optional[str] = None checkpoint: bool = False **model_card_args  )   Parameters   output_dir (str) —
The output directory where the model predictions and checkpoints will be written and synced with the
repository on the Hub.   save_strategy (str or IntervalStrategy, optional, defaults to "epoch") —
The checkpoint save strategy to adopt during training. Possible values are:
"no": Save is done at the end of training.
"epoch": Save


PushToHubCallback

is done at the end of each epoch.
"steps": Save is done every save_steps
   save_steps (int, optional) —
The number of steps between saves when using the “steps” save_strategy.   tokenizer (PreTrainedTokenizerBase, optional) —
The tokenizer used by the model. If supplied, will be uploaded to the repo alongside the weights.   hub_model_id (str, optional) —
The name of the repository to keep in sync with the local output_dir. It can be a simple model ID in
which case the model will be pushed in your namespace. Otherwise it should be the whole repository name,
for instance "user_name/model", which allows you to push to an organization you are a member of with
"organization_name/model".
Will default to the name of output_dir.   hub_token (str, optional) —
The token to use to push the model to the Hub. Wil


PushToHubCallback

l default to the token in the cache folder obtained with
huggingface-cli login.   checkpoint (bool, optional, defaults to False) —
Whether to save full training checkpoints (including epoch and optimizer state) to allow training to be
resumed. Only usable when save_strategy is "epoch".    Callback that will save and push the model to the Hub regularly. By default, it pushes once per epoch, but this can
be changed with the save_strategy argument. Pushed models can be accessed like any other model on the hub, such
as with the from_pretrained method.    Copied from transformers.keras_callbacks import PushToHubCallback
push_to_hub_callback = PushToHubCallback(
    output_dir="./model_save",
    tokenizer=tokenizer,
    hub_model_id="gpt5-7xlarge",
)
model.fit(train_dataset, callbacks=[push_to_hub_callback


DeepSpeed Integration

DeepSpeed implements everything described in the ZeRO paper. Currently it provides full support for:
Optimizer state partitioning (ZeRO stage 1) Gradient partitioning (ZeRO stage 2) Parameter partitioning (ZeRO stage 3) Custom mixed precision training handling A range of fast CUDA-extension-based optimizers ZeRO-Offload to CPU and NVMe
ZeRO-Offload has its own dedicated paper: ZeRO-Offload: Democratizing Billion-Scale Model Training. And NVMe-support is described in the paper ZeRO-Infinity: Breaking the GPU
Memory Wall for Extreme Scale Deep Learning.
DeepSpeed ZeRO-2 is primarily used only for training, as its features are of no use to inference.
DeepSpeed ZeRO-3 can be used for inference as well, since it allows huge models to be loaded on multiple GPUs, which
won’t be possible on a single GPU.
🤗 Transformers integrates DeepSpeed via 2 options:
Integration of the core DeepSpeed features via Trainer. This is an everything-done-for-you type
of integration - just supp


DeepSpeed Integration

ly your custom config file or use our template and you have nothing else to do. Most of
this document is focused on this feature. If you don’t use Trainer and want to use your own Trainer where you integrated DeepSpeed
yourself, core functionality functions like from_pretrained and from_config include integration of essential
parts of DeepSpeed like zero.Init for ZeRO stage 3 and higher. To tap into this feature read the docs on
non-Trainer DeepSpeed Integration.
What is integrated:
Training:
DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 with ZeRO-Infinity (CPU and NVME offload).
Inference:
DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
it doesn’t use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
zero-inference.
There is also DeepSpeed Inference - this is a totally different technology which uses Tensor Parallelism instead of
ZeRO (coming soon)


Installation

Install the library via pypi:
  Copied pip install deepspeed
or via transformers’ extras:
  Copied pip install transformers[deepspeed]
or find more details on the DeepSpeed’s GitHub page and
advanced install.
If you’re still struggling with the build, first make sure to read CUDA Extension Installation Notes.
If you don’t prebuild the extensions and rely on them to be built at run time and you tried all of the above solutions
to no avail, the next thing to try is to pre-build the modules before installing them.
To make a local build for DeepSpeed:
  Copied git clone https://github.com/microsoft/DeepSpeed/
cd DeepSpeed
rm -rf build
TORCH_CUDA_ARCH_LIST="8.6" DS_BUILD_CPU_ADAM=1 DS_BUILD_UTILS=1 pip install . \
--global-option="build_ext" --global-


Installation

option="-j8" --no-cache -v \
--disable-pip-version-check 2>&1 | tee build.log
If you intend to use NVMe offload you will also need to include DS_BUILD_AIO=1 in the instructions above (and also
install libaio-dev system-wide).
Edit TORCH_CUDA_ARCH_LIST to insert the code for the architectures of the GPU cards you intend to use. Assuming all
your cards are the same you can get the arch via:
  Copied CUDA_VISIBLE_DEVICES=0 python -c "import torch; print(torch.cuda.get_device_capability())"
So if you get 8, 6, then use TORCH_CUDA_ARCH_LIST="8.6". If you have multiple different cards, you can list all
of them like so TORCH_CUDA_ARCH_LIST="6.1;8.6"
If you need to use the same setup on multiple machines, make a binary wheel:
  Copied git clone https://g


Installation

ithub.com/microsoft/DeepSpeed/
cd DeepSpeed
rm -rf build
TORCH_CUDA_ARCH_LIST="8.6" DS_BUILD_CPU_ADAM=1 DS_BUILD_UTILS=1 \
python setup.py build_ext -j8 bdist_wheel
it will generate something like dist/deepspeed-0.3.13+8cd046f-cp38-cp38-linux_x86_64.whl which now you can install
as pip install deepspeed-0.3.13+8cd046f-cp38-cp38-linux_x86_64.whl locally or on any other machine.
Again, remember to ensure to adjust TORCH_CUDA_ARCH_LIST to the target architectures.
You can find the complete list of NVIDIA GPUs and their corresponding Compute Capabilities (same as arch in this
context) here.
You can check the archs pytorch was built with using:
  Copied python -c "import torch; print(torch.cuda.get_arch_list())"
Here is how to find out the arch for on


Installation

e of the installed GPUs. For example, for GPU 0:
  Copied CUDA_VISIBLE_DEVICES=0 python -c "import torch; \
print(torch.cuda.get_device_properties(torch.device('cuda')))"
If the output is:
  Copied _CudaDeviceProperties(name='GeForce RTX 3090', major=8, minor=6, total_memory=24268MB, multi_processor_count=82)
then you know that this card’s arch is 8.6.
You can also leave TORCH_CUDA_ARCH_LIST out completely and then the build program will automatically query the
architecture of the GPUs the build is made on. This may or may not match the GPUs on the target machines, that’s why
it’s best to specify the desired archs explicitly.
If after trying everything suggested you still encounter build issues, please, proceed with the GitHub Issue of
Deepspeed,


Deployment with multiple GPUs

To deploy the DeepSpeed integration adjust the Trainer command line arguments to include a new argument --deepspeed ds_config.json, where ds_config.json is the DeepSpeed configuration file as
documented here. The file naming is up to you.
It’s recommended to use DeepSpeed’s add_config_arguments utility to add the necessary command line arguments to your code.
For more information please see DeepSpeed’s Argument Parsing doc.
You can use a launcher of your choice here. You can continue using the pytorch launcher:
  Copied torch.distributed.run --nproc_per_node=2 your_program.py <normal cl args> --deepspeed ds_config.json
or use the launcher provided by deepspeed:
  Copied deepspeed --num_gpus=2 your_program.py <normal cl args> --deepspeed ds_config.json
As you can see the arguments aren’t the same, but for most needs either of them works. The
full details on how to configure various nodes and GPUs can be found here.
W


Deployment with multiple GPUs

hen you use the deepspeed launcher and you want to use all available gpus you can just omit the --num_gpus flag.
Here is an example of running run_translation.py under DeepSpeed deploying all available GPUs:
  Copied deepspeed examples/pytorch/translation/run_translation.py \
--deepspeed tests/deepspeed/ds_config_zero3.json \
--model_name_or_path t5-small --per_device_train_batch_size 1 \
--output_dir output_dir --overwrite_output_dir --fp16 \
--do_train --max_train_samples 500 --num_train_epochs 1 \
--dataset_name wmt16 --dataset_config "ro-en" \
--source_lang en --target_lang ro
Note that in the DeepSpeed documentation you are likely to see --deepspeed --deepspeed_config ds_config.json - i.e.
two DeepSpeed-related arguments, but for the sake of simplicity, and since there are already so many arguments to deal
with, we combined the two into a single argument.
For some practical usage examples, please, see this post


Deployment with one GPU

To deploy DeepSpeed with one GPU adjust the Trainer command line arguments as follows:
  Copied deepspeed --num_gpus=1 examples/pytorch/translation/run_translation.py \
--deepspeed tests/deepspeed/ds_config_zero2.json \
--model_name_or_path t5-small --per_device_train_batch_size 1 \
--output_dir output_dir --overwrite_output_dir --fp16 \
--do_train --max_train_samples 500 --num_train_epochs 1 \
--dataset_name wmt16 --dataset_config "ro-en" \
--source_lang en --target_lang ro
This is almost the same as with multiple-GPUs, but here we tell DeepSpeed explicitly to use just one GPU via
--num_gpus=1. By default, DeepSpeed deploys all GPUs it can see on the given node. If you have only 1 GPU to start
with, then you don’t need this argument. The following documentation discusses the launcher options.
Why would you want to use Dee


Deployment with one GPU

pSpeed with just one GPU?
It has a ZeRO-offload feature which can delegate some computations and memory to the host’s CPU and RAM, and thus
leave more GPU resources for model’s needs - e.g. larger batch size, or enabling a fitting of a very big model which
normally won’t fit. It provides a smart GPU memory management system, that minimizes memory fragmentation, which again allows you to fit
bigger models and data batches.
While we are going to discuss the configuration in details next, the key to getting a huge improvement on a single GPU
with DeepSpeed is to have at least the following configuration in the configuration file:
  Copied {
  "zero_optimization": {
     "stage": 2,
     "offload_optimizer": {
         "device": "cpu",
         "pin_memory": true
     },
     "allgather_partitions": true,
     "allgather_bucke


Deployment with one GPU

t_size": 2e8,
     "reduce_scatter": true,
     "reduce_bucket_size": 2e8,
     "overlap_comm": true,
     "contiguous_gradients": true
  }
}
which enables optimizer offload and some other important features. You may experiment with the buffer sizes, you will
find more details in the discussion below.
For a practical usage example of this type of deployment, please, see this post.
You may also try the ZeRO-3 with CPU and NVMe offload as explained further in this document.
Notes:
if you need to run on a specific GPU, which is different from GPU 0, you can’t use CUDA_VISIBLE_DEVICES to limit
the visible scope of available GPUs. Instead, you have to use the following syntax:   Copied deepspeed --include localhost:1 examples/pytorch/translation/run_translation.py ... In this example, we tell DeepSpeed to use GPU 1 (second gpu)


Deployment with multiple Nodes

The information in this section isn’t not specific to the DeepSpeed integration and is applicable to any multi-node program. But DeepSpeed provides a deepspeed launcher that is easier to use than other launchers unless you are in a SLURM environment.
For the duration of this section let’s assume that you have 2 nodes with 8 gpus each. And you can reach the first node with ssh hostname1 and second node with ssh hostname2, and both must be able to reach each other via ssh locally without a password. Of course, you will need to rename these host (node) names to the actual host names you are working with.
 The torch.distributed.run launcher
For example, to use torch.distributed.run, you could do:
  Copied python -m torch.distributed.run --nproc_per_node=8 --nnode=2 --node_rank=0 --master_addr=hostname1 \
--master_port=9901 your_program.


Deployment with multiple Nodes

py <normal cl args> --deepspeed ds_config.json
You have to ssh to each node and run this same command on each one of them! There is no rush, the launcher will wait until both nodes will synchronize.
For more information please see torchrun. Incidentally, this is also the launcher that replaced torch.distributed.launch a few pytorch versions back.
 The deepspeed launcher
To use the deepspeed launcher instead, you have to first create a hostfile file:
  Copied hostname1 slots=8
hostname2 slots=8
and then you can launch it as:
  Copied deepspeed --num_gpus 8 --num_nodes 2 --hostfile hostfile --master_addr hostname1 --master_port=9901 \
your_program.py <normal cl args> --deepspeed ds_config.json
Unlike the torch.distributed.run launcher, deepspeed will automatically launch this command on both nodes!
For more information please see Reso


Deployment with multiple Nodes

urce Configuration (multi-node).
 Launching in a SLURM environment
In the SLURM environment the following approach can be used. The following is a slurm script launch.slurm which you will need to adapt it to your specific SLURM environment.
  Copied #SBATCH --job-name=test-nodes        # name
#SBATCH --nodes=2                    # nodes
#SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
#SBATCH --cpus-per-task=10           # number of cores per tasks
#SBATCH --gres=gpu:8                 # number of gpus
#SBATCH --time 20:00:00              # maximum execution time (HH:MM:SS)
#SBATCH --output=%x-%j.out           # output file name
export GPUS_PER_NODE=8
export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_PORT=9901
srun --jobid $SLURM_JOBID bash -c 'python -m torch.d


Deployment with multiple Nodes

istributed.run \
 --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
 --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
your_program.py <normal cl args> --deepspeed ds_config.json'
All is left is to schedule it to run:
  Copied sbatch launch.slurm
srun will take care of launching the program simultaneously on all nodes.
 Use of Non-shared filesystem
By default DeepSpeed expects that a multi-node environment uses a shared storage. If this is not the case and each node can only see the local filesystem, you need to adjust the config file to include a  checkpoint_section with the following setting:
  Copied {
  "checkpoint": {
    "use_node_local_storage": true
  }
}
Alternatively, you can also use the Trainer’s --save_on_each_node argument, and the above config will be added automatically for you.


Deployment in Notebooks

The problem with running notebook cells as a script is that there is no normal deepspeed launcher to rely on, so
under certain setups we have to emulate it.
If you’re using only 1 GPU, here is how you’d have to adjust your training code in the notebook to use DeepSpeed.
  Copied # DeepSpeed requires a distributed environment even when only one process is used.
# This emulates a launcher in the notebook
import os
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "9994"  # modify if RuntimeError: Address already in use
os.environ["RANK"] = "0"
os.environ["LOCAL_RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
# Now proceed as normal, plus pass the deepspeed config file
training_args = TrainingArguments(..., deepspeed="ds_config_zero3.json")
trainer = Trainer(...)
trainer.train()
Note: ... stands for the normal arguments that you’d pass to the functions.
If you want to


Deployment in Notebooks

use more than 1 GPU, you must use a multi-process environment for DeepSpeed to work. That is, you have
to use the launcher for that purpose and this cannot be accomplished by emulating the distributed environment presented
at the beginning of this section.
If you want to create the config file on the fly in the notebook in the current directory, you could have a dedicated
cell with:
  Copied %%bash
cat <<'EOT' > ds_config_zero3.json
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "par


Deployment in Notebooks

ams": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_bat


Deployment in Notebooks

ch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
EOT
If the training script is in a normal file and not in the notebook cells, you can launch deepspeed normally via
shell from a cell. For example, to use run_translation.py you would launch it with:
  Copied !git clone https://github.com/huggingface/transformers
!cd transformers; deepspeed examples/pytorch/translation/run_translation.py ...
or with %%bash magic, where you can write a multi-line code for the shell program to run:
  Copied %%bash
git clone https://github.com/huggingface/transformers
cd transformers
deepspeed examples/pytorch/translation/run_translation.py ...
In such case you don’t need any of the code presented at the beginning of this section.
Note: While %%bash magic is neat, but currently it buffers the output so you won’t see the logs until the process
complete


Configuration

For the complete guide to the DeepSpeed configuration options that can be used in its configuration file please refer
to the following documentation.
You can find dozens of DeepSpeed configuration examples that address various practical needs in the DeepSpeedExamples
repo:
  Copied git clone https://github.com/microsoft/DeepSpeedExamples
cd DeepSpeedExamples
find . -name '*json'
Continuing the code from above, let’s say you’re looking to configure the Lamb optimizer. So you can search through the
example .json files with:
  Copied grep -i Lamb $(find . -name '*json')
Some more examples are to be found in the main repo as well.
When using DeepSpeed you always need to supply a DeepSpeed configuration file, yet some configuration parameters have
to be configured via the command line. You will fi


Configuration

nd the nuances in the rest of this guide.
To get an idea of what DeepSpeed configuration file looks like, here is one that activates ZeRO stage 2 features,
including optimizer states cpu offload, uses AdamW optimizer and WarmupLR scheduler and will enable mixed
precision training if --fp16 is passed:
  Copied {
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr":


Configuration

"auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
}
When you execute the program, DeepSpeed will log the configuration it received from the Trainer
to the console, so you can see exactly what was the final configuration passed to it


Passing Configuration

As discussed in this document normally the DeepSpeed configuration is passed as a path to a json file, but if you’re
not using the command line interface to configure the training, and instead instantiate the
Trainer via TrainingArguments then for the deepspeed argument you can
pass a nested dict. This allows you to create the configuration on the fly and doesn’t require you to write it to
the file system before passing it to TrainingArguments.
To summarize you can do:
  Copied TrainingArguments(..., deepspeed="/path/to/ds_config.json")
or:
  Copied ds_config_dict = dict(scheduler=scheduler_params, optimizer=optimizer_params)
TrainingArguments(..., deepspeed=ds_config_dict)


Shared Configuration

This section is a must-read
Some configuration values are required by both the Trainer and DeepSpeed to function correctly,
therefore, to prevent conflicting definitions, which could lead to hard to detect errors, we chose to configure those
via the Trainer command line arguments.
Additionally, some configuration values are derived automatically based on the model’s configuration, so instead of
remembering to manually adjust multiple values, it’s the best to let the Trainer do the majority
of configuration for you.
Therefore, in the rest of this guide you will find a special configuration value: auto, which when set will be
automatically replaced with the correct or most efficient value. Please feel free to choose to ignore this
recommendation and set the values explicitly, in which case be very careful that your the
Trainer a


Shared Configuration

rguments and DeepSpeed configurations agree. For example, are you using the same
learning rate, or batch size, or gradient accumulation settings? if these mismatch the training may fail in very
difficult to detect ways. You have been warned.
There are multiple other values that are specific to DeepSpeed-only and those you will have to set manually to suit
your needs.
In your own programs, you can also use the following approach if you’d like to modify the DeepSpeed config as a master
and configure TrainingArguments based on that. The steps are:
Create or load the DeepSpeed configuration to be used as a master configuration Create the TrainingArguments object based on these values
Do note that some values, such as scheduler.params.total_num_steps are calculated by
Trainer during train, but you can of course do the math yourself


ZeRO

Zero Redundancy Optimizer (ZeRO) is the workhorse of DeepSpeed. It
supports 3 different levels (stages) of optimization. The first one is not quite interesting for scalability purposes,
therefore this document focuses on stages 2 and 3. Stage 3 is further improved by the latest addition of ZeRO-Infinity.
You will find more indepth information in the DeepSpeed documentation.
The zero_optimization section of the configuration file is the most important part (docs), since that is where you define
which ZeRO stages you want to enable and how to configure them. You will find the explanation for each parameter in the
DeepSpeed docs.
This section has to be configured exclusively via DeepSpeed configuration - the Trainer provides
no equivalent command line arguments.
Note: currently DeepSpeed doesn’t validate parameter names, so if you misspell any, it’ll use the default setting for
the parameter that got misspelled. You can watch the DeepSpeed engine start up


ZeRO

log messages to see what values it is
going to use.
 ZeRO-2 Config
The following is an example of configuration for ZeRO stage 2:
  Copied {
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "contiguous_gradients": true
    }
}
Performance tuning:
enabling offload_optimizer should reduce GPU RAM usage (it requires "stage": 2) "overlap_comm": true trades off increased GPU RAM usage to lower all-reduce latency. overlap_comm uses 4.5x
the allgather_bucket_size and reduce_bucket_size values. So if they are set to 5e8, this requires a 9GB
footprint (5e8 x 2Bytes x 2 x 4.5). Therefore, if you have a GPU with 8GB or less RAM, to avoid getting
OOM-errors you will need to reduce those pa


ZeRO

rameters to about 2e8, which would require 3.6GB. You will want to do
the same on larger capacity GPU as well, if you’re starting to hit OOM. when reducing these buffers you’re trading communication speed to avail more GPU RAM. The smaller the buffer size is,
the slower the communication gets, and the more GPU RAM will be available to other tasks. So if a bigger batch size is
important, getting a slightly slower training time could be a good trade.
Additionally, deepspeed==0.4.4 added a new option round_robin_gradients which you can enable with:
  Copied {
    "zero_optimization": {
        "round_robin_gradients": true
    }
}
This is a stage 2 optimization for CPU offloading that parallelizes gradient copying to CPU memory among ranks by fine-grained gradient partitioning. Performance benefit grows with gradient accumulation steps (more copying between optimizer steps) or GPU count (increased parallelism).
 ZeRO-3 Config
The following is an example of


ZeRO

 configuration for ZeRO stage 3:
  Copied {
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
If you are getting OOMs, because your model or activations don’t fit into the GPU memory and you have unutilized CPU
memory offloading the optimizer states and parameters to CPU memory with "device": "cpu" may solve this limitation.
If you don’t want to offload to CPU memory, use none i


ZeRO

nstead of cpu for the device entry. Offloading to
NVMe is discussed further down.
Pinned memory is enabled with pin_memory set to true. This feature can improve the throughput at the cost of
making less memory available to other processes. Pinned memory is set aside to the specific process that requested it
and its typically accessed much faster than normal CPU memory.
Performance tuning:
stage3_max_live_parameters: 1e9 stage3_max_reuse_distance: 1e9
If hitting OOM reduce stage3_max_live_parameters and stage3_max_reuse_distance. They should have minimal impact
on performance unless you are doing activation checkpointing. 1e9 would consume ~2GB. The memory is shared by
stage3_max_live_parameters and stage3_max_reuse_distance, so it’s not additive, it’s just 2GB total.
stage3_max_live_parameters is the upper limit on how many full parameters you want to keep on the GPU at any given
time. “reuse distance” is a metric we are using to figure out when will a


ZeRO

parameter be used again in the future, and we
use the stage3_max_reuse_distance to decide whether to throw away the parameter or to keep it. If a parameter is
going to be used again in near future (less than stage3_max_reuse_distance) then we keep it to reduce communication
overhead. This is super helpful when you have activation checkpointing enabled, where we do a forward recompute and
backward passes a single layer granularity and want to keep the parameter in the forward recompute till the backward
The following configuration values depend on the model’s hidden size:
reduce_bucket_size: hidden_size*hidden_size stage3_prefetch_bucket_size: 0.9 * hidden_size * hidden_size stage3_param_persistence_threshold: 10 * hidden_size
therefore set these values to auto and the Trainer will automatically assign the recommended
values. But, of course, feel free to set these explicitly as well.
stage3_gather_16bit_weights_on_model_save enables model fp16 weights co


ZeRO

nsolidation when model gets saved. With large
models and multiple GPUs this is an expensive operation both in terms of memory and speed. It’s currently required if
you plan to resume the training. Watch out for future updates that will remove this limitation and make things more
flexible.
If you’re migrating from ZeRO-2 configuration note that allgather_partitions, allgather_bucket_size and
reduce_scatter configuration parameters are not used in ZeRO-3. If you keep these in the config file they will just
be ignored.
sub_group_size: 1e9
sub_group_size controls the granularity in which parameters are updated during optimizer steps. Parameters are
grouped into buckets of sub_group_size and each buckets is updated one at a time. When used with NVMe offload in
ZeRO-Infinity, sub_group_size therefore controls the granularity in which model states are moved in and out of CPU
memory from NVMe during the optimizer step. This prevents running out of CPU memory fo


ZeRO

r extremely large models.
You can leave sub_group_size to its default value of 1e9 when not using NVMe offload. You may want to change its
default value in the following cases:
Running into OOM during optimizer step: Reduce sub_group_size to reduce memory utilization of temporary buffers Optimizer Step is taking a long time: Increase sub_group_size to improve bandwidth utilization as a result of
the increased data buffers.
 ZeRO-0 Config
Note that we’re listing Stage 0 and 1 last since they are rarely used.
Stage 0 is disabling all types of sharding and just using DeepSpeed as DDP. You can turn it on with:
  Copied {
    "zero_optimization": {
        "stage": 0
    }
}
This will essentially disable ZeRO without you needing to change anything else.
 ZeRO-1 Config
Stage 1 is Stage 2 minus gradient sharding. You can always try it to speed things a tiny bit to only shard the optimizer states with:
  Copied {
    "zero_optimization": {
        "stage": 1


NVMe Support

ZeRO-Infinity allows for training incredibly large models by extending GPU and CPU memory with NVMe memory. Thanks to
smart partitioning and tiling algorithms each GPU needs to send and receive very small amounts of data during
offloading so modern NVMe proved to be fit to allow for an even larger total memory pool available to your training
process. ZeRO-Infinity requires ZeRO-3 enabled.
The following configuration example enables NVMe to offload both optimizer states and the params:
  Copied {
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/local_nvme",
            "pin_memory": true,
            "buffer_count": 4,
            "fast_init": false
        },
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/local_nvme",
            "pin_memory": true,
            "buffer_count": 5,
            "buffer_size": 1e8,
            "max_in_cpu": 1e9


NVMe Support

},
        "aio": {
            "block_size": 262144,
            "queue_depth": 32,
            "thread_count": 1,
            "single_submit": false,
            "overlap_events": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
}
You can choose to offload both optimizer states and params to NVMe, or just one of them or none. For example, if you
have copious amounts of CPU memory available, by all means offload to CPU memory only as it’d be faster (hint:
“device”: “cpu”).
Here is the full documentation for offloading optimizer states and parameters.
Make sure that your nvme_path is actually an NVMe, since it will work w


NVMe Support

ith the normal hard drive or SSD, but it’ll
be much much slower. The fast scalable training was designed with modern NVMe transfer speeds in mind (as of this
writing one can have ~3.5GB/s read, ~3GB/s write peak speeds).
In order to figure out the optimal aio configuration block you must run a benchmark on your target setup, as
explained here.
 ZeRO-2 vs ZeRO-3 Performance
ZeRO-3 is likely to be slower than ZeRO-2 if everything else is configured the same because the former has to gather
model weights in addition to what ZeRO-2 does. If ZeRO-2 meets your needs and you don’t need to scale beyond a few GPUs
then you may choose to stick to it. It’s important to understand that ZeRO-3 enables a much higher scalability capacity
at a cost of speed.
It’s possible to adjust ZeRO-3 configuration to make it perform closer to ZeRO-2:
set stage3_param_persistence_threshold to a very large number - larger than the largest parameter, e.g., 6 * hidden_size * hidden_size. This will kee


NVMe Support

p the parameters on the GPUs. turn off offload_params since ZeRO-2 doesn’t have that option.
The performance will likely improve significantly with just offload_params turned off, even if you don’t change
stage3_param_persistence_threshold. Of course, these changes will impact the size of the model you can train. So
these help you to trade scalability for speed depending on your needs.
 ZeRO-2 Example
Here is a full ZeRO-2 auto-configuration file ds_config_zero2.json:
  Copied {
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",


NVMe Support

          "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
Here is a full ZeRO-2 all-enabled manually set configuration file. It is here mainly for you to see what the typical
values look like, but we highly recommend using the one with multiple auto settings in it.
  Copied {
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_windo


NVMe Support

w": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 3e-5,
            "betas": [0.8, 0.999],
            "eps": 1e-8,
            "weight_decay": 3e-7
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 3e-5,
            "warmup_num_steps": 500
        }
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "steps_per_print": 2000,
    "wall_clock_breakdown": false
}
 ZeRO-3 Example
Here is a full ZeRO-3 auto-configuration file d


NVMe Support

s_config_zero3.json:
  Copied {
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce


NVMe Support

_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
Here is a full ZeRO-3 all-enabled manually set configuration file. It is here mainly for you to see what the typical
values look like, but we highly recommend using the one with multiple auto settings in it.
  Copied {
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "l


NVMe Support

r": 3e-5,
            "betas": [0.8, 0.999],
            "eps": 1e-8,
            "weight_decay": 3e-7
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 3e-5,
            "warmup_num_steps": 500
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": 1e6,
        "stage3_prefetch_bucket_size": 0.94e6,
        "stage3_param_persistence_threshold": 1e4,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "steps_per_print": 2000,
    "wall_clock_bre


NVMe Support

akdown": false
}
 How to Choose Which ZeRO Stage and Offloads To Use For Best Performance
So now you know there are all these different stages. How to decide which of them to use? This section will attempt to address this question.
In general the following applies:
Speed-wise (left is faster than right)
Stage 0 (DDP) > Stage 1 > Stage 2 > Stage 2 + offload > Stage 3 > Stage 3 + offloads
GPU Memory usage-wise (right is more GPU memory efficient than left)
Stage 0 (DDP) < Stage 1 < Stage 2 < Stage 2 + offload < Stage 3 < Stage 3 + offloads
So when you want to get the fastest execution while fitting into minimal number of GPUs, here is the process you could follow. We start with the fastest approach and if running into GPU OOM we then go to the next slower approach, but which will use less GPU memory. And so on and so forth.
First of all set batch size to 1 (you can always use gradient accumulation for any desired effective batch size).
Enable --gradient_checkpointing 1 (H


NVMe Support

F Trainer) or directly model.gradient_checkpointing_enable() - if OOM then Try ZeRO stage 2 first. if OOM then Try ZeRO stage 2 + offload_optimizer - if OOM then Switch to ZeRO stage 3 - if OOM then Enable offload_param to cpu - if OOM then Enable offload_optimizer to cpu - if OOM then If you still can’t fit a batch size of 1 first check various default values and lower them if you can. For example, if you use generate and you don’t use a wide search beam make it narrower as it’d take a lot of memory. Definitely use mixed half-precision over fp32 - so bf16 on Ampere and higher GPUs and fp16 on older gpu architectures. If you still OOM you could add more hardware or enable ZeRO-Infinity - that is switch offloads offload_param and  offload_optimizer to nvme. You need to make sure it’s a very fast nvme. As an anecdote I was able to infer BLOOM-176B on a tiny GPU using ZeRO-Infinity except it was extremely slow. But it worked!
You can, of course, work through these steps in