Repository: flairNLP/flair

Issue:
Keep word embedding tensors on pinned cpu memory during training
If `embeddings_storage_mode = 'cpu'` is selected during training, all word embedding tensors are moved to CPU memory and are at each epoch copied over to GPU memory when processing a mini-batch. It's possible to pin tensors so that the GPU copy operation becomes faster. The drawback is that pinning tensors is costly, so this would increase the time needed for the first epoch (which creates and pins the tensors) but would decrease the time needed for all subsequent epochs (since copy operations are now faster). 

Task: Add tensor pinning to training routine and investigate how this impacts training speed.


Discussion:

Resolving patch:
diff --git a/flair/data.py b/flair/data.py
--- a/flair/data.py
+++ b/flair/data.py
@@ -153,7 +153,7 @@ def embedding(self):
         pass
 
     @abstractmethod
-    def to(self, device: str):
+    def to(self, device: str, pin_memory: bool = False):
         pass
 
     @abstractmethod
@@ -166,9 +166,9 @@ def __init__(self, first: DataPoint, second: DataPoint):
         self.first = first
         self.second = second
 
-    def to(self, device: str):
-        self.first.to(device)
-        self.second.to(device)
+    def to(self, device: str, pin_memory: bool = False):
+        self.first.to(device, pin_memory)
+        self.second.to(device, pin_memory)
 
     def clear_embeddings(self, embedding_names: List[str] = None):
         self.first.clear_embeddings(embedding_names)
@@ -239,9 +239,15 @@ def set_embedding(self, name: str, vector: torch.tensor):
             device = next(iter(self._embeddings.values())).device
         self._embeddings[name] = vector.to(device, non_blocking=True)
 
-    def to(self, device: str):
+    def to(self, device: str, pin_memory: bool = False):
         for name, vector in self._embeddings.items():
-            self._embeddings[name] = vector.to(device, non_blocking=True)
+            if str(vector.device) != str(device):
+                if pin_memory:
+                    self._embeddings[name] = vector.to(
+                        device, non_blocking=True
+                    ).pin_memory()
+                else:
+                    self._embeddings[name] = vector.to(device, non_blocking=True)
 
     def clear_embeddings(self, embedding_names: List[str] = None):
         if embedding_names is None:
@@ -580,15 +586,21 @@ def get_embedding(self) -> torch.tensor:
 
         return torch.Tensor()
 
-    def to(self, device: str):
+    def to(self, device: str, pin_memory: bool = False):
 
         # move sentence embeddings to device
         for name, vector in self._embeddings.items():
-            self._embeddings[name] = vector.to(device, non_blocking=True)
+            if str(vector.device) != str(device):
+                if pin_memory:
+                    self._embeddings[name] = vector.to(
+                        device, non_blocking=True
+                    ).pin_memory()
+                else:
+                    self._embeddings[name] = vector.to(device, non_blocking=True)
 
         # move token embeddings to device
         for token in self:
-            token.to(device)
+            token.to(device, pin_memory)
 
     def clear_embeddings(self, embedding_names: List[str] = None):
 
@@ -795,9 +807,15 @@ def set_embedding(self, name: str, vector: torch.tensor):
             device = next(iter(self._embeddings.values())).device
         self._embeddings[name] = vector.to(device, non_blocking=True)
 
-    def to(self, device: str):
+    def to(self, device: str, pin_memory: bool = False):
         for name, vector in self._embeddings.items():
-            self._embeddings[name] = vector.to(device, non_blocking=True)
+            if str(vector.device) != str(device):
+                if pin_memory:
+                    self._embeddings[name] = vector.to(
+                        device, non_blocking=True
+                    ).pin_memory()
+                else:
+                    self._embeddings[name] = vector.to(device, non_blocking=True)
 
     def clear_embeddings(self, embedding_names: List[str] = None):
         if embedding_names is None:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -1867,7 +1867,7 @@ def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
                     token.set_embedding(self.name, embedding.clone())
 
             all_hidden_states_in_lm = all_hidden_states_in_lm.detach()
-            all_hidden_states_in_lm = None
+            del all_hidden_states_in_lm
 
         return sentences
 
diff --git a/flair/training_utils.py b/flair/training_utils.py
--- a/flair/training_utils.py
+++ b/flair/training_utils.py
@@ -5,6 +5,8 @@
 from enum import Enum
 from pathlib import Path
 from typing import List
+
+import flair
 from flair.data import Dictionary, Sentence
 from functools import reduce
 from sklearn.metrics import mean_squared_error, mean_absolute_error
@@ -358,7 +360,8 @@ def store_embeddings(sentences: List[Sentence], storage_mode: str):
         for sentence in sentences:
             sentence.clear_embeddings(delete_keys)
 
-    # memory management - option 1: send everything to CPU
+    # memory management - option 1: send everything to CPU (pin to memory if we train on GPU)
     if storage_mode == "cpu":
+        pin_memory = False if str(flair.device) == "cpu" else True
         for sentence in sentences:
-            sentence.to("cpu")
+            sentence.to("cpu", pin_memory=pin_memory)


Test patch: