ZeRO-Inference refresh

[tjruwase]

Two new optimizations

• Weight quantization
• KV cache offloading to CPU memory

[delock]

@tjruwase are there examples using weight quantization and kv_cache offloading? Wonder what it takes to support int4 on other accelerators as well.

[tjruwase]

@tjruwase are there examples using weight quantization and kv_cache offloading? Wonder what it takes to support int4 on other accelerators as well.

Sorry, that release won't be ready until next week. It will be great to get your feedback.

[tjruwase]

FYI, @donglinz @cli99

[ftian1]

@tjruwase May I get some inputs from your side about contributing more weight only quantization algorithms like WeightOnly RTN/AWQ/GPTQ to DeepSpeed inference? is it ok for DeepSpeed inference to accept other compression algos except for ZeRO? and some of such algos may require calibration on some datasets which may have impact on inference API. do you think it's valuable to add? or shall we stick on those data-free compression at first?

[ftian1]

and some questions/comments about current PR:

1. will _init_group_wise_weight_quantization() be the final interface exposing to user to enable weight quantization? it's a little weird as its name convention is more like internal function naming and the user experiment is not quite good.

2. will this feature be enabled in DeepSpeed Inference path? that's user can enable this feature by invoking deepspeed.init_inference() but not below

        model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16)
        model = model.eval()

        model = _init_group_wise_weight_quantization(model, ds_config)
        ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]

3. may I know when this PR will be merged? or is there any special concerns on not merging such feature to DeepSpeed?

[tjruwase]

@tjruwase May I get some inputs from your side about contributing more weight only quantization algorithms like WeightOnly RTN/AWQ/GPTQ to DeepSpeed inference? is it ok for DeepSpeed inference to accept other compression algos except for ZeRO? and some of such algos may require calibration on some datasets which may have impact on inference API. do you think it's valuable to add? or shall we stick on those data-free compression at first?

@ftian1, thanks for the questions. We will share a response asap.

[tjruwase]

and some questions/comments about current PR:

1. will _init_group_wise_weight_quantization() be the final interface exposing to user to enable weight quantization? it's a little weird as its name convention is more like internal function naming and the user experiment is not quite good.

_init_group_wise_weight_quantization() is still experimental. We hope to finalize the interface soon.

2. will this feature be enabled in DeepSpeed Inference path? that's user can enable this feature by invoking deepspeed.init_inference() but not below

        model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16)
        model = model.eval()

        model = _init_group_wise_weight_quantization(model, ds_config)
        ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]

Yes, this is the plan.

3. may I know when this PR will be merged? or is there any special concerns on not merging such feature to DeepSpeed?

This week.