Fix label truncation for per-sample nested structures in Trainer
Author: raimbekovm
Pull request: yes

## What does this PR do?

This PR fixes incorrect label handling in `Trainer.evaluation_loop` for models that use per-sample nested label structures like `tuple[list[Tensor], list[Tensor]]` (e.g., Mask2Former for instance segmentation).

### The Problem

When labels are structured as `(mask_labels, class_labels)` where each is a list of tensors (one per image with varying sizes), `gather_for_metrics` truncates incorrectly:

- **With `use_gather_object=False`**: truncates tensor dimensions instead of list length → instances are lost
- **With `use_gather_object=True`**: truncates the tuple itself → when `remainder=1`, `class_labels` is completely dropped

This caused **metrics degradation of 3-12+ mAP points** depending on batch configuration, and potential crashes when `class_labels` was lost entirely.

### Detailed Illustration

**Mask2Former label structure:**
```python
labels = (
    [mask_img0, mask_img1, mask_img2, mask_img3],  # mask_labels: list of tensors
    [class_img0, class_img1, class_img2, class_img3]  # class_labels: list of tensors
)
# Each tensor has shape [num_instances, ...] where num_instances varies per image
# Example: mask_img0.shape = [5, 256, 256]  (5 instances)
#          mask_img1.shape = [3, 256, 256]  (3 instances)
```

**Scenario:** `batch_size=8`, last batch has 4 images, so `remainder=4`

**Bug with `use_gather_object=False`:**
```python
# gather_for_metrics does: recursively_apply(lambda t: t[:remainder], labels)
# This truncates EACH TENSOR to 4 elements in first dimension:
mask_img0[:4]  # Shape [5, 256, 256] → [4, 256, 256] — LOST 1 INSTANCE!
mask_img1[:4]  # Shape [3, 256, 256] → [3, 256, 256] — ok (less than 4)
```

**Bug with `use_gather_object=True`:**
```python
# gather_for_metrics does: labels[:remainder]
labels[:4]  # Returns full tuple (ok when remainder >= 2)
labels[:1]  # Returns (mask_labels,) — CLASS_LABELS COMPLETELY LOST!
```

**What we need:**
```python
# Truncate at IMAGE level (list length), not at instance level (tensor dim)
result = (
    mask_labels[:num_samples],   # Keep first N images
    class_labels[:num_samples]   # Keep first N images
)
```

### The Solution

Added two helper functions in `trainer_pt_utils.py`:
- `is_per_sample_nested()` — detects `tuple[list[Tensor], ...]` structures
- `flatten_per_sample_nested_batches()` — correctly flattens and truncates at list level

Modified `Trainer.evaluation_loop` to:
1. Detect per-sample nested labels with `is_per_sample_nested()`
2. Use `gather_object` directly (bypassing `gather_for_metrics` truncation)
3. Accumulate labels from all processes with `extend()`
4. Flatten and truncate correctly at the end with `flatten_per_sample_nested_batches()`

This also properly handles **distributed training** where `gather_object` returns labels from all GPU processes.

Fixes #43388

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request), Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link to it if that's the case. → #43388
- [ ] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests?

## Who can review?

@SunMarc @muellerzr

Comments:
Thanks for the fix — the overall approach makes sense and aligns well with the issue.

I’m working on a small follow-up PR that builds on this change by:
- adding regression tests for per-sample nested label structures
- tightening guards to avoid applying truncation logic to non-label data (e.g. processors)
- covering edge cases like gradient_state.remainder == 1

I’ll reference this PR once the follow-up is submitted.



Thanks for the review @SunMarc!

I checked the list conversion idea — unfortunately it doesn't help. The issue is that `data[:remainder]` truncates the outer tuple/list, so with `remainder=1` we get `(mask_labels,)` and lose `class_labels` entirely. Converting to list first gives the same result. We need to truncate the inner lists (images), not the outer structure (label types).

As for fixing `gather_for_metrics` directly — yeah, that's definitely an option. It would need the same `is_per_sample_nested` detection though. I went with the transformers-only approach since it's a pretty model-specific pattern and doesn't require waiting for an accelerate release, but I'm open to opening a PR there instead if you think that's cleaner.

Happy to add a more detailed illustration in the description if that would help. Let me know what you'd prefer!

> Happy to add a more detailed illustration in the description if that would help. Let me know what you'd prefer!

If you can do that, that would be for the best ! 

> I checked the list conversion idea — unfortunately it doesn't help. The issue is that data[:remainder] truncates the outer tuple/list, so with remainder=1 we get (mask_labels,) and lose class_labels entirely. Converting to list first gives the same result. We need to truncate the inner lists (images), not the outer structure (label types).

Oh indeed, if the images for in the inner list, then it won't work. The issue if that in distributed setup, we won't be able to correctly gather the metrics with your solution no ? 



View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43395&sha=87af30

Can you try again @J-Bracke and check if it also works in a distributed case ?