Fix: forbid repeated deepspeed.initialize on training objects

[traincheck-team]

Previously closed PR:
#6848

Partially Fixes: #6772 #6771 #6770 by forbidding repeated initialization.

What are changed:

1. Marking 'model', 'optimizer' and 'lr_scheduler' in the arguments of deepspeed.initialize with the flag ds_is_inited = True.
2. Marking 'engine', 'engine.optimizer' and 'engine.lr_scheduler' in the return values of deepspeed.initialize with the flag ds_is_inited = True.
3. When calling deepspeed.initialize, raise an exception if detected ds_is_inited == True in the input model, optimizer or lr_scheduler

Expected Behavior:
Forbid repeated deepspeed.initialize invocations on model, optimizer and lr_scheduler objects.

[traincheck-team]

@microsoft-github-policy-service agree

[traincheck-team]

This fix still has interference with existing unit tests. Let me double check before we proceed.

[traincheck-team]

The unit tests in tests/unit/runtime/test_ds_initialize.py all passed.
The PR is ready for review @tjruwase.

I am not able to check other unit tests due to GPU memory constraint.

[tjruwase]

Note that optimizer could be a Callable, not an object

DeepSpeed/deepspeed/__init__.py

Line 71 in 4cd1d97

optimizer: Optional[Union[Optimizer, DeepSpeedOptimizerCallable]] = None,

[tjruwase]

Ditto for lr_scheduler

DeepSpeed/deepspeed/__init__.py

Line 74 in 4cd1d97

lr_scheduler: Optional[Union[_LRScheduler, DeepSpeedSchedulerCallable]] = None,

[tjruwase]

I think this call should be moved into `_mark_trainobjs_initialized()

[traincheck-team]

Thanks for the review @tjruwase.

• I added handling for callable types for optimizer and lr_scheduler. The handling is to only mark is_ds_inited for objects instead of callables, as the callables are not stateful and reuse should be allowed.

Regarding,

I think this call should be moved into _mark_trainobjs_initialized()

I think _assert_trainobjs_not_inited should still be separated from _mark_trainobjs_initialized as _mark_trainobjs_initialized is also called on the wrapped model and optimizers before exiting from deepspeed.initialize. The wrapped models may already have is_ds_inited being True since in DeepSpeedEngine all model flags will be passed through on the wrapper.

If we still want to keep _assert_trainobjs_not_inited inside _mark_trainobjs_initialized, we can do either of the three:

1. add a flag to _mark_trainobjs_initialized to indicate whether the *not_inited` check should be performed
2. add/check a flag using __dict__ rather than setattr/getattr
3. check whether inited use type information as well, i.e. for types of DeepSpeedEngine we directly return inited == True instead of checking for flags.