Fix Gradient Accumulation issue

[ArthurZucker]

What does this PR do?

First draft

End goal is to make it easy for anyone to:

• change the loss for his model
• contribute a new loss for a model (like vision model, ENCODEC etc)
• allow passing arbitrary kwargs, interfacing

TODO:

• Fix deformable detr loss computation

[BenjaminBossan]

Thanks for coming forward with this fix so quickly. There is probably not much I can help with, but I took a look and added some comments.

[paulcx]

As hiyouga/LLaMA-Factory#5747 (comment) suggested, maybe a few changes will make this fix work again?

[iridescentee]

After using this fix, the following problem occurred: the grad norm and loss (custom loss from trl cpo trainer) increased dramatically
green：llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix orange: llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix

Can confirm the same issue
same. encountered when fine-tuning using qwen2.5.

[muellerzr]

Fix is being worked on here: #34283

[thusinh1969]

Oh my....

[taehyunzzz]

Hello all, I could not follow the whole context of the PR here, but I am guessing a fix is underway regarding the loss becoming huge due to this code. The code scales the loss by gradient accumulation step in the latest trainer (not in previous version). The snippet is included in a single training iteration step, for which I think scaling the loss by grad_acc_step is not appropriate, if the loss is not de-scaled by grad_acc_step inside the accelerator.backward call in the following line.

Does the latest accelerator backward de-scale by grad_acc_step (or is being fixed to do so)? If so, maybe the output loss of training_step() should be de-scaled one more time by grad_acc_step?
For now, is it correct to scale down the learning rate by the grad_acc_step for the original functionality?

[muellerzr]

Yes accelerate de-scales the loss when calling backward() which is why we do so. (Accelerate has always done this)

[taehyunzzz]

@muellerzr Thank you for the reply. I have just one more question to clear things up.

You've mentioned that accelerate had always internally de-scaled the loss by gradient accumulation steps.
However, the previous trainer.py implementations do not scale (scale up) the loss by grad_acc_steps before accelerator.backward.
If accelerator.backward had been internally de-scaling the loss by grad_acc_steps, does that mean the previous trainer had been training a model on 1/grad_acc_steps * sum(loss) instead of sum(loss) when grad_acc_steps > 1?

[cphillippi-stripe]

Hi, @ArthurZucker,

I'm seeing issues w/ multi-gpu sequence classification training after this change (using Mistral). I believe it is due to the removal of this line (which I'm having a very hard time monkeypatching via loss_function for some reason). I also see this line very often in quite a few different models here. Were multi-gpu setups tested here for training in sequence classification mode? I'm really curious how this passed if so, because I don't see any lines like this in the new ForSequenceClassificationLoss function that self.loss_function now seems to resolve to.

I'm seeing tracebacks like this (edited for privacy):

