Function reshape

deepspeed/checkpoint/reshape_meg_2d.py:182–208

reshape([tp_size_src, pp_size_src, dp_size_src], [tp_size_tgt, pp_size_tgt, dp_size_tgt])

(src, tgt)

def reshape(src, tgt):
    """
    reshape([tp_size_src, pp_size_src, dp_size_src],
            [tp_size_tgt, pp_size_tgt, dp_size_tgt])
    """

    print(f"\n\n*** Reshaping: {src} => {tgt}")

    tp_size_src, pp_size_src, dp_size_src = src
    tp_size_tgt, pp_size_tgt, dp_size_tgt = tgt

    tp_ranks1, pp_ranks1, dp_ranks1 = get_mpu_ranks(tp_size=tp_size_src, pp_size=pp_size_src, dp_size=dp_size_src)
    tp_ranks2, pp_ranks2, dp_ranks2 = get_mpu_ranks(tp_size=tp_size_tgt, pp_size=pp_size_src, dp_size=dp_size_src)
    tp_ranks3, pp_ranks3, dp_ranks3 = get_mpu_ranks(tp_size=tp_size_tgt, pp_size=pp_size_tgt, dp_size=dp_size_src)

    # handle tp contraction first
    print("\n*** TP contraction:")

    for i, r in enumerate(tp_ranks1):
        print(f'{tp_ranks1[i]} => {tp_ranks2[i]}')

    # handle pp contraction next

    print("\n*** PP contraction:")

    for i, r in enumerate(pp_ranks1):
        print(f'{pp_ranks2[i]} => {pp_ranks3[i]}')


# easy