[Bugfix] Consolidate params with orig size

[comaniac]

Description

Previously we use sch.metadata.base_params to maintain the original parameter shape before sharding. However, the scope of schedule metadata is limited to the current schedule and won't be transferred if .replace is used, so sch.metadata.base_params won't be correct anymore in this case.

In this PR, we removed sch.metadata.base_params and directly set an attribute to sharded parameters to record their original shapes. In future, we should revisit the sch.metadata design.

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @szhengac @chhzh123

[chhzh123]

LGTM

[comaniac]

Thanks @szhengac @chhzh123