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 



98
99
100
101
 
# File 'lib/torch/distributed.rb', line 98

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

.barrier(group: nil) ⇒ Object 



93
94
95
96
 
# File 'lib/torch/distributed.rb', line 93

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

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



103
104
105
106
 
# File 'lib/torch/distributed.rb', line 103

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

.default_process_group ⇒ Object 



79
80
81
 
# File 'lib/torch/distributed.rb', line 79

def default_process_group
  _default_process_group
end

.destroy_process_group ⇒ Object 



75
76
77
 
# File 'lib/torch/distributed.rb', line 75

def destroy_process_group
  _destroy_process_group
end

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



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
202
203
 
# File 'lib/torch/distributed.rb', line 136

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) 


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

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 



205
206
207
208
209
210
 
# File 'lib/torch/distributed.rb', line 205

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) 


120
121
122
123
124
 
# File 'lib/torch/distributed.rb', line 120

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 



88
89
90
91
 
# File 'lib/torch/distributed.rb', line 88

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

.get_world_size(group = nil) ⇒ Object 



83
84
85
86
 
# File 'lib/torch/distributed.rb', line 83

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
72
73
 
# 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)
  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
  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 



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

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