From 95261f85dd72ab579d1d7e12c59d3b7e7538e4c9 Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Wed, 12 Aug 2026 17:09:17 +0000 Subject: [PATCH] perf!: remove the zero-gradient-ratio metric get_zero_gradient_ratio launched a device-scalar transfer and a count_nonzero kernel per parameter plus two all-reduces and a sync every step, serializing the post-backward path on models with many tensors: on Qwen3-30B-A3B (8xH200, seq 8K) it cost 1.4 s of a 4.6 s step. The metric also read 1.0 on every validation run, so it was likely not measuring what it intended on FSDP2 DTensor gradients. Removes the metric, its trainer call sites, and the Prometheus gauge. Co-Authored-By: Claude Fable 5 --- src/prime_rl/trainer/rl/train.py | 5 --- src/prime_rl/trainer/sft/train.py | 4 -- src/prime_rl/trainer/utils.py | 60 +--------------------------- src/prime_rl/utils/metrics_server.py | 7 ---- 4 files changed, 1 insertion(+), 75 deletions(-) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 5946fcb3b7..e9df77f9d6 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -50,7 +50,6 @@ Tensors, export_benchmark_json, filter_rl_trainer_tensor_stats_for_wandb, - get_zero_gradient_ratio, get_ckpt_disk_metrics, setup_torch_distributed, print_benchmark, @@ -544,8 +543,6 @@ def train(config: TrainerConfig): if grad_norm.device.type == "cpu": grad_norm = grad_norm.to(torch.device("cuda")) - zero_grad_ratio = get_zero_gradient_ratio(model.parameters(), parallel_dims.dp_replicate) - # Update the model parameters optimizer.step() optimizer.zero_grad() @@ -650,7 +647,6 @@ def train(config: TrainerConfig): # Log optimizer metrics optim_metrics = { "optim/lr": current_lr, - "optim/zero_grad_ratio": zero_grad_ratio, "step": progress.step, } if grad_norm is not None: @@ -695,7 +691,6 @@ def train(config: TrainerConfig): mfu=mfu, entropy=tensor_stats.get("entropy/all/mean", 0.0), mismatch_kl=tensor_stats.get("mismatch_kl/all/mean", 0.0), - zero_grad_ratio=zero_grad_ratio, ) # Send heartbeat if configured diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 725659c12a..450aac55d6 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -40,7 +40,6 @@ GarbageCollection, MemoryProfiler, export_benchmark_json, - get_zero_gradient_ratio, get_ckpt_disk_metrics, print_sample, setup_torch_distributed, @@ -470,8 +469,6 @@ def run_validation(step: int) -> None: ) if grad_norm.device.type == "cpu": grad_norm = grad_norm.to(torch.device("cuda")) - zero_grad_ratio = get_zero_gradient_ratio(model.parameters(), parallel_dims.dp_replicate) - logger.debug("Optimizer step") optimizer.step() optimizer.zero_grad() @@ -578,7 +575,6 @@ def run_validation(step: int) -> None: # Log optimizer metrics optim_metrics = { "optim/lr": current_lr, - "optim/zero_grad_ratio": zero_grad_ratio, "step": progress.step, } if grad_norm is not None: diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 66cbe22f4a..fe0315d72e 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -4,7 +4,6 @@ import shutil import time from collections import defaultdict -from collections.abc import Iterable from datetime import timedelta from pathlib import Path from typing import Any @@ -16,8 +15,7 @@ from rich.console import Console from rich.table import Table from rich.text import Text -from torch import Tensor, nn -from torch.distributed.tensor import DTensor +from torch import Tensor from transformers.tokenization_utils import PreTrainedTokenizer from prime_rl.trainer.world import get_world @@ -55,62 +53,6 @@ def _collect(self, generation: int = 1): get_logger().info(f"[GC] collection took {time.monotonic() - begin:.2f}s") -def _to_local_tensor(tensor: Tensor | DTensor) -> Tensor: - if isinstance(tensor, DTensor): - return tensor.to_local() - return tensor - - -def count_zero_gradient_elements(parameters: Iterable[nn.Parameter]) -> tuple[Tensor, Tensor]: - """Count zero-gradient parameter elements on the local distributed shards. - - Parameters that require gradients but did not receive one in the current step - are counted as fully zero. This makes inactive MoE experts visible in the - metric instead of silently dropping them from the count. - """ - - device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") - num_zeros = torch.zeros((), dtype=torch.long, device=device) - num_tracked = torch.zeros((), dtype=torch.long, device=device) - - for param in parameters: - if not param.requires_grad: - continue - - local_param = _to_local_tensor(param.detach()) - if local_param.numel() == 0: - continue - - if local_param.device != num_zeros.device: - num_zeros = num_zeros.to(local_param.device) - num_tracked = num_tracked.to(local_param.device) - - local_numel = torch.tensor(local_param.numel(), dtype=torch.long, device=local_param.device) - num_tracked += local_numel - - if param.grad is None: - num_zeros += local_numel - continue - - local_grad = _to_local_tensor(param.grad.detach()) - if local_grad.numel() != local_param.numel(): - raise ValueError("Local gradient shape does not match the local parameter shape") - - num_zeros += local_numel - torch.count_nonzero(local_grad) - - return num_zeros, num_tracked - - -def get_zero_gradient_ratio(parameters: Iterable[nn.Parameter], dp_replicate: int = 1) -> float: - num_zero_grad, num_grad_elements = count_zero_gradient_elements(parameters) - dist.all_reduce(num_zero_grad, op=dist.ReduceOp.SUM) - dist.all_reduce(num_grad_elements, op=dist.ReduceOp.SUM) - if dp_replicate > 1: - num_zero_grad = torch.div(num_zero_grad, dp_replicate, rounding_mode="floor") - num_grad_elements = torch.div(num_grad_elements, dp_replicate, rounding_mode="floor") - return (num_zero_grad.float() / num_grad_elements.clamp_min(1).float()).item() - - def get_ckpt_disk_metrics(output_dir: Path) -> dict[str, float]: """ Disk usage metrics for the checkpoint directory (/checkpoints). diff --git a/src/prime_rl/utils/metrics_server.py b/src/prime_rl/utils/metrics_server.py index d3036bde3b..8807eddb56 100644 --- a/src/prime_rl/utils/metrics_server.py +++ b/src/prime_rl/utils/metrics_server.py @@ -103,11 +103,6 @@ def __init__(self, config: "MetricsServerConfig"): "trainer_mismatch_kl", "KL divergence between trainer and inference model", registry=self._registry ) self._kl_ent_ratio = Gauge("trainer_kl_ent_ratio", "Ratio of mismatch KL to entropy", registry=self._registry) - self._zero_grad_ratio = Gauge( - "trainer_zero_grad_ratio", - "Fraction of tracked parameter elements with zero gradient", - registry=self._registry, - ) def _make_handler(self) -> type[BaseHTTPRequestHandler]: """Create handler with /metrics and /health endpoints.""" @@ -164,7 +159,6 @@ def update( mfu: float = 0.0, entropy: float = 0.0, mismatch_kl: float = 0.0, - zero_grad_ratio: float = 0.0, ) -> None: """Update metrics after a training step.""" self._step.set(step) @@ -177,7 +171,6 @@ def update( self._mfu.set(mfu) self._entropy.set(entropy) self._mismatch_kl.set(mismatch_kl) - self._zero_grad_ratio.set(zero_grad_ratio) if entropy > 0: self._kl_ent_ratio.set(mismatch_kl / entropy) self._last_step_ts.set(time.time())