[Discussion] Better way to support DeepSpeed, Apex, and more

[C1rN09]

This thread welcomes community discussions on MMEngine about integration with other training frameworks such as DeepSpeed, Apex, FairScale, etc.

Some related PRs and issues:

• [Feature Request] Support deepspeed integration #627
• [Feature] Add ApexOptimWrapper #742

[nijkah]

We have our own implementation to support DeepSpeed in MMEngine and willingness to post the PR.

One of the crucial changes to support DeepSpeed is that deepspeed inlcude deepspeed.initialize(model=model, optimizer=optimizer) which is a very similar interface as apex.amp.initialize or colossalai.initialize.

These initialize interfaces should be called before Runner.wrap_model and it is not possible with OptimWrapper.optim_context.

So we had to write the new DeepSpeedRunner to support it.
To avoid writing a new custom Runner to support these interfaces, we may have to refactor some of Runner's logic.
I think it is a quite key feature to several related frameworks such as Colossal-AI or Apex.

It would be nice to support these interfaces in MMEngine.

[nijkah]

Hi @C1rN09 , we are ready to share our DeepSpeed base code that provides some of the important features of DeepSpeed. We can start the discussion from the code.
Which way do you prefer to check the code?

[C1rN09]

Hi @C1rN09 , we are ready to share our DeepSpeed base code that provides some of the important features of DeepSpeed. We can start the discussion from the code. Which way do you prefer to check the code?

Can you create a branch on your github repo? We can first have some discussions before you post a PR ^_^

[C1rN09]

One of the crucial changes to support DeepSpeed is that deepspeed inlcude deepspeed.initialize(model=model, optimizer=optimizer) which is a very similar interface as apex.amp.initialize or colossalai.initialize.

I'm recently working on some related features and find the same issue. MMEngine is now providing low-level operations seperately on models, optimizers, etc. while some other libraries are providing high-level APIs to handle ordering and arguments issues.

I'm still wondering if we have to refactor and even bring some BC-breakings, or we can deal with it under current design. But anyway, I think a new Runner should be a good start 😸

Just notify me if you have any ideas on this!

[nijkah]

One naive idea is supporting wrap_model_and_optim instead of wrap_model. Then, some related frameworks can provide integrated a interface with ModelOptimWrapper to support *.initialize(model, optimizer).

[nijkah]

This is our code. @C1rN09
nijkah#1

It supports

• ZeRO stage 1
• ZeRO stage 2
  -- offload_optimizer
  -- overlap_comm
• ZeRO stage 3
• monolithic state_dict file for saving and loading (deepspeed saves optim_state and model_state separately and for all ranks)
  -- training with deepspeed, and resume with deepspeed
  -- training with deepspeed, and load it

cannot support

• FP 16 (there is a compatibility issue between mmengine and deepspeed)
• offload_parameters in ZeRO Stage 3
• it cannot support training with deepspeed, but resume with DDP

Need to improve it

• inject_basemodel_methods can be deleted by wrapping DeepSpeedEngine with BaseModel and do not use deepspeed.initialize.

[nijkah]

It would be better if we can separate OptimWrapper and ModelWrapper logic, but currently wrapping optimizer logic is heavily dependent on DeepSpeedEngine.

• https://github.com/microsoft/DeepSpeed/blob/211055216792cbb52ab6d355f698c194f9c55efb/deepspeed/runtime/engine.py#L327
• https://github.com/microsoft/DeepSpeed/blob/211055216792cbb52ab6d355f698c194f9c55efb/deepspeed/runtime/engine.py#L1153

[nijkah]

If you have any questions, feel free to write comments in the lined PR.
I'll try to continuously add some comments on why we had to choose this kind of implementation.

[nijkah]

Hi @C1rN09 Can we continue the progress of this feature?

[C1rN09]

Hi, @nijkah Sorry for my late response. Have just recovered from the flu... But I'm glad to share you with our plans!

The Refactoring

We are planing a refactoring on the Runner. Currently the design is still under review, but we will focus on it to make it come out ASAP. The core idea is to separate the initialization, wrapping, checkpoint loading & saving, etc. from the Runner into a new distinct module. We hope this refactoring will make the Runner more suitable for integration with other frameworks like DeepSpeed, ColossalAI, etc.

Your great work on DeepSpeed integration

A sad news is that you will have to modify your codes after we release our refactoring 😞 . We appreciate your work a lot, but currently writing a new Runner without refactoring will surely make it hard to maintain. You can still work on new features (fp16, zero-v3, etc.) in your PR if you are interested. After the release of our refactoring, we'll help you migrate to the new codebase 😄

[nijkah]

@C1rN09 Thank you for your explanation! I understand the current direction. Actually, I'd like to convey the message 'the current Runner has limitations to support DeepSpeed, ColossalAI, etc. I also agree that our code needs to be refactored further.

I'll keep my eye on the development of Runner. 👀 Thank you!