I am training a model that consists of some nn.Embeddings, and LSTM, a Feed Forward layer, and torch.distributions.negative_binomial.NegativeBinomial.
Around 90 epochs (~7 hours with num_workers=7, ~13 hours with num_workers=0), the job stalls, stops outputting logs, and GPU usage spikes to 100%. I am running the job with a P100 on Google Cloud AI Platform, and also experience the problem when running on Google Cloud Compute Engine with a P100.
This is the result of ps -aux:
USER       PID %CPU %MEM    VSZ   RSS TTY      STAT START   TIME COMMAND
root         1  0.0  0.0  18376  2968 ?        Ss   Jul17   0:00 /bin/bash ./execute.sh train production 2020-07-10
root         6 96.9 35.6 40669080 22036460 ?   Rl   Jul17 4154:35 python -m sales_forecasting.train --config production --ts 2020-07-10
root     22525  0.0 35.0 40641548 21673036 ?   Sl   Jul17   1:03 python -m sales_forecasting.train --config production --ts 2020-07-10
root     22526  0.0 35.0 40641560 21673048 ?   Sl   Jul17   0:57 python -m sales_forecasting.train --config production --ts 2020-07-10
root     22527  0.0 35.0 40638804 21670476 ?   Sl   Jul17   0:57 python -m sales_forecasting.train --config production --ts 2020-07-10
root     22528  0.0 35.0 40638816 21670476 ?   Sl   Jul17   0:57 python -m sales_forecasting.train --config production --ts 2020-07-10
root     22529  0.0 35.0 40641596 21673092 ?   Sl   Jul17   0:58 python -m sales_forecasting.train --config production --ts 2020-07-10
root     22530  0.0 35.0 40641608 21673148 ?   Sl   Jul17   0:56 python -m sales_forecasting.train --config production --ts 2020-07-10
root     22531  0.0 35.0 40641620 21673236 ?   Sl   Jul17   0:57 python -m sales_forecasting.train --config production --ts 2020-07-10

As mentioned, there is no stack trace with SIGINT or SIGQUIT, but when forcing a stack trace using the signal library it outputs:
Current thread 0x00007f021d6f3740 (most recent call first):
  File "/opt/cpml/sales_forecasting/sales_forecasting/DeepAR/loss.py", line 32 in neg_log_likelihood
  File "/opt/cpml/sales_forecasting/sales_forecasting/DeepAR/estimator.py", line 156 in train_batch
  File "/opt/cpml/sales_forecasting/sales_forecasting/train.py", line 53 in train_epoch
  File "/opt/cpml/sales_forecasting/sales_forecasting/train.py", line 96 in train
  File "/opt/cpml/sales_forecasting/sales_forecasting/train.py", line 246 in <module>
  File "/opt/conda/lib/python3.7/runpy.py", line 85 in _run_code
  File "/opt/conda/lib/python3.7/runpy.py", line 193 in _run_module_as_main

The neg_log_likelihood function is here:
def neg_log_likelihood(
        distribution: Distribution,
        target: torch.Tensor
) -> Tuple[torch.Tensor, int]:
    log_likelihood = distribution.log_prob(target)

    non_zero_mask = target != 0.0

    if not any(non_zero_mask):
        return -torch.sum(log_likelihood), len(target)

    return -torch.sum(log_likelihood[non_zero_mask]), torch.sum(non_zero_mask)

Finally, resource usage changes can be seen here:
image2830×1294 219 KB
Since the job takes 8 hours to fail, I haven’t yet isolated the code to a minimum example that fails. Has anyone seen anything similar before? Any idea what the issue could be?
Also worth noting is that the issue only occurs when using GPU, not CPU.