torch-ddp 0.1.4 → 0.2.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 341ecc522b0ed689920f0841ef56512531aa68706fc90768ee32f0234100a3d6
-  data.tar.gz: f6ae0221f15cc9d793c3151865ed9d0189a74efba3122cff08b4f63aa829c754
+  metadata.gz: 01c3dafea6b9ee6b691d59448137a325d54bd794c4004ac6cf519ac6d2961bc3
+  data.tar.gz: fbb84030108e52d0f54ce67ac43bb07ab1608d3878c8988bfcf91073e9e3f292
 SHA512:
-  metadata.gz: afb85d047d017eb59f1001eac7bb73bd81a33001c0ec9351d4e578f898530acc015de83cb8fe9137580ba93399abe910652fefa1118c103724ee314f5f0cb2f4
-  data.tar.gz: 30e2153309523ef5aaab848317f2e60b2fce16ab90950e0c9f7327e726bf8e3c47af4c99dc52a7c80719aad3680dcc0cc2a79ed040654d11148d7b2a0de12dd9
+  metadata.gz: f378128687e1c9db4deb53d01b5e10918d76456f586f1d9c35cca9ccdd3ac0f3effebf98724ad52d64c1ee38a4019e2996057791ddcc0cb6232e99a3263dc893
+  data.tar.gz: 2f028ae611ba90f2f664dbca9b37168f725d5bbe6743c2bd633e40e6d53e9083264e51407bb1365a699b17336a24454995866a492e01e16fb8a96cac0d89ce25

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

@@ -0,0 +1,68 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+#if defined(WITH_CUDA)
+#include <cuda_runtime_api.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+#endif
+
+namespace {
+
+void register_cuda_helpers(Rice::Module& m) {
+  auto rb_mDDP = Rice::define_module_under(m, "DDP");
+
+  rb_mDDP.define_singleton_function(
+      "_cuda_set_device",
+      [](int device_id) {
+#if defined(WITH_CUDA)
+        int count = 0;
+        auto status = cudaGetDeviceCount(&count);
+        if (status != cudaSuccess) {
+          rb_raise(
+              rb_eRuntimeError,
+              "cudaGetDeviceCount failed with code %d",
+              static_cast<int>(status));
+        }
+        if (device_id < 0 || device_id >= count) {
+          rb_raise(
+              rb_eArgError,
+              "Invalid device_id %d for CUDA (available devices: %d)",
+              device_id,
+              count);
+        }
+        status = cudaSetDevice(device_id);
+        if (status != cudaSuccess) {
+          rb_raise(
+              rb_eRuntimeError,
+              "cudaSetDevice(%d) failed with code %d",
+              device_id,
+              static_cast<int>(status));
+        }
+#else
+        rb_raise(
+            rb_eRuntimeError,
+            "Torch::DDP._cuda_set_device requires CUDA support");
+#endif
+        return Rice::Object(Qnil);
+      });
+
+  rb_mDDP.define_singleton_function(
+      "_cuda_empty_cache",
+      []() {
+#if defined(WITH_CUDA)
+        c10::cuda::CUDACachingAllocator::emptyCache();
+#else
+        rb_raise(
+            rb_eRuntimeError,
+            "Torch::DDP._cuda_empty_cache requires CUDA support");
+#endif
+        return Rice::Object(Qnil);
+      });
+}
+
+} // namespace
+
+void init_cuda_helpers(Rice::Module& m) {
+  register_cuda_helpers(m);
+}

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
@@ -254,7 +254,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 +267,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 +293,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/ext.cpp

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

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,12 +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")
@@ -69,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"
@@ -79,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>
 
@@ -102,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>
@@ -116,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

@@ -1,5 +1,3 @@
-require "fiddle"
-
 module Torch
   module DDP
     module MonkeyPatch
@@ -14,7 +12,8 @@ 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_device_helpers
+          patch_cuda_empty_cache if missing.include?(:cuda_empty_cache)
+          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 +24,19 @@ module Torch
         def missing_features
           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)
@@ -56,48 +63,33 @@ module Torch
         end
 
         def cuda_set_device!(device_id)
-          cuda_set_device_proc.call(Integer(device_id))
+          unless Torch.const_defined?(:DDP) && Torch::DDP.respond_to?(:_cuda_set_device)
+            raise Torch::Error, "Torch::CUDA.set_device is unavailable; ensure torch is built with CUDA or upgrade torch."
+          end
+
+          Torch::DDP._cuda_set_device(Integer(device_id))
         end
         public :cuda_set_device!
 
-        def cuda_set_device_proc
-          @cuda_set_device_proc ||= begin
-            candidates = [
-              ENV["LIBCUDART_PATH"],
-              "/usr/local/cuda/lib64/libcudart.so",
-              "/usr/local/cuda/lib/libcudart.so",
-              "/usr/local/cuda/lib/libcudart.dylib",
-              "libcudart.so.12",
-              "libcudart.so.11",
-              "libcudart.so",
-              "libcudart.dylib"
-            ].compact
-
-            function = nil
-            candidates.each do |path|
-              begin
-                handle = Fiddle.dlopen(path)
-                function = Fiddle::Function.new(handle["cudaSetDevice"], [Fiddle::TYPE_INT], Fiddle::TYPE_INT)
-                break
-              rescue Fiddle::DLError
-                next
-              end
-            end
+        def patch_cuda_empty_cache
+          return unless Torch.const_defined?(:CUDA)
 
-            if function
-              ->(device_id) do
-                result = function.call(device_id)
-                raise Torch::Error, "cudaSetDevice(#{device_id}) failed with code #{result}" unless result.zero?
-                nil
-              end
-            else
-              ->(device_id) do
-                raise Torch::Error, "Torch::CUDA.set_device is unavailable; ensure torch is built with CUDA or upgrade torch."
-              end
+          Torch::CUDA.singleton_class.class_eval do
+            define_method(:empty_cache) do
+              Torch::DDP::MonkeyPatch.cuda_empty_cache!
             end
           end
         end
 
+        def cuda_empty_cache!
+          unless Torch.const_defined?(:DDP) && Torch::DDP.respond_to?(:_cuda_empty_cache)
+            raise Torch::Error, "Torch::CUDA.empty_cache is unavailable; ensure torch is built with CUDA or upgrade torch."
+          end
+
+          Torch::DDP._cuda_empty_cache
+        end
+        public :cuda_empty_cache!
+
         def patch_device_helpers
           Torch::Device.class_eval do
             alias_method :_torch_ddp_original_to_s, :to_s unless method_defined?(:_torch_ddp_original_to_s)

data/lib/torch/ddp/version.rb

@@ -1,5 +1,5 @@
 module Torch
   module DDP
-    VERSION = "0.1.4"
+    VERSION = "0.2.5"
   end
 end

data/lib/torch/distributed.rb

@@ -264,17 +264,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 +283,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.1.4
+  version: 0.2.5
 platform: ruby
 authors:
 - Ivan Razuvaev
-autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-12-16 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
@@ -51,12 +49,14 @@ files:
 - bin/torchrun
 - examples/benchmark/training.rb
 - examples/mnist/distributed.rb
+- ext/torch_ddp/cuda.cpp
 - ext/torch_ddp/distributed.cpp
 - ext/torch_ddp/ext.cpp
 - ext/torch_ddp/extconf.rb
 - 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
@@ -70,7 +70,6 @@ homepage: https://github.com/ankane/torch.rb
 licenses:
 - BSD-3-Clause
 metadata: {}
-post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -78,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: []