Question about deepspeed/fsdp compatiblity about badam HOT 9 CLOSED

Thanks for your watch on our project. We now support deepspeed ZeRO-3 when using BAdam, please see our detailed instruction: https://github.com/Ledzy/BAdam?tab=readme-ov-file#partition-by-module-model-parallel.

from badam.

Could you provide a bit more detail, for example on how to use BAdamCallback? The examples in the Readme are not complete, are they?

from badam.

Could you provide a bit more detail, for example on how to use BAdamCallback? The examples in the Readme are not complete, are they?

Hi @paulcx , please follow the provided example:

from badam.utils import BAdamCallback

callbacks = original_callbacks.append(BAdamCallback) # add the callback
trainer = YourTrainerClass(
    ...,
    callbacks=callbacks
)

If there is no original_callbacks, just pass callbacks=[BAdamCallback] when initializing your Trainer.

from badam.

Could you provide a bit more detail, for example on how to use BAdamCallback? The examples in the Readme are not complete, are they?

Hi @paulcx , please follow the provided example:

from badam.utils import BAdamCallback

callbacks = original_callbacks.append(BAdamCallback) # add the callback
trainer = YourTrainerClass(
    ...,
    callbacks=callbacks
)

If there is no original_callbacks, just pass callbacks=[BAdamCallback] when initializing your Trainer.

My understanding is that BAdam doesn’t need to be initialized? How about the zero3 config. Is anything need to be done except these codes above?

from badam.

To use ZeRO-3, you should create a configuration file which specifies ZeRO stage and other configurations; see https://github.com/Ledzy/BAdam/blob/master/sample_ds_zero3.json for an example. You script should be launched using deepspeed or accelerate launch or other distributed launching command, depending on your circumstance.

The aforementioned part is the setup of Deepspeed and is decouple with the optimizers, e.g. BAdam. Using BAdam under ZeRO-3 is almost the same as using it in single GPU mode, except that one needs to

1. pass ds_zero3_enabled=True when creating BlockOptimizer, and
2. create a reference to deepspeed optimizer, which can be achieved by

• approach 1: simply set block_optim.ds_optimizer = ds_optimizer after deepspeed.initialize, or
• approach 2: add an additional BAdamCallback when instancing your Trainer, which actually does the same thing as approach 1; this is cleaner when using huggingface Trainer.

Hope my response is helpful for you. Feel free to reach us if you have any other questions.

from badam.

In approach 2, step 1 (pass ds_zero3_enabled=True when creating BlockOptimizer) is not required?

from badam.

It's required for the current version.

from badam.

In deepspeed+huggingface Trainer approach, the optimizer parameters are saved in ds_config.json. For example,
"optimizer": { "type": "AdamW", "params": { "lr": "auto", "betas": [ 0.9, 0.999 ], "eps": 1e-8, "weight_decay": "auto" }
Is an example of BlockOptimizer?

from badam.

BlockOptimizer inherit configurations in ds_config.json

from badam.