[ZeRO Infinity] Allow Init to take a dict for the deepspeed config

[SeanNaren]

Hey guys!

I needed this fix in order to run ZeRO Infinity with Lightning, as the deepspeed config is already loaded into memory as a dictionary. When passing to Init and remote_device=nvme, passing the config as a dictionary would crash without this fix!

[tjruwase]

I think typing json_file as a dict is confusing, rather we should pass dict as param_dict. So ideally this file should be unchanged for this PR. Instead the following should be modified to optionally pass param_dict as appropriate.
Call 1
Call 2

[SeanNaren]

Much cleaner, just made the change!

[stas00]

@SeanNaren, this is great - we need it too! thank you!

Just one suggestion - deepspeed.initialize names this same functionally arg differently. Would you be open to rename s/param_dict/config_params/? to match deepspeed.initialize? I guess it's a bit late to rename deepspeed.initialize's existing args .

see it here:

DeepSpeed/deepspeed/__init__.py

Line 90 in cf5ea89

config_params: Optional: Instead of requiring args.deepspeed_config you can pass your deepspeed config

I am fine if it remains as you proposed, as yours is more intuitive naming than config_params - the latter says twice the same thing and doesn't say it's a dict.

If I had a choice I'd rename both to config_dict but it's too late :)

[SeanNaren]

@SeanNaren, this is great - we need it too! thank you!

Just one suggestion - deepspeed.initialize names this same functionally arg differently. Would you be open to rename s/param_dict/config_params/? to match deepspeed.initialize? I guess it's a bit late to rename deepspeed.initialize's existing args .

see it here:

DeepSpeed/deepspeed/__init__.py

Line 90 in cf5ea89

config_params: Optional: Instead of requiring args.deepspeed_config you can pass your deepspeed config

I am fine if it remains as you proposed, as yours is more intuitive naming than config_params - the latter says twice the same thing and doesn't say it's a dict.

If I had a choice I'd rename both to config_dict but it's too late :)

I'm glad finally I'm able to assist you out stas! Should be straight forward, I can make a PR :)

[stas00]

or should we just make it a single arg that can take path or dict? This is what we do in transformers:

def deepspeed_parse_config(ds_config):
    """
    If ``ds_config`` isn't already a dict, read it from the config file.

    If it's already a dict, return a copy of it, so that we can freely modify it.
    """
    dep_version_check("deepspeed")

    if isinstance(ds_config, dict):
        # Don't modify user's data should they want to reuse it (e.g. in tests), because once we
        # modified it, it will not be accepted here again, since some config params must be not set by users
        config = deepcopy(ds_config)
    elif isinstance(ds_config, str):
        with io.open(ds_config, "r", encoding="utf-8") as f:
            config = json.load(f)
    else:
        raise ValueError("expecting either a path to a config file or a pre-populated dict")

    return config

minus copy and dep_checking.

[SeanNaren]

or should we just make it a single arg that can take path or dict? This is what we do in transformers:

def deepspeed_parse_config(ds_config):
    """
    If ``ds_config`` isn't already a dict, read it from the config file.

    If it's already a dict, return a copy of it, so that we can freely modify it.
    """
    dep_version_check("deepspeed")

    if isinstance(ds_config, dict):
        # Don't modify user's data should they want to reuse it (e.g. in tests), because once we
        # modified it, it will not be accepted here again, since some config params must be not set by users
        config = deepcopy(ds_config)
    elif isinstance(ds_config, str):
        with io.open(ds_config, "r", encoding="utf-8") as f:
            config = json.load(f)
    else:
        raise ValueError("expecting either a path to a config file or a pre-populated dict")

    return config

minus copy and dep_checking.

Also how we set it up in PL, which I think is natural! If @tjruwase is at agreement I can make the change here, such that deepspeed_config can be an dict object rather than a path. If we're worried about BW compatibility we can keep config_params still as an argument?

[tjruwase]

@SeanNaren, I am fine with making the change. I think it is early days that BW compatibility is less priority given the long-term benefits of the change.

[stas00]

There was no new release yet since your addition, so I agree with @tjruwase, it should not be an issue of BC.

Thank you!

[SeanNaren]

I'm at cross-roads on how far I should go in this refactor:

1. For DeepSpeed Init, remove param_dict, making deepspeed_config either a JSON file or config dict. This introduces the initial logic in the PR
2. make deepspeed_config an input from initialize which can also be a JSON file or a config dict, and continues to take priority over args.
3. Unify references to config_params or params_dict to deepspeed_config

I think 1/2 are probably the most crucial but am curious what you guys think. 3 is more involved as it refactors a lot of variables defined across various config functions and files. cc @stas00 @tjruwase

[tjruwase]

@SeanNaren, thanks for sharing these thoughts. I agree that we should 1 & 2 for now and defer 3 for later. To my understanding 3 has minimal if any user-facing impact.

[stas00]

Thank you, @SeanNaren for this proposal.

The crucial thing for HF is that Init won't save the config and try to use it later, but only pluck a few params it absolutely needs and drop it. This is because during model init we don't yet have the full config ready. We need the model object to get the hidden size to automatically fill out auto-params that rely on hidden size and we don't have the number of training steps figured out until much later. So our config happens in 2 stages, one immediate upon argparse and we are ready to feed Init, and the 2nd stage is at the start of train or eval.

That said I'm concerned that if we use unification and it's hard to tell there is a chance someone will try to save the config at Init stage and disregard the updated version during initialize, leading to errors.

So for me I think that Init should only extract immediately the 3-4 params it needs and drop the passed config, in which case it doesn't matter what you call it internally.

Therefore 1 for sure, 2 I'm impartial, but definitely is a good idea, and 3 instead of unifying I'd actually rename config_params to init_params so it's clear that we can only refer to it for Init-specific params and not as a normal config. In fact I propose that a user should be able to pass to Init a trimmed version of config that has only 4 params (or however many Init might ever need).

I hope this anti-unification logic is sensible.

If we do 2, then perhaps args should be deprecated and in time removed.

[SeanNaren]

Thanks Stas! we can increment on this; I agree with being able to pass just the parameters required for NVMe, but also it would be nice to keep the option of just passing the entire config for extensibility in the future + for user simplicity

[stas00]

Oh, I'm not against passing the entire config, future-proofing is exactly that - I just want to make sure it's not saved and used later in deepspeed's normal initialize.