Closed Jiaxin-Wen closed 10 months ago
Oh I notice that input_ids
is wrapped into BatchedTensor
. problems are solved.
Yep, this is vmap
parallelizing over the batch dimension.
I wonder if there is a proper way to get batch_size
and sequence_length
from input_ids
after it is parallelized since some models in transformers still use such code batch_size, seq_len = input_ids.shape[:2]
For example, when I adopt llama as the backbone model with transformers==4.30.2
, here are some lines of its source code:
And I'm confused about how to adapt the input_ids
, which are already wrapped into BatchedTensor
to these lines of code.
You can think of the BatchedTensor
as a simple torch
tensor with batch dimension equal to one. If there's no batch dimension, you can .unsqueeze(0)
to create a "dummy" batch dimension. Closing for now, but feel free to re-open if this doesn't resolve the issue.
Take trak in qnli for example, https://github.com/MadryLab/trak/blob/76b13ca55f1a16a243a23aba312cf6aad57b84d0/examples/qnli.py#L138 We first specify the batch size as 16. However, as I inspect the input_ids, I find its data type is
GradTrackingTensor
, with a shape oftorch.Size([1, 128])
, instead of torch.Size([16, 128]) https://github.com/MadryLab/trak/blob/76b13ca55f1a16a243a23aba312cf6aad57b84d0/examples/qnli.py#L74How to understand it?