perf!: remove the zero-gradient-ratio metric

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

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:

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 (<output_dir>/checkpoints).

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())