palm_62b_2d.py

from colossalai.zero.shard_utils import TensorShardStrategy

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

parallel = dict(
    tensor=dict(mode="2d", size=4),
)

model = dict(
    type="palm_62b",
    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