From ab839bef26ce73687ff0c7e79f57eac31c04127f Mon Sep 17 00:00:00 2001 From: yszhli Date: Thu, 13 Aug 2026 14:30:14 +0800 Subject: [PATCH] Fix O(segments) per-comm memory probe that slows weight sync & training The pre-call memory check in _wrap_low_level_call runs available_memory() before every wrapped collective. available_memory() calls torch.cuda.memory_reserved()/ memory_allocated(), which build the full memory_stats() dict -- O(number of CUDA caching-allocator segments). Under PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True the segment count grows across steps, so this probe gets progressively more expensive; it also calls cudaMemGetInfo, which synchronizes the driver and thus serializes async collectives. Weight sync issues one dist.broadcast(async_op=True) per parameter and one dist.all_gather per expert bucket; training issues per-micro-batch TP all-gather (_allgather_base) and grad all-reduce. All of these went through the probe because: (a) they were registered with get_new_comm_function(fn) i.e. op_name=None, which forces check_memory=True; and (b) the low-level c10d method names (_allgather_base, allreduce, ...) dispatched by ReloadableProcessGroup._fwd were not in the skip set. Measured on Qwen3.5-35B-A3B (MoE + gated-delta-net) RL, 8xB300: - weight sync (perf/update_weights_time) grew 143s -> 363s -> 966s -> 1596s across steps; after the fix it is flat at ~7-20s; - each training micro-batch dropped from ~14s to ~1.5s (a single ~700-token fwd/bwd on a 3B-active MoE), i.e. the step went from ~57min to ~9min. py-spy showed the actor pinned in mem_get_info / memory_reserved via _wrap_low_level_call on those collectives. Free memory during both phases is tens of GB, so the clear_memory() the probe guards never fires -- it is pure overhead. Fix: 1. _wrap_low_level_call: use torch.cuda.mem_get_info() directly instead of available_memory(), dropping the O(segments) memory_stats() call. Any comm op still subject to the check no longer pays the expandable_segments blow-up. 2. Register the async collectives (all_reduce/all_gather/broadcast/reduce/all_to_all) with explicit op_names, and list both the dist.* spellings and the c10d method names (_allgather_base, allgather, allreduce, ...) in _COMM_MEMORY_CHECK_SKIP_OPS, so the hot weight-sync / training collectives skip the probe entirely. The check is retained for any op not listed. Co-Authored-By: Claude Opus 5 --- slime/utils/reloadable_process_group.py | 62 ++++++++++++++++++++----- 1 file changed, 50 insertions(+), 12 deletions(-) diff --git a/slime/utils/reloadable_process_group.py b/slime/utils/reloadable_process_group.py index 3e79281d4d..f616bd9740 100644 --- a/slime/utils/reloadable_process_group.py +++ b/slime/utils/reloadable_process_group.py @@ -10,7 +10,7 @@ from torch.distributed.distributed_c10d import PrefixStore, _get_default_group, _get_default_store from slime.utils.distributed_utils import get_gloo_group, init_gloo_group, set_gloo_group -from slime.utils.memory_utils import available_memory, clear_memory, print_memory +from slime.utils.memory_utils import clear_memory, print_memory logger = logging.getLogger(__name__) @@ -125,15 +125,47 @@ def _reload_default_process_group() -> None: ) +# Communication ops that skip the pre-call memory probe in _wrap_low_level_call. +# +# The probe (available_memory() -> cudaMemGetInfo + torch.cuda.memory_reserved()) +# runs before the collective. It is pure overhead on the hot path: +# - cudaMemGetInfo synchronizes the driver, which serializes async collectives +# (weight sync issues one dist.broadcast(async_op=True) per parameter); +# - the clear_memory() it guards only fires when free < 3GB, which never happens +# on real runs (weight sync / training have tens of GB free), so it is inert. +# +# Weight sync and training each issue one collective *per parameter / per +# micro-batch*, so the skip set must cover BOTH the dist.* monkey-patch spellings +# AND the low-level c10d method names dispatched by ReloadableProcessGroup._fwd +# (e.g. _allgather_base for sequence-parallel TP all-gather, allreduce for grad +# sync). See PR/issue: this made weight sync grow to >1500s and every training +# micro-batch ~2x slower on a large MoE + expandable_segments. _COMM_MEMORY_CHECK_SKIP_OPS = { - "all_gather_into_tensor", - "allgather_into_tensor_coalesced", + # point-to-point / barrier "barrier", + "isend", + "irecv", + # broadcast (dist.* and c10d spellings) + "broadcast", "broadcast_object_list", + # all-gather family + "all_gather", + "allgather", + "_allgather_base", + "all_gather_into_tensor", + "allgather_coalesced", + "allgather_into_tensor_coalesced", + # all-reduce family + "all_reduce", + "allreduce", + "allreduce_coalesced", + # reduce / reduce-scatter + "reduce", + "reduce_scatter", "reduce_scatter_tensor", + # all-to-all + "all_to_all", "all_to_all_single", - "isend", - "irecv", } @@ -223,15 +255,15 @@ def new_function(*args, **kwargs): dist.get_group_rank = get_new_query_function(dist.get_group_rank) dist.get_process_group_ranks = get_new_query_function(dist.get_process_group_ranks) - dist.all_reduce = get_new_comm_function(dist.all_reduce) - dist.all_gather = get_new_comm_function(dist.all_gather) + dist.all_reduce = get_new_comm_function(dist.all_reduce, "all_reduce") + dist.all_gather = get_new_comm_function(dist.all_gather, "all_gather") dist.all_gather_into_tensor = get_new_comm_function(dist.all_gather_into_tensor, "all_gather_into_tensor") dist.all_gather_object = get_new_comm_function(dist.all_gather_object) - dist.all_to_all = get_new_comm_function(dist.all_to_all) + dist.all_to_all = get_new_comm_function(dist.all_to_all, "all_to_all") dist.all_to_all_single = get_new_comm_function(dist.all_to_all_single, "all_to_all_single") - dist.broadcast = get_new_comm_function(dist.broadcast) + dist.broadcast = get_new_comm_function(dist.broadcast, "broadcast") dist.broadcast_object_list = get_new_comm_function(dist.broadcast_object_list, "broadcast_object_list") - dist.reduce = get_new_comm_function(dist.reduce) + dist.reduce = get_new_comm_function(dist.reduce, "reduce") dist.reduce_scatter = get_new_comm_function(dist.reduce_scatter) dist.reduce_scatter_tensor = get_new_comm_function(dist.reduce_scatter_tensor, "reduce_scatter_tensor") dist.scatter = get_new_comm_function(dist.scatter) @@ -475,8 +507,14 @@ def reload_process_groups(): def _wrap_low_level_call(check_memory=True): try: if check_memory: - mem_info = available_memory() - if mem_info["free_GB"] < 3: + # Use torch.cuda.mem_get_info() directly rather than available_memory(): + # the latter also calls torch.cuda.memory_reserved()/memory_allocated(), + # which build the full memory_stats() dict -- O(number of allocator + # segments). Under expandable_segments the segment count grows across + # steps, so this probe gets progressively more expensive. The free/clear + # decision only needs the driver-level free byte count. + free_bytes, _ = torch.cuda.mem_get_info() + if free_bytes < 3 * 1024**3: clear_memory() yield except Exception as e: