Fix O(#segments) per-comm memory probe that slows weight sync & training under expandable_segments

[yszhli]

Fixes #2268

Problem

slime/utils/reloadable_process_group.py::_wrap_low_level_call runs a memory probe before every wrapped collective:

def _wrap_low_level_call(check_memory=True):
    if check_memory:
        mem_info = available_memory()          # calls torch.cuda.memory_reserved()/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(number of CUDA caching-allocator segments). Under PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True the segment count grows across steps, so the probe gets progressively more expensive. It also calls cudaMemGetInfo, which synchronizes the driver and serializes async collectives.

This is on the hot path because weight sync issues one dist.broadcast(param, async_op=True) per parameter (update_weights_from_distributed) plus one dist.all_gather per expert bucket, and training issues per-micro-batch TP all-gather (_allgather_base) and grad all-reduce (allreduce). All of them hit the probe because:

1. they are registered with get_new_comm_function(fn) i.e. op_name=None, which forces check_memory=True (True if op_name is None else _should_check_memory_for_comm(op_name)); and
2. 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.

Impact (measured)

Qwen3.5-35B-A3B (MoE + gated-delta-net) async GRPO, 8×B300, expandable_segments:True:

	before	after
perf/update_weights_time	143 → 363 → 966 → 1596s (grows every step)	flat ~7–20s
training micro-batch (one ~700-token fwd/bwd, 3B active)	~14s	~1.5s
step time	~57 min	~9 min

py-spy on the training actor showed it pinned in mem_get_info / memory_reserved via _wrap_low_level_call, first on update_weights broadcast and (after only fixing broadcast) on _allgather_base. 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.

Fix

1. _wrap_low_level_call: use torch.cuda.mem_get_info() directly instead of available_memory(), dropping 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 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.

Both changes are safe: none of these collectives allocate significant new memory (they operate on existing tensors), so none benefit from a pre-call clear_memory(), and the error-path memory dump in the except branch is unchanged.