Does DeepSpeed Inference support offload params to cpu

[TigerYang414]

I just want to run llama 30b inference on A10, for which params exceed gpu memory capacity.
It seems InferenceEngine does not process zero related arguments.

Then I try to use deepspeed.initialize with zero offload, it will always trigger cuda oom.

[devkaranjoshi]

Exactly same error i am getting whenever i am trying to take inference from vicuna33b.

[TigerYang414]

It seems DS Inference only support tensor parallel. I just try use training mode with zero offload. But it runs very slow.
Does anyone know how I can check whether it works as expected? Like prefetch param during other layer calculation?

Here is my ds config used:
ds_config = {
"max_tokens": args.max_tokens,
"fp16": {
"enabled": True
},
"zero_optimization": {
"stage": 3,
"overlap_comm": True,
"contiguous_gradients": True,
"stage3_param_persistence_threshold": 1 << 30, # disable tensor partition
"stage3_max_live_parameters": 1 << 30,
"offload_param":{"device":"cpu","pin_memory":True}
},
"steps_per_print": 2000,
"train_batch_size": 2,
}