def __init__(self, num_tokentypes=0, parallel_output=True, topology=None):
args = get_args()
self.parallel_output = parallel_output
self.hidden_size = args.hidden_size
self.num_tokentypes = num_tokentypes
self.init_method = init_method_normal(args.init_method_std)
self.output_layer_init_method = scaled_init_method_normal(args.init_method_std, args.num_layers)
weight_tying = not args.no_weight_tying
if args.pos_emb == 'rpe':
rpe_emb = ParallelRelativePositionBias(causal=True, num_buckets=args.rpe_num_buckets,
max_distance=args.rpe_max_distance,
heads=args.num_attention_heads)
self.fp16_lm_cross_entropy = args.fp16_lm_cross_entropy
#
# forward() prototype
#
self.specs = []
# Embedding layer
if weight_tying:
self.specs.append(TiedLayerSpec('embed',
EmbeddingPipe,
self.hidden_size,
args.padded_vocab_size,
args.max_position_embeddings,
args.hidden_dropout,
self.init_method,
self.num_tokentypes,
tied_weight_attr='word_embeddings_weight'))
else:
self.specs.append(LayerSpec(EmbeddingPipe,
self.hidden_size,
args.padded_vocab_size,
args.max_position_embeddings,
args.hidden_dropout,
self.init_method,
self.num_tokentypes))
# outputs are now (hidden_states, attention_mask)
# data format change to avoid explicit tranposes : [b s h] --> [s b h]
self.specs.append(lambda x: (x[0].transpose(0, 1).contiguous(), *x[1:]))
# Transformer layers
for x in range(args.num_layers):
if args.sparsity == 'none':
sparse = False
elif args.sparsity == 'all':
sparse = True
elif args.sparsity == 'interspersed':
sparse = not x % 2 == 0
self.specs.append(
LayerSpec(ParallelTransformerLayerPipe,
attention_mask_func=gpt2_attention_mask_func,
init_method=self.init_method,
output_layer_init_method=self.output_layer_init_method,
layer_number=x,
sparse=sparse,
rpe=rpe_emb if args.pos_emb == 'rpe' else None,
rotary=args.pos_emb == 'rotary'))
# Undo data format change and drop mask
self.specs.append(lambda x: x[0].transpose(0, 1).contiguous())
# Final layernorm after transformer layers
if args.norm == "rmsnorm":
norm = RMSNorm
eps = args.rms_norm_epsilon
elif args.norm == "layernorm":
eps = args.layernorm_epsilon
norm = LayerNorm
elif args.norm == "scalenorm":
eps = args.scalenorm_epsilon
norm = ScaleNorm
self.specs.append(
LayerSpec(norm,
args.hidden_size,
eps=eps))
# XXX forward_method_parallel_output is assumed to be None, but we're not in a
# fwd method to assert
def _logits_helper(embedding, lm_output):
"""Just a wrapper to massage inputs/outputs from pipeline. """
return parallel_lm_logits(
lm_output,
embedding.word_embeddings_weight,
self.parallel_output)
if weight_tying:
self.specs.append(
TiedLayerSpec('embed',
EmbeddingPipe,
self.hidden_size,
args.padded_vocab_size,
args.max_position_embeddings,
args.hidden_dropout,
self.init_method,
self.num_tokentypes,
forward_fn=_logits_helper,
tied_weight_attr='word_embeddings_weight')
)
else:
# TODO: not sure whether to use RowParallelLinear's default scatter to mp region here, or copy, which is
# the default of parallel_lm_logits. Should investigate benefits of both
self.specs.append(
LayerSpec(
mpu.RowParallelLinear,
args.hidden_size,
args.padded_vocab_size,
bias=False,
input_is_parallel=False,
parallel_output=self.parallel_output,
skip_bias_add=False
)
)
self.specs.append(lambda x: x[0]) # drop bias
loss_fn = partial(cross_entropy, _fp16=self.fp16_lm_cross_entropy)
if args.checkpoint_activations:
interval = args.checkpoint_num_layers
else:
interval = 0
super().__init__(layers=self.specs,
loss_fn=loss_fn,
topology=topology,
activation_checkpoint_interval=interval,
partition_method=args.pipe_partition_method) # 'type:transformer' / 'parameters'
def init_specs(self):
weight_tying = not self.neox_args.no_weight_tying
self.specs = []
# Embedding layer
# input will be (input_ids, position_ids, attention_mask)
if weight_tying:
self.specs.append(
TiedLayerSpec(
"embed",
EmbeddingPipe,
self.neox_args,
self.hidden_size,
self.neox_args.padded_vocab_size,
self.neox_args.max_position_embeddings,
self.neox_args.hidden_dropout,
self.init_method,
self.num_tokentypes,
tied_weight_attr="word_embeddings_weight",
))
else:
self.specs.append(
LayerSpec(
EmbeddingPipe,
self.neox_args,
self.hidden_size,
self.neox_args.padded_vocab_size,
self.neox_args.max_position_embeddings,
self.neox_args.hidden_dropout,
self.init_method,
self.num_tokentypes,
))
# NB: the attention mask always needs to be the *last* item in the args when being passed from
# one stage to the next, because deepspeed is hacks on top of hacks.
#
# outputs are now (hidden_states, attention_mask)
self.specs.append(_pre_transformer_block)
# T5 RPE positional embedding
if self.neox_args.pos_emb == "rpe":
hidden_size_per_attention_head = mpu.divide(
self.neox_args.hidden_size, self.neox_args.num_attention_heads)
rpe_scale = math.sqrt(hidden_size_per_attention_head)
rpe_emb = ParallelRelativePositionBias(
neox_args=self.neox_args,
scale=rpe_scale,
causal=True,
num_buckets=self.neox_args.rpe_num_buckets,
max_distance=self.neox_args.rpe_max_distance,
heads=self.neox_args.num_attention_heads,
)
# Transformer layers
for i in range(self.neox_args.num_layers):
layer_type = self.neox_args.attention_config[i]
if layer_type in ["gmlp", "amlp"]:
self.specs.append(
LayerSpec(
GMLPBlock,
init_method=self.init_method,
layer_number=i,
output_layer_init_method=self.output_layer_init_method,
neox_args=self.neox_args,
mask_fn=gpt2_attention_mask_func,
))
else:
self.specs.append(
LayerSpec(
ParallelTransformerLayerPipe,
neox_args=self.neox_args,
attention_mask_func=gpt2_attention_mask_func,
init_method=self.init_method,
output_layer_init_method=self.output_layer_init_method,
layer_number=i,
rpe=rpe_emb
if self.neox_args.pos_emb == "rpe" else None,
rotary=self.neox_args.pos_emb == "rotary",
use_cache=self.use_cache,
))
# used to drop attention mask + reshape hidden states
self.specs.append(_post_transformer_block)
# NormPipe is a (deprecated) helper class that used to be used to pass presents along the pipeline - since presents are now cached to the `TransformerLayer` class this is no longer needed
norm, eps = get_norm(self.neox_args)
self.specs.append(
LayerSpec(NormPipe, norm, self.neox_args.hidden_size, eps=eps))
# outputs are now a single tensor: hidden_states
def _logits_helper(embedding, lm_output):
"""Just a wrapper to massage inputs/outputs from pipeline."""
logits = parallel_lm_logits(lm_output,
embedding.word_embeddings_weight,
self.parallel_output)
return logits
if weight_tying:
self.specs.append(
TiedLayerSpec(
"embed",
EmbeddingPipe,
self.neox_args,
self.hidden_size,
self.neox_args.padded_vocab_size,
self.neox_args.max_position_embeddings,
self.neox_args.hidden_dropout,
self.init_method,
self.num_tokentypes,
forward_fn=_logits_helper,
tied_weight_attr="word_embeddings_weight",
))
else:
self.specs.append(
LayerSpec(
ParallelLinearPipe,
neox_args=self.neox_args,
init_method=self.init_method,
parallel_output=self.parallel_output,
))
def __init__(self, attention_mask_func,
init_method, output_layer_init_method):
super(ParallelTransformer, self).__init__()
args = get_args()
# Store activation checkpoiting flag.
self.checkpoint_activations = args.checkpoint_activations
self.checkpoint_num_layers = args.checkpoint_num_layers
# Number of layers:
self.num_layers = args.num_layers
self.num_unique_layers = args.num_unique_layers
if self.num_unique_layers is None:
self.num_unique_layers = self.num_layers
assert self.num_layers % self.num_unique_layers == 0, \
'number of layers should be divisible by number of unique layers'
self.param_sharing_style = args.param_sharing_style
if args.pos_emb == 'rpe':
rpe_emb = ParallelRelativePositionBias(causal=True, num_buckets=args.rpe_num_buckets,
max_distance=args.rpe_max_distance,
heads=args.num_attention_heads)
# Transformer layers.
sparsity = args.sparsity
def build_layer(layer_number):
if sparsity == 'none':
sparse = False
elif sparsity == 'all':
sparse = True
elif sparsity == 'interspersed':
sparse = not layer_number % 2 == 0
else:
raise ValueError(f'Sparsity type {sparsity} not recognized')
return ParallelTransformerLayer(
attention_mask_func, init_method,
output_layer_init_method, layer_number, sparse=sparse, rpe=rpe_emb if args.pos_emb == 'rpe' else None)
self.layers = torch.nn.ModuleList(
[build_layer(i + 1) for i in range(self.num_unique_layers)])
# Print layer ordering.
if self.num_layers != self.num_unique_layers:
if torch.distributed.get_rank() == 0:
print('> will be using the following layer ordering:')
for i in range(self.num_layers):
print(' layer id: {:3d} --> unique layer id: '
'{:3d}'.format(i, self._get_layer_index(i)),
flush=True)
# Final layer norm before output.
if args.norm == "rmsnorm":
norm = RMSNorm
eps = args.rms_norm_epsilon
elif args.norm == "layernorm":
eps = args.layernorm_epsilon
norm = LayerNorm
elif args.norm == "scalenorm":
eps = args.scalenorm_epsilon
norm = ScaleNorm
self.final_layernorm = norm(
args.hidden_size,
eps=eps)
if deepspeed.checkpointing.is_configured():
global get_cuda_rng_tracker, checkpoint
get_cuda_rng_tracker = deepspeed.checkpointing.get_cuda_rng_tracker
checkpoint = deepspeed.checkpointing.checkpoint
def init_specs(self):
weight_tying = not self.neox_args.no_weight_tying
if self.embedding_type == 'rpe':
rpe_emb = ParallelRelativePositionBias(
neox_args=self.neox_args,
causal=True,
num_buckets=self.neox_args.rpe_num_buckets,
max_distance=self.neox_args.rpe_max_distance,
heads=self.neox_args.num_attention_heads)
self.specs = []
# Embedding layer
# input will be (input_ids, position_ids, attention_mask) in Training
# and (input_ids, position_ids, attention_mask, layer_past) in Inference
if weight_tying:
self.specs.append(
TiedLayerSpec('embed',
EmbeddingPipe,
self.neox_args,
self.hidden_size,
self.neox_args.padded_vocab_size,
self.neox_args.max_position_embeddings,
self.neox_args.hidden_dropout,
self.init_method,
self.num_tokentypes,
tied_weight_attr='word_embeddings_weight'))
else:
self.specs.append(
LayerSpec(EmbeddingPipe, self.neox_args, self.hidden_size,
self.neox_args.padded_vocab_size,
self.neox_args.max_position_embeddings,
self.neox_args.hidden_dropout, self.init_method,
self.num_tokentypes))
# NB: in inference, the attention mask always needs to be the *last* item in the args when being passed from
# one stage to the next, because deepspeed is hacks on top of hacks.
#
# outputs are now
# Train: (hidden_states, attention_mask)
# Inference: (hidden_states, layer_past, attention_mask)
self.specs.append(_pre_transformer_block)
# Transformer layers
for i in range(self.neox_args.num_layers):
layer_type = self.neox_args.attention_config[i]
if layer_type in ["gmlp", "amlp"]:
self.specs.append(
LayerSpec(
GMLPBlock,
init_method=self.init_method,
layer_number=i,
output_layer_init_method=self.output_layer_init_method,
neox_args=self.neox_args,
mask_fn=gpt2_attention_mask_func))
else:
self.specs.append(
LayerSpec(
ParallelTransformerLayerPipe,
neox_args=self.neox_args,
attention_mask_func=gpt2_attention_mask_func,
init_method=self.init_method,
output_layer_init_method=self.output_layer_init_method,
layer_number=i,
rpe=rpe_emb
if self.neox_args.pos_emb == 'rpe' else None,
rotary=self.neox_args.pos_emb == 'rotary',
get_key_value=self.get_key_value))
self.specs.append(_post_transformer_block)
# NormPipe is a helper class to pass presents through to the output when doing inference
norm, eps = get_norm(self.neox_args)
self.specs.append(
LayerSpec(NormPipe, norm, self.neox_args.hidden_size, eps=eps))
# outputs are now
# Train: hidden_states
# Inference: (hidden_states, presents)
def _logits_helper(embedding, lm_output):
"""Just a wrapper to massage inputs/outputs from pipeline. """
if self._inference and len(lm_output) == 2:
hidden_states, presents = lm_output
logits = parallel_lm_logits(hidden_states,
embedding.word_embeddings_weight,
self.parallel_output)
return logits, presents
else:
logits = parallel_lm_logits(lm_output,
embedding.word_embeddings_weight,
self.parallel_output)
return logits
if weight_tying:
self.specs.append(
TiedLayerSpec('embed',
EmbeddingPipe,
self.neox_args,
self.hidden_size,
self.neox_args.padded_vocab_size,
self.neox_args.max_position_embeddings,
self.neox_args.hidden_dropout,
self.init_method,
self.num_tokentypes,
forward_fn=_logits_helper,
tied_weight_attr='word_embeddings_weight'))
else:
self.specs.append(
LayerSpec(ParallelLinearPipe,
neox_args=self.neox_args,
init_method=self.init_method,
parallel_output=self.parallel_output))