...
215       File ".../training_workflow.py", line 426, in train
216     trainer.train()
217       File ".../transformers/transformers/transformers/trainer.py", line 2123, in train
218     return inner_training_loop(
219            ^^^^^^^^^^^^^^^^^^^^
220       File ".../transformers/transformers/transformers/trainer.py", line 2481, in _inner_training_loop
221     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
222                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
223       File ".../transformers/transformers/transformers/trainer.py", line 3579, in training_step
224     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
225            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
226       File ".../transformers/transformers/transformers/trainer.py", line 3633, in compute_loss
227     outputs = model(**inputs)
228               ^^^^^^^^^^^^^^^
229       File ".../torch/torch/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
230     return self._call_impl(*args, **kwargs)
231            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
232       File ".../torch/torch/torch/nn/modules/module.py", line 1541, in _call_impl
233     return forward_call(*args, **kwargs)
234            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
235       File ".../accelerate/accelerate/accelerate/utils/operations.py", line 823, in forward
236     return model_forward(*args, **kwargs)
237            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
238       File ".../accelerate/accelerate/accelerate/utils/operations.py", line 811, in __call__
239     return convert_to_fp32(self.model_forward(*args, **kwargs))
240                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
241       File ".../torch/torch/torch/amp/autocast_mode.py", line 16, in decorate_autocast
242     return func(*args, **kwargs)
243            ^^^^^^^^^^^^^^^^^^^^^
244       File ".../peft/peft/peft/peft_model.py", line 1446, in forward
245     return self.base_model(
246            ^^^^^^^^^^^^^^^^
247       File ".../torch/torch/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
248     return self._call_impl(*args, **kwargs)
249            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
250       File ".../torch/torch/torch/nn/modules/module.py", line 1541, in _call_impl
251     return forward_call(*args, **kwargs)
252            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
253       File ".../peft/peft/peft/tuners/tuners_utils.py", line 197, in forward
254     return self.model.forward(*args, **kwargs)
255            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
256       File ".../accelerate/accelerate/accelerate/hooks.py", line 170, in new_forward
257     output = module._old_forward(*args, **kwargs)
258              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
259       File ".../transformers/transformers/transformers/models/mistral/modeling_mistral.py", line 1200, in forward
260     loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
261            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
262       File ".../transformers/transformers/transformers/loss/loss_utils.py", line 67, in ForSequenceClassificationLoss
263     loss = fixed_cross_entropy(pooled_logits.view(-1, num_labels), labels.view(-1), **kwargs)
264            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
265       File ".../transformers/transformers/transformers/loss/loss_utils.py", line 26, in fixed_cross_entropy
266     loss = nn.functional.cross_entropy(source, target, ignore_index=ignore_index, reduction=reduction)
267            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
268       File ".../torch/torch/torch/nn/functional.py", line 3086, in cross_entropy
269     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
270            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
271 
272 Message:
273 
274     RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:7 and cuda:0! (when checking argument for argument target in method wrapper_CUDA_nll_loss_forward)

For now, if you think you have a good way I can just patch this, I'd appreciate it, but replacing PretrainedModel.loss_function doesn't seem to do the trick...

[bauwenst]

FYI: As also noted by @cphillippi-stripe, the @property PreTrainedModel.loss_function that was added by this PR breaks all software packages that have models which (1) inherit from PreTrainedModel while (2) setting a loss_function beforehand.

I maintain such a package, and indeed, training is now broken and my IDE complains that loss_function cannot be set.

The default behaviour of setting the loss to causal LM loss when no loss_type is defined in the config seems highly unwarranted to me. It is assumed now that either (1) models define this brand new loss_type field in the config and otherwise (2) the loss type can be inferred from the ForXYZ suffix:

transformers/src/transformers/modeling_utils.py

Lines 5110 to 5120 in 63766ab

	if getattr(self.config, "loss_type", None) is not None:
	loss_type = self.config.loss_type
	else:
	loss_type = self.__class__.__name__
	if loss_type not in LOSS_MAPPING:
	loss_groups = f"({'|'.join(LOSS_MAPPING)})"
	loss_type = re.findall(loss_groups, self.__class__.__name__)
	if len(loss_type) > 0:
	loss_type = loss_type[0]
	else:
	loss_type = None

Yet, this fails to take into account the following cases where transformers is used as a dependency:

1. Model classes defined for tasks that don't have a known auto class in transformers. For example: if your package defines a ModelForDependencyParsing, the above regex search will not find its loss and default to causal loss, which is wrong.

  File ".../ArchIt/src/archit/instantiation/tasks.py", line 328, in computeLoss
    return self.loss_function(arc_scores, arc_labels) + self.loss_function(rel_scores, rel_labels)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: ForCausalLMLoss() missing 1 required positional argument: 'vocab_size'

2. Model classes that have a known task, but it's not in the name (e.g. ModelForSentenceClassifying). Unless the config is from after October 2024, the config does not define its loss type.

3. Loss is not inherent to a PreTrainedModel. There are plenty of papers that start from the model weights of one model and then change its loss function for continued pretraining (take this one as a recent example).

TL;DR: Seems like a design flaw to make loss_function a @property and seems like a mistake to overwrite it to ForCausalLMLoss for all models that don't fit the very narrow range of tasks supported by transformers.

[cphillippi-stripe]

FYI: As also noted by @cphillippi-stripe, the @property PreTrainedModel.loss_function that was added by this PR breaks all software packages that have models which (1) inherit from PreTrainedModel while (2) setting a loss_function beforehand.

I maintain such a package, and indeed, training is now broken and my IDE complains that loss_function cannot be set.

The default behaviour of setting the loss to causal LM loss when no loss_type is defined in the config seems highly unwarranted to me. It is assumed now that either models define this brand new loss_type field in the config and that otherwise the loss type can be inferred from the ForXYZ suffix:

transformers/src/transformers/modeling_utils.py

Lines 5110 to 5120 in 63766ab

	if getattr(self.config, "loss_type", None) is not None:
	loss_type = self.config.loss_type
	else:
	loss_type = self.__class__.__name__
	if loss_type not in LOSS_MAPPING:
	loss_groups = f"({'|'.join(LOSS_MAPPING)})"
	loss_type = re.findall(loss_groups, self.__class__.__name__)
	if len(loss_type) > 0:
	loss_type = loss_type[0]
	else:
	loss_type = None

Yet, this fails to take into account the following cases where transformers is used as a dependency:

1. Model classes defined for tasks that don't have a known auto class in transformers. For example: if your package defines a ModelForDependencyParsing, the above regex search will not find its loss and default to causal loss, which is wrong.

  File ".../ArchIt/src/archit/instantiation/tasks.py", line 328, in computeLoss
    return self.loss_function(arc_scores, arc_labels) + self.loss_function(rel_scores, rel_labels)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: ForCausalLMLoss() missing 1 required positional argument: 'vocab_size'

2. Model classes that have a known task, but it's not in the name (e.g. ModelForSentenceClassifying). Unless the config is from pre-October 2024, the config does not define its loss type.
3. Loss is not inherent to a PreTrainedModel. There are plenty of papers that start from the model weights of one model and then change its loss function for continued pretraining (take this one as a recent example).

TL;DR: Seems like a design flaw to make loss_function a @property and seems like a mistake to overwrite it to ForCausalLMLoss for all models that don't fit the very narrow range of tasks supported by transformers.

Not a fix but a temporary workaround for me was this patch after loading a mdoel:

model = AutoModelForSequenceClassification.from_pretrained(...)

old_loss_function = model.loss_function
@lru_cache
def get_loss_function(self):
    def fixed_loss(labels, pooled_logits, config, **kwargs):
        labels = labels.to(pooled_logits.device)
        return old_loss_function(labels, pooled_logits, config, **kwargs)
    return fixed_loss
    
model.__class__.loss_function = property(get_loss_function)

Also, @ArthurZucker, I believe this fix confirms removing the labels = labels.to(pooled_logits.device) line breaks the multi-gpu setup. I would be very surprised if only Mistral is affected here given similar lines were removed from other models. Initially, I tried patching this at the PreTrainedModel.loss_function level, but for whatever reason the loaded models don't seem to pick up the change, and I'm not sure why.

[ArthurZucker]

Feedback taken, will open a PR to get to something better. It's hard to take everything into account, and support important to get your feedbacks! 🤗
Will merge the fix to loss on multi-GPU (I don't think we test all tasks for training and that is a flaw in our test environnement for sure! ) cc @ydshieh as for a TODO 😉