Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Modify efficient GPU training doc with now-available adamw_bnb_8bit optimizer #25807

Merged
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
22 changes: 11 additions & 11 deletions docs/source/en/perf_train_gpu_one.md
Original file line number Diff line number Diff line change
Expand Up @@ -237,10 +237,11 @@ For example if you have [NVIDIA/apex](https://github.com/NVIDIA/apex) installed,
fastest training experience among all supported AdamW optimizers.

[`Trainer`] integrates a variety of optimizers that can be used out of box: `adamw_hf`, `adamw_torch`, `adamw_torch_fused`,
`adamw_apex_fused`, `adamw_anyprecision` or `adafactor`. More optimizers can be plugged in via a third-party implementation.
`adamw_apex_fused`, `adamw_anyprecision`, `adafactor`, or `adamw_bnb_8bit`. More optimizers can be plugged in via a third-party implementation.

Let's take a closer look at two alternatives to AdamW optimizer - Adafactor (available in Trainer), and 8bit BNB quantized
optimizer (third-party implementation).
Let's take a closer look at two alternatives to AdamW optimizer:
1. `Adafactor`, available in Trainer
veezbo marked this conversation as resolved.
Show resolved Hide resolved
2. `8bit BNB quantized optimizer` which is [now](https://github.com/huggingface/transformers/issues/14819) available in Trainer, but for which a third-party integration is provided below for demonstration.
veezbo marked this conversation as resolved.
Show resolved Hide resolved

For comparison, for a 3B-parameter model, like “t5-3b”:
* A standard AdamW optimizer will need 24GB of GPU memory because it uses 8 bytes for each parameter (8*3 => 24GB)
Expand Down Expand Up @@ -269,7 +270,13 @@ Instead of aggregating optimizer states like Adafactor, 8-bit Adam keeps the ful
means that it stores the state with lower precision and dequantizes it only for the optimization. This is similar to the
idea behind mixed precision training.

To use the 8-bit optimizer, you need to install it separately and then pass it as a custom optimizer to the [`Trainer`].
To use the 8-bit optimizer, similar to Adafactor, you simply need to set `optim="adamw_bnb_8bit"` in [`TrainingArguments`]:
veezbo marked this conversation as resolved.
Show resolved Hide resolved

```py
training_args = TrainingArguments(per_device_train_batch_size=4, optim="adamw_bnb_8bit", **default_args)
```

However, for demonstration purposes, we can also use a third-party implementation of the 8-bit optimizer, to see how that can be integrated. In that case, we need to install it separately and then pass it as a custom optimizer to the [`Trainer`].
veezbo marked this conversation as resolved.
Show resolved Hide resolved

First, follow the installation guide in the GitHub [repo](https://github.com/TimDettmers/bitsandbytes) to install the `bitsandbytes` library
that implements the 8-bit Adam optimizer.
Expand Down Expand Up @@ -311,13 +318,6 @@ adam_bnb_optim = bnb.optim.Adam8bit(
)
```

<Tip>
younesbelkada marked this conversation as resolved.
Show resolved Hide resolved

To use the 8-bit optimizer with an existing pretrained model, you need to make a change to the embedding layer.
Read [this issue](https://github.com/huggingface/transformers/issues/14819) for more information.

</Tip>

Finally, pass the custom optimizer as an argument to the `Trainer`:

```py
Expand Down