Add support for models using final_logits_softcapping to AsyncGRPOTrainer

[mlarnouhet]

Feature request

Add support for models using final_logits_softcapping to AsyncGRPOTrainer

Motivation

Some models (like gemma 2) use final_logits_softcapping and currently cannot be run with AsyncGRPOTrainer.

Your contribution

I would be happy to collaborate on this feature request by creating a potential PR.

[rehan243]

Ran into a similar issue with logits softcapping on our production models at Reallytics.ai - we were using an older version of trl and had to patch it ourselves. Looks like you're trying to add support for final_logits_softcapping to AsyncGRPOTrainer. One thing that might work is modifying the loss_fn to accept an additional softcapping arg, like this:

def loss_fn(logits, labels, softcapping=None):
    if softcapping is not None:
        logits = logits / softcapping
    # rest of the loss calculation

Are you using the latest gemma 2 model config, and did you try passing final_logits_softcapping as a trainer arg?