[REQUEST] Partial weight load on demand

Is your feature request related to a problem? Please describe.
I'm seeking to deploy DeepSpeed (DS) on CPU platforms, such as Intel's SapphireRapid or Arm CPU with SME/SVE, and utilize AutoTP. The current issue arises when DS loads weights using AutoTP: a single rank loads all weights, splits the tensor, and keeps the required part. When deployed on a CPU with potentially dozens of ranks, this results in an equal number of copies of all weight tensors needing to be loaded into memory. This creates a significantly large memory requirement.

Describe the solution you'd like
My ideal solution would allow DS to load weights on demand, meaning each rank would load a specified part in conjunction with AutoTP.

Describe alternatives you've considered
There are several alternatives I've considered, and I would appreciate input from the DS maintainers to determine the most effective solution.

• Partial load from saved weight
  Using AutoTP, the model is initialized with metadata (fast_init) and the weight shape is adapted to the TP number. Each rank only reads the part of the weight it needs from .bin or .pt. However, I'm unsure if pt or other weight formats could support this. If not, we could potentially convert to a different format that does support partial load.
• Shared memory among multiple ranks in one node/server
  Rather than loading all weights per rank, all ranks in a node would maintain a single copy of the weight and load the required part from it. However, I've yet to find a suitable method to achieve this. As I have switched the DS coll-comm backend to OpenMPI, each rank is set up independently.
• Split weight offline and load split checkpoint with kernel injection
  Load all weights with a single rank, manually split them according to the TP number, and save each part for each rank. During inference, each rank loads the required part with the kernel injection code path (replace_module(checkpoint=...)).

Additional context
The crucial point is identifying which of these options would be most feasible, require the least amount of code development, and be easiest to merge into the master branch.

Start by tracing DeepSpeed's AutoTP weight-loading path and compare the three approaches described: partial reads from .bin or .pt files, shared memory across ranks, and offline-split checkpoints with kernel injection. Done would mean selecting and implementing a maintainable way for each rank to load only its required weight partition, with the approach validated on the target CPU deployment scenarios.