torch-ddp 0.1.0

data/examples/mnist/distributed.rb

@@ -0,0 +1,240 @@
+# Distributed MNIST training with Torch::Distributed + DistributedDataParallel
+# Run with: ruby examples/mnist/distributed.rb --gpus 2
+
+require "bundler/setup"
+require "optparse"
+require "torch"
+require "torch/distributed"
+require "torch/nn/parallel/distributed_data_parallel"
+require "torchvision"
+require "tmpdir"
+
+unless Torch::Distributed.available?
+  abort "torch.distributed was not built in this binary"
+end
+
+DEFAULT_CHECKPOINT_PATH = File.join(Dir.tmpdir, "mnist_ddp_checkpoint.pt")
+DEFAULT_BACKEND = if Torch.const_defined?(:CUDA) && Torch::CUDA.respond_to?(:available?) && Torch::CUDA.available?
+  "nccl"
+else
+  Torch::Distributed.get_default_backend_for_device(Torch::Accelerator.current_accelerator) || "gloo"
+end
+
+class MyNet < Torch::NN::Module
+  def initialize
+    super()
+    @conv1 = Torch::NN::Conv2d.new(1, 32, 3, stride: 1)
+    @conv2 = Torch::NN::Conv2d.new(32, 64, 3, stride: 1)
+    @dropout1 = Torch::NN::Dropout2d.new(p: 0.25)
+    @dropout2 = Torch::NN::Dropout2d.new(p: 0.5)
+    @fc1 = Torch::NN::Linear.new(9216, 128)
+    @fc2 = Torch::NN::Linear.new(128, 10)
+  end
+
+  def forward(x)
+    x = Torch::NN::F.relu(@conv1.call(x))
+    x = Torch::NN::F.relu(@conv2.call(x))
+    x = Torch::NN::F.max_pool2d(x, 2)
+    x = @dropout1.call(x)
+    x = Torch.flatten(x, start_dim: 1)
+    x = Torch::NN::F.relu(@fc1.call(x))
+    x = @dropout2.call(x)
+    Torch::NN::F.log_softmax(@fc2.call(x), 1)
+  end
+end
+
+def parse_options
+  defaults = {
+    epochs: 5,
+    batch_size: 64,
+    lr: 1.0,
+    gamma: 0.7,
+    backend: DEFAULT_BACKEND,
+    gpus: Torch::CUDA.available? ? [Torch::CUDA.device_count, 1].max : 1,
+    log_interval: 20,
+    data_dir: File.join(__dir__, "data"),
+    checkpoint_path: DEFAULT_CHECKPOINT_PATH,
+    resume: false
+  }
+
+  OptionParser.new do |opts|
+    opts.banner = "Usage: ruby distributed.rb [options]"
+    opts.on("--epochs N", Integer, "Number of epochs (default: #{defaults[:epochs]})") { |v| defaults[:epochs] = v }
+    opts.on("--batch-size N", Integer, "Batch size per process (default: #{defaults[:batch_size]})") { |v| defaults[:batch_size] = v }
+    opts.on("--lr FLOAT", Float, "Learning rate (default: #{defaults[:lr]})") { |v| defaults[:lr] = v }
+    opts.on("--gamma FLOAT", Float, "LR scheduler gamma (default: #{defaults[:gamma]})") { |v| defaults[:gamma] = v }
+    opts.on("--backend NAME", String, "Process group backend (default: #{defaults[:backend]})") { |v| defaults[:backend] = v }
+    opts.on("--gpus N", Integer, "Number of GPUs/processes to use") { |v| defaults[:gpus] = v }
+    opts.on("--log-interval N", Integer, "Batches between log statements") { |v| defaults[:log_interval] = v }
+    opts.on("--data-dir PATH", String, "Directory for cached MNIST data") { |v| defaults[:data_dir] = v }
+    opts.on("--checkpoint PATH", String, "Checkpoint file to save to (default: #{defaults[:checkpoint_path]})") { |v| defaults[:checkpoint_path] = v }
+    opts.on("--resume", "Load checkpoint weights before training if the file exists") { defaults[:resume] = true }
+  end.parse!(ARGV)
+
+  defaults
+end
+
+def load_datasets(rank, data_dir)
+  transforms = TorchVision::Transforms::Compose.new([
+    TorchVision::Transforms::ToTensor.new,
+    TorchVision::Transforms::Normalize.new([0.1307], [0.3081])
+  ])
+
+  if rank.zero?
+    train = TorchVision::Datasets::MNIST.new(data_dir, train: true, download: true, transform: transforms)
+    test = TorchVision::Datasets::MNIST.new(data_dir, train: false, download: true, transform: transforms)
+    Torch::Distributed.barrier
+  else
+    Torch::Distributed.barrier
+    train = TorchVision::Datasets::MNIST.new(data_dir, train: true, download: false, transform: transforms)
+    test = TorchVision::Datasets::MNIST.new(data_dir, train: false, download: false, transform: transforms)
+  end
+
+  [train, test]
+end
+
+def subset_for_rank(dataset, rank, world_size)
+  indices = rank.step(dataset.size - 1, world_size).to_a
+  Torch::Utils::Data::Subset.new(dataset, indices)
+end
+
+def checkpoint_map_location(device, rank)
+  accelerator_device = Torch::Accelerator.current_accelerator
+  return nil unless accelerator_device
+
+  accelerator_type = accelerator_device.type
+  target_index = device.index
+  if target_index.nil? && Torch::Accelerator.respond_to?(:device_count)
+    count = Torch::Accelerator.device_count
+    target_index = count.positive? ? rank % count : 0
+  end
+  { "#{accelerator_type}:0" => "#{accelerator_type}:#{target_index}" }
+end
+
+def load_checkpoint_if_present(ddp, device, rank, path)
+  return false unless path && File.exist?(path)
+
+  Torch::Distributed.barrier
+  kwargs = { weights_only: true }
+  map_location = checkpoint_map_location(device, rank)
+  kwargs[:map_location] = map_location if map_location
+  state_dict = Torch.load(path, **kwargs)
+  ddp.module.load_state_dict(state_dict)
+  true
+end
+
+def save_checkpoint(ddp, path, rank)
+  return unless path
+
+  Torch.save(ddp.module.state_dict, path) if rank.zero?
+  Torch::Distributed.barrier
+  puts "Saved checkpoint to #{path}" if rank.zero?
+end
+
+def train_epoch(model, device, loader, optimizer, epoch, rank, log_interval)
+  model.train
+  loader.each_with_index do |(data, target), batch_idx|
+    data = data.to(device)
+    target = target.to(device)
+
+    optimizer.zero_grad
+    loss = Torch::NN::F.nll_loss(model.call(data), target)
+    loss.backward
+    optimizer.step
+
+    next unless rank.zero? && (batch_idx % log_interval).zero?
+
+    processed = batch_idx * data.size(0)
+    total = loader.dataset.size
+    percent = 100.0 * processed / total
+    puts "Rank #{rank} | Epoch #{epoch} [#{processed}/#{total} (#{percent.round})%] Loss: #{'%.4f' % loss.item}"
+  end
+end
+
+def evaluate(model, device, loader)
+  model.eval
+  loss = 0.0
+  correct = 0
+  Torch.no_grad do
+    loader.each do |data, target|
+      data = data.to(device)
+      target = target.to(device)
+      output = model.call(data)
+      loss += Torch::NN::F.nll_loss(output, target, reduction: "sum").item
+      pred = output.argmax(1, keepdim: true)
+      correct += pred.eq(target.view_as(pred)).sum.item
+    end
+  end
+
+  loss /= loader.dataset.size
+  acc = 100.0 * correct / loader.dataset.size
+  puts "Test set: Average loss: #{format('%.4f', loss)}, Accuracy: #{correct}/#{loader.dataset.size} (#{format('%.1f', acc)}%)"
+end
+
+def run_worker(rank, world_size, port, options)
+  store = Torch::Distributed::TCPStore.new("127.0.0.1", port, world_size, rank.zero?)
+  accelerator = Torch::Accelerator.current_accelerator
+  backend = options[:backend] || Torch::Distributed.get_default_backend_for_device(accelerator) || DEFAULT_BACKEND
+  Torch::Distributed.init_process_group(backend, store: store, rank: rank, world_size: world_size)
+
+  device = if Torch::CUDA.available? && options[:gpus] > 0
+    Torch.device("cuda:#{rank % Torch::CUDA.device_count}")
+  else
+    Torch.device("cpu")
+  end
+
+  model = MyNet.new.to(device)
+  ddp = Torch::NN::Parallel::DistributedDataParallel.new(model, device_ids: device.type == "cuda" ? [device.index] : nil)
+  optimizer = Torch::Optim::Adadelta.new(ddp.module.parameters, lr: options[:lr])
+  scheduler = Torch::Optim::LRScheduler::StepLR.new(optimizer, step_size: 1, gamma: options[:gamma])
+
+  train_dataset, test_dataset = load_datasets(rank, options[:data_dir])
+  train_subset = subset_for_rank(train_dataset, rank, world_size)
+  train_loader = Torch::Utils::Data::DataLoader.new(train_subset, batch_size: options[:batch_size], shuffle: true)
+  test_loader = Torch::Utils::Data::DataLoader.new(test_dataset, batch_size: options[:batch_size], shuffle: false) if rank.zero?
+  checkpoint_path = options[:checkpoint_path]
+
+  if options[:resume]
+    loaded = load_checkpoint_if_present(ddp, device, rank, checkpoint_path)
+    if rank.zero?
+      if loaded
+        puts "Loaded checkpoint weights from #{checkpoint_path}"
+      else
+        puts "No checkpoint found at #{checkpoint_path}, starting from random initialization"
+      end
+    end
+  end
+
+  options[:epochs].times do |epoch_idx|
+    epoch = epoch_idx + 1
+    train_epoch(ddp, device, train_loader, optimizer, epoch, rank, options[:log_interval])
+    if rank.zero?
+      evaluate(ddp.module, device, test_loader)
+    end
+    save_checkpoint(ddp, checkpoint_path, rank) if checkpoint_path
+  end
+
+  Torch::Distributed.destroy_process_group
+end
+
+options = parse_options
+world_size = options[:gpus]
+raise "Number of GPUs requested must be >= 1" if world_size < 1
+if Torch::CUDA.available?
+  max_devices = Torch::CUDA.device_count
+  if world_size > max_devices
+    raise "Requested #{world_size} GPUs but only #{max_devices} visible"
+  end
+else
+  puts "CUDA not available, running #{world_size} CPU workers"
+end
+
+Torch.manual_seed(1)
+
+if world_size == 1
+  run_worker(0, 1, Torch::Distributed.free_port, options)
+else
+  Torch::Distributed.fork_world(world_size, start_method: :spawn) do |rank, port|
+    run_worker(rank, world_size, port, options)
+  end
+end

