TypeError: DeepSpeedZeroOptimizer.check_overflow() got an unexpected keyword argument 'external'

I tried to enable the '--align_overflow' option during step3 training (with zero stage 2), which was updated in PR #624 . However, I encountered the following error during gradient updates:

File "/mapping-data/rlhf_training/step3_rlhf_finetuning/ppo_trainer.py", line 322, in train_rlhf

actor_overflow = self.actor_model.optimizer.check_overflow(

TypeError: DeepSpeedZeroOptimizer.check_overflow() got an unexpected keyword argument 'external'

I checked the recent versions of [deepspeed/runtime/zero/stage_1_and_2.py] (https://github.com/microsoft/DeepSpeed/blob/v0.10.0/deepspeed/runtime/zero/stage_1_and_2.py#L1904) and did not find the 'external' parameter except for 'partition_gradients'. I am wandering which version of deepspeed I need to use to apply this option?

Also, I noticed that the aforementioned function does not have a return value and I have to modify the code to: 'actor_overflow = self.actor_model.optimizer.overflow' to get actor_overflow value.

If I could get some help, I would be very grateful.

Start at rlhf_training/step3_rlhf_finetuning/ppo_trainer.py around line 322 and compare the check_overflow call with deepspeed/runtime/zero/stage_1_and_2.py in v0.10.0. Reproduce step3 training with --align_overflow and verify which supported DeepSpeed version or API makes the call work and provides the actor overflow value without a TypeError.