const AutogradMetadata& autogradMetadata,
ContextPtr autogradContext,
rpc::worker_id_t fromWorkerId,
- std::unordered_map<c10::Device, c10::Device> deviceMap)
+ rpc::DeviceMap deviceMap)
: autogradMetadata_(autogradMetadata),
// NOLINTNEXTLINE(performance-move-const-arg)
autogradContext_(std::move(autogradContext)),
const AutogradMetadata& autogradMetadata,
std::shared_ptr<DistAutogradContext> autogradContext,
rpc::worker_id_t fromWorkerId,
- std::unordered_map<c10::Device, c10::Device> deviceMap);
+ rpc::DeviceMap deviceMap);
torch::autograd::variable_list apply(
torch::autograd::variable_list&& grads) override;
rpc::worker_id_t fromWorkerId_;
// Device mapping for tensors sent over RPC.
- const std::unordered_map<c10::Device, c10::Device> deviceMap_;
+ const rpc::DeviceMap deviceMap_;
};
} // namespace autograd
MessageType messageType,
const AutogradMetadata& autogradMetadata,
c10::intrusive_ptr<rpc::Message> wrappedMessage,
- std::unordered_map<c10::Device, c10::Device> deviceMap)
+ rpc::DeviceMap deviceMap)
: fromWorkerId_(fromWorkerId),
messageType_(messageType),
autogradMetadata_(autogradMetadata),
std::unique_ptr<RpcCommandBase> wrappedRpc,
MessageType wrappedMessageType,
std::vector<torch::Tensor> tensors,
- std::unordered_map<c10::Device, c10::Device> deviceMap)
+ rpc::DeviceMap deviceMap)
: fromWorkerId_(fromWorkerId),
messageType_(messageType),
autogradMetadata_(autogradMetadata),
auto c10DeviceMap = tupleElements[4].to<c10::Dict<std::string, std::string>>();
// Convert to regular map.
- std::unordered_map<c10::Device, c10::Device> deviceMap;
+ rpc::DeviceMap deviceMap;
for (const auto& mapEntry : c10DeviceMap) {
deviceMap.insert({mapEntry.key(), mapEntry.value()});
}
return fromWorkerId_;
}
-const std::unordered_map<c10::Device, c10::Device>& RpcWithAutograd::
+const rpc::DeviceMap& RpcWithAutograd::
deviceMap() {
return deviceMap_;
}
rpc::MessageType messageType,
const AutogradMetadata& autogradMetadata,
c10::intrusive_ptr<rpc::Message> wrappedMessage,
- std::unordered_map<c10::Device, c10::Device> deviceMap = {});
+ rpc::DeviceMap deviceMap = {});
// Used when receiving an RPC over the wire.
RpcWithAutograd(
std::unique_ptr<rpc::RpcCommandBase> wrappedRpc,
rpc::MessageType wrappedMessageType,
std::vector<torch::Tensor> tensors,
- std::unordered_map<c10::Device, c10::Device> deviceMap = {});
+ rpc::DeviceMap deviceMap = {});
c10::intrusive_ptr<rpc::Message> toMessageImpl() && override;
rpc::worker_id_t fromWorkerId() const;
// Retrieve the device map.
- const std::unordered_map<c10::Device, c10::Device>& deviceMap();
+ const rpc::DeviceMap& deviceMap();
private:
// WorkerId from which this RPC originated. This is necessary for knowing
std::vector<torch::Tensor> tensors_;
// Device mapping for tensors that are sent across an RPC to another node.
- std::unordered_map<c10::Device, c10::Device> deviceMap_;
+ rpc::DeviceMap deviceMap_;
};
} // namespace autograd
const AutogradMetadata& autogradMetadata,
std::vector<torch::Tensor>& tensors,
rpc::worker_id_t fromWorkerId,
- const std::unordered_map<c10::Device, c10::Device>& deviceMap) {
+ const rpc::DeviceMap& deviceMap) {
// Initialize autograd context if necessary.
auto& autogradContainer = DistAutogradContainer::getInstance();
auto autogradContext =
c10::intrusive_ptr<torch::distributed::rpc::Message> wrappedRpcMsg,
MessageType msgType,
bool forceGradRecording,
- const std::unordered_map<c10::Device, c10::Device>& deviceMap) {
+ const rpc::DeviceMap& deviceMap) {
auto& autogradContainer = DistAutogradContainer::getInstance();
// If there is no valid context and no tensor requires grads, send original
const AutogradMetadata& autogradMetadata,
std::vector<torch::Tensor>& tensors,
rpc::worker_id_t fromWorkerId,
- const std::unordered_map<c10::Device, c10::Device>& deviceMap);
+ const rpc::DeviceMap& deviceMap);
// This method is a wrapper utility used internally to wrap autograd info
// and attach autograd function for each type of rpc call if it has valid
c10::intrusive_ptr<rpc::Message> wrappedRpcMsg,
rpc::MessageType msgType,
bool forceGradRecording = false,
- const std::unordered_map<c10::Device, c10::Device>& deviceMap =
+ const rpc::DeviceMap& deviceMap =
{});
// Send message after autograd checking
// Need to reverse the device map for the backward pass of distributed
// autograd.
- std::unordered_map<c10::Device, c10::Device> reverseDeviceMap;
+ DeviceMap reverseDeviceMap;
for (const auto& mapEntry : rpcWithAutograd.deviceMap()) {
reverseDeviceMap.insert({mapEntry.second, mapEntry.first});
}
const WorkerInfo& to,
c10::intrusive_ptr<Message> message,
const float rpcTimeoutSeconds = kUnsetRpcTimeout,
- const std::unordered_map<c10::Device, c10::Device>& deviceMap = {}) = 0;
+ const DeviceMap& deviceMap = {}) = 0;
// Retries sending the message up to maxRetries times until an ACK is
// receieved. The duration between consecutive sends is increased over
const WorkerInfo& to,
c10::intrusive_ptr<Message> message,
const float rpcTimeoutSeconds,
- const std::unordered_map<c10::Device, c10::Device>& /* unused */) {
+ const DeviceMap& /* unused */) {
// We only fail control messages that have been specified by the test case.
// For all other messages, we just send them without any failures.
if (!shouldFailMessage(message->type())) {
const WorkerInfo& to,
c10::intrusive_ptr<Message> message,
const float rpcTimeoutSeconds = torch::distributed::rpc::kUnsetRpcTimeout,
- const std::unordered_map<c10::Device, c10::Device>& deviceMap = {})
+ const DeviceMap& deviceMap = {})
override;
// Add delay to writes
// Need to reverse the device map for the backward pass of distributed
// autograd.
- std::unordered_map<c10::Device, c10::Device> reverseDeviceMap;
+ DeviceMap reverseDeviceMap;
for (const auto& mapEntry : rpcWithAutograd.deviceMap()) {
reverseDeviceMap.insert({mapEntry.second, mapEntry.first});
}