palm_30b_2p5d.py

from colossalai.zero.shard_utils import TensorShardStrategy

SEQ_LENGTH = 2048
BATCH_SIZE = 2
NUM_EPOCHS = 10
WARMUP_EPOCHS = 1

parallel = dict(
    tensor=dict(mode="2.5d", size=4, depth=1),
)

model = dict(
    type="palm_30b",
    use_grad_checkpoint=True,
    use_act_offload=True,
)

zero = dict(
    model_config=dict(
        tensor_placement_policy="cuda",
        shard_strategy=TensorShardStrategy()
    ),
    optimizer_config=dict(
        initial_scale=2**5,
    )
)

clip_grad_norm = 1.0

LOG_PATH = "./palm_30b_2p5d/"