[Data][1/N] iter_torch_batches: proper h2d pipelining

[kyuds]

Description

Properly pipeline h2d transfers with compute.

Related issues

N/A

Additional information

N/A

[gemini-code-assist]

Code Review

This pull request introduces PipelinedFinalizeFn to overlap host-to-device transfers with downstream GPU compute using a dedicated copy stream and CUDA events, replacing the previous synchronous default_finalize_fn in iter_torch_batches. A critical issue was identified in the recursive _record_stream helper, where passing a string literal "torch.Tensor" to isinstance will cause a runtime TypeError instead of checking the type correctly.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit fc318d4. Configure here.}

[justinvyu]

could we use pytree for this? I'd prefer to not have the "Any" input type; we should have an explicit input type contract.

[kyuds]

I am confused on where we would actually use pytree though. In the traversal?

[kyuds]

also improved typing as batch is a TensorBatchReturnType, which is

TensorBatchReturnType = Union[
    "torch.Tensor",
    Tuple["torch.Tensor", ...],
    Dict[str, "torch.Tensor"],
]

[kyuds]

ps: I don't think pytree will be placed in the public library component anytime soon (discussion in torch repo shows that maintainers don't want it because pytree is not optimized and inefficient). There is another optree repo that we can use, but this is an external dependency for something that is not really that necessary.