From b3b003b437891137d18440f347590781d26d4f22 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:28:15 +0000 Subject: [PATCH] Critic FP32 masters: gather into one buffer per leaf; single-rank DP gathers on the CPU s19-compaction-g3-20261010a (CompactionRL, two islands, strict-avg): after round 0 the actor was still resident on the GPU (16k-token segments) when the critic woke for export_critic_tensors; full_masters built a FP32 copy per parameter on the GPU and then a second torch.cat copy (2 x 3.25 GiB) -> CUDA OOM with 1.46 GiB free, CrossChannelCommitError, island 0 exited. * full_masters: one flat FP32 buffer per optimizer leaf, each master a view of it (no concatenated copy); optional device for the buffers. * _critic_masters: when torch.distributed is up and every DP group has one rank (no collective needed), gather on the CPU: export/import no longer need a full FP32 critic copy on the GPU. Multi-rank groups keep the GPU buffer + all-reduce; a CPU buffer with a multi-rank group is refused. * _write_critic_masters: targets stay on their device (copy_ moves them); the stale-parameter check casts each master to the parameter's device. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_critic_fp32_masters.py | 37 ++++++++++++ yeto/rl/adapters/miles/state_plugin.py | 84 ++++++++++++++++++-------- 2 files changed, 96 insertions(+), 25 deletions(-) diff --git a/tests/test_rl_critic_fp32_masters.py b/tests/test_rl_critic_fp32_masters.py index b3d91acda..90c7b84a5 100644 --- a/tests/test_rl_critic_fp32_masters.py +++ b/tests/test_rl_critic_fp32_masters.py @@ -323,3 +323,40 @@ def sleep(self): out = state_plugin.critic_state_summary(Actor()) assert calls == ["wake", "sleep"] assert len(out["specs"]) == 2 and len(out["weights_sha256"]) == 64 + + +def test_single_rank_dp_gathers_on_cpu_in_one_buffer(monkeypatch): + """s19-compaction-g3-20261010a: the critic export ran out of GPU memory on a + second (torch.cat) copy of the FP32 masters while the actor was still + resident. One flat buffer per leaf now; a single-rank DP group (dist up) + gathers on the CPU with no collective, and export/import still round-trip.""" + import torch.distributed as dist + + torch.manual_seed(5) + body = torch.nn.Module() + body.embedding = torch.nn.Linear(5, 3, bias=False) + body.output_layer = torch.nn.Linear(3, 1, bias=False) + body = body.to(torch.bfloat16) + actor = _island(3, dp_rank=0, dp_size=1, body=body) + ref = sp._critic_masters(_island(3, body=body), sp._critic_parameters(actor)) + calls = [] + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda group=None: 1) + monkeypatch.setattr(sp, "_dp_all_reduce_sum", lambda leaf, flat: calls.append(flat)) + monkeypatch.setattr(sp, "_rank", lambda: 0) + params = sp._critic_parameters(actor) + keys = sorted(params) + masters = sp.full_masters(actor.optimizer, [params[k] for k in keys], device="cpu") + sharded = [m for k, m in zip(keys, masters) if not sp.has_complete_master(params[k])] + assert sharded and len({m.untyped_storage().data_ptr() for m in sharded}) == 1 # one buffer + assert not calls # single-rank group: no collective + e = sp._export_critic_tensors(actor) + assert not calls and all(t.device.type == "cpu" for t in e["tensors"].values()) + for k, m in ref.items(): + assert torch.equal(e["tensors"][k], m) + target = {k: v + 3e-5 for k, v in e["tensors"].items()} + h = cs.critic_weights_sha256(target) + assert sp._import_critic_tensors(actor, by_rank={0: {"tensors": target, "sha256": h}}) == { + "rank": 0, "weights_sha256": h} + diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index 9bbb1d1f3..bd570b621 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -226,53 +226,80 @@ def needs_gather(optimizer: Any, parameters: Sequence[Any]) -> bool: return any(not has_complete_master(p) for p in parameters) -def full_masters(optimizer: Any, parameters: Sequence[Any], *, all_reduce_sum=None, reduce=None) -> list[Any]: +def _single_rank_dp(leaf: Any) -> bool: + """True only when torch.distributed is up and the leaf's DP group has one rank + (no collective needed). Not initialised: False, the caller keeps the + ``_dp_all_reduce_sum`` path (a no-op there, patched in tests).""" + import torch.distributed as dist + + if not (dist.is_available() and dist.is_initialized()): + return False + try: + return int(dist.get_world_size(group=getattr(leaf, "data_parallel_group", None))) == 1 + except (RuntimeError, ValueError, TypeError, AttributeError): + return False # unknown group: keep the collective path + + +def full_masters(optimizer: Any, parameters: Sequence[Any], *, all_reduce_sum=None, reduce=None, + device: Any = None) -> list[Any]: """FP32 masters in each parameter's shape; DP-sharded ones are gathered. Collective over each owning leaf's DP group when any parameter is DP-sharded (every rank must call it with the same parameters in the same order). ``reduce(flat, leaf)`` / ``all_reduce_sum(flat)`` replace the all-reduce (tests). + + One flat FP32 buffer per leaf holds the gathered masters (each returned + master is a view of it); no second concatenated copy is made + (s19-compaction-g3-20261010a: the critic export ran out of GPU memory on + that copy). ``device``: where the gather buffers live (default: each + parameter's device). A leaf whose DP group has one rank needs no + collective, so a CPU buffer is allowed there; a CPU buffer with a + multi-rank group is refused (NCCL reduces GPU tensors only). """ import torch has_dist, owned = distributed_ranges(optimizer) - out: list[Any] = [] - pending: dict[int, tuple[Any, list[tuple[int, Any]]]] = {} + out: list[Any] = [None] * len(parameters) + pending: dict[int, tuple[Any, list[int]]] = {} for i, param in enumerate(parameters): if has_complete_master(param): - out.append(master_of(param)) + out[i] = master_of(param) continue if not has_dist: - out.append(master_of(param)) # raises: no master anywhere + out[i] = master_of(param) # raises: no master anywhere continue - full = torch.zeros(param.numel(), dtype=torch.float32, device=param.device) - if id(param) in owned: - _, _, start, end = owned[id(param)] - leaf = owned[id(param)][0] - shard = leaf._get_main_param_and_optimizer_states(param)["param"] - if shard.dtype != torch.float32 or shard.numel() != end - start: - raise StatePluginError("distributed-optimizer main shard does not match its range") - full[start:end] = shard.detach().reshape(-1) leaf = _leaf_for(optimizer, param, owned) - out.append(full) - pending.setdefault(id(leaf), (leaf, []))[1].append((i, full)) + pending.setdefault(id(leaf), (leaf, []))[1].append(i) # Leaf order is the same on every rank of a DP group, so the per-group # collectives are issued in the same order. order = {id(leaf): k for k, leaf in enumerate(_optimizer_leaves(optimizer))} if pending else {} for _, (leaf, items) in sorted(pending.items(), key=lambda kv: order.get(kv[0], 0)): - flat = torch.cat([f for _, f in items]) + hooked = reduce is not None or all_reduce_sum is not None + where = device if device is not None else parameters[items[0]].device + if (not hooked and torch.device(where).type == "cpu" + and parameters[items[0]].device.type != "cpu" and not _single_rank_dp(leaf)): + raise StatePluginError("CPU gather buffer with a multi-rank DP group") + flat = torch.zeros(sum(parameters[i].numel() for i in items), dtype=torch.float32, device=where) + offset = 0 + for i in items: + param = parameters[i] + n = param.numel() + full = flat[offset : offset + n] + if id(param) in owned: + _, _, start, end = owned[id(param)] + shard = owned[id(param)][0]._get_main_param_and_optimizer_states(param)["param"] + if shard.dtype != torch.float32 or shard.numel() != end - start: + raise StatePluginError("distributed-optimizer main shard does not match its range") + full[start:end].copy_(shard.detach().reshape(-1)) + out[i] = full.view(param.shape) + offset += n if reduce is not None: reduce(flat, leaf) elif all_reduce_sum is not None: all_reduce_sum(flat) - else: + elif not _single_rank_dp(leaf): _dp_all_reduce_sum(leaf, flat) - offset = 0 - for i, full in items: - n = full.numel() - out[i] = flat[offset : offset + n].view(parameters[i].shape) - offset += n return out @@ -1288,7 +1315,13 @@ def _critic_masters(actor: Any, params: Mapping[str, Any]) -> dict[str, Any]: """ keys = sorted(params) - masters = full_masters(getattr(actor, "optimizer", None), [params[k] for k in keys]) + optimizer = getattr(actor, "optimizer", None) + leaves = [leaf for leaf in _optimizer_leaves(optimizer) if _is_distributed(leaf)] if optimizer is not None else [] + # Single-rank DP groups need no collective: gather on the CPU, so the + # export/import does not need a full FP32 critic copy on the GPU + # (s19-compaction-g3-20261010a: OOM with the actor still resident). + device = "cpu" if leaves and all(_single_rank_dp(leaf) for leaf in leaves) else None + masters = full_masters(optimizer, [params[k] for k in keys], device=device) return dict(zip(keys, masters, strict=True)) @@ -1300,7 +1333,8 @@ def _write_critic_masters(actor: Any, params: Mapping[str, Any], targets: Mappin keys = sorted(params) plist = [params[k] for k in keys] - tlist = [targets[k].to(device=params[k].device, dtype=torch.float32) for k in keys] + # Targets stay where they are (CPU from the syncer); copy_ moves each one. + tlist = [targets[k].to(dtype=torch.float32) for k in keys] write_masters(getattr(actor, "optimizer", None), plist, tlist) for param, target in zip(plist, tlist, strict=True): if param.dtype != torch.float32: @@ -1361,7 +1395,7 @@ def _import_critic_tensors(actor: Any, *, by_rank: dict[int, dict[str, Any]]) -> masters = _critic_masters(actor, params) written = _masters_cpu(masters) stale = [k for k, p in params.items() if p.dtype != torch.float32 - and not torch.equal(p.detach(), masters[k].to(p.dtype).view(p.shape))] + and not torch.equal(p.detach(), masters[k].to(device=p.device, dtype=p.dtype).view(p.shape))] except StatePluginError as error: return {"refused": f"critic masters: {error}", "rank": rank} got = critic_weights_sha256(written)