From 44b6f5800bdefcdd7ebae53ea97e9ff2c34d107c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 24 Aug 2026 18:36:59 +0000 Subject: [PATCH 1/2] Batch SignSGD updates with foreach ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-param loop launched thousands of tiny kernels per step; that long asynchronous tail overlapped the gradient frees and collectives that follow the step and hit a stream-ordering race under FSDP2 + CP (async CUDA illegal memory access, deterministic after step 1, on both the plain and optim_cpu_offload paths; full offload was immune because it runs no GPU optimizer kernels). CUDA_LAUNCH_BLOCKING=1 makes the crash vanish, confirming ordering rather than indexing. Foreach batching collapses the step into a few fused launches per (device, dtype) bucket — the same kernel pattern AdamW uses, which has never triggered the race. Update semantics are unchanged: decoupled decay p *= 1 - lr*wd (algebraically identical to the previous self-aliased add) then p -= lr * sign(g). --- src/prime_rl/trainer/sign_sgd.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/trainer/sign_sgd.py b/src/prime_rl/trainer/sign_sgd.py index 2fe1002e90..fcc22794f8 100644 --- a/src/prime_rl/trainer/sign_sgd.py +++ b/src/prime_rl/trainer/sign_sgd.py @@ -39,15 +39,21 @@ def step(self, closure: Callable = None): loss = closure() for group in self.param_groups: + # Batch per (device, dtype): foreach ops need homogeneous lists, + # and a few fused launches (vs thousands of per-param kernels) + # keep the step's kernel tail from overlapping the grad frees and + # collectives that follow it + buckets: dict[tuple, list] = {} for p in group["params"]: if p.grad is None: continue + buckets.setdefault((p.device, p.dtype), []).append(p) - sign_grad = torch.sign(p.grad) - + for params in buckets.values(): + grads = [p.grad for p in params] if group["weight_decay"] > 0.0: - p.add_(p, alpha=-group["lr"] * group["weight_decay"]) - - p.add_(sign_grad, alpha=-group["lr"]) + torch._foreach_mul_(params, 1 - group["lr"] * group["weight_decay"]) + signs = torch._foreach_sign(grads) + torch._foreach_add_(params, signs, alpha=-group["lr"]) return loss From 06b60c83ea96743b7cf9d70cb4c448b95d8ee25a Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 24 Aug 2026 22:07:43 +0000 Subject: [PATCH 2/2] Synchronize before the pre-broadcast empty_cache empty_cache returns cached blocks to the driver, so any kernel still in flight that references a cached block faults with an illegal memory access once the block is freed under it. The pre-broadcast empty_cache (added for FP8 gather headroom) ran right after optimizer.step(); optimizers with long kernel tails (SignSGD's per-param loop) deterministically crashed after step 1 under FSDP2 + CP, while AdamW's compact foreach step happened to finish in time. Drain all pending work before releasing blocks. --- src/prime_rl/trainer/rl/train.py | 13 +++++++++++-- src/prime_rl/trainer/sign_sgd.py | 22 ++++++++++++++-------- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index dab11298d6..fe63de8d47 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -84,7 +84,11 @@ def train(config: TrainerConfig): # Setup the monitors asyncio.run( monitors.setup( - wandb=config.monitors.wandb, file=config.monitors.file, output_dir=config.output_dir, run_config=config + producer="trainer", + wandb=config.monitors.wandb, + file=config.monitors.file, + output_dir=config.output_dir, + run_config=config, ) ) @@ -600,7 +604,12 @@ def train(config: TrainerConfig): broadcast_weights_start_time = time.perf_counter() # The per-layer gather + fp8 conversion peaks ~50 GiB above the # resident weights; release cached blocks (incl. offload-stream - # pools) so the broadcast gets the full headroom. + # pools) so the broadcast gets the full headroom. Drain all + # pending work first: empty_cache returns blocks to the driver, + # so a still-running kernel holding a cached block (e.g. the + # optimizer step's tail) faults with an illegal memory access + # once its block is freed under it. + torch.cuda.synchronize() torch.cuda.empty_cache() weight_sender.broadcast(model, step=progress.step) broadcast_weights_time = time.perf_counter() - broadcast_weights_start_time diff --git a/src/prime_rl/trainer/sign_sgd.py b/src/prime_rl/trainer/sign_sgd.py index fcc22794f8..d1d798411e 100644 --- a/src/prime_rl/trainer/sign_sgd.py +++ b/src/prime_rl/trainer/sign_sgd.py @@ -1,6 +1,7 @@ from typing import Callable import torch +from torch.distributed.tensor import DTensor from torch.optim import Optimizer @@ -39,10 +40,9 @@ def step(self, closure: Callable = None): loss = closure() for group in self.param_groups: - # Batch per (device, dtype): foreach ops need homogeneous lists, - # and a few fused launches (vs thousands of per-param kernels) - # keep the step's kernel tail from overlapping the grad frees and - # collectives that follow it + # Batch per (device, dtype): a few fused launches (vs thousands of + # per-param kernels) keep the step's kernel tail from overlapping + # the grad frees and collectives that follow it buckets: dict[tuple, list] = {} for p in group["params"]: if p.grad is None: @@ -50,10 +50,16 @@ def step(self, closure: Callable = None): buckets.setdefault((p.device, p.dtype), []).append(p) for params in buckets.values(): - grads = [p.grad for p in params] + # Update the local shards directly: every op below is + # pointwise and grads share the param's placement at step + # time, so the update commutes with sharding — and plain + # tensors keep foreach dispatch available + # (aten._foreach_sign has no DTensor sharding strategy) + local_params = [p.to_local() if isinstance(p, DTensor) else p for p in params] + local_grads = [p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad for p in params] if group["weight_decay"] > 0.0: - torch._foreach_mul_(params, 1 - group["lr"] * group["weight_decay"]) - signs = torch._foreach_sign(grads) - torch._foreach_add_(params, signs, alpha=-group["lr"]) + torch._foreach_mul_(local_params, 1 - group["lr"] * group["weight_decay"]) + signs = torch._foreach_sign(local_grads) + torch._foreach_add_(local_params, signs, alpha=-group["lr"]) return loss