[GRPO] completion lengths are incorrectly logged when mask_truncated_completions=True

[edbeeching]

Reproduction

When the mask_truncated_completions is True, the values logged for the completion lengths are incorrect, for the max they are equal to the terminated completions and for in they are often zero. This is because we use the attn masks to calculate the lengths and these are zeroed for the truncated sequences.

The issue is obvious when you look at the min completion lengths:

It should be a simple change to make a copy of the mask here:

trl/trl/trainer/grpo_trainer.py

Line 1110 in a528b9c

if self.mask_truncated_completions:

and then gathering the copied masks here:

trl/trl/trainer/grpo_trainer.py

Line 1215 in a528b9c

agg_completion_mask = self.accelerator.gather_for_metrics(completion_mask.sum(1))

System Info

• Platform: Linux-5.15.0-1049-aws-x86_64-with-glibc2.31
• Python version: 3.11.11
• TRL version: 0.18.0.dev0+a528b9c
• PyTorch version: 2.6.0
• accelerator(s): cpu
• Transformers version: 4.51.3
• Accelerate version: 1.3.0
• Accelerate config: not found
• Datasets version: 3.4.1
• HF Hub version: 0.30.2
• bitsandbytes version: 0.45.3
• DeepSpeed version: 0.16.4
• Diffusers version: 0.32.2
• Liger-Kernel version: 0.5.8
• LLM-Blender version: 0.0.2
• OpenAI version: 1.66.5
• PEFT version: 0.14.0
• vLLM version: 0.8.4

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete

[shirinyamani]

Thanks @edbeeching for bringing this to our attention.
So I think the best place to make the copy of the completion_mask tensor is outside of the if block.
This is because if we put it inside the if block, sth like this:

if self.mask_truncated_completions:
    original_completion_mask = completion_mask.clone()  
    truncated_completions = ~is_eos.any(dim=1)
    completion_mask = completion_mask * (~truncated_completions).unsqueeze(1).int()

Then original_completion_mask would only be created when mask_truncated_completions is True.

Therefore, I think the best place to put it is where this PR represent it, to make sure we have a copy of it before any modification inside the mask_truncated_completions block. Please let me know what you think about it?

[edbeeching]

@shirinyamani

This is fine, I understand that you only want to make the copy if mask_truncated_completions is being used, however this will mean that when the statistics are gathered and reported on L1215 you will need to have another if statement in order to gather the right mask tensor.

[shirinyamani]

@edbeeching

Got it, yes I agree. So I moved it inside the block, Therefore;

When mask_truncated_completions is True:

• We make a copy of the mask before modifying it
• The copy is used for metrics
• The original is modified for loss computation

When mask_truncated_completions is False:

• We don't modify the mask
• We just use the original mask directly for both metrics and loss computation

Okay, the tests are all passed and should be good now in the associated PR.