saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Better error message when ORTModule used with torch.DataParallel (#7287)

Browse source 
* Better error message when ORTModule used with torch.DataParallel
This commit is contained in:
baijumeswani 2021-04-09 10:07:22 -07:00 committed by GitHub
parent c22963c23d
commit b221a4fd86
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 19 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
orttraining/orttraining/python/training/ortmodule.py
View file

@ -133,3 +133,22 @@ class ORTModule(torch.nn.Module):
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]]:
"""Override original method to delegate execution to the base module"""
yield from self._flattened_module._base_module.named_buffers(prefix=prefix, recurse=recurse)
def _replicate_for_data_parallel(self):
"""Raises a NotImplementedError exception since ORTModule is not compatible with torch.nn.DataParallel
torch.nn.DataParallel requires the model to be replicated across multiple devices, and
in this process, ORTModule tries to export the model to onnx on multiple devices with the same
sample input. Because of this multiple device export with the same sample input, torch throws an
exception that reads: "RuntimeError: Input, output and indices must be on the current device"
which can be vague to the user since they might not be aware of what happens behind the scene.
We therefore try to preemptively catch use of ORTModule with torch.nn.DataParallel and throw a
more meaningful exception.
Users must use torch.nn.parallel.DistributedDataParallel instead of torch.nn.DataParallel
which does not need model replication and is also recommended by torch to use instead.
"""
raise NotImplementedError("ORTModule is not compatible with torch.nn.DataParallel. "
"Please use torch.nn.parallel.DistributedDataParallel instead.")