class DenseLm12kWide41BAdam16x16(DenseLm128B16x16):
"""41B params LM model with 2D split and ADAM optimizer on v3-512."""
# Each layer has 1.6875B parameters.
SEQUENCE_LENGTH = 2048
NUM_DEVICES_PER_SPLIT = 512
BATCH_DIM_PER_DEVICE = 0.5 # Total batch size 256
DEVICE_MESH_SHAPE = [16, 32]
DEVICE_MESH = gshard_utils.GetNonPod2dMesh(DEVICE_MESH_SHAPE, [16, 16, 2])
NUM_TRANSFORMER_LAYERS = 24
HIDDEN_DIM = 48 * 1024
MODEL_DIM = 12 * 1024
NUM_HEADS = 96
ATTENTION_KEY_VALUE_DIM = 128
GATED_GELU = False
POSITIONAL_EMBEDDING = True
NUM_MICRO_BATCHES = 1
def Task(self):
p = super().Task()
p.train.optimizer = ShardedAdam.Params().Set(
beta1=0.9,
beta2=0.999,
epsilon=1e-6,
num_micro_batches=self.NUM_MICRO_BATCHES)
return p
class DenseLm12kWide41BAdam8x8(DenseLm12kWide41BAdam16x16):
# IF OOM, try 0.25 BATCH_DIM_PER_DEVICE and 8 NUM_MICRO_BATCHES
BATCH_DIM_PER_DEVICE = 0.5 # Total micro-batch size 64
NUM_MICRO_BATCHES = 4 # Total batch size 256
NUM_DEVICES_PER_SPLIT = 128
DEVICE_MESH_SHAPE = [8, 16]
DEVICE_MESH = gshard_utils.GetNonPod2dMesh(DEVICE_MESH_SHAPE, [8, 8, 2])
class DenseLm128B16x16(DenseLm128B8x8):
"""128B params LM model with 2D split on v3-512."""
SEQUENCE_LENGTH = 1024
NUM_DEVICES_PER_SPLIT = 512
BATCH_DIM_PER_DEVICE = 0.25 # Total batch size 128
DEVICE_MESH_SHAPE = [16, 32]
DEVICE_MESH = gshard_utils.GetNonPod2dMesh(DEVICE_MESH_SHAPE, [16, 16, 2])
class DenseLm12kWide162BAdam16x16(DenseLm12kWide41BAdam16x16): """162B params LM model with 2D split and ADAM optimizer on v3-512.""" BATCH_DIM_PER_DEVICE = 0.125 # Total batch size 64 NUM_TRANSFORMER_LAYERS = 96 DEVICE_MESH_SHAPE = [16, 32] DEVICE_MESH = gshard_utils.GetNonPod2dMesh(DEVICE_MESH_SHAPE, [16, 16, 2])
class DenseLm175B8x8Decode2D(DenseLm175B32x32):
"""175B params LM model decoding on v3-128.
2D logical mesh. It can load a checkpoint from DenseLm175B32x32.
"""
BATCH_DIM_PER_DEVICE = 0.125
NUM_DEVICES_PER_SPLIT = 128
# NUM_HEADS is not a multiple of 128 so we use 2D sharding on M and H.
DEVICE_MESH_SHAPE = [8, 16]
DEVICE_MESH = gshard_utils.GetNonPod2dMesh(DEVICE_MESH_SHAPE, [8, 8, 2])
def Task(self):
p = super().Task()
# relative_attention_use_universal_1d_position should be set to False in
# decoding.
p.builder.relative_attention_use_universal_1d_position = False
p.builder.model_dim_reshape_segments = self.DEVICE_MESH_SHAPE[0]
p.builder.emb_w_split = [1, 0]
p.builder.emb_out_split = [-1, -1, 0]
p.builder.blm_split = [-1, -1, 0]
p.builder.blh_split = [-1, -1, 1]
p.builder.qkv_split = [0, -1, 1,
-1] # [-1, -1, 1, -1] for global batch 1.
p.builder.logits_split = [-1, -1, 1]
return p
class DenseLm175B32x32DP(DenseLm175B32x32):
"""175B model running on v3-2048 with 2-way data parallelism."""
NUM_DEVICES_PER_SPLIT = 1024
TRAIN_STEPS_PER_LOOP = 20
DEVICE_MESH_SHAPE = [64, 16]
DEVICE_MESH = gshard_utils.GetNonPod2dMesh([16, 64],
[16, 32, 2]).transpose()
MODEL_DIM_RESHAPE_SEGMENTS = [16]
class DenseLm128B8x8(DenseLmTemplate):
"""128B params LM model with 2D split."""
SEQUENCE_LENGTH = 1024
NUM_DEVICES_PER_SPLIT = 128
BATCH_DIM_PER_DEVICE = 0.125
NUM_TRANSFORMER_LAYERS = 64 # 64 blocks of [DecSelfAttention, DenseReluDense]
DEVICE_MESH_SHAPE = [8, 16]
DEVICE_MESH = gshard_utils.GetNonPod2dMesh(DEVICE_MESH_SHAPE, [8, 8, 2])
def Task(self):
p = super().Task()
p.train.tpu_device_order_mode = 2 # DeviceOrderMode.MESH
p.builder.model_dim_reshape_segments = self.DEVICE_MESH_SHAPE[1]
p.builder.emb_w_split = [-1, 1]
p.builder.emb_out_split = [0, -1, 1]
p.builder.blm_split = [0, -1, 1]
p.builder.logits_split = [0, -1, 1]
return p