Refactor autoTP inference for HE

[molly-smith]

Refactor replace_wo_policy with autoTP class functions. AutoTP class functions will be reused in Hybrid Engine in follow-up PR.

Remove new tensor allocation for weights and biases to reduce memory.

Some gpt2 bug fixes.

[chunyuan-w]

Hi @molly-smith, before this PR, the weight and bias of LinearLayer and LinearAllreduce are created by mp_replace.copy, the output of which is a nn.Parameter():
https://github.com/microsoft/DeepSpeed/blob/85dc854b11a66ec18e0f54e23eb72eeeddf9c81a/deepspeed/module_inject/auto_tp.py#L109

So in the __init__ function of LinearLayer, for both the if and else branch, weight and bias are always Parameter.
https://github.com/microsoft/DeepSpeed/blob/85dc854b11a66ec18e0f54e23eb72eeeddf9c81a/deepspeed/module_inject/layers.py#L36-L47

With the change in the current PR, when weight and bias are provided as input to the __init__ function, they are no longer Parameter of the module (mp_replace.copy isn't called anymore and the current code won't wrap weight and bias into Parameter anymore).
May I know if this change is intended?

[molly-smith]

Hi @molly-smith, before this PR, the weight and bias of LinearLayer and LinearAllreduce are created by mp_replace.copy, the output of which is a nn.Parameter():

https://github.com/microsoft/DeepSpeed/blob/85dc854b11a66ec18e0f54e23eb72eeeddf9c81a/deepspeed/module_inject/auto_tp.py#L109

So in the __init__ function of LinearLayer, for both the if and else branch, weight and bias are always Parameter.

https://github.com/microsoft/DeepSpeed/blob/85dc854b11a66ec18e0f54e23eb72eeeddf9c81a/deepspeed/module_inject/layers.py#L36-L47

With the change in the current PR, when weight and bias are provided as input to the __init__ function, they are no longer Parameter of the module (mp_replace.copy isn't called anymore and the current code won't wrap weight and bias into Parameter anymore). May I know if this change is intended?

Hi @chunyuan-w, thanks for bringing this to my attention. This was not intentional. Will have a fix soon.

[chunyuan-w]

Hi @molly-smith, before this PR, the weight and bias of LinearLayer and LinearAllreduce are created by mp_replace.copy, the output of which is a nn.Parameter():
https://github.com/microsoft/DeepSpeed/blob/85dc854b11a66ec18e0f54e23eb72eeeddf9c81a/deepspeed/module_inject/auto_tp.py#L109

So in the __init__ function of LinearLayer, for both the if and else branch, weight and bias are always Parameter.
https://github.com/microsoft/DeepSpeed/blob/85dc854b11a66ec18e0f54e23eb72eeeddf9c81a/deepspeed/module_inject/layers.py#L36-L47

With the change in the current PR, when weight and bias are provided as input to the __init__ function, they are no longer Parameter of the module (mp_replace.copy isn't called anymore and the current code won't wrap weight and bias into Parameter anymore). May I know if this change is intended?

Hi @chunyuan-w, thanks for bringing this to my attention. This was not intentional. Will have a fix soon.

Thanks!