I tried a few different settings. It seems the easiest thing to do is to ignore the pin_memory flag and embed everything on the CPU before calling DataParallel.
More or less this:
embed = torch.nn.Embedding.from_pretrained(embeddings, freeze=True)
model = Model()
model.cuda()

x = torch.LongTensor(batch_size, dim).random_(0, vocab_size-1)
# If you pin at this point, doesn't impact performance.
emb = embed(x)
# If you pin at this point, slightly slowed performance.
out = torch.nn.parallel.data_parallel(model, (emb, ))

Here’s some example code I used to try various settings:

  
      gist.github.com
  
  
    https://gist.github.com/mrdrozdov/3bcb412ff60151f0cb6caf95d0fcccaa 10
demo_embeddings.py
import argparse
import json

import torch
import torch.nn as nn

from tqdm import tqdm


class Model(nn.Module):
This file has been truncated. show original
results.txt
# Embeddings on Each GPU / Pin Memory before embedding

python demo_embeddings.py --pin_early --embedding_in_model
time: > 30m
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.130                Driver Version: 384.130                   |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
This file has been truncated. show original