DeepSpeed integration breaks existing DeepSpeed logic

[muellerzr]

What's the issue, what's expected?:

There are attributes inside of regular deepspeed.runtime that are missing in this repo, and the monkey-patch doesn't cover, such as:

from deepspeed.runtime.lr_schedules import VALID_LR_SCHEDULES

How to reproduce it?:

from msamp import deepspeed
deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES

Log message or shapshot?:

Additional information:

Still trying to figure out a good solution here to make sure old namespaces can remain intact (as calling normal deepspeed later has consequences during training)

[muellerzr]

Generally I think this repo might want to recommend users do something like:

# Order doesn't really matter
import deepspeed
from msamp import deepspeed as msamp_deepspeed

Since otherwise they may face issues like this.

Not per-se a problem, but just something users should know about

[muellerzr]

For completeness, here's what my solution wound up being in the end:

import deepspeed
ds_initialize = deepspeed.initialize
if use_msamp:
    # MS-AMP requires DeepSpeed patches
    from msamp import deepspeed as msamp_deepspeed
    ds_initialize = msamp_deepspeed.initialize