Repository: huggingface/transformers

Issue:
BatchEncoding interacts poorly with apex.amp
## Environment info
- `transformers` version: 3.0.0
- Platform: Linux
- Python version: 3.6
- PyTorch version (GPU?): 1.6.0 GPU

### Who can help
tokenizers: @mfuntowicz 

## Information

We are building a multi-task transformer on top of the transformers API, using Nvidia's Apex to get the best speed out of the GPUs. When using Apex with `O2` optimisations (which are necessary to get speedups on our model), it inserts casts around the inputs to `forward` which in our case is a `BatchEncoding`. This cast blindly calls `BatchEncoding.to` with a `dtype` argument, which in turn is passed through to the values `to` method because there is only a type hint `str` on the `BatchEncoding.to` input argument, not a `isinstance` check. This causes our tensor of BPE ids to be cast from a `LongTensor` into `HalfTensor` and then the embedding layer raises an error because it's not expecting floating point values (also we rounded away most of our vocabulary terms).

## To reproduce

Run amp.initialize with the `O2` opt level on a model which's forward method accepts BatchEncoding as the input.

This [line](https://github.com/NVIDIA/apex/blob/4ef930c1c884fdca5f472ab2ce7cb9b505d26c1a/apex/amp/_initialize.py#L46) in apex calls the `to` method blindly on all the input types, which in this case calls `BatchEncoding.to` which doesn't support the full scope of a pytorch `to` because it's designed for moving things between devices and not for casting them.

We're currently guarding it this way:
```python
def patched_to(self, device: Union[str,torch.device]):
    """Send all values to device by calling v.to(device)"""
    if isinstance(device, str) or isinstance(device, torch.device):
        self.data = {k: v.to(device=device) for k, v in self.data.items()}
    return self
BatchEncoding.to = patched_to
```

## Expected behavior

The `BatchEncoding.to` method to not to pass the cast through (as the documentation indicates).

Though it's not clear if this is because the `to` method is underspecified, or because Apex is being quite aggressive (and a little silly) in it's blind casting of everything. Either way, if the `BatchEncoding.to` method isn't supposed to be able to cast things (and it doesn't look like it) then it probably shouldn't pass casts down to the tensors it contains.


Discussion:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

Any interest in fixing this?
Hello! Indeed, there doesn't seem to be any type cast. Do you want to open a PR with a fix?
Do you think the fix we have is appropriate? I'm worried it'll break something elsewhere which is expecting the to method to do casting.
I think here we could do something better than what is currently implemented. We could do it like it's done in the PyTorch implementation, as it takes all `args` and `kwargs`. Something like this:

```py
    def to(self, *args, **kwargs) -> "BatchEncoding":
        """
        *Appropriate docstring here*

        Args:
            *Appropriate Args here*

        Returns:
            :class:`~transformers.BatchEncoding`:
            The same instance of :class:`~transformers.BatchEncoding` after modification.
        """
        self.data = {k: v.to(*args, **kwargs) for k, v in self.data.items()}
        return self

```
What do you think?
Won't that pass the casts down? It looks like it would directly pass everything through, but we wouldn't usually want to cast the `LongTensor` that lives in a `BatchEncoding`.
Ah, then I misunderstood. Let me try to understand the issue better to see how to work it out:  as Apex casts from `LongTensor` to `HalfTensor` when using `BatchEncoding` which raises this issue, doesn't the issue happen if the inputs are `torch.Tensor`s instead of a dict/`BatchEncoding`? These inputs would be converted to `HalfTensor`s as well, resulting in the error, wouldn't it?
Apex.amp blindly calls a `to()` method on anything that gets passed into a module if it's a type it doesn't know about, passing it arguments that cast a tensor to `HalfTensor`. `BatchEncoding` exposes a `to()` method which is designed for moving things between devices, and not for casting, yet it passes the casts down to the tensors despite the docs saying it's only for moving between devices. We don't want to cast the tensors inside `BatchEncoding` to `HalfTensor` as that truncates and rounds the vocabulary to 65k and every 16th word in the high end, even if we cast the indices back to `LongTensor` at the point they are looked up in the embedding. This is a mostly silent failure as if you blindly insert the cast back in that the embedding error suggests you lose a lot of vocabulary items.

The workaround I posted makes `BatchEncoding` ignore any casts that are passed into it.

Honestly the issue is that pytorch overloaded the `to` method with two completely different functions, casting and device transfer, and python doesn't have a type system that allows you to select an appropriate method overloading based on the inbound method arguments.
Thanks for taking the time to write such a detailed explanation. I think your proposal makes complete sense. The way the `to` method is implemented in `BatchEncoding` is only intended for use with devices, not casting, so it won't break the current implementation.

On top of what you have proposed, I think logging a warning in the case that it's not a device would be appropriate (if it doesn't get caught in the `if` statement). We would want to warn users that their cast did not go through if they attempt doing something like this, rather than silently doing nothing.
Sure. I can log a warning in there and make up a PR.

Resolving patch:
diff --git a/src/transformers/tokenization_utils_base.py b/src/transformers/tokenization_utils_base.py
--- a/src/transformers/tokenization_utils_base.py
+++ b/src/transformers/tokenization_utils_base.py
@@ -776,7 +776,16 @@ def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding":
             :class:`~transformers.BatchEncoding`: The same instance of :class:`~transformers.BatchEncoding` after
             modification.
         """
-        self.data = {k: v.to(device) for k, v in self.data.items()}
+
+        # This check catches things like APEX blindly calling "to" on all inputs to a module
+        # Otherwise it passes the casts down and casts the LongTensor containing the token idxs
+        # into a HalfTensor
+        if isinstance(device, str) or isinstance(device, torch.device):
+            self.data = {k: v.to(device=device) for k, v in self.data.items()}
+        else:
+            logger.warning(
+                f"Attempting to cast a BatchEncoding to another type, {str(device)}. This is not supported."
+            )
         return self
 
 


Test patch: