torch-ddp 0.2.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c04c1f358a671d251826b7bf9db798bd5d9f11e279b639cff382f9f8f07d4b5f
-  data.tar.gz: 14c9db6913aaf75f98242f06808db8a3b696f37e6d779e092c8ed59cd398a644
+  metadata.gz: 9f9b645358a4641e4276ff48adb382c869feb024882d84c4835a39025bda6e62
+  data.tar.gz: 8f141183e503f72f0ec1622209e4ca17dcf8611471f68a45e34e0e2c146e4631
 SHA512:
-  metadata.gz: 0c62affe04041abca2dc56d6e82cc511d0338c3a4460c076b2fde13f219b889903b376521c92a38a6df0f0447bbd058081075ec1fcd6d3fcfb914063a0338691
-  data.tar.gz: 11634e1c29274033be29c37f51eae9bcc62bab902afedb469bdfde39e74f9d9230a93295a1c0b5fb2ee8e303af0df20e659335ee5e21c3c9e6b89608c9600dde
+  metadata.gz: fe63d5d9cbc27128b15166363c6a1ae84f342a33c66211e43f6464b9fc58a2a70e497d053561bdf176826307ef089a616fd102db2dd8da135b1ac4e8ee9b781c
+  data.tar.gz: 733d6befc6151f55452b3834d2b8c4e174184adda5e8cacc99bc510b65576b26f6e721454d7e88d6309ff2a39290488cc75926f5e137c4ce697ccc0dd53e8827

data/examples/benchmark/training.rb

@@ -13,7 +13,7 @@ require "torchvision"
 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"
+  Torch::Distributed.get_default_backend_for_device(nil) || "gloo"
 end
 SPAWN_BACKEND_ENV = "TORCH_RB_BENCH_BACKEND".freeze
 SPAWN_GROUP_ENV = "TORCH_RB_BENCH_GROUP_SIZE".freeze
@@ -181,8 +181,7 @@ def benchmark_worker(rank, world_size, port, options)
   raise ArgumentError, "Unsupported architecture #{arch.inspect}" unless config
 
   distributed = world_size > 1
-  accelerator = Torch::Accelerator.current_accelerator
-  selected_backend = options[:backend] || Torch::Distributed.get_default_backend_for_device(accelerator) || DEFAULT_BACKEND
+  selected_backend = options[:backend] || Torch::Distributed.get_default_backend_for_device(nil) || DEFAULT_BACKEND
   if distributed
     store = Torch::Distributed::TCPStore.new("127.0.0.1", port, world_size, rank.zero?)
     Torch::Distributed.init_process_group(selected_backend, store: store, rank: rank, world_size: world_size)
@@ -350,7 +349,7 @@ results = []
 
 backends.each do |backend|
   unless backend_supported?(backend)
-    warn "Skipping backend=#{backend} because required accelerator support is unavailable."
+    warn "Skipping backend=#{backend} because required CUDA support is unavailable."
     next
   end
 

data/examples/mnist/distributed.rb

@@ -17,7 +17,7 @@ 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"
+  Torch::Distributed.get_default_backend_for_device(nil) || "gloo"
 end
 
 class MyNet < Torch::NN::Module
@@ -99,16 +99,18 @@ def subset_for_rank(dataset, rank, world_size)
 end
 
 def checkpoint_map_location(device, rank)
-  accelerator_device = Torch::Accelerator.current_accelerator
-  return nil unless accelerator_device
-
-  accelerator_type = accelerator_device.type
+  device_type = device&.type
+  return nil unless device_type && device_type != "cpu"
   target_index = device.index
-  if target_index.nil? && Torch::Accelerator.respond_to?(:device_count)
-    count = Torch::Accelerator.device_count
+  if target_index.nil?
+    count = if Torch.const_defined?(:CUDA) && Torch::CUDA.respond_to?(:device_count)
+      Torch::CUDA.device_count
+    else
+      0
+    end
     target_index = count.positive? ? rank % count : 0
   end
-  { "#{accelerator_type}:0" => "#{accelerator_type}:#{target_index}" }
+  { "#{device_type}:0" => "#{device_type}:#{target_index}" }
 end
 
 def load_checkpoint_if_present(ddp, device, rank, path)
@@ -173,8 +175,7 @@ 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
+  backend = options[:backend] || Torch::Distributed.get_default_backend_for_device(nil) || DEFAULT_BACKEND
   Torch::Distributed.init_process_group(backend, store: store, rank: rank, world_size: world_size)
 
   device = if Torch::CUDA.available? && options[:gpus] > 0

data/ext/torch_ddp/cuda.cpp

@@ -44,7 +44,7 @@ void register_cuda_helpers(Rice::Module& m) {
             rb_eRuntimeError,
             "Torch::DDP._cuda_set_device requires CUDA support");
 #endif
-        return Rice::Nil;
+        return Rice::Object(Qnil);
       });
 
   rb_mDDP.define_singleton_function(
@@ -57,7 +57,7 @@ void register_cuda_helpers(Rice::Module& m) {
             rb_eRuntimeError,
             "Torch::DDP._cuda_empty_cache requires CUDA support");
 #endif
-        return Rice::Nil;
+        return Rice::Object(Qnil);
       });
 }
 

data/ext/torch_ddp/distributed.cpp

@@ -17,7 +17,7 @@
 #endif
 
 static_assert(
-    TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 9,
+    TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 13,
     "Incompatible LibTorch version");
 
 #ifdef USE_C10D
@@ -54,22 +54,44 @@ struct StoreWrapper {
 
 struct ProcessGroupWrapper {
   ProcessGroupWrapper() = default;
-  explicit ProcessGroupWrapper(ProcessGroupPtr pg) : pg_(std::move(pg)) {}
+  explicit ProcessGroupWrapper(ProcessGroupPtr pg)
+      : pg_(std::move(pg)), preserve_on_destroy_(pg_ && pg_->getBackendName() == "nccl") {}
+
+  ~ProcessGroupWrapper() {
+    if (preserve_on_destroy_) {
+      pg_.release();
+    }
+  }
 
   ProcessGroupPtr pg_;
+  bool preserve_on_destroy_ = false;
 };
 
 ProcessGroupPtr default_process_group;
 std::once_flag default_pg_cleanup_once;
 
 void shutdown_default_process_group() {
-  if (default_process_group) {
+  if (!default_process_group) {
+    return;
+  }
+
+  auto process_group = std::move(default_process_group);
+  if (process_group->getBackendName() == "nccl") {
     try {
-      default_process_group->shutdown();
+      process_group->abort();
     } catch (...) {
-      // best effort; ensure reset still happens
+      // best effort; avoid throwing from Ruby/C++ exit handlers
     }
-    default_process_group.reset();
+    // LibTorch 2.13 may double-free NCCL resources from the intrusive_ptr
+    // destructor after abort; leave this final owning reference leaked.
+    process_group.release();
+    return;
+  }
+
+  try {
+    process_group->shutdown();
+  } catch (...) {
+    // best effort; the local intrusive_ptr still releases the Gloo group
   }
 }
 
@@ -254,7 +276,7 @@ void init_distributed(Rice::Module& m) {
       "_destroy_process_group",
       []() {
         shutdown_default_process_group();
-        return Rice::Nil;
+        return Rice::Object(Qnil);
       });
 
   rb_mDistributed.define_singleton_function(
@@ -267,7 +289,7 @@ void init_distributed(Rice::Module& m) {
       "_default_process_group",
       [rb_cProcessGroup]() -> Rice::Object {
         if (!default_process_group) {
-          return Rice::Nil;
+          return Rice::Object(Qnil);
         }
         return Rice::Data_Object<ProcessGroupWrapper>(new ProcessGroupWrapper(default_process_group), true, rb_cProcessGroup);
       });
@@ -293,7 +315,7 @@ void init_distributed(Rice::Module& m) {
         ::c10d::BarrierOptions opts;
         auto work = pg->barrier(opts);
         work->wait();
-        return Rice::Nil;
+        return Rice::Object(Qnil);
       });
 
   rb_mDistributed.define_singleton_function(

data/ext/torch_ddp/extconf.rb

@@ -1,6 +1,6 @@
 require "mkmf-rice"
 
-$CXXFLAGS += " -std=c++17 $(optflags)"
+$CXXFLAGS += " -std=c++20 $(optflags)"
 
 # change to 0 for Linux pre-cxx11 ABI version
 $CXXFLAGS += " -D_GLIBCXX_USE_CXX11_ABI=1"
@@ -38,7 +38,6 @@ cuda_inc ||= "/usr/include"
 
 cudnn_inc, cudnn_lib = dir_config("cudnn")
 cudnn_lib ||= "/usr/local/cuda/lib"
-abort "cuda.h not found" unless find_header("cuda.h")
 
 gloo_inc, _ = dir_config("gloo")
 gloo_inc ||= "./vendor/gloo"
@@ -56,13 +55,25 @@ if Dir["#{lib}/*torch_cuda*"].any?
   $LDFLAGS += " -L#{cudnn_lib}" if Dir.exist?(cudnn_lib) && cudnn_lib != cuda_lib
   with_cuda = have_library("cuda") && have_library("cudnn")
 end
+abort "cuda.h not found" if with_cuda && !find_header("cuda.h")
 $defs << "-DWITH_CUDA" if with_cuda
 
 $INCFLAGS += " -I#{inc}"
 $INCFLAGS += " -I#{inc}/torch/csrc/api/include"
 
 CONFIG["CC"] = CONFIG["CXX"]
+CONFIG["CXXFLAGS"] = $CXXFLAGS
+RbConfig::CONFIG["CXXFLAGS"] = $CXXFLAGS
 $CFLAGS = $CXXFLAGS
+$CFLAGS += " -std=c++20"
+
+def try_link_with_cppflags(source, flags)
+  original_cppflags = $CPPFLAGS
+  $CPPFLAGS = [original_cppflags, flags].compact.join(" ")
+  try_link(source)
+ensure
+  $CPPFLAGS = original_cppflags
+end
 
 $LDFLAGS += " -Wl,-rpath,#{lib}"
 if RbConfig::CONFIG["host_os"] =~ /darwin/i && RbConfig::CONFIG["host_cpu"] =~ /arm|aarch64/i && Dir.exist?("/opt/homebrew/opt/libomp/lib")
@@ -70,7 +81,7 @@ if RbConfig::CONFIG["host_os"] =~ /darwin/i && RbConfig::CONFIG["host_cpu"] =~ /
 end
 $LDFLAGS += ":#{cuda_lib}/stubs:#{cuda_lib}" if with_cuda
 
-# https://github.com/pytorch/pytorch/blob/v2.9.0/torch/utils/cpp_extension.py#L1351-L1364
+# https://github.com/pytorch/pytorch/blob/v2.13.0/torch/utils/cpp_extension.py#L1351-L1364
 $LDFLAGS += " -lc10 -ltorch_cpu -ltorch"
 if with_cuda
   $LDFLAGS += " -lcuda -lnvrtc"
@@ -80,7 +91,10 @@ if with_cuda
   $LDFLAGS += " -Wl,--no-as-needed,#{lib}/libtorch.so"
 end
 
-supports_c10d = try_link(<<~CPP, "-DUSE_C10D")
+supports_c10d = try_link_with_cppflags(<<~CPP, "-DUSE_C10D")
+  // Ruby's headers define `_` as a compatibility macro, which collides with
+  // LibTorch 2.13's TensorAccessor headers.
+  #undef _
   #include <torch/torch.h>
   #include <torch/csrc/distributed/c10d/FileStore.hpp>
 
@@ -103,7 +117,8 @@ else
   puts "Building without distributed support"
 end
 
-supports_c10d_gloo = supports_c10d && try_link(<<~CPP, "-DUSE_C10D -DUSE_C10D_GLOO")
+supports_c10d_gloo = supports_c10d && try_link_with_cppflags(<<~CPP, "-DUSE_C10D -DUSE_C10D_GLOO")
+  #undef _
   #include <torch/torch.h>
   #include <torch/csrc/distributed/c10d/ProcessGroupGloo.hpp>
   #include <torch/csrc/distributed/c10d/FileStore.hpp>
@@ -117,7 +132,8 @@ supports_c10d_gloo = supports_c10d && try_link(<<~CPP, "-DUSE_C10D -DUSE_C10D_GL
   }
 CPP
 
-supports_c10d_nccl = with_cuda && try_link(<<~CPP, "-DUSE_C10D -DUSE_C10D_NCCL")
+supports_c10d_nccl = with_cuda && try_link_with_cppflags(<<~CPP, "-DUSE_C10D -DUSE_C10D_NCCL")
+  #undef _
   #include <torch/torch.h>
   #include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
 

data/lib/torch/ddp/monkey_patch.rb

@@ -13,7 +13,7 @@ module Torch
           warn("#{WARNING_PREFIX} Applying torch compatibility patch for: #{missing.join(', ')}. Please upgrade the torch gem for native support.")
           patch_cuda_set_device if missing.include?(:cuda_set_device)
           patch_cuda_empty_cache if missing.include?(:cuda_empty_cache)
-          patch_device_helpers
+          patch_device_helpers if missing.include?(:device_helpers)
           patch_load if missing.include?(:load_keywords)
           patch_tensor_item if missing.include?(:tensor_item_scalar)
           @applied = true
@@ -25,11 +25,18 @@ module Torch
           missing = []
           missing << :cuda_set_device unless Torch.const_defined?(:CUDA) && Torch::CUDA.respond_to?(:set_device)
           missing << :cuda_empty_cache unless Torch.const_defined?(:CUDA) && Torch::CUDA.respond_to?(:empty_cache)
+          missing << :device_helpers unless native_device_helpers?
           missing << :load_keywords unless load_supports_map_location_and_weights_only?
           missing << :tensor_item_scalar unless tensor_item_returns_scalar?
           missing
         end
 
+        def native_device_helpers?
+          Torch.tensor([0]).device.is_a?(Torch::Device)
+        rescue StandardError
+          false
+        end
+
         def load_supports_map_location_and_weights_only?
           params = Torch.method(:load).parameters
           keyword_names = params.select { |kind, _| [:key, :keyreq].include?(kind) }.map(&:last)

data/lib/torch/ddp/version.rb

@@ -1,5 +1,5 @@
 module Torch
   module DDP
-    VERSION = "0.2.0"
+    VERSION = "0.3.0"
   end
 end

data/lib/torch/distributed.rb

@@ -62,12 +62,10 @@ module Torch
 
         timeout_ms = (timeout * 1000).to_i
         bound_device_id = device_id.nil? ? -1 : Integer(device_id)
-        if backend == "nccl" && bound_device_id >= 0 && Torch.const_defined?(:CUDA) && Torch::CUDA.respond_to?(:set_device)
-          device_count = Torch::CUDA.device_count if Torch::CUDA.respond_to?(:device_count)
-          # Only attempt to switch devices when the requested id exists to avoid
-          # raising on hosts with fewer GPUs than the provided local rank.
-          Torch::CUDA.set_device(bound_device_id) if device_count.nil? || bound_device_id < device_count
-        end
+        # The native extension selects the NCCL device after creating the process group.
+        # Calling a Ruby compatibility cudaSetDevice before construction triggers a
+        # LibTorch 2.13/CUDA teardown double free.
+
         pg = _init_process_group(backend, store, rank, world_size, timeout_ms, bound_device_id)
         warmup_process_group(pg, backend)
       end
@@ -264,17 +262,17 @@ module Torch
         when Torch::Device
           device.type
         when NilClass
-          accelerator_type || "cpu"
+          device_type_fallback || "cpu"
         when String
           Torch.device(device).type
         when Integer
-          return accelerator_type || "cpu" if device.negative?
+          return device_type_fallback || "cpu" if device.negative?
           if Torch.const_defined?(:CUDA) && Torch::CUDA.respond_to?(:device_count)
             max = Torch::CUDA.device_count
-            return accelerator_type || "cpu" if max <= 0 || device >= max
+            return device_type_fallback || "cpu" if max <= 0 || device >= max
             return Torch.device("cuda:#{device}").type
           end
-          accelerator_type || "cpu"
+          device_type_fallback || "cpu"
         else
           return device.type if device.respond_to?(:type)
           Torch.device(device).type
@@ -283,9 +281,17 @@ module Torch
         raise ArgumentError, "Invalid device #{device.inspect}: #{e.message}"
       end
 
-      def accelerator_type
-        acc = Torch::Accelerator.current_accelerator
-        acc.type if acc && acc.respond_to?(:type)
+      def device_type_fallback
+        if Torch.const_defined?(:CUDA) && Torch::CUDA.respond_to?(:available?) && Torch::CUDA.available?
+          return "cuda"
+        end
+
+        if Torch.const_defined?(:Backends) && Torch::Backends.const_defined?(:MPS) &&
+            Torch::Backends::MPS.respond_to?(:available?) && Torch::Backends::MPS.available?
+          return "mps"
+        end
+
+        "cpu"
       rescue
         nil
       end

metadata

@@ -1,14 +1,13 @@
 --- !ruby/object:Gem::Specification
 name: torch-ddp
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
 platform: ruby
 authors:
 - Ivan Razuvaev
-autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-12-19 00:00:00.000000000 Z
+date: 1980-01-02 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: torch-rb
@@ -16,29 +15,28 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 0.22.2
+        version: 0.25.0
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 0.22.2
+        version: 0.25.0
 - !ruby/object:Gem::Dependency
   name: rice
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '4.7'
+        version: '4.8'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '4.7'
-description:
+        version: '4.8'
 email: i@orlando-labs.com
 executables:
 - torchrun
@@ -58,6 +56,7 @@ files:
 - lib/torch-ddp.rb
 - lib/torch/ddp/monkey_patch.rb
 - lib/torch/ddp/version.rb
+- lib/torch/ddp_ext.so
 - lib/torch/distributed.rb
 - lib/torch/nn/parallel/distributed_data_parallel.rb
 - lib/torch/torchrun.rb
@@ -71,7 +70,6 @@ homepage: https://github.com/ankane/torch.rb
 licenses:
 - BSD-3-Clause
 metadata: {}
-post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -79,15 +77,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '3.2'
+      version: '3.3'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.5.22
-signing_key:
+rubygems_version: 4.0.1
 specification_version: 4
 summary: Distributed data parallel support for torch-rb
 test_files: []