Replace shape-based empty batch handling inside DPDataLoader with structure-aware approach (#806)

[iden-kalemaj]

Summary:

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Docs change / refactoring / dependency upgrade

Motivation and Context / Related issue

Replaces unstable shape-based empty batch handling with a stateful approach that learns and replicates the actual output structure from collate_fn. This fixes a critical bug where custom collate functions returning non-list structures (dicts, custom classes) were incompatible with Poisson sampling.

The old implementation inspected dataset[0] to pre-compute shapes, then hardcoded empty batches as lists:

def collate(batch, collate_fn, sample_empty_shapes, dtypes):
    if len(batch) > 0:
        return collate_fn(batch)  # Could return dict, custom class, etc.
    else:
        return [torch.zeros(shape, dtype=dtype) for ...]  # Always list!

Bug -> if collate_fn returns a dict, non-empty batches are dicts but empty batches are lists -> type mismatch crash

Existing, related issue: #534

Solution:

New CollateFnWithEmpty learns the structure from the first non-empty batch:

class CollateFnWithEmpty:
    def __call__(self, batch):
        if len(batch) > 0:
            output = self.wrapped_collator_fn(batch)
            if self.first_batch is None:
                self.first_batch = copy.deepcopy(output)  # Learn structure
        else:
            output = self._make_empty_batch(self.first_batch)  # Replicate structure
        return output

Now empty batches match the structure of non-empty batches, regardless of what collate_fn returns.

If the first non-empty batch is actually the first batch, then it returns an error:

if self.first_batch is None:
    raise ValueError(
        "First sampled batch cannot be empty. Please ensure your dataset "
        "has sufficient samples or increase sample_rate."
    )

Key Changes

• Removed: shape_safe(), dtype_safe(), hardcoded list return
• Added: CollateFnWithEmpty class with recursive structure replication
• Changed: wrap_collate_with_empty() signature: (collate_fn, sample_empty_shapes, dtype) -> (collate_fn, batch_first, rand_on_empty)

It is compatible with existing API.
A small disclosure: for small percentage of users who hacked around empty batches handling, it might cause problems but in majority of cases it should be compatible.

How Has This Been Tested (if it applies)

• We used this approach to fine-tune Qwen 7B model using trl library for model alignment
• Tested on Mellum 5B parameter model fine-tuning

Checklist

• The documentation is up-to-date with the changes I made.
• [] I have read the CONTRIBUTING document and completed the CLA (see CONTRIBUTING).
• All tests passed, and additional code has been covered with new tests.

Test Plan:
Imported from GitHub, without a Test Plan: line.
Unit tests

Differential Revision: D98312879

Pulled By: iden-kalemaj

[meta-codesync]

@iden-kalemaj has exported this pull request. If you are a Meta employee, you can view the originating Diff in D98312879.