feat(isp): support layer fused isp communication

[mwiacx]

Motivation

support fused isp-wp communication layer by layer, we have drawn on the successful experience of FSDP in aggregating AG/RS, see ag/rs part in pytorch/pytorch#114299.

Usage

add layer_fuse_isp_comm=True into your config when you already enable isp communication overlap, for example:

parallel = dict(
    zero1=dict(size=-1),
    tensor=dict(size=1, mode="isp"),
    pipeline=dict(size=1, interleaved_overlap=True, mode="1f1b"),
    weight=dict(size=4, overlap=True, layer_fuse_isp_comm=True),
)

Result

a 7B internlm2 case.

cuda memory

communication overlap

loss

layer_fuse_isp_comm=False v.s. layer_fuse_isp_comm=True

BC

no bc when layer_fuse_isp_comm=False