Module: Torch::Distributed

Defined in:
lib/torch/distributed.rb

Defined Under Namespace

Modules: ReduceOp Classes: FileStore, HashStore, TCPStore

Constant Summary collapse

DEFAULT_DEVICE_BACKENDS = 
{
  "cpu" => "gloo",
  "cuda" => "nccl",
  "xpu" => "xccl",
  "mps" => "gloo"
}.freeze
DEFAULT_TIMEOUT = 
30 * 60
SPAWN_ENV_KEY = 
"TORCH_DISTRIBUTED_SPAWNED".freeze
SPAWN_RANK_ENV_KEY = 
"TORCH_DISTRIBUTED_SPAWN_RANK".freeze
SPAWN_WORLD_SIZE_ENV_KEY = 
"TORCH_DISTRIBUTED_SPAWN_WORLD_SIZE".freeze
SPAWN_PORT_ENV_KEY = 
"TORCH_DISTRIBUTED_SPAWN_PORT".freeze
SPAWN_PIPE_ENV_KEY = 
"TORCH_DISTRIBUTED_SPAWN_PIPE".freeze
SPAWN_SCRIPT_ENV_KEY = 
"TORCH_DISTRIBUTED_SPAWN_SCRIPT".freeze
SPAWN_TEST_ENV_KEY = 
"TORCH_DISTRIBUTED_SPAWN_TEST".freeze
SPAWN_ARGV = 
ARGV.dup.freeze

Class Method Summary collapse

Class Method Details

.all_reduce(tensor, op: ReduceOp::SUM, group: nil) ⇒ Object 



96
97
98
99
 
# File 'lib/torch/distributed.rb', line 96

def all_reduce(tensor, op: ReduceOp::SUM, group: nil)
  ensure_process_group!(group)
  _all_reduce(tensor, op, group)
end

.barrier(group: nil) ⇒ Object 



91
92
93
94
 
# File 'lib/torch/distributed.rb', line 91

def barrier(group: nil)
  ensure_process_group!(group)
  _barrier(group)
end

.broadcast(tensor, src:, group: nil) ⇒ Object 



101
102
103
104
 
# File 'lib/torch/distributed.rb', line 101

def broadcast(tensor, src:, group: nil)
  ensure_process_group!(group)
  _broadcast(tensor, src, group)
end

.default_process_group ⇒ Object 



77
78
79
 
# File 'lib/torch/distributed.rb', line 77

def default_process_group
  _default_process_group
end

.destroy_process_group ⇒ Object 



73
74
75
 
# File 'lib/torch/distributed.rb', line 73

def destroy_process_group
  _destroy_process_group
end

.fork_spawn_world(world_size, host:, start_method:, &block) ⇒ Object 



134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
 
# File 'lib/torch/distributed.rb', line 134

def fork_spawn_world(world_size, host:, start_method:, &block)
  port = free_port(host: host)
  readers = []
  pids = []
  pgid = nil
  completed = false

  begin
    world_size.times do |rank|
      reader, writer = IO.pipe
      begin
        case start_method
        when :fork
          pids << fork_worker(reader, writer, rank, port, world_size, &block)
        when :spawn
          pid, pgid = spawn_worker(reader, writer, rank, port, host: host, world_size: world_size, pgid: pgid)
          pids << pid
        else
          raise ArgumentError, "Unsupported start_method: #{start_method.inspect}"
        end
        readers << reader
        writer.close unless writer.closed?
      rescue Exception
        reader.close unless reader.closed?
        writer.close unless writer.closed?
        raise
      end
    end

    read_failure = Object.new

    outputs = readers.map do |reader|
      begin
        Marshal.load(reader)
      rescue EOFError
        read_failure
      ensure
        reader.close unless reader.closed?
      end
    end

    statuses = pids.each_with_index.map do |pid, idx|
      _pid, status = Process.wait2(pid)
      [idx, pid, status]
    end

    statuses.each do |idx, pid, status|
      output = outputs[idx]
      if output.equal?(read_failure)
        raise Torch::Error, "Child #{pid} closed pipe before sending result (status #{status.exitstatus})"
      end
      if !status.success? || (output.is_a?(Hash) && output[:error])
        message = if output.is_a?(Hash) && output[:error]
          "Child #{pid} failed: #{output[:error]}\n#{Array(output[:backtrace]).join("\n")}"
        else
          "Child #{pid} exited with status #{status.exitstatus}"
        end
        raise Torch::Error, message
      end
    end

    completed = true
    outputs
  ensure
    # Ensure child workers are cleaned up if an interrupt or error occurs.
    terminate_processes(pids, pgid: pgid) unless completed
  end
