Fix cpu_offloaded_metric_module type-check for new_group return (#4457)

Summary:
Pull Request resolved: #4457

Removing the TorchComms `new_group` -> `split_group` delegation in
`distributed_c10d.py` (D113127841 / pytorch#189071) leaves `_new_group_with_tag`
as `new_group`'s only return path. `new_group` has no explicit return
annotation, so Pyrefly now infers `ProcessGroup | int | None` (the `int` is the
`GroupMember.NON_GROUP_MEMBER == -100` sentinel) instead of the narrower type it
used to see.

`cpu_offloaded_metric_module.py` assigns `dist.new_group(backend="gloo")`
straight into an attribute typed `Optional[ProcessGroup]`, which no longer
type-checks:

  cpu_offloaded_metric_module.py:303 `ProcessGroup | int | None` is not
  assignable to attribute `cpu_process_group` with type `ProcessGroup | None`

Both call sites pass `ranks=None`, so every rank is a member and the result is
always a real `ProcessGroup` (never `NON_GROUP_MEMBER`); narrow the value with
`cast`. This is an internal-only, monorepo-coexistence break: OSS pytorch CI does
not type-check torchrec against the new signature, which is why pytorch#189071
was green on GitHub. Kept as a separate diff so the DiffTrain import stays
byte-identical to GitHub.

Reviewed By: dboyda

Differential Revision: D113289735

fbshipit-source-id: 0f3dba9c0fd80fced2902dd0bc507f4704af4291

Author: tushar00jain

torchrec/metrics/cpu_offloaded_metric_module.py

@@ -15,7 +15,7 @@
 import threading
 import time
 import traceback
-from typing import Any, Dict, Mapping, Optional, Union
+from typing import Any, cast, Dict, Mapping, Optional, Union
 
 import torch
 from torch import distributed as dist
@@ -298,9 +298,12 @@ def __init__(
 
         # Created lazily on first compute so the module can be constructed without
         # an initialized process group (e.g. single-process model validation).
-        # pyrefly: ignore
+        # new_group returns a real ProcessGroup here (ranks=None => every rank is a
+        # member, never NON_GROUP_MEMBER); narrow away the int/None sentinels.
         self.cpu_process_group: Optional[dist.ProcessGroup] = (
-            dist.new_group(backend="gloo") if dist.is_initialized() else None
+            cast(Optional[dist.ProcessGroup], dist.new_group(backend="gloo"))
+            if dist.is_initialized()
+            else None
         )
         self.comms_module: CPUCommsRecMetricModule = CPUCommsRecMetricModule(
             *args,
@@ -755,7 +758,9 @@ def _process_metric_compute_job(
                 # Manual distributed sync (replaces TorchMetrics.metric.Metric.sync())
                 all_gather_start_ms = time.time()
                 if self.cpu_process_group is None:
-                    self.cpu_process_group = dist.new_group(backend="gloo")
+                    self.cpu_process_group = cast(
+                        dist.ProcessGroup, dist.new_group(backend="gloo")
+                    )
                 aggregated_states = self.comms_module.get_pre_compute_states(
                     self.cpu_process_group
                 )