[trainer] fix the GA model_accepts_loss_kwargs

[ArthurZucker]

What does this PR do?

Fixes #34577
model_accepts_loss_kwargs was wrongly looking at kwarg names, while you should only need kwargs (since the name can vary for FlashAttentionKwargs, LossKwargs etc)

[muellerzr]

TIL! However, as you can see by the failing test, this doesn't always work 😅 (If we can get it to that's great, I think that's originally why I went with explicit rather than implicit)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[techkang]

I think this pr introduced a new bug that if user use user defined loss funcion and model_accepts_loss_kwargs is False, compute_loss function cannot get the num_items_in_batch argument. Finally, user defined compute_loss_func will not receive this argument either.

[ArthurZucker]

model_accepts_loss_kwargs is supposed to be determined at init time and only depends on the forward pass of the model

[techkang]

I understand now, I misinterpreted the if condition earlier. However, in this PR, when model_accepts_loss_kwargs is True, it won't pass the num_items_in_batch parameter, which would make the GA loss modification functionality ineffective. Why is that?

[techkang]

In newest code, run

export RUN_SLOW=True
pytest tests/trainer/test_trainer.py::TrainerIntegrationPrerunTest::test_gradient_accumulation_loss_alignment

will cause error:

======================================================= short test summary info =======================================================
FAILED tests/trainer/test_trainer.py::TrainerIntegrationPrerunTest::test_gradient_accumulation_loss_alignment - AssertionError: 0.9038000000000004 not less than 0.1 : Difference -0.9038000000000004 is not within 0.1
============================================== 1 failed, 2 warnings in 102.43s (0:01:42) ==============================================

[ArthurZucker]

Ah shit the if condition is reversed

[ArthurZucker]

Opened a PR for a fix, thanks!