* support sharding parallel

* fix name

* fix

* update

* test amp for sharding

---------

Co-authored-by: pangengzheng <pangengzheng.baidu.com>