Fix(scheduler): WarmupLR inherits optimizer lr when not specified

[Flink-ddd]

This PR fixes issue #7303.

1. Description of the Bug

Currently, when using the WarmupLR scheduler, if warmup_max_lr is not explicitly set in the scheduler's parameters, it incorrectly falls back to its internal default value (0.001), ignoring the learning rate set in the optimizer's parameters. This can lead to unexpected training behavior and diverges from user expectations.

2. Description of the Fix

This fix modifies the __init__ method of the WarmupLR scheduler in deepspeed/runtime/lr_schedules.py.

• The default value for the warmup_max_lr argument in the function signature is changed from 0.001 to None.
• Logic is added to check if warmup_max_lr is None upon initialization. If it is, the scheduler now correctly inherits the learning rate from the optimizer's parameter groups.

This change ensures that the optimizer's learning rate is respected as the default warmup_max_lr, aligning the scheduler's behavior with the user's configuration intent.

3. Verification

The fix has been verified using a minimal reproduction script that clearly demonstrates the behavioral change.

Before Fix:
Without warmup_max_lr in the scheduler config, the learning rate incorrectly defaults to 0.001.

After Workaround (Demonstrating the Mechanism):
By explicitly adding warmup_max_lr to the scheduler config, the learning rate behaves as expected. My code change makes this the default behavior.

[Flink-ddd]

Hi @tjruwase , thanks again for the earlier approval!

After rebasing my branch onto the latest master to resolve the previous test failures, it seems a new, unrelated test test_autofill_dsconfig_from_ds_plugin is now failing with a ValueError. This seems to be related to recent changes on the master branch.

Could you please take a look when you have a moment? My changes are confined to lr_schedules.py and should not affect this config autofill logic. Thank you!

[Flink-ddd]

Hi @tjruwase , thank you so much for merging the master branch and helping to get the CI to pass! It looks like all checks have passed now. Is there anything else needed from my end for this PR to be merged?

Thank you again for your time and guidance!

[tjruwase]

@Flink-ddd, thanks for your contribution. It is greatly appreciated.

[Flink-ddd]

Hi @tjruwase , Thank you! My pleasure! It's an honor to contribute to the DeepSpeed project. I've learned a lot from the process and look forward to contributing more in the future. Thanks again for your response on time!