Can't load GPT Model: ImportError: 'TP_REPLICATED_PARAMETER_PATTERNS' from 'deepspeed.checkpoint'

[V3RGANz]

When I am trying to load GPT Model it is fails, because it requires non-existing deepspeed import :

  File "/workspace/megatron/megatron/model/gpt_model.py", line 27, in <module>
    from deepspeed.checkpoint import (
ImportError: cannot import name 'TP_REPLICATED_PARAMETER_PATTERNS' from 'deepspeed.checkpoint' (/usr/local/lib/python3.10/dist-packages/deepspeed/checkpoint/__init__.py)

This is the line with constant import:
https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/model/gpt_model.py#L30

However, deepspeed package does not contain this constant:
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/checkpoint/constants.py
I could not find where this constant come from, because it never appears in VCS history of deepspeed package

Megatron-Deepspeed version: main branch
deepspeed version: 0.11.1

[conglongli]

@tjruwase I believe this is related to your PR #265. This issue probably will be solved after the companion DS PR: microsoft/DeepSpeed#4516 is merged, but I feel it'd be better to have a bit more backward compatibility (what if user is using an older DeepSpeed with Megatron-DeepSpeed).

[tjruwase]

@V3RGANz, sorry about this bc-breakage. Will fix asap.

[tjruwase]

@conglongli and @V3RGANz, after a second look. This issue is less of bc-breakage but more of a bug fix. The previous DS and Meg-DS behavior was a bug pointed out by #265 (comment). So, an older DS is not recommended. My mistake was failing to generate a new DS release. I will do that immediately.

[conglongli]

So, an older DS is not recommended.

@tjruwase Instead of completely stop supporting older DS, could we add DS version check in Megatron-DeepSpeed side, and only trigger the universal checkpointing feature when matching DS version detected?

[tjruwase]

@conglongli, that is a good point. I can work on that.

[tjruwase]

@V3RGANz, please check if fixed by #271.

[V3RGANz]

@tjruwase, checked #271. Yes, this fixes my issue