[Unity] Add LoadParamOnWorker0 function in shard loader

[jinhongyii]

In DistIR compilation flow, shard loading is implemented in a Relax func, like

a = LoadParamOnWorker0(loader, index=0)
b = broadcast(a)
c = LoadParamOnWorker0(loader, index=1)
d = scatter_from_worker0(c)

LoadWholeParamOnWorker0 loads the unsharded param on worker0, and for other workers returns an empty array.

This PR implements LoadWholeParamOnWorker0

[jinhongyii]

cc: @junrushao

[junrushao]

The implementation looks good, but not sure what the best name tbh. @tqchen maybe you could suggest?

[tqchen]

How about LoadParamOnWorker0. And for future sharded version we can say LoadShardedParams

[jinhongyii]

sounds good

[junrushao]

Sounds good! Let’s do that then