torch-ddp 0.2.5 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 01c3dafea6b9ee6b691d59448137a325d54bd794c4004ac6cf519ac6d2961bc3
-  data.tar.gz: fbb84030108e52d0f54ce67ac43bb07ab1608d3878c8988bfcf91073e9e3f292
+  metadata.gz: 9f9b645358a4641e4276ff48adb382c869feb024882d84c4835a39025bda6e62
+  data.tar.gz: 8f141183e503f72f0ec1622209e4ca17dcf8611471f68a45e34e0e2c146e4631
 SHA512:
-  metadata.gz: f378128687e1c9db4deb53d01b5e10918d76456f586f1d9c35cca9ccdd3ac0f3effebf98724ad52d64c1ee38a4019e2996057791ddcc0cb6232e99a3263dc893
-  data.tar.gz: 2f028ae611ba90f2f664dbca9b37168f725d5bbe6743c2bd633e40e6d53e9083264e51407bb1365a699b17336a24454995866a492e01e16fb8a96cac0d89ce25
+  metadata.gz: fe63d5d9cbc27128b15166363c6a1ae84f342a33c66211e43f6464b9fc58a2a70e497d053561bdf176826307ef089a616fd102db2dd8da135b1ac4e8ee9b781c
+  data.tar.gz: 733d6befc6151f55452b3834d2b8c4e174184adda5e8cacc99bc510b65576b26f6e721454d7e88d6309ff2a39290488cc75926f5e137c4ce697ccc0dd53e8827

data/ext/torch_ddp/distributed.cpp

@@ -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
   }
 }
 

data/lib/torch/ddp/version.rb

@@ -1,5 +1,5 @@
 module Torch
   module DDP
-    VERSION = "0.2.5"
+    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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: torch-ddp
 version: !ruby/object:Gem::Version
-  version: 0.2.5
+  version: 0.3.0
 platform: ruby
 authors:
 - Ivan Razuvaev