Lightning creates two DeepSpeedEngine instances for the same model

[HeyangQin]

Bug description

Hello Lightning team!

We got serval user reports (e.g. deepspeedai/DeepSpeed#3068) about errors when using Lighning with DeepSpeed. The issue is that Lightning creates two DeepSpeedEngine instances for the same model at https://github.com/Lightning-AI/lightning/blob/6ec9a6bd9e792f505ebc931742d4235f311eb289/src/lightning/pytorch/strategies/deepspeed.py#L447-L450
Yet neither of the DeepSpeedEngine is aware of the existence of the other. So when it comes to zero3 optimization, these two DeepSpeedEngines are going to operate on the same set of parameters on their own management, which leads to the crash.
We tried to tackle this issue from our end by bounding the parameters management to the model so they can be shared among DeepSpeedEngine instances, yet we realize the Lightning creates different wrapper instances for the model before passing it to DeepSpeed so from the DeepSpeed end it looks like different models.
DeepSpeed can do both training and validation on the same DeepSpeedEngine instance. Thus we want to reach out to understand more about the intuition behind using multiple DeepSpeedEngines (or wrappers) and also to check if there is anything we can do on our end to make the same DeepSpeedEngine usable for both the training and validation in your use case.

What version are you seeing the problem on?

master

How to reproduce the bug

There is a pretty nice reproduction script from the user https://github.com/microsoft/DeepSpeed/issues/3068#issuecomment-1486539136

Error messages and logs

No response

Environment

No response

More info

No response

cc @awaelchli

[awaelchli]

Hello @HeyangQin
Thanks for providing the reproducible script. I was able to reproduce the error reported on the linked issue. However, if I look at the model at the lines you pointed out above, none of the submodules is wrapped in a deepspeed engine. So calling validate() and then fit() always gets the pure LightningModule as input:

This holds at line 447:

assert not [name for name, mod in self.model.named_modules() if isinstance(mod, deepspeed.DeepSpeedEngine)]

So I think the problem is not that there are multiple DeepSpeed engines, but rather that the first call to trainer shards the model, then the user passes the (unwrapped) LightningModule again into the next Trainer.fit() call. When doing so, something seems to go wrong.

In pseudo steps, I think this is what happens:

1. model = MyLightningModule()
2. trainer.validate(model) wraps model in engine = DeepSpeedEngine(model)
3. trainer runs with engine (zero 3), parameters are sharded
4. user takes original (unwrapped) model (LightningModule) whose parameters are sharded already and passes it again into trainer.fit()
5. trainer.fit(model) again initializes deepspeed engine = DeepSpeedEngine(model) but this time, original model already has parameters sharded (I think?)

Can I ask you, conceptually what does this error mean?

assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
AssertionError: {'id': 0, 'status': 'INFLIGHT', 'numel': 23440896, 'ds_numel': 23440896, 'shape': (30522, 768), 'ds_shape': (30522, 768), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {233}}

what does inflight mean here?

[HeyangQin]

Hello @awaelchli. Thank you for the investigation and detailed explanation! Yes. I think the pseudo steps you outlined are precisely how the error is triggered.

To give more context about the error message, in zero3 optimization all the model weights will be partitioned (sharded) to all GPUs. When a certain weight is needed for computation, the DeepSpeedEngine is going to fetch it from all GPUs. inflight means this parameter is currently being fetched (fetching started but not necessarily completed). Due to the prefetch mechanism, there could be multiple parameters in the inflight status at the same time.

At runtime, trainer.validate and trainer.fit each create a DeepSpeedEngine respectively. Thus when DeepSpeedEngine A put some parameters inflight, DeepSpeedEngine B will get confused at why there are parameters that it never fetched inflight and error out.

We cannot simply have the parameters status shared by all DeepSpeedEngines because it would stop DeepSpeed from working when there are multiple models at the same time (e.g. RLHF).
A straightforward solution from the DeepSpeed side would be: bounding the parameters management to the model so they can be shared among DeepSpeedEngine instances for the same model, yet the Lightning creates new wrapper instance for the model before passing it to DeepSpeed so from the DeepSpeed end it looks like different models. I wonder if it is possible to

1. Either implement a check on the Lightning side to see if DeepSpeedEngine has been initialized for the model to reuse existing DeepSpeedEngine to prevent reinstantiation.
2. Or implement a reuse mechanism for Lightning wrapper instance so DeepSpeed can correctly recognize this is the same model. We will update the DeepSpeed framework to make parameter status tied to the model.

I think either way would fix this issue

[awaelchli]

@HeyangQin I think I understand now, thanks for the explanation. So if I understand correctly, if this special wrapper
https://github.com/Lightning-AI/lightning/blob/f6af74bf158a064673b5db284490b0da1f6c6852/src/lightning/pytorch/strategies/deepspeed.py#L445

wouldn't exist, and the raw LightningModule is passed to the deepspeed.initialize(), then deepspeed could detect that this was the same model already used before and it would work properly, correct? And are saying that this detection is not yet implemented.

I might know a way to avoid the wrapper in Lightning, but I have to verify that it is possible. Given that this could be a solution, the with the changes both in Lightning and deepspeed, we could resolve this limitation. So I see solution 2 being the most flexible, if possible.

[HeyangQin]

Hello @awaelchli. Yes. We have implemented a patch (not merged into master yet) so that if multiple DeepSpeedEngine instances work on the same model, they will share the parameter management. For the model itself, it doesn't matter if it is a wrapper or the raw model, as long as deepspeed.initialize() gets the same model, the patch will work.

cc my colleague @tjruwase

[awaelchli]

Sounds good. I'll work on a draft so I can verify your patch together with the modifications in Lightning.

[awaelchli]

@HeyangQin Is this the branch for the patch you mentioned? deepspeedai/DeepSpeed#3380