end

.fork_world(world_size, host: "127.0.0.1", start_method: :fork, &block) ⇒ Object

Raises:

  • (ArgumentError) 


124
125
126
127
128
129
130
131
132
 
# File 'lib/torch/distributed.rb', line 124

def fork_world(world_size, host: "127.0.0.1", start_method: :fork, &block)
  raise ArgumentError, "world_size must be positive" unless world_size.to_i.positive?
  raise ArgumentError, "block required" unless block

  start_method = normalize_start_method(start_method)
  return run_spawn_worker(&block) if start_method == :spawn && spawn_worker?

  fork_spawn_world(world_size, host: host, start_method: start_method, &block)
end

.free_port(host: "127.0.0.1") ⇒ Object 



203
204
205
206
207
208
 
# File 'lib/torch/distributed.rb', line 203

def free_port(host: "127.0.0.1")
  server = TCPServer.new(host, 0)
  port = server.addr[1]
  server.close
  port
end

.get_default_backend_for_device(device) ⇒ Object

Raises:

  • (ArgumentError) 


118
119
120
121
122
 
# File 'lib/torch/distributed.rb', line 118

def get_default_backend_for_device(device)
  backend = DEFAULT_DEVICE_BACKENDS[device_type_from(device)]
  raise ArgumentError, "Default backend not registered for device: #{device.inspect}" unless backend
  backend
end

.get_rank(group = nil) ⇒ Object 



86
87
88
89
 
# File 'lib/torch/distributed.rb', line 86

def get_rank(group = nil)
  ensure_process_group!(group)
  _get_rank(group)
end

.get_world_size(group = nil) ⇒ Object 



81
82
83
84
 
# File 'lib/torch/distributed.rb', line 81

def get_world_size(group = nil)
  ensure_process_group!(group)
  _get_world_size(group)
end

.init_process_group(backend = nil, init_method: "env://", store: nil, rank: nil, world_size: nil, timeout: DEFAULT_TIMEOUT, wait_for_workers: true, device_id: nil) ⇒ Object

Raises:

  • (Torch::Error) 


37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
 
# File 'lib/torch/distributed.rb', line 37

def init_process_group(backend = nil, init_method: "env://", store: nil, rank: nil, world_size: nil, timeout: DEFAULT_TIMEOUT, wait_for_workers: true, device_id: nil)
  raise Torch::Error, "torch.distributed is not available" unless available?

  backend ||= default_backend_for(device_id)

  if store.nil?
    case init_method
    when "env://"
      rank = Integer(ENV.fetch("RANK")) if rank.nil?
      world_size = Integer(ENV.fetch("WORLD_SIZE")) if world_size.nil?
      master_addr = ENV.fetch("MASTER_ADDR", "127.0.0.1")
      master_port = Integer(ENV.fetch("MASTER_PORT", "29500"))
      raise ArgumentError, "rank is required" if rank.nil?
      raise ArgumentError, "world_size is required" if world_size.nil?
      is_master = rank.zero?
      store = TCPStore.new(master_addr, master_port, world_size, is_master, wait_for_workers: wait_for_workers, timeout: timeout)
    else
      raise ArgumentError, "store is required when using init_method=#{init_method.inspect}"
    end
  end

  raise ArgumentError, "rank is required" if rank.nil?
  raise ArgumentError, "world_size is required" if world_size.nil?

  device_id ||= default_device_id_for_backend(backend, rank, world_size)

  timeout_ms = (timeout * 1000).to_i
  bound_device_id = device_id.nil? ? -1 : Integer(device_id)
  # 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

.initialized? ⇒ Boolean

Returns:

  • (Boolean) 


33
34
35
 
# File 'lib/torch/distributed.rb', line 33

def initialized?
  _initialized?
end

.register_ddp_hook(tensor, process_group, world_size) ⇒ Object 



106
107
108
109
110
111
112
113
114
115
116
 
# File 'lib/torch/distributed.rb', line 106

def register_ddp_hook(tensor, process_group, world_size)
  ensure_process_group!(process_group)
  _register_ddp_hook(tensor, process_group, Integer(world_size))
rescue NoMethodError
  # Fallback for environments built without the native helper; this may
  # still call back into Ruby from autograd threads.
  tensor.register_hook do |grad|
    all_reduce(grad, group: process_group)
    grad.div!(world_size.to_f)
  end
end