William Lin
7d61590213
Ray sets CUDA_VISIBLE_DEVICES of each GPU actor to the GPU of that actor. The import of fastvideo initializes CUDA, so the executor cannot change the device list later. On a node with more than one GPU, each worker after the first worker failed with "invalid device ordinal". On NVIDIA GPUs, each worker actor now keeps the device list of its raylet (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1). Each worker uses the position of its GPU in that list as its local rank. The safetensors loader now reads each file on node-group rank 0, which is the rank that sends the broadcast.
531 lines
23 KiB
Python
531 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapt from https://github.com/vllm-project/vllm/blob/releases/v0.11.0/vllm/executor/ray_distributed_executor.py
|
|
|
|
import asyncio
|
|
from collections import defaultdict
|
|
from queue import Queue
|
|
import os
|
|
import cloudpickle
|
|
|
|
import fastvideo.envs as envs
|
|
from dataclasses import dataclass
|
|
|
|
from typing import Any, TYPE_CHECKING
|
|
from collections.abc import Callable
|
|
from fastvideo.utils import (DEPRECATED_HF_TOKEN_ENV_VARS, HF_TOKEN_ENV_VARS, get_ip, get_distributed_init_method,
|
|
get_open_port, get_loopback_ip)
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
|
from fastvideo.worker.executor import Executor
|
|
from fastvideo.worker.ray_utils import (
|
|
initialize_ray_cluster,
|
|
RayWorkerWrapper,
|
|
ray,
|
|
)
|
|
from fastvideo.worker.ray_env import get_env_vars_to_copy
|
|
from fastvideo.logger import init_logger
|
|
|
|
if ray is not None:
|
|
from ray.actor import ActorHandle
|
|
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
|
else:
|
|
ActorHandle = None
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.util.placement_group import PlacementGroup
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def should_use_gloo_loopback(worker_ips: list[str]) -> bool:
|
|
"""Loopback is only safe when every worker shares one host IP.
|
|
|
|
Single-node Ray (one or many GPUs on the same box) can dial the Gloo store
|
|
on loopback. Two Sparks already have distinct worker IPs, so this returns
|
|
False and Gloo stays on the fabric address. Per-node NIC names are a
|
|
separate issue: do not copy ``NCCL_SOCKET_IFNAME`` / ``GLOO_SOCKET_IFNAME``
|
|
from the driver onto those workers.
|
|
"""
|
|
return len(set(worker_ips)) <= 1
|
|
|
|
|
|
# Ray narrows each GPU actor's CUDA_VISIBLE_DEVICES to the actor's own GPU unless this is set.
|
|
RAY_NOSET_CUDA_VISIBLE_DEVICES = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"
|
|
|
|
|
|
def keep_raylet_cuda_devices(ray_remote_kwargs: dict[str, Any]) -> dict[str, Any]:
|
|
"""Return actor options under which Ray leaves CUDA_VISIBLE_DEVICES as the raylet's.
|
|
|
|
Ray merges an actor's ``env_vars`` into the job's runtime environment, so
|
|
only this one variable is added.
|
|
"""
|
|
runtime_env = dict(ray_remote_kwargs.get("runtime_env") or {})
|
|
runtime_env["env_vars"] = {**(runtime_env.get("env_vars") or {}), RAY_NOSET_CUDA_VISIBLE_DEVICES: "1"}
|
|
return {**ray_remote_kwargs, "runtime_env": runtime_env}
|
|
|
|
|
|
def ray_worker_device_ordinal(node_gpu_ids: list[int], node_local_index: int, visible_devices: str | None) -> int:
|
|
"""Return the CUDA ordinal of the ``node_local_index``-th worker on a node.
|
|
|
|
The worker uses ``node_gpu_ids[node_local_index]`` (the sorted Ray GPU IDs
|
|
of the node's workers), addressed in the device list it inherited from its
|
|
raylet. Ray reports GPU IDs as entries of the raylet's CUDA_VISIBLE_DEVICES,
|
|
or as ordinals when that variable is unset.
|
|
"""
|
|
gpu_id = node_gpu_ids[node_local_index]
|
|
if not visible_devices:
|
|
return gpu_id
|
|
visible = [entry.strip() for entry in visible_devices.split(",")]
|
|
if str(gpu_id) not in visible:
|
|
raise RuntimeError(f"Ray assigned GPU {gpu_id}, which is not in the worker's "
|
|
f"CUDA_VISIBLE_DEVICES={visible_devices!r}.")
|
|
return visible.index(str(gpu_id))
|
|
|
|
|
|
@dataclass
|
|
class RayWorkerMetaData:
|
|
"""
|
|
Metadata for a Ray worker.
|
|
The order of ray worker creation can be random,
|
|
and we need to reset the rank after creating all workers.
|
|
"""
|
|
|
|
worker: ActorHandle
|
|
created_rank: int
|
|
adjusted_rank: int = -1
|
|
ip: str = ""
|
|
|
|
|
|
class RayDistributedExecutor(Executor):
|
|
"""Ray-based distributed executor"""
|
|
|
|
# These env vars are worker-specific, therefore are NOT copied
|
|
# from the driver to the workers
|
|
WORKER_SPECIFIC_ENV_VARS = {
|
|
"FASTVIDEO_HOST_IP",
|
|
"LOCAL_RANK",
|
|
"CUDA_VISIBLE_DEVICES",
|
|
}
|
|
|
|
# Per-node fabric names. spark_pair_env.sh / ibdev2netdev can differ across
|
|
# boxes; pushing the driver's value overwrites the export set before ray start.
|
|
WORKER_LOCAL_NIC_ENV_VARS = {
|
|
"NCCL_SOCKET_IFNAME",
|
|
"NCCL_IB_HCA",
|
|
"GLOO_SOCKET_IFNAME",
|
|
}
|
|
|
|
# These non-vLLM env vars are copied from the driver to workers.
|
|
# NCCL_* knobs present on the driver are added dynamically in
|
|
# ``_env_vars_to_copy_from_driver``, except the per-node NIC trio above.
|
|
ADDITIONAL_ENV_VARS = {
|
|
*HF_TOKEN_ENV_VARS,
|
|
*DEPRECATED_HF_TOKEN_ENV_VARS,
|
|
"NCCL_IB_DISABLE",
|
|
"NCCL_P2P_DISABLE",
|
|
"NCCL_CUMEM_ENABLE",
|
|
"NCCL_NVLS_ENABLE",
|
|
"NCCL_DEBUG",
|
|
"NCCL_DEBUG_SUBSYS",
|
|
"LD_LIBRARY_PATH",
|
|
# MiniMax-H3 encoder split: the driver resolves these into
|
|
# fastvideo_args, but keep the environment aligned on the workers so
|
|
# code paths that consult envs directly agree with the driver.
|
|
"FASTVIDEO_H3_ENCODER_SPLIT",
|
|
"FASTVIDEO_H3_ENCODER_NODES",
|
|
}
|
|
|
|
def _init_executor(self) -> None:
|
|
initialize_ray_cluster(self.fastvideo_args)
|
|
placement_group = self.fastvideo_args.ray_placement_group
|
|
|
|
# Disable Ray usage stats collection.
|
|
ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0")
|
|
if ray_usage != "1":
|
|
envs.set_external("RAY_USAGE_STATS_ENABLED", "0")
|
|
|
|
self._init_workers_ray(placement_group)
|
|
|
|
# child class could overwrite this to return actual env vars.
|
|
def _get_env_vars_to_be_updated(self) -> list[dict[str, str]]:
|
|
return self._env_vars_for_all_workers
|
|
|
|
def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs):
|
|
from fastvideo.platforms import current_platform
|
|
|
|
num_gpus = envs.FASTVIDEO_RAY_PER_WORKER_GPUS.get()
|
|
|
|
# On NVIDIA GPUs every actor keeps its raylet's CUDA_VISIBLE_DEVICES.
|
|
# Otherwise Ray narrows it to the actor's own GPU before the actor
|
|
# imports fastvideo, and that import initializes CUDA (diffusers and
|
|
# Triton query the device at import time). CUDA then keeps the one-GPU
|
|
# view, a later CUDA_VISIBLE_DEVICES update has no effect, and
|
|
# cuda:<local_rank> is an invalid ordinal for every worker but the first
|
|
# on a multi-GPU node. Each worker addresses its GPU by position in the
|
|
# inherited list instead (ray_worker_device_ordinal).
|
|
keep_raylet_devices = current_platform.is_cuda()
|
|
if keep_raylet_devices:
|
|
ray_remote_kwargs = keep_raylet_cuda_devices(ray_remote_kwargs)
|
|
|
|
# The remaining workers are the actual ray actors.
|
|
self.workers: list[RayWorkerWrapper] = []
|
|
|
|
# Create the workers.
|
|
# use the first N bundles that have GPU resources.
|
|
bundle_indices: list[int] = []
|
|
for bundle_id, bundle in enumerate(placement_group.bundle_specs):
|
|
if bundle.get(current_platform.ray_device_key, 0):
|
|
bundle_indices.append(bundle_id)
|
|
bundle_indices = bundle_indices[:self.fastvideo_args.num_gpus]
|
|
|
|
worker_metadata: list[RayWorkerMetaData] = []
|
|
driver_ip = get_ip()
|
|
for rank, bundle_id in enumerate(bundle_indices):
|
|
scheduling_strategy = PlacementGroupSchedulingStrategy(
|
|
placement_group=placement_group,
|
|
placement_group_capture_child_tasks=True,
|
|
placement_group_bundle_index=bundle_id,
|
|
)
|
|
|
|
if current_platform.ray_device_key == "GPU":
|
|
# NV+AMD GPUs, and Intel XPUs
|
|
worker = ray.remote(
|
|
num_cpus=0,
|
|
num_gpus=num_gpus,
|
|
scheduling_strategy=scheduling_strategy,
|
|
**ray_remote_kwargs,
|
|
)(RayWorkerWrapper).remote(fastvideo_args=self.fastvideo_args, rpc_rank=rank)
|
|
else:
|
|
worker = ray.remote(
|
|
num_cpus=0,
|
|
num_gpus=0,
|
|
resources={current_platform.ray_device_key: num_gpus},
|
|
scheduling_strategy=scheduling_strategy,
|
|
**ray_remote_kwargs,
|
|
)(RayWorkerWrapper).remote(fastvideo_args=self.fastvideo_args, rpc_rank=rank)
|
|
worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank))
|
|
|
|
worker_ips = ray.get([
|
|
each.worker.get_node_ip.remote() # type: ignore[attr-defined]
|
|
for each in worker_metadata
|
|
])
|
|
|
|
for each, ip in zip(worker_metadata, worker_ips, strict=False):
|
|
each.ip = ip
|
|
|
|
logger.debug("workers: %s", worker_metadata)
|
|
|
|
ip_counts: dict[str, int] = {}
|
|
for ip in worker_ips:
|
|
ip_counts[ip] = ip_counts.get(ip, 0) + 1
|
|
|
|
def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
|
|
"""
|
|
Sort the workers based on 3 properties:
|
|
1. If the worker is on the same node as the driver (vllm engine),
|
|
it should be placed first.
|
|
2. Then, if the worker is on a node with fewer workers, it should
|
|
be placed first.
|
|
3. Finally, if the work is on a node with smaller IP address, it
|
|
should be placed first.
|
|
"""
|
|
ip = item.ip
|
|
return (0 if ip == driver_ip else 1, ip_counts[ip], ip)
|
|
|
|
# After sorting, the workers on the same node will be
|
|
# close to each other, and the workers on the driver
|
|
# node will be placed first.
|
|
sorted_worker_metadata = sorted(worker_metadata, key=sort_by_driver_then_worker_ip)
|
|
start_rank = 0
|
|
for i, item in enumerate(sorted_worker_metadata):
|
|
item.adjusted_rank = i + start_rank
|
|
self.workers = [item.worker for item in sorted_worker_metadata]
|
|
rerank_mapping = {item.created_rank: item.adjusted_rank for item in sorted_worker_metadata}
|
|
self._run_ray_workers("adjust_rank", rerank_mapping)
|
|
|
|
if getattr(self.fastvideo_args, "h3_encoder_split", False):
|
|
# The sort keeps every node's workers contiguous, so the first
|
|
# ``h3_encoder_nodes`` distinct IPs are exactly the encoder group.
|
|
# Stamped before init_worker pickles the args to every worker.
|
|
self._setup_h3_encoder_split([item.ip for item in sorted_worker_metadata])
|
|
|
|
# Get the set of GPU IDs used on each node.
|
|
worker_node_and_gpu_ids = self._run_ray_workers("get_node_and_gpu_ids")
|
|
node_workers = defaultdict(list) # node id -> list of worker ranks
|
|
node_gpus = defaultdict(list) # node id -> list of gpu ids
|
|
|
|
for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
|
|
node_workers[node_id].append(i)
|
|
# `gpu_ids` can be a list of strings or integers.
|
|
# convert them to integers for consistency.
|
|
# NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs),
|
|
# string sorting is not sufficient.
|
|
# see https://github.com/vllm-project/vllm/issues/5590
|
|
gpu_ids = [int(x) for x in gpu_ids]
|
|
node_gpus[node_id].extend(gpu_ids)
|
|
for node_id, gpu_ids in node_gpus.items():
|
|
node_gpus[node_id] = sorted(gpu_ids)
|
|
|
|
all_ips = set(worker_ips + [driver_ip])
|
|
n_ips = len(all_ips)
|
|
n_nodes = len(node_workers)
|
|
|
|
if n_nodes != n_ips:
|
|
raise RuntimeError(f"Every node should have a unique IP address. Got {n_nodes}"
|
|
f" nodes with node ids {list(node_workers.keys())} and "
|
|
f"{n_ips} unique IP addresses {all_ips}. Please check your"
|
|
" network configuration. If you set `FASTVIDEO_HOST_IP`"
|
|
" environment variable, make sure it is unique for"
|
|
" each node.")
|
|
|
|
# Set environment variables for the driver and workers.
|
|
all_args_to_update_environment_variables: list[dict[str, str]]
|
|
if keep_raylet_devices:
|
|
# The inherited device list stays; each worker gets its ordinal in it.
|
|
worker_visible_devices = self._run_ray_workers("get_cuda_visible_devices")
|
|
all_args_to_update_environment_variables = [{} for _ in worker_node_and_gpu_ids]
|
|
else:
|
|
all_args_to_update_environment_variables = [{
|
|
current_platform.device_control_env_var:
|
|
",".join(map(str, node_gpus[node_id])),
|
|
} for (node_id, _) in worker_node_and_gpu_ids]
|
|
|
|
# Environment variables to copy from driver to workers
|
|
extra_nccl = {k for k in os.environ if k.startswith("NCCL_") and k not in self.WORKER_LOCAL_NIC_ENV_VARS}
|
|
env_vars_to_copy = get_env_vars_to_copy(
|
|
exclude_vars=self.WORKER_SPECIFIC_ENV_VARS | self.WORKER_LOCAL_NIC_ENV_VARS,
|
|
additional_vars=set(current_platform.additional_env_vars).union(self.ADDITIONAL_ENV_VARS).union(extra_nccl),
|
|
destination="workers",
|
|
)
|
|
|
|
# Copy existing env vars to each worker's args
|
|
for args in all_args_to_update_environment_variables:
|
|
# TODO: refactor platform-specific env vars
|
|
for name in env_vars_to_copy:
|
|
if name in os.environ:
|
|
args[name] = os.environ[name]
|
|
|
|
self._env_vars_for_all_workers: list[dict[str, str]] = (all_args_to_update_environment_variables)
|
|
|
|
self._run_ray_workers("update_environment_variables", self._get_env_vars_to_be_updated())
|
|
|
|
_lb = should_use_gloo_loopback(worker_ips)
|
|
if _lb:
|
|
driver_ip = get_loopback_ip()
|
|
distributed_init_method = get_distributed_init_method(driver_ip, get_open_port())
|
|
|
|
# Initialize the actual workers inside worker wrapper.
|
|
all_kwargs = []
|
|
for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids):
|
|
local_rank = node_workers[node_id].index(rank)
|
|
if keep_raylet_devices:
|
|
local_rank = ray_worker_device_ordinal(node_gpus[node_id], local_rank, worker_visible_devices[rank])
|
|
kwargs = dict(
|
|
fastvideo_args=self.fastvideo_args,
|
|
local_rank=local_rank,
|
|
rank=rank,
|
|
distributed_init_method=distributed_init_method,
|
|
)
|
|
all_kwargs.append(kwargs)
|
|
self._run_ray_workers("init_worker", all_kwargs)
|
|
self._run_ray_workers("init_device")
|
|
|
|
# This is the list of workers that are rank 0 of each TP group EXCEPT
|
|
# global rank 0. These are the workers that will broadcast to the
|
|
# rest of the workers.
|
|
self.tp_driver_workers: list[RayWorkerWrapper] = []
|
|
# This is the list of workers that are not drivers and not the first
|
|
# worker in a TP group. These are the workers that will be
|
|
# broadcasted to.
|
|
self.non_driver_workers: list[RayWorkerWrapper] = []
|
|
|
|
# Enforce rank order for correct rank to return final output.
|
|
for index, worker in enumerate(self.workers):
|
|
# The driver worker is rank 0 and not in self.workers.
|
|
rank = index + 1
|
|
if rank % self.fastvideo_args.tp_size == 0:
|
|
self.tp_driver_workers.append(worker)
|
|
else:
|
|
self.non_driver_workers.append(worker)
|
|
|
|
def _setup_h3_encoder_split(self, sorted_worker_ips: list[str]) -> None:
|
|
"""Reserve the leading encoder nodes' workers and stamp the count onto the args.
|
|
|
|
``sorted_worker_ips`` follows the driver-then-node sort, so workers of one
|
|
node are contiguous; the first ``h3_encoder_nodes`` distinct IPs form the
|
|
encoder group (ranks ``0..n-1``). Workers must be one-per-node style
|
|
placements; multi-GPU nodes simply contribute all of their workers.
|
|
"""
|
|
args = self.fastvideo_args
|
|
# ``_fold_h3_encoder_split_env`` already normalized 0/unset to 1.
|
|
nodes_needed = int(args.h3_encoder_nodes)
|
|
ordered_ips: list[str] = []
|
|
for ip in sorted_worker_ips:
|
|
if ip not in ordered_ips:
|
|
ordered_ips.append(ip)
|
|
if len(ordered_ips) < nodes_needed + 1:
|
|
raise RuntimeError(f"MiniMax-H3 encoder split needs at least {nodes_needed + 1} distinct nodes "
|
|
f"({nodes_needed} encoder + 1 denoise), got {len(ordered_ips)}.")
|
|
encoder_ips = set(ordered_ips[:nodes_needed])
|
|
encoder_workers = sum(1 for ip in sorted_worker_ips if ip in encoder_ips)
|
|
if not 1 <= encoder_workers < len(sorted_worker_ips):
|
|
raise RuntimeError(f"MiniMax-H3 encoder split resolved {encoder_workers} encoder workers out of "
|
|
f"{len(sorted_worker_ips)}; need at least one denoise worker.")
|
|
from fastvideo.fastvideo_args import h3_split_sp_error, probe_h3_attention_heads
|
|
|
|
denoise_workers = len(sorted_worker_ips) - encoder_workers
|
|
heads = probe_h3_attention_heads(args.model_path)
|
|
if heads and heads % denoise_workers:
|
|
raise RuntimeError(h3_split_sp_error(heads, denoise_workers, len(sorted_worker_ips), unit="workers"))
|
|
args.h3_encoder_nodes = nodes_needed
|
|
args.h3_encoder_workers = encoder_workers
|
|
logger.info(
|
|
"MiniMax-H3 encoder split: ranks 0..%d on nodes %s run the Qwen3-VL conditioner only; ranks %d..%d "
|
|
"run DiT/VAE with sp_size=%d", encoder_workers - 1, sorted(encoder_ips), encoder_workers,
|
|
len(sorted_worker_ips) - 1,
|
|
len(sorted_worker_ips) - encoder_workers)
|
|
|
|
def execute_streaming_reset(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> dict[str, Any]:
|
|
responses: list[dict[str, Any]] = self.collective_rpc(
|
|
"execute_streaming_reset",
|
|
kwargs={
|
|
"forward_batch": forward_batch,
|
|
"fastvideo_args": fastvideo_args,
|
|
},
|
|
)
|
|
return responses[0]
|
|
|
|
def execute_streaming_step(self, keyboard_action=None, mouse_action=None) -> ForwardBatch:
|
|
responses: list[ForwardBatch] = self.collective_rpc(
|
|
"execute_streaming_step",
|
|
kwargs={
|
|
"keyboard_action": keyboard_action,
|
|
"mouse_action": mouse_action,
|
|
},
|
|
)
|
|
return responses[0]
|
|
|
|
async def execute_streaming_step_async(self, keyboard_action=None, mouse_action=None) -> ForwardBatch:
|
|
kwargs = {
|
|
"keyboard_action": keyboard_action,
|
|
"mouse_action": mouse_action,
|
|
}
|
|
futures = [w.execute_method.remote("execute_streaming_step", **kwargs) for w in self.workers]
|
|
responses = await asyncio.gather(*futures)
|
|
return responses[0]
|
|
|
|
def execute_streaming_clear(self) -> None:
|
|
self.collective_rpc("execute_streaming_clear")
|
|
|
|
def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
|
responses: list[ForwardBatch] = self.collective_rpc(
|
|
"execute_forward",
|
|
kwargs={
|
|
"forward_batch": forward_batch,
|
|
"fastvideo_args": fastvideo_args,
|
|
},
|
|
)
|
|
# Under the MiniMax-H3 encoder split the encoder ranks answer with a
|
|
# stub batch; the decoded output lives on the first denoise rank.
|
|
src = 0
|
|
# The executor's own args carry h3_encoder_workers (stamped when the
|
|
# workers were placed); fall back to the per-call args when absent.
|
|
split_args = getattr(self, "fastvideo_args", fastvideo_args)
|
|
if getattr(split_args, "h3_encoder_split", False):
|
|
src = int(split_args.h3_encoder_workers)
|
|
output = responses[src].output.cpu()
|
|
|
|
logging_info = None
|
|
if envs.FASTVIDEO_STAGE_LOGGING.get():
|
|
logging_info = responses[src].logging_info
|
|
|
|
result_batch = ForwardBatch(
|
|
data_type=forward_batch.data_type,
|
|
output=output,
|
|
logging_info=logging_info,
|
|
extra=responses[src].extra,
|
|
trajectory_latents=responses[src].trajectory_latents,
|
|
trajectory_timesteps=responses[src].trajectory_timesteps,
|
|
)
|
|
return result_batch
|
|
|
|
def set_lora_adapter(self,
|
|
lora_nickname: str,
|
|
lora_path: str | None = None,
|
|
strength: float = 1.0,
|
|
accumulate: bool = False) -> None:
|
|
responses = self.collective_rpc("set_lora_adapter",
|
|
kwargs={
|
|
"lora_nickname": lora_nickname,
|
|
"lora_path": lora_path,
|
|
"strength": strength,
|
|
"accumulate": accumulate
|
|
})
|
|
for i, response in enumerate(responses):
|
|
if response["status"] != "lora_adapter_set":
|
|
raise RuntimeError(f"Worker {i} failed to set LoRA adapter to {lora_path}")
|
|
|
|
def unmerge_lora_weights(self) -> None:
|
|
responses = self.collective_rpc("unmerge_lora_weights", kwargs={})
|
|
for i, response in enumerate(responses):
|
|
if response["status"] != "lora_adapter_unmerged":
|
|
raise RuntimeError(f"Worker {i} failed to unmerge LoRA weights")
|
|
|
|
def merge_lora_weights(self) -> None:
|
|
responses = self.collective_rpc("merge_lora_weights", kwargs={})
|
|
for i, response in enumerate(responses):
|
|
if response["status"] != "lora_adapter_merged":
|
|
raise RuntimeError(f"Worker {i} failed to merge LoRA weights")
|
|
|
|
def set_log_queue(self, log_queue: Queue | None) -> None:
|
|
"""Keep the driver-side queue locally.
|
|
|
|
``multiprocessing.Queue`` is not picklable across Ray nodes, so worker
|
|
logs stay in the Ray session log dir instead of being forwarded.
|
|
"""
|
|
self._log_queue = log_queue
|
|
|
|
def clear_log_queue(self) -> None:
|
|
self._log_queue = None
|
|
|
|
def collective_rpc(self,
|
|
method: str | Callable,
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict | None = None) -> list[Any]:
|
|
if timeout is not None:
|
|
raise NotImplementedError("RayDistributedExecutor.collective_rpc does not support per-call timeouts; "
|
|
"omit the timeout and rely on Ray's own task failure handling instead.")
|
|
return self._run_ray_workers(method, *args, **(kwargs or {}))
|
|
|
|
def _run_ray_workers(
|
|
self,
|
|
method: str | Callable,
|
|
*args,
|
|
**kwargs,
|
|
) -> Any:
|
|
sent_method = method if isinstance(method, str) else cloudpickle.dumps(method)
|
|
del method
|
|
|
|
# Start the ray workers first.
|
|
ray_workers = self.workers
|
|
ray_worker_outputs = [worker.execute_method.remote(sent_method, *args, **kwargs) for worker in ray_workers]
|
|
|
|
# Get the results of the ray workers.
|
|
ray_worker_outputs = ray.get(ray_worker_outputs)
|
|
|
|
return ray_worker_outputs
|
|
|
|
def shutdown(self) -> None:
|
|
logger.info("Shutting down Ray distributed executor. If you see error log "
|
|
"from logging.cc regarding SIGTERM received, please ignore because "
|
|
"this is the expected termination process in Ray.")
|
|
import ray
|
|
for worker in self.workers:
|
|
ray.kill(worker)
|
|
|
|
self.workers = []
|
|
|
|
def __del__(self):
|
|
self.shutdown()
|