Per-comm memory probe (available_memory/memory_stats) is O(#allocator segments) and slows weight sync & training under expandable_segments

Summary

slime/utils/reloadable_process_group.py::_wrap_low_level_call runs a memory probe before every wrapped collective, and that probe is O(number of CUDA caching-allocator segments). Under PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True the segment count grows across steps, so weight sync and training get progressively slower. On a large MoE RL run this made perf/update_weights_time grow 143s → 363s → 966s → 1596s across steps, and made every training micro-batch ~2x slower.

Where

# slime/utils/reloadable_process_group.py
def _wrap_low_level_call(check_memory=True):
    if check_memory:
        mem_info = available_memory()          # <-- calls memory_reserved()/memory_allocated()
        if mem_info["free_GB"] < 3:
            clear_memory()
    yield

available_memory() calls torch.cuda.memory_reserved()/memory_allocated(), which build the full memory_stats() dict — O(#segments). It also calls cudaMemGetInfo, which synchronizes the driver and thus serializes async collectives.

This runs on the hot path because:

1. Weight sync issues one dist.broadcast(param, async_op=True) per parameter (update_weights_from_distributed), plus one dist.all_gather per expert bucket.
2. Training issues per-micro-batch TP all-gather (_allgather_base) and grad all-reduce (allreduce).
3. These collectives hit the probe because (a) they are registered with get_new_comm_function(fn) i.e. op_name=None, which forces check_memory=True in get_new_comm_function (True if op_name is None else _should_check_memory_for_comm(op_name)); and (b) the low-level c10d method names (_allgather_base, allreduce, …) dispatched by ReloadableProcessGroup._fwd are not in _COMM_MEMORY_CHECK_SKIP_OPS.

Free memory during both phases is tens of GB, so the clear_memory() the probe guards never fires — the probe is pure overhead here.

Evidence

py-spy dump on the training actor (6/6 samples, weight-sync phase):

memory_stats (torch/cuda/memory.py)
memory_reserved
available_memory (slime/utils/memory_utils.py)
_wrap_low_level_call (slime/utils/reloadable_process_group.py)
update_weights_from_distributed  # dist.broadcast(param, async_op=True)

Training phase, the residual moved to _allgather_base (same _wrap_low_level_call → mem_get_info).

The geometric growth (143→1596s) is the tell: a constant path cost can't grow like that; it's memory_stats() walking an ever-growing segment list under expandable_segments.

Reproduction

• Any MoE model, distributed weight update, PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True.
• Observed on Qwen3.5-35B-A3B (MoE + gated-delta-net) async GRPO, 8×B300, slime + Megatron-LM + SGLang.
• Watch perf/update_weights_time climb every step; py-spy the actor during update_weights.

Impact (measured)

Proposed fix

1. In _wrap_low_level_call, use torch.cuda.mem_get_info() directly instead of available_memory() — drop the O(#segments) memory_stats() call; the free/clear decision only needs the driver-level free byte count.
2. Register the async collectives (all_reduce/all_gather/broadcast/reduce/all_to_all) with explicit op_names, and add both the dist.* spellings and the c10d method names (_allgather_base, allgather, allreduce, …) to _COMM_MEMORY_CHECK_SKIP_OPS.

Reference implementation (against a fork of this repo): https://github.com/yszhli/slime/pull/1 — happy to open a PR here if the approach looks right.

Start in slime/utils/reloadable_process_group.py at _wrap_low_level_call, get_new_comm_function, _fwd, and _COMM_MEMORY_CHECK_SKIP_OPS, then inspect slime/utils/memory_utils.py for available_memory(). The change is complete when the probe avoids the O(#segments) memory_stats path and the listed async collectives skip unnecessary checks; validate with the expandable_segments reproduction and the reported update/training timings.