`NotImplemtedError` when using 'torchaudio::rnnt_loss' on CUDA

username: Didac
cooked: <p>When I train the RNN-T on my local machine <strong>without CUDA</strong>, everything seems to work. However, when I use CUDA, I have this error when we are computing the loss:</p>
<blockquote>
<p>NotImplementedError: Could not run ‘torchaudio::rnnt_loss’ with arguments from the ‘CUDA’ backend. This could be because the operator doesn’t exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit <a href="https://fburl.com/ptmfixes" class="inline-onebox" rel="noopener nofollow ugc">Internal Login</a> for possible resolutions. ‘torchaudio::rnnt_loss’ is only available for these backends: [CPU, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMeta, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradNestedTensor, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].</p>
</blockquote>
<p>I’ve seen that some people are facing the same issue with <code>torchvision</code> and the solution was to compute the loss on the CPU instead of using the GPU. So instead of using the <code>rnnt_loss</code> as:</p>
<pre><code class="lang-auto">        rnnt_loss = torchaudio.functional.rnnt_loss(
            logits,
            transcript,
            output_len,
            transcript_len,
        )
</code></pre>
<p>I’m doing it as:</p>
<pre><code class="lang-auto">        rnnt_loss = torchaudio.functional.rnnt_loss(
            torch.from_numpy(logits.detach().cpu().numpy()),
            torch.from_numpy(transcript.detach().cpu().numpy()),
            torch.from_numpy(output_len.detach().cpu().numpy()),
            torch.from_numpy(transcript_len.detach().cpu().numpy()),
        ) 
</code></pre>
<p>This solution worked on my local machine <strong>using CUDA and just one GPU</strong>. However, when I run it on an instance with multiple GPU’s then I have this <code>RuntimeError</code>:</p>
<blockquote>
<p>RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument <code>find_unused_parameters=True</code> to <code>torch.nn.parallel.DistributedDataParallel</code>, and by making sure all <code>forward</code> function outputs participate in calculating loss. If you already have done the above, then the distributed data parallel module wasn’t able to locate the output tensors in the return value of your module’s <code>forward</code> function. Please include the loss function and the structure of the return value of <code>forward</code> of your module when reporting this issue (e.g. list, dict, iterable).</p>
</blockquote>
<p>To try to find the “unused parameters”, I set <code>find_unused_parameters=True</code> as suggested in the error. However, I had no useful information and all I had it was a warning message saying that everything seems to be OK and I don’t need to set <code>find_unused_parameters=True</code> :</p>
<blockquote>
<p>[W reducer.cpp:1298] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration, which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())</p>
</blockquote>
<p>The ideal solution would be to run the <code>rnnt_loss</code> with <code>distributed</code> GPU’s but I don’t know how. Any help or comments will be highly appreciate it. Thank you.</p>

username: Didac
cooked: <p>Just to let you know, the problem was solved after updating <code>torchaudio</code> from <code>1.0.1+cpu</code> to <code>1.0.2</code>. Hope this helps.</p>