diff --git a/examples/v1/config/reasoning_rl_qwen3p5vl_mtp_ep.py b/examples/v1/config/reasoning_rl_qwen3p5vl_mtp_ep.py index 92d038ace..f4ebffbd2 100644 --- a/examples/v1/config/reasoning_rl_qwen3p5vl_mtp_ep.py +++ b/examples/v1/config/reasoning_rl_qwen3p5vl_mtp_ep.py @@ -30,6 +30,7 @@ import ray # export LMDEPLOY_FP32_MAMBA_SSM_DTYPE=1 +# export XTUNER_ACTIVATION_OFFLOAD=1 def _as_list(value): return value if isinstance(value, list) else [value] @@ -61,6 +62,7 @@ def _as_list(value): checkpoint_interval = 50 evaluate_step = 5 enable_initial_evaluate = os.environ.get("ENABLE_INITIAL_EVALUATE", False) +train_ep_size = 4 # 1. resources resources = AcceleratorResourcesConfig( @@ -89,7 +91,7 @@ def _as_list(value): lmdeploy_uvicorn_log_level="INFO", lmdeploy_speculative_algorithm='qwen3_5_mtp', # MTP draft tokens trade throughput for extra activation memory; try 3 if still tight. - lmdeploy_speculative_num_draft_tokens=4, + lmdeploy_speculative_num_draft_tokens=3, ), health_check_interval_seconds=300, health_check_failure_threshold=3, @@ -207,15 +209,21 @@ def _as_list(value): }, ) +from xtuner.v1.float8 import Float8Config, ScalingGranularity # 5. train worker model_cfg = Qwen3_5_VLMoE35BA3Config(freeze_vision=True, freeze_projector=True) -model_cfg.float8_cfg = None -model_cfg.text_config.ep_size = 1 +float8_cfg = Float8Config( + scaling_granularity_gemm=None, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, +) +model_cfg.float8_cfg = float8_cfg +model_cfg.text_config.ep_size = train_ep_size model_cfg.text_config.z_loss_cfg = None model_cfg.text_config.balancing_loss_cfg = None model_cfg.text_config.freeze_routers = True +model_cfg.compile_cfg = None model_cfg.text_config.mtp_config = MTPConfig( - num_layers=4, + num_layers=3, loss_scaling_factor=1.0, detach_mtp_lm_head_weight=True, detach_mtp_inputs=True, @@ -246,7 +254,7 @@ def _as_list(value): ), ) lr_cfg = LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6) -fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=1, fp32_lm_head=True) +fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=train_ep_size, fp32_lm_head=True) train_worker_cfg = WorkerConfig( model_cfg=model_cfg, load_from=model_path, diff --git a/setup.cfg b/setup.cfg index 73bb498ec..94d343f0e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [codespell] -ignore-words-list = nd, ba, warmup, ans, dout +ignore-words-list = nd, ba, warmup, ans, dout, ptd [flake8] max-line-length = 119 diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index b87da88ce..ea339b1cd 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -1,6 +1,6 @@ import math import os -from typing import Literal, TypedDict +from typing import Any, Literal, TypedDict import ray import torch @@ -24,6 +24,23 @@ class ColateItem(TypedDict): rollout_logprobs: torch.Tensor | None +def _summarize_process_group_results(results: list[dict[str, Any]]) -> str: + if not results: + return "ranks=0" + + count_key = "destroyed" if "destroyed" in results[0] else "reloaded" + counts = [result.get(count_key, 0) for result in results] + count_summary = f"{counts[0]} on all ranks" if len(set(counts)) == 1 else f"by_rank={counts}" + result_errors = [error for result in results for error in result.get("errors", [])] + unwrapped = sum(len(result.get("unwrapped_nccl_groups", [])) for result in results) + summary = f"ranks={len(results)}, {count_key}={count_summary}" + if unwrapped: + summary += f", unwrapped_nccl_groups={unwrapped}" + if result_errors: + summary += f", errors={len(result_errors)}, first_error={result_errors[0]}" + return summary + + class TrainingController: def __init__(self, workers: list[TrainingWorker]) -> None: self.workers = workers @@ -246,8 +263,6 @@ def fit(self, data_batches: list[ColateItem], pack_max_length: int, rollout_idx: pad_data_samples = [pad_data for _ in range(pad_num)] packed_data_batches = packed_data_batches + pad_data_samples - print(f"len(packed_data_batches): {len(packed_data_batches)}") - handles = [] data_batch_refs = {} for worker_idx, worker in enumerate(self.workers): @@ -323,6 +338,26 @@ def update_weights(self): ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) return + def destroy_train_nccl_process_groups(self): + """Destroy train-side NCCL process groups after weight sync.""" + handles = [ + worker.destroy_train_nccl_process_groups.remote() # type: ignore[attr-defined] + for worker in self.workers + ] + results = ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) + self.logger.info(f"Destroyed train NCCL process groups: {_summarize_process_group_results(results)}") + return results + + def reload_train_nccl_process_groups(self): + """Reload train-side reloadable NCCL process groups before training.""" + handles = [ + worker.reload_train_nccl_process_groups.remote() # type: ignore[attr-defined] + for worker in self.workers + ] + results = ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) + self.logger.info(f"Reloaded train NCCL process groups: {_summarize_process_group_results(results)}") + return results + def save_hf(self, hf_dir: str, save_dtype: torch.dtype = torch.bfloat16): handles = [worker.save_hf.remote(hf_dir, save_dtype) for worker in self.workers] # type: ignore ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index f7e0bc487..9f4d96a2a 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -57,6 +57,14 @@ ray_method, set_deterministic, ) +from xtuner.v1.utils.activation_offload import OffloadManager +from xtuner.v1.utils.fsdp import release_deferred_fsdp_all_gathers +from xtuner.v1.utils.reloadable_process_group import ( + destroy_reloadable_process_groups, + monkey_patch_reloadable_process_groups, + process_group_name, + reload_process_groups, +) from ..rollout_is import merge_rollout_is_metrics @@ -221,6 +229,7 @@ def __init__( self.config = cast(WorkerConfig, self.config) torch.accelerator.set_device_index(int(os.environ["LOCAL_RANK"])) self.rank = rank + monkey_patch_reloadable_process_groups() # TODO: add lr scheduler log_dir = worker_cfg.log_dir @@ -278,6 +287,74 @@ def __init__( engine=self._engine, ) + @ray_method + def destroy_train_nccl_process_groups(self) -> dict[str, object]: + """Destroy reloadable train-side NCCL PG inners after weight sync and + train offload.""" + + details = destroy_reloadable_process_groups() + errors: list[str] = [] + unwrapped_nccl_groups: list[dict[str, object]] = [] + + try: + import torch.distributed.distributed_c10d as c10d + except Exception as exc: + errors.append(f"import distributed_c10d {type(exc).__name__}: {exc}") + c10d = None # type: ignore + + groups: list[tuple[dist.ProcessGroup, list[int], str]] = [] + default_group = dist.group.WORLD if dist.is_initialized() else None + if c10d is not None: + for group, rank_map in list(c10d._world.pg_group_ranks.items()): + if group is default_group: + continue + try: + backend = str(dist.get_backend(group)).lower() + except ValueError: + # The c10d registry may still contain entries for groups + # just destroyed by the reloadable wrapper. They are not + # actionable unwrapped NCCL groups. + continue + except Exception as exc: + errors.append(f"get_backend {type(exc).__name__}: {exc}") + continue + if "nccl" not in backend: + continue + ranks = [global_rank for global_rank, _ in sorted(rank_map.items(), key=lambda item: item[1])] + groups.append((group, ranks, backend)) + + for group, ranks, backend in groups: + if group.__class__.__name__ == "ReloadableProcessGroup": + continue + unwrapped_nccl_groups.append({"ranks": ranks, "backend": backend, "group_name": process_group_name(group)}) + + DEVICE_MODULE.empty_cache() + result = { + "rank": self.rank, + "destroyed": len(details), + "unwrapped_nccl_groups": unwrapped_nccl_groups, + "errors": errors, + } + return result + + @ray_method + def reload_train_nccl_process_groups(self) -> dict[str, object]: + """Reload train-side NCCL PG inners before FSDP2 runs again.""" + + errors: list[str] = [] + try: + details = reload_process_groups() + except Exception as exc: + details = [] + errors.append(f"reload {type(exc).__name__}: {exc}") + DEVICE_MODULE.empty_cache() + result = { + "rank": self.rank, + "reloaded": len(details), + "errors": errors, + } + return result + @ray_method def bind_rollout_weight_update(self, *args, **kwargs): return self.update_weighter.bind_rollout_weight_update(*args, **kwargs) @@ -679,6 +756,13 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo only_calc_mismatch_ratio = os.environ.get("ONLY_CALC_MISMATCH_RATIO", "0") == "1" if only_calc_mismatch_ratio: + # This early-return path skips the MTP forward. FSDP may have + # prefetched the MTP module already, leaving its all-gather result + # cached until we release it explicitly. The normal training branch + # can leave the same residue, but it is outside the peak-memory + # path and can be ignored. + self._release_deferred_fsdp_all_gathers(f"rollout_{rollout_idx}/only_calc_mismatch") + DEVICE_MODULE.empty_cache() return worker_log_item # compute reference logprobs @@ -772,6 +856,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo train_step_info = self._engine.train_step( data_batches=engine_input, ) + OffloadManager().clear() self.logger.debug( f"Rank{self.rank} Rollout {rollout_idx} GlobalStep {global_train_step} " f"train_step[{i}].engine_train_step elapsed={time.perf_counter() - train_step_begin:.4f}s" @@ -954,10 +1039,30 @@ def get_model_cfg(self): model_cfg = self._engine.model_cfg return model_cfg + def _clear_cublas_workspaces(self) -> None: + clear_fn = getattr(torch._C, "_cuda_clearCublasWorkspaces", None) + if clear_fn is None: + self.logger.warning("torch._C._cuda_clearCublasWorkspaces is unavailable") + return + try: + clear_fn() + DEVICE_MODULE.empty_cache() + except Exception as e: + self.logger.warning(f"Failed to clear cuBLAS workspaces: {e}") + + def _release_deferred_fsdp_all_gathers(self, log_tag: str) -> None: + comm_states, param_groups = release_deferred_fsdp_all_gathers(self._engine.model) + if comm_states or param_groups: + self.logger.debug( + f"[{log_tag}] released deferred FSDP all-gathers: " + f"comm_states={comm_states}, param_groups={param_groups}" + ) + @ray_method def offload_model(self): self._engine.put_model_to_device("cpu") DEVICE_MODULE.empty_cache() + self._clear_cublas_workspaces() self.logger.info( f"Offloaded model to CPU. Current allocate {DEVICE_MODULE.memory_allocated() / (1024**2)} MB, reserved: {DEVICE_MODULE.memory_reserved() / (1024**2)} MB" ) diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index ad9f4689c..70f57ed3d 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -913,6 +913,12 @@ def _train_one_batch( ) ray.get(self.rollout_controller.offload.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) if onload_train_before_train: + if getattr(self, "_train_nccl_destroyed", False): + # FSDP2 keeps DeviceMesh/PG references alive. Recreate the + # wrapped NCCL inners before loading train weights again. + with timer("reload_train_nccl", step_timer_dict): + self.train_controller.reload_train_nccl_process_groups() + self._train_nccl_destroyed = False with timer("onload", step_timer_dict): self.train_controller.onload(target="all") self.logger.info("Training controller loaded") @@ -1772,6 +1778,15 @@ def _sync_weights_and_save(self, train_step: int, step_timer_dict: dict) -> bool self.train_controller.update_weights() self.logger.info("Rollout workers update weights successfully in colocate mode") self.train_controller.offload(target="model") + destroy_train_nccl = os.getenv("XTUNER_DESTROY_TRAIN_NCCL_AFTER_SYNC", "0") == "1" + already_destroyed = getattr(self, "_train_nccl_destroyed", False) + if destroy_train_nccl and not already_destroyed: + # This runs only after rollout weights have been updated + # and the train model has been offloaded. Only replayable + # non-default NCCL groups are destroyed. + with timer("destroy_train_nccl", step_timer_dict): + self.train_controller.destroy_train_nccl_process_groups() + self._train_nccl_destroyed = True else: self.train_controller.offload(target="model") ray.get(self.rollout_controller.onload_weights.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) diff --git a/xtuner/v1/utils/activation_offload.py b/xtuner/v1/utils/activation_offload.py index ddff7f9ea..eac495b12 100644 --- a/xtuner/v1/utils/activation_offload.py +++ b/xtuner/v1/utils/activation_offload.py @@ -287,12 +287,42 @@ def prefetch_get(self, block_idx, tensor_idx, h2d_stream, d2h_stream, group="def def empty(self): return len(self.items) == 0 - def clear(self, key=None): - if key is None: - self.items.clear() - else: - self.assert_exist(key) - self.items.pop(key) + def clear(self, key=None, group=None, clear_pin_memory_cache: bool = False): + """Clear runtime offload state. + + ``pin_memory_cache`` is kept by default because it is a reusable CPU + pinned-memory cache, not a live GPU activation reference. + """ + if key is not None and group is not None: + raise ValueError("Only one of key or group can be specified") + + if key is not None: + self.items.pop(key, None) + self.may_npu_tensors.pop(key, None) + if clear_pin_memory_cache: + self.pin_memory_cache.pop(key, None) + return + + if group is not None: + prefix = f"{group}_" + for item_key in list(self.items.keys()): + if item_key.startswith(prefix): + self.items.pop(item_key, None) + for item_key in list(self.may_npu_tensors.keys()): + if item_key.startswith(prefix): + self.may_npu_tensors.pop(item_key, None) + self.getcnt.pop(group, None) + if clear_pin_memory_cache: + for item_key in list(self.pin_memory_cache.keys()): + if item_key.startswith(prefix): + self.pin_memory_cache.pop(item_key, None) + return + + self.items.clear() + self.may_npu_tensors.clear() + self.getcnt.clear() + if clear_pin_memory_cache: + self.pin_memory_cache.clear() # event interface # diff --git a/xtuner/v1/utils/fsdp.py b/xtuner/v1/utils/fsdp.py new file mode 100644 index 000000000..8f1096b71 --- /dev/null +++ b/xtuner/v1/utils/fsdp.py @@ -0,0 +1,70 @@ +import torch + +from xtuner.v1.utils.device import get_torch_device_module + + +DEVICE_MODULE = get_torch_device_module() + + +def release_deferred_fsdp_all_gathers(model: torch.nn.Module) -> tuple[int, int]: + """Release FSDP2 all-gather buffers that were prefetched but not + consumed.""" + + try: + from torch.distributed._composable_state import _get_module_state + except Exception: + return 0, 0 + + released_comm_states = 0 + released_param_groups = 0 + seen_comm_ctx: set[int] = set() + seen_param_groups: set[int] = set() + + def wait_all_gather_result(all_gather_result) -> None: + event = getattr(all_gather_result, "all_gather_event", None) + if event is not None: + try: + torch.accelerator.current_stream().wait_event(event) + except Exception: + DEVICE_MODULE.synchronize() + work = getattr(all_gather_result, "all_gather_work", None) + if work is not None and hasattr(work, "wait"): + try: + work.wait() + except Exception: + DEVICE_MODULE.synchronize() + + for module in model.modules(): + try: + state = _get_module_state(module) + except Exception: + continue + if state is None: + continue + + comm_ctx = getattr(state, "_comm_ctx", None) + if comm_ctx is not None and id(comm_ctx) not in seen_comm_ctx: + seen_comm_ctx.add(id(comm_ctx)) + all_gather_state = getattr(comm_ctx, "all_gather_state", None) + if all_gather_state is not None: + event = getattr(all_gather_state, "event", None) + if event is not None: + try: + event.synchronize() + except Exception: + DEVICE_MODULE.synchronize() + comm_ctx.all_gather_state = None + released_comm_states += 1 + + param_group = getattr(state, "_fsdp_param_group", None) + if param_group is None or id(param_group) in seen_param_groups: + continue + seen_param_groups.add(id(param_group)) + all_gather_result = getattr(param_group, "_all_gather_result", None) + if all_gather_result is None: + continue + wait_all_gather_result(all_gather_result) + param_group._all_gather_result = None + released_param_groups += 1 + + return released_comm_states, released_param_groups diff --git a/xtuner/v1/utils/reloadable_process_group.py b/xtuner/v1/utils/reloadable_process_group.py new file mode 100644 index 000000000..8a10098bb --- /dev/null +++ b/xtuner/v1/utils/reloadable_process_group.py @@ -0,0 +1,491 @@ +"""Reload selected train-side NCCL process groups without rebuilding FSDP2. + +FSDP2 and DeviceMesh keep Python references to process groups. Destroying those groups directly leaves stale objects in +FSDP hooks and functional collective registries, so this module gives PyTorch a stable wrapper object and swaps only +the wrapped NCCL process group when train memory is released/reloaded. + +The default/world process group is intentionally not managed here. It is used by Ray/torch distributed control paths +and is not replayed safely from this layer. +""" + +import os +from collections.abc import Callable +from typing import Any, cast + +import torch +import torch.distributed as dist + + +class ReloadableProcessGroup(torch.distributed.ProcessGroup): + """Stable ProcessGroup proxy whose inner NCCL group can be destroyed.""" + + _GROUPS_BY_PID: dict[int, list["ReloadableProcessGroup"]] = {} + _GROUPS_BY_NAME: dict[int, dict[str, "ReloadableProcessGroup"]] = {} + + def __init__(self, group: dist.ProcessGroup, ranks: list[int], backend: str | None = None): + super().__init__(rank=group.rank(), size=group.size()) # type: ignore[call-arg] + self.group: dist.ProcessGroup | None = group + self._stable_group_name = group.group_name + self._stable_group_desc = getattr(group, "group_desc", "") + self.group_info = { + "ranks": list(ranks), + "backend": backend or _backend_for_reload(group), + } + # DeviceMesh reads these attributes directly instead of going through + # methods, so expose stable Python properties on the wrapper itself. + try: + self._set_group_name(self._stable_group_name) + self._set_group_desc(self._stable_group_desc) + except RuntimeError: + pass + self._copy_public_pg_attrs(group) + + pid = os.getpid() + self._GROUPS_BY_PID.setdefault(pid, []).append(self) + self._GROUPS_BY_NAME.setdefault(pid, {})[self._stable_group_name] = self + _register_wrapper_in_world(self, group) + + @property + def group_name(self) -> str: + return self._stable_group_name + + @property + def group_desc(self) -> str: + return self._stable_group_desc + + @property + def _device_types(self): + # all_gather_object uses this private field to select a CPU fallback. + # Proxying it avoids "No backend type associated with device type cpu". + return self._inner()._device_types + + def _copy_public_pg_attrs(self, group: dist.ProcessGroup) -> None: + if hasattr(group, "bound_device_id"): + self.bound_device_id = group.bound_device_id # type: ignore[attr-defined] + + def __getattr__(self, name: str) -> Any: + group = self._inner() + return getattr(group, name) + + def _inner(self) -> dist.ProcessGroup: + if self.group is None: + raise RuntimeError("ReloadableProcessGroup inner process group is destroyed; reload it before use.") + return self.group + + def _fwd(self, method: str, *args, **kwargs): + return getattr(self._inner(), method)(*args, **kwargs) + + def rank(self) -> int: + return self._inner().rank() + + def size(self) -> int: + return self._inner().size() + + def name(self) -> str: + return self._inner().name() + + def get_group_store(self): + return self._inner().get_group_store() + + def _get_backend(self, *args, **kwargs): + return self._inner()._get_backend(*args, **kwargs) + + def _get_backend_name(self, *args, **kwargs): + return self._inner()._get_backend_name(*args, **kwargs) + + def _get_sequence_number_for_group(self, *args, **kwargs): + return self._inner()._get_sequence_number_for_group(*args, **kwargs) + + def _set_sequence_number_for_group(self, *args, **kwargs): + return self._inner()._set_sequence_number_for_group(*args, **kwargs) + + def shutdown(self) -> None: + if self.group is not None: + getattr(self.group, "shutdown")() + + def abort(self) -> None: + if self.group is not None: + getattr(self.group, "abort")() + + def split_group(self, *args, **kwargs): + return self._inner().split_group(*args, **kwargs) + + def barrier(self, *args, **kwargs): + return self._fwd("barrier", *args, **kwargs) + + def broadcast(self, *args, **kwargs): + return self._fwd("broadcast", *args, **kwargs) + + def allreduce(self, *args, **kwargs): + return self._fwd("allreduce", *args, **kwargs) + + def allreduce_coalesced(self, *args, **kwargs): + return self._fwd("allreduce_coalesced", *args, **kwargs) + + def reduce(self, *args, **kwargs): + return self._fwd("reduce", *args, **kwargs) + + def allgather(self, *args, **kwargs): + return self._fwd("allgather", *args, **kwargs) + + def allgather_coalesced(self, *args, **kwargs): + return self._fwd("allgather_coalesced", *args, **kwargs) + + def allgather_into_tensor_coalesced(self, *args, **kwargs): + return self._fwd("allgather_into_tensor_coalesced", *args, **kwargs) + + def _allgather_base(self, *args, **kwargs): + return self._fwd("_allgather_base", *args, **kwargs) + + def gather(self, *args, **kwargs): + return self._fwd("gather", *args, **kwargs) + + def scatter(self, *args, **kwargs): + return self._fwd("scatter", *args, **kwargs) + + def reduce_scatter(self, *args, **kwargs): + return self._fwd("reduce_scatter", *args, **kwargs) + + def reduce_scatter_tensor_coalesced(self, *args, **kwargs): + return self._fwd("reduce_scatter_tensor_coalesced", *args, **kwargs) + + def _reduce_scatter_base(self, *args, **kwargs): + return self._fwd("_reduce_scatter_base", *args, **kwargs) + + def alltoall(self, *args, **kwargs): + return self._fwd("alltoall", *args, **kwargs) + + def alltoall_base(self, *args, **kwargs): + return self._fwd("alltoall_base", *args, **kwargs) + + def send(self, *args, **kwargs): + return self._fwd("send", *args, **kwargs) + + def recv(self, *args, **kwargs): + return self._fwd("recv", *args, **kwargs) + + def recv_anysource(self, *args, **kwargs): + return self._fwd("recv_anysource", *args, **kwargs) + + def monitored_barrier(self, *args, **kwargs): + return self._fwd("monitored_barrier", *args, **kwargs) + + +_ORIGINALS_BY_PID: dict[int, dict[str, Any]] = {} +_SPECS_BY_PID: dict[int, list[dict[str, Any]]] = {} + + +def monkey_patch_reloadable_process_groups() -> None: + """Install per-process patches before DeviceMesh/FSDP2 creates + subgroups.""" + + pid = os.getpid() + if pid in _ORIGINALS_BY_PID: + return + + import torch.distributed.device_mesh as device_mesh + import torch.distributed.distributed_c10d as c10d + + originals = { + "c10d_new_group": c10d.new_group, + "c10d_split_group": c10d.split_group, + "c10d_resolve_process_group": c10d._resolve_process_group, + } + if hasattr(dist, "split_group"): + originals["dist_split_group"] = dist.split_group + _ORIGINALS_BY_PID[pid] = originals + + def new_group(*args, **kwargs): + group = originals["c10d_new_group"](*args, **kwargs) + wrapper_or_group = _wrap_group_if_needed(group, _extract_new_group_ranks(args, kwargs), kwargs.get("backend")) + if _should_record_spec(group, args, kwargs): + # Every rank must replay every subgroup creation call in the same + # order, including nonmember calls that return NON_GROUP_MEMBER. + _SPECS_BY_PID.setdefault(os.getpid(), []).append( + { + "kind": "new_group", + "args": args, + "kwargs": dict(kwargs), + "wrapper": wrapper_or_group if isinstance(wrapper_or_group, ReloadableProcessGroup) else None, + } + ) + return wrapper_or_group + + def split_group(*args, **kwargs): + group = originals["c10d_split_group"](*args, **kwargs) + wrapper_or_group = _wrap_group_if_needed(group, None, None) + _SPECS_BY_PID.setdefault(os.getpid(), []).append( + { + "kind": "split_group", + "args": args, + "kwargs": dict(kwargs), + "wrapper": wrapper_or_group if isinstance(wrapper_or_group, ReloadableProcessGroup) else None, + } + ) + return wrapper_or_group + + def resolve_process_group(group_name: str): + # Functional collectives store only the group name. Keep that name + # resolving to the stable wrapper even after the inner PG is recreated. + wrapper = ReloadableProcessGroup._GROUPS_BY_NAME.get(os.getpid(), {}).get(group_name) + if wrapper is not None: + return wrapper + resolve_process_group_fn = cast(Callable[[str], Any], originals["c10d_resolve_process_group"]) + return resolve_process_group_fn(group_name) + + dist.new_group = new_group + c10d.new_group = new_group + device_mesh.new_group = new_group + c10d.split_group = split_group + device_mesh.split_group = split_group + c10d._resolve_process_group = resolve_process_group + device_mesh._resolve_process_group = resolve_process_group # type: ignore[attr-defined] + if "dist_split_group" in originals: + dist.split_group = split_group # type: ignore[attr-defined] + + +def destroy_reloadable_process_groups() -> list[dict[str, object]]: + """Destroy wrapped NCCL groups while leaving the wrapper objects cached.""" + + details: list[dict[str, object]] = [] + for wrapper in reversed(ReloadableProcessGroup._GROUPS_BY_PID.get(os.getpid(), [])): + if wrapper.group is None: + continue + inner = wrapper.group + ranks = list(wrapper.group_info["ranks"]) + backend = str(wrapper.group_info["backend"]) + try: + dist.destroy_process_group(inner) + wrapper.group = None + details.append({"ranks": ranks, "backend": backend, "group_name": wrapper.group_name}) + except ValueError: + wrapper.group = None + details.append( + {"ranks": ranks, "backend": backend, "group_name": wrapper.group_name, "already_gone": True} + ) + return details + + +def reload_process_groups() -> list[dict[str, object]]: + """Replay recorded subgroup creation and attach fresh inners to + wrappers.""" + + pid = os.getpid() + originals = _ORIGINALS_BY_PID.get(pid) + if originals is None: + return [] + if not any(group.group is None for group in ReloadableProcessGroup._GROUPS_BY_PID.get(pid, [])): + return [] + + details: list[dict[str, object]] = [] + for spec in _SPECS_BY_PID.get(pid, []): + wrapper = spec.get("wrapper") + if wrapper is not None and wrapper.group is not None: + continue + if spec["kind"] == "new_group": + group = originals["c10d_new_group"]( + *_resolve_args_for_reload(spec["args"]), + **_resolve_kwargs_for_reload(spec["kwargs"]), + ) + elif spec["kind"] == "split_group": + group = originals["c10d_split_group"]( + *_resolve_args_for_reload(spec["args"]), + **_resolve_kwargs_for_reload(spec["kwargs"]), + ) + else: + continue + if wrapper is None: + continue + if group is None or group is dist.GroupMember.NON_GROUP_MEMBER: + continue + # New inner PGs get new private names/tags. FSDP2 caches the wrapper, + # while functional collectives resolve by name, so both registries must + # be moved back to the wrapper's original identity. + _restore_inner_group_name(wrapper, group) + wrapper.group = group + wrapper._copy_public_pg_attrs(group) + _register_wrapper_in_world(wrapper, group) + _restore_inner_tag(wrapper, group) + details.append( + { + "ranks": list(wrapper.group_info["ranks"]), + "backend": wrapper.group_info["backend"], + "group_name": wrapper.group_name, + } + ) + return details + + +def reloadable_process_group_status() -> dict[str, object]: + groups = ReloadableProcessGroup._GROUPS_BY_PID.get(os.getpid(), []) + return { + "total": len(groups), + "alive": sum(group.group is not None for group in groups), + "destroyed": sum(group.group is None for group in groups), + "specs": len(_SPECS_BY_PID.get(os.getpid(), [])), + } + + +def _wrap_group_if_needed( + group: dist.ProcessGroup | None, + ranks: list[int] | None, + backend_arg: object, +) -> dist.ProcessGroup | None: + if group is None or group is dist.GroupMember.NON_GROUP_MEMBER: + return group + if group is dist.group.WORLD: + return group + backend = str(backend_arg or _backend_for_reload(group)).lower() + if "gloo" in backend and "nccl" not in backend: + return group + if "nccl" not in backend and "cuda" not in backend: + return group + ranks = ranks or _ranks_from_world(group) + if len(ranks) <= 1: + return group + return ReloadableProcessGroup(group, ranks, _backend_for_reload(group)) + + +def _should_record_spec(group: dist.ProcessGroup | None, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: + # Nonmember calls must still be recorded to keep future reload calls + # collectively ordered across ranks. + backend = str(kwargs.get("backend", "")).lower() + if "gloo" in backend and "nccl" not in backend and "cuda" not in backend: + return False + if group is None: + return False + if group is dist.group.WORLD: + return False + if group is dist.GroupMember.NON_GROUP_MEMBER: + return True + try: + backend = str(dist.get_backend(group)).lower() + except Exception: + return False + return "nccl" in backend or "cuda" in backend + + +def _resolve_args_for_reload(args: tuple[Any, ...]) -> tuple[Any, ...]: + return tuple(_resolve_obj_for_reload(arg) for arg in args) + + +def _resolve_kwargs_for_reload(kwargs: dict[str, Any]) -> dict[str, Any]: + return {key: _resolve_obj_for_reload(value) for key, value in kwargs.items()} + + +def _resolve_obj_for_reload(value: Any) -> Any: + # Replay creation with current inner groups, not stale wrappers whose + # inners may have been destroyed. + if isinstance(value, ReloadableProcessGroup): + return value._inner() + if isinstance(value, list): + return [_resolve_obj_for_reload(item) for item in value] + if isinstance(value, tuple): + return tuple(_resolve_obj_for_reload(item) for item in value) + if isinstance(value, dict): + return {key: _resolve_obj_for_reload(item) for key, item in value.items()} + return value + + +def _extract_new_group_ranks(args: tuple[Any, ...], kwargs: dict[str, Any]) -> list[int] | None: + if args and args[0] is not None: + return list(args[0]) + if kwargs.get("ranks") is not None: + return list(kwargs["ranks"]) + if dist.is_initialized(): + return list(range(dist.get_world_size())) + return None + + +def _ranks_from_world(group: dist.ProcessGroup) -> list[int]: + import torch.distributed.distributed_c10d as c10d + + rank_map = c10d._world.pg_group_ranks.get(group) + if rank_map is None: + return list(range(group.size())) + return [global_rank for global_rank, _ in sorted(rank_map.items(), key=lambda item: item[1])] + + +def _backend_for_reload(group: dist.ProcessGroup) -> str: + backend = str(dist.get_backend(group)).lower() + if "nccl" in backend or "cuda" in backend: + return "nccl" + return backend + + +def _register_wrapper_in_world(wrapper: ReloadableProcessGroup, inner: dist.ProcessGroup) -> None: + import torch.distributed.distributed_c10d as c10d + + # c10d helpers query these Python maps directly. Registering the wrapper + # makes dist.* calls and DeviceMesh code treat it like the original PG. + world = c10d._world + if inner in world.pg_map: + world.pg_map[wrapper] = world.pg_map[inner] + if inner in world.pg_names: + world.pg_names[wrapper] = wrapper.group_name + if inner in world.pg_group_ranks: + world.pg_group_ranks[wrapper] = dict(world.pg_group_ranks[inner]) + if inner in world.pg_backend_config: + world.pg_backend_config[wrapper] = world.pg_backend_config[inner] + if inner in world.pg_to_tag: + tag = world.pg_to_tag[inner] + world.pg_to_tag[wrapper] = tag + world.tags_to_pg.setdefault(tag, []) + if wrapper not in world.tags_to_pg[tag]: + world.tags_to_pg[tag].append(wrapper) + + +def _restore_inner_group_name(wrapper: ReloadableProcessGroup, inner: dist.ProcessGroup) -> None: + import torch.distributed.distributed_c10d as c10d + + # The newly created inner receives a fresh c10d name. Restore the original + # name so functional collectives using that name keep resolving correctly. + old_name = process_group_name(inner) + stable_name = wrapper.group_name + if old_name == stable_name: + return + + world = c10d._world + if old_name != "unnamed": + try: + c10d._unregister_process_group(old_name) + except Exception: + pass + try: + inner._set_group_name(stable_name) + except RuntimeError: + pass + if inner in world.pg_names: + world.pg_names[inner] = stable_name + try: + c10d._register_process_group(stable_name, inner) + except Exception: + pass + + +def _restore_inner_tag(wrapper: ReloadableProcessGroup, inner: dist.ProcessGroup) -> None: + import torch.distributed.distributed_c10d as c10d + + # all_to_all_single_autograd and similar paths resolve by tag/name instead + # of by object identity, so the tag has to move to the fresh inner PG. + world = c10d._world + stable_tag = world.pg_to_tag.get(wrapper) + if stable_tag is None: + stable_tag = f"ptd:{wrapper.group_name}" + old_tag = world.pg_to_tag.get(inner) + if old_tag is not None and old_tag in world.tags_to_pg: + try: + world.tags_to_pg[old_tag].remove(inner) + except ValueError: + pass + world.pg_to_tag[inner] = stable_tag + world.tags_to_pg.setdefault(stable_tag, []) + if inner not in world.tags_to_pg[stable_tag]: + world.tags_to_pg[stable_tag].append(inner) + + +def process_group_name(group: dist.ProcessGroup) -> str: + try: + return str(group.group_name) + except RuntimeError: + return "unnamed"