DDP fails when a HF module is present

rkd1137: <p>I ran into a strange issue while trying to use DDP + logging, with a single main process handling the logging. I went a bit deeper and it appears to have <em>something</em> to do with the presence of huggingface modules, but I’m not entirely certain. I’ve distilled it into an MWE (pasted below) with some following commentary.</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group, get_rank
from transformers import AutoModel

torch.manual_seed(1137)

bert = AutoModel.from_pretrained("prajjwal1/bert-small")


class SampleModule(nn.Module):
    def __init__(self):
        super().__init__()
        # random model
        self.small_model = bert
        self.W = nn.Linear(128, 32, bias=False)

    def forward(self, sample, debug=False):
        # Fake loss
        return self.W(sample[0] + sample[1] + sample[2]).mean()


batch_size = 32


def get_loader():
    while True:
        yield (
            torch.randn(batch_size, 128),
            torch.randn(batch_size, 128),
            torch.randn(batch_size, 128),
        )


init_process_group(backend="nccl")
rank = get_rank()
main_process = str(rank) == "0"
device = f"cuda:{rank}"

loader = iter(get_loader())
model = SampleModule().to(device)

for name, param in model.named_parameters():
    if name.startswith("small_model"):
        param.requires_grad = False

model = DDP(model, device_ids=[rank])
optimizer = optim.Adam(model.parameters())

# Block A
# batch = next(loader)
# batch = (batch[0].to(device), batch[1].to(device), batch[2].to(device))
# loss = model(batch)
# loss.backward()

for m, batch in enumerate(loader):
    print(f"Process {rank} on iter {m} is ready.")

    optimizer.zero_grad()
    model.train()
    batch = (batch[0].to(device), batch[1].to(device), batch[2].to(device))
    loss = model(batch)
    loss.backward()
    optimizer.step()

    # somehow torch.no_grad + the HF models cause problems...
    # but not if I "initialize" them by running a bakcwards pass frist?
    if main_process:
        print(f"logging on process {rank}")
        # can this cause problems? that batch is already on device but includes...
        model.eval()
        with torch.no_grad():
            batch = next(loader)
            batch = (batch[0].to(device), batch[1].to(device), batch[2].to(device))
            test_loss = model(batch)

destroy_process_group()
</code></pre>
<p>To run this code, I use a command like <code>CUDA_VISIBLE_DEVICES=6,7 torchrun --nproc-per-node 2 mwe.py</code>. If I run this, I get an error where the program will hang after a couple of iterations. Now, the program runs fine with any of the following changes:</p>
<ol>
<li>Uncomment block A, performing a single backwards pass before doing anything else</li>
<li>Remove <code>small_model</code> from the <code>SampleModule</code> (this was meant to be a generic HuggingFace module, and not the specific ones I’m using for my task).</li>
<li>Remove <code>torch.no_grad</code>. It still hangs, but it hangs at a later iteration.</li>
<li>Allow both processes to access the “logging” section (so no <code>if main_process</code>).</li>
</ol>
<p>I know I may get better luck on the HuggingFace forums, but this seems to be a strange interaction between DDP and HF. If anyone has any insight, would be much appreciated!</p>
<p>UPDATE: Though I’m not quite sure why, initializing DDP with <code>find_unused_parameters=True</code> (even without any unused parameters) seems to solve the problem, with some performance cost.</p>