Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Option to set 'non_blocking' for to(device) in BatchEncoding and BatchFeature #34883

Merged
merged 5 commits into from
Dec 9, 2024
Merged

Option to set 'non_blocking' for to(device) in BatchEncoding and BatchFeature #34883

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
88eddf8
Option to set 'non_blocking' for to(device) operation for performance…
Nov 22, 2024
08f5d4b
Enabling non_blocking in to() operation of BatchFeature.
Nov 22, 2024
9f465d7
Improved docstring on utilization of non_blocking
Nov 22, 2024
4e8e786
Force non_blocking as keyword argument
daniel-bogdoll Nov 25, 2024
2d2fe86
Merge branch 'main' into batch-encoding-with-non_blocking-option
daniel-bogdoll Dec 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/transformers/feature_extraction_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,7 @@ def to(self, *args, **kwargs) -> "BatchFeature":
Will be passed to the `to(...)` function of the tensors.
kwargs (`Dict`, *optional*):
Will be passed to the `to(...)` function of the tensors.
To enable asynchronous data transfer, set the `non_blocking` flag in `kwargs` (defaults to `False`).

Returns:
[`BatchFeature`]: The same instance after modification.
Expand All @@ -222,6 +223,7 @@ def to(self, *args, **kwargs) -> "BatchFeature":

new_data = {}
device = kwargs.get("device")
non_blocking = kwargs.get("non_blocking", False)
# Check if the args are a device or a dtype
if device is None and len(args) > 0:
# device should be always the first argument
Expand All @@ -241,7 +243,7 @@ def to(self, *args, **kwargs) -> "BatchFeature":
# cast and send to device
new_data[k] = v.to(*args, **kwargs)
elif isinstance(v, torch.Tensor) and device is not None:
new_data[k] = v.to(device=device)
new_data[k] = v.to(device=device, non_blocking=non_blocking)
else:
new_data[k] = v
self.data = new_data
Expand Down
10 changes: 7 additions & 3 deletions src/transformers/tokenization_utils_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -799,12 +799,13 @@ def as_tensor(value, dtype=None):

return self

def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding":
def to(self, device: Union[str, "torch.device"], *, non_blocking: bool = False) -> "BatchEncoding":
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we need * ?

Copy link
Contributor Author

@daniel-bogdoll daniel-bogdoll Dec 5, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@qubvel suggested this to enforce it as a keyword argument for future backwards compatability. All arguments after the * are forced to be passed as keyword arguments: #34883 (comment)

Copy link
Member

@qubvel qubvel Dec 5, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, only device can be passed as a positional argument with * introduced. This way, we will prevent anyone from using batch_feature.to("cuda", True) instead of batch_feature.to("cuda", non_blocking=True). This would be useful in case we introduce more positional arguments in the future or need to change order, for example, with adding dtype.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for explaining, good decision @qubvel ! 🤗

"""
Send all values to device by calling `v.to(device)` (PyTorch only).
Send all values to device by calling `v.to(device, non_blocking=non_blocking)` (PyTorch only).

Args:
device (`str` or `torch.device`): The device to put the tensors on.
non_blocking (`bool`): Whether to perform the copy asynchronously.

Returns:
[`BatchEncoding`]: The same instance after modification.
Expand All @@ -816,7 +817,10 @@ def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding":
# Otherwise it passes the casts down and casts the LongTensor containing the token idxs
# into a HalfTensor
if isinstance(device, str) or is_torch_device(device) or isinstance(device, int):
self.data = {k: v.to(device=device) if isinstance(v, torch.Tensor) else v for k, v in self.data.items()}
self.data = {
k: v.to(device=device, non_blocking=non_blocking) if isinstance(v, torch.Tensor) else v
for k, v in self.data.items()
}
else:
logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.")
return self
Expand Down