fix: stream-ordering crash in SignSGD under FSDP2

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

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,15 +40,26 @@ def step(self, closure: Callable = None):
                 loss = closure()
 
         for group in self.param_groups:
+            # 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:
                     continue
+                buckets.setdefault((p.device, p.dtype), []).append(p)
 
-                sign_grad = torch.sign(p.grad)
-
+            for params in buckets.values():
+                # 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:
-                    p.add_(p, alpha=-group["lr"] * group["weight_decay"])
-
-                p.add_(sign_grad, 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