Frozen dataclass as a batch

[mazurkin]

Description & Motivation

So I have a frozen immutable dataclass as a batch (this is my way to control complexity), and with bf16-true precision I am having this error in the HalfPrecision which calls apply_to_collection without allow_frozen parameter:

[rank0]: ValueError: A frozen dataclass was passed to `apply_to_collection` but this is not allowed.

I would like to suggest modification of the _apply_to_collection_slow, so in the is_dataclass_instance(data): you could check for the to() method in the dataclass with hasattr(data, 'to') so you could call just data.to(dtype=dtype) and let the user to handle the situation with the conversion themselves.

So I could add to my batch the new method:

import dataclasses
import typing

@dataclasses.dataclass(frozen=True)
class MyBatch:
    
    input: torch.Tensor

    def to(self, device: torch.device, dtype: torch.dtype, non_blocking: bool = False, dataset_idx: int = 0) -> typing.Self:
        # I am handling conversion myself
        return MyBatch(
            input=input.to(device=device, dtype=dtype, non_blocking=non_blocking),
        )

This feature also allow to get rid of the transfer_batch_to_device() callback, as the batch allows to move to the other device itself.

Pitch

The feature allows to provide more flexibility and keep the code neat.

Alternatives

• make dataclass mutable (unfrozen) - it is possible, but makes the code fragile
• convert dataclass to dict and back - two additional operations

Additional context

No response

cc @lantiga

[mazurkin]

same issue #10873

[mazurkin]

I tried to make a workaround with:

    @t.override
    def on_before_batch_transfer(self, batch: ContrastiveBatch, dataloader_idx: int) -> dict:
        return dataclasses.asdict(batch)

    @t.override
    def on_after_batch_transfer(self, batch: dict, dataloader_idx: int) -> ContrastiveBatch:
        return ContrastiveBatch(**batch)

unfortunately precision conversion in happens before the first hook in evaluation_loop.py:

if not (using_dataloader_iter := isinstance(data_fetcher, _DataLoaderIterDataFetcher)):
    batch = trainer.precision_plugin.convert_input(batch)
    batch = trainer.lightning_module._on_before_batch_transfer(batch, dataloader_idx=dataloader_idx)
    batch = call._call_strategy_hook(trainer, "batch_to_device", batch, dataloader_idx=dataloader_idx)

so it still fails

[mazurkin]

as a positive example - you already allow this for moving the batch to the device:

    @t.override
    def transfer_batch_to_device(
        self,
        batch: ContrastiveBatch,
        device: torch.device,
        dataloader_idx: int,
    ) -> ContrastiveBatch:
        return batch.to(device=device, non_blocking=self.on_gpu)

[mazurkin]

or the other way - just to split the dataclass right in _apply_to_collection_slow:

dict = {
  f.name: getattr(batch, f.name)
  for f in dataclasses.fields(batch)
}

and then assemble again back:

result = batch.__class__(**dict)