data/ext/torch_ddp/distributed.cpp

@@ -0,0 +1,348 @@
+#include <algorithm>
+#include <chrono>
+#include <cctype>
+#include <cstdlib>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <vector>
+
+#include <torch/torch.h>
+#if defined(USE_C10D) && defined(USE_C10D_NCCL)
+#include <torch/cuda.h>
+#include <c10/cuda/CUDAFunctions.h>
+#endif
+
+#include <rice/rice.hpp>
+#include <rice/stl.hpp>
+
+static_assert(
+    TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 9,
+    "Incompatible LibTorch version");
+
+#ifdef USE_C10D
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/distributed/c10d/ProcessGroupGloo.hpp>
+#include <torch/csrc/distributed/c10d/TCPStore.hpp>
+#include <torch/csrc/distributed/c10d/FileStore.hpp>
+#include <torch/csrc/distributed/c10d/Work.hpp>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#endif
+
+#if defined(USE_C10D) && defined(USE_C10D_NCCL)
+#include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
+#endif
+
+#if defined(USE_C10D) && !defined(_WIN32)
+#include <torch/csrc/distributed/c10d/HashStore.hpp>
+#endif
+
+namespace {
+
+#ifdef USE_C10D
+
+using StorePtr = c10::intrusive_ptr<::c10d::Store>;
+using ProcessGroupPtr = c10::intrusive_ptr<::c10d::Backend>;
+
+struct StoreWrapper {
+  StoreWrapper() = default;
+  explicit StoreWrapper(StorePtr store) : store_(std::move(store)) {}
+
+  StorePtr store_;
+};
+
+struct ProcessGroupWrapper {
+  ProcessGroupWrapper() = default;
+  explicit ProcessGroupWrapper(ProcessGroupPtr pg) : pg_(std::move(pg)) {}
+
+  ProcessGroupPtr pg_;
+};
+
+ProcessGroupPtr default_process_group;
+std::once_flag default_pg_cleanup_once;
+
+void shutdown_default_process_group() {
+  if (default_process_group) {
+    try {
+      default_process_group->shutdown();
+    } catch (...) {
+      // best effort; ensure reset still happens
+    }
+    default_process_group.reset();
+  }
+}
+
+void register_default_pg_cleanup() {
+  std::call_once(default_pg_cleanup_once, []() {
+    std::atexit([]() { shutdown_default_process_group(); });
+  });
+}
+
+ProcessGroupPtr resolve_process_group(Rice::Object pg_obj) {
+  if (pg_obj.is_nil()) {
+    if (!default_process_group) {
+      rb_raise(rb_eRuntimeError, "Distributed process group not initialized");
+    }
+    return default_process_group;
+  }
+  auto& wrapper = Rice::detail::From_Ruby<ProcessGroupWrapper&>().convert(pg_obj.value());
+  if (!wrapper.pg_) {
+    rb_raise(rb_eRuntimeError, "Invalid process group");
+  }
+  return wrapper.pg_;
+}
+
+int reduce_op_from_int(int code) {
+  if (code < 0 || code > static_cast<int>(::c10d::ReduceOp::UNUSED)) {
+    rb_raise(rb_eArgError, "Unknown reduce op code");
+  }
+  return code;
+}
+
+#endif
+
+} // namespace
+
+void init_distributed(Rice::Module& m) {
+  auto rb_mDistributed = Rice::define_module_under(m, "Distributed");
+#ifdef USE_C10D
+  register_default_pg_cleanup();
+  rb_mDistributed.define_singleton_function("available?", []() { return true; });
+
+  auto rb_cStore = Rice::define_class_under<StoreWrapper>(rb_mDistributed, "Store");
+  rb_cStore.define_method(
+      "_native?",
+      [](StoreWrapper& self) {
+        return static_cast<bool>(self.store_);
+      });
+
+  auto rb_cProcessGroup = Rice::define_class_under<ProcessGroupWrapper>(rb_mDistributed, "ProcessGroup")
+    .define_method(
+      "rank",
+      [](ProcessGroupWrapper& self) {
+        return self.pg_ ? self.pg_->getRank() : -1;
+      })
+    .define_method(
+      "size",
+      [](ProcessGroupWrapper& self) {
+        return self.pg_ ? self.pg_->getSize() : 0;
+      })
+    .define_method(
+      "backend",
+      [](ProcessGroupWrapper& self) {
+        if (!self.pg_) {
+          return std::string();
+        }
+        return self.pg_->getBackendName();
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_create_tcp_store",
+      [rb_cStore](const std::string& host,
+                  int port,
+                  int world_size,
+                  bool is_master,
+                  int64_t timeout_millis,
+                  bool wait_for_workers) -> Rice::Object {
+        ::c10d::TCPStoreOptions opts;
+        opts.port = static_cast<uint16_t>(port);
+        opts.isServer = is_master;
+        opts.numWorkers = world_size;
+        opts.waitWorkers = wait_for_workers;
+        opts.timeout = std::chrono::milliseconds(timeout_millis);
+        auto store = c10::make_intrusive<::c10d::TCPStore>(host, opts);
+        // Pass ownership first, then the Ruby class so Rice doesn't treat the class as the owner flag
+        return Rice::Data_Object<StoreWrapper>(new StoreWrapper(store), true, rb_cStore);
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_create_file_store",
+      [rb_cStore](const std::string& path, int world_size) -> Rice::Object {
+        auto store = c10::make_intrusive<::c10d::FileStore>(path, world_size);
+        return Rice::Data_Object<StoreWrapper>(new StoreWrapper(store), true, rb_cStore);
+      });
+
+#if !defined(_WIN32)
+  rb_mDistributed.define_singleton_function(
+      "_create_hash_store",
+      [rb_cStore]() -> Rice::Object {
+        auto store = c10::make_intrusive<::c10d::HashStore>();
+        return Rice::Data_Object<StoreWrapper>(new StoreWrapper(store), true, rb_cStore);
+      });
+#endif
+
+  rb_mDistributed.define_singleton_function(
+      "_init_process_group",
+      [rb_cProcessGroup](const std::string& backend,
+                         StoreWrapper& store_wrapper,
+                         int rank,
+                         int world_size,
+                         int64_t timeout_millis,
+                         int device_id) -> Rice::Object {
+        StorePtr store = store_wrapper.store_;
+        if (!store) {
+          rb_raise(rb_eArgError, "Store is required for init_process_group");
+        }
+
+        std::string backend_lower = backend;
+        std::transform(backend_lower.begin(), backend_lower.end(), backend_lower.begin(), ::tolower);
+
+        ProcessGroupPtr pg;
+        if (backend_lower == "gloo") {
+#ifdef USE_C10D_GLOO
+          auto options = ::c10d::ProcessGroupGloo::Options::create();
+          options->timeout = std::chrono::milliseconds(timeout_millis);
+          options->devices.push_back(::c10d::ProcessGroupGloo::createDefaultDevice());
+          pg = c10::make_intrusive<::c10d::ProcessGroupGloo>(store, rank, world_size, options);
+#else
+          rb_raise(rb_eRuntimeError, "Gloo backend is not available in this build");
+#endif
+        } else if (backend_lower == "nccl") {
+#if defined(USE_C10D_NCCL)
+          auto options = c10::make_intrusive<::c10d::ProcessGroupNCCL::Options>();
+          options->timeout = std::chrono::milliseconds(timeout_millis);
+          pg = c10::make_intrusive<::c10d::ProcessGroupNCCL>(store, rank, world_size, options);
+#else
+          rb_raise(rb_eRuntimeError, "NCCL backend is not available in this build");
+#endif
+        } else {
+          rb_raise(rb_eArgError, "Unsupported backend: %s", backend.c_str());
+        }
+
+        if (device_id >= 0 && backend_lower == "nccl") {
+#if defined(USE_C10D_NCCL)
+          if (!torch::cuda::is_available()) {
+            rb_raise(rb_eRuntimeError, "CUDA is not available for NCCL backend");
+          }
+          auto device_count = torch::cuda::device_count();
+          if (device_id >= static_cast<int>(device_count)) {
+            rb_raise(
+                rb_eArgError,
+                "Invalid device_id %d for NCCL backend (available devices: %d)",
+                device_id,
+                static_cast<int>(device_count));
+          }
+          c10::cuda::set_device(device_id);
+          pg->setBoundDeviceId(c10::Device(c10::kCUDA, device_id));
+#endif
+        }
+
+        default_process_group = pg;
+        return Rice::Data_Object<ProcessGroupWrapper>(new ProcessGroupWrapper(pg), true, rb_cProcessGroup);
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_destroy_process_group",
+      []() {
+        shutdown_default_process_group();
+        return Rice::Nil;
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_initialized?",
+      []() {
+        return static_cast<bool>(default_process_group);
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_default_process_group",
+      [rb_cProcessGroup]() -> Rice::Object {
+        if (!default_process_group) {
+          return Rice::Nil;
+        }
+        return Rice::Data_Object<ProcessGroupWrapper>(new ProcessGroupWrapper(default_process_group), true, rb_cProcessGroup);
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_get_world_size",
+      [](Rice::Object pg_obj) {
+        auto pg = resolve_process_group(pg_obj);
+        return pg->getSize();
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_get_rank",
+      [](Rice::Object pg_obj) {
+        auto pg = resolve_process_group(pg_obj);
+        return pg->getRank();
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_barrier",
+      [](Rice::Object pg_obj) {
+        auto pg = resolve_process_group(pg_obj);
+        ::c10d::BarrierOptions opts;
+        auto work = pg->barrier(opts);
+        work->wait();
+        return Rice::Nil;
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_all_reduce",
+      [](torch::Tensor& tensor, int op_code, Rice::Object pg_obj) {
+        auto pg = resolve_process_group(pg_obj);
+        ::c10d::AllreduceOptions opts;
+        opts.reduceOp = ::c10d::ReduceOp(static_cast<::c10d::ReduceOp::RedOpType>(reduce_op_from_int(op_code)));
+        std::vector<at::Tensor> tensors{tensor};
+        auto work = pg->allreduce(tensors, opts);
+        work->wait();
+        return tensor;
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_broadcast",
+      [](torch::Tensor& tensor, int src, Rice::Object pg_obj) {
+        auto pg = resolve_process_group(pg_obj);
+        ::c10d::BroadcastOptions opts;
+        opts.rootRank = src;
+        std::vector<at::Tensor> tensors{tensor};
+        auto work = pg->broadcast(tensors, opts);
+        work->wait();
+        return tensor;
+      });
+
+  rb_mDistributed.define_singleton_function(
+      "_register_ddp_hook",
+      [](torch::Tensor& tensor, ProcessGroupWrapper& pg_wrapper, int world_size) -> unsigned {
+        if (!pg_wrapper.pg_) {
+          rb_raise(rb_eArgError, "Process group is required for DDP hook registration");
+        }
+        if (world_size <= 0) {
+          rb_raise(rb_eArgError, "world_size must be positive");
+        }
+
+        auto pg = pg_wrapper.pg_;
+        // Register a native autograd hook that all-reduces gradients and scales
+        // them by the world size. This avoids calling back into Ruby from
+        // autograd worker threads.
+        unsigned handle = tensor.register_hook([pg, world_size](const at::Tensor& grad) {
+          ::c10d::AllreduceOptions opts;
+          opts.reduceOp = ::c10d::ReduceOp::SUM;
+          std::vector<at::Tensor> tensors{grad};
+          auto work = pg->allreduce(tensors, opts);
+          work->wait();
+          grad.div_(static_cast<double>(world_size));
+          return grad;
+        });
+
+        return handle;
+      });
+
+  auto rb_mReduceOp = Rice::define_module_under(rb_mDistributed, "ReduceOp");
+  rb_mReduceOp.const_set("SUM", INT2NUM(static_cast<int>(::c10d::ReduceOp::SUM)));
+  rb_mReduceOp.const_set("AVG", INT2NUM(static_cast<int>(::c10d::ReduceOp::AVG)));
+  rb_mReduceOp.const_set("PRODUCT", INT2NUM(static_cast<int>(::c10d::ReduceOp::PRODUCT)));
+  rb_mReduceOp.const_set("MIN", INT2NUM(static_cast<int>(::c10d::ReduceOp::MIN)));
+  rb_mReduceOp.const_set("MAX", INT2NUM(static_cast<int>(::c10d::ReduceOp::MAX)));
+  rb_mReduceOp.const_set("BAND", INT2NUM(static_cast<int>(::c10d::ReduceOp::BAND)));
+  rb_mReduceOp.const_set("BOR", INT2NUM(static_cast<int>(::c10d::ReduceOp::BOR)));
+  rb_mReduceOp.const_set("BXOR", INT2NUM(static_cast<int>(::c10d::ReduceOp::BXOR)));
+  rb_mReduceOp.const_set("PREMUL_SUM", INT2NUM(static_cast<int>(::c10d::ReduceOp::PREMUL_SUM)));
+
+  rb_mDistributed.const_set("DEFAULT_TIMEOUT", INT2NUM(::kProcessGroupDefaultTimeout.count() / 1000));
+#else
+  rb_mDistributed.define_singleton_function("available?", []() { return false; });
+  rb_mDistributed.const_set("DEFAULT_TIMEOUT", INT2NUM(30 * 60));
+#endif
+}

data/ext/torch_ddp/ext.cpp

@@ -0,0 +1,11 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+void init_distributed(Rice::Module& m);
+
+extern "C"
+void Init_ddp_ext() {
+  auto m = Rice::define_module("Torch");
+  init_distributed(m);
+}