Class DistConfig

lmdeploy/pytorch/config.py:130–221

@dataclass
class DistConfig:
    dp: int = 1
    ep: int = 1
    dp_rank: int = 0
    enable_microbatch: bool = False
    enable_eplb: bool = False
    world_size: int = 1

    # tp
    tp: int = 1  # default tp, equal to attn_tp
    attn_tp: int = None  # tp for attention
    mlp_tp: int = None  # tp for mlp
    moe_tp: int = None  # tp for moe

    # tp mode
    mlp_tp_mode: TPMode = TPMode.DEFAULT
    moe_tp_mode: TPMode = TPMode.DEFAULT

    def __post_init__(self):
        """Post init."""
        assert self.dp_rank < self.dp
        assert self.dp >= 1

        dp = self.dp
        tp = self.tp
        ep = self.ep

        # ignore layer to for dp==1
        if dp == 1:
            self.mlp_tp = None
            self.attn_tp = None
            self.moe_tp = None

        # mlp and moe tp
        self.mlp_tp = self.mlp_tp or tp
        self.moe_tp = self.moe_tp or (1 if ep > 1 else self.mlp_tp)

        # world_size
        world_size = ep if ep > 1 else max(self.mlp_tp, self.moe_tp)
        self.world_size = world_size
        assert (world_size >= dp and world_size % dp == 0), (f'world_size {world_size}, dp {dp}')
        assert (world_size >= ep and world_size % ep == 0), (f'world_size {world_size}, ep {ep}')
        assert (world_size >= self.mlp_tp
                and world_size % self.mlp_tp == 0), (f'world_size {world_size}, mlp_tp {self.mlp_tp}')
        assert (world_size >= self.moe_tp
                and world_size % self.moe_tp == 0), (f'world_size {world_size}, moe_tp {self.moe_tp}')

        # attn tp
        self.attn_tp = self.attn_tp or self.world_size // dp
        self.tp = self.attn_tp
        if self.mlp_tp > 1:
            assert (self.mlp_tp >= self.attn_tp
                    and self.mlp_tp % self.attn_tp == 0), (f'mlp_tp {self.mlp_tp}, attn_tp {self.attn_tp}')
        if self.moe_tp > 1:
            assert (self.moe_tp >= self.attn_tp
                    and self.moe_tp % self.attn_tp == 0), (f'moe_tp {self.moe_tp}, attn_tp {self.attn_tp}')
        assert (world_size >= self.attn_tp
                and world_size % self.attn_tp == 0), (f'world_size {world_size}, attn_tp {self.attn_tp}')