Experiment/dcp save ordering

[IsaevIlya]

Description

Improved Performance for Partial Checkpoint Loading

Background

By default, PyTorch sorts checkpoint data based on size, which distributes tensors/weights randomly across checkpoint shards. While this approach works well with local storage, it can impact performance when working with cloud storage. Currently, PyTorch's checkpoint loading process doesn't follow the same order used during saving, resulting in non-sequential file access patterns.

Changes

1. Sequential Read Optimization

• Modified the ordering of items in LocalPlan based on their actual offset in checkpoint shards
• Ensures sequential reading of data, improving I/O efficiency

2. Custom Sorting for Partial Loading

• Added ability to provide a custom sorting key function during checkpoint saving
• Allows users to group specific data at the beginning of checkpoints
• Optimizes partial loading scenarios by reducing the amount of data that needs to be read

Usage Example

If you need to load only model layers that start with "model.model.layers.1", you can group these tensors at the beginning of checkpoint shards:

def sort_key_func(item):
    return not item.index.fqn.startswith("model.model.layers.1"), item.index.fqn

def save_checkpoint(state_dict, region, uri):
    writer = S3StorageWriter(region, uri, sort_key=sort_key_func)
    dcp.save(state_dict, storage_writer=storage_writer)

Results

The resulting checkpoint shard layout will prioritize tensors starting with "model.model.layers.1":

File name: __0_0.distcp
        key: model.model.layers.10.self_attn.o_proj.weight     ,        offset: 0
        key: model.model.layers.11.mlp.down_proj.weight        ,        offset: 67110441
        key: model.model.layers.11.self_attn.k_proj.weight     ,        offset: 247467090
        key: model.model.layers.13.mlp.gate_proj.weight        ,        offset: 314577531
        key: model.model.layers.15.input_layernorm.weight      ,        offset: 494934180
        key: model.model.layers.15.post_attention_layernorm.weight,     offset: 494952141
        ...
        key: model.lm_head.weight                              ,        offset: 1124125164
        key: model.model.embed_tokens.weight                   ,        offset: 1950404629
        key: model.model.layers.21.mlp.gate_proj.weight        ,        offset: 2776684094
        ....

• I have updated the CHANGELOG or README if appropriate

---

By submitting this pull request, I confirm that my contribution is made under the terms of BSD 3-Clause License and I agree to the terms of the LICENSE.

[jet-tong]

Tested the PR with "model.model.layers.1" sort key function, and can confirm that these changes speed up partial checkpoint loading.