@kashif and/or@qgallouedec wondering if either of you could take a look at this PR since you're familiar with how FA2 works with custom position ids. In particular I'm kind of lost on whytest_forward_pass_with_packing is failing. I've validated that the logic for packing the inputs and unpacking the logits is accurate through my test cases and debugging.

Just for some context I'm trying to pack all the responses for a given prompt/query into a single row in the batch, before running a forward pass through the reference, old and current policy models.

The thing I noticed is that despite sending in the right position ids FA2 doesn't seem to give me the same outputs with and without packing.

For example, I tried this out trying to simulate how a packed GRPO group would look like. Tokens[646,647,648] areprompt tokens here.

pad_token_id = trainer.processing_class.pad_token_id            sample_input_ids = torch.tensor(    [        [646, 647, 648, 649, pad_token_id],        [646, 647, 648, 650, 651],    ],    device=trainer.model.device,)sample_attention_mask = torch.tensor(    [        [1, 1, 1, 1, 0],        [1, 1, 1, 1, 1],    ],    device=trainer.model.device,)sample_packed_input_ids = torch.tensor(    [        [646, 647, 648, 649, 650, 651],    ],    device=trainer.model.device,)sample_position_ids = torch.tensor(    [        [0, 1, 2, 3, 3, 4],    ],    device=trainer.model.device,)reg_logs = trainer.model(input_ids=sample_input_ids, attention_mask=sample_attention_mask).logitsreg_logs = torch.gather(reg_logs, -1, sample_input_ids.unsqueeze(-1))packed_logs = trainer.model(    input_ids=sample_packed_input_ids,    position_ids=sample_position_ids,).logitspacked_logs = torch.gather(packed_logs, -1, sample_packed_input_ids.unsqueeze(-1))

I think that the logit scores for all the token ids[646,647,648,649,650, 651] should match forreg_logs andpacked_logs but I'm seeing that the logit scores for tokens650 and 651 do not match. So it seems like FA2 isn't correctly accounting for packing or I'm doing something wrong. Any pointers would be much appreciated. I think the attention_mask isn't correctly being generated for this use case inside FA2, so wondering if explicitly passing aattention_mask might help.