torch-dl 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 395b034959ba88dc532f1b9903dfcbf5404902117d5b833c112ee7e3851b20e5
+  data.tar.gz: f7dce1d7003eafc711e172f0877d7929b01055d879a4057314ba22b26cbbb53e
+SHA512:
+  metadata.gz: 2bb333d0eb7bc3d2e9785d7a66e82d1340cc25547d760296cb0ae47f20ba5517f4fa2c03709c81868612d1313ac6ed4e4e4808041d20de6a66a3e4e00c9f8454
+  data.tar.gz: 1c78b7ada3a7f669db54ea24fe385c0e2811418ce1546836db6fa31549979e14fee9c1466b5dd34b79eae0541705a4db7cf4de59d66ccc19a4a221ad17cf26f6

data/CHANGELOG.md

@@ -0,0 +1,9 @@
+# Changelog
+
+## 0.1.0 (2025-01-01)
+
+- Initial release
+- DataParallel for multi-GPU training
+- Scatter/gather operations
+- Thread-based parallel execution
+- Replica caching for performance

data/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Chris Hasinski
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1,115 @@
+# torch-dl
+
+Multi-GPU training for [torch.rb](https://github.com/ankane/torch.rb). Split batches across multiple GPUs automatically.
+
+## Installation
+
+Add to your Gemfile:
+
+```ruby
+gem "torch-dl"
+```
+
+## Usage
+
+### Basic Usage
+
+```ruby
+require "torch_dl"
+
+# Create your model on the first GPU
+model = MyModel.new.to("cuda:0")
+
+# Wrap with DataParallel
+dp_model = DataParallel.new(model, device_ids: [0, 1])
+
+# Training loop
+optimizer.zero_grad
+output = dp_model.call(input)
+loss = criterion.call(output, target)
+loss.backward
+optimizer.step
+```
+
+### Models That Return Loss
+
+If your model returns a scalar loss (e.g., GPT models returning `[logits, loss]`), use `dp_model.backward` instead of `loss.backward`:
+
+```ruby
+optimizer.zero_grad
+logits, loss = dp_model.call(input, targets: targets)
+dp_model.backward(scale: 1.0)
+optimizer.step
+```
+
+This is necessary because gathering scalar tensors across devices breaks the autograd graph. The `backward` method calls backward on each replica's loss separately, then reduces gradients to the original module.
+
+### Gradient Accumulation
+
+For gradient accumulation, scale the backward pass:
+
+```ruby
+gradient_accumulation_steps = 4
+
+(0...gradient_accumulation_steps).each do |step|
+  logits, loss = dp_model.call(input_batch, targets: targets_batch)
+  dp_model.backward(scale: 1.0 / gradient_accumulation_steps)
+end
+optimizer.step
+optimizer.zero_grad
+```
+
+## API Reference
+
+### DataParallel
+
+```ruby
+DataParallel.new(model, device_ids: nil, output_device: nil, dim: 0)
+```
+
+- `model` - The module to parallelize (must be on `cuda:0`)
+- `device_ids` - Array of GPU indices to use (default: all available)
+- `output_device` - GPU index for output (default: first device)
+- `dim` - Dimension to scatter inputs along (default: 0, batch dimension)
+
+#### Methods
+
+- `call(*inputs, **kwargs)` - Forward pass with automatic scattering/gathering
+- `backward(scale: 1.0)` - Backward pass for loss-returning models
+- `module` / `wrapped_module` - Access the underlying model
+- `parameters` - Access model parameters (for optimizer)
+- `state_dict` / `load_state_dict` - Save/load model state
+- `train` / `eval` - Set training/evaluation mode
+
+### Low-level Functions
+
+```ruby
+# Split tensor across devices
+TorchDL.scatter(tensor, ["cuda:0", "cuda:1"], dim)
+
+# Gather tensors to a single device
+TorchDL.gather(tensors, "cuda:0", dim)
+```
+
+## How It Works
+
+1. **Scatter**: Input batch is split across GPUs
+2. **Replicate**: Model is copied to each GPU (cached for performance)
+3. **Parallel Apply**: Forward pass runs on each GPU in parallel using threads
+4. **Gather**: Outputs are collected back to the output device
+
+## Requirements
+
+- Ruby 3.1+
+- torch.rb 0.17+
+- Multiple CUDA GPUs
+
+## Notes
+
+- Works with stock torch.rb from RubyGems
+- Uses FFI (fiddle) to call CUDA runtime directly for `synchronize`, `current_device`, and `set_device`
+- No C extension required - pure Ruby gem
+
+## License
+
+MIT

data/lib/torch_dl/cuda.rb

@@ -0,0 +1,76 @@
+require "fiddle"
+
+module TorchDL
+  module CUDA
+    class << self
+      def synchronize
+        return unless Torch::CUDA.available?
+
+        @cudart ||= load_cudart
+        return unless @cudart
+
+        # cudaDeviceSynchronize returns cudaError_t (0 = success)
+        @cuda_device_synchronize ||= Fiddle::Function.new(
+          @cudart["cudaDeviceSynchronize"],
+          [],
+          Fiddle::TYPE_INT
+        )
+        @cuda_device_synchronize.call
+      end
+
+      def current_device
+        return -1 unless Torch::CUDA.available?
+
+        @cudart ||= load_cudart
+        return -1 unless @cudart
+
+        @cuda_get_device ||= Fiddle::Function.new(
+          @cudart["cudaGetDevice"],
+          [Fiddle::TYPE_VOIDP],
+          Fiddle::TYPE_INT
+        )
+
+        device_ptr = Fiddle::Pointer.malloc(Fiddle::SIZEOF_INT)
+        @cuda_get_device.call(device_ptr)
+        device_ptr[0, Fiddle::SIZEOF_INT].unpack1("i")
+      end
+
+      def set_device(device_id)
+        return unless Torch::CUDA.available?
+
+        @cudart ||= load_cudart
+        return unless @cudart
+
+        @cuda_set_device ||= Fiddle::Function.new(
+          @cudart["cudaSetDevice"],
+          [Fiddle::TYPE_INT],
+          Fiddle::TYPE_INT
+        )
+        @cuda_set_device.call(device_id)
+      end
+
+      private
+
+      def load_cudart
+        # Try common CUDA runtime library paths
+        paths = [
+          "libcudart.so",
+          "libcudart.so.12",
+          "libcudart.so.11",
+          "/usr/local/cuda/lib64/libcudart.so",
+          "/usr/lib/x86_64-linux-gnu/libcudart.so"
+        ]
+
+        paths.each do |path|
+          begin
+            return Fiddle.dlopen(path)
+          rescue Fiddle::DLError
+            next
+          end
+        end
+
+        nil
+      end
+    end
+  end
+end

data/lib/torch_dl/ext.rb

@@ -0,0 +1,20 @@
+module TorchDL
+  # Pure Ruby implementation of scatter/gather operations
+  module Ext
+    class << self
+      # Splits a tensor across devices
+      def scatter(input, devices, dim = 0)
+        chunks = input.chunk(devices.size, dim)
+        chunks.each_with_index.map do |chunk, i|
+          chunk.to(devices[i])
+        end
+      end
+
+      # Gathers tensors from multiple devices onto a single device
+      def gather(inputs, target_device, dim = 0)
+        on_target = inputs.map { |t| t.to(target_device) }
+        Torch.cat(on_target, dim)
+      end
+    end
+  end
+end

data/lib/torch_dl/nn/parallel/data_parallel.rb

@@ -0,0 +1,273 @@
+module TorchDL
+  module NN
+    # Implements data parallelism at the module level.
+    #
+    # This container parallelizes the application of the given module by
+    # splitting the input across the specified devices by chunking in the
+    # batch dimension. In the forward pass, the module is replicated on each
+    # device, and each replica handles a portion of the input. During the
+    # backwards pass, gradients from each replica are summed into the
+    # original module.
+    #
+    # @note Backward Pass for Models Returning Loss
+    #   When your model returns a scalar loss (e.g., GPT models that return
+    #   [logits, loss]), you must use the {#backward} method instead of
+    #   calling loss.backward directly. This is because gathering scalar
+    #   tensors across devices breaks the autograd graph in torch.rb.
+    #
+    #   @example Training loop with loss-returning model
+    #     dp_model = TorchDL::NN::DataParallel.new(model, device_ids: [0, 1])
+    #     optimizer.zero_grad
+    #     logits, loss = dp_model.call(input, targets: targets)
+    #     dp_model.backward(scale: 1.0 / gradient_accumulation_steps)
+    #     optimizer.step
+    #
+    # @example Basic usage (model returns output only)
+    #   model = MyModel.new.to("cuda:0")
+    #   dp_model = TorchDL::NN::DataParallel.new(model, device_ids: [0, 1])
+    #   output = dp_model.call(input)
+    #   loss = criterion.call(output, target)
+    #   loss.backward  # Standard backward works when loss computed after gather
+    #
+    class DataParallel < Torch::NN::Module
+      attr_reader :module, :device_ids, :output_device, :dim
+      alias_method :wrapped_module, :module
+
+      # @param mod [Torch::NN::Module] Module to parallelize
+      # @param device_ids [Array<Integer>, nil] CUDA devices (default: all available)
+      # @param output_device [Integer, nil] Device for output (default: device_ids[0])
+      # @param dim [Integer] Dimension to scatter inputs along (default: 0)
+      def initialize(mod, device_ids: nil, output_device: nil, dim: 0)
+        super()
+        @module = mod
+        @device_ids = device_ids || (0...Torch::CUDA.device_count).to_a
+        @output_device = output_device || @device_ids.first
+        @dim = dim
+        @replica_losses = nil
+
+        if @device_ids.empty?
+          raise ArgumentError, "device_ids cannot be empty"
+        end
+
+        # Convert to device strings for internal use
+        @device_strings = @device_ids.map { |id| "cuda:#{id}" }
+        @output_device_string = "cuda:#{@output_device}"
+      end
+
+      def forward(*inputs, **kwargs)
+        # Empty input check
+        if inputs.empty?
+          return @module.call(**kwargs)
+        end
+
+        # Single GPU fast path
+        if @device_ids.size == 1
+          return @module.call(*inputs.map { |i| i.to(@device_strings.first) }, **kwargs)
+        end
+
+        # Scatter inputs across devices
+        scattered_inputs = scatter(inputs, @device_strings, @dim)
+        scattered_kwargs = scatter_kwargs(kwargs, @device_strings, @dim)
+
+        # Get or create replicas, sync weights
+        num_replicas = scattered_inputs.size
+        devices = @device_strings[0...num_replicas]
+        @replicas = get_replicas(devices)
+        sync_replica_weights
+
+        # Apply in parallel
+        outputs = parallel_apply(@replicas, scattered_inputs, scattered_kwargs)
+
+        # Ensure all CUDA operations complete before gathering
+        TorchDL::CUDA.synchronize if Torch::CUDA.available?
+
+        # Gather outputs back to output device
+        gather(outputs, @output_device_string, @dim)
+      end
+
+      # Performs backward pass on all replica losses and reduces gradients.
+      # This is needed because gather creates a new tensor that breaks the
+      # autograd connection across devices. By calling backward on each
+      # replica's loss separately, gradients flow properly.
+      #
+      # @param scale [Float] Scale factor for gradients (e.g., 1.0/gradient_accumulation_steps)
+      def backward(scale: 1.0)
+        if @replica_losses && @replica_losses.size > 1
+          # Each replica's loss contributes equally to total loss
+          # Scale by 1/N to average gradients across replicas
+          replica_scale = scale / @replica_losses.size
+
+          @replica_losses.each do |replica_loss|
+            # Scale the loss before backward to get properly scaled gradients
+            scaled_loss = replica_loss * replica_scale
+            scaled_loss.backward
+          end
+
+          @replica_losses = nil
+        end
+
+        # Reduce gradients from all replicas to the original module
+        reduce_gradients
+      end
+
+      # Reduce gradients from replicas back to the original module.
+      # Called automatically by backward(), but can be called manually if needed.
+      def reduce_gradients
+        return if @replicas.nil? || @replicas.size <= 1
+
+        # Get original model's parameters (first replica is the original)
+        original_params = @module.parameters.to_a
+
+        # Accumulate gradients from other replicas
+        @replicas[1..].each do |replica|
+          replica_params = replica.parameters.to_a
+          original_params.zip(replica_params).each do |orig, repl|
+            next unless repl.grad
+
+            # Move replica gradient to original device and add
+            if orig.grad
+              orig.grad.add!(repl.grad.to(orig.device))
+            else
+              orig.grad = repl.grad.to(orig.device).clone
+            end
+          end
+        end
+      end
+
+      # Delegate training mode to wrapped module
+      def train(mode = true)
+        super
+        @module.train(mode)
+        self
+      end
+
+      def eval
+        train(false)
+      end
+
+      # Delegate parameter access to wrapped module
+      def parameters
+        @module.parameters
+      end
+
+      def named_parameters(prefix: "", recurse: true)
+        @module.named_parameters(prefix: prefix, recurse: recurse)
+      end
+
+      def state_dict(destination: nil, prefix: "")
+        @module.state_dict(destination: destination, prefix: prefix)
+      end
+
+      def load_state_dict(state_dict, strict: true)
+        @module.load_state_dict(state_dict, strict: strict)
+      end
+
+      def extra_inspect
+        format("device_ids: %s, output_device: %s, dim: %d", @device_ids, @output_device, @dim)
+      end
+
+      private
+
+      # Scatter a single value across devices
+      def scatter_value(value, devices, dim)
+        if value.is_a?(Torch::Tensor)
+          TorchDL::Ext.scatter(value, devices, dim).map(&:contiguous)
+        else
+          devices.map { value }
+        end
+      end
+
+      def scatter(inputs, devices, dim)
+        # Scatter each input, then transpose to group by device
+        scattered = inputs.map { |input| scatter_value(input, devices, dim) }
+        scattered.first.size.times.map { |i| scattered.map { |s| s[i] } }
+      end
+
+      def scatter_kwargs(kwargs, devices, dim)
+        return devices.map { {} } if kwargs.empty?
+        scattered = kwargs.transform_values { |v| scatter_value(v, devices, dim) }
+        devices.size.times.map { |i| scattered.transform_values { |v| v[i] } }
+      end
+
+      def get_replicas(devices)
+        # Return cached replicas if they match the requested devices
+        if @replicas && @replica_devices == devices
+          return @replicas
+        end
+
+        # Create new replicas and cache them
+        @replica_devices = devices
+        Replicate.replicate(@module, devices)
+      end
+
+      def sync_replica_weights
+        return if @replicas.nil? || @replicas.size <= 1
+
+        # Sync parameters in-place to avoid memory allocation
+        original_params = @module.parameters.to_a
+        @replicas[1..].each do |replica|
+          replica.parameters.to_a.each_with_index do |param, i|
+            param.data.copy!(original_params[i].data)
+          end
+        end
+      end
+
+      def parallel_apply(replicas, inputs, kwargs_list)
+        ParallelApply.parallel_apply(replicas, inputs, kwargs_list)
+      end
+
+      def gather(outputs, target_device, dim)
+        # Handle different output types
+        first = outputs.first
+
+        case first
+        when Torch::Tensor
+          gather_tensor(outputs, target_device, dim)
+        when Array
+          # Tuple/array of tensors - gather each element
+          first.size.times.map do |i|
+            tensors = outputs.map { |o| o[i] }
+            if tensors.first.is_a?(Torch::Tensor)
+              gather_tensor(tensors, target_device, dim)
+            elsif tensors.first.nil?
+              nil
+            else
+              tensors.first # Non-tensor, just return first
+            end
+          end
+        when Hash
+          # Dict of tensors - gather each value
+          first.keys.map do |key|
+            tensors = outputs.map { |o| o[key] }
+            if tensors.first.is_a?(Torch::Tensor)
+              [key, gather_tensor(tensors, target_device, dim)]
+            else
+              [key, tensors.first]
+            end
+          end.to_h
+        else
+          # Scalar or other - return first output
+          first
+        end
+      end
+
+      def gather_tensor(tensors, target_device, dim)
+        first = tensors.first
+        # Handle scalar tensors (0-dim) - store for backward, return average for display
+        if first.dim == 0
+          # Store individual losses for backward() method
+          @replica_losses = tensors
+
+          # Return mean of losses (moved to same device) for logging/display
+          # Note: This returned tensor should NOT be used for backward() - use backward() instead
+          sum = tensors.reduce(Torch.tensor(0.0, device: target_device)) do |acc, t|
+            acc + t.to(target_device).detach
+          end
+          sum / tensors.size
+        else
+          TorchDL::Ext.gather(tensors, target_device, dim)
+        end
+      end
+    end
+  end
+end

data/lib/torch_dl/nn/parallel/parallel_apply.rb

@@ -0,0 +1,60 @@
+module TorchDL
+  module NN
+    module ParallelApply
+      class << self
+        # Applies modules to inputs in parallel across devices.
+        #
+        # @param modules [Array<Torch::NN::Module>] List of module replicas
+        # @param inputs [Array] List of inputs, one per module
+        # @param kwargs_list [Array<Hash>, nil] Optional list of kwargs, one per module
+        # @return [Array] List of outputs, one per module
+        def parallel_apply(modules, inputs, kwargs_list = nil)
+          kwargs_list ||= modules.map { {} }
+
+          unless modules.size == inputs.size && modules.size == kwargs_list.size
+            raise ArgumentError, "modules, inputs, and kwargs_list must have the same length"
+          end
+
+          # Single module - no parallelism needed
+          if modules.size == 1
+            return [apply_module(modules[0], inputs[0], kwargs_list[0])]
+          end
+
+          parallel_apply_threads(modules, inputs, kwargs_list)
+        end
+
+        private
+
+        def parallel_apply_threads(modules, inputs, kwargs_list)
+          results = Array.new(modules.size)
+          errors = Array.new(modules.size)
+
+          threads = modules.each_with_index.map do |mod, i|
+            Thread.new(mod, inputs[i], kwargs_list[i], i) do |m, inp, kw, idx|
+              results[idx] = apply_module(m, inp, kw)
+            rescue => e
+              errors[idx] = e
+            end
+          end
+
+          threads.each(&:join)
+
+          # Re-raise first error if any
+          errors.each_with_index do |err, i|
+            raise err if err
+          end
+
+          results
+        end
+
+        def apply_module(mod, input, kwargs)
+          if input.is_a?(Array)
+            mod.call(*input, **kwargs)
+          else
+            mod.call(input, **kwargs)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/torch_dl/nn/parallel/replicate.rb

@@ -0,0 +1,178 @@
+module TorchDL
+  module NN
+    module Replicate
+      class << self
+        # Replicates a module on multiple devices.
+        #
+        # @param network [Torch::NN::Module] The module to replicate
+        # @param devices [Array<String>] List of device strings (e.g., ["cuda:0", "cuda:1"])
+        # @return [Array<Torch::NN::Module>] List of module replicas, one per device
+        #
+        # Note: The first device uses the original network (not a copy) to ensure
+        # gradients flow back to the original parameters during backward pass.
+        def replicate(network, devices)
+          devices = devices.map { |d| d.is_a?(String) ? d : "cuda:#{d}" }
+
+          # Single device - just return the network (already on correct device)
+          return [network] if devices.size == 1
+
+          # Get the state dict once for creating replicas
+          state_dict = network.state_dict
+
+          # Create replicas - first device uses original network for gradient flow
+          devices.each_with_index.map do |device, idx|
+            if idx == 0
+              # First device: use the original network to maintain gradient connection
+              network
+            else
+              # Other devices: create independent replicas
+              replica = deep_copy_module(network)
+
+              # Copy state dict tensors to the target device
+              # Filter to only include keys that exist in the replica
+              replica_keys = replica.state_dict.keys
+              device_state = state_dict.select { |k, _| replica_keys.include?(k) }
+                                       .transform_values { |t| t.to(device) }
+              replica.load_state_dict(device_state)
+              replica.to(device)
+              replica.train(network.instance_variable_get(:@training))
+              replica
+            end
+          end
+        end
+
+        private
+
+        # Creates a deep copy of a module structure
+        def deep_copy_module(mod)
+          # Check for custom replication hook first
+          return mod.class._replicate(mod) if mod.class.respond_to?(:_replicate)
+
+          # Handle container modules
+          return copy_sequential(mod) if mod.is_a?(Torch::NN::Sequential)
+          return copy_module_list(mod) if mod.is_a?(Torch::NN::ModuleList)
+
+          # Try built-in module copiers
+          copier = module_copiers[mod.class]
+          return copier.call(mod) if copier
+
+          # Fallback to generic copy for custom modules
+          copy_custom_module(mod)
+        end
+
+        # Lazily initialized registry of module copiers
+        def module_copiers
+          @module_copiers ||= {
+            Torch::NN::Linear => ->(m) { Torch::NN::Linear.new(m.in_features, m.out_features, bias: has_bias?(m)) },
+            Torch::NN::Conv1d => ->(m) { copy_conv(Torch::NN::Conv1d, m) },
+            Torch::NN::Conv2d => ->(m) { copy_conv(Torch::NN::Conv2d, m) },
+            Torch::NN::Conv3d => ->(m) { copy_conv(Torch::NN::Conv3d, m) },
+            Torch::NN::BatchNorm1d => ->(m) { copy_batch_norm(Torch::NN::BatchNorm1d, m) },
+            Torch::NN::BatchNorm2d => ->(m) { copy_batch_norm(Torch::NN::BatchNorm2d, m) },
+            Torch::NN::BatchNorm3d => ->(m) { copy_batch_norm(Torch::NN::BatchNorm3d, m) },
+            Torch::NN::LayerNorm => ->(m) {
+              Torch::NN::LayerNorm.new(m.instance_variable_get(:@normalized_shape),
+                               eps: m.instance_variable_get(:@eps),
+                               elementwise_affine: m.instance_variable_get(:@elementwise_affine))
+            },
+            Torch::NN::Embedding => ->(m) {
+              Torch::NN::Embedding.new(m.instance_variable_get(:@num_embeddings),
+                               m.instance_variable_get(:@embedding_dim),
+                               padding_idx: m.instance_variable_get(:@padding_idx))
+            },
+            Torch::NN::Dropout => ->(m) { Torch::NN::Dropout.new(p: m.instance_variable_get(:@p)) },
+            Torch::NN::Dropout2d => ->(m) { Torch::NN::Dropout2d.new(p: m.instance_variable_get(:@p)) },
+            Torch::NN::Dropout3d => ->(m) { Torch::NN::Dropout3d.new(p: m.instance_variable_get(:@p)) },
+            Torch::NN::LSTM => ->(m) { copy_rnn(Torch::NN::LSTM, m) },
+            Torch::NN::GRU => ->(m) { copy_rnn(Torch::NN::GRU, m) },
+            Torch::NN::RNN => ->(m) { copy_rnn(Torch::NN::RNN, m) },
+            Torch::NN::ReLU => ->(_) { Torch::NN::ReLU.new },
+            Torch::NN::GELU => ->(_) { Torch::NN::GELU.new },
+            Torch::NN::Tanh => ->(_) { Torch::NN::Tanh.new },
+            Torch::NN::Sigmoid => ->(_) { Torch::NN::Sigmoid.new },
+            Torch::NN::Identity => ->(_) { Torch::NN::Identity.new },
+            Torch::NN::Softmax => ->(m) { Torch::NN::Softmax.new(dim: m.instance_variable_get(:@dim)) },
+            Torch::NN::LogSoftmax => ->(m) { Torch::NN::LogSoftmax.new(dim: m.instance_variable_get(:@dim)) },
+          }
+        end
+
+        def has_bias?(mod)
+          !mod.instance_variable_get(:@bias).nil?
+        end
+
+        def copy_conv(klass, mod)
+          klass.new(mod.in_channels, mod.out_channels, mod.kernel_size,
+                    stride: mod.stride, padding: mod.padding, dilation: mod.dilation,
+                    groups: mod.groups, bias: has_bias?(mod), padding_mode: mod.padding_mode)
+        end
+
+        def copy_batch_norm(klass, mod)
+          klass.new(mod.num_features, eps: mod.eps, momentum: mod.momentum,
+                    affine: mod.affine, track_running_stats: mod.track_running_stats)
+        end
+
+        def copy_rnn(klass, mod)
+          klass.new(mod.input_size, mod.hidden_size, num_layers: mod.num_layers,
+                    bias: mod.bias, batch_first: mod.batch_first,
+                    dropout: mod.dropout, bidirectional: mod.bidirectional)
+        end
+
+        def copy_sequential(mod)
+          Torch::NN::Sequential.new(*mod.children.map { |child| deep_copy_module(child) })
+        end
+
+        def copy_module_list(mod)
+          Torch::NN::ModuleList.new(mod.map { |child| deep_copy_module(child) })
+        end
+
+        def copy_custom_module(mod)
+          klass = mod.class
+          children = mod.named_children.to_h
+
+          if children.any?
+            # Module has submodules - create structural copy
+            replica = klass.allocate
+            replica.send(:initialize_module_state)
+            copy_instance_state(mod, replica)
+            children.each do |name, child|
+              replica.instance_variable_set("@#{name}", deep_copy_module(child))
+            end
+            replica
+          else
+            # Leaf module - try clone
+            mod.clone
+          end
+        rescue => e
+          raise ArgumentError, "Cannot replicate #{klass}. " \
+            "Implement #{klass}._replicate(mod) class method. Error: #{e.message}"
+        end
+
+        def copy_instance_state(src, dst)
+          src.instance_variables.each do |ivar|
+            next if %i[@parameters @buffers @modules @training @non_persistent_buffers_set].include?(ivar)
+            val = src.instance_variable_get(ivar)
+            next if val.is_a?(Torch::Tensor) || val.is_a?(Torch::NN::Module)
+            dst.instance_variable_set(ivar, val)
+          end
+        end
+      end
+    end
+  end
+end
+
+# Add helper method to Module for initializing state (if not already defined)
+module Torch
+  module NN
+    class Module
+      private
+
+      def initialize_module_state
+        @training = true
+        @parameters = {}
+        @buffers = {}
+        @modules = {}
+        @non_persistent_buffers_set = Set.new
+      end unless private_method_defined?(:initialize_module_state)
+    end
+  end
+end

data/lib/torch_dl/version.rb

@@ -0,0 +1,3 @@
+module TorchDL
+  VERSION = "0.1.0"
+end

data/lib/torch_dl.rb

@@ -0,0 +1,24 @@
+require "torch"
+require_relative "torch_dl/version"
+require_relative "torch_dl/cuda"
+require_relative "torch_dl/ext"
+require_relative "torch_dl/nn/parallel/replicate"
+require_relative "torch_dl/nn/parallel/parallel_apply"
+require_relative "torch_dl/nn/parallel/data_parallel"
+
+module TorchDL
+  class Error < StandardError; end
+
+  class << self
+    def scatter(input, devices, dim = 0)
+      TorchDL::Ext.scatter(input, devices, dim)
+    end
+
+    def gather(inputs, target_device, dim = 0)
+      TorchDL::Ext.gather(inputs, target_device, dim)
+    end
+  end
+end
+
+# Convenience alias
+DataParallel = TorchDL::NN::DataParallel

metadata

@@ -0,0 +1,84 @@
+--- !ruby/object:Gem::Specification
+name: torch-dl
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Chris Hasinski
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: torch-rb
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.17'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.17'
+- !ruby/object:Gem::Dependency
+  name: fiddle
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+description: DataParallel and distributed training utilities for torch.rb. Split batches
+  across multiple GPUs automatically.
+email:
+- krzysztof.hasinski@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- LICENSE.txt
+- README.md
+- lib/torch_dl.rb
+- lib/torch_dl/cuda.rb
+- lib/torch_dl/ext.rb
+- lib/torch_dl/nn/parallel/data_parallel.rb
+- lib/torch_dl/nn/parallel/parallel_apply.rb
+- lib/torch_dl/nn/parallel/replicate.rb
+- lib/torch_dl/version.rb
+homepage: https://github.com/khasinski/torch-rb-dl
+licenses:
+- MIT
+metadata:
+  allowed_push_host: https://rubygems.org
+  homepage_uri: https://github.com/khasinski/torch-rb-dl
+  source_code_uri: https://github.com/khasinski/torch-rb-dl
+  changelog_uri: https://github.com/khasinski/torch-rb-dl/blob/master/CHANGELOG.md
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.1'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.2
+specification_version: 4
+summary: Multi-GPU training for torch.rb
+test_files: []