[shardformer] support layer state_dict and load_state_dict

[FrankLeeeee]

We need to ensure that the weights of a distributed layer and that of a normal pytorch layer can be interchangeably loaded and saved in shardformer.

[FrankLeeeee]

Closed by #4062 and #4073