[Disco] Loading-time sharding support

[LeshengJin]

In our previous implementation, parameter sharding relies on pre-quantization weight processing, meaning each set of quantized weights corresponds strictly to a hardcoded constant num_shards, and re-quantization is strictly required upon each change of #GPUs, e.g. from 4-GPU to 8-GPU setting. This PR makes it possible to move parameter sharding to post-quantization loading-time. During loading, we iterate over all parameters and apply the sharding operation based on the provided sharding information.

To make this happen, this PR makes an enhancement to the existing shard_info.json to include the sharding function being used at loading time. Each parameter is attached to a list of loading-time preprocessing methods that are serially applied to it to transform this parameter to the desired shape, as shown in the example below:

shard_info = {
  "x_0": [ # name of the parameter
    [ # a list of preprocessing functions to be applied
      "tests.disco.shard_dim_1",  # name of the sharding function
      [(num_shards, 64, 64), "float16"],  # output shape/dtype of `tests.disco.shard_dim_1`
      num_shards,  # extra inputs to `tests.disco.shard_dim_1`
    ],
  ],
  "x_1": [...],
}

To parameter x_0, it means we will call method tests.disco.shard_dim_1 which has the signature:

def shard_dim_1(
  input: NDArray,
  num_shards, # extra inputs
  output: NDArray, # and its shape is (num_shards, 64, 64), and dtype is "float16"
) -> None: ...

This approach simplifies parameter sharding for users and ensures correctness.

[junrushao]

CC: @jinhongyii

[junrushao]

Let's unify sharding and reordering into a single function - no need for shard3d any more :)

[jinhongyii]

Can we separate reorganizing function and sharding function again? I'm thinking of some situation where there can be 2 reorganizing function ( doing nothing and combine qkv), and n different sharding functions( different tensor shape and different sharding dimension). In this case, we will need at most 2n functions in total in IRModule. Another benefit of separating them is that DistIR will not need to handle merging reorganizing function and sharding function.