Add MLP/lm_head tp grain size setting.

[Yejing-Lai]

This PR aims to add MLP/lm_head tp size granularity setting to deepspeed.init_inference() API. It will be more flexible to set the MLP/lm_head sharding grain size.

DNN library favors tensor size in granularity of power of 2, we pick 64 as a default size.

We aim to be able to set the MLP/lm_head tp grain size flexibly. This is a preliminary solution. If there is a better solution, we can discuss it together. Thanks~

[Yejing-Lai]

@delock Please review, thanks~

[delock]

Hi @tjruwase @awan-10 , this PR from @Yejing-Lai adds a new parameter to init_inference to set the MLP/lm_head granularity. The reason we need this parameter is different library will need different granularity to perform optimal. i.e. oneDNN needs granularity of 64. In quantization, granularity better set to the size of quantization group. Other library may need different granularity setting.

We know it extends init_inference where we think it's the right place to set autotp inference havior. Want's to hear your thoughts on this, thanks!

[tjruwase]

@delock, @Yejing-Lai, thanks for the PR. I agree that updating init_inference is right choice. However, I am curious whether it is better to make grain_size a sub-field of the existing tensor_parallel field.

DeepSpeed/deepspeed/inference/config.py

Lines 132 to 136 in 1b58ba5

	tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp")
	"""
	Configuration for tensor parallelism used to split the model across several
	GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`.
	"""

Did you guys consider this choice? Either way, I have approved the PR.

[Yejing-Lai]

@delock, @Yejing-Lai, thanks for the PR. I agree that updating init_inference is right choice. However, I am curious whether it is better to make grain_size a sub-field of the existing tensor_parallel field.

DeepSpeed/deepspeed/inference/config.py

Lines 132 to 136 in 1b58ba5

	tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp")
	"""
	Configuration for tensor parallelism used to split the model across several
	GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`.
	"""

Did you guys consider this choice? Either way, I have approved the PR.

Thank you for your advice. I updated the code. Please review~

[tjruwase]

Thank you for your advice. I updated the code. Please review~

@Yejing-Lai, thanks for updating the code. LGTM.