[wip] Entmax loss

megatron/arguments.py

@@ -83,7 +83,7 @@ def validate_args(args):
  if args.no_pipeline_parallel:
  assert args.pipeline_model_parallel_size == 1, \
  "pipeline_model_parallel_size must be 1 if pipeline parallel is disabled"
- 
+
  if args.ds_sequence_parallel_size > 1:
  assert version.parse(deepspeed.__version__) >= version.parse("0.10.2"), "sequence parallelism requires DeepSpeed version 0.10.2+"
 
@@ -432,6 +432,10 @@ def validate_args(args):
  assert not args.mos, 'GQA currently does not support args.mos'
  assert not args.kd, 'GQA currently does not support args.kd'
 
+ # entmax loss
+ if args.loss_function != "cross_entropy":
+ assert not args.fp16_lm_cross_entropy, "entmax loss only supports fp32"
+
  # Print arguments.
  _print_args("arguments", args)
  retro_args = get_retro_args()
@@ -477,7 +481,7 @@ def core_transformer_config_from_args(args):
  kw_args['bias_gelu_fusion'] = False
  if args.init_method_xavier_uniform:
  kw_args['init_method'] = torch.nn.init.xavier_uniform_
- kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_
+ kw_args['output_layer_init_method'] = torch.nn.init.xavier_uniform_
 
  return TransformerConfig(**kw_args)
 
@@ -632,7 +636,7 @@ def _add_network_size_args(parser):
  group.add_argument('--apply-layernorm-1p', action='store_true',
  help='Adjust LayerNorm weights such that they are centered '
  'around zero. This improves numerical stability.')
- group.add_argument('--disable-mem-efficient-ln', action='store_false', 
+ group.add_argument('--disable-mem-efficient-ln', action='store_false',
  help='Disable the memory-efficient fused LayerNorm optimization '
  'introduced in https://github.com/NVIDIA/apex/pull/1715')
  group.add_argument('--apply-residual-connection-post-layernorm',
@@ -848,7 +852,7 @@ def _add_training_args(parser):
  'training runs.')
  group.add_argument('--random-ltd',
  action='store_true',
- help='enable random layer token drop') 
+ help='enable random layer token drop')
  group.add_argument('--log-interval', type=int, default=100,
  help='Report loss and timing interval.')
  group.add_argument('--exit-interval', type=int, default=None,
@@ -940,6 +944,15 @@ def _add_training_args(parser):
  dest='gradient_accumulation_fusion')
  group.add_argument('--use-dataset-only', type=bool, required=False, default=False,
  help='If set to True, only use the megatron dataset for external trainer ')
+ group.add_argument('--loss-function', default='cross_entropy',
+ choices=['cross_entropy', 'entmax15', 'sparsemax', 'entmax_bisect'],
+ help='Loss function for model training')
+ group.add_argument('--entmax-alpha', type=float, default=1.5,
+ help='Entmax alpha for entmax_bisect (unused otherwise)')
+ group.add_argument('--entmax-topk', type=int, default=512,
+ help='Top k for computation of exact entmax loss (for entmax15 and sparsemax)')
+ group.add_argument('--entmax-n-iter', type=int, default=30,
+ help='Number of bisection interations for entmax_bisect')
  return parser
 
 
@@ -1034,7 +1047,7 @@ def _add_checkpointing_args(parser):
  group.add_argument('--no-load-rng', action='store_true', default=None,
  help='Do not load rng state when loading checkpoint.')
  group.add_argument('--no-load-lr-state', action='store_true',
- help='Do not load lr state when loading checkpoint.') 
+ help='Do not load lr state when loading checkpoint.')
  group.add_argument('--finetune', action='store_true',
  help='Load model for finetuning. Do not load optimizer '
  'or rng state from checkpoint and set iteration to 0. '
@@ -1210,7 +1223,7 @@ def _add_data_args(parser):
  'form: dataset1-weight dataset1-path dataset2-weight '
  'dataset2-path ...')
  group.add_argument('--multiple-valid-sets', action='store_true',
- help='multiple separated validation steps') 
+ help='multiple separated validation steps')
  group.add_argument('--test-data-path', nargs='*', default=None,
  help='Path to the test dataset. Accepted format:'
  '1) a single data path, 2) multiple datasets in the'
@@ -1490,15 +1503,15 @@ def _add_activation_checkpoint_args(parser):
 def _add_distillation_args(parser):
  group = parser.add_argument_group('Knowledge distillation',
  'Distillation Configurations')
- 
+
  group.add_argument('--num-layers-teacher', type=int, default=None,
- help='Number of the teacher transformer layers.') 
+ help='Number of the teacher transformer layers.')
  group.add_argument('--num-experts-teacher', type=int, nargs='+', default=[1,],
  help='number of teacher experts list, MoE related.')
  group.add_argument('--hidden-size-teacher', type=int, default=None,
  help='Tansformer teacher hidden size.')
  group.add_argument('--num-attention-heads-teacher', type=int, default=None,
- help='Number of teacher transformer attention heads.') 
+ help='Number of teacher transformer attention heads.')
 
  group.add_argument('--mos', action='store_true',
  help='Enable Mixture-of-Students via knolwedge distillation.')
@@ -1509,7 +1522,7 @@ def _add_distillation_args(parser):
  group.add_argument('--kd-temp', default=1.0, type=float)
  group.add_argument('--reset-iteration', action='store_true',
  help='Reset the iteration count.')
- 
+
  group.add_argument('--load-teacher', type=str, default=None,
  help='Directory containing a teacher model checkpoint.')
 

megatron/model/gpt_model.py

@@ -2,7 +2,10 @@
 
 """GPT-2 model."""
 
+from functools import partial
+
 import torch
+import entmax
 
 from megatron import get_args
 from megatron.core import mpu, tensor_parallel, sequence_parallel
@@ -24,32 +27,36 @@
 except ImportError:
  MixedFusedRMSNorm = None
 
-try: 
+try:
  from deepspeed.checkpoint import (
  VOCABULARY_PARAMETER_PATTERNS,
  PIPELINE_REPLICATED_PARAMETER_PATTERNS,
  TP_REPLICATED_PARAMETER_PATTERNS,
  PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
  )
- DS_UNIVERSAL_CHECKPOINT_INFO = True 
+ DS_UNIVERSAL_CHECKPOINT_INFO = True
 except ImportError:
- DS_UNIVERSAL_CHECKPOINT_INFO = False 
+ DS_UNIVERSAL_CHECKPOINT_INFO = False
 
 
 def post_language_model_processing(lm_output, labels, logit_weights,
  parallel_output,
- fp16_lm_cross_entropy):
+ fp16_lm_cross_entropy,
+ loss_function, alpha, topk, n_iter, return_support_size=False):
 
  # Output. Format [s b h]
  output = parallel_lm_logits(
  lm_output,
  logit_weights,
  parallel_output)
 
+ # should it return a None support size in this case? let's say no for now
  if labels is None:
  # [s b h] => [b s h]
  return output.transpose(0,1).contiguous()
- else:
+
+ if loss_function == "cross_entropy":
+ # cross entropy
  # [b s] => [s b]
  labels = labels.transpose(0,1).contiguous()
  cross_entropy = sequence_parallel.vocab_sequence_parallel_cross_entropy if mpu.get_sequence_parallel_world_size() > 1 \
@@ -62,6 +69,42 @@ def post_language_model_processing(lm_output, labels, logit_weights,
 
  # [s b] => [b, s]
  loss = loss.transpose(0,1).contiguous()
+ support = None
+ else:
+ # now: the loss function is "entmax15", "sparsemax", or "entmax_bisect"
+ loss_funcs = {
+ "entmax15": partial(entmax.entmax15_loss, k=topk, return_support_size=True),
+ "sparsemax": partial(entmax.sparsemax_loss, k=topk, return_support_size=True),
+ "entmax_bisect": partial(entmax.entmax_bisect_loss, alpha=alpha, n_iter=n_iter)
+ }
+ f = loss_funcs[loss_function]
+ b, s = labels.size()
+ output = output.transpose(0, 1).contiguous()
+ vocab_size = output.size(-1)
+
+ # currently entmax_bisect_loss always returns a None support size,
+ # which is not ideal. This is a stopgap until entmax_bisect_loss is
+ # fixed
+ loss = f(output.float().view(-1, vocab_size), labels.view(-1))
+ if isinstance(loss, tuple):
+ loss, support = loss
+ else:
+ support = None
+
+ # old version which breaks because entmax_bisect unexpectedly returned
+ # a tuple:
+ '''
+ if loss_function != "entmax_bisect":
+ loss, support = f(output.float().view(-1, vocab_size), labels.view(-1))
+ else:
+ loss = f(output.float().view(-1, vocab_size), labels.view(-1))
+ support = None
+ '''
+ loss = loss.view(b, s)
+
+ if return_support_size:
+ return loss, support
+ else:
  return loss
 
 
@@ -84,6 +127,10 @@ def __init__(self,
  self.fp16_lm_cross_entropy = args.fp16_lm_cross_entropy
  self.return_moe_loss = return_moe_loss
  self.untie_embeddings_and_output_weights = args.untie_embeddings_and_output_weights
+ self.loss_function = args.loss_function
+ self.entmax_alpha = args.entmax_alpha
+ self.entmax_topk = args.entmax_topk
+ self.entmax_n_iter = args.entmax_n_iter
 
  self.language_model, self._language_model_key = get_language_model(
  config=config,
@@ -106,7 +153,8 @@ def forward(self, input_ids, position_ids, attention_mask,
  retriever_position_ids=None,
  retriever_attn_mask=None,
  labels=None, tokentype_ids=None, inference_params=None,
- curriculum_seqlen=None):
+ curriculum_seqlen=None,
+ return_support_size=False):
  args = get_args()
  if curriculum_seqlen is not None:
  args.curriculum_seqlen = curriculum_seqlen
@@ -135,13 +183,25 @@ def forward(self, input_ids, position_ids, attention_mask,
  inference_params=inference_params)
 
  if self.post_process:
- lm_output = post_language_model_processing(
+ post_lm_out = post_language_model_processing(
  lm_output, labels,
  self.language_model.output_layer.weight if self.untie_embeddings_and_output_weights else self.shared_embedding_or_output_weight(),
  self.parallel_output,
- self.fp16_lm_cross_entropy)
-
- return lm_output, moe_losses if self.return_moe_loss else lm_output
+ self.fp16_lm_cross_entropy,
+ self.loss_function,
+ self.entmax_alpha,
+ self.entmax_topk,
+ self.entmax_n_iter,
+ return_support_size=True)
+ if labels is not None:
+ lm_output, support_size = post_lm_out
+ else:
+ lm_output = post_lm_out
+ # now...what do do about support_size?
+ if return_support_size:
+ return lm_output, moe_losses if self.return_moe_loss else lm_output, support_size
+ else:
+ return lm_output, moe_losses if self.return_moe_loss else lm_output
 
  def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False):
 
@@ -210,7 +270,7 @@ def universal_checkpoint_info(self):
  ]
 
  return info
- 
+
 def CrossEntropy(output, labels):
  labels, loss_mask = labels[0], labels[1]
 

pretrain_gpt.py

@@ -211,7 +211,7 @@ def get_batch_pipe(data):
  return (tokens, position_ids, attention_mask), (labels, loss_mask)
 
 
-def loss_func(loss_mask, moe_loss, mos_loss, output_tensor):
+def loss_func(loss_mask, moe_loss, mos_loss, support_size, output_tensor):
  args = get_args()
  losses = output_tensor.float()
  loss_mask = loss_mask.view(-1).float()
@@ -229,11 +229,68 @@ def loss_func(loss_mask, moe_loss, mos_loss, output_tensor):
  print_rank_0('>>> total loss: {}, lm loss {}, kd loss {}'.format(loss, averaged_loss[0], mos_loss))
  else:
  if max(args.num_experts) <= 1:
- return loss, {'lm loss': averaged_loss[0]}
+ if support_size is not None:
+ support_size = support_size.float()
+ n_tokens = support_size.size(0)
+
+ # how often is support size 1?
+ fully_peaked = support_size.eq(1).sum()
+
+ # what is the max support size in the batch?
+ max_support = support_size.max()
+ min_support = support_size.min()
+
+ # sum support stats across groups
+ current_group = mpu.get_data_parallel_group()
+ torch.distributed.all_reduce(support_size, group=current_group)
+ torch.distributed.all_reduce(max_support, group=current_group)
+ torch.distributed.all_reduce(min_support, group=current_group)
+ torch.distributed.all_reduce(fully_peaked, group=current_group)
+
+ # find number of groups
+ world_size = torch.distributed.get_world_size(
+ group=mpu.get_data_parallel_group()
+ )
+
+ # compute mean support size
+ support_size = support_size / world_size
+
+ # compute mean max support
+ max_support = max_support / world_size
+ min_support = min_support / world_size
+
+ # compute how often support is fully peaked
+ fully_peaked = fully_peaked / (world_size * n_tokens)
+
+ # the stats for support size might be slightly wrong, but still
+ # illustrative
+ support_mean = support_size.mean()
+ support_std = support_size.std()
+ quantiles = torch.linspace(
+ 0, 1, 5, dtype=support_size.dtype, device=support_size.device
+ )
+ support_quantiles = torch.quantile(support_size, quantiles)
+
+ loss_dict = {
+ 'lm loss': averaged_loss[0],
+ 'support size': support_mean,
+ "support std": support_std,
+ "support max": max_support,
+ "support min": min_support,
+ "support 25%": support_quantiles[1],
+ "support 50%": support_quantiles[2],
+ "support 75%": support_quantiles[3],
+ "fully peaked": fully_peaked
+ }
+
+ return loss, loss_dict
+ else:
+ return loss, {'lm loss': averaged_loss[0]}
  else:
  loss = loss + moe_loss
  return loss, {'lm loss': averaged_loss[0], 'moe loss': moe_loss}
 
+
 def calculate_mos_loss(args, stu_output, teacher_model, tokens, position_ids, attention_mask):
  mos_loss = 0
  alpha = args.kd_alpha_ce
@@ -277,6 +334,7 @@ def forward_step(data_iterator, model):
  args.data_efficiency_curriculum_learning_seqlen_type == 'seqlen_reshape':
  args.data_efficiency_curriculum_learning_numel = torch.numel(tokens)
 
+ support_size = None
  if args.mos or args.kd:
  # The forward func can return either the loss or the logits, depending on whether passing in the labels or not.
  stu_output, other_losses = model(tokens, position_ids, attention_mask)
@@ -285,8 +343,9 @@ def forward_step(data_iterator, model):
  labels = labels[:, :args.curriculum_seqlen].contiguous()
  output_tensor = tensor_parallel.vocab_parallel_cross_entropy(stu_output.contiguous().float(), labels)
  else:
- output_tensor, other_losses = model(tokens, position_ids, attention_mask,
- labels=labels)
+ output_tensor, other_losses, support_size = model(
+ tokens, position_ids, attention_mask, labels=labels, return_support_size=True
+ )
  if args.curriculum_learning_legacy and args.curriculum_seqlen < args.seq_length:
  loss_mask = loss_mask[:, :args.curriculum_seqlen].contiguous()
 
@@ -304,7 +363,7 @@ def forward_step(data_iterator, model):
  args.teacher_model[0], tokens, position_ids, attention_mask)
 
  # Output_tensor stores the standard loss, loos_func calculates the total loss.
- return output_tensor, partial(loss_func, loss_mask, moe_loss, mos_loss)
+ return output_tensor, partial(loss_func, loss_mask, moe_loss, mos_loss, support_size)
 
 
 def train_valid_test_datasets_provider(train_val_test_num_samples):

tasks/main.py

@@ -70,7 +70,11 @@ def get_tasks_args(parser):
  group.add_argument('--val-av-rank-other-neg', type=int, default=30,
  help='Av.rank validation: how many other negatives to'
  ' take from each question pool')
-
+ group.add_argument('--eval-metric', default=None,
+ help='Eval metric to use other than a task-specific'
+ 'default')
+ group.add_argument('--acc-k', default=5, type=int,
+ help='k for force-decoded accuracy at k')
 
  return parser
 
@@ -79,7 +83,7 @@ def get_tasks_args(parser):
 
  initialize_megatron(extra_args_provider=get_tasks_args)
 
- args = get_args()
+ args = get_args() # will the task args be included here?
 
  if args.num_layers_per_virtual_pipeline_stage is not None:
  print("Interleaved pipeline schedule is not yet supported for downstream tasks.")