Add option to accumulate train loss over tokens.

[aadyotb]

What does this PR do?

Adds an option to accumulate train loss over the number of tokens in a batch instead of the number of samples.

What issue(s) does this change relate to?

Currently, losses in the trainer are accumulated in a way that weights each sample in a batch equally. However, for NLP use cases where batches contain padding tokens, it makes more sense to accumulate the loss in a way that instead weights every (non-padding) token equally.

Before submitting

• Have you read the contributor guidelines?
• Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[mvpatel2000]

@dakinggg @irene-dea can you please look?

I agree we should have an option for this. I'm not sure if it's necessary to pass to Composer vs. check if its an attribute/property on a dataloader

[dakinggg]

@mvpatel2000 I think trainer arg is right for this...code looks fine at a glance, would want to test a bit more before merging.

[mvpatel2000]

@mvpatel2000 I think trainer arg is right for this...code looks fine at a glance, would want to test a bit more before merging.

You will own testing?

[mvpatel2000]

Can we add a unit test please?

[aadyotb]

Can we add a unit test please?

Is there a good template I can base this test on? I'm not sure how to isolate the impact of this change from the rest of the trainer.

[dakinggg]

@aadyotb Here is a test that exercises the full NLP pipeline (

composer/tests/test_full_nlp.py

Lines 231 to 338 in 01eec3a

	@device('cpu', 'gpu')
	# Note: the specificity of these settings are due to incompatibilities (e.g. the simpletransformer model is not traceable)
	@pytest.mark.parametrize(
	'model_type,algorithms,save_format',
	[
	('tinybert_hf', [GatedLinearUnits], 'onnx'),
	('simpletransformer', [], 'torchscript'),
	],
	)
	@pytest.mark.parametrize('onnx_opset_version', [13, None])
	def test_full_nlp_pipeline(
	model_type,
	algorithms,
	save_format,
	tiny_bert_tokenizer,
	onnx_opset_version,
	tmp_path,
	request,
	device,
	):
	"""This test is intended to exercise our full pipeline for NLP.
	To this end, it performs pretraining, loads the pretrained model with a classification head for finetuning
	and finetunes it, exports the model for inference, and loads it back in to make predictions.
	"""
	pytest.importorskip('libcloud')
	pytest.importorskip('transformers')
	if onnx_opset_version == None and version.parse(torch.__version__) < version.parse('1.13'):
	pytest.skip("Don't test prior PyTorch version's default Opset version.")
	algorithms = [algorithm() for algorithm in algorithms]
	device = get_device(device)
	tiny_bert_model = None
	if model_type == 'tinybert_hf':
	tiny_bert_model = request.getfixturevalue('tiny_bert_model')
	# pretraining
	if model_type == 'tinybert_hf':
	assert tiny_bert_model is not None
	pretraining_metrics = [LanguageCrossEntropy(ignore_index=-100), MaskedAccuracy(ignore_index=-100)]
	pretraining_model = HuggingFaceModel(
	tiny_bert_model,
	tiny_bert_tokenizer,
	use_logits=True,
	metrics=pretraining_metrics,
	)
	elif model_type == 'simpletransformer':
	pretraining_model = SimpleTransformerMaskedLM(vocab_size=tiny_bert_tokenizer.vocab_size)
	else:
	raise ValueError('Unsupported model type')
	pretraining_output_path = pretraining_test_helper(
	tiny_bert_tokenizer,
	pretraining_model,
	algorithms,
	tmp_path,
	device,
	)
	# finetuning
	if model_type == 'tinybert_hf':
	finetuning_metric = MulticlassAccuracy(num_classes=3, average='micro')
	hf_finetuning_model, _ = HuggingFaceModel.hf_from_composer_checkpoint(
	pretraining_output_path,
	model_instantiation_class='transformers.AutoModelForSequenceClassification',
	model_config_kwargs={'num_labels': 3},
	)
	finetuning_model = HuggingFaceModel(
	model=hf_finetuning_model,
	tokenizer=tiny_bert_tokenizer,
	use_logits=True,
	metrics=[finetuning_metric],
	)
	elif model_type == 'simpletransformer':
	finetuning_model = SimpleTransformerClassifier(vocab_size=tiny_bert_tokenizer.vocab_size, num_classes=3)
	else:
	raise ValueError('Unsupported model type.')
	finetuning_model_copy = copy.deepcopy(finetuning_model)
	finetuning_trainer, finetuning_dataloader, rud, finetuning_output_path = finetuning_test_helper(
	tiny_bert_tokenizer,
	finetuning_model,
	algorithms,
	pretraining_output_path,
	pretraining_model,
	tmp_path,
	device,
	)
	# inference
	batch = next(iter(finetuning_dataloader))
	finetuning_trainer.state.model.to('cpu')
	finetuning_trainer.state.model.eval()
	original_output = finetuning_trainer.state.model(batch)
	inference_test_helper(
	finetuning_output_path,
	rud,
	finetuning_model_copy,
	algorithms,
	batch,
	original_output,
	onnx_opset_version,
	tmp_path,
	save_format,
	device,
	)

). I think to test this we probably want to (for just the training part of the code that I linked) construct a model that has deterministic loss (based on num padding tokens maybe?) and then test that the results are different in the expected way between sample weighting and token weighting.

[dakinggg]

So basically make a trainer with a dummy model and a dummy dataset, and then call it with sample weighting and token weighting (with microbatching on), and assert the losses are different in the expected way.

[aadyotb]

@dakinggg I've added a unit test that requires sample-based and token-based weighting result in different outcomes when padding is present.

[dakinggg]

@aadyotb awesome, thank you!! Will take a look soon.

[aadyotb]

Thanks @mvpatel2000. For the time being, I've implemented these changes by overriding the Trainer class in our local repo, so we will be okay for now. Happy to get further review once Daniel returns.

[aadyotb]

@dakinggg bumping this PR for review.

[mvpatel2000]

LGTM but requiring sign off from @dakinggg as well

[mvpatel2000]

will let daniel review the unit test

[dakinggg]

Hey @aadyotb taking a look now, mostly convincing myself that the math is correct :)

[dakinggg]

Other than waiting for Mihir's thought on changing the behavior for the sample weighting case, looks good to me! Thanks so much for the contribution!

[mvpatel2000]

Approving

[dakinggg]

Ok cool, LGTM. I'm running a before and after loss curve just to double check (in the normal case, with even number of samples per device batch), and will post that graph here when done.

[dakinggg]

For posterity, I validated that finetuning behavior is unchanged before and after this PR (for a case with constant samples per device batch), and does change if you specify the new flag.

[dakinggg]

@aadyotb I think one more run of precommit should do it, and we should be good to merge.