Table to RAM with 64 MB sequential reads (page-faulting the mmap over the mount was far too slow)

app.py

@@ -98,9 +98,15 @@ def _table_to_ram():
     if not isinstance(PLE.ngram_embedding, MmapNGramTable):
         return
     t0 = time.time()
-    table = torch.cat(PLE.ngram_embedding._parts)                       # sequential read of the pieces from the mount
-    PLE.ngram_embedding = torch.nn.Embedding.from_pretrained(table, freeze=True)
-    print(f"[table] n-gram table in RAM: {tuple(table.shape)} in {time.time() - t0:.0f}s", flush=True)
+    print("[table] reading the n-gram table into RAM (64 MB sequential reads)", flush=True)
+    try:
+        emb = PLE.ngram_embedding.to_ram()
+    except Exception as e:
+        print(f"[table] staying memory-mapped: {e!r}", flush=True)
+        return
+    PLE.ngram_embedding = emb
+    dt = time.time() - t0
+    print(f"[table] n-gram table in RAM: {tuple(emb.weight.shape)} in {dt:.0f}s ({emb.weight.numel() * emb.weight.element_size() / 2**30 / dt:.2f} GiB/s)", flush=True)
 
 
 threading.Thread(target=_table_to_ram, daemon=True).start()

whittle_load.py

@@ -55,7 +55,7 @@ class MmapNGramTable(torch.nn.Module):
 
     def __init__(self, pieces, dim: int, dtype):
         super().__init__()
-        self._parts, starts, row = [], [], 0
+        self._pieces, self._parts, starts, row = list(pieces), [], [], 0
         for path, offset, rows, stored in pieces:
             with open(path, "rb") as fh:
                 buf = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
@@ -65,6 +65,27 @@ class MmapNGramTable(torch.nn.Module):
         self.register_buffer("_bounds", torch.tensor(starts[1:], dtype=torch.long), persistent=False)
         self.register_buffer("weight", torch.empty(0, dim, dtype=dtype), persistent=False)   # device/dtype marker only
 
+    def to_ram(self, chunk: int = 64 << 20) -> torch.nn.Embedding:
+        """Copy the whole table into one in-memory nn.Embedding with large sequential file reads (much faster than
+        faulting the memory map in page by page, especially on a network mount)."""
+        stored = self._pieces[0][3]
+        assert all(p[3] == stored for p in self._pieces), "pieces with mixed dtypes"
+        table = torch.empty((self.num_embeddings, self.embedding_dim), dtype=stored)
+        raw = memoryview(table.view(torch.uint8).numpy().reshape(-1))
+        pos = 0
+        for path, offset, rows, _ in self._pieces:
+            nbytes = rows * self.embedding_dim * table.element_size()
+            with open(path, "rb", buffering=0) as fh:
+                fh.seek(offset)
+                done = 0
+                while done < nbytes:
+                    n = fh.readinto(raw[pos + done: pos + min(nbytes, done + chunk)])
+                    if not n:
+                        raise IOError(f"short read in {path} at {offset + done}")
+                    done += n
+            pos += nbytes
+        return torch.nn.Embedding.from_pretrained(table.to(self.weight.dtype), freeze=True)
+
     def forward(self, ids: torch.Tensor) -> torch.Tensor:
         flat = ids.reshape(-1).cpu()
         which = torch.bucketize(flat, self._bounds, right